# Path: utility/input.py
# class AtomicInput(object):
#     def __init__(self, energy, projectile, source, current, param_name):
#         self.energy = energy
#         self.projectile = projectile
#         self.data_source = source
#         self.current = current
#         self.param_name = param_name
#         self.components = None
#         self.param = None
#         self.component_col = ['q', 'Z', 'A', 'Molecule']
#         self.component_index = []
#         self.component_data = []
#         self._build_param()
# 
#     def _build_param(self):
#         xml = etree.Element('xml')
#         head = etree.SubElement(xml, 'head')
#         id_tag = etree.SubElement(head, 'id')
#         id_tag.text = self.param_name
#         body_tag = etree.SubElement(xml, 'body')
#         beamlet_energy = etree.SubElement(body_tag, 'beamlet_energy', {'unit': 'keV'})
#         beamlet_energy.text = str(self.energy)
#         beamlet_species = etree.SubElement(body_tag, 'beamlet_species')
#         beamlet_species.text = self.projectile
#         beamlet_source = etree.SubElement(body_tag, 'beamlet_source')
#         beamlet_source.text = self.data_source
#         beamlet_current = etree.SubElement(body_tag, 'beamlet_current', {'unit': 'A'})
#         beamlet_current.text = str(self.current)
#         self.param = etree.ElementTree(element=xml)
# 
#     def _build_components(self):
#         if len(self.component_data) == 0:
#             raise ValueError('The component data base was not initiated.')
#         else:
#             self.components = pd.DataFrame(np.array(self.component_data), columns=self.component_col,
#                                            index=self.component_index)
# 
#     def _next_element(self, name):
#         return name + str(len([component for component in self.component_index if name in component])+1)
# 
#     def add_target_component(self, charge, atomic_number, mass_number, molecule_name=None):
#         element = [charge, atomic_number, mass_number, molecule_name]
#         if charge == -1:
#             if 'electron' in self.component_index:
#                 print('Already a component.')
#             else:
#                 self.component_index.append('electron')
#                 self.component_data.append(element)
#         elif charge == 0:
#             if element in self.component_data:
#                 print('Already a component.')
#             else:
#                 self.component_index.append(self._next_element('neutral'))
#                 self.component_data.append(element)
#         elif charge >= 1:
#             if element in self.component_data:
#                 print('Already a component.')
#             else:
#                 self.component_index.append(self._next_element('ion'))
#                 self.component_data.append(element)
# 
#     def get_atomic_db_input(self):
#         self._build_components()
#         return self.param, self.components
# 
# Path: utility/exceptions.py
# class RenateNotValidTransitionError(Exception):
#     pass
# 
# Path: crm_solver/atomic_db.py
# class RenateDB:
# class AtomicDB(RenateDB):
#     def __init__(self, param, rate_type, data_path):
#     def __set_impurity_mass_scaling_dictionary(self):
#     def __projectile_parameters(self):
#     def __get_atomic_mass(self):
#     def __get_projectile_velocity(self):
#     def __set_atomic_dictionary(self):
#     def __set_rates_path(self, rate_type):
#     def __set_charge_state_lib(self):
#     def load_rate_data(path, tag_name):
#     def set_default_atomic_levels(self):
#     def get_from_renate_atomic(self, source):
#     def __init__(self, atomic_source='renate', param=None, rate_type='default', resolution=None,
#                  data_path='beamlet/testimp0001.xml', components=None, atomic_ceiling=False):
#     def __set_neutral_db(self, param, resolution):
#     def __set_ceiling_for_atomic_levels(self, atomic_ceiling):
#     def __generate_rate_function_db(self):
#     def __set_temperature_axis(self):
#     def __set_einstein_coefficient_db(self):
#     def __set_impact_loss_functions(self):
#     def __set_electron_impact_transition_functions(self):
#     def __set_ion_impact_transition_functions(self):
#     def __interp1d_scaled_ion(self, rates, target):
#     def plot_rates(self, *args, temperature=None, external_density=1.):


import unittest
from utility.input import AtomicInput
from utility.exceptions import RenateNotValidTransitionError
from crm_solver.atomic_db import NeutralDB


class NeutralDBTest(unittest.TestCase):

    INPUT_ENERGY = 50
    INPUT_PROJECTILE = 'H'
    INPUT_CURRENT = 0.001
    INPUT_SOURCE = 'Internal'
    INPUT_PARAM_NAME = 'NeutralDB test'
    INPUT_H = [0, 1, 1, None]
    INPUT_H2 = [0, 1, 1, 'H2']
    INPUT_TEST = 'test'

    EXPECTED_KEY_2 = 2
    EXPECTED_KEY_1 = 1
    EXPECTED_KEY_NAME_1 = 'neutral1'
    EXPECTED_KEY_NAME_2 = ['neutral1', 'neutral2']
    EXPECTED_TEST_ATOMIC_LEVELS = 4

    INPUT_FROM_LEVELS = 0
    INPUT_TO_LEVELS = [1, 2, 3]

    EXPECTED_LOSS = [0.0001, 0.0002, 0.0003, 0.0004]
    EXPECTED_TRANS = [0.0012, 0.0013, 0.0014]

    def setUp(self):

