'''
General data models for metabolomics.
Simple is good, and complex and specialized classes can be derived from these.

A mass spectrum is a list of masses. 
MS1 is direct scan of ions to generate intensity values.
MS^2 and MS^n refers to spectrum to measure fragmentation products of a precursor ion from lower MS level.

LC-MS is a composite of many spectra, each being a scan at a specific elution time. 
A peak can refer to either peak on the m/z axis, or peak on the LC axis.
The former is less common now as high-resolution mass spectrometry data 
require little attention in identifying (centroiding) the m/z peaks.
A peak on the LC axis is a chromatographic peak (elution peak) in LC-MS, 
but people may use 3-D detection or higher dimension e.g. IM-LC-MS.

Empirical compound is a key concept proposed by Shuzhao Li lab to connect 
experimental measurements and database records, including metabolic models.

Not all concepts have to be explicitly modeled in a project (e.g. expt, peak, network).
Use derived/inherited classes for more explict or specialized data.

We try to be explicit in source code, and Python supports introspection.
Therefore, getters and setters are avoided.
A serialize function is made available for easy JSON export. 
For simple cases, these classes can be simplified as Python NamedTuples, 
which allow easy access to attributes but have no additional methods.
'''

#
# Experimental concepts: experiment, peak, feature, empirical compound; massTrace, MSnSpectrum
# only considering mass spec not NMR data here
#

class Study:
    '''
    A study can include multiple experiments and datasets by different methods.
    '''
    def __init__(self, id=''):
        '''
        Equivalent to a study found in public repositories, e.g. MetaboLights or Metabolomics Workbench.
        '''
        self.id = id
        self.url = ''
        self.time_retrieval = None

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'url': self.url,
                'time_retrieval': self.time_retrieval,
                }

class Experiment:
    '''
    An experiment of LC-MS, LC-MS/MS, GC-MS, LC-IMS, etc.
    An experiment can include multiple methods and various types of samples,
    which may require separate data processing.

    The XCMSnExp class in the XCMS R package, or MSExperiment in the OpenMS software,
    may refer to a subset of experiment by the same method.
    We can use `monoExperiment` to refer a subset of experiment by the same method.
    Then the `Experiment` contains multiple monoExperiments.

    `Method` has its own class only if needed.
    In most cases, it's simpler to have some str references in `Experiment` or `monoExperiment`.

    Measurement data are attached to an Experiment, 
    in the form of a list of features and 
    a list of empCpds (the latter generated by annotation).
    Flexibility is given by any type of data can be attached.
    No need to be extensive when pre-processing is not the focus.

    For LC-MS, the feature-level data is a DataFrame,
    features in rows and observations (samples) in columns, similar to gene express data matrix.
    On disk, the data can follow the convention of ANNdata and HiCoNet,
    the 3-file-society Data Strucutre: DataMatrix, FeatureAnnotation and ObservationAnnotation. 
    The DataMatrix in file format uses a single row for observation IDs and a single column for feature IDs.
    Ref: https://github.com/shuzhao-li/hiconet

    The empCpd-level data can be in JSON or other formats.
    '''
    def __init__(self, id=''):
        '''
        Try to use long str to be unique in the world.
        Additional fields can be added to the dictionaries.
        '''
        self.id = id                    # e.g. 'EXP00001234'
        self.parent_study = ''
        self.number_samples = None
        self.species = 'hsa'
        self.tissue = 'plasma'

        self.provenance = {
            'generated_time': '',       # day of experiment
            'generated_by': '',         # operator of experiment
            'input_filename': '',
            'preprocess_software': '',
            'preprocess_parameters': {
                'version': '0.0',
                'ppm': 1,
                'SNR': 10,
            },
        }
        self.instrumentation = {
            'type': 'LC-MS',
            'spectrometer': '',         # mass spectrometer model
            'method_file': '',          # method file used
            'ionization': 'pos',        # positive or negative ionization
        }
        self.chromatography = {
            'system': '',               # chromatograph model
            'total_time': 300,          # seconds
            'method_file': '',
            'column_model': '',
            'column_length': '',
        }
        
        # data 
        self.feature_DataFrame = None
        self.FeatureAnnotation = {} 
        self.ObservationAnnotation = {
            'sample_list': [],
            'file_sample_mapper': {}
        }
        # immutable ordered version of sample_list
        self.ordered_samples = ()

        # EmpiricalCompounds, after annotation
        self.List_of_empCpds = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'study_id': self.parent_study,
                'number_samples': self.number_samples,
                'species': self.species,
                'tissue': self.tissue,
                'provenance': self.provenance,
                'instrumentation': self.instrumentation,
                'chromatography': self.chromatography,
                }
        

