# Path: helita/sim/rh15d.py
# class Rh15dout:
# class HDF5Atmos:
# class DataHolder:
# class AtomFile:
#     def __init__(self, fdir='.', verbose=True, autoread=True):
#     def read_groups(self, infile):
#     def read_ray(self, infile=None):
#     def close(self):
#     def __del__(self):
#     def __init__(self, infile):
#     def close(self):
#     def read(self, infile):
#     def write_multi(self, outfile, xi, yi, nti=0, writeB=False,
#                     write_dscale=False, zcut=0, depth_optimise=False):
#     def write_multi_3d(self, outfile, nti=0, sx=None, sy=None, sz=None,
#                        big_endian=False):
#     def __init__(self):
#     def __init__(self, filename, format='RH'):
#     def read_atom_levels(data, format='RH'):
#     def read_atom(self, filename, format='RH'):
# def read_hdf5(inclass, infile):
# def make_xarray_atmos(outfile, T, vz, z, nH=None, x=None, y=None, Bz=None, By=None,
#                       Bx=None, rho=None, ne=None, vx=None, vy=None, vturb=None,
#                       desc=None, snap=None, boundary=None, append=False):
# def depth_optim(height, temp, ne, vz, rho, nh=None, bx=None, by=None, bz=None,
#                 tmax=5e4):
# def make_wave_file(outfile, start=None, end=None, step=None, new_wave=None,
#                    ewave=None, air=True):
# def read_wave_file(infile):
# def clean_var(data, only_positive=True):
#                 OUTFILE = os.path.join(self.fdir, "%s.hdf5" % (outfile))
#             RAYFILE = os.path.join(self.fdir, "output_ray.hdf5")
#         GROUPS = netCDF4.Dataset(infile).groups.keys()
#         M_TO_CM3 = (units.m**-3).to('1 / cm3')
#         M_TO_KM = units.m.to('km')
#         COLLISION_KEYS_TEMP = ['OHMEGA', 'OMEGA', 'CE', 'CI', 'CP', 'CH',
#                                'CH0', 'CH+', 'CR', 'TEMP']
#         COLLISION_KEYS_LINE = ['AR85-CEA', 'AR85-CHP', 'AR85-CHH', 'SHULL82',
#                                'BURGESS', 'SUMMERS']
#         COLLISION_KEYS_OTHER = ['AR85-CDI', 'BADNELL']
#         ALL_KEYS = (COLLISION_KEYS_TEMP + COLLISION_KEYS_LINE +
#                         COLLISION_KEYS_OTHER)
#         SINGLE_KEYS = ['GENCOL', 'END']
#     VARS4D = ['temperature', 'B_x', 'B_y', 'B_z', 'density', 'velocity_x',
#               'velocity_y', 'velocity_z', 'velocity_turbulent', 'density',
#               'electron_density']


import numpy as np
from helita.sim import rh15d
   60.0   1.8302E-23
   55.0   1.7001E-23
   50.0   1.5539E-23
   45.0   1.3944E-23
   40.0   1.2248E-23
   35.0   1.0486E-23
 TEMP    6          3000.0     5000.0     7000.0    15000.0    50000.0   100000.0
 OMEGA   1  0    2.378E+00  2.284E+00  2.203E+00  1.920E+00  1.961E+00  1.846E+00
 OMEGA   2  0    3.568E+00  3.426E+00  3.304E+00  2.879E+00  2.942E+00  2.770E+00
 TEMP    6          3000.0     5000.0     7000.0    15000.0    50000.0   100000.0
 CI      0  5    4.580E-18  4.580E-18  4.580E-18  4.580E-18  4.580E-18  4.580E-18
SUMMERS  1.0
SHULL82 0 3 0.00e+00 1.31e+05 4.70e-13 6.24e-01 0.00e-03 4.42e-02 1.57e+05 3.74e+05
AR85-CDI   0   3   2
    11.30     3.60    -9.60     7.20    -9.06
    16.60    14.58    -4.68     1.50   -14.40
AR85-CEA   0   3  6.00e-01

