import numpy as np
import scipy as scp
from time import sleep
import pyvisa

import logging
log = logging.getLogger(__name__)


from prsctrl.devices.lamp import Lamp
from prsctrl.devices.shutter import Shutter
from prsctrl.devices.monochromator import Monochromator
from .update_funcs import Monitor
from prsctrl.devices.lock_in import Lock_In_Amp
from prsctrl.devices.lock_in.impl.sr830 import SR830

def set_measurement_params(lockin: SR830, p: dict={}, **kwargs):
    params = p | kwargs
    key_to_setter = {
        "time_constant_s": lockin.set_time_constant_s,
        "filter_slope": lockin.set_filter_slope,
        "sync_filter": lockin.set_sync_filter,
        "reserve": lockin.set_reserve,
        "sensitivity_volt": lockin.set_sensitivity_volt,
        "frequency_Hz": lockin.set_frequency_Hz,
        "reference": lockin.set_reference,
        "reference_trigger": lockin.set_reference_trigger,
    }
    for k, v in params.items():
        if k not in key_to_setter.keys():
            raise KeyError(f"Invalid parameter {k}")
        key_to_setter[k](v)


def set_offsets_laser_only(lockin: SR830, shutter: Shutter, wait_time_s, R=True, phase=True):
    """
    Set the R offset from the signal when only the laser is on.
    This signal should be stray laser light and laser induced PL
    :param phase: If True, use the Auto-Phase function to offset the phase
    :param R: If True, use the Auto-Offset function to offset R
    :return: Offset as percentage of the full scale R, Phase offset in degrees
    """
    log.info("Setting offset when the lamp is off.")
    shutter.close()
    sleep(wait_time_s + 10)
    lockin.run("AOFF 3")  # auto offset R
    # R must come before phase, because after auto-phase the signal needs to stabilize again
    if R:
        R_offset_fs = float(lockin.query("OEXP? 3").split(",")[0])  # returns R offset and expand
    if phase:
        lockin.run("APHS")
    phase_offset_deg = float(lockin.query("PHAS? 3"))  # returns R offset and expand
    return R_offset_fs, phase_offset_deg


