#!/usr/bin/env python
# coding: utf-8

# # MyRadex usage example

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')

def lower_keys(p):
    return {k.lower(): p[k] for k in p}

def print_e(en, f):
    for i in range(len(en)):
        print('{:9.2f} {:9.2e}'.format(en[i], f[i]))
    return

def cast_into_dic(col_names, arr):
    '''col_names is column_info, and arr is data_transitions (see below).'''
    names = col_names.split()
    return {names[i]: arr[i,:] for i in range(len(names))}

import os


# # Start using the wrapper

# ## Import the wrapper

# In[2]:


import wrapper_my_radex
wrapper = wrapper_my_radex.myradex_wrapper
about_info = wrapper.about.tobytes()
column_info = wrapper.column_names.tobytes()


# ## Load the molecule data
# 
# **The returned ```n_levels,n_item,n_transitions, n_partners``` are needed for the next step.**

# In[5]:


n_levels, n_item, n_transitions, n_partners = \
    wrapper.config_basic(os.path.expanduser('~/_c/my_radex/'),
                         'ph2co-h2.dat', 2.73, True)


# In[6]:


n_levels, n_item, n_transitions, n_partners


# ## Do the statistical equilibrium calculation

# ### Using the default geometry type

# In[25]:


params = {'Tkin': 25.0,
          'dv_CGS': 1e5,
          'dens_X_CGS': 1e6,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e5,
          'HI_density_CGS': 1e1,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 1e6,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions}

"""The keywords to the wrapper function have to be in
lower case, so I have to lower the keys of params.
Of course you can use lower case letters from the beginning."""

params = lower_keys(params)


# In[26]:


"""Do the calculation.
Return a four element tuple, whose meanings are self-evident."""

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)


# ### Calculate the critical densities
# 
# The `config_basic` function accepts two optional inputs:
# `recalculatefreqwitheupelow` and `ilevel_subtract_one`.
# By default they are both set to False.
# - `recalculatefreqwitheupelow`: If True, the code will recalculate the transition frequencies based on the energy differences, otherwise it will simply adopt the values in the input data file.
# - `ilevel_subtract_one`: If True, the level numbers in the output will be subtracted by 1, otherwise it will not be the same as in the input data file.
# 
# **You must run `run_one_params` once before calling `calc_critical_density`** to feed in the physical parameters.

# In[9]:


n_levels, n_item, n_transitions, n_partners = \
    wrapper.config_basic(os.path.expanduser('~/_c/my_radex/'),
                         '12C16O_H2.dat', 2.73, True)
print(n_levels, n_item, n_transitions, n_partners)

# A dummy call to run_one_params is necessary to feed in the parameters.
# This of course could be done better.
params.update({'n_levels': n_levels,
               'n_item': n_item,
               'n_transitions': n_transitions})
_ = wrapper.run_one_params(**params)


# In[10]:


tau = 1e-10 # optically thin
n_cri, iup, ilow = wrapper.calc_critical_density(tau, n_partners, n_transitions)

ntran_show = 10

for j in range(n_partners):
    print('Collisional partner', j)
    for i in range(ntran_show):
        print(iup[i], ilow[i], '{:.3e}'.format(n_cri[j,i]))


# **Calculate the column density with the old definition, aka.**
# $$n_\text{crit,old}\equiv \frac{A_{ij}}{\gamma_{ij}}$$

# In[11]:


tau = 1e-10 # optically thin
n_cri, iup, ilow = wrapper.calc_critical_density_old_def(tau, n_partners, n_transitions)

ntran_show = 10

for j in range(n_partners):
    print('Collisional partner', j)
    for i in range(ntran_show):
        print(iup[i], ilow[i], '{:.3e}'.format(n_cri[j,i]))


# In[12]:


n_levels, n_item, n_transitions, n_partners = \
    wrapper.config_basic(os.path.expanduser('~/_c/my_radex/'),
                         'hcn.dat', 2.73, True)
print(n_levels, n_item, n_transitions, n_partners)
params.update({'n_levels': n_levels,
               'n_item': n_item,
               'n_transitions': n_transitions})

_ = wrapper.run_one_params(**params)

tau = 1e-10 # optically thin
n_cri, iup, ilow = wrapper.calc_critical_density(tau, n_partners, n_transitions)

ntran_show = 15

for j in range(n_partners):
    print('Collisional partner', j)
    for i in range(ntran_show):
        print(iup[i], ilow[i], '{:.3e}'.format(n_cri[j,i]))


# In[13]:


_ = wrapper.run_one_params(**params)


# In[14]:


tau = 1e-10 # optically thin
n_cri, iup, ilow = wrapper.calc_critical_density(tau, n_partners, n_transitions)

ntran_show = 5

for j in range(n_partners):
    print('Collisional partner', j)
    for i in range(ntran_show):
        print(iup[i], ilow[i], '{:.3e}'.format(n_cri[j,i]))


