__author__ = 'Yuri E. Corilo'
__date__ = 'Jul 02, 2019'

import dataclasses
import os
from typing import List, Dict

from corems.encapsulation.constant import Atoms, Labels

@dataclasses.dataclass
class TransientSetting:
    '''
    Transient processing settings class
    Attributes
    ----------
    implemented_apodization_function : tuple
        Available apodization functions
    apodization_method : str
        Apodization function to use. Hanning is a good default. For absorption mode processing, Half-Sine or Half-Kaiser may be more appropriate.
    number_of_truncations : int
        How many times to truncate the transient prior to Fourier transform
    number_of_zero_fills : int
        How many times to zero fille the transient prior to Fourier transform.
    next_power_of_two : bool
        If True, zero fill to the next power of two after the new length of len(transient)+(number_of_zero_fills*len(transient)). 
    kaiser_beta : float
        Beta parameter for Kaiser or Half-Kaiser apodisation function. 0 is rectangular,  5 is similar to Hamming,
        6 is similar to hanning, and 8.6 is similar to Blackman (from numpy docs)

    '''
    implemented_apodization_function: tuple = ('Hamming', 'Hanning', 'Blackman','Full-Sine','Half-Sine','Kaiser','Half-Kaiser')
    apodization_method: str = 'Hanning'
    number_of_truncations: int = 0
    number_of_zero_fills: int = 1
    next_power_of_two: bool = False
    kaiser_beta: float = 8.6

    def __post_init__(self):

        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

@dataclasses.dataclass
class DataInputSetting:

    # add to this dict the VALUES to match your labels, THE ORDER WON"T MATTER
    # "column_translate" : {"m/z":"m/z", "Resolving Power":"Resolving Power", "Abundance":"Abundance" , "S/N":"S/N"}
    header_translate: dict = dataclasses.field(default_factory=dict)

    def __post_init__(self):

        self.header_translate = {'m/z': Labels.mz, 
                                 'mOz': Labels.mz,
                                 'Mass': Labels.mz,
                                 'Resolving Power': 'Resolving Power',
                                 'Res.': Labels.rp,
                                 'Intensity': Labels.abundance,
                                 'I': Labels.abundance,
                                 'Abundance': Labels.abundance,
                                 'Signal/Noise': Labels.s2n,
                                 'S/N': Labels.s2n,
                                 'abs_abu': Labels.abundance,
                                 'sn': Labels.s2n,
                                 'resolution': Labels.rp}

    def add_mz_label(self, label):

        self.header_translate[label] = Labels.mz

    def add_peak_height_label(self, label):

        self.header_translate[label] = Labels.abundance

    def add_sn_label(self, label):

        self.header_translate[label] = Labels.s2n

    def add_resolving_power_label(self, label):

        self.header_translate[label] = Labels.rp

@dataclasses.dataclass
class LiquidChromatographSetting:
    '''
    enforce_target_ms2: bool
            only perform EIC for target_mz if the m/z was selected as precursor for ms2
    start_scan: int
        default -1 will select the lowest available
    end_scan: int
        default -1 will select the highest available
    peak_height_max_percent: float
        1-100 % used for baseline detection use 0.1 for second_derivative and 10 for other methods    
    peak_max_prominence_percent: float
        1-100 % used for baseline detection
    peak_height_min_percent: float
        0-100 % used for peak detection
    eic_signal_threshold: 
        0-100 % used for extracted ion chromatogram peak detection
    '''
    
    start_scan: int = -1
    
    end_scan: int = -1
        
    eic_tolerance_ppm: float = 5
    
    correct_eic_baseline = True

    enforce_target_ms2: bool = True

    average_target_mz: bool = True

    smooth_window: int = 5

    smooth_method: str = 'savgol'

    implemented_smooth_method: tuple = ('savgol', 'hanning', 'blackman', 'bartlett', 'flat', 'boxcar')

    savgol_pol_order: int = 2

    peak_height_max_percent: float = 10  

    peak_max_prominence_percent: float = 1

    peak_derivative_threshold:float = 0.0005

    min_peak_datapoints: float = 5

    noise_threshold_method: str = 'manual_relative_abundance'

    implemented_noise_threshold_methods: tuple = ('auto_relative_abundance', 'manual_relative_abundance', 'second_derivative')

    std_noise_threshold: int = 3

    peak_height_min_percent: float = 0.1  

    eic_signal_threshold: float = 0.01  

    def __post_init__(self):
        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

