"""
run this before using this library:
ipython -i prctrl_interactive.py
"""
from toolz import frequencies

version = "0.1"

import numpy as np
import matplotlib.pyplot as plt
import time

from datetime import datetime as dtime
from os import path, makedirs
import threading as mt
import multiprocessing as mp

import argparse

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__ = "prsctrl"
        from os import path
        filepath = path.realpath(path.abspath(__file__))
        sys.path.insert(0, path.dirname(path.dirname(filepath)))


# import device modules
from .devices import shutter as mod_shutter
from .devices import lock_in as mod_lock_in
from .devices import lamp as mod_lamp
from .devices import monochromator as mod_monochromator
# import base classes
from .devices.lock_in import LockInAmp
from .devices.shutter import Shutter
from .devices.lamp import Lamp
from .devices.monochromator import Monochromator

from .devices.lock_in.impl.sr830 import set_measurement_params
from .measurement import measure_spectrum as _measure_spectrum, set_offsets_laser_only, get_wavelengths_values
from .utility.prsdata import PrsData, plot_spectrum
from .utility.config_file import ConfigFile
from .utility.device_select import select_device_interactive, connect_device_from_config_or_interactive
from .update_funcs import Monitor

import logging
log = logging.getLogger(__name__)

# CONFIGURATION
_runtime_vars = {
    "last-measurement": ""
}
# defaults, these may be overridden by a config file
settings = {
    "datadir":      path.expanduser("~/Desktop/PR/data"),
    "name":         "interactive-test",
    "flush_after":  3000,
    "use_buffer":   False,
}

cfilename: str = "photoreflectance.yaml"
config_path: str = ""
config_file: ConfigFile = ConfigFile("")

test = False

# DEVICES
# global variable for the instrument/client returned by pyvisa/bleak
lockin: LockInAmp | None = None
shutter: Shutter|None = None
lamp: Lamp|None = None
mcm: Monochromator|None = None

t0 = 0
data = None
md = None

from .test_measurement import  _measure_both_sim
def measure_both_sim(**kwargs):
    return _measure_both_sim(mcm, lockin, shutter, **kwargs)

default_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,
}
default_measurement_params = {
    "measurement_time_s": 30,
    "sample_rate_Hz": 512,
    "wait_time_s": 0,
    "wavelengths_nm": (550, 555, 1),
}


def get_time_estimate(lockin_params: dict, measurement_params: dict, offset_with_laser_only=True, extra_wait_time_s=10) -> float:
    global lockin

    lockin_params = default_lockin_params | lockin_params
    measurement_params = default_measurement_params | {"wait_time_s":  lockin.get_wait_time_s(lockin_params["time_constant_s"], lockin_params["filter_slope"]) + extra_wait_time_s} | measurement_params
    wls = get_wavelengths_values(measurement_params["wavelengths_nm"])

    wait_time_s = measurement_params["wait_time_s"]
    measurement_time_s = measurement_params["measurement_time_s"]

    t = 0
    if offset_with_laser_only:
        t += 2 * (wait_time_s + 10)   # 2 * (at beginning and end), 10 for auto functions
    t += len(wls) * (5 + wait_time_s + measurement_time_s + 10)     # 5 for setting wavelength, 10 for data transfers
    return t

