"""
===========
NH2D fitter: ortho- and para- in the same file, but not modeled together
===========
Reference for line params:

New line parameters are taken from the recent laboratory work:
Melosso et al. (2021) 
Journal of Molecular Spectroscopy, vol. 377, March 2021, 111431
https://doi.org/10.1016/j.jms.2021.111431

These superseed the frequencies from F. Daniel et al. (2016)
http://adsabs.harvard.edu/abs/2016A%26A...586L...4D

"""
from . import hyperfine
import astropy.units as u

freq_dict_cen = {
                'o-1_11-1_01':  85.926263e9,
                'p-1_11-1_01': 110.153599e9,
                'o-1_01-0_00': 332.82251e9,
                'p-1_01-0_00': 332.78189e9,
                }

freq_dict = {
    ####### ortho-NH2D J=1_11-1_01
    'o-1_11-1_01_01': 85.9246990e9,
    'o-1_11-1_01_02': 85.9247505e9,
    'o-1_11-1_01_03': 85.9247772e9,
    'o-1_11-1_01_04': 85.9252797e9,
    'o-1_11-1_01_05': 85.9253576e9,
    'o-1_11-1_01_06': 85.9256502e9,
    'o-1_11-1_01_07': 85.9256695e9,
    'o-1_11-1_01_08': 85.9256886e9,
    'o-1_11-1_01_09': 85.9256962e9,
    'o-1_11-1_01_10': 85.9257017e9,
    'o-1_11-1_01_11': 85.9257285e9,
    'o-1_11-1_01_12': 85.9261987e9,
    'o-1_11-1_01_13': 85.9262178e9,
    'o-1_11-1_01_14': 85.9262310e9,
    'o-1_11-1_01_15': 85.9262475e9,
    'o-1_11-1_01_16': 85.9262517e9,
    'o-1_11-1_01_17': 85.9262708e9,
    'o-1_11-1_01_18': 85.9262766e9,
    'o-1_11-1_01_19': 85.9262824e9,
    'o-1_11-1_01_20': 85.9262914e9,
    'o-1_11-1_01_21': 85.9262990e9,
    'o-1_11-1_01_22': 85.9263088e9,
    'o-1_11-1_01_23': 85.9263091e9,
    'o-1_11-1_01_24': 85.9263257e9,
    'o-1_11-1_01_25': 85.9268116e9,
    'o-1_11-1_01_26': 85.9268282e9,
    'o-1_11-1_01_27': 85.9268645e9,
    'o-1_11-1_01_28': 85.9268718e9,
    'o-1_11-1_01_29': 85.9268895e9,
    'o-1_11-1_01_30': 85.9269061e9,
    'o-1_11-1_01_31': 85.9271096e9,
    'o-1_11-1_01_32': 85.9271418e9,
    'o-1_11-1_01_33': 85.9277048e9,
    'o-1_11-1_01_34': 85.9277224e9,
    'o-1_11-1_01_35': 85.9277391e9,
    ####### ortho-NH2D J=1_01-0_00
    'o-1_01-0_00_01': 332.7808404e9,
    'o-1_01-0_00_02': 332.7808404e9,
    'o-1_01-0_00_03': 332.7808404e9,
    'o-1_01-0_00_04': 332.7816734e9,
    'o-1_01-0_00_05': 332.7816734e9,
    'o-1_01-0_00_06': 332.7816734e9,
    'o-1_01-0_00_07': 332.7816983e9,
    'o-1_01-0_00_08': 332.7817512e9,
    'o-1_01-0_00_09': 332.7817512e9,
    'o-1_01-0_00_10': 332.7822538e9,
    'o-1_01-0_00_11': 332.7822538e9,
    'o-1_01-0_00_12': 332.7822538e9,
    'o-1_01-0_00_13': 332.7822805e9,
    'o-1_01-0_00_14': 332.7822805e9,
    'o-1_01-0_00_15': 332.7823320e9,
    ####### para-NH2D J=1_11-1_01
    'p-1_11-1_01_01': 110.1520107e9,
    'p-1_11-1_01_02': 110.1520630e9,
    'p-1_11-1_01_03': 110.1520892e9,
    'p-1_11-1_01_04': 110.1525927e9,
    'p-1_11-1_01_05': 110.1526707e9,
    'p-1_11-1_01_06': 110.1529617e9,
    'p-1_11-1_01_07': 110.1529820e9,
    'p-1_11-1_01_08': 110.1530013e9,
    'p-1_11-1_01_09': 110.1530081e9,
    'p-1_11-1_01_10': 110.1530141e9,
    'p-1_11-1_01_11': 110.1530402e9,
    'p-1_11-1_01_12': 110.1535116e9,
    'p-1_11-1_01_13': 110.1535309e9,
    'p-1_11-1_01_14': 110.1535437e9,
    'p-1_11-1_01_15': 110.1535596e9,
    'p-1_11-1_01_16': 110.1535645e9,
    'p-1_11-1_01_17': 110.1535838e9,
    'p-1_11-1_01_18': 110.1535897e9,
    'p-1_11-1_01_19': 110.1535945e9,
    'p-1_11-1_01_20': 110.1536042e9,
    'p-1_11-1_01_21': 110.1536119e9,
    'p-1_11-1_01_22': 110.1536207e9,
    'p-1_11-1_01_23': 110.1536217e9,
    'p-1_11-1_01_24': 110.1536381e9,
    'p-1_11-1_01_25': 110.1541241e9,
    'p-1_11-1_01_26': 110.1541415e9,
    'p-1_11-1_01_27': 110.1541770e9,
    'p-1_11-1_01_28': 110.1541857e9,
    'p-1_11-1_01_29': 110.1542022e9,
    'p-1_11-1_01_30': 110.1542196e9,
    'p-1_11-1_01_31': 110.1544234e9,
    'p-1_11-1_01_32': 110.1544555e9,
    'p-1_11-1_01_33': 110.1550195e9,
    'p-1_11-1_01_34': 110.1550360e9,
    'p-1_11-1_01_35': 110.1550534e9,
    ####### para-NH2D J=1_01-0_00
    'p-1_01-0_00_01': 332.8215574e9,
    'p-1_01-0_00_02': 332.8215574e9,
    'p-1_01-0_00_03': 332.8215574e9,
    'p-1_01-0_00_04': 332.8223912e9,
    'p-1_01-0_00_05': 332.8223912e9,
    'p-1_01-0_00_06': 332.8223912e9,
    'p-1_01-0_00_07': 332.8224164e9,
    'p-1_01-0_00_08': 332.8224692e9,
    'p-1_01-0_00_09': 332.8224692e9,
    'p-1_01-0_00_10': 332.8229727e9,
    'p-1_01-0_00_11': 332.8229727e9,
    'p-1_01-0_00_12': 332.8229727e9,
    'p-1_01-0_00_13': 332.8229989e9,
    'p-1_01-0_00_14': 332.8229989e9,
    'p-1_01-0_00_15': 332.8230512e9,
    }

