import h3ppy
import matplotlib.pyplot as plt
import numpy as np

# Create the H3+ object
h3p = h3ppy.h3p()

# Define a wavelength range, e.g. typical of an observation of the H3+ Q branch
# Specify the start and end wavelengths, and the number of wavelength elements
wave = h3p.wavegen(3.94, 4.03, 1024)

# Create a H3+ model spectrum for a set of physical parameters 
# Spectral resolution R = 1200, T = 1000, N = 1e14
# This is the minimum set of parameters required for generating a model
model = h3p.model(density = 1e14, temperature = 1000, R = 1000, wavelength = wave)

# Plot the model
fig, ax = plt.subplots()
ax.plot(wave, model)
# Automagically set the labels 
ax.set_xlabel(h3p.xlabel())
ax.set_ylabel(h3p.ylabel())
plt.savefig('example_model.png')
plt.close() 

# Generate some noise to add to the model  
noise = np.random.normal(size = model.size) * np.max(model) / 50
pretend_data = model + noise

# Set the initial guess for the data. I'm making it different from the model input
# to show that the fit actually can converge on 'real' values
h3p.set(density = 1e13, temperature = 1300, data = pretend_data)

# Let h3ppy make an initial guess at the density
h3p.guess_density()

# Fit temperature and density to the pretend data
fit = h3p.fit()

# Get the fit variables and associated errors
vars, errs = h3p.get_results()

# Plot the model
fig, ax = plt.subplots()
ax.plot(wave, pretend_data, label='Pretend data')
ax.plot(wave, fit, label = '$H_3^+$ spectrum fit')
ax.legend()
# Automagically set the labels 
ax.set_xlabel(h3p.xlabel())
ax.set_ylabel(h3p.ylabel())
plt.savefig('../img/example_fit.png')
plt.close() 

import matplotlib.pyplot as plt
import numpy as np
import h3ppy

# Read the UKIRT CGS4 data from Trafton et al., (1993)
data_file = 'cgs4_uranus_u01apr92.txt'
types = {'names' : ( 'w', 'i' ), 'formats' : ('f', 'f') }
dat = np.loadtxt(data_file, skiprows=4, dtype = types)

# Need to convert the instrument FOV to units of sterradian
# The pixel width is 3.1 arcsec with a slit widht of 3.1 arcsec
# Note - there are 4.2545e10 arceconds in a sterradian
spectrum = dat['i'] * 4.2545e10 / (3.1 * 3.1)
wave = dat['w']

# Make our h3ppy object :-) 
h3p = h3ppy.h3p()

# Set the wavelength and data, and use the spectral resulution to input the 
# expected line width
h3p.set(wavelength = wave, data = spectrum, R = 1300)

# We need to guess a temperature
h3p.set(temperature = 1000)

# Let h3ppy try and guess a wavelength offset
guess = h3p.guess_offset()

# Guess the density - this'll effectively scale the spectrum to the observed spectrum
# It really is not a measure of the actual density!  
guess = h3p.guess_density()

# Let h3ppy do the fitting - this will do a full five parameter fit
fit = h3p.fit(verbose = False)    

# Get the results
vars, errs = h3p.get_results()

# Plot the results! 
fig, ax = plt.subplots()
ax.set_title('First ever H$_3^+$ spectrum from Uranus: Trafton et al. (1993)')
ax.plot(wave, spectrum * 1e6, 'o', label = 'Original CGS4 Uranus spectrum')
ax.plot(wave, fit * 1e6, label = 'h3ppy fit to data')
ax.legend(frameon = False)

# Use the h3ppy helper functions for the labels
ax.set_xlabel(h3p.xlabel())
ax.set_ylabel(h3p.ylabel(prefix = '$\mu$'))
plt.savefig('../img/cgs4_uranus_fit.png')
plt.close()

import matplotlib.pyplot as plt
import numpy as np
import h3ppy

# Read the NASA IRTF CGS4 data from Trafton et al., (1993)
data_file = 'nirspec_jupiter.txt'
types = {'names' : ( 'w', 'i' ), 'formats' : ('f', 'f') }
dat = np.loadtxt(data_file, skiprows=4, dtype = types)
wave = dat['w']
spec = dat['i']

# Create the h3ppy object feed data into it
# The spectral resolution of NIRSPEC is ~20k
h3p = h3ppy.h3p()

# Plot the observation
title = 'Keck NIRSPEC H$_3^+$ spectrum of the southern aurora of Jupiter'
fig, ax = plt.subplots()
ax.plot(wave, spec * 1e3)
ax.set(xlabel = h3p.xlabel(), ylabel = h3p.ylabel(prefix = 'm'), title = title)
plt.tight_layout()
plt.savefig('../img/nirspec_jupiter_data.png')
# plt.show()
plt.close()

# This function sub-divides data centered on a list of wavelengths
def subdivide(wave, spec, middles, width = 20) : 
    ret = []
    for m in middles : 
        centre = np.abs(wave - m).argmin()
        for i in range(centre - width, centre + width) : 
            ret.append(spec[i])
    return np.array(ret)
    
# The H3+ line centeres contained withing this spectral band
centers = [3.953, 3.971, 3.986, 3.9945]
cpos = np.arange(4) * 41 + 20

# Create sub-arrays, focusing on where the H3+ lines are
subspec = subdivide(wave, spec, centers)
subwave = subdivide(wave, wave, centers)

# Set the wavelength and the data
h3p.set(wavelength = subwave, data = subspec, R = 20000)

# Create a x scale for plotting 
xx      = range(len(subspec))

# Guess the density and proceed with a five parameter fit
h3p.guess_density()
fit = h3p.fit()
vars, errs = h3p.get_results()

# Plot the fit
fig, ax = plt.subplots()
ax.plot(xx, subspec * 1e3, '.', label = 'Observation')
ax.plot(xx, fit * 1e3, label = 'h3ppy H$_3^+$ fit')
ax.set(xlabel = h3p.xlabel(), ylabel = h3p.ylabel(prefix = 'm'), xticks = cpos, title=title)
ax.set_xticklabels(centers)
ax.legend(frameon = False)
plt.tight_layout()
plt.savefig('../img/nirspec_jupiter_fit.png')
plt.close()

import h3ppy
import matplotlib.pyplot as plt
import numpy as np

# Instrument resolution
R = 300

# Temperature of the thermosphere
T = 900

# Set up the H3+ model
h3p = h3ppy.h3p()
h3p.set(temperature = T, density = 2e15, R = R)
wave = h3p.wavegen(1.8, 4.2, 1000)

# Set up the H2 model
amagat = 2.76e25
h2 = h3ppy.h2()
h2.set(temperature = T, density = amagat, R = R, wavelength = wave)

# Generate models
model_h3p = h3p.model()
model_h2  = h2.model()

# Plot the result
fig, ax = plt.subplots()
ax.plot(wave, model_h3p * 1e6, label = 'H$_3^+$ model')
ax.plot(wave, model_h2 * 1e6, label = 'H$_2$ model')
ax.set(ylabel = h3p.ylabel(prefix = '$\mu$'), xlabel = h3p.xlabel())
ax.legend(frameon = False)
plt.save('img/h2_h3p_spectrum.png')

h3p.set(nsigma = 3, sigma_0 = 0.1, sigma_1 = 0.01, sigma_2 = 0.001) 

FWHM = 2 * np.sqrt(2 * np.log(2)) * sigma = 2.35482 * sigma