def measure_spectrum(metadata:dict={},
        lockin_params={},
        measurement_params={},
        aux_DC="Aux In 4",
        offset_with_laser_only=True,
        extra_wait_time_s=10,
        name:str|None=None,
        dirname=None,
        save_spectrum=True
    ):
    global _runtime_vars
    global data, md
    global lockin, shutter, lamp, mcm
    _runtime_vars["last_measurement"] = dtime.now().isoformat()

    lockin_params = default_lockin_params | lockin_params
    set_measurement_params(lockin, lockin_params)

    # must come after lockin settings, since its using get_wait_time_s
    measurement_params = default_measurement_params | {"wait_time_s":  lockin.get_wait_time_s() + extra_wait_time_s} | measurement_params

    # get the frequency, if not given
    refkey = "frequency_Hz"
    if not refkey in measurement_params:
        measurement_params[refkey] = lockin.get_frequency_Hz()

    # 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")

    # set metadata
    metadata["lock-in_settings"] = lockin_params
    metadata["measurement_parameters"] = measurement_params
    if name is None:
        name = settings["name"]
    metadata["name"] = name
    print(f"Starting measurement with: Use <C-c> to stop. Save the data using 'data.save_csv()' afterwards.")
    plt.ion()
    plt_monitor = 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="aqua"),
        dict(ax=7, ylabel=r"$\sigma_{\theta}$", color="aqua"),
    ])
    plt_monitor.set_fig_title(f"Turn on laser and plug detector into A and {aux_DC} ")
    data = PrsData(data={}, metadata=metadata, write_data_path=settings["datadir"], write_data_name=settings["name"], write_dirname=dirname)

    # measure/set offset
    full_scale_voltage = lockin_params["sensitivity_volt"]
    def set_offsets(name):
        shutter.close()
        plt_monitor.set_fig_title(f"Measuring baseline with lamp off")
        R_offset_fs, phase_offset_deg = set_offsets_laser_only(lockin, shutter, measurement_params["wait_time_s"])
        R_offset_volt = R_offset_fs * full_scale_voltage
        data.metadata[f"R_offset_volt_{name}"] = R_offset_volt
        data.metadata[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_collector.clear()
    data_queue = mp.Queue()
    command_queue = mp.Queue()
    # Argument order must match the definition
    proc_measure = mt.Thread(target=_measure_spectrum, args=(
        mcm,
        lockin,
        shutter,
        data,
        measurement_params,
        aux_DC,
        command_queue,
        data_queue,
        True,  # add metadata
    ))
    proc_measure.start()
    try:
        while proc_measure.is_alive():
            while not data_queue.empty():
                # print(data_queue.qsize(), "\n\n")
                msg = data_queue.get(block=False)
                if msg[0] == "data":
                    # data is as returned by PrsData.get_spectrum_data()
                    plt_monitor.update(*msg[1][0,:])
                elif msg[0] == "set_wavelength":
                    plt_monitor.set_ax_title("Setting wavelength...")
                    plt_monitor.set_fig_title(f"$\\lambda = {msg[1]}$ nm")
                elif msg[0] == "wait_stable":
                    plt_monitor.set_ax_title("Waiting until signal is stable...")
                elif msg[0] == "measuring":
                    plt_monitor.set_ax_title("Measuring...")
                else:
                    log.error(f"Invalid tuple received from measurement thread: {msg[0]}")
            time.sleep(0.01)

    except KeyboardInterrupt:
        pass
    command_queue.put("stop")
    proc_measure.join()
    plt_monitor.set_fig_title("")
    plt_monitor.set_ax_title("")
    try:
        if offset_with_laser_only: set_offsets("before")
        data["reference_freq_Hz_after"] = lockin.get_frequency_Hz()
    except Exception as e:
        print(e)
    print("Measurement stopped" + " "*50)
    if save_spectrum:
        plt.ioff()
        fig = plot_spectrum(data, title=name, what=["dR_R", "theta"])
        fig_path = path.join(data.path, data.dirname + ".pdf")
        fig.savefig(fig_path)
        return fig


def sweep_ref():
    wavelenghts = [500, 550, 650, 660, 670, 680]
    frequencies = list(range(27, 500, 5))
    frequencies = [111, 444]
    lockin_params = {
        "time_constant_s": 10,
    }
    measurement_params = {
        "wavelengths_nm": wavelenghts,
    }
    time_est = len(frequencies) * get_time_estimate(lockin_params=lockin_params, measurement_params=measurement_params, offset_with_laser_only=True, extra_wait_time_s=10)
    print(f"Estimated time: {time_est}")
    return
    for f in frequencies:
        dirname = f"2025-05-07_f-scan_f={f}_Hz"
        lockin_params["frequency"] = f
        measure_spectrum(lockin_params=lockin_params, measurement_params=measurement_params, dirname=dirname, name="Frequency scan $f = {f}$ Hz")
        plt.close('all')

# DATA
def data_load(dirname:str) -> tuple[np.ndarray, dict]:
    """
    Load data in directory <dirname> in the data directory as numpy array.
    Sets the `data` and `md` variables

    Parameters
    ----------
    dirname : str
        Absolute path to the directory containing the measurement or directory name in the data directory.
    """
    global data, md
    if path.isabs(dirname):
        dirpath = dirname
    else:
        dirpath = path.join(settings["datadir"], dirname)
    data, md = PrsData.load_data_from_dir(dirpath, verbose=True)
    
# SETTINGS
def set(setting, value):
    global settings, config_path
    if setting in settings:
        if type(value) != type(settings[setting]):
            print(f"set: setting '{setting}' currently holds a value of type '{type(settings[setting])}'")
            return
    settings[setting] = value
    config_file.set(setting, value)

def name(s:str):
    global settings
    settings["name"] = s

def save_settings():
    global settings
    for k, v in settings.items():
        config_file.set(k, v)
    config_file.save()

def load_settings():
    global settings, config_path
    settings = config_file.get_values()
    settings["datadir"] = path.expanduser(settings["datadir"])  # replace ~

def help(topic=None):
    if topic == None:
        print("""
Functions:
    monitor         - take measurements with live monitoring in a matplotlib window
    data_load       - load data from a directory
    data_plot       - plot a data array
    Run 'help(function)' to see more information on a function

Available topics:
    imports
    settings
Run 'help("topic")' to see more information on a topic""")

    elif topic in [settings, "settings"]:
        print(f"""Settings:
    name: str       - name of the measurement, determines filename
    led: str        - name/model of the LED that is being used
    datadir: str    - output directory for the csv files
    interval: int   - interval (inverse frequency) of the measurements, in seconds
    beep: bool      - whether the device should beep or not


Functions:
    name("<name>")         - short for set("name", "<name>")
    set("setting", value)  - set a setting to a value
    save_settings()        - store the settings in the config file
    load_settings()        - load settings from a file

Upon startup, settings are loaded from the config file.
The global variable 'config_path' determines the path used by save/load_settings. Use -c '<path>' to set another path.
The search path is:
    <working-dir>/{cfilename}
    $XDG_CONFIG_HOME/{cfilename}
    ~/.config/photoreflectance/{cfilename}
    
The current file path is: 
    {config_path}
              """)
    elif topic == "imports":
        print("""Imports:
    numpy as np
    pandas as pd
    matplotlib.pyplot as plt
    os.path """)
    else:
        print(topic.__doc__.strip(" ").strip("\n"))


def connect_devices():
    global lockin, shutter, lamp, mcm
    lockin  = mod_lock_in.connect_device(*select_device_interactive(mod_lock_in.list_devices(), "Select Lock-In-Amplifier: "))
    shutter = mod_shutter.connect_device(*select_device_interactive(mod_shutter.list_devices(), "Select Shutter: "))
    lamp    = mod_lamp.connect_device(*select_device_interactive(mod_lamp.list_devices(), "Select Lamp: "))
    mcm     = mod_monochromator.connect_device(*select_device_interactive(mod_monochromator.list_devices(), "Select Monochromator: "))

def init():
    global lockin, shutter, lamp, mcm, settings, config_path, config_file
    print(r""" __         .__   
_____________  ______ _____/  |________|  |  
\____ \_  __ \/  ___// ___\   __\_  __ \  |  
|  |_> >  | \/\___ \\  \___|  |  |  | \/  |__
|   __/|__|  /____  >\___  >__|  |__|  |____/
|__|              \/     \/                  """ + f"""{version}
Interactive Shell for Photoreflectance measurements
---
Enter 'help()' for a list of commands""")
    parser = argparse.ArgumentParser(
        prog="prsctrl",
        description="measure photoreflectance",
    )
    backend_group = parser.add_mutually_exclusive_group(required=False)
    parser.add_argument("-c", "--config", action="store", help="alternate path to config file")
    args = vars(parser.parse_args())
    from os import environ

    # Load config file
    if path.isfile(cfilename):
        config_path = cfilename
    elif 'XDG_CONFIG_HOME' in environ.keys():
        config_path = path.join(environ["XDG_CONFIG_HOME"], "prsctrl", cfilename)
    else:
        config_path = path.join(path.expanduser("~/.config/prsctrl"), cfilename)
    if args["config"]:
        config_path = args["config"]

    config_file = ConfigFile(config_path, init_values=settings)
    load_settings()

    # setup logging
    log_path = path.expanduser(config_file.get_or("path_log", "~/.cache/prsctrl-interactive.log"))
    makedirs(path.dirname(log_path), exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
        handlers=[
            logging.FileHandler(log_path),
            logging.StreamHandler()
        ]
    )

    if not path.isdir(settings["datadir"]):
        makedirs(settings["datadir"])

    # init the devices
    shutter = connect_device_from_config_or_interactive(config_file, "shutter", "Shutter", mod_shutter, log=log)
    lockin = connect_device_from_config_or_interactive(config_file, "lock-in", "Lock-In Amplifier", mod_lock_in, log=log)
    lamp = connect_device_from_config_or_interactive(config_file, "lamp", "Lamp", mod_lamp, log=log)
    mcm = connect_device_from_config_or_interactive(config_file, "monochromator", "Monochromator", mod_monochromator, log=log)

    # atexit.register(_backend.exit, dev)


if __name__ == "__main__":
    init()