import pyvisa

if __name__ == "__main__":
    import sys
    if __package__ is None:
        # make relative imports work as described here: https://peps.python.org/pep-0366/#proposed-change
        __package__ = "photoreflectance"
        from os import path
        filepath = path.realpath(path.abspath(__file__))
        sys.path.insert(0, 'C:\\Users\Administrator\Desktop\Software\Python\Python\github')

from time import sleep

import numpy as np
import scipy as scp
from Bentham import Bentham, DummyBentham
from devices.Xenon import Xenon
from devices.Shutter import ShutterProbe, DummyShutter

from .update_funcs import Monitor
from .measurement_device.impl.sr830 import SR830
from .measurement_device.impl.model7260 import Model7260
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
    handlers=[
        # logging.FileHandler(log_path),
        logging.StreamHandler()
    ]
)
log = logging.getLogger(__name__)

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,
    }
    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_offset_laser_only(lockin: SR830, shutter: ShutterProbe, wait_time_s):
    """
    Set the R offset from the signal when only the laser is on.
    This signal should be stray laser light and laser induced PL
    :return: Offset as percentage of the full scale R
    """
    log.info("Setting offset when the lamp is off.")
    shutter.close_()
    sleep(wait_time_s + 10)
    lockin.run("AOFF 3")  # auto offset R
    R_offset_fs = float(lockin.query("OEXP? 3").split(",")[0])  # returns R offset and expand
    return R_offset_fs


def _measure_both_sim(monochromator: Bentham, lockin: SR830, shutter: ShutterProbe, wl_range=(400, 750, 25), aux_DC="Aux In 4", offset_with_laser_only=True, monitor=None):
    data = {}
    lockin_params = {
        "time_constant_s": 10,
        # "time_constant_s": 100e-3,
        "sensitivity_volt": 50e-6,
        "filter_slope": 12,
        "sync_filter": 1,
        "reserve": "Normal",
    }
    measurement_params = {
        "measurement_time_s": 30,
        "sample_rate_Hz": 512,
    }

    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 + 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"),
    ])
    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"]
    if offset_with_laser_only:
        mon.set_fig_title(f"Measuring baseline with lamp off")
        R_offset_volt = set_offset_laser_only(lockin, shutter, wait_time_s) * full_scale_voltage
        data["R_offset_volt_before"] = R_offset_volt
        print(f"R_offset_volt_before {R_offset_volt}")
    data["reference_freq_Hz_before"] = lockin.get_frequency()

    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.drive(wl)
        mon.set_fig_title(f"Waiting for signal to stabilize")
        # wait the wait time
        sleep(wait_time_s)
        mon.set_fig_title(f"Measuring...")
        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")
        arr = run_lockin_cmd(lambda: lockin.buffer_get_data(CH1=True, CH2=True))
        data[wl] = {}
        data[wl]["raw"] = arr
        # calculate means
        means = np.mean(arr, axis=1)
        errs = np.std(arr, axis=1)
        dR = means[0]
        R = means[1]
        sdR = errs[1]
        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)
    # if it fails, we still want the data returned
    try:
        if offset_with_laser_only:
            mon.set_fig_title(f"Measuring baseline with lamp off")
            R_offset_volt = set_offset_laser_only(lockin, shutter, wait_time_s) * full_scale_voltage
            data["R_offset_volt_after"] = R_offset_volt
            print(f"R_offset_volt_after {R_offset_volt}")
        data["reference_freq_Hz_after"] = lockin.get_frequency()
    except Exception as e:
        print(e)
    mon.set_fig_title("Photoreflectance")
    mon.set_ax_title("")
    return data, mon