END
"""

def test_read_atom_levels():
    assert np.array_equal(rh15d.AtomFile.read_atom_levels(TEST_LEVELS, format='RH'),
                          TEST_LEVELS_DATA)

def test_AtomFile():
    temp_file = open(TMP_ATOM_FILENAME, 'w')
    temp_file.write(TEST_ATOM_RH)
    temp_file.close()


# Path: helita/sim/rh15d.py
# class Rh15dout:
# class HDF5Atmos:
# class DataHolder:
# class AtomFile:
#     def __init__(self, fdir='.', verbose=True, autoread=True):
#     def read_groups(self, infile):
#     def read_ray(self, infile=None):
#     def close(self):
#     def __del__(self):
#     def __init__(self, infile):
#     def close(self):
#     def read(self, infile):
#     def write_multi(self, outfile, xi, yi, nti=0, writeB=False,
#                     write_dscale=False, zcut=0, depth_optimise=False):
#     def write_multi_3d(self, outfile, nti=0, sx=None, sy=None, sz=None,
#                        big_endian=False):
#     def __init__(self):
#     def __init__(self, filename, format='RH'):
#     def read_atom_levels(data, format='RH'):
#     def read_atom(self, filename, format='RH'):
# def read_hdf5(inclass, infile):
# def make_xarray_atmos(outfile, T, vz, z, nH=None, x=None, y=None, Bz=None, By=None,
#                       Bx=None, rho=None, ne=None, vx=None, vy=None, vturb=None,
#                       desc=None, snap=None, boundary=None, append=False):
# def depth_optim(height, temp, ne, vz, rho, nh=None, bx=None, by=None, bz=None,
#                 tmax=5e4):
# def make_wave_file(outfile, start=None, end=None, step=None, new_wave=None,
#                    ewave=None, air=True):
# def read_wave_file(infile):
# def clean_var(data, only_positive=True):
#                 OUTFILE = os.path.join(self.fdir, "%s.hdf5" % (outfile))
#             RAYFILE = os.path.join(self.fdir, "output_ray.hdf5")
#         GROUPS = netCDF4.Dataset(infile).groups.keys()
#         M_TO_CM3 = (units.m**-3).to('1 / cm3')
#         M_TO_KM = units.m.to('km')
#         COLLISION_KEYS_TEMP = ['OHMEGA', 'OMEGA', 'CE', 'CI', 'CP', 'CH',
#                                'CH0', 'CH+', 'CR', 'TEMP']
#         COLLISION_KEYS_LINE = ['AR85-CEA', 'AR85-CHP', 'AR85-CHH', 'SHULL82',
#                                'BURGESS', 'SUMMERS']
#         COLLISION_KEYS_OTHER = ['AR85-CDI', 'BADNELL']
#         ALL_KEYS = (COLLISION_KEYS_TEMP + COLLISION_KEYS_LINE +
#                         COLLISION_KEYS_OTHER)
#         SINGLE_KEYS = ['GENCOL', 'END']
#     VARS4D = ['temperature', 'B_x', 'B_y', 'B_z', 'density', 'velocity_x',
#               'velocity_y', 'velocity_z', 'velocity_turbulent', 'density',
#               'electron_density']


# Path: helita/sim/tests/test_rh15d.py
import numpy as np
from helita.sim import rh15d

   60.0   1.8302E-23
   55.0   1.7001E-23
   50.0   1.5539E-23
   45.0   1.3944E-23
   40.0   1.2248E-23
   35.0   1.0486E-23
 TEMP    6          3000.0     5000.0     7000.0    15000.0    50000.0   100000.0
 OMEGA   1  0    2.378E+00  2.284E+00  2.203E+00  1.920E+00  1.961E+00  1.846E+00
 OMEGA   2  0    3.568E+00  3.426E+00  3.304E+00  2.879E+00  2.942E+00  2.770E+00
 TEMP    6          3000.0     5000.0     7000.0    15000.0    50000.0   100000.0
 CI      0  5    4.580E-18  4.580E-18  4.580E-18  4.580E-18  4.580E-18  4.580E-18
SUMMERS  1.0
SHULL82 0 3 0.00e+00 1.31e+05 4.70e-13 6.24e-01 0.00e-03 4.42e-02 1.57e+05 3.74e+05
AR85-CDI   0   3   2
    11.30     3.60    -9.60     7.20    -9.06
    16.60    14.58    -4.68     1.50   -14.40
AR85-CEA   0   3  6.00e-01

END
"""

def test_read_atom_levels():
    assert np.array_equal(rh15d.AtomFile.read_atom_levels(TEST_LEVELS, format='RH'),
                          TEST_LEVELS_DATA)

def test_AtomFile():
    temp_file = open(TMP_ATOM_FILENAME, 'w')
    temp_file.write(TEST_ATOM_RH)
    temp_file.close()


    data = rh15d.AtomFile(TMP_ATOM_FILENAME)