# Path: DiscEvolution/chemistry/CO_chem.py
# class SimpleCOAtomAbund(ChemicalAbund):
#     """Class to hold the raw atomic abundaces of C/O/Si for the CO chemistry"""
#     def __init__(self, *sizes):
#         atom_ids = ['C', 'O', 'Si']
#         masses   = [ 12., 16., 28. ]
#         
#         super(SimpleCOAtomAbund, self).__init__(atom_ids, masses, *sizes)
# 
# 
#     def set_solar_abundances(self, muH=1.28):
#         """Solar mass fractions of C, O and Si.
# 
#         args:
#             muH : mean atomic mass, default = 1.28
#         """
#         self._data[:] = np.outer(np.array([12*2.7e-4, 16*4.9e-4, 28*3.2e-5]),
#                                  np.ones(self.size)) / muH
# 
# class SimpleCOChemOberg(COChemOberg, StaticChem):
#     def __init__(self, **kwargs):
#         COChemOberg.__init__(self)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCOChemOberg(COChemOberg, EquilibriumChem):
#     def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
#         COChemOberg.__init__(self, fix_grains)
#         EquilibriumChem.__init__(self, 
#                                  fix_ratios=fix_ratios, 
#                                  **kwargs)
# 
# class SimpleCOChemMadhu(COChemMadhu, StaticChem):
#     def __init__(self, **kwargs):
#         COChemMadhu.__init__(self)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCOChemMadhu(COChemMadhu, EquilibriumChem):
#     def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
#         COChemMadhu.__init__(self, fix_grains)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios, 
#                                  **kwargs)
# 
# Path: DiscEvolution/chemistry/CNO_chem.py
# class SimpleCNOAtomAbund(ChemicalAbund):
#     """Class to hold the raw atomic abundaces of C/O/Si for the CNO chemistry"""
# 
#     def __init__(self, *sizes):
#         atom_ids = ['C', 'N', 'O', 'Si']
#         masses = [12., 14., 16., 28.]
# 
#         super(SimpleCNOAtomAbund, self).__init__(atom_ids, masses, *sizes)
# 
#     def set_solar_abundances(self, muH=1.28):
#         """Solar mass fractions of C, O and Si.
# 
#         args:
#             muH : mean atomic mass, default = 1.28
#         """
#         m_abund = np.array([12 * 2.7e-4, 14 * 6.8e-5, 16 * 4.9e-4, 28 * 3.2e-5])
#         self._data[:] = np.outer(m_abund, np.ones(self.size)) / muH
# 
# class SimpleCNOChemOberg(CNOChemOberg, StaticChem):
#     def __init__(self, fNH3=None, **kwargs):
#         CNOChemOberg.__init__(self, fNH3)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCNOChemOberg(CNOChemOberg, EquilibriumChem):
#     def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
#                  fix_N=False, **kwargs):
#         CNOChemOberg.__init__(self, fNH3, fix_grains, fix_N)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios,
#                                  **kwargs)
# 
# class SimpleCNOChemMadhu(CNOChemMadhu, StaticChem):
#     def __init__(self, fNH3=None, **kwargs):
#         CNOChemMadhu.__init__(self, fNH3)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCNOChemMadhu(CNOChemMadhu, EquilibriumChem):
#     def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
#                  fix_N=False, **kwargs):
#         CNOChemMadhu.__init__(self, fNH3, fix_grains, fix_N)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios,
#                                  **kwargs)


import numpy as np
from DiscEvolution.chemistry import (
    SimpleCOAtomAbund,
    SimpleCOChemMadhu, SimpleCOChemOberg,
    EquilibriumCOChemMadhu, EquilibriumCOChemOberg
)
from DiscEvolution.chemistry import (
    SimpleCNOAtomAbund,
    SimpleCNOChemMadhu, SimpleCNOChemOberg,
    EquilibriumCNOChemMadhu, EquilibriumCNOChemOberg
)
# test_simple_chem.py
#
# Author: R. Booth
# Date: 22 - May - 2018
#
# Check behaviour of simple chemistry functions
###############################################################################

# CO chemistry

# CNO chemistry

def _test_chemical_model(ChemModel, Abundances):
    T = np.logspace(0.5, 3, 6)
    
    Xi = Abundances(len(T))
    Xi.set_solar_abundances()

    Chem = ChemModel()
    
    mol = Chem.equilibrium_chem(T, 1e-10, 0.01, Xi)
        
    T *= 3
    Chem.update(0, T, 1e-10, 0.01, mol)