line_strength_dict = {
    ####### ortho-NH2D J=1_11-1_01
    'o-1_11-1_01_01': 0.01295,
    'o-1_11-1_01_02': 0.06163,
    'o-1_11-1_01_03': 0.03552,
    'o-1_11-1_01_04': 0.00033,
    'o-1_11-1_01_05': 0.00068,
    'o-1_11-1_01_06': 0.01377,
    'o-1_11-1_01_07': 0.00730,
    'o-1_11-1_01_08': 0.06915,
    'o-1_11-1_01_09': 0.02936,
    'o-1_11-1_01_10': 0.00037,
    'o-1_11-1_01_11': 0.02030,
    'o-1_11-1_01_12': 0.09879,
    'o-1_11-1_01_13': 0.01726,
    'o-1_11-1_01_14': 0.01881,
    'o-1_11-1_01_15': 0.01031,
    'o-1_11-1_01_16': 0.02297,
    'o-1_11-1_01_17': 0.17283,
    'o-1_11-1_01_18': 0.02660,
    'o-1_11-1_01_19': 0.03280,
    'o-1_11-1_01_20': 0.00662,
    'o-1_11-1_01_21': 0.01394,
    'o-1_11-1_01_22': 0.05729,
    'o-1_11-1_01_23': 0.01174,
    'o-1_11-1_01_24': 0.00757,
    'o-1_11-1_01_25': 0.01545,
    'o-1_11-1_01_26': 0.03456,
    'o-1_11-1_01_27': 0.06345,
    'o-1_11-1_01_28': 0.01973,
    'o-1_11-1_01_29': 0.00112,
    'o-1_11-1_01_30': 0.00568,
    'o-1_11-1_01_31': 0.00018,
    'o-1_11-1_01_32': 0.00057,
    'o-1_11-1_01_33': 0.01069,
    'o-1_11-1_01_34': 0.06063,
    'o-1_11-1_01_35': 0.03905,
    ####### ortho-NH2D J=1_01-0_00
    'o-1_01-0_00_01': 0.01579,
    'o-1_01-0_00_02': 0.03206,
    'o-1_01-0_00_03': 0.06326,
    'o-1_01-0_00_04': 0.04784,
    'o-1_01-0_00_05': 0.05919,
    'o-1_01-0_00_06': 0.00409,
    'o-1_01-0_00_07': 0.25925,
    'o-1_01-0_00_08': 0.14819,
    'o-1_01-0_00_09': 0.03699,
    'o-1_01-0_00_10': 0.04748,
    'o-1_01-0_00_11': 0.01986,
    'o-1_01-0_00_12': 0.04376,
    'o-1_01-0_00_13': 0.03699,
    'o-1_01-0_00_14': 0.14819,
    'o-1_01-0_00_15': 0.03704,
    ####### para-NH2D J=1_11-1_01
    'p-1_11-1_01_01': 0.01296,
    'p-1_11-1_01_02': 0.06162,
    'p-1_11-1_01_03': 0.03552,
    'p-1_11-1_01_04': 0.00033,
    'p-1_11-1_01_05': 0.00068,
    'p-1_11-1_01_06': 0.01378,
    'p-1_11-1_01_07': 0.00730,
    'p-1_11-1_01_08': 0.06914,
    'p-1_11-1_01_09': 0.02937,
    'p-1_11-1_01_10': 0.00037,
    'p-1_11-1_01_11': 0.02030,
    'p-1_11-1_01_12': 0.09872,
    'p-1_11-1_01_13': 0.01729,
    'p-1_11-1_01_14': 0.01882,
    'p-1_11-1_01_15': 0.01029,
    'p-1_11-1_01_16': 0.02300,
    'p-1_11-1_01_17': 0.17283,
    'p-1_11-1_01_18': 0.02660,
    'p-1_11-1_01_19': 0.03282,
    'p-1_11-1_01_20': 0.00661,
    'p-1_11-1_01_21': 0.01395,
    'p-1_11-1_01_22': 0.01173,
    'p-1_11-1_01_23': 0.05728,
    'p-1_11-1_01_24': 0.00757,
    'p-1_11-1_01_25': 0.01547,
    'p-1_11-1_01_26': 0.03455,
    'p-1_11-1_01_27': 0.06342,
    'p-1_11-1_01_28': 0.01973,
    'p-1_11-1_01_29': 0.00113,
    'p-1_11-1_01_30': 0.00570,
    'p-1_11-1_01_31': 0.00018,
    'p-1_11-1_01_32': 0.00057,
    'p-1_11-1_01_33': 0.01069,
    'p-1_11-1_01_34': 0.06062,
    'p-1_11-1_01_35': 0.03906,
    ####### para-NH2D J=1_01-0_00
    'p-1_01-0_00_01': 0.01580,
    'p-1_01-0_00_02': 0.03207,
    'p-1_01-0_00_03': 0.06324,
    'p-1_01-0_00_04': 0.04784,
    'p-1_01-0_00_05': 0.05919,
    'p-1_01-0_00_06': 0.00409,
    'p-1_01-0_00_07': 0.25927,
    'p-1_01-0_00_08': 0.14814,
    'p-1_01-0_00_09': 0.03704,
    'p-1_01-0_00_10': 0.04747,
    'p-1_01-0_00_11': 0.01985,
    'p-1_01-0_00_12': 0.04379,
    'p-1_01-0_00_13': 0.03704,
    'p-1_01-0_00_14': 0.14814,
    'p-1_01-0_00_15': 0.03704,
    }


