import numpy as np
from astropy import units as u
import itertools
from operator import itemgetter
import pyspeckit
import scipy.stats

from pyspeckit.spectrum.models import ammonia, ammonia_constants

import pylab as pl
pl.close('all')
pl.figure(1).clf()

oneonefreq = ammonia_constants.freq_dict['oneone']
twotwofreq = ammonia_constants.freq_dict['twotwo']
# create an axis that covers the 1-1 and 2-2 inversion lines
xaxis = pyspeckit.units.SpectroscopicAxis(np.linspace(oneonefreq*(1-50/3e5),
                                                      twotwofreq*(1+50/3e5),
                                                      1000.),
                                          unit=u.Hz)
sigma = 2.
center = 0.
trot = 35.
ntot = 15.
# Adopting an NLTE model: T_ex < T_rot (this case is better for the default fitter)
# pyradex.Radex(species='p-nh3', column=10e14, collider_densities={'h2':10000}, temperature=35)()[8:10]
tex = {'oneone':16.4, 'twotwo':12.6}
synth_data = ammonia.ammonia(xaxis, trot=trot, tex=tex, width=sigma,
                             xoff_v=center, ntot=ntot)

# Add noise
stddev = 0.1
noise = np.random.randn(xaxis.size)*stddev
error = stddev*np.ones_like(synth_data)
data = noise+synth_data

# this will give a "blank header" warning, which is fine
sp = pyspeckit.Spectrum(data=data, error=error, xarr=xaxis,
                        xarrkwargs={'unit':'km/s'},
                        unit=u.K)

sp.plotter(figure=pl.figure(1), errstyle='fill')

# fit with some vague initial guesses and a fixed ortho/para
# fraction of 1 (it is unconstrained in our data)
sp.specfit(fittype='ammonia',
           #guesses=[20, 15, 15.5, 3, 2, 1],
           guesses=[30, 25, 15.5, 3, 2, 0],
           fixed=[False,False,False,False,False,True])

# do it again to set the errors neatly
sp.specfit(fittype='ammonia',
           #guesses=[20, 15, 15.5, 3, 2, 1],
           guesses=[30, 25, 15.5, 3, 2, 0],
           fixed=[False,False,False,False,False,True])


sp.plotter(errstyle='fill')
sp.specfit.plot_fit()


sp.plotter.savefig('oned_ammonia_nLTE_fit_example.pdf')


# do a deeper examination of the likelihood function

def chi2(sp, pars):
    """
    Given a spectrum and some parameters, calculate the chi^2 value
    """
    return ((sp.specfit.get_model_frompars(sp.xarr, pars) -
             sp.specfit.spectofit)**2 /
            (sp.specfit.errspec**2)
           ).sum()

def replace(lst, eltid, value):
    """
    Helper function to make new parameter value lists below
    """
    lstcopy = list(lst)
    lstcopy[eltid] = value
    return lstcopy


fig = pl.figure(2, figsize=(12,12))
fig.clf()

for ii,par in enumerate(sp.specfit.parinfo):
    if par.fixed:
        continue
    par_vals = np.linspace(par.value - par.error * 2,
                           par.value + par.error * 2,
                           10)

    par_likes = np.array([chi2(sp, replace(sp.specfit.parinfo.values, ii, val))
                          for val in par_vals])

    pct68 = scipy.stats.chi2.cdf(1, 1)

    # this is for the full 3-parameter case
    delta_chi2 = scipy.stats.chi2.ppf(pct68, sp.specfit.fitter.npars - sum(sp.specfit.parinfo.fixed))

    # this is the one we're actually using (which is just 1) because we're
    # marginalizing over the other parameters
    delta_chi2 = scipy.stats.chi2.ppf(pct68, 1)

    ax = fig.add_subplot(5,5,1+ii*6)
    ax.plot(par_vals, par_likes)
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    ax.hlines(par_likes.min() + delta_chi2, xmin, xmax, 'k', '--')
    ax.vlines([par.value - par.error,
               par.value + par.error,], ymin, ymax, 'k', '--')
    ax.set_ylabel("$\Delta \chi^2$")
    ax.set_xlabel("{0} Value".format(par.parname))
    pl.setp(ax.get_xticklabels(), rotation=90)

    ax.set_xlim(xmin, xmax,)
    ax.set_ylim(ymin, ymax,)

fig.tight_layout()


nsigma_range = 2
nsteps = 15.

