import time
import argparse

import fn_gen.errors as fg_err
from common import close_output, abs, get_preamplified
from fn_gen import DG2052
from fn_gen.enums import SweepSignalType, SweepSpacing, SweepTriggerSource


def sweep_over_signal(
    signal: str,
    v_min: float,
    v_max: float,
    delay: int,
    duration: int,
    freq_start: int,
    freq_stop: int,
    phase: int,
    spacing: str,
):
    fg = DG2052("TCPIP::192.168.1.11::INSTR")
    v_min = get_preamplified(13.0, v_min)
    v_max = get_preamplified(13.0, v_max)
    v = (v_max - v_min)/2
    channel = 2
    # signal_type = SweepSignalType.SINE
    match signal:
        case "sine":
            fg.set_sine_wave(channel, freq_start, v, 0, 0)
        case "square":
            fg.set_square_wave(channel, freq_start, v, 0, 0)
        case "ramp":
            fg.set_ramp(channel, freq_start, v, 0, 0)
    spacing_type = SweepSpacing.LOG
    match spacing:
        case "lin":
            spacing_type = SweepSpacing.LIN
        case "log":
            spacing_type = SweepSpacing.LOG
    try:
        print(fg.whoami())
        print(f"\nOutput{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")
        print(
            f"Output{channel}: {fg.get_output_signal(channel)} | {fg.get_output_state(channel)}"
        )
        fg.set_output(channel, True)
        time.sleep(delay)            
        # fg.trigger_sweep(channel)
        t0 = time.time()
        t1 = time.time()
        freq = freq_start
        freq_delta = (freq_stop - freq_start) / (duration*10_000)
        while (t1 - t0) < duration:
            print(f"Current Frequency: {fg.get_output_signal(channel)}")
            if spacing_type == SweepSpacing.LIN:
                freq = freq + freq_delta
                match signal:
                    case "sine":
                        fg.set_sine_wave(channel, freq, v, 0, 0)
                    case "square":
                        fg.set_square_wave(channel, freq, v, 0, 0)
                    case "ramp":
                        fg.set_ramp(channel, freq, v, 0, 0)
            time.sleep(0.0001)
            t1 = time.time()
        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:
        pass
    finally:
        close_output(fg, channel)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="This program is for testing the DG2052 function genrator library. It sweeps over a signal with the supplied parameters."
    )
    parser.add_argument(
        "-s",
        "--signal",
        type=str,
        choices=["sine", "square", "ramp"],
        default="sine",
        help='The type of signal being sweeped',
    )
    parser.add_argument(
        "--vmin", type=float, default=0, help="The minimum voltage supplied"
    )
    parser.add_argument(
        "--vmax", type=float, default=1, help="The maximum voltage supplied"
    )
    parser.add_argument(
        "-d",
        "--delay",
        type=int,
        default=0,
        help="The buffer time before the sweep starts",
    )
    parser.add_argument(
        "--duration", type=int, default=5 * 60, help="The duration of the sweep"
    )
    parser.add_argument(
        "--freq-start", type=int, default=10, help="The start frequency of the sweep"
    )
    parser.add_argument(
        "--freq-stop", type=int, default=1000, help="The stop frequency of the sweep"
    )
    parser.add_argument(
        "-p",
        "--phase",
        type=int,
        default=0,
        help="The phase shift of the signal generated (must be between 0 and 360)",
    )
    parser.add_argument(
        "--spacing",
        type=str,
        choices=["lin", "log"],
        default="log",
        help='The spacing of the sweep',
    )
    args = parser.parse_args()
    if (
        args.phase not in range(0, 360)
        or args.spacing not in ["lin", "log"]
        or args.signal not in ["sine", "square", "ramp"]
    ):
        parser.print_help()
        exit(1)
    sweep_over_signal(args.signal, args.vmin, args.vmax, args.delay, args.duration, args.freq_start, args.freq_stop, args.phase, args.spacing)