from __future__ import annotations

from os.path import isfile
from typing import Tuple, List, Optional, Union, Callable
from datetime import datetime
from warnings import warn
from scipy.interpolate import make_interp_spline, BSpline

import numpy as np
import matplotlib.pyplot as plt


class UVVisSpectrum:
    """
    The UVVisSpectrum class allows the manipulation of UV-Visible data originated from experimental
    measurements. The class provides a standard constructor returning an empty object and a series
    of classmethods designed to parse specific file formats.
    """

    def __init__(self) -> None:
        self.title: str = None
        self.instrument: Optional[str] = None
        self.__timestamp: Optional[datetime] = None
        self.__wavelength: List[float] = []
        self.__absorbance: List[float] = []

    def __str__(self) -> str:
        msg = f"UV-Visible: {self.title}\n"
        msg += "-----------------------------------------------------------------\n"
        msg += f"Date: {self.__timestamp}\n"
        msg += f"Instrument: {self.instrument}\n"
        msg += f"Wavelength: {max(self.__wavelength)} - {min(self.__wavelength)} nm\n"
        msg += f"Max absorbance: {max(self.__absorbance)}\n"
        msg += f"Min absorbance: {min(self.__absorbance)}\n"
        return msg

    def __repr__(self) -> str:
        return str(self)

    @property
    def timestamp(self) -> datetime:
        """
        The time at which the measurement has been started.

        Returns
        -------
        datetime
            The datetime object encoding the time at which the measurement has been started.
        """
        return self.__timestamp

    @property
    def wavelength(self) -> List[float]:
        """
        The wavelength values in nanometers associated with each datapoint.

        Returns
        -------
        List[float]
            The list of float values encoding the wavelength, in nanometers, associated with each datapoint.
        """
        return self.__wavelength

    @property
    def absorbance(self) -> List[float]:
        """
        The absorbance values associated with each datapoint.

        Returns
        -------
        List[float]
            The list of float values encoding absorbance associated with each datapoint.
        """
        return self.__absorbance

    @property
    def transmittance(self) -> List[float]:
        """
        The transmittance, expressed as a percentage value, associated with each datapoint.

        Returns
        -------
        List[float]
            The list of float values encoding transmittance associated with each datapoint.
        """
        return [10 ** (2 - A) for A in self.__absorbance]

    @property
    def pitch(self) -> float:
        """
        The wavelength separation existing between two subsequent data-points.

        Returns
        -------
        float
            The wavelength separation in nanometers existing between to subsequent datapoints. If the points are not
            equally spaced a warning will be raised and the average pitch value will be returned
        """
        delta = [self.wavelength[i] - self.wavelength[i - 1] for i in range(1, len(self))]

        if any([value != delta[0] for value in delta]):
            warn("The spectral points are not equally spaced. The pitch property will return the average pitch value.")
            return abs(sum(delta) / len(delta))

        else:
            return abs(delta[0])

    @classmethod
    def from_JASCO_ASCII(cls, path: str) -> UVVisSpectrum:
        """
        The classmethod designed to parse ASCII data files generated by JASCO instruments.

        Arguments
        ---------
        path: str
            The path to the ASCII file encoding the experimental measurements.

        Raises:
        -------
        RuntimeError
            Exception raised if the file does not exist or if it cannot be properly parsed.
        """

        if not isfile(path):
            raise RuntimeError(f"The specified file '{path}' does not exist.")

        obj = cls()

        with open(path, "r") as file:
            npt, xunits, yunits = None, None, None

            for line in file:
                if "TITLE" in line:
                    obj.title = line.split("\t")[-1].strip("\n")

                if "SPECTROMETER" in line:
                    obj.instrument = line.split("\t")[-1].strip("\n")

                if "XUNITS" in line:
                    xunits = line.split("\t")[-1].strip("\n")

                if "YUNITS" in line:
                    yunits = line.split("\t")[-1].strip("\n")

                if "DATE" in line:
                    date_str = line.split("\t")[-1].strip("\n")
                    date_str += " "
                    date_str += file.readline().split("\t")[-1].strip("\n")
                    obj.__timestamp = datetime.strptime(date_str, "%y/%m/%d %H:%M:%S")

                if "NPOINTS" in line:
                    npt = int(line.split("\t")[-1])

                if "XYDATA" in line:
                    for _ in range(npt):
                        data = file.readline().split("\t")

                        xvalue, yvalue = float(data[0]), float(data[1])

                        if xunits == "NANOMETERS":
                            obj.__wavelength.append(xvalue)
                        else:
                            raise RuntimeError(f"Cannot parse unit {xunits}.")

                        if yunits == "ABSORBANCE":
                            obj.__absorbance.append(yvalue)
                        else:
                            raise RuntimeError(f"Cannot parse unit {yunits}.")

        obj.__sort()
        return obj

    def __getitem__(self, i: int) -> Tuple[float, float]:
        if i < 0 or i >= len(self):
            raise ValueError("Index out of bounds")

        return self.__wavelength[i], self.__absorbance[i]

    def __iter__(self) -> Tuple[float, float]:
        for w, a in zip(self.__wavelength, self.__absorbance):
            yield w, a

    def __len__(self) -> int:
        return len(self.__wavelength)

    def __sort(self) -> None:
        """
        Function sorting the spectal datapoint from the smaller wavelength to the highest.
        """
        self.__absorbance = [A for _, A in sorted(zip(self.__wavelength, self.__absorbance), key=lambda pair: pair[0])]
        self.__wavelength.sort()

    def __check_binary_operation(self, obj: UVVisSpectrum) -> None:
        """
        The function checks whether a given UVVisSpectrum object can be used in a binary operation involving
        this UVVisSpectum.

        Raises
        ------
        RuntimeError
            Exception raised if the lenght of the spectrum objects, their pitch or their spectral range is different.
        """
        if len(self) != len(obj):
            raise RuntimeError("Cannot perform binary operation between spectra of different lengths.")

        if self.pitch != obj.pitch:
            raise RuntimeError("Cannot perform binary operation between spectra of different pitch.")

        if not all([w1 == w2 for w1, w2 in zip(self.wavelength, obj.wavelength)]):
            raise RuntimeError("Cannot perform binary operation between spectra with different wavelength ranges")

    def __add__(self, obj: UVVisSpectrum) -> UVVisSpectrum:
        self.__check_binary_operation(obj)

        result = UVVisSpectrum()
        result.title = self.title + " + " + obj.title
        result.__timestamp = None
        result.__wavelength = self.__wavelength
        result.__absorbance = [A1 + A2 for A1, A2 in zip(self.__absorbance, obj.__absorbance)]

        return result

    def __sub__(self, obj: UVVisSpectrum) -> UVVisSpectrum:
        self.__check_binary_operation(obj)

        result = UVVisSpectrum()
        result.title = self.title + " - " + obj.title
        result.__timestamp = None
        result.__wavelength = self.__wavelength
        result.__absorbance = [A1 - A2 for A1, A2 in zip(self.__absorbance, obj.__absorbance)]

        return result

    def __mul__(self, obj: UVVisSpectrum) -> UVVisSpectrum:
        self.__check_binary_operation(obj)

        result = UVVisSpectrum()
        result.title = self.title + " * " + obj.title
        result.__timestamp = None
        result.__wavelength = self.__wavelength
        result.__absorbance = [A1 * A2 for A1, A2 in zip(self.__absorbance, obj.__absorbance)]

        return result

    def __truediv__(self, obj: UVVisSpectrum) -> UVVisSpectrum:
        self.__check_binary_operation(obj)

        result = UVVisSpectrum()
        result.title = self.title + " / " + obj.title
        result.__timestamp = None
        result.__wavelength = self.__wavelength
        result.__absorbance = [A1 / A2 for A1, A2 in zip(self.__absorbance, obj.__absorbance)]

        return result

    def scale(self, value: float, inplace: bool = False) -> Optional[UVVisSpectrum]:
        """
        Scale the absorbance of the spectrum according to a float scalar value.

        Arguments
        ---------
        value: float
            The factor to be used in the multiplication
        inplace: bool
            If set to False (default) will return the spectrum scaled by the specified value. Else, it will update the
            absorbance list of the current spectrum with the updated values.

        Returns
        -------
        Optional[UVVisSpectrum]
            If inplace is set to False (default) will return the scaled spectrum.
        """

        if inplace is True:
            self.title = f"{value}*{self.title}"
            self.__absorbance = [value * A for A in self.__absorbance]

        else:
            result = UVVisSpectrum()
            result.title = f"{value}*{self.title}"
            result.__timestamp = self.__timestamp
            result.__wavelength = self.__wavelength
            result.__absorbance = [value * A for A in self.__absorbance]

            return result

    def subspectrum(self, lower: float, upper: float) -> UVVisSpectrum:
        """
        Generate from the spectrum a new UVVisSpectrum object containing only the datapoints between the user specified
        lower and upper limits.

        Arguments
        ---------
        lower: float
            The lower limit of the new spectrum (included if available)
        upper: float
            The upper limit of the new spectrum (included if available)

        Raises
        ------
        ValueError
            Exception raised if the lower limit is not smaller than the upper one.

        Returns
        -------
        UVVisSpectrum
            The spectrum object containing the experimental point withing the user specified wavelength bounaries.
        """
        if lower >= upper:
            raise ValueError("The lower limit must be smaller than the upper one.")

        subspectrum = UVVisSpectrum()
        subspectrum.title = self.title
        subspectrum.instrument = self.instrument
        subspectrum.__timestamp = self.__timestamp

        for wavelength, absorbance in self:
            if wavelength >= lower and wavelength <= upper:
                subspectrum.__wavelength.append(wavelength)
                subspectrum.__absorbance.append(absorbance)

        return subspectrum

    def interpolate(self, k: int = 3) -> BSpline:
        """
        Compute a k-th order interpolating B-spline for the spectrum.

        Arguments
        ---------
        k: int
            The B-spline degree (default: cubic B-spline k=3)

        Returns
        -------
        BSpline
            The interpolating B-spline function as a `scipy.interpolate.BSpline` object.
        """
        bspline = make_interp_spline(self.__wavelength, self.__absorbance, k=k)
        return bspline

    def resample(self, lower: float, upper: float, pitch: float, k: int = 3) -> UVVisSpectrum:
        """
        Using a B-spline interpolation, generate a new UVVisSpectum object responding the the user specified requirements.

        Arguments
        ---------
        lower: float
            The minimum value of the wavelength scale in nanometers.
        upper: float
            The maximum value of the wavelength scale in nanometers.
        pitch: float
            The value of the wavelength sparation, in nanometers, existing between two subsequent datapoints.
        k: int
            The B-spline degree used during the resample (default: cubic B-spline k=3)

        Raises
        ------
        ValueError
            Exception raised if the `lower` is bigger or equal to the `upper` value or if `pitch` is invalid.
        RuntimeError
            Exception raised if the `lower` and `upper` values specified by the user are outside the original spectum range.

        Returns
        -------
        UVVisSpectrum
            The UVVisSpectrum generated according to the user specified set of parameters.
        """
        if lower >= upper:
            raise ValueError("The lower wavelength value must be lower than the upper one.")

        if pitch <= 0:
            raise ValueError("The pitch value must be a positive float.")

        if lower < min(self.wavelength) or upper > max(self.wavelength):
            raise RuntimeError("Cannot perform interpolation. The required range exceed the range of available data.")

        spectrum = UVVisSpectrum()
        spectrum.title = self.title + " (resampled)"
        spectrum.instrument = f"B-spline (k={k}) interpolation <- original: {self.instrument}"
        spectrum.__timestamp = self.__timestamp

        spectrum.__wavelength = np.arange(start=lower, stop=upper + pitch, step=pitch)

        bspline = self.interpolate(k=k)
        spectrum.__absorbance = bspline(spectrum.__wavelength)
        spectrum.__sort()

        return spectrum