class Method:
    '''
    A study can include multiple experiments and datasets by different methods.
    '''
    def __init__(self, id=''):
        '''
        Equivalent to a study found in public repositories, e.g. MetaboLights or Metabolomics Workbench.
        '''
        self.id = id
        self.url = ''
        self.citation = ''
        self.instrumentation = {
            'type': 'LC-MS',
            'spectrometer': '',         # mass spectrometer model
            'method_file': '',          # method file used
            'ionization': 'pos',        # positive or negative ionization
        }
        self.chromatography = {
            'system': '',               # chromatograph model
            'total_time': 300,          # seconds
            'method_file': '',
            'column_model': '',
            'column_length': '',
        }

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'instrumentation': self.instrumentation,
                'chromatography': self.chromatography,
                }


class Sample:
    '''
    This is an analytical sample subjected to a single metabolomic method and associated with a single data file.
    Elution peaks are defined at sample level. A feature is defined at experiment level.
    One can extend to many attributes via registry dictionary, e.g.
        self.data_location = registry['data_location']
        self.track_mzs = registry['track_mzs']
        self.max_scan_number = registry['max_scan_number']
        self.anchor_mz_pairs = registry['anchor_mz_pairs']
        self.rt_numbers = registry['list_scan_numbers']
    '''
    def __init__(self, registry={}, experiment=None, mode='pos'):
        self.experiment = experiment    # parent Experiment instance
        self.mode = mode                # pos or neg, but can be mixed in a Method
        self.sample_type = ''           # QC, blank, study_sample, pooled_study_sample, ...

        # A number of attributes can be passed via registry dictionary
        self.input_file = registry['input_file']
        self.name = registry['name']
        self.id = registry['sample_id']
        self.list_retention_time = registry['list_retention_time']
        
        self.list_MassTracks = []
        self.list_peaks = []

    def serialize(self):
        '''
        return dictionary of key variables. Extend as needed.
        '''
        return {'id': self.id, 
                'list_peaks': self.list_peaks,
                'sample_type': self.sample_type
                }


class Spectrum:
    '''
    A list of values on a property in analytical chemistry.
    A mass spectrum is a list of m/z values with corresponding intensity values. 
    The "Spectrum" here is generic enough for LC-MS, GC-MS, LC-IMS-MS, etc. 
    It can be used for NMR and other technologies with minor modifications. 
    
    Mass Spectrum provide data points as measured on instrument, 
    This can be MS level 1, 2 or n.
    
    Example of a MS2 spectrum from MONA:
            {"instrument": "",
            "precursor type": "M+H[1+]",
            "precursor m/z": 169,
            "collision energy": "30V",
            "score": 5.5,
            "spectrum": "59.000:0.615142 72.600:0.031546 74.600:0.015773 78.900:0.086751 85.200:1.490536 150.500:0.055205 166.000:0.055205 167.200:100.000000",
            },
            {},
            {}
    '''
    def __init__(self, id=''):
        '''
        Default to mass spectrum level 2, but easy to be used for level 1 and other data types.
        '''
        self.id = id
        self.ms_level = 2
        self.precursor_ion = self.precursor_ion_mz = 0
        self.ionization = "pos"
        self.list_mz = []
        self.list_intensity = []
        self.retention_time = self.rtime = 0

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'precursor_ion': self.precursor_ion, 
                'rtime': self.rtime, 
                'ms_level': self.ms_level,
                'ionization': self.ionization,
                'list_intensity': self.list_intensity,
                }


class ArrayOfSpectra:
    '''
    Metabolomic experiments usually employ some chromatography as separation technique. 
    Therefore, analysis of a "Sample" by a "Method" generates a series of spectra.
    The "Array of Spectra" is composed by linking separation parameters with spectra.
    '''
    def __init__(self, id=''):
        '''
        This class is a conceptual abstraction, but can be wrapped into a "Sample" in practice. 
        '''
        self.id = id
        self.sample = None
        self.parameters = []        # typically for LC-MS: m/z, retention time, intensity
        self.list_values = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'sample': self.sample, 
                'list_values': self.list_values,
                }