# Get offset velocity dictionary in km/s based on the lines frequencies and rest frequency
conv_o1_1 = u.doppler_radio(freq_dict_cen['o-1_11-1_01']*u.Hz)
conv_p1_1 = u.doppler_radio(freq_dict_cen['p-1_11-1_01']*u.Hz)
conv_o1_0 = u.doppler_radio(freq_dict_cen['o-1_01-0_00']*u.Hz)
conv_p1_0 = u.doppler_radio(freq_dict_cen['p-1_01-0_00']*u.Hz)

voff_lines_dict = {
    name: ((freq_dict[name]*u.Hz).to(u.km/u.s, equivalencies=conv_o1_1).value) for name in freq_dict.keys() if "o-1_11-1_01" in name
    }
voff_lines_dict.update({
	name: ((freq_dict[name]*u.Hz).to(u.km/u.s, equivalencies=conv_p1_1).value) for name in freq_dict.keys() if "p-1_11-1_01" in name
    })
voff_lines_dict.update({
	name: ((freq_dict[name]*u.Hz).to(u.km/u.s, equivalencies=conv_o1_0).value) for name in freq_dict.keys() if "o-1_01-0_00" in name
    })
voff_lines_dict.update({
	name: ((freq_dict[name]*u.Hz).to(u.km/u.s, equivalencies=conv_p1_0).value) for name in freq_dict.keys() if "p-1_01-0_00" in name
    })

