"""
ADDED CONSTRAINT: density > 10^4
(actually, density > dust-derived mean density, or 10^3, whichever is higher)

For each of the fitted spectra from individual_spectra.py, use the fitted ratio
(and appropriate ancillary data, e.g. h2 column) to derive the best fit
temperature, etc.

Store them in an output table and plot by calling 
execfile(paths.pcpath('parameter_comparisons.py'))
"""
import paths
import os
import matplotlib
#matplotlib.rc_file(paths.pcpath('pubfiguresrc'))
import pylab as pl
import numpy as np
from scipy import stats
from astropy import log
from astropy import units as u
from astropy import constants
from astropy import coordinates
from paths import analysispath, tpath
from pyspeckit_fitting import (texgrid303, taugrid303, texgrid321, taugrid321,
                               texgrid322, taugrid322, hdr)

# use the local constrain_parameters; h2co_modeling's version is... not exactly up to date.
from constrain_parameters import paraH2COmodel
from h2co_modeling import grid_fitter
from astropy import table
import despotic_heating as heating


pl.rcParams['font.size'] = 16.0
pl.close(4)

# mf means modelfitter
mf = paraH2COmodel()

fittable = table.Table.read(tpath("fitted_line_parameters.ipac"),
                            format='ascii.ipac')
fittable.add_columns([table.Column(name=name, dtype='float', length=len(fittable))
                      for name in ['temperature_chi2','tmin1sig_chi2','tmax1sig_chi2',
                                   'expected_temperature',
                                   'column_chi2','cmin1sig_chi2','cmax1sig_chi2',
                                   'expected_column',
                                   'density_chi2','dmin1sig_chi2','dmax1sig_chi2',
                                   'expected_density',
                                   'logh2column','elogh2column',
                                   'logabundance','elogabundance',
                                   'tkin_turb', 'reff_pc',
                                  ]])

if not os.path.exists(paths.fpath('param_fits')):
    os.makedirs(paths.fpath('param_fits'))

nlevs = 4
#levels = [stats.norm.cdf(ii)-stats.norm.cdf(-ii)
#           for ii in range(1,nlevs)]
ndof = 3
levels = ([0]+
          [stats.chi2.ppf(stats.norm.cdf(ii)-stats.norm.cdf(-ii),
                          ndof)
           for ii in range(1,nlevs)])

density_label = 'Density $n(\mathrm{H}_2)$ [log cm$^{-3}$]'
column_label = 'Column p-H$_2$CO [log cm$^{-2}$/(km s$^{-1}$ pc)]'
density_label_short = "$n(\mathrm{H}_2) (\mathrm{cm}^{-3})$"
column_label_short = "$N(\mathrm{p-H}_2\mathrm{CO}) (\mathrm{cm}^{-2})$"
temperature_label = 'Temperature (K)'
prevname = ''
num = 0