class Peak:
    '''
    This refers to an elution peak in chromatography.
    The m/z peaks are handled by centroiding algorithms these days and are not part of this.

    Most software, e.g. XCMS, MZmine, OpenMS, does peak detection per sample then align them.
    Asari detects peaks on composite signals not on individual samples.

    Example peak in asari:    
        {"apex": 315,
        "peak_area": 3456351,
        "height": 471023,
        "left_base": 311,
        "right_base": 323,
        "cSelectivity": 1.0,
        "parent_masstrack_id": 902,
        "mz": 134.10852813720703,
        "snr": 74,
        "goodness_fitting": 0.8357564804888509,
        "id_number": "F97",
        "rtime": 121.099636272,
        "rtime_left_base": 119.5951254229998,
        "rtime_right_base": 124.12197382300019,}
    '''
    def __init__(self, id='', mode='pos'):
        '''
        EIC and peak_shape are defined by intensity as the the function of rtime.
        If mz is of little variation, it's not always necesssary to show list_mz.
        # other attributes of interest
        # e.g. for IM data
        # collision_cross_section = 0
        '''
        self.id = id
        self.ms_level = 1                   # MS levle 1, 2. 3, etc.
        self.mode = mode                    # ionization mode
        self.sample = ''                    # if back track to sample

        self.list_mz = []
        self.list_retention_time = []
        self.list_intensity = []

        # if RT aligned/adjusted
        self.list_retention_time_corrected = []

        # derivative to XIC
        self.mz, self.min_mz, self.max_mz = 0, 0, 0
        self.rtime, self.min_rtime, self.max_rtime = 0, 0, 0

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'mz': self.mz, 
                'rtime': self.rtime, 
                'ms_level': self.ms_level,
                'mode': self.mode,
                'list_mz': self.list_mz,
                'list_retention_time': self.list_retention_time,
                'list_intensity': self.list_intensity,
                }

class MassTrack:
    '''
    Same as extracted ion chromatogram. This is used in place of EIC (or XIC, mass trace), 
    defined by m/z, list_retention_time, list_intensity.
    Computationally equivalent to EIC for LC-MS data, but spanning for full RT range in asari.
    See https://github.com/shuzhao-li/asari for application in data preprocessing.
    '''
    def __init__(self, id=''):
        self.id = id 
        self.mz = 0
        self.list_retention_time = []
        self.list_intensity = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'mz': self.mz, 
                'list_retention_time': self.list_retention_time,
                'list_intensity': self.list_intensity,
                }


class Feature:
    '''
    A feature is a set of peaks that are aligned across samples.
    This is experiment specific.
    The m/z and retention_time of a feature is summarized on the member peaks.
    The variation between samples is reflected in data at peak level.

    Selectivity is defined in asari (https://github.com/shuzhao-li/asari):
        mSelectivity, how distinct are m/z measurements
        cSelectivity, how distinct are chromatograhic elution peaks
        dSelectivity, how distinct are database records

    The default is LC-MS feature. Derivative classes can be MS2feature, etc.
    '''
    def __init__(self, id=''):
        self.id = id                # e.g. 'F00001234'
        self.ms_level = 1           # MS levle - 1, 2. 3, etc.
        # These attributes are used in asari, but some are optional
        self.mz = 0
        self.parent_masstrack_id = None
        self.rtime = 0
        self.left_base = None
        self.right_base = None
        self.height = 0
        self.peak_area = 0
        self.goodness_fitting = 0
        self.snr = 0
        self.mSelectivity = 0
        self.cSelectivity = 0
        self.dSelectivity = 0

        # other attributes of interest
        self.list_peaks = self.including_peaks = []
        self.experiment_belonged = ''

        # place holder. Will have separate annotation class/method
        self.annotation = {
        }

        # statistics across samples
        self.statistics = {
            'intensity_sample_mean': None,
            'intensity_sample_std': None,
            'intensity_sample_cv': None,
            'intensity_replicate_cv': None,
            # statistic_score and p_value depend on the statistical test
            'statistic_score': None,
            'p_value': None,
        }

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'mz': self.mz, 
                'rtime': self.rtime, 
                'apex': self.rtime,
                'parent_masstrack_id': self.parent_masstrack_id,
                'peak_area': self.peak_area,
                'height': self.height,
                'left_base': self.left_base, 
                'right_base': self.right_base,   
                'mSelectivity': self.mSelectivity,
                'cSelectivity': self.cSelectivity,
                'dSelectivity': self.dSelectivity,
                'list_peaks': self.list_peaks,
                }