plotinds = {0:6, 1:11, 2:12, 3:16, 4:17, 5:18, 6:21, 7:22, 8:23, 9:24}
for ii,(par1,par2) in enumerate(itertools.combinations(sp.specfit.parinfo[:5],2)):

    if par1.limited[0] and par1.value-par1.error*nsigma_range < par1.limits[0]:
        start1 = par1.limits[0]+par1.limits[0]*1e-3
    else:
        start1 = par1.value - par1.error*nsigma_range
    if par1.limited[1] and par1.value+par1.error*nsigma_range > par1.limits[1]:
        end1 = par1.limits[1]-par1.limits[1]*1e-3
    else:
        end1 = par1.value + par1.error*nsigma_range

    step1 = (end1 - start1) / nsteps

    if par2.limited[0] and par2.value-par2.error*nsigma_range < par2.limits[0]:
        start2 = par2.limits[0]+par2.limits[0]*1e-3
    else:
        start2 = par2.value - par2.error*nsigma_range
    if par2.limited[1] and par2.value+par2.error*nsigma_range > par2.limits[1]:
        end2 = par2.limits[1]-par2.limits[1]*1e-3
    else:
        end2 = par2.value + par2.error*nsigma_range

    step2 = (end2 - start2) / nsteps

    assert end1 > start1
    assert end2 > start2

    p1vals, p2vals = np.mgrid[start1:end1:step1,
                              start2:end2:step2]


    par_likes = np.array([chi2(sp, replace(replace(sp.specfit.parinfo.values, par1.n, val1), par2.n, val2))
                          for val1,val2 in zip(p1vals.ravel(), p2vals.ravel())]).reshape(p1vals.shape)

    pct68 = 1-scipy.stats.norm.sf(1)*2
    pct95 = 1-scipy.stats.norm.sf(2)*2
    pct997 = 1-scipy.stats.norm.sf(3)*2

    # marginalize over only one parameter; we now have two free parameters
    delta_chi2_68 = scipy.stats.chi2.ppf(pct68, 2)
    delta_chi2_95 = scipy.stats.chi2.ppf(pct95, 2)
    delta_chi2_997 = scipy.stats.chi2.ppf(pct997, 2)

    ax = fig.add_subplot(5,5,plotinds[ii])
    ax.contour(p1vals, p2vals, par_likes,
               levels=[par_likes.min()+delta_chi2_68,
                       par_likes.min()+delta_chi2_95,
                       par_likes.min()+delta_chi2_997])
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    ax.plot(par1.value, par2.value, 'kx')
    ax.hlines([par2.value - par2.error,
               par2.value + par2.error,], xmin, xmax, 'k', '--')
    ax.vlines([par1.value - par1.error,
               par1.value + par1.error,], ymin, ymax, 'k', '--')
    ax.set_xlabel("{0} Value".format(par1.parname))
    ax.set_ylabel("{0} Value".format(par2.parname))
    pl.setp(ax.get_xticklabels(), rotation=90)

    ax.set_xlim(xmin, xmax,)
    ax.set_ylim(ymin, ymax,)

fig.tight_layout()

fig.savefig("ammonia_nLTE_default_error_estimate_demonstration.pdf")


# fit with restricted delta-T instead of restricted T

sp.specfit.Registry.add_fitter('ammonia_dtex',
                               ammonia.ammonia_model_restricted_tex(), 7,
                               multisingle='multi')

sp.specfit(fittype='ammonia_dtex',
           guesses=[30, 25, 15.5, 3, 2, 0, 5],
           fixed=[False,False,False,False,False,True,False])

sp.plotter.savefig('oned_ammonia_nLTE_fit_example_deltaT.pdf')

# do a deeper examination of the likelihood function

def chi2(sp, pars):
    """
    Given a spectrum and some parameters, calculate the chi^2 value
    """
    return ((sp.specfit.get_model_frompars(sp.xarr, pars) -
             sp.specfit.spectofit)**2 /
            (sp.specfit.errspec**2)
           ).sum()

def replace(lst, eltid, value):
    """
    Helper function to make new parameter value lists below
    """
    lstcopy = list(lst)
    lstcopy[eltid] = value
    return lstcopy


fig = pl.figure(2, figsize=(12,12))
fig.clf()