@dataclasses.dataclass
class MassSpectrumSetting:

    threshold_method: str = 'log'

    implemented_noise_threshold_methods: tuple = ('minima', 'signal_noise', 'relative_abundance', 'absolute_abundance', 'log')

    noise_threshold_std: int = 6

    s2n_threshold: float = 4

    relative_abundance_threshold: float = 6  # from 0-100

    absolute_abundance_threshold: float = 1_000_000

    log_nsigma: int = 6
    log_nsigma_corr_factor: float = 0.463 #mFT is 0.463, aFT is 1.0
    log_nsigma_bins: int = 500 # bins for the histogram for the noise

    min_noise_mz: float = 50
    max_noise_mz: float = 1200.0

    min_picking_mz: float = 50
    max_picking_mz: float = 1200.0

    calib_minimize_method: str = 'Powell'
    calib_pol_order: int = 2
    max_calib_ppm_error: float = 1.0
    min_calib_ppm_error: float = -1.0
    calib_sn_threshold: float = 2.0

    do_calibration: bool = True

    def __post_init__(self):
        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

@dataclasses.dataclass
class MassSpecPeakSetting:

    kendrick_base: Dict = dataclasses.field(default_factory=dict)
    
    kendrick_rounding_method: str = 'floor' # 'floor', 'ceil' or 'round' are valid methods for calculating nominal kendrick mass
    
    implemented_kendrick_rounding_methods : tuple = ('floor','ceil','round')

    # kendrick_base : Dict =  {'C': 1, 'H':2}
    peak_derivative_threshold: float =  0.0 # define derivative crossing threshould 0-1

    peak_min_prominence_percent: float = 0.1  # 1-100 % used for peak detection

    min_peak_datapoints: float = 5 # 0-inf used for peak detection

    peak_max_prominence_percent: float = 0.1  # 1-100 % used for baseline detection

    peak_height_max_percent: float = 10  # 1-100 % used for baseline detection

    legacy_resolving_power: bool = True # Use the legacy (CoreMS v1) resolving power calculation (True)
    
    def __post_init__(self):

        # default to CH2
        if not self.kendrick_base:
            self.kendrick_base = {'C': 1, 'H': 2}
        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

@dataclasses.dataclass
class GasChromatographSetting:

    use_deconvolution: bool = False

    implemented_smooth_method: tuple = ('savgol', 'hanning', 'blackman', 'bartlett', 'flat', 'boxcar')

    smooth_window: int = 5

    smooth_method: str = 'savgol'

    savgol_pol_order: int = 2

    peak_derivative_threshold:float = 0.0005

    peak_height_max_percent: float = 10  # 1-100 % used for baseline detection use 0.1 for second_derivative and 10 for other methods

    peak_max_prominence_percent: float = 1  # 1-100 % used for baseline detection

    min_peak_datapoints: float = 5

    max_peak_width: float = 0.1

    noise_threshold_method: str = 'manual_relative_abundance'

    implemented_noise_threshold_methods: tuple = ('auto_relative_abundance', 'manual_relative_abundance', 
                                                  'second_derivative')

    std_noise_threshold: int = 3

    peak_height_min_percent: float = 0.1  # 0-100 % used for peak detection

    peak_min_prominence_percent: float = 0.1  # 0-100 % used for peak detection

    eic_signal_threshold: float = 0.01  # 0-100 % used for extracted ion chromatogram peak detection

    max_rt_distance: float = 0.025  # minutes, max distance allowance hierarchical clutter

    def __post_init__(self):

        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