def _measure_both(monochromator: Bentham, lockin: SR830, shutter: ShutterProbe, wl_range=(400, 750, 25), AC=True, DC=True, monitor=None):
    mon = monitor if monitor is not None else Monitor(r"$\lambda$ [nm]", [
        # dict(ax=0, ylabel="Wavelength [nm]", color="red"),
        # dict(ax=1, ylabel="Ref", color="blue", lim=(0, 5)),
        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=3, ylabel="R", color="blue"),
        dict(ax=5, ylabel=r"$\theta$", color="orange", lim=(-180, 180)),
        dict(ax=6, ylabel=r"$\sigma_\theta$", color="orange")
    ])
    shutter.open_()
    data_raw = []
    data_wl = {}
    # TODO these are only printed, not set!
    time_constant = 30e-3
    filter_slope = 24
    sensitivity = 1.0

    SYNC = 1
    sample_rate_AC = 64
    sample_rate_DC = 512
    n_bins_AC = 3 * sample_rate_AC   # x sec messen mit <sample_rate> werte pro sekunde
    n_bins_DC = 10 * sample_rate_DC
    timeout_s = 60
    timeout_interval = 0.5
    # lockin.run("SENS 17")  # 1 mV
    # lockin.run("SENS 20")   # 10 mV
    # lockin.run("SENS 21")  # 20 mV
    lockin.run("SENS 26")  # 1 V
    # 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
    # lockin.run("OFLT 5")   # 3 ms
    lockin.run("OFLT 7")   # 30 ms
    # lockin.run("OFLT 8")   # 100 ms
    # lockin.run("OFLT 10")   # 1 s
    lockin.run("RMOD 2")   # low noise (small reserve)
    # lockin.run("OFSL 0")   # 6dB/Oct
    lockin.run("OFSL 3")   # 24dB/Oct
    lockin.run(f"SYNC {SYNC}")   # sync filter

    print(f"Time estimate {40 * (wl_range[1]-wl_range[0])/(wl_range[2]*60)} minutes")
    if AC:
        input("Plug the detector into lock-in port 'A/I' (front panel) and press enter > ")
        input("Make sure the laser is turned on and press enter > ")
        mon.set_fig_title("Turn on laser and plug detector into A")
        for i_wl, wl in enumerate(range(*wl_range)):
            mon.set_ax_title(f"$\\lambda = {wl}$ nm")
            lockin.buffer_setup(CH1="R", CH2="Theta", length=n_bins_AC, sample_rate=sample_rate_AC)
            mon.set_fig_title(f"Setting wavelength to {wl} nm")
            monochromator.drive(wl)
            mon.set_fig_title(f"Waiting for signal to stabilize")
            # wait time depends on filter and time constant, for 24dB/Oct and Sync on the minimum is ~12 time constants
            sleep(1.0)
            mon.set_fig_title(f"Measuring...")
            lockin.buffer_start_fill()
            t = timeout_s
            while t > 0:
                t -= timeout_interval
                sleep(timeout_interval)
                if lockin.buffer_is_done():
                    break
            if t < 0: raise RuntimeError("Timed out waiting for buffer measurement to finish")
            arr = lockin.buffer_get_data(CH1=True, CH2=True)
            data_raw.append([wl, arr])
            # calculate means, for theta use circular mean
            dR = np.mean(arr[0,:])
            sdR = np.std(arr[0,:])
            theta = scp.stats.circmean(arr[1,:], low=-180, high=180)
            stheta = scp.stats.circstd(arr[1,:], low=-180, high=180)
            data_wl[wl] = {"dR": dR, "Theta": theta, "sdR": sdR, "sTheta": stheta}
            # wl - dR, sdR, R, sR, dR/R, Theta
            mon.update(wl, dR, sdR, None, None, None, theta, stheta)
    if DC:
        mon.set_ax_title("")
        mon.set_fig_title("Turn off laser and plug detector into Aux 1")
        input("Turn off the laser and press enter > ")
        input("Plug the detector into lock-in port 'Aux In 1' (rear panel) and press enter > ")
        for i_wl, wl in enumerate(range(*wl_range)):
            mon.set_ax_title(f"$\\lambda = {wl}$ nm")
            lockin.buffer_setup(CH1="Aux In 1", CH2="Theta", length=n_bins_DC, sample_rate=sample_rate_DC)
            mon.set_fig_title(f"Setting wavelength to {wl} nm")
            monochromator.drive(wl)
            sleep(0.5)
            mon.set_fig_title(f"Measuring...")
            lockin.buffer_start_fill()
            t = timeout_s
            while t > 0:
                t -= timeout_interval
                sleep(timeout_interval)
                if lockin.buffer_is_done():
                    break
            if t < 0: raise RuntimeError("Timed out waiting for buffer measurement to finish")
            arr = lockin.buffer_get_data(CH1=True, CH2=False)
            if AC:
                data_raw[i_wl].append(arr)
            else:
                data_raw.append([wl, arr])
            means = np.mean(arr, axis=1)
            errs = np.std(arr, axis=1)
            if not wl in data_wl: data_wl[wl] = {}
            data_wl[wl] |= {"R": means[0], "sR": errs[0]}
            # wl - dR, sdR, R, sR, dR/R, Theta
            if AC:
                dR_R = data_wl[wl]["dR"] / data_wl[wl]["R"]
                mon.override(wl, None, None, data_wl[wl]["R"], data_wl[wl]["sR"], dR_R, None, None)
            else:
                mon.update(wl, None, None, data_wl[wl]["R"], data_wl[wl]["sR"], None, None, None)
    mon.set_fig_title(f"Time constant = {time_constant} s\nFilter slope = {filter_slope} dB/oct\nSync Filter = {SYNC}\nSensitivity = {sensitivity*1e3} mV")
    mon.set_ax_title("")
    return data_wl, data_raw, mon