for row in fittable:
    if row['Source_Name'] == prevname:
        num += 1
    else:
        num = 0
        prevname = row['Source_Name']
    log.info("Fitting {0}_{1}".format(row['Source_Name'],num))
    logh2column = np.log10(row['higalcolumndens'])
    elogh2column = 1.0
    linewidth = row['spline_width']
    elinewidth = row['espline_width']

    par1 = row['spline_ampH2CO']
    epar1 = row['espline_ampH2CO']
    par2 = row['spline_ampH2CO']*row['spline_h2coratio321303']
    epar2 = row['spline_ampH2CO']*row['espline_h2coratio321303']
    #match,indbest,chi2b = grid_fitter.grid_2p_getmatch(par1, epar1, tline303,
    #                                                   par2, epar2, tline321)
    ratio = row['spline_h2coratio321303']
    eratio = row['espline_h2coratio321303']
    ratio2 = row['spline_h2coratio322321']
    eratio2 = row['espline_h2coratio322321']

    if ratio == 0 or np.isnan(ratio):
        print("Skipped {0} because it had no fit ratio".format(row['Source_Name']))
        continue

    # We can impose a "loose" abundance constraint
    # Given that we know the H2 density, and the line width is ~5-10 km/s,
    # abundance = column / pc / density
    # We'll say abundance = 1.2e9 with error 0.6e9
    # Or, log(abundance) = log(1.2e9) +/- 1
    logabundance = np.log10(1.2e-9)
    elogabundance = 1.0

    # Constraint from density
    r_deg = (row['area']/np.pi)**0.5 * u.deg
    reff = (r_deg*(8.5*u.kpc)).to(u.pc, u.dimensionless_angles())
    mass = ((10**logh2column*u.cm**-2)*np.pi*reff**2*2.8*constants.m_p).to(u.M_sun)
    # reff = mean of major and minor axes of elliptical.  This is not a bad
    # assumption... we COULD use the HWHM instead here, though.
    density = (mass/(4/3.*np.pi*reff**3)/constants.m_p/2.8).to(u.cm**-3)
    mindens = np.log10(density.value)
    if mindens < 3:
        mindens = 3

    # Combined abundance + total column constraint
    # N(H2CO) * dv * X = N(H2)
    # We are effectively ignoring errors in the linewidth here:
    # (noop - see chi2_h2)

    # Even though the lines are subject to filling-factor uncertainty, we can
    # set a *minimum* brightness in the models.  Given that we observe a line
    # brightness T_A, the true brightness is T_B = T_A/ff, where ff<1 by
    # definition
    # We therefore *increase* the chi^2 value wherever the model is fainter
    # than the line, enforcing a soft lower limit

    mf.set_constraints(ratio321303=ratio, eratio321303=eratio,
                       ratio321322=ratio2, eratio321322=eratio2,
                       logh2column=logh2column, elogh2column=elogh2column,
                       logabundance=logabundance, elogabundance=elogabundance,
                       taline303=par1, etaline303=epar1,
                       taline321=par2, etaline321=epar2,
                       mindens=mindens,
                       linewidth=10) # for consistency with dendro


    chi2r = mf.chi2_r321303
    chi2r2 = mf.chi2_r321322
    chi2_h2 = mf.chi2_h2
    chi2X = mf.chi2_X
    chi2_1 = mf.chi2_ff1
    chi2_2 = mf.chi2_ff2
    chi2_ff = chi2_1+chi2_2
    chi2_dens = mf.chi2_dens
    chi2b = chi2r + chi2_ff + chi2X + chi2_h2 + chi2_dens

    match = chi2b < 1
    indbest,match = grid_fitter.getmatch(chi2b, match)

    sh = match.shape
    (zz,yy,xx) = np.unravel_index(indbest, sh)

    fig1 = pl.figure(1, figsize=(12,12))
    fig1.clf()

    vmin = np.max([mf.tline303.min(), 0.1])
    vmax = np.min([mf.tline303.max(), par1+10])
    ax1 = pl.subplot(3,3,1)
    im1 = pl.imshow(mf.tline303[zz,:,:], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              extent=mf.crange+mf.drange, vmin=vmin, vmax=vmax)
    pl.contour(mf.carr, mf.darr, chi2b[zz,:,:], levels=chi2b.min()+levels)
    # only label middle row
    #pl.ylabel(column_label)
    pl.xlabel(column_label_short)
    ax1.set_xticks(np.arange(11,16))

    ax2 = pl.subplot(3,3,2)
    im2 = pl.imshow(mf.tline303[:,yy,:], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              aspect=np.diff(mf.crange)/np.diff(mf.trange),
              extent=mf.crange+mf.trange, vmin=vmin, vmax=vmax)
    pl.contour(mf.carr, mf.tarr, chi2b[:,yy,:], levels=chi2b.min()+levels)
    pl.xlabel(column_label_short)
    ax2.set_xticks(np.arange(11,16))
    pl.ylabel(temperature_label)
    #ax2.set_title("p-H$_2$CO $3_{0,3}-2_{0,2}$")

    ax3 = pl.subplot(3,3,3)
    im3 = pl.imshow(mf.tline303[:,:,xx], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              aspect=np.diff(mf.drange)/np.diff(mf.trange),
              extent=mf.drange+mf.trange, vmin=vmin, vmax=vmax)
    pl.contour(mf.darr, mf.tarr, chi2b[:,:,xx], levels=chi2b.min()+levels)
    pl.xlabel(density_label_short)
    ax3.xaxis.set_ticks(np.arange(mf.darr.min(), mf.darr.max()))
    pl.ylabel(temperature_label)
    cax = fig1.add_axes([0.91,0.68,0.02,0.22])
    cb = fig1.colorbar(mappable=im3, cax=cax, ax=ax2)
    cb.set_label("$T_B$ (p-H$_2$CO $3_{0,3}-2_{0,2}$)")

    vmin = np.max([mf.tline321.min(), 0.1])
    vmax = np.min([mf.tline321.max(), par2+10])
    ax4 = pl.subplot(3,3,4)
    pl.imshow(mf.tline321[zz,:,:], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              extent=mf.crange+mf.drange, vmin=vmin, vmax=vmax)
    pl.contour(mf.carr, mf.darr, chi2b[zz,:,:], levels=chi2b.min()+levels)
    pl.xlabel(column_label_short)
    ax4.set_xticks(np.arange(11,16))
    pl.ylabel(density_label_short)

    ax5 = pl.subplot(3,3,5)
    im5 = pl.imshow(mf.tline321[:,yy,:], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              aspect=np.diff(mf.crange)/np.diff(mf.trange),
              extent=mf.crange+mf.trange, vmin=vmin, vmax=vmax)
    pl.contour(mf.carr, mf.tarr, chi2b[:,yy,:], levels=chi2b.min()+levels)
    pl.xlabel(column_label_short)
    ax5.set_xticks(np.arange(11,16))
    pl.ylabel(temperature_label)
    #ax5.set_title("p-H$_2$CO $3_{2,1}-2_{2,0}$")

    ax6 = pl.subplot(3,3,6)
    im6 = pl.imshow(mf.tline321[:,:,xx], cmap=pl.cm.gray_r, interpolation='spline36',
              norm=pl.matplotlib.colors.LogNorm(),
              aspect=np.diff(mf.drange)/np.diff(mf.trange),
              extent=mf.drange+mf.trange, vmin=vmin, vmax=vmax)