# ### Specify the geometry type

# In[28]:


params = {'Tkin': 1e3,
          'dv_CGS': 1e5,
          'dens_X_CGS': 1e3,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e4,
          'HI_density_CGS': 1e1,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 1e1,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions,
          'geotype': 'lvg'
         }

params = lower_keys(params)

"""
The geotype parameter can take the following values:
spherical
lvg
slab
default
"""

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params_geometry(**params)


# # Play with the results
# 
# 1. You may want to transpose ```data_transitions``` to match the expected layout, but you don't have to.
# 1. You may also cast the the results into a python dict using ```cast_into_dic``` defined above.

# In[20]:


figure(figsize=(6,3), dpi=200)
plot(energies, f_occupations, marker='o')
ax = gca()
ax.set_xlabel('E (K)')
ax.set_ylabel('Occupation')
ax.set_ylim((1e-15,1e0))
ax.set_xscale('log')
ax.set_yscale('log')
#set_axis_format(ax, graygrid=True)


# In[29]:


cooling_rate


# In[22]:


data_transitions.shape


# In[23]:


wrapper.flag_good


# In[24]:


wrapper.n_item_column


# **You can stop reading here.  I forgot what I did in the following.**
# 
# ---

# # Work with other molecules

# In[ ]:


n_levels,n_item,n_transitions = \
    wrapper.config_basic(os.path.expanduser('~/not_synced/work/from_others/CHIANTI/toLAMDA/'),
                         'O.dat', 5.73, True)

params = {'Tkin': 1e3,
          'dv_CGS': 1e5,
          'dens_X_CGS': 1e6,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e9,
          'HI_density_CGS': 1e1,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 1e6,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions}

params = lower_keys(params)

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)

print_e(energies, f_occupations)


# In[11]:


n_levels,n_item,n_transitions = \
    wrapper.config_basic('~/not_synced/work/from_others/CHIANTI/toLAMDA/',
                         'N+_1e5K.dat', True)

params = {'Tkin': 1e3,
          'dv_CGS': 1e5,
          'dens_X_CGS': 1e6,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e9,
          'HI_density_CGS': 1e1,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 1e6,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions}

lower_keys(params)

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)

print_e(energies, f_occupations)


# In[55]:


n_levels,n_item,n_transitions = \
    wrapper.config_basic('~/not_synced/work/from_others/CHIANTI/toLAMDA/',
                         'N+_1e5K.dat', True)

params = {'Tkin': 1e3,
          'dv_CGS': 1e5,
          'dens_X_CGS': 1e6,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e9,
          'HI_density_CGS': 1e1,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 1e6,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions}

lower_keys(params)

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)

print_e(energies, f_occupations)


# In[3]:


n_levels,n_item,n_transitions,n_partners = \
    wrapper.config_basic('./', '12C16O_H2.dat', 2.73, True)


# In[4]:


params = {'Tkin': 1e2,
          'dv_CGS': 10e5,
          'dens_X_CGS': 1e0,
          'Ncol_X_CGS': 1e15,
          'H2_density_CGS': 1e2,
          'HI_density_CGS': 0.0,
          'oH2_density_CGS': 0.0,
          'pH2_densty_CGS': 0.0,
          'HII_density_CGS': 0.0,
          'Electron_density_CGS': 0.0,
          'n_levels': n_levels,
          'n_item': n_item,
          'n_transitions': n_transitions}

params = lower_keys(params)

energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)
d1 = cast_into_dic(column_info, data_transitions)
print(d1.keys())

params['H2_density_CGS'.lower()] = 1e4
energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)
d2 = cast_into_dic(column_info, data_transitions)

params['H2_density_CGS'.lower()] = 1e5
energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)
d3 = cast_into_dic(column_info, data_transitions)

params['H2_density_CGS'.lower()] = 1e6
energies,f_occupations,data_transitions,cooling_rate = \
    wrapper.run_one_params(**params)
d4 = cast_into_dic(column_info, data_transitions)


# In[5]:


figure(figsize=(10,7))
plot(d1[b'Eup'], d1[b'Tex'], marker='o', label='n(H$_2$)=$10^2$ cm$^{-3}$')
plot(d2[b'Eup'], d2[b'Tex'], marker='o', label='n(H$_2$)=$10^4$ cm$^{-3}$')
plot(d3[b'Eup'], d3[b'Tex'], marker='o', label='n(H$_2$)=$10^5$ cm$^{-3}$')
plot(d3[b'Eup'], d4[b'Tex'], marker='o', label='n(H$_2$)=$10^6$ cm$^{-3}$')
ax = gca()
ax.set_xlabel('E (K)')
ax.set_ylabel(r'$T_{\rm ex}$ (K)')
ax.set_xscale('log')
ax.set_yscale('symlog', linthreshy=10)
_ = legend(fontsize=15, framealpha=0.2, loc='best')
ax.grid(which='both')


# In[ ]: