from sys import argv
from backscattering_analyzer.settings import Settings
from backscattering_analyzer.signal import Signal
from backscattering_analyzer import (
    compute_light,
    fit_compute_light,
    opt_compute_light,
    interpolate,
)
from numpy import loadtxt, logspace, pi
from scipy.io.matlab import loadmat
from scipy.optimize import Bounds, minimize


class Analyzer:
    """
    Utility class to study backscattering light in VIRGO
    """

    def __init__(
        self, arguments: None | list[str] | str = None
    ) -> None:
        if arguments is None:
            options = []
        elif isinstance(arguments, str):
            options = arguments.split(" ")
        else:
            options = arguments
        self.settings = Settings(argv[1:] + options)
        self.load()
        self.process_movement()
        self.compute_light()

    def load(self):
        """
        Load all data
        """
        self.settings.log("Load all data")
        try:
            self.load_bench()
            self.load_mirror()
            self.load_data()
            self.load_reference()
            self.load_coupling()
        except Exception:
            raise Exception("An error occured during data loading")

    def load_bench(self):
        """
        Load bench movement
        """
        file = self.settings.bench_file()
        self.settings.log("loading bench movement")
        try:
            data = loadtxt(file).T
            self.bench_movement = (
                Signal(data[0], data[1], self.settings) * 1e-6
            )  # um
        except OSError:
            raise Exception("{file} does not exist".format(file=file))

    def load_mirror(self):
        """
        Load mirror movement
        """
        file = self.settings.mirror_file()
        self.settings.log("loading mirror movement")
        try:
            data = loadtxt(file).T
            self.mirror_movement = (
                Signal(data[0], data[1], self.settings) * 1e-6
            )
        except OSError:
            raise Exception("{file} does not exist".format(file=file))

    def load_data(self):
        """
        Load excited h(t)
        """
        file = self.settings.data_file()
        self.settings.log("loading excited h(t)")
        try:
            data = loadtxt(file).T
            self.data_signal = Signal(data[0], data[1], self.settings)
        except OSError:
            raise Exception("{file} does not exist".format(file=file))

    def load_reference(self):
        """
        Load reference h(t)
        """
        file = self.settings.reference_file()
        self.settings.log("loading reference h(t)")
        try:
            data = loadtxt(file).T
            self.reference_signal = Signal(
                data[0], data[1], self.settings
            )
        except OSError:
            raise Exception("{file} does not exist".format(file=file))

    def load_coupling(self):
        """
        Load modelisation coupling data
        """
        file = self.settings.modelisation_file()
        self.settings.log("loading matlab modelisation data")
        try:
            self.modelisation = loadmat(file)
            coupling_values = self.modelisation[
                self.settings.coupling_name()
            ]
            self.coupling = [
                Signal(
                    self.modelisation["freq"][0],
                    coupling,
                    self.settings,
                )
                for coupling in coupling_values
            ]
        except OSError:
            raise Exception("{file} does not exist".format(file=file))

    def process_movement(self):
        """
        Clean and compute relative movement between the bench and the
        miror
        """
        self.settings.log("computing relative movement")

        self.movement = (
            (
                self.bench_movement * self.settings.calib_bench
                - self.mirror_movement * self.settings.calib_mirror
            )
            .detrend("linear")
            .low_pass_filter(5 * self.settings.vibration_frequency)
        )

    def compute_light(self) -> None:
        """
        Compute projection with current bench excitation
        """
        self.projection = compute_light(
            scatter_factor=self.settings.scattering_factor,
            coupling=self.coupling,
            movement=self.movement,
            wavelength=self.settings.wavelength,
            power_in=self.settings.power_in,
            power_out=self.settings.power_out,
        )

    def fit_scatter_factor(self, guess: None | float = None) -> float:
        """
        Find the best scatter factor (first order only) in the given
        range
        """
        if guess is None:
            guess = self.settings.scattering_factor[0]
        phase = 4 * pi / self.settings.wavelength
        factor_n = (self.movement * phase).sin().psd().sqrt()
        coupling_n = self.coupling[0].abs()
        factor_d = (self.movement * phase).cos().psd().sqrt()
        coupling_d = self.coupling[1].abs()

        coupling_d = coupling_d.cut_x(
            10, 200
        )  # cut signal between 10 and 200 Hz

        factor_n, coupling_n, factor_d, coupling_d, data, reference = (
            interpolate(
                (
                    factor_n,
                    coupling_n,
                    factor_d,
                    coupling_d,
                    self.data_signal.psd().sqrt(),
                    self.reference_signal.psd().sqrt(),
                )
            )
        )

        bounds = Bounds(1e-30, 1e-3)
        min_result = minimize(
            fit_compute_light,
            guess,
            (
                factor_n,
                coupling_n,
                factor_d,
                coupling_d,
                self.settings.power_in,
                self.settings.power_out,
                data,
                reference,
            ),
            method="Powell",
            bounds=bounds,
        )

        if not min_result.success:
            raise Exception(min_result.message)

        self.settings.log(
            "found the best scattering factor in {} iterations".format(
                min_result.nit
            )
        )

        import matplotlib.pyplot as plt

        projection = Signal(
            factor_n.x,
            opt_compute_light(
                scatter_factor=min_result.x,
                factor_n=factor_n,
                coupling_n=coupling_n,
                factor_d=factor_d,
                coupling_d=coupling_d,
                power_in=self.settings.power_in,
                power_out=self.settings.power_out,
            ),
            self.settings,
        )

        _ = plt.loglog(projection.x, projection.y, label="projection")
        _ = plt.loglog(reference.x, reference.y, label="référence")
        _ = plt.loglog(data.x, data.y, label="data")
        _ = plt.loglog(
            reference.x, projection.y + reference.y, label="somme"
        )
        _ = plt.legend()
        _ = plt.show()

        """
        as the minimization does not work (why ?), visual help
        """
        x = logspace(-50, 0, 1000)
        y = [
            fit_compute_light(
                x[i],
                factor_n,
                coupling_n,
                factor_d,
                coupling_d,
                self.settings.power_in,
                self.settings.power_out,
                data,
                reference,
            )
            for i in range(len(x))
        ]
        _ = plt.loglog(x, y)
        _ = plt.show()

        return min_result.x