import numpy as np
import os
import xarray as xr

from .ts4cor import TS4COR
from .ts4cor_sigma import TS4COR_SIGMA


class ACSTSCor:
    """A wrapper class for the TS4COR and TS4COR_SIGMA dictionaries that are included with acspype."""

    def __init__(self):
        self.wavelengths = tuple(TS4COR.keys())

        self.psi_s_a = list(map(lambda x: x[1]['psi_s_a'], TS4COR.items()))
        self.psi_s_c = list(map(lambda x: x[1]['psi_s_c'], TS4COR.items()))
        self.psi_t = list(map(lambda x: x[1]['psi_t'], TS4COR.items()))

        self.sigma_psi_s_a= list(map(lambda x:x[1]['sigma_psi_s_a'], TS4COR_SIGMA.items()))
        self.sigma_psi_s_c= list(map(lambda x:x[1]['sigma_psi_s_c'], TS4COR_SIGMA.items()))
        self.sigma_psi_t= list(map(lambda x:x[1]['sigma_psi_t'], TS4COR_SIGMA.items()))

        self.method = 'Sullivan et al., 2006'

    def to_xarray(self):
        """
        Convert the TS4COR class attributes to an xarray.Dataset.

        :return: A dimensioned xarray.Dataset.
        """
        ds = xr.Dataset()
        ds = ds.assign_coords({'wavelength': np.array(self.wavelengths)})
        ds['psi_t'] = (['wavelength'], np.array(self.psi_t))
        ds['psi_s_c'] = (['wavelength'], np.array(self.psi_s_c))
        ds['psi_s_a'] = (['wavelength'], np.array(self.psi_s_a))
        ds['sigma_psi_t'] = (['wavelength'], np.array(self.sigma_psi_t))
        ds['sigma_psi_s_c'] = (['wavelength'], np.array(self.sigma_psi_s_c))
        ds['sigma_psi_s_a'] = (['wavelength'], np.array(self.sigma_psi_s_a))
        ds.attrs['tscor_data'] = self.method
        return ds


class ACSTS4CorReader:
    """
    A class for parsing ACS TS4.cor files and putting them into a format that is easier to work with for larger or
    multiple file datasets.

    This class can be used if users want to import the TS4.cor file themselves and don't want to use values
    interpolated by acspype.
    """

    def __init__(self, filepath: str) -> None:
        """
        Parse the .cor file and assign data as attributes.

        :param filepath: The filepath of the TS4.cor file.
        """

        self.filepath = os.path.normpath(filepath)
        self.__read_cor()
        self.__parse_lines()
        self.method = 'Sullivan et al., 2006'

    def __read_cor(self) -> None:
        """
        Read .cor file and store lines as a class attribute.

        :return: None
        """

        with open(self.filepath, 'r') as _file:
            self._lines = _file.readlines()

    def __parse_lines(self) -> None:
        """
        Parse the lines of the .cor file to get correction information.

        :return: None
        """

        wavelengths = []
        psi_t = []
        psi_s_c = []
        psi_s_a = []
        for line in self._lines:
            line_data = line.split('\t')
            line_data = [v.replace('\n', '') for v in line_data]
            line_data = [v.replace(' ', '') for v in line_data]
            if line_data == ['']:
                break
            line_data = [float(v) for v in line_data]
            wavelengths.append(line_data[0])
            psi_t.append(line_data[1])
            psi_s_c.append(line_data[2])
            psi_s_a.append(line_data[3])
        if len(wavelengths) != len(psi_t) != len(psi_s_c) != len(psi_s_a):
            raise ValueError('Mismatch in length of TS4cor file.')
        else:
            self.wavelengths = np.array(wavelengths)
            self.psi_t = np.array(psi_t)
            self.psi_s_c = np.array(psi_s_c)
            self.psi_s_a = np.array(psi_s_a)

    def to_xarray(self):
        """
        Convert the TS4COR class attributes to an xarray.Dataset.

        :return: A dimensioned xarray.Dataset.
        """
        ds = xr.Dataset()
        ds = ds.assign_coords({'wavelength': np.array(self.wavelengths)})
        ds['psi_t'] = (['wavelength'], np.array(self.psi_t))
        ds['psi_s_c'] = (['wavelength'], np.array(self.psi_s_c))
        ds['psi_s_a'] = (['wavelength'], np.array(self.psi_s_a))

        ds.attrs['tscor_data'] = self.method
        return ds