# Path: utility/input.py
# class AtomicInput(object):
#     def __init__(self, energy, projectile, source, current, param_name):
#         self.energy = energy
#         self.projectile = projectile
#         self.data_source = source
#         self.current = current
#         self.param_name = param_name
#         self.components = None
#         self.param = None
#         self.component_col = ['q', 'Z', 'A', 'Molecule']
#         self.component_index = []
#         self.component_data = []
#         self._build_param()
# 
#     def _build_param(self):
#         xml = etree.Element('xml')
#         head = etree.SubElement(xml, 'head')
#         id_tag = etree.SubElement(head, 'id')
#         id_tag.text = self.param_name
#         body_tag = etree.SubElement(xml, 'body')
#         beamlet_energy = etree.SubElement(body_tag, 'beamlet_energy', {'unit': 'keV'})
#         beamlet_energy.text = str(self.energy)
#         beamlet_species = etree.SubElement(body_tag, 'beamlet_species')
#         beamlet_species.text = self.projectile
#         beamlet_source = etree.SubElement(body_tag, 'beamlet_source')
#         beamlet_source.text = self.data_source
#         beamlet_current = etree.SubElement(body_tag, 'beamlet_current', {'unit': 'A'})
#         beamlet_current.text = str(self.current)
#         self.param = etree.ElementTree(element=xml)
# 
#     def _build_components(self):
#         if len(self.component_data) == 0:
#             raise ValueError('The component data base was not initiated.')
#         else:
#             self.components = pd.DataFrame(np.array(self.component_data), columns=self.component_col,
#                                            index=self.component_index)
# 
#     def _next_element(self, name):
#         return name + str(len([component for component in self.component_index if name in component])+1)
# 
#     def add_target_component(self, charge, atomic_number, mass_number, molecule_name=None):
#         element = [charge, atomic_number, mass_number, molecule_name]
#         if charge == -1:
#             if 'electron' in self.component_index:
#                 print('Already a component.')
#             else:
#                 self.component_index.append('electron')
#                 self.component_data.append(element)
#         elif charge == 0:
#             if element in self.component_data:
#                 print('Already a component.')
#             else:
#                 self.component_index.append(self._next_element('neutral'))
#                 self.component_data.append(element)
#         elif charge >= 1:
#             if element in self.component_data:
#                 print('Already a component.')
#             else:
#                 self.component_index.append(self._next_element('ion'))
#                 self.component_data.append(element)
# 
#     def get_atomic_db_input(self):
#         self._build_components()
#         return self.param, self.components
# 
# Path: utility/exceptions.py
# class RenateNotValidTransitionError(Exception):
#     pass
# 
# Path: crm_solver/atomic_db.py
# class RenateDB:
# class AtomicDB(RenateDB):
#     def __init__(self, param, rate_type, data_path):
#     def __set_impurity_mass_scaling_dictionary(self):
#     def __projectile_parameters(self):
#     def __get_atomic_mass(self):
#     def __get_projectile_velocity(self):
#     def __set_atomic_dictionary(self):
#     def __set_rates_path(self, rate_type):
#     def __set_charge_state_lib(self):
#     def load_rate_data(path, tag_name):
#     def set_default_atomic_levels(self):
#     def get_from_renate_atomic(self, source):
#     def __init__(self, atomic_source='renate', param=None, rate_type='default', resolution=None,
#                  data_path='beamlet/testimp0001.xml', components=None, atomic_ceiling=False):
#     def __set_neutral_db(self, param, resolution):
#     def __set_ceiling_for_atomic_levels(self, atomic_ceiling):
#     def __generate_rate_function_db(self):
#     def __set_temperature_axis(self):
#     def __set_einstein_coefficient_db(self):
#     def __set_impact_loss_functions(self):
#     def __set_electron_impact_transition_functions(self):
#     def __set_ion_impact_transition_functions(self):
#     def __interp1d_scaled_ion(self, rates, target):
#     def plot_rates(self, *args, temperature=None, external_density=1.):


# Path: crm_solver/neutral_dbtest.py
import unittest
from utility.input import AtomicInput
from utility.exceptions import RenateNotValidTransitionError
from crm_solver.atomic_db import NeutralDB



class NeutralDBTest(unittest.TestCase):

    INPUT_ENERGY = 50
    INPUT_PROJECTILE = 'H'
    INPUT_CURRENT = 0.001
    INPUT_SOURCE = 'Internal'
    INPUT_PARAM_NAME = 'NeutralDB test'
    INPUT_H = [0, 1, 1, None]
    INPUT_H2 = [0, 1, 1, 'H2']
    INPUT_TEST = 'test'

    EXPECTED_KEY_2 = 2
    EXPECTED_KEY_1 = 1
    EXPECTED_KEY_NAME_1 = 'neutral1'
    EXPECTED_KEY_NAME_2 = ['neutral1', 'neutral2']
    EXPECTED_TEST_ATOMIC_LEVELS = 4

    INPUT_FROM_LEVELS = 0
    INPUT_TO_LEVELS = [1, 2, 3]

    EXPECTED_LOSS = [0.0001, 0.0002, 0.0003, 0.0004]
    EXPECTED_TRANS = [0.0012, 0.0013, 0.0014]

    def setUp(self):

        self.input = AtomicInput(energy=self.INPUT_ENERGY, projectile=self.INPUT_PROJECTILE, current=self.INPUT_CURRENT,