@dataclasses.dataclass
class CompoundSearchSettings:

    url_database: str = "postgresql+psycopg2://coremsappdb:coremsapppnnl@localhost:5432/lowres" # 'postgresql://postgres:labthomson0102@172.22.113.27:5432/GCMS' # 'sqlite:///db/pnnl_lowres_gcms_compounds.sqlite'

    ri_search_range: float = 35

    rt_search_range: float = 1.0  # used for retention index calibration

    correlation_threshold: float = 0.5  # used for calibration, spectral similarity 

    score_threshold: float = 0.0

    ri_spacing: float = 200

    ri_std: float = 3  # in standard deviation

    ri_calibration_compound_names: List = dataclasses.field(default_factory=list)

    # calculates and export all spectral similarity methods
    exploratory_mode: bool = False

    score_methods: tuple = ('highest_sim_score', 'highest_ss')

    output_score_method: str = 'All'

    

    def __post_init__(self):
        # enforce datatype
        self.url_database = os.getenv('SPECTRAL_GCMS_DATABASE_URL', 'sqlite:///db/pnnl_lowres_gcms_compounds.sqlite')

        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

        self.ri_calibration_compound_names = ['Methyl Caprylate', 
                                              'Methyl Caprate', 
                                              'Methyl Pelargonate', 
                                              'Methyl Laurate', 
                                              'Methyl Myristate', 
                                              'Methyl Palmitate', 
                                              'Methyl Stearate', 
                                              'Methyl Eicosanoate', 
                                              'Methyl Docosanoate', 
                                              'Methyl Linocerate', 
                                              'Methyl Hexacosanoate', 
                                              'Methyl Octacosanoate', 
                                              'Methyl Triacontanoate']

@dataclasses.dataclass
class MolecularLookupDictSettings:

    '''
    ### DO NOT CHANGE IT! These are used to generate the database entries 

    ### DO change when creating a new application database 

    ### FOR search settings runtime and database query check use the MolecularFormulaSearchSettings class below

    ### C, H, N, O, S and P atoms are ALWAYS needed at usedAtoms
    ### if you don't want to include one of those atoms set the max and min at 0
    ### you can include any atom listed at Atoms class inside encapsulation.settings.constants module
    ### make sure to include the selected covalence at the used_atoms_valences when adding new atoms 
    ### NOTE : Adducts atoms have zero covalence
    ### NOTE : Not using static variable because this class is distributed using multiprocessing
    '''
    def __init__(self):

        self.usedAtoms = {'C': (1, 90),
                          'H': (4, 200),
                          'O': (0, 12),
                          'N': (0, 0),
                          'S': (0, 0),
                          'P': (0, 0),
                          'Cl': (0, 0),
                          }

        self.min_mz = 50

        self.max_mz = 1200

        self.min_dbe = 0

        self.max_dbe = 50

        # overwrites the dbe limits above to DBE = (C + heteroatoms) * 0.9
        self.use_pah_line_rule = False

        self.isRadical = True

        self.isProtonated = True

        self.min_op_filter = 2

        self.min_hc_filter = 0.3

        self.min_oc_filter = 1.2

        self.max_oc_filter = 0.0

        self.url_database = None

        self.db_jobs = 1

        self.used_atom_valences = {'C': 4,
                                   '13C': 4,
                                   'H': 1,
                                   'O': 2,
                                   '18O': 2,
                                   'N': 3,
                                   'S': 2,
                                   '34S': 2,
                                   'P': 3,
                                   'Cl': 1,
                                   '37Cl': 1,
                                   'Br': 1,
                                   'Na': 1,
                                   'F': 1,
                                   'K': 0,
                                   }

