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 pyvisa
import numpy as np
import scipy as scp

from Bentham import Bentham
from prsctrl.devices.lamp.impl.xenon import Xenon
from prsctrl.devices.shutter import Shutter
from prsctrl.devices.lamp.impl.xenon import Xenon

from .update_funcs import Monitor
from prsctrl.devices.lock_in.impl.sr830 import SR830
from prsctrl.devices.lock_in.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: Shutter, 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: Shutter, 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")
        # ToDo Phase messen
        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[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)
    # 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: Shutter, 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: Shutter, 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 = module_shutter.connect_device(module_shutter.TYPENAME_DAQ, "TAS Lamp Shutter")
    # mcm.park()
    lamp = Xenon()
    lockin = SR830.connect_device(SR830.enumerate_devices()[0])
    # lockin = Model7260.connect_device(Model7260.enumerate_devices()[0])
    # mcm = DummyBentham()
    # shutter = DummyShutter()