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, time as now

import numpy as np
import matplotlib.pyplot as plt
from Bentham import Bentham
from prsctrl.devices.lamp.impl.xenon import Xenon
from prsctrl.devices.shutter import ShutterProbe

from .update_funcs import Monitor
from prsctrl.devices.lock_in.impl.sr830 import SR830
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=550, aux_DC="Aux In 4", 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
    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
    input("Make sure the laser is turned off and press enter > ")
    mon = monitor if monitor is not None else Monitor(r"$t$ [s]", [
        dict(ax=0, ylabel=r"$R [V]$", color="green"),
    ])
    data["lock-in-params"] = lockin_params
    data["measurement-params"] = measurement_params

    N_runs = 60

    shutter.open_()
    monochromator.drive(wl)
    sleep(10)

    t0 = now()
    j = 0
    data['R'] = []
    data['t'] = []
    for i in range(N_runs):
        mon.set_ax_title(f"{i+1}/{N_runs}")
        run_lockin_cmd(lambda: lockin.buffer_setup(CH1="R", CH2=aux_DC, length=n_bins_AC, sample_rate=sample_rate_AC))
        ti = now() - t0
        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")
        R = run_lockin_cmd(lambda: lockin.buffer_get_data(CH1=False, CH2=True))[0]
        data['R'].append(R)
        tdata = np.arange(n_bins_AC) * 1/sample_rate_AC * measurement_time_s + ti
        data['t'].append(tdata)
        mon.update_array(tdata, R)
    mon.set_fig_title("Background")
    mon.set_ax_title("")
    return data, mon

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

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


def _load_process():
    import os; import pickle
    with open(os.path.expanduser("~/Desktop/PR/2025-05-07_Xenon_1.pkl"), "rb") as file:
        data = pickle.load(file)
    process_data(data['R'], data['t'])


def process_data(data_R: list[np.ndarray], data_t: list[np.ndarray]) -> np.ndarray:
    assert len(data_R) == len(data_t)
    n_points_2 = 512 * 5

    n_points_averaged = data_R[0].shape[0]-2*n_points_2
    l = len(data_R) * n_points_averaged
    # Rs = np.empty(l, dtype=float)
    # ts = np.empty(l, dtype=float)
    Rs = []
    ts = []

    fig, ax = plt.subplots()
    ax.set_xlabel("$t$ [s]")
    ax.set_ylabel(r"$R [V]$")
    for i in range(len(data_R)):
        Rs_i = moving_average(data_R[i], n_points_2)
        new_t_data = np.arange(data_R[i].shape[0], dtype=int) /512 + (i * 30)
        ts_i = new_t_data[n_points_2:-n_points_2]
        # ts_i = data_t[i][n_points_2:-n_points_2]
        # Rs[i*n_points_averaged:(i+1)*n_points_averaged] = Rs_i
        # ts[i*n_points_averaged:(i+1)*n_points_averaged] = ts_i
        Rs.append(Rs_i)
        ts.append(ts_i)
        ax.plot(ts_i, Rs_i, color="blue")
    plt.show()
    return fig

def moving_average(a: np.ndarray, n_points_2=2):
    """

    :param a:
    :param n_points_2: Number of points to take to the left and right of the averaged point, leading to 2*n_points+1 per value
    :return:
    """
    l = a.shape[0]
    lret = l-2*n_points_2
    ret = np.empty(lret)
    for i in range(lret):
        around_i = n_points_2 + i
        begin = around_i - n_points_2
        end = around_i + n_points_2 + 1
        ret[i] = np.sum(a[begin:end]) / (2*n_points_2 + 1)
        # print(i, begin, end, ret[i])
    return ret


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()