import numpy as np
from pathlib import PosixPath
from pydantic import (BaseModel,
                      confloat,
                      conint,
                      conlist,
                      constr,
                      FilePath,
                      validator)
from scipy.interpolate import interp1d
from typing import Optional


class MCACeriaCalibrationConfig(BaseModel):
    '''Class representing metadata required to perform a Ceria calibration for an
    MCA detector.

    :ivar spec_file: Path to the SPEC file containing the CeO2 scan
    :ivar scan_number: Number of the CeO2 scan in `spec_file`
    :ivar scan_step_index: Index of the scan step to use for calibration,
        optional. If not specified, the calibration routine will be performed on
        the average of all MCA spectra for the scan.

    :ivar flux_file: csv file containing station beam energy in eV (column 0)
        and flux (column 1)

    :ivar detector_name: name of the MCA to calibrate
    :ivar num_bins: number of channels on the MCA to calibrate
    :ivar max_energy_kev: maximum channel energy of the MCA in keV

    :ivar hexrd_h5_material_file: path to a HEXRD materials.h5 file containing an
        entry for the material properties.
    :ivar hexrd_h5_material_name: Name of the material entry in
        `hexrd_h5_material_file`, defaults to `'CeO2'`.
    :ivar lattice_parameter_angstrom: lattice spacing in angstrom to use for
        the cubic CeO2 crystal, defaults to `5.41153`.

    :ivar tth_max: detector rotation about hutch x axis, defaults to `90`.
    :ivar hkl_tth_tol: minimum resolvable difference in 2&theta between two
        unique HKL peaks, defaults to `0.15`.

    :ivar fit_include_bin_ranges: list of MCA channel index ranges whose data
        will be included in the calibration routine
    :ivar fit_hkls: list of unique HKL indices to fit peaks for in the
        calibration routine

    :ivar tth_initial_guess: initial guess for 2&theta
    :ivar slope_initial_guess: initial guess for detector channel energy
        correction linear slope, defaults to `1.0`.
    :ivar intercept_initial_guess: initial guess for detector channel energy
        correction y-intercept, defaults to `0.0`.

    :ivar tth_calibrated: calibrated value for 2&theta, defaults to None
    :ivar slope_calibrated: calibrated value for detector channel energy
        correction linear slope, defaults to `None`
    :ivar intercept_calibrated: calibrated value for detector channel energy
        correction y-intercept, defaluts to None

    :ivar max_iter: maximum number of iterations of the calibration routine,
        defaults to `10`.
    :ivar tune_tth_tol: stop iteratively tuning 2&theta when an iteration
        produces a change in the tuned value of 2&theta that is smaller than this
        value, defaults to `1e-8`.
    '''

    spec_file: FilePath
    scan_number: conint(gt=0)
    scan_step_index: Optional[conint(ge=0)]

    flux_file: FilePath

    detector_name: constr(strip_whitespace=True, min_length=1)
    num_bins: conint(gt=0)
    max_energy_kev: confloat(gt=0)

    hexrd_h5_material_file: FilePath
    hexrd_h5_material_name: constr(strip_whitespace=True, min_length=1) = 'CeO2'
    lattice_parameter_angstrom: confloat(gt=0) = 5.41153

    tth_max: confloat(gt=0, allow_inf_nan=False) = 90.0
    hkl_tth_tol: confloat(gt=0, allow_inf_nan=False) = 0.15

    fit_include_bin_ranges: conlist(min_items=1,
                                    item_type=conlist(item_type=conint(ge=0),
                                                      min_items=2,
                                                      max_items=2))
    fit_hkls: conlist(item_type=conint(ge=0), min_items=1)

    tth_initial_guess: confloat(gt=0, le=tth_max, allow_inf_nan=False)
    slope_initial_guess: float = 1.0
    intercept_initial_guess: float = 0.0
    tth_calibrated: Optional[confloat(gt=0, allow_inf_nan=False)]
    slope_calibrated: Optional[confloat(allow_inf_nan=False)]
    intercept_calibrated: Optional[confloat(allow_inf_nan=False)]

    max_iter: conint(gt=0) = 10
    tune_tth_tol: confloat(ge=0) = 1e-8

    @validator('fit_include_bin_ranges', each_item=True)
    def validate_include_bin_range(cls, value, values):
        '''Ensure no bin ranges are outside the boundary of the detector'''

        num_bins = values.get('num_bins')
        value[1] = min(value[1], num_bins)
        return(value)

    def mca_data(self):
        '''Get the 1D array of MCA data to use for calibration.
        
        :return: MCA data
        :rtype: np.ndarray
        '''

        from CHAP.common.utils.scanparsers import SMBMCAScanParser as ScanParser
        scanparser = ScanParser(self.spec_file, self.scan_number)
        if self.scan_step_index is None:
            data = scanparser.get_all_detector_data(self.detector_name)
            if scanparser.spec_scan_npts > 1:
                data = np.average(data, axis=1)
            else:
                data = data[0]
        else:
            data = scanparser.get_detector_data(self.detector_name, self.scan_step_index)

        return(np.array(data))

    def mca_mask(self):
        '''Get a boolean mask array to use on MCA data before fitting.

        :return: boolean mask array
        :rtype: numpy.ndarray
        '''

        mask = np.asarray([False]*self.num_bins)
        bin_indices = np.arange(self.num_bins)
        for min_, max_ in self.fit_include_bin_ranges:
            _mask =  np.logical_and(bin_indices > min_, bin_indices < max_)
            mask = np.logical_or(mask, _mask)

        return(mask)

    def flux_correction_interpolation_function(self):
        '''Get an interpolation function to correct MCA data for relative energy
        flux of the incident beam.

        :return: energy flux correction interpolation function
        :rtype: scipy.interpolate._polyint._Interpolator1D
        '''

        flux = np.loadtxt(self.flux_file)
        energies = flux[:,0]/1.e3
        relative_intensities = flux[:,1]/np.max(flux[:,1])
        interpolation_function = interp1d(energies, relative_intensities)
        return(interpolation_function)

    def material(self):
        '''Get CeO2 as a `CHAP.common.utils.Material` object.

        :return: CeO2 material
        :rtype: CHAP.common.utils.Material
        '''

        from CHAP.common.utils import Material
        material = Material(material_name=self.hexrd_h5_material_name,
                            material_file=self.hexrd_h5_material_file,
                            lattice_parameters_angstroms=self.lattice_parameter_angstrom)
        # The following kwargs will be needed if we allow the material to be
        # built using xrayutilities (for now, we only allow hexrd to make the
        # material):
                            # sgnum=225,
                            # atoms=['Ce4p', 'O2mdot'],
                            # pos=[(0.,0.,0.), (0.25,0.75,0.75)],
                            # enrgy=50000.) # Why do we need to specify an energy to get HKLs when using xrayutilities?
        return(material)

    def unique_ds(self):
        '''Get a list of unique HKLs and their lattice spacings
        
        :return: unique HKLs and their lattice spacings in angstroms
        :rtype: np.ndarray, np.ndarray
        '''
        
        unique_hkls, unique_ds = self.material().get_unique_ds(tth_tol=self.hkl_tth_tol, tth_max=self.tth_max)

        return(unique_hkls, unique_ds)

    def fit_ds(self):
        '''Get a list of HKLs and their lattice spacings that will be fit in the
        calibration routine
        
        :return: HKLs to fit and their lattice spacings in angstroms
        :rtype: np.ndarray, np.ndarray
        '''
        
        unique_hkls, unique_ds = self.unique_ds()

        fit_hkls = np.array([unique_hkls[i] for i in self.fit_hkls])
        fit_ds = np.array([unique_ds[i] for i in self.fit_hkls])

        return(fit_hkls, fit_ds)

    def dict(self):
        '''Return a representation of this configuration in a dictionary that is
        suitable for dumping to a YAML file (one that converts all instances of
        fields with type `PosixPath` to `str`).

        :return: dictionary representation of the configuration.
        :rtype: dict
        '''

        d = super().dict()
        for k,v in d.items():
            if isinstance(v, PosixPath):
                d[k] = str(v)
        return(d)
