# -*- coding: utf-8 -*-
"""
Created on Fri Jan 24 15:18:31 2025

@author: Matthias Quintern
"""
from .devices.lock_in.base import LockInAmp
from .devices.shutter.base import Shutter
from .devices.monochromator import Monochromator
from .utility.prsdata import PrsData

import time
import datetime
from queue import Queue
import numpy as np

import logging
log = logging.getLogger(__name__)

def get_wavelengths_values(wl_range: tuple | list):
    """
    :param wl_range:
        if tuple, return list(range(*wl_range))
        if list, return copy of wl_range
    """
    if isinstance(wl_range, tuple):
        wavelengths = list(range(*wl_range))
    elif isinstance(wl_range, list) or isinstance(wl_range, np.ndarray):
        wavelengths = wl_range.copy()
    else:
        raise ValueError(f"Invalid type for 'wavelengths_nm': {type(wl_range)}")
    return wavelengths


def measure_spectrum(
        monochromator: Monochromator,
        lockin: LockInAmp,
        shutter: Shutter,
        data: PrsData,
        measurement_params:dict,
        aux_DC="Aux In 4",
        command_queue: None | Queue = None,
        data_queue: None | Queue = None,
        add_measurement_info_to_metadata=True
):

    import pyvisa
    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:
                # TODO: retry if status bit is set
                lockin.try_recover_from_communication_error(e)
            com_success -= 1
        raise e

    default_measurement_params = {
        "measurement_time_s": 30,
        "sample_rate_Hz": 512,
        "wait_time_s": 1,
        "wavelengths_nm": (390, 720, 1),
    }
    measurement_params = default_measurement_params | measurement_params
    wait_time_s = measurement_params["wait_time_s"]
    sample_rate_Hz = measurement_params["sample_rate_Hz"]
    measurement_time_s = measurement_params["measurement_time_s"]
    n_bins = sample_rate_Hz * measurement_time_s
    wavelengths = get_wavelengths_values(measurement_params["wavelengths_nm"])
    print(wavelengths)

    timeout_s = 3 * measurement_time_s
    timeout_interval = 0.5

    get_time = lambda: datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
    if add_measurement_info_to_metadata:
        data.metadata["device_lock-in"] = str(lockin)
        data.metadata["device_monochromator"] = str(monochromator)
        data.metadata["measurement_time_start"] = get_time()
    # write metadata to disk
    data.write_metadata()
    print(wait_time_s)
    data.metadata["messages"] = []

    try:
        shutter.open()
        for i_wl, wl in enumerate(wavelengths):
            log.info(f"Measuring at lambda={wl} nm")
            run_lockin_cmd(lambda: lockin.buffer_setup(CH1="R", CH2=aux_DC, length=n_bins, sample_rate=sample_rate_Hz))

            data_queue.put(("set_wavelength", wl))
            monochromator.set_wavelength_nm(wl)
            data_queue.put(("wait_stable", ))
            # wait the wait time
            time.sleep(wait_time_s)
            overload = run_lockin_cmd(lambda: lockin.check_overloads())
            if overload:
                msg = f"Overload of {overload} at {wl} nm"
                log.warning(msg)
                data.metadata["messages"].append(msg)
            theta = []
            measure_phase = lambda: theta.append(run_lockin_cmd(lambda: lockin.read_value("theta")))
            data_queue.put(("measuring", ))
            measure_phase()
            run_lockin_cmd(lambda: lockin.buffer_start_fill())
            # check if its done
            t = timeout_s
            while t > 0:
                t -= timeout_interval
                time.sleep(timeout_interval)
                measure_phase()
                if run_lockin_cmd(lambda: lockin.buffer_is_done()):
                    break
            if t < 0: raise RuntimeError("Timed out waiting for buffer measurement to finish")
            # done
            dR_raw, R_raw = run_lockin_cmd(lambda: lockin.buffer_get_data(CH1=True, CH2=True))
            data[wl] = {}
            data[wl]["dR_raw"] = dR_raw * 2 * np.sqrt(2)  # convert RMS to Peak-Peak
            data[wl]["R_raw"] = R_raw
            data[wl]["theta_raw"] = np.array(theta)
            spec_data = data.get_spectrum_data([wl])
            data.write_partial_file(wl)
            data_queue.put(("data", spec_data))

            # if a pipe was given, check for messages
            if command_queue is not None and command_queue.qsize() > 0:
                recv = command_queue.get(block=False)
                if recv == "stop":
                    log.info(f"Received 'stop', stopping measurement")
                    break
                elif type(recv) == tuple and recv[0] == "metadata":
                    log.info(f"Received 'metadata', updating metadata")
                    data.metadata |= recv[1]
                    data.write_metadata()
                else:
                    log.error(f"Received invalid message: '{recv}'")
    except KeyboardInterrupt:
        log.info("Keyboard interrupt, stopping measurement")
    except Exception as e:
        log.critical(f"Unexpected error, stopping measurement. Error: {e}")
        if command_queue is not None:
            command_queue.put(("exception", e))
        raise e

    if add_measurement_info_to_metadata:
        data.metadata["measurement_time_stop"] = get_time()
    # Write again after having updated the stop time
    data.write_metadata()
    data.write_full_file()


def set_offsets_laser_only(
        lockin: LockInAmp,
        shutter: Shutter,
        wait_time_s,
        R=True,
        phase=True,
        data_queue: None | Queue = None,
    ):
    """
    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()
    if data_queue:
        data_queue.put(("wait_stable", ))
    time.sleep(wait_time_s + 10)
    # TODO: generalize for other lock-ins
    if data_queue:
        data_queue.put(("set_offsets", ))
    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
    if data_queue: data_queue.put(("offsets", R_offset_fs, phase_offset_deg))
    return R_offset_fs, phase_offset_deg