class EmpiricalCompound:
    '''
    EmpiricalCompound is a tentative compound/metabolite,
    a computational unit to represent the result of annotation on mass spec experiment.
    It should have reference to multiple ions (isotopes/adducts) that belong to the same metabolite,
    and can be a mixture of isobaric/isomeric metabolites when they are not distinguished by the mass spec data.

    This is used because the identification of compounds is not definitve at many stages of a project, 
    and this allows probablistic annotation on experimental data. 
    The probablity ranges between [0, 1]. 
    This unit then enables approaches to factor the probablistic models into biological interpretation (e.g. mummichog). 
    If an annotation method only provides scores (e.g. from MS2 search), mummichog will use them.

    EmpiricalCompound is experiment specific,
    and can combine multiple methods, including pos and neg ESI, MS^n.

    Similar concepts are 'pseudo spectrum' in CAMERA, and 'feature group' in mz.Unity.
    '''

    def __init__(self, id=''):
        '''
        An empCpd is the result of annotation.
        It has one and only one base neutral mass.
        Many attributes are optional.

        In a specific Experiment, an EmpiricalCompound consists of a set of features across a set of samples.
        They can be list of MS1 features, either using pointers to Features in the database.
        Features can include MS1 and MSn data.
        How to group Features into empCpd depends on annotation method, e.g. Annotation, AnnotationResult from mass2chem
        After annotation, not ruling out an empCpd can be mixture (isomers, etc)

        '''
        self.id = id                            # e.g. 'E00001234'
        self.interim_id = ''
        # Experiment specific.
        self.experiment_belonged = ''
        self.annotation_method = ''
        
        self.neutral_base_mass = self.neutral_formula_mass = 0.0000
        self.neutral_formula = ''
        self.charge = 0
        self.charged_formula = ''
        self.Database_referred = []

        self.MS1_pseudo_Spectra = self.list_features = []            # list of features that belong to this empCpd
        self.MS2_Spectra = []                   # MS2 identifiers can be universal (e.g. hashed ids)
        self.identity = self.annotation = []    # see desired serialize() output; also in README

    def read_json_model(self, jmodel):
        '''Modify as needed
        '''
        self.interim_id = jmodel['interim_id']
        self.neutral_formula_mass = jmodel['neutral_formula_mass']
        self.neutral_formula = jmodel['neutral_formula']
        self.Database_referred = jmodel['Database_referred']
        self.identity = jmodel['identity']
        self.MS1_pseudo_Spectra  = jmodel['MS1_pseudo_Spectra']
        self.MS2_Spectra = jmodel['MS2_Spectra']    

    def serialize(self):
        '''
        return dictionary of key variables. See README for an example.
        '''
        features = []
        for peak in self.MS1_pseudo_Spectra:
                features.append(        # this is given as example; one may need to modify the mapping variable names
                   {"feature_id": peak['id'], "mz": peak['mz'], "rtime": peak['rtime'], "charged_formula": "",  
                        "ion_relation": peak['ion_relation'],}
                )
        return {'interim_id': self.interim_id, 
                'neutral_formula_mass': self.neutral_formula_mass,
                'neutral_formula': self.neutral_formula,
                'Database_referred': self.Database_referred,
                'identity': self.write_identity(),
                'MS1_pseudo_Spectra': features,
                'MS2_Spectra': self.MS2_Spectra,
                }

    def write_identity(self):
        '''Place holder'''
        return self.identity

    def get_intensities(self):
        ''' Representative intensity values, can base on the MS1 feature of highest intensity
        self.intensities = { "sample1": 0, "sample2": 0, ... }
        # more efficient version of self.intensities
        self.intensities_by_ordered_samples = []
        '''
        pass

    def mummichog_annotation(self):
        '''
        Updated identity table by mummichog
        '''
        self.identity_probability_mummichog = [
            # updated probability after mummichog analysis
        ]



#
# Theoretical concepts (metabolic model): compound, reaction, pathway, network; enzyme, gene
#

