"""
This module contains the ExpData class as part of the QuaCCAToo package.
"""

import matplotlib.pyplot as plt
import numpy as np


class ExpData:
    """
    Class to load experimental data and perform basic data processing.

    Attributes
    ----------
    variable : np.ndarray
        Array containing the variable data
    results : np.ndarray or list of np.ndarray
        Array or list of arrays containing the results data
    variable_name : str
        Name of the variable
    result_name : str
        Name of the results

    Methods
    -------
    subtract_results_columns:
        Subtracts the results of the negative column from the positive column
    offset_correction:
        Substracts a background value from the results
    rescale_correction:
        Multiplies the results by a rescale value
    poly_baseline_correction:
        Performs a polynomial baseline correction to the results
    plot_exp_data:
        Plots the experimental data
    """

    def __init__(
        self,
        file_path,
        variable_column=0,
        results_columns=1,
        variable_name="Time",
        result_name="Expectation Value",
        plot=False,
        figsize=(6, 4),
        figtitle="Experimental Data",
        **loadtxt_kwargs,
    ):
        """
        Constructor of the ExpData class.
        It loads experimental data from a file and sets the variable and results attributes according with the specified column arguments.

        Parameters
        ----------
        file_path  : str
            Path to the file containing the experimental data
        variable_column : int
            Column index of the variable
        results_columns : int or list of int
            Column index of the results
        variable_name : str
            Name of the variable
        result_name : str
            Name of the results
        plot : bool
            Boolean indicating whether to plot the experimental data or not
        figsize : tuple
            Size of the figure for the plot
        figtitle : str
            Title of the figure for the plot
        **loadtxt_kwargs : dict
            Additional arguments for the np.loadtxt function
        """
        if not isinstance(file_path, str):
            raise ValueError("file_path must be a string")

        if not isinstance(variable_column, int):
            raise ValueError("variable_column must be an integer")

        # the results columns needs to be an integer or a list of integers
        if not isinstance(results_columns, int) and not (
            isinstance(results_columns, list) and all(isinstance(col, int) for col in results_columns)
        ):
            raise ValueError("results_columns must be an integer or a list of two integers")

        if not isinstance(variable_name, str) or not isinstance(result_name, str):
            raise ValueError("variable_name and result_name must be strings")

        if not isinstance(loadtxt_kwargs, dict):
            raise ValueError("loadtxt_kwargs must be a dictionary for the np.loadtxt function")

        # loads experimental data from a file with the specified arguments
        exp_data = np.loadtxt(file_path, **loadtxt_kwargs)

        # sets the results and variable attributes of the ExpData object
        self.variable = exp_data[:, variable_column]

        if isinstance(results_columns, int):
            self.results = exp_data[:, results_columns]
        else:
            self.results = [exp_data[:, column] for column in results_columns]

        self.variable_name = variable_name
        self.result_name = result_name

        if not isinstance(plot, bool):
            raise ValueError("plot must be a boolean")
        # plots the experimental data
        elif plot:
            self.plot_exp_data(figsize=figsize, figtitle=figtitle)

    def subtract_results_columns(
        self, pos_col=0, neg_col=1, plot=False, figsize=(6, 4), figtitle="Subtracted Expt. Data"
    ):
        """
        Overwrites the results attribute substracting the results of the negative column from the positive column.

        Parameters
        ----------
        pos_col: int
            Index of the positive column
        neg_col: int
            Index of the negative column
        plot: bool
            Boolean indicating whether to plot the experimental data or not
        figsize: tuple
            Size of the figure for the plot
        figtitle: str
            Title of the figure for the plot
        """
        if not isinstance(self.results[pos_col], np.ndarray) or not isinstance(
            self.results[neg_col], np.ndarray
        ):
            raise ValueError(f"pos_col={pos_col} and neg_col={neg_col} where not found in the results.")

        self.results = self.results[pos_col] - self.results[neg_col]

        if not isinstance(plot, bool):
            raise ValueError("plot must be a boolean")
        elif plot:
            self.plot_exp_data(figsize=figsize, figtitle=figtitle)

    def offset_correction(
        self, background_value, plot=False, figsize=(6, 4), figtitle="Expt. Data with Offset Correction"
    ):
        """
        Overwrites the results attribute substracting the background value from the results.

        Parameters
        ----------
        background_value : int or float
            Value to be substracted from the results
        plot : bool
            Boolean indicating whether to plot the experimental data or not
        figsize : tuple
            Size of the figure for the plot
        figtitle : str
            Title of the figure for the plot
        """
        if not isinstance(background_value, (int, float)):
            raise ValueError("background_value must be a number.")

        if isinstance(self.results, np.ndarray):
            self.results = self.results - background_value
        elif isinstance(self.results, list) and all(
            isinstance(result, np.ndarray) for result in self.results
        ):
            self.results = [result - background_value for result in self.results]
        else:
            raise ValueError("Results must be a numpy array or a list of numpy arrays")

        if not isinstance(plot, bool):
            raise ValueError("plot must be a boolean")
        elif plot:
            self.plot_exp_data(figsize=figsize, figtitle=figtitle)

    def rescale_correction(
        self, rescale_value, plot=False, figsize=(6, 4), figtitle="Expt. Data with Rescale Correction"
    ):
        """
        Overwrites the results attribute multiplying the results by the rescale value.

        Parameters
        ----------
        rescale_value : int or float
            Value to be multiplied by the results
        plot : bool
            Boolean indicating whether to plot the experimental data or not
        figsize : tuple
            Size of the figure for the plot
        figtitle : str
            Title of the figure for the plot
        """
        if not isinstance(rescale_value, (int, float)):
            raise ValueError("rescale_value must be a number.")

        if isinstance(self.results, np.ndarray):
            self.results = self.results * rescale_value
        elif isinstance(self.results, list) and all(
            isinstance(result, np.ndarray) for result in self.results
        ):
            self.results = [result * rescale_value for result in self.results]
        else:
            raise ValueError("Results must be a numpy array or a list of numpy arrays")

        if not isinstance(plot, bool):
            raise ValueError("plot must be a boolean")
        elif plot:
            self.plot_exp_data(figsize=figsize, figtitle=figtitle)

    def poly_base_correction(
        self,
        x_start=None,
        x_end=None,
        poly_order=2,
        plot=False,
        figsize=(6, 4),
        figtitle="Expt. Data with Polynomial Baseline Correction",
    ):
        """
        Overwrites the results attribute performing a polynomial baseline correction.
        The baseline is fitted to the data between x_start and x_end, representing the start and end of the xaxis index.

        Parameters
        ----------
        x_start : int or list of int
            Start index of the x axis for the baseline fit
        x_end : int or list of int
            End index of the x axis for the baseline fit
        poly_order : int
            Order of the polynomial to fit the baseline
        plot : bool
            Boolean indicating whether to plot the experimental data or not
        figsize : tuple
            Size of the figure for the plot
        figtitle : str
            Title of the figure for the plot

        """
        # check all variables
        if x_start is None:
            x_start = 0
        elif not isinstance(x_start, int) and not (
            isinstance(x_start, list) and all(isinstance(x, int) for x in x_start)
        ):
            raise ValueError("x_start must be a integer index or a list of integer indexes.")

        if x_end is None:
            x_end = -1
        elif not isinstance(x_end, int) and not (
            isinstance(x_end, list) and all(isinstance(x, int) for x in x_end)
        ):
            raise ValueError("x_end must be a integer index or a list of integer indexes.")

        if not isinstance(poly_order, int):
            raise ValueError("poly_order must be an integer.")

        # crops the x and y axis for performing the baseline fit
        if isinstance(x_start, int) and isinstance(x_end, int):
            baseline_xaxis = self.variable[x_start:x_end]
            baseline_yaxis = self.results[x_start:x_end]
        elif isinstance(x_start, list) and isinstance(x_end, list) and len(x_start) == len(x_end):
            baseline_xaxis = np.concatenate(
                [self.variable[x_start[i] : x_end[i]] for i in range(len(x_start))]
            )
            baseline_yaxis = np.concatenate(
                [self.results[x_start[i] : x_end[i]] for i in range(len(x_start))]
            )
        else:
            raise ValueError("x_start and x_end must int or a list of the same length.")

        if isinstance(self.results, np.ndarray):
            poly_fit = np.polyfit(baseline_xaxis, baseline_yaxis, poly_order)
            self.results -= np.polyval(poly_fit, self.variable)

        elif isinstance(self.results, list) and all(
            isinstance(result, np.ndarray) for result in self.results
        ):
            poly_fit = [
                np.polyfit(baseline_xaxis[i], baseline_yaxis[i], poly_order)
                for i in range(len(baseline_xaxis))
            ]
            self.results = [
                self.results[i] - np.polyval(poly_fit[i], self.variable) for i in range(len(self.results))
            ]

        if not isinstance(plot, bool):
            raise ValueError("plot must be a boolean")
        elif plot:
            self.plot_exp_data(figsize=figsize, figtitle=figtitle)

    def plot_exp_data(self, figsize=(6, 4), figtitle="Experimental Data"):
        """
        Plots the experimental data.

        Parameters
        ----------
        figsize : tuple
            Size of the figure for the plot
        figtitle : str
            Title of the figure for the plot
        """
        if not (isinstance(figsize, tuple) or len(figsize) == 2):
            raise ValueError("figsize must be a tuple of two positive floats")

        if not isinstance(figtitle, str):
            raise ValueError("figtitle must be a string")

        fig, ax = plt.subplots(1, 1, figsize=figsize)

        # check if the results is a list of results or a single result
        if isinstance(self.results, np.ndarray):
            ax.scatter(self.variable, self.results, lw=2, alpha=0.7, label="Observable", s=15)

        elif isinstance(self.results, list) and all(
            isinstance(result, np.ndarray) for result in self.results
        ):
            for itr in range(len(self.results)):
                ax.scatter(self.variable, self.results[itr], label=f"Observable {itr}", alpha=0.7, s=15)

        else:
            raise ValueError("Results must be a numpy array or a list of numpy arrays")

        # set the x-axis limits to the variable of the experiment
        ax.set_xlim(self.variable[0], self.variable[-1])

        ax.set_xlabel(self.variable_name)
        ax.set_ylabel(self.result_name)
        ax.legend(loc="upper right", bbox_to_anchor=(1.2, 1))
        ax.set_title(figtitle)
