import pytest

from lab_control.test.mock.mock_lab import MockLab
from lab_control.frequency_response import FrequencyResponseMeasurement
from lab_control.measurement import getLinearRange
from lab_control.function_generator import FunctionGenerator

@pytest.fixture
def mockLab():
    return MockLab()

@pytest.fixture
def uut(mockLab):
    return FrequencyResponseMeasurement()

def prepareRampResponse(uut, mockLab):
    # Expect a ramp response from 0.5 to 1.5 * input amplitude
    minScale = 0.5
    maxScale = 1.5

    def testFunction(f: float) -> float:
        assert f >= uut.minFrequency and f <= uut.maxFrequency
        frequencyPu = (f - uut.minFrequency) / (uut.maxFrequency - uut.minFrequency)
        return minScale + maxScale * frequencyPu

    mockLab.connectChannels(uut.functionGeneratorChannel, uut.oscilloscopeChannel)
    mockLab.setTestFunction(uut.oscilloscopeChannel, testFunction)

    return [(f, testFunction(f) * uut.inputAmplitude) for f in getLinearRange(uut.minFrequency, uut.maxFrequency, uut.steps)]

def test_frequencyResponseDefaults(uut):
    assert uut.minFrequency == 20e0
    assert uut.maxFrequency == 16e3
    assert uut.steps == 50
    assert uut.inputAmplitude == 0.1
    assert uut.functionGeneratorChannel == 1
    assert uut.oscilloscopeChannel == 1

def test_frequencyResponseRamp(uut, mockLab):
    uut.minFrequency = 100.0
    uut.maxFrequency = 200.0
    uut.steps = 11
    uut.functionGeneratorChannel = 1
    uut.oscilloscopeChannel = 1
    uut.inputAmplitude = 2.0

    fg = mockLab.getFunctionGeneratorChannel(uut.functionGeneratorChannel)
    osc = mockLab.getOscilloscopeChannel(uut.oscilloscopeChannel)

    expectedData = prepareRampResponse(uut, mockLab)

    assert fg.amplitude is None
    assert fg.function is None
    assert not fg.on
    assert osc.voltsPerDiv is None
    assert not uut.measurementDone
    assert uut.data is None
    
    uut.measure(mockLab, mockLab)

    assert fg.amplitude == uut.inputAmplitude
    assert fg.function == FunctionGenerator.SINE
    assert fg.on
    assert osc.voltsPerDiv == uut.inputAmplitude / (mockLab.getDivisionsDisplayed() - 2)
    assert uut.measurementDone
    assert uut.data == expectedData

def test_dumpToCSV(uut, mockLab, tmp_path):
    uut.minFrequency = 10.0
    uut.maxFrequency = 30.0
    uut.steps = 5
    csv = tmp_path / "data.csv"

    assert not csv.exists()
    uut.dumpToCSV(csv)
    assert not csv.exists()

    expectedData = prepareRampResponse(uut, mockLab)
    uut.measure(mockLab, mockLab)

    uut.dumpToCSV(csv)
    savedData = csv.read_text().split("\n")

    assert savedData[0] == "Frequency (Hz),Amplitude (V)"
    for i, l in enumerate(savedData[1:]):
        assert l == f"{expectedData[i][0]},{expectedData[i][1]}"