import time
import argparse
import fn_gen.errors as fg_err
from common import (
    close_output,
    get_preamplified,
    get_postamplified,
    AMPLIFICATION,
    N_CYCLES,
    V_STEPS,
    F_STEPS,
    FREQ_START,
    FREQ_STOP,
    STEP_DURATION,
)
from fn_gen import DG2052
import numpy as np


def calculate_n_samples(freq: float, ts: float, n_cycles: int) -> int:
    """
    Calculates the number of samples based on the frequency of the signal, the sampling time and the number of cycles within one step.

    Parameters
    ---------
    freq: int
        The frequency of the signal
    ts: float
        The sampling time
    n_cycles: int
        The number of cycles within a step

    Returns
    ------
    n_samples: int
        The number of samples (can be multiplied by the sampling time (ts) to get the total duration)
    """
    return int(round(n_cycles / (ts * freq)))


def discrete_sweep(
    v_min: float,
    v_max: float,
    v_steps: int,
    adaptive: bool,
    step_duration: int,
    freq_start: int,
    freq_stop: int,
    f_steps: int,
):
    ##################### PROGRAM START ###########
    freqs = np.logspace(np.log10(freq_start), np.log10(freq_stop), f_steps)
    v_min = get_preamplified(AMPLIFICATION, v_min)
    v_max = get_preamplified(AMPLIFICATION, v_max)
    volts = np.linspace(v_min, v_max, v_steps)
    print(freqs)
    fg = DG2052("TCPIP::192.168.1.11::INSTR")
    channel = 2
    try:
        print(fg.whoami())
        print("")
        print(f"Output{channel} Impedance: {fg.get_output_impedance(channel)} Ohm")
        print(f"Output{channel} Load: {fg.get_output_load(channel)} Ohm")
        print(f"Output{channel} Voltage Limits: {fg.get_output_volt_limits(channel)} V")
        for v in volts:
            for freq in freqs:
                sampling_rate = 10 * freq
                n_samples = calculate_n_samples(freq, 1 / sampling_rate, N_CYCLES)
                if adaptive:
                    step_duration = n_samples * (1 / sampling_rate)
                print(f"V: {get_postamplified(AMPLIFICATION, v)} V")
                print(f"Freq: {freq} Hz")
                print(f"Duration: {step_duration} s")
                print(f"N_Samples: {n_samples} samples")
                fg.set_output(channel, False)
                fg.set_sine_wave(channel, freq, v, 0, 0)
                fg.set_output(channel, True)
                print(
                    f"Output{channel}: {fg.get_output_signal(channel)} | {fg.get_output_state(channel)}"
                )
                time.sleep(step_duration)
                # fg.set_output(channel, False)
        print(f"Output{channel} State: {fg.get_output_state(channel)}")
    except fg_err.ValueOutOfBoundsError as err:
        print(err)
    except fg_err.UndefinedValueError as err:
        print(err)
    except KeyboardInterrupt:
        close_output(fg, channel)
    finally:
        close_output(fg, channel)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="This program is for testing the DG2052 function genrator library. It does a discrete sweep with the supplied parameters."
    )
    parser.add_argument(
        "--vmin", type=float, required=True, help="The minimum voltage supplied"
    )
    parser.add_argument(
        "--vmax", type=float, required=True, help="The maximum voltage supplied"
    )
    parser.add_argument(
        "--v-steps", type=int, default=V_STEPS, help="The number of voltage steps"
    )
    parser.add_argument(
        "--freq-start", type=int, default=FREQ_START, help="The starting frequency"
    )
    parser.add_argument(
        "--freq-stop", type=int, default=FREQ_STOP, help="The stop frequency"
    )
    parser.add_argument(
        "--f-steps", type=int, default=F_STEPS, help="The number of steps"
    )
    parser.add_argument(
        "--step-duration",
        type=int,
        default=STEP_DURATION,
        help="The duration of each step",
    )
    parser.add_argument(
        "--adaptive-step",
        action="store_true",
        help="Adapts the step duration to the frequency of the step",
    )
    args = parser.parse_args()
    discrete_sweep(
        v_min=args.vmin,
        v_max=args.vmax,
        v_steps=args.v_steps,
        adaptive=args.adaptive_step,
        step_duration=args.step_duration,
        freq_start=args.freq_start,
        freq_stop=args.freq_stop,
        f_steps=args.f_steps,
    )