def _measure(monochromator: Bentham, lamp: Xenon, lockin: SR830, shutter: ShutterProbe, wl_range=(400, 750, 25)):
    data = []
    mon = Monitor(r"$\lambda$ [nm]", [
        # dict(ax=0, ylabel="Wavelength [nm]", color="red"),
        # dict(ax=1, ylabel="Ref", color="blue", lim=(0, 5)),
        dict(ax=0, ylabel=r"$\Delta R$", color="green"),
        dict(ax=1, ylabel=r"$\sigma_{\Delta R}$", color="green"),
        dict(ax=2, ylabel="Phase", color="orange", lim=(-180, 180))
        # dict(ax=3, ylabel="R", color="blue"),
    ])
    N_bins = 100
    dt = 0.01
    i = 0
    shutter.open_()
    if isinstance(lockin, SR830):
        lockin.run("SENS 17")  # 1 mV/nA
        lockin.run("OFLT 5")
        for wl in range(*wl_range):
            arr = np.empty((N_bins, 2))
            monochromator.drive(wl)

            # lockin.auto_gain()
            for j in range(N_bins):
                dR, theta = lockin.snap(what="3,4")
                arr[j,:] = (dR, theta)
                i += 1
                # sleep(dt)
            means = np.mean(arr, axis=0)
            errs = np.std(arr, axis=0)
            mon.update(wl, means[0], errs[0], means[1])
            data.append((wl, arr))
    elif isinstance(lockin, Model7260):
        # lockin.run("SEN 21")  # 10 mV/nA
        lockin.run("SEN 24")  # 100 mV/nA
        for wl in range(*wl_range):
            monochromator.drive(wl)

            lockin.buffer_setup(MAG=1, PHASE=1, ADC1=1, ADC2=1, interval_ms=5, length=1000)
            lockin.buffer_start_fill()
            timeout_s = 60
            timeout_interval = 0.5
            while timeout_s > 0:
                timeout_s -= timeout_interval
                sleep(timeout_interval)
                if lockin.buffer_is_done():
                    break
            arr = lockin.buffer_get_data()
            length = arr.shape[0]
            # for j in range(length):
            #     # wl, ref, dR, R, theta
            #     mon.update(i*length+j, wl, arr[j][3], arr[j][0], arr[j][2], arr[j][1])
            mon.update_array(range(i * length, (i+1)*length), [wl for _ in range(length)], arr[:,3], arr[:,0], arr[:,2], arr[:,1])
            data.append((wl, arr))
            i += 1

    shutter.close_()
    return data, mon

lockin = None
lamp = None
mcm = None
shutter = None

def measure(wl_range=(400, 500, 2)):
    return _measure(mcm, lamp, lockin, shutter, wl_range=wl_range)

def measure_both(**kwargs):
    return _measure_both(mcm, lockin, shutter, **kwargs)

def measure_both_sim(**kwargs):
    return _measure_both_sim(mcm, lockin, shutter, **kwargs)

if __name__ == "__main__":
    mcm = Bentham()
    shutter = ShutterProbe()
    # mcm.park()
    lamp = Xenon()
    lockin = SR830.connect_device(SR830.enumerate_devices()[0])
    # lockin = Model7260.connect_device(Model7260.enumerate_devices()[0])
    # mcm = DummyBentham()
    # shutter = DummyShutter()