# relative_strength_total_degeneracy is not used in the CLASS implementation
# of the hfs fit. It is the sum of the degeneracy values for all hyperfines
# for a given line; it gives the relative weights between lines.
# Hyperfine weights are treated as normalized within one rotational transition.
wo1_1 = sum(val for name,val in line_strength_dict.items() if 'o-1_11-1_01' in name)
wp1_1 = sum(val for name,val in line_strength_dict.items() if 'p-1_11-1_01' in name)
wo1_0 = sum(val for name,val in line_strength_dict.items() if 'o-1_01-0_00' in name)
wp1_0 = sum(val for name,val in line_strength_dict.items() if 'p-1_01-0_00' in name)
relative_strength_total_degeneracy = {
    name : wo1_1 for name  in line_strength_dict.keys() if "o-1_11-1_01" in name
    }
relative_strength_total_degeneracy.update({
    name : wp1_1 for name  in line_strength_dict.keys() if "p-1_11-1_01" in name
    })
relative_strength_total_degeneracy.update({
    name : wo1_0 for name  in line_strength_dict.keys() if "o-1_01-0_00" in name
    })
relative_strength_total_degeneracy.update({
    name : wp1_0 for name  in line_strength_dict.keys() if "p-1_01-0_00" in name
    })

# Get the list of line names from the previous lists
line_names = [name for name in voff_lines_dict.keys()]

nh2d_vtau = hyperfine.hyperfinemodel(line_names, voff_lines_dict, freq_dict,
                                     line_strength_dict,
                                     relative_strength_total_degeneracy)
nh2d_vtau_fitter = nh2d_vtau.fitter
nh2d_vtau_vheight_fitter = nh2d_vtau.vheight_fitter
nh2d_vtau_tbg_fitter = nh2d_vtau.background_fitter