@dataclasses.dataclass
class MolecularFormulaSearchSettings:

    use_isotopologue_filter: bool = False

    isotopologue_filter_threshold: float = 33

    isotopologue_filter_atoms: tuple = ('Cl', 'Br')

    use_runtime_kendrick_filter: bool = False

    use_min_peaks_filter: bool = True

    min_peaks_per_class: int = 15

    url_database: str = 'postgresql+psycopg2://coremsappdb:coremsapppnnl@localhost:5432/coremsapp'

    db_jobs: int = 3

    db_chunk_size: int = 300

    '''query setting'''
    ion_charge: int = -1

    min_hc_filter: float = 0.3

    max_hc_filter: float = 3

    min_oc_filter: float = 0.0

    max_oc_filter: float = 1.2

    min_op_filter: float = 2

    use_pah_line_rule: bool = False

    min_dbe: float = 0

    max_dbe: float = 40

    mz_error_score_weight: float = 0.6

    isotopologue_score_weight: float = 0.4

    # look for close shell ions [M + Adduct]+ only considers metal set in the list adduct_atoms  
    adduct_atoms_neg: tuple = ('Cl', 'Br')

    adduct_atoms_pos: tuple = ('Na', 'K')

    score_methods: tuple = ('S_P_lowest_error', 'N_S_P_lowest_error', 'lowest_error', 'prob_score',
                            'air_filter_error', 'water_filter_error', 'earth_filter_error')

    score_method: str = 'prob_score'

    output_min_score: float = 0.1

    output_score_method: str = 'All Candidates'

    output_min_score: float = 0.1

    # depending on the polarity mode it looks for [M].+ , [M].-
    # query and automatically compile add entry if it doesn't exist

    isRadical: bool = False

    # depending on the polarity mode it looks for [M + H]+ , [M - H]+
    # query and automatically compile and push options if it doesn't exist
    isProtonated: bool = True

    isAdduct: bool = False

    usedAtoms: dict = dataclasses.field(default_factory=dict)

    ''' search setting '''

    ionization_type: str = 'ESI'

    # empirically set / needs optimization
    min_ppm_error: float = -10.0  # ppm

    # empirically set / needs optimization
    max_ppm_error: float = 10.0  # ppm

    # empirically set / needs optimization set for isotopologue search
    min_abun_error: float = -100.0  # percentage

    # empirically set / needs optimization set for isotopologue search
    max_abun_error: float = 100.0  # percentage

    # empirically set / needs optimization
    mz_error_range: float = 1.5

    # 'distance', 'lowest', 'symmetrical','average' 'None'
    error_method: str = 'None'

    mz_error_average: float = 0.0

    # used_atom_valences: {'C': 4, 'H':1, etc} = dataclasses.field(default_factory=dict)
    used_atom_valences: dict = dataclasses.field(default_factory=dict)

    def __post_init__(self):

        self.url_database = os.getenv('COREMS_DATABASE_URL', 'sqlite:///db/molformula.db')
        # enforce datatype
        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            if not isinstance(value, field.type):

                value = field.type(value)
                setattr(self, field.name, value)

        # enforce C and H if either do not exists
        if 'C' not in self.usedAtoms.keys():
            self.usedAtoms['C'] = (1, 100)
        if 'H' not in self.usedAtoms.keys():
            self.usedAtoms['H'] = (1, 200)

        # add cummon values
        current_used_atoms = self.used_atom_valences.keys()
        for atom in Atoms.atoms_covalence.keys():
            if atom not in current_used_atoms:
                self.used_atom_valences[atom] = Atoms.atoms_covalence.get(atom)

        '''
        self.used_atom_valences.update({'13C': 4,
                                        '18O': 2,
                                        '34S': 2,
                                        'Cl': 1,
                                        '37Cl': 1,
                                        'Br': 1,
                                        'Na': 1,
                                        'F': 1,
                                        'K': 1,
                                        })
        '''

if __name__ == '__main__':
    a = DataInputSetting()
    print(a.__dict__)
