"""The LinxProcessor for digital lock-in."""

import logging
import os
from pathlib import Path
from typing import Literal, Self

import numpy as np
from scipy import optimize, signal

from ..core import BaseProcessor, sp
from ..log import configure_logger
from ._config import LinxConfig

PICKUP_NAME = "pickup"
SERIE_NAME = ""

logger = logging.getLogger(__name__)


class LinxProcessor(BaseProcessor):
    """A Processor class for digital lock-in."""

    cfg: LinxConfig

    def __init__(self, *args, **kwargs):
        configure_logger(logger, "linx.log")
        self._config_cls = LinxConfig

        # Serie name for the NeXus data processed group name
        self._serie_name = SERIE_NAME

        super().__init__(*args, **kwargs)

        # Initialize attributes
        if not self.is_nexus_file:
            self._prepare_metadata()
        else:
            # Fill missing attributes
            self.npoints = self.get_data_processed("magfield").size
            self.nchannels = len(self.cfg.measurements)
            self._bin_data = np.array([])

        # Store names
        self._pu_name = PICKUP_NAME
        # Read signal and reference available names
        self._signal_names = self._get_signal_names()
        # Check those names are set correctly in the settings
        self.ref_name = self.cfg.settings.ref_name
        self.signal_name = self.cfg.settings.signal_name

    @property
    def idx_serie(self) -> str:
        """
        Track the current signal.

        The returned value is read from the "signal_name" parameter in the "settings"
        section of the configuration.
        """
        return self.cfg.settings.signal_name

    @idx_serie.setter
    def idx_serie(self, value: str):
        """
        Setter for `idx_serie`.

        The "signal_name" parameter in the "settings" section of the configuration is
        updated and the NeXus group for this serie is created.
        """
        if value not in self._signal_names:
            raise ValueError(
                f"'{value}' was not found in the measurements, check your settings"
            )
        self.cfg.settings.signal_name = value
        self.metadata["phase_offset"][value] = self.cfg.settings.phase_offset
        self.create_data_serie()

    @property
    def signal_name(self) -> str:
        """
        Track the current signal. Alias for `idx_serie`.

        The returned value is read from the "signal_name" parameter in the "settings"
        section of the configuration.
        """
        return self.idx_serie

    @signal_name.setter
    def signal_name(self, value: str):
        """Setter for `signal_name`, the alias for `idx_serie`."""
        self.idx_serie = value
        logger.info(f"Using signal : '{value}'")

    @property
    def ref_name(self) -> str:
        """
        Track the current reference.

        The returned value is read from the "ref_name" parameter in the "settings"
        section of the configuration.
        """
        return self.cfg.settings.ref_name

    @ref_name.setter
    def ref_name(self, value: str):
        """
        Setter for `ref_name`.

        The "ref_name" parameter in the "settings" section of the configuration is
        updated.
        """
        if value not in self._signal_names:
            raise ValueError(
                f"'{value}' was not found in the measurements, check your settings"
            )
        self.cfg.settings.ref_name = value
        logger.info(f"Using reference : '{value}'")

    @property
    def f0(self) -> float:
        """
        Get the center frequency of the current reference signal.

        It is read from the metadata dictionary.
        """
        if not self._check_f0_computed():
            logger.error(f"f0 was not computed for '{self.signal_name}'")
            return -1.0
        return self.metadata["f0"][self.signal_name]

    @f0.setter
    def f0(self, value: float):
        """
        Set the center frequency of the current reference signal.

        It is set in the metadata dictionary and in the configuration.
        """
        self.metadata["f0"][self.signal_name] = value
        self.cfg.settings.f0 = value

    @property
    def phase_offset(self) -> float:
        """
        Get the phase offset of the current reference signal.

        It is read from the metadata dictionary.
        """
        return self.metadata["phase_offset"][self.signal_name]

    @phase_offset.setter
    def phase_offset(self, value: float):
        """
        Set the phase offset of the current reference signal.

        It is set in the metadata dictionary, and in the configuration.
        """
        self.metadata["phase_offset"][self.signal_name] = value
        self.cfg.settings.phase_offset = value

    @property
    def amp0(self) -> float:
        """
        Get the amplitude of the current reference signal.

        It is read from the metadata dictionary.
        """
        if not self._check_demodulated():
            logger.warning(f"amp0 was not computed for '{self.ref_name}'")
            return 0
        return self.metadata["amp0"][self.signal_name]

    @amp0.setter
    def amp0(self, value: float):
        """
        Set the amplitude of the current reference signal.

        It is set in the metadata dictionary.
        """
        self.metadata["amp0"][self.signal_name] = value

    def _reinit(self) -> None:
        """Reinitialize data objects."""
        super()._reinit()

        self._prepare_metadata()

    def load_pickup(self) -> Self:
        """Read pickup time serie from the binary file."""
        if self._bin_data.size == 0:
            self._load_bin_file()

        # Get the correct index
        pickup_idx = self.cfg.measurements[self._pu_name]

        # Set data
        logger.info("Reading pickup data...")
        self.set_data_raw("pickup", self._bin_data[:, pickup_idx])
        if self.get_data_raw("time_exp", checkonly=True):
            self.set_data_raw("pickup_time", self.get_data_raw("time_exp"))
        else:
            self.set_data_raw(
                "pickup_time",
                self.get_time_exp(),
            )
        logger.info("Pickup signal loaded")

        return self

    def compute_field(self, method: str = "trapz") -> Self:
        """Compute magnetic field."""
        super().compute_field()

        # Apply time offset
        self.apply_time_offset()

        return self

    def load_data(self) -> Self:
        """Load raw data."""
        self.load_oscillo()
        return self

    def preprocess(self) -> Self:
        """Find reference center frequency f0."""
        self.find_f0()
        return self

    def align_field(self, on_pickup: bool = True, **kwargs) -> Self:
        """
        Align magnetic field on the `time_exp` vector.

        **kwargs are ignored.
        """
        if not self.get_data_processed("time_exp", checkonly=True):
            logger.warning("Experiment time vector not built, load data first")
            return self

        # Align the pickup signal
        super().align_field(on_pickup=on_pickup, left=0, right=0)

        return self

    def run_analysis(self) -> Self:
        """Run the analysis : perform the demodulation and computes the resistance."""
        self.demodulate().compute_resistance()
        return self

    def export_csv(
        self,
        fname: str | Path | None = None,
        sep: str = "\t",
        downsample: int = 0,
        average: bool = False,
    ) -> None:
        r"""
        Export results to a CSV file.

        The CSV will contain the following columns : time, pickup, B, R_in_phase, R_out_phase.

        Optionally, downsample the data with or without averaging.

        Parameters
        ----------
        fname : str or Path or None, optional
            Path to the output file. If None, the filename is generated automatically.
            Default is None.
        sep : str, optional
            Delimiter in the CSV file. Default is "\t".
        downsample : int, optional
            Downsample the data before exporting. Disable with 0 (default).
        average : bool, optional
            If downsampling is enabled, use the average instead of direct indexing.
            Default is True.
        """
        if fname is None:
            fname = self.get_csv_filename()

        logger.info(f"Saving results at {fname}...")

        # Collect data
        if self.get_data_processed("pickup", checkonly=True):
            pickup = self.get_data_processed("pickup")
        else:
            pickup = self.get_data_raw("pickup")
        data = np.concatenate(
            (
                self.get_data_processed("time_exp")[..., np.newaxis],
                pickup[..., np.newaxis],
                self.get_data_processed("magfield")[..., np.newaxis],
                self.get_data_serie("resistance_in_phase")[..., np.newaxis],
                self.get_data_serie("resistance_out_phase")[..., np.newaxis],
            ),
            axis=1,
        )
        if downsample > 0:
            data = sp.subsample_array(
                data, downsample, start=0, average=average, axis=0
            )

        # Configure output file
        header = sep.join(("time", "pickup", "field", "R_in_phase", "R_out_phase"))

        # Save
        np.savetxt(fname, data, delimiter=sep, header=header, comments="")

        logger.info("Results saved")

    def get_bin_file_properties(self) -> Self:
        """Determine binary file number columns and number of samples."""
        # Get parameters
        filepath = self.cfg.filenames["oscillo"]
        precision = self.cfg.files["oscillo"].precision

        self.nchannels, self.nsamples = self._get_bin_file_properties(
            filepath, precision=precision
        )

        return self

    def load_oscillo(self) -> Self:
        """Read the reference and signal data for the current measurement."""
        if self._bin_data.size == 0:
            self._load_bin_file()

        # Get the correct indices
        ref_idx = self.cfg.measurements[self.ref_name]
        signal_idx = self.cfg.measurements[self.signal_name]

        # Set the data
        logger.info(f"Reading {self.ref_name}...")
        self.set_data_raw(self.ref_name, self._bin_data[:, ref_idx])
        logger.info(f"Reading {self.signal_name}...")
        self.set_data_raw(self.signal_name, self._bin_data[:, signal_idx])
        if self.get_data_raw("pickup_time", checkonly=True):
            time_exp = self.get_data_raw("pickup_time")
        else:
            time_exp = self.get_time_exp()
        self.set_data_raw("time_exp", time_exp)
        # Set in data_processed as well
        self.set_data_processed("time_exp", time_exp)

        logger.info("Data loaded")

        return self

    def get_time_exp(self) -> np.ndarray:
        """Build the experiment time vector."""
        return np.linspace(
            0,
            self.nsamples / self.cfg.parameters.samplerate,
            self.nsamples,
            endpoint=False,
        )

    def apply_time_offset(self) -> Self:
        """Apply a time offset on the magnetic field time vector."""
        time_offset = self.cfg.settings.time_offset

        if not self._check_field_computed():
            logger.warning("Field not computed, cannot apply time offset")
            return self

        # Get the original time vector in case an offset was already applied
        if not self.get_data_raw("time_exp", checkonly=True):
            # Build the time vector in case it was not done before
            self.set_data_raw("time_exp", self.get_time_exp())
        time_vector = self.get_data_raw("time_exp")
        logger.info(f"Applying a {time_offset * 1e3}ms offset on magnetic field...")
        self.set_data_processed("magfield_time", time_vector - time_offset)
        logger.info("Time offset applied on magnetic field")

        # Align field if it was created
        if self.get_data_processed("time_exp", checkonly=True):
            self.align_field(on_pickup=True)

        return self

    def find_f0(self) -> Self:
        """Find the center frequency in the reference signal."""
        if not self._check_data_loaded():
            self.load_oscillo()

        logger.info(f"Finding center frequency in '{self.ref_name}'...")
        self.f0 = sp.find_f0(
            self.get_data_raw(self.ref_name), self.cfg.parameters.samplerate
        ).mean()
        logger.info(f"Found {self.f0 / 1000:3.3f}kHz")

        return self

    def build_lowpass_filter(
        self,
    ) -> (
        tuple[np.ndarray, Literal["sos", "ba"]]
        | tuple[tuple[np.ndarray, np.ndarray], Literal["sos", "ba"]]
    ):
        """Build low-pass filter."""
        filter_order = self.cfg.settings.filter_order
        filter_fc = self.cfg.settings.filter_fc
        filter_edge = self.cfg.settings.filter_edge
        fs = self.cfg.parameters.samplerate
        pb_mult = self.cfg.settings.filter_passband_mult
        sb_mult = self.cfg.settings.filter_stopband_mult
        pb_att = self.cfg.settings.filter_passband_att
        sb_att_mult = self.cfg.settings.filter_stopband_att_mult

        # Filter properties
        wp = filter_fc * pb_mult / fs  # passband
        ws = wp * sb_mult  # stopband
        gpass = pb_att  # max loss in passband
        gstop = filter_order * sb_att_mult  # min attenuation in stopband
        order, wn = signal.buttord(wp, ws, gpass, gstop)

        # Create filter
        if filter_edge == "pad":
            # Using padding to apply the filter, so we can use second order sections
            res = signal.butter(order, wn, btype="lowpass", output="sos")
            filt_type = "sos"
        elif filter_edge == "gust":
            # Using Gustafsson’s method, we need numerator/denominator
            res = signal.butter(order, wn, btype="lowpass", output="ba")
            filt_type = "ba"
        else:
            raise ValueError(
                f"{filter_edge} is not a valid method, choose 'pad' or 'gust'"
            )

        return res, filt_type

    def demodulate(self) -> Self:
        """Digital demodulation."""
        # Checks
        if not self._check_data_loaded():
            logger.warning("Data not loaded")
            return self
        if not self._check_f0_computed():
            self.find_f0()

        # Get reference phase fit parameters
        fs = self.cfg.parameters.samplerate
        ncycles = self.cfg.settings.phase_fit_ncycles
        f0 = self.f0

        # Define cos function to fit reference phase
        def cos_func(
            x: np.ndarray, A: float | np.ndarray, phi: float | np.ndarray
        ) -> np.ndarray:
            return A * np.cos(2 * np.pi * f0 * x + phi)

        # Get signals
        fit_npoints = int(ncycles * fs / f0)  # number of points to use for the fit
        reftime = self.get_data_raw("time_exp")[:fit_npoints]
        ref = self.get_data_raw(self.ref_name)[:fit_npoints]
        sigtime = self.get_data_raw("time_exp")
        sig = self.get_data_raw(self.signal_name)

        # Fit reference signal to find the phase
        logger.info(f"Fitting '{self.ref_name}' phase...")
        amp, phi = sp.fit_signal(
            cos_func, reftime, ref, bounds=([0, -np.inf], [np.inf, np.inf])
        )
        amp = float(amp)  # for type checking
        phi = float(phi)  # for type checking
        logger.info(f"Found amplitude: {amp:2.4f}V and phase: {phi:0.4f}rad")

        # Build filter
        logger.info("Building low-pass filter...")
        filt, filt_type = self.build_lowpass_filter()
        logger.info("Low-pass fitler built")

        # Demodulation
        phase_offset = self.phase_offset
        logger.info(
            f"Demodulation at {self.f0 / 1e3:2.4f}kHz (using phase: "
            f"{phi + phase_offset:0.4f}rad)..."
        )
        in_phase, out_phase = self._demodulate(
            filt_type,
            None,
            None,
            sigtime,
            sig,
            cos_func,
            filt,
            extra_phase=phase_offset,
            refas=amp,
            refphis=phi,
        )
        logger.info("Signal demodulated")

        # Store
        self.amp0 = amp
        self.set_data_serie("in_phase", in_phase.ravel())
        self.set_data_serie("out_phase", out_phase.ravel())

        # Prepare placeholders that will be replaced by actual NXlink when the
        # object has a NXroot entry
        time_exp = f"!link to:processed/analysis/{self._results_name}/time_exp"
        magfield = f"!link to:processed/analysis/{self._results_name}/magfield"
        # Store the link in the 'serie' group
        self.set_data_serie("time_exp", time_exp)
        self.set_data_serie("magfield", magfield)

        return self

    def find_phase_offset(self) -> Self:
        """Find the phase shift that maximises signal in I."""
        if not self._check_demodulated():
            logger.warning("Cannot find phase, demodulation was not performed")
            return self
        # Get parameters
        method = self.cfg.settings.auto_phase_method
        if method == "at peak field":
            if not self._check_field_computed():
                logger.warning(
                    f"Field was not computed, cannot find phase with method {method}"
                )
                return self
            idx_max_field = int(np.argmax(self.get_data_processed("magfield")))
        else:
            idx_max_field = 0

        # Get signals
        in_phase = self.get_data_serie("in_phase")
        out_phase = self.get_data_serie("out_phase")

        # Compute
        logger.info(f"Finding phase with the '{method}' method...")
        phase = self._find_phase(in_phase, out_phase, method, idx_max=idx_max_field)
        logger.info(f"Found {phase:1.3f} rad ({np.rad2deg(phase):3.3f}deg)")

        # Store
        self.phase_offset = phase
        return self

    def compute_resistance(self) -> Self:
        """Derive resistance in ohms from demodulated signal."""
        # Checks
        if not self._check_demodulated():
            logger.warning("Cannot compute resistance, demodulation was not performed")
            return self

        # Get parameters
        current = self.cfg.parameters.current
        gain = self.cfg.parameters.gain
        amp0 = self.amp0
        rms = self.cfg.parameters.gain_rms

        logger.info("Computing resistance...")
        intgain = 0.5 * amp0 * gain * current
        if rms:
            intgain *= np.sqrt(2)

        res_in = self.get_data_serie("in_phase") / intgain
        res_out = self.get_data_serie("out_phase") / intgain

        logger.info(f"Done, found max. resistance {res_in.max():2.3f}ohm")

        # Store
        self.set_data_serie("resistance_in_phase", res_in)
        self.set_data_serie("resistance_out_phase", res_out)

        return self

    def get_csv_filename(self) -> str:
        """
        Generate output CSV file name for the current signal.

        The file is placed in the data directory set in the configuration. Its name
        contains the signal name.
        """
        fout = os.path.join(self.cfg.data_directory, self.cfg.expid + "-results")

        return f"{fout}_{self.signal_name}.csv"

    def _prepare_metadata(self) -> None:
        """Prepare metadata and data objects."""
        # Metadata
        self.metadata = dict()
        self.metadata["f0"] = dict()  # will be signal_name: f0
        self.metadata["phase_offset"] = dict()  # will be signal_name: phase
        self.metadata["amp0"] = dict()  # will be signal_name: amplitude

        self.nchannels = 0
        self.nsamples = 0

        self._bin_data = np.array([])

        # Read binary file properties
        self.get_bin_file_properties()
        # Prepare data groups
        self.create_data_serie()

        # Fill metadata from configuration
        self.phase_offset = self.cfg.settings.phase_offset

    def _get_signal_names(self) -> set[str]:
        """Determine signal and reference names from configuration file measurements."""
        return {*self.cfg.measurements.keys()} - {self._pu_name}

    def _get_bin_file_properties(
        self, filepath: str | Path, precision: int = 4
    ) -> tuple[int, int]:
        """
        Determine binary file number columns and number of samples.

        Parameters
        ----------
        filepath : str or Path
            Path to the binary file.
        precision : int, optional
            Precision of floats in the binary file (typically, 4, i.e. float32). Default
            is 4.

        Returns
        -------
        nchannels : int
            Number of columns in the binary file.
        nsamples : int
            Number of data points per channel.
        """
        if isinstance(filepath, str):
            filepath = Path(filepath)

        ncols = len(self.cfg.measurements.keys())  # number of expected columns
        filesize = filepath.stat().st_size  # size in bytes
        nsamples = int(filesize / (ncols * precision))

        return ncols, nsamples

    def _load_bin_file(self) -> None:
        """
        Load binary the binary file in a temporary variable.

        Since all channels, including the pickup signal, are stored in the same binary
        file and in a not-contiguous manner, it is not possible to load a single channel
        from the file. Therefore, we cache the whole file here and access and store its
        data on-demand.
        """
        # Get parameters
        filepath = self.cfg.filenames["oscillo"]
        precision = self.cfg.files["oscillo"].precision
        endian = self.cfg.files["oscillo"].endian
        order = self.cfg.files["oscillo"].order

        # Read the time series
        self._bin_data = (
            np.fromfile(
                filepath,
                dtype=f"{endian}f{precision}",
                count=self.nchannels * self.nsamples,
            )
            .reshape((self.nsamples, self.nchannels), order=order)
            .astype(float)
        )

    def _demodulate(
        self,
        filt_type: Literal["sos", "ba"],
        reftime: np.ndarray | None,
        ref: np.ndarray | None,
        sigtime,
        sig,
        cos_func,
        filt: np.ndarray | tuple[np.ndarray, np.ndarray],
        extra_phase: float = 0.0,
        refas: float | None = None,
        refphis: float | None = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Perform digital demodulation.

        Wrapper around `sp.demodulate()`, determining filter type and edge handling.
        """
        # Determine how to apply low-pass filter
        fs = self.cfg.parameters.samplerate
        filter_fc = self.cfg.settings.filter_fc

        if filt_type == "sos":
            # Second order sections, with padding
            padding_type = self.cfg.settings.filter_padding_type
            padding_length = self.cfg.settings.filter_padding_length
            if padding_length == -1:
                # Automatic padding length
                padlen = int(10 * fs / (2 * np.pi * filter_fc))
            filt_arg = (filt,)  # convert sos to a tuple (sos, ) to unpack it later
            filt_kwargs = dict(padtype=padding_type, padlen=padlen)
            filt_edge_method = "pad"
        elif filt_type == "ba":
            # Not-recommended scipy.signal.filtfilt() with Gustafsson’s method
            filter_ir_length = self.cfg.settings.filter_ir_length
            if filter_ir_length == -1:
                # Automatic impulse response length
                irlen = int(20 * fs / (2 * np.pi * filter_fc))
            filt_arg = filt  # already a tuple (b, a)
            filt_kwargs = dict(irlen=irlen)
            filt_edge_method = "gust"
        else:
            raise ValueError(
                f"{filt_type} is not a valid filter type, it must be 'sos' or 'ba'"
            )

        # Run the demodulation
        return sp.demodulate(
            reftime,
            ref,
            sigtime,
            sig,
            cos_func,
            *filt_arg,
            extra_phase=extra_phase,
            refas=refas,
            refphis=refphis,
            filter_edge_method=filt_edge_method,
            filter_kws=filt_kwargs,
        )

    @staticmethod
    def _find_phase_minoop(in_phase: np.ndarray, out_phase: np.ndarray) -> float:
        """
        Find phase that minimizes the integral of `out_phase`, keeping `in_phase` > 0.

        Parameters
        ----------
        in_phase, out_phase : np.ndarray
            I and Q signal vectors.

        Returns
        -------
        phase : float
            Phase.
        """

        # Define function to minimize
        def rot_oop(th: np.ndarray) -> float:
            return np.abs(np.sin(-th) * in_phase + np.cos(-th) * out_phase).sum() / len(
                in_phase
            )

        res = optimize.minimize_scalar(
            rot_oop, method="bounded", bounds=(-np.pi, np.pi)
        )
        ang = res.x
        if (np.cos(-ang) * in_phase - np.sin(-ang) * out_phase).sum() < 0:
            ang += np.pi
        # Wrap to -pi, +pi
        ang = np.arctan2(np.sin(ang), np.cos(ang))

        return ang

    def _find_phase(
        self,
        in_phase: np.ndarray,
        out_phase: np.ndarray,
        method: Literal["at peak field", "minimum oop"],
        idx_max: int = 0,
    ) -> float:
        """
        Find the phase shift that maximises signal in I.

        Parameters
        ----------
        in_phase, out_phase : np.ndarray
            I and Q signal vectors.
        method : {"at peak field", "minimum oop"}
            Method to find the phase. "at peak field" maximises the in-phase signal at
            peak field. "minimum oop" minimizes the integral of the out of phase signal
            while keeping a positive in-phase signal.
        idx_max : int, optional
            Index at maximum field. Required for "at peak field" method.
        """
        match method:
            case "at peak field":
                return np.arctan2(out_phase[idx_max], in_phase[idx_max])
            case "minimum oop":
                return self._find_phase_minoop(in_phase, out_phase)
            case _:
                logging.error(
                    f"Unkown method to find phase : {method}. Choose 'at peak field' or"
                    " 'minimum oop'"
                )

    def _check_data_loaded(self) -> bool:
        """Check whether the data was loaded."""
        return self.get_data_raw(self.signal_name, checkonly=True)

    def _check_demodulated(self) -> bool:
        """Check whether the signal has been demodulated."""
        return self.get_data_serie("in_phase", checkonly=True) and self.get_data_serie(
            "out_phase", checkonly=True
        )

    def _check_resistance_computed(self) -> bool:
        """Check whether the resistance was computed."""
        return self.get_data_serie(
            "resistance_in_phase", checkonly=True
        ) and self.get_data_serie("resistance_out_phase", checkonly=True)

    def _check_f0_computed(self) -> bool:
        """Check whether center frequency was found."""
        if self.signal_name in self.metadata["f0"]:
            if self.metadata["f0"][self.signal_name] > 0:
                return True
        return False