nsteps = 10
dd = 0
for ii,par in enumerate(sp.specfit.parinfo):
    if par.fixed or par.tied:
        dd = dd + 1
        continue
    par_vals = np.linspace(par.value - par.error * 2,
                           par.value + par.error * 2,
                           nsteps)

    if par.parname == 'delta0':
        start, step, end = 0, 5/nsteps, 25
        par_vals = np.linspace(start, end, nsteps)

    par_likes = np.array([chi2(sp, replace(sp.specfit.parinfo.values, ii, val))
                          for val in par_vals])

    pct68 = scipy.stats.chi2.cdf(1, 1)

    # this is for the full 3-parameter case
    delta_chi2 = scipy.stats.chi2.ppf(pct68, sp.specfit.fitter.npars - sum(sp.specfit.parinfo.fixed))

    # this is the one we're actually using (which is just 1) because we're
    # marginalizing over the other parameters
    delta_chi2 = scipy.stats.chi2.ppf(pct68, 1)

    ax = fig.add_subplot(5,5,1+(ii-dd)*6)
    ax.plot(par_vals, par_likes)
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    ax.hlines(par_likes.min() + delta_chi2, xmin, xmax, 'k', '--')
    ax.vlines([par.value - par.error,
               par.value + par.error,], ymin, ymax, 'k', '--')
    ax.set_ylabel("$\Delta \chi^2$")
    ax.set_xlabel("{0} Value".format(par.parname))
    pl.setp(ax.get_xticklabels(), rotation=90)

    ax.set_xlim(xmin, xmax,)
    ax.set_ylim(ymin, ymax,)

fig.tight_layout()


nsigma_range = 2
nsteps = 15.


plotinds = {0:6, 1:11, 2:12, 3:16, 4:17, 5:18, 6:21, 7:22, 8:23, 9:24}
for ii,(par1,par2) in enumerate(itertools.combinations(itemgetter(0,2,3,4,6)(sp.specfit.parinfo),2)):

    if par1.limited[0] and par1.value-par1.error*nsigma_range < par1.limits[0]:
        start1 = par1.limits[0]+par1.limits[0]*1e-3
    else:
        start1 = par1.value - par1.error*nsigma_range
    if par1.limited[1] and par1.value+par1.error*nsigma_range > par1.limits[1]:
        end1 = par1.limits[1]-par1.limits[1]*1e-3
    else:
        end1 = par1.value + par1.error*nsigma_range

    step1 = (end1 - start1) / nsteps

    if par2.limited[0] and par2.value-par2.error*nsigma_range < par2.limits[0]:
        start2 = par2.limits[0]+par2.limits[0]*1e-3
    else:
        start2 = par2.value - par2.error*nsigma_range
    if par2.limited[1] and par2.value+par2.error*nsigma_range > par2.limits[1]:
        end2 = par2.limits[1]-par2.limits[1]*1e-3
    else:
        end2 = par2.value + par2.error*nsigma_range

    step2 = (end2 - start2) / nsteps

    if par1.parname == 'delta0':
        start1, step1, end1 = 0, 5/nsteps, 25
    if par2.parname == 'delta0':
        start2, step2, end2 = 0, 5/nsteps, 25

    assert end1 > start1
    assert end2 > start2

    p1vals, p2vals = np.mgrid[start1:end1:step1,
                              start2:end2:step2]


    par_likes = np.array([chi2(sp, replace(replace(sp.specfit.parinfo.values, par1.n, val1), par2.n, val2))
                          for val1,val2 in zip(p1vals.ravel(), p2vals.ravel())]).reshape(p1vals.shape)

    pct68 = 1-scipy.stats.norm.sf(1)*2
    pct95 = 1-scipy.stats.norm.sf(2)*2
    pct997 = 1-scipy.stats.norm.sf(3)*2

    # marginalize over only one parameter; we now have two free parameters
    delta_chi2_68 = scipy.stats.chi2.ppf(pct68, 2)
    delta_chi2_95 = scipy.stats.chi2.ppf(pct95, 2)
    delta_chi2_997 = scipy.stats.chi2.ppf(pct997, 2)

    ax = fig.add_subplot(5,5,plotinds[ii])
    ax.contour(p1vals, p2vals, par_likes,
               levels=[par_likes.min()+delta_chi2_68,
                       par_likes.min()+delta_chi2_95,
                       par_likes.min()+delta_chi2_997])
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    ax.plot(par1.value, par2.value, 'kx')
    ax.hlines([par2.value - par2.error,
               par2.value + par2.error,], xmin, xmax, 'k', '--')
    ax.vlines([par1.value - par1.error,
               par1.value + par1.error,], ymin, ymax, 'k', '--')
    ax.set_xlabel("{0} Value".format(par1.parname))
    ax.set_ylabel("{0} Value".format(par2.parname))
    pl.setp(ax.get_xticklabels(), rotation=90)

    ax.set_xlim(xmin, xmax,)
    ax.set_ylim(ymin, ymax,)

fig.tight_layout()

fig.savefig("ammonia_nLTE_restricteddelta_error_estimate_demonstration.pdf")