def plot_spectrum(
    spectra: Union[List[UVVisSpectrum], UVVisSpectrum],
    transmittance: bool = False,
    xrange: Optional[Tuple[float, float]] = None,
    yrange: Optional[Tuple[float, float]] = None,
    figsize: Tuple[float, float] = (12.0, 8.0),
    savepath: Optional[str] = None,
    show: bool = True,
):
    """
    Function capable of plotting one or more spectra.

    Arguments
    ---------
    spectra: Union[List[UVVisSpectrum], UVVisSpectrum]
        The spectrum or the list of spectra to plot.
    transmittance: bool
        If set to True wil switch the y-axis to transmittance mode.
    xrange: Optional[Tuple[float, float]]
        The range of values to be shown on the x-axis. The meaning of the values depends on the display mode selected.
    yrange: Optional[Tuple[float, float]]
        The range of values to be shown on the y-axis. The meaning of the values depends on the display mode selected.
    figsize: Tuple[float, float]
        The size of the matplotlib figure to be used in plotting the spectrum. (default: (12, 8))
    savepath: Optional[str]
        If set to a value different from None, will specify the path of the file to be saved.
    show: bool
        If set to True (default) will show an interactive window where the plot is displayed.
    """
    plt.rc("font", **{"size": 18})

    fig = plt.figure(figsize=figsize)

    spectra: List[UVVisSpectrum] = [spectra] if type(spectra) == UVVisSpectrum else spectra

    for spectrum in spectra:
        plt.plot(
            spectrum.wavelength,
            spectrum.transmittance if transmittance else spectrum.absorbance,
            label=f"{spectrum.title}",
        )

    if xrange is not None:
        plt.xlim(xrange)

    if yrange is not None:
        plt.ylim(yrange)

    plt.xlabel("Wavelength [nm]", size=22)
    plt.ylabel(r"Transmittance [$\%$]" if transmittance else "Absorbance [a.u.]", size=22)

    plt.grid(which="major", c="#DDDDDD")
    plt.grid(which="minor", c="#EEEEEE")

    plt.legend()

    plt.tight_layout()

    if savepath is not None:
        plt.savefig(savepath, dpi=600)

    if show is True:
        plt.show()