def _measure_both_sim(monochromator: Monochromator, lockin: SR830, shutter: Shutter, wl_range=(400, 750, 25), aux_DC="Aux In 4", offset_with_laser_only=True, monitor=None, laser_power_mW=None):
    data = {}
    lockin_params = {
        "time_constant_s": 10,
        # "time_constant_s": 300e-3,
        "sensitivity_volt": 500e-6,
        "filter_slope": 12,
        "sync_filter": 1,
        "reserve": "Normal",
        "reference": "Internal",
        "reference_trigger": "Falling Edge",
        "frequency_Hz": 173,
    }
    measurement_params = {
        "measurement_time_s": 30,
        "sample_rate_Hz": 512,
    }
    if laser_power_mW:
        measurement_params["laser_power_mW"] = laser_power_mW

    set_measurement_params(lockin, lockin_params)

    measurement_time_s = measurement_params["measurement_time_s"]
    sample_rate_AC = measurement_params["sample_rate_Hz"]
    n_bins_AC = measurement_time_s * sample_rate_AC   # x sec messen mit <sample_rate> werte pro sekunde
    timeout_s = 60
    timeout_interval = 0.5
    # trigger on the falling edge, since the light comes through when the ref signal is low
    # could of course also trigger on rising and apply 180° shift
    lockin.run("RSLP 2")
    # since we dont expect changes in our signal, we can use larger time constants and aggressive filter slope
    # for better signal to noise
    wait_time_s = lockin.get_wait_time_s()

    def run_lockin_cmd(cmd, n_try=2):
        com_success = n_try
        e = None
        while com_success > 0:
            try:
                return cmd()
            except pyvisa.VisaIOError as e:
                lockin.try_recover_from_communication_error(e)
            com_success -= 1
        raise e

    # 5s for setting buffer,
    # 5s for get values and plot
    print(f"Time estimate {(measurement_time_s + wait_time_s + 10 + 5 + 5)/60 * ((wl_range[1]-wl_range[0])/wl_range[2])} minutes")
    input("Make sure the laser is turned on and press enter > ")
    mon = monitor if monitor is not None else Monitor(r"$\lambda$ [nm]", [
        dict(ax=0, ylabel=r"$\Delta R$", color="green"),
        dict(ax=1, ylabel=r"$\sigma_{\Delta R}$", color="green"),
        dict(ax=2, ylabel=r"$R$", color="blue"),
        dict(ax=3, ylabel=r"$\sigma_R$", color="blue"),
        dict(ax=4, ylabel=r"$\Delta R/R$", color="red"),
        dict(ax=5, ylabel=r"$\sigma_{\Delta R/R}$", color="red"),
        dict(ax=6, ylabel=r"$\theta$", color="pink"),
    ])
    mon.set_fig_title(f"Turn on laser and plug detector into A and {aux_DC} ")

    data["lock-in-params"] = lockin_params
    data["measurement-params"] = measurement_params
    full_scale_voltage = lockin_params["sensitivity_volt"]
    def set_offsets(name):
        shutter.close()
        mon.set_fig_title(f"Measuring baseline with lamp off")
        R_offset_fs, phase_offset_deg = set_offsets_laser_only(lockin, shutter, wait_time_s)
        R_offset_volt = R_offset_fs * full_scale_voltage
        data[f"R_offset_volt_{name}"] = R_offset_volt
        data[f"phase_offset_deg_{name}"] = phase_offset_deg
        print(f"R_offset_volt_{name} {R_offset_volt}")
        print(f"phase_offset_deg_{name}: {phase_offset_deg}")

    if offset_with_laser_only: set_offsets("before")
    data["reference_freq_Hz_before"] = lockin.get_frequency_Hz()

    data["info"] = []

    shutter.open()
    for i_wl, wl in enumerate(range(*wl_range)):
        mon.set_ax_title(f"$\\lambda = {wl}$ nm")
        run_lockin_cmd(lambda: lockin.buffer_setup(CH1="R", CH2=aux_DC, length=n_bins_AC, sample_rate=sample_rate_AC))
        mon.set_fig_title(f"Setting wavelength to {wl} nm")
        monochromator.set_wavelength_nm(wl)
        mon.set_fig_title(f"Waiting for signal to stabilize")
        # wait the wait time
        sleep(wait_time_s + 10)
        overload = run_lockin_cmd(lambda: lockin.check_overloads())
        if overload:
            msg = f"Overload of {overload} at {wl} nm"
            log.warning(msg)
            data["info"].append(msg)
        theta = []
        mon.set_fig_title(f"Measuring...")
        theta.append(run_lockin_cmd(lambda: float(lockin.query("OUTP? 4"))))
        run_lockin_cmd(lambda: lockin.buffer_start_fill())
        t = timeout_s
        while t > 0:
            t -= timeout_interval
            sleep(timeout_interval)
            if run_lockin_cmd(lambda: lockin.buffer_is_done()):
                break
        if t < 0: raise RuntimeError("Timed out waiting for buffer measurement to finish")
        theta.append(run_lockin_cmd(lambda: float(lockin.query("OUTP? 4"))))
        arr = run_lockin_cmd(lambda: lockin.buffer_get_data(CH1=True, CH2=True))
        data[wl] = {}
        data[wl]["raw"] = arr
        data[wl]["theta"] = theta
        # calculate means
        means = np.mean(arr, axis=1)
        errs = np.std(arr, axis=1)
        dR = means[0]
        R = means[1]
        sdR = errs[0]
        sR = errs[1]
        data[wl] |= {"dR": dR, "sdR": sdR, "R": R, "sR": sR}
        dR_R = dR / R
        sdR_R = np.sqrt((sdR / R) + (dR * sR/R**2))
        data[wl] |= {"dR_R": dR_R, "sdR_R": sdR_R}
        mon.update(wl, dR, sdR, R, sR, dR_R, sdR_R, theta[0])
    # if it fails, we still want the data returned
    try:
        if offset_with_laser_only: set_offsets("before")
        data["reference_freq_Hz_after"] = lockin.get_frequency_Hz()
    except Exception as e:
        print(e)
    mon.set_fig_title("Photoreflectance")
    mon.set_ax_title("")
    return data, mon