# Path: DiscEvolution/chemistry/CO_chem.py
# class SimpleCOAtomAbund(ChemicalAbund):
#     """Class to hold the raw atomic abundaces of C/O/Si for the CO chemistry"""
#     def __init__(self, *sizes):
#         atom_ids = ['C', 'O', 'Si']
#         masses   = [ 12., 16., 28. ]
#         
#         super(SimpleCOAtomAbund, self).__init__(atom_ids, masses, *sizes)
# 
# 
#     def set_solar_abundances(self, muH=1.28):
#         """Solar mass fractions of C, O and Si.
# 
#         args:
#             muH : mean atomic mass, default = 1.28
#         """
#         self._data[:] = np.outer(np.array([12*2.7e-4, 16*4.9e-4, 28*3.2e-5]),
#                                  np.ones(self.size)) / muH
# 
# class SimpleCOChemOberg(COChemOberg, StaticChem):
#     def __init__(self, **kwargs):
#         COChemOberg.__init__(self)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCOChemOberg(COChemOberg, EquilibriumChem):
#     def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
#         COChemOberg.__init__(self, fix_grains)
#         EquilibriumChem.__init__(self, 
#                                  fix_ratios=fix_ratios, 
#                                  **kwargs)
# 
# class SimpleCOChemMadhu(COChemMadhu, StaticChem):
#     def __init__(self, **kwargs):
#         COChemMadhu.__init__(self)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCOChemMadhu(COChemMadhu, EquilibriumChem):
#     def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
#         COChemMadhu.__init__(self, fix_grains)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios, 
#                                  **kwargs)
# 
# Path: DiscEvolution/chemistry/CNO_chem.py
# class SimpleCNOAtomAbund(ChemicalAbund):
#     """Class to hold the raw atomic abundaces of C/O/Si for the CNO chemistry"""
# 
#     def __init__(self, *sizes):
#         atom_ids = ['C', 'N', 'O', 'Si']
#         masses = [12., 14., 16., 28.]
# 
#         super(SimpleCNOAtomAbund, self).__init__(atom_ids, masses, *sizes)
# 
#     def set_solar_abundances(self, muH=1.28):
#         """Solar mass fractions of C, O and Si.
# 
#         args:
#             muH : mean atomic mass, default = 1.28
#         """
#         m_abund = np.array([12 * 2.7e-4, 14 * 6.8e-5, 16 * 4.9e-4, 28 * 3.2e-5])
#         self._data[:] = np.outer(m_abund, np.ones(self.size)) / muH
# 
# class SimpleCNOChemOberg(CNOChemOberg, StaticChem):
#     def __init__(self, fNH3=None, **kwargs):
#         CNOChemOberg.__init__(self, fNH3)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCNOChemOberg(CNOChemOberg, EquilibriumChem):
#     def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
#                  fix_N=False, **kwargs):
#         CNOChemOberg.__init__(self, fNH3, fix_grains, fix_N)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios,
#                                  **kwargs)
# 
# class SimpleCNOChemMadhu(CNOChemMadhu, StaticChem):
#     def __init__(self, fNH3=None, **kwargs):
#         CNOChemMadhu.__init__(self, fNH3)
#         StaticChem.__init__(self, **kwargs)
# 
# class EquilibriumCNOChemMadhu(CNOChemMadhu, EquilibriumChem):
#     def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
#                  fix_N=False, **kwargs):
#         CNOChemMadhu.__init__(self, fNH3, fix_grains, fix_N)
#         EquilibriumChem.__init__(self,
#                                  fix_ratios=fix_ratios,
#                                  **kwargs)


# Path: tests/test_simple_chem.py
import numpy as np
from DiscEvolution.chemistry import (
    SimpleCOAtomAbund,
    SimpleCOChemMadhu, SimpleCOChemOberg,
    EquilibriumCOChemMadhu, EquilibriumCOChemOberg
)
from DiscEvolution.chemistry import (
    SimpleCNOAtomAbund,
    SimpleCNOChemMadhu, SimpleCNOChemOberg,
    EquilibriumCNOChemMadhu, EquilibriumCNOChemOberg
)

# test_simple_chem.py
#
# Author: R. Booth
# Date: 22 - May - 2018
#
# Check behaviour of simple chemistry functions
###############################################################################

# CO chemistry

# CNO chemistry

def _test_chemical_model(ChemModel, Abundances):
    T = np.logspace(0.5, 3, 6)
    
    Xi = Abundances(len(T))
    Xi.set_solar_abundances()

    Chem = ChemModel()
    
    mol = Chem.equilibrium_chem(T, 1e-10, 0.01, Xi)
        
    T *= 3
    Chem.update(0, T, 1e-10, 0.01, mol)


    atom  = mol.gas.atomic_abundance()