class Compound:
    def __init__(self):
        '''
        All metabolites are compounds, but the reverse is not true.
        The identifiers from databases are a list of lists in self.db_ids. Because multiple identifiers
        can be found in one DB: [['KEGG', 'C0000'], ['HMDB', 'HMDB01858'], ['HMDB', 'HMDB13762'], ...].
        Use lists not tuple for JSON compatibility.
        '''
        self.internal_id = self.id = ''
        self.name = ''          # common name
        self.db_ids = []        # use list of lists, 
                                # e.g. [['KEGG', ''], ['HMDB', ''], ['Azimuth', 'HMDB13762'], ['PubChem', ''], ]
        self.neutral_formula = ''
        self.neutral_mono_mass = 0.0000
        # Often in metabolic models, compounds are in charged form
        self.charge = 0
        self.charged_formula = ''

        self.SMILES = ''
        self.inchi = ''

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'name': self.name, 
                'identifiers': self.db_ids,
                'neutral_formula': self.neutral_formula,
                'charge': self.charge,
                'charged_formula': self.charged_formula,
                'neutral_mono_mass': self.neutral_mono_mass,
                'SMILES': self.SMILES,
                'inchi': self.inchi,
                }


class Reaction:
    '''
    A reaction is defined by reactants and products, each a list of compounds.
    There is directionality of a reaction. A forward reaction is different from reverse reaction.
    We can treat the reactions catalyzed by different enzymes as the same reactions.

    Reactions are species specific, 
    because genes are species specific.
    '''
    def __init__(self):
        self.azimuth_id = self.id = ''
        self.name = ''    # common name of the reaction; ontologies should be higher level
        self.source = []
        self.version = ''
        # status, one of ['active', 'under review', 'obsolete']
        self.status = ''

        self.reactants = []
        self.products = []

        # below this line are optional
        self.enzymes = []
        self.genes = []

        # belong to
        self.pathways = []
        # still looking for good ontology for reactions. Maybe notes like "Glucuronidation" for now.
        self.ontologies = []

        self.species = ''
        self.compartments = []
        self.cell_types = []
        self.tissues = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'name': self.name,
                'reactants': self.reactants, 
                'products': self.products, 
                'genes': self.genes,
                'enzymes': self.enzymes,
                }


class Pathway:
    '''
    A pathway is defined by connected biochemical reactions, according to human definition.
    '''
    def __init__(self):
        self.azimuth_id = self.id = ''
        self.name = ''
        self.source = []
        self.list_of_reactions = []
        self.status = ''

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'name': self.name, 
                'list_of_reactions': self.list_of_reactions, 
                }


class Network:
    '''
    This refers to a metabolic network, defined by connected biochemical reactions.

    Network is mathematically identical to pathway, but not limited by pathway definition.
    Edges and nodes are computed based on reactions.

    Based on prior knowledge. This class does not include correlation networks etc.
    '''
    def __init__(self):
        self.azimuth_id = self.id = ''
        self.name = ''
        self.source = []
        self.list_of_reactions = []
        self.status = ''

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'list_of_reactions': self.list_of_reactions, 
                }


class MetabolicModel:
    '''
    A metabolic model, minimal information is list_of_reactions.
    Pathway definition isn't always available.
    Compounds are union of all reactants and products in reactions.
    Genes and proteins correspond to enzymes in reactions.
    The JMS package (https://github.com/shuzhao-li/JMS) handles the conversion of genome scale metabolic models.
    '''
    def __init__(self):
        self.id = ''
        self.meta_data = {
            'species': '',
            'version': '',
            'sources': [],
            'status': '',
            'last_update': '',
        }
        self.list_of_reactions = []
        self.list_of_pathways = []
        # list_of_compounds can be inferred from reactions, but good to keep annotation in cpds
        self.list_of_compounds = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'list_of_reactions': self.list_of_reactions, 
                'list_of_compounds': self.list_of_compounds,
                'list_of_pathways': self.list_of_pathways,
                'meta_data': self.meta_data,
                }


# ---------------------------------------------------------
# To extend later
#

class Enzyme:
    '''
    An enzyme is a protein that catalyzes biochemical reactions.
    '''
    def __init__(self):
        self.id = ''
        self.name = ''
        self.ec_num = ''
        self.url = ''
        self.genes = []
        self.reactions = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'name': self.name,
                'reactions': self.reactions
                }


class Gene:
    '''
    A gene is defined by polynucleotide sequence in a genome.
    Not a detailed model fo gene structure here. Main objective is to link to enzyme and biochemistry.
    '''
    def __init__(self):
        self.id = ''
        self.name = ''
        self.symbol = ''
        self.ensembl_id = ''
        self.description = ''
        self.proteins = []              # can be enzymes

        self.linked_metabolites = []
        self.linked_dieases = []

    def serialize(self):
        '''
        return dictionary of key variables.
        '''
        return {'id': self.id, 
                'name': self.name,
                'symbol': self.symbol,
                'proteins': self.proteins
                }

