# -*- coding: utf-8 -*-
"""
copyright 2024 Peak Design
this file is part of hsr1, which is distributed under the GNU Lesser General Public License v3 (LGPL)
"""
import importlib_resources

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from hsr1.utils.spectrum import spectrumUtils as SpectrumUtils
from hsr1.utils.spectrum.pixelSpectrum import PixelSpectrum
from hsr1.utils.spectrum.weightedSpectrum import WeightedSpectrum


class RawDataset:
    def __init__(self, dataset:pd.DataFrame, 
                 calibrations:dict=None, calibrations_filepath=None, 
                 amplitude_calibration=None, amplitude_calibration_filepath=None, 
                 spectrum_type="wavelength"):
        """stores a raw dataset containing all the channels of raw data
        
        params:
            dataset: dataframe with columns "pc_time_end_measurement", "channel_0", "channel_1" etc.
            calibrations: dict containing polynomials with keys "channel_0", "channel_1" etc.
            calibrations_filepath: filepath to the current wavelength calibration file
            amplitude_cal: an amplitude calibration to apply to found spectrum, only used for graphs
            amplitude_calibration_filepath
        """
        self.dataset = dataset
        self.calibrations = calibrations
        self.amplitude_calibration = amplitude_calibration
        self.spectrum_type = spectrum_type
        
        if spectrum_type != "wavelength" and spectrum_type != "pixel":
            raise ValueError("spectrum_type must be either \"wavelength\" or \"pixel\"")
        
        if calibrations is None and spectrum_type == "wavelength":
            if calibrations_filepath is None:
                raise ValueError("""original calibrations must be passed, either as a dict or as the location of the calibration file. If this is uncalibrated data, pass spectrum_type=\"pixel\"""")
            self.calibrations = SpectrumUtils.read_wavelength_polys(calibrations_filepath)
        
        if spectrum_type == "pixel":
            if calibrations is not None or calibrations_filepath is not None:
                print("there is no point passing existing calibrations to a pixel dataset, pixel readings are assumed to be uncalibrated. \
                      \n if this is a calibration that is used on this spectrometer, you can change the default_calibration parameter in find_all_calibrations.")
        
        if amplitude_calibration is None and amplitude_calibration_filepath is not None:
            self.amplitude_calibration = SpectrumUtils.read_amplitude_file(amplitude_calibration_filepath)
        
        if amplitude_calibration is None and amplitude_calibration_filepath is None:
            self.amplitude_calibration = np.ones(801)
        
    
    
    def plot_calibration(self, polys, names=None):
        """plots a line graph of pixel against wavelength for each channel
        params:
            polys: list of lists or dict containing the polynommials for each channel's new calibration
            names: list of names to display on the legend
        """
        ps = PixelSpectrum()
        
        
        if isinstance(polys, dict):
            poly_names = []
            poly_list = []
            for poly in polys:
                poly_list.append(polys[poly])
                poly_names.append(poly)
            polys = poly_list 
            names = poly_names
        
        if names is None:
            names = [""]*len(polys)
        
        plt.figure(figsize=(20,10))
        for poly, name in zip(polys, names):
            plt.plot(ps.apply_wavelength_calibration(np.arange(0, 1200), poly, inverse=True), np.arange(300, 1101), linewidth=0.5, label=name + ": "+SpectrumUtils.poly_to_string(poly))
        plt.xlim(0, 1200)
        plt.suptitle("individual channel calibrations")
        plt.xlabel("pixels")
        plt.ylabel("wavelength")
        plt.legend()
        plt.show()
    
    def plot_calibration_wide(self, polys, names=None, title=""):
        """plots a line graph of pixel against wavelength for each channel
        params:
            polys: list of lists containing the polynommials for each channel's new calibration
            names: list of names to display on the legend
        """
        if names is None:
            names = [""]*len(polys)
        
        plt.figure(figsize=(20,10))
        for poly, name in zip(polys, names):
            x = np.linspace(-10000, 10000, 1000)
            y = [SpectrumUtils.apply_polynomial(i, poly[0], poly[1], poly[2], poly[3]) for i in x]
            plt.plot(x, y, label=name + ": " + SpectrumUtils.poly_to_string(poly))

        plt.axvline(0)
        plt.axvline(1200)
        plt.suptitle(title)
        plt.legend()
        plt.show()
    
    def plot_all_errors(self, polys, dips, names):
        plt.figure(figsize=(20, 10))
        for i, poly, _dips, name in zip(list(range(len(polys))), polys, dips, names):
            measured = _dips[0]
            ref = _dips[1]
            error = measured-ref
            
            plt.plot(ref, error-(i/50), label=name)
        
        plt.legend()
        plt.show()
    
    
    def find_all_calibrations(self, file_output=None, 
                              reference_spectrum=None, reference_filepath=None, 
                              plot_total_output=True, plot_individual_results=False, plot_debug=False, 
                              **kwargs):
        """loops through each channel in the dataset and finds the correct calibration
        params:
            file_output: for outputting results, which text file to use. leave as None to not output
            reference_spectrum: a spectrum in wavelength space that the measured data will be matched to
            plot_total_output: plots the calibration for each channel on one graph
            plot_individual_results: plots the results from each channel, passes plot_results to each calibration
            plot_debug: plots the result from each segment of each channel, passed to each calibration
            kwargs: arguments to pass to each channel's find_calibration method, cleaner than having all arguments as parameters here
        """
        if reference_filepath is None:
            res = importlib_resources.files("hsr1.data").joinpath("smarts.txt")
            file = importlib_resources.as_file(res)
            with file as f:
                reference_filepath = f
            
        
        
        if reference_spectrum is None:
            try:
                reference_spectrum = SpectrumUtils.read_simple_file(reference_filepath)
            except ValueError:
                reference_spectrum = SpectrumUtils.read_reference_file(reference_filepath)
        
        
        all_dips = []
        polys = []
        names = []
        new_calibrations = {}
        for channel_name in self.dataset.columns[1:]:
            
            
            ##### extract the polynomial that has already been applied to this channel
            applied_poly = None
            if self.spectrum_type == "wavelength":
                if channel_name in self.calibrations:
                    applied_poly = self.calibrations[channel_name]
                else:
                    print("error while calibrating "+channel_name+": couldnt find an existing calibration for this channel")
                    continue
            
            ps = PixelSpectrum(cal_poly=applied_poly)
            
            channel = self.dataset[["pc_time_end_measurement", channel_name]].copy()
            channel = channel.rename(columns={channel_name: "global_spectrum"})
            
            
            
            ##### calculate one pixel spectrum, for graphing
            pixel_spectrum = None
            if self.spectrum_type == "wavelength":
                pixel_data = ps.calculate_from_wavelengths(channel, ["global_spectrum"])
                max_index = 0
                if len(channel) > 1:
                    integrals = pixel_data["global_spectrum"].apply(np.sum)
                    max_index = np.argmax(integrals)
                pixel_spectrum = pixel_data.loc[max_index, "global_spectrum"]
            elif self.spectrum_type == "pixel":
                max_index = 0
                if len(channel) > 1:
                    integrals = channel.apply(np.sum)
                    max_index = np.argmax(integrals)
                pixel_spectrum = channel.loc[max_index, "global_spectrum"]
            
            if max(pixel_spectrum) == 0:
                print("invalid spectrum selected for graphing")
            
            ##### calculate the weighted spectrum and calibrate
            ws = None
            if self.spectrum_type == "wavelength":
                ws = WeightedSpectrum(wavelength_spectra=channel, reference_spectrum=reference_spectrum, 
                                      applied_poly=applied_poly, 
                                      filter_low_values=False, 
                                      amplitude_calibration=self.amplitude_calibration, 
                                      pixel_spectrum=pixel_spectrum, 
                                      name=channel_name, **kwargs)
            elif self.spectrum_type == "pixel":
                ws = WeightedSpectrum(pixel_spectra=channel, reference_spectrum=reference_spectrum, 
                                      applied_poly=applied_poly, 
                                      filter_low_values=False, 
                                      amplitude_calibration=self.amplitude_calibration, 
                                      pixel_spectrum=pixel_spectrum, 
                                      name=channel_name, **kwargs)
                
            wavelength_cal, dips = (None, None)
            if plot_total_output:
                wavelength_cal, dips = ws.find_calibration(plot_debug=plot_debug, plot_results=plot_individual_results, 
                                                            return_dips=plot_total_output, 
                                                            **kwargs)
            else:
                wavelength_cal = ws.find_calibration(plot_debug=plot_debug, plot_results=plot_individual_results, 
                                                            return_dips=plot_total_output, 
                                                            **kwargs)
            if wavelength_cal is not None:
                all_dips.append(dips)
                new_calibrations[channel_name] = wavelength_cal
                polys.append(wavelength_cal)
                names.append(channel_name)
        
        if plot_total_output:
            self.plot_calibration(polys, names)
            self.plot_calibration_wide(polys, names)
            self.plot_all_errors(polys, all_dips, names)
        
        if file_output is not None:
            SpectrumUtils.format_output(new_calibrations, file_output)
        
        return new_calibrations