import time
import argparse

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

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)
    channel = 2
    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")
        curr_freq = freq_start
        df = 1
        dt = 0.001
        match signal:
            case "sine":
                fg.set_sine_wave(channel, curr_freq, v_min)
            case "square":
                fg.set_square_wave(channel, curr_freq, v_min)
            case "ramp":
                fg.set_ramp(channel, curr_freq, v_min)
        while curr_freq < freq_stop:
            match spacing:
                case "lin":
                    curr_freq += df
                    fg.set_frequency(channel, curr_freq)
                case "log":
                    pass # TODO: setup log sweep
            time.sleep(dt)
            close_output(fg, 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)