import os
if os.path.exists('gbrammer' in os.environ['HOME']) & False:
    
    # Extract summary from database
    from grizli.aws import db
    import pyperclip

    # db queries require DB credentials
    nre = db.SQL("""select SUBSTR(MIN(dataset), 4, 4) as pid,
                           array_agg(DISTINCT(root)) as root,
                           count(*) as N,
                           array_agg(DISTINCT(grating || '-' || filter)) as Grating
    FROM nirspec_extractions WHERE ROOT LIKE '%%v2' AND dataset not like 'jw01208%%'
    GROUP BY SUBSTR(dataset, 4, 4)
    ORDER BY MAX(root)""")
    
    nre['root'] = ['\n'.join(r) for r in nre['root']]
    nre['grating'] = ['\n'.join(r) for r in nre['grating']]
    nre.rename_column('root','DJA root')
    url = '[{pid}](https://www.stsci.edu/cgi-bin/get-proposal-info?id={pid}&observatory=JWST)'
    prog = [url.format(**row) for row in nre]
    nre['Program'] = prog
    
    pyperclip.copy('## Observing programs\n\n' +
                   nre['Program','DJA root','n','grating'].to_pandas(index=False).to_markdown(index=False))

%matplotlib inline

import os
import yaml

import numpy as np
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

import astropy.io.fits as pyfits

import grizli
import grizli.catalog
from grizli import utils

import eazy
import msaexp
import msaexp.spectrum

print(f'grizli version: {grizli.__version__}')
print(f'eazy-py version: {eazy.__version__}')
print(f'msaexp version: {msaexp.__version__}')

z, v2_file = 3.7616, 'https://s3.amazonaws.com/msaexp-nirspec/extractions/ceers-ddt-v2/ceers-ddt-v2_prism-clear_2750_1598.spec.fits'

v1_file = v2_file.replace('ceers-v2','ceers-lr-v1').replace('-v2','-v1')
print(v1_file + '\n' + v2_file)

sp1 = msaexp.spectrum.SpectrumSampler(v1_file)
sp2 = msaexp.spectrum.SpectrumSampler(v2_file)

f1 = sp1.drizzled_hdu_figure(z=z, unit='fnu')
f2 = sp2.drizzled_hdu_figure(z=z, unit='fnu')

import msaexp.slit_combine
from msaexp.slit_combine import slit_prf_fraction

fig, axes = plt.subplots(2,1,figsize=(8,4), sharex=True)

ax = axes[0]
ax.plot(sp1['wave'], sp1['flux'], label='v1')
ax.plot(sp2['wave'], sp2['flux'], label='v2')
leg = ax.legend(loc='upper left')
leg.set_title(sp2.meta['SRCNAME'])
ax.grid()

# Ratios
axes[1].plot(sp2['wave'], sp2['flux'] / sp1['flux'], label=r'v2 / v1', color='pink')

path_corr = slit_prf_fraction(sp2.spec['wave'].astype(float),
                              sigma=sp2.spec.meta['SIGMA'],
                              x_pos=sp2.spec.meta['SRCXPOS'],
                              slit_width=0.2,
                              pixel_scale=msaexp.slit_combine.PIX_SCALE,
                              verbose=False)

axes[1].plot(sp2.spec['wave'], 1./path_corr, color='plum', label='v2 path loss correction')
axes[1].grid()
axes[1].legend(loc='upper left')

fig.tight_layout(pad=1)

z, prism_file = 2.0611, 'https://s3.amazonaws.com/msaexp-nirspec/extractions/ceers-v2/ceers-v2_prism-clear_1345_3506.spec.fits'

sp = {}
for grating in ['prism-clear','g140m-f100lp','g235m-f170lp','g395m-f290lp']:
    file = prism_file.replace('prism-clear',grating)
    print(file)
    sp[grating] = msaexp.spectrum.SpectrumSampler(file)

    fig = sp[grating].drizzled_hdu_figure(z=z, unit='fnu')

fig, axes = plt.subplots(4,1,figsize=(8,8))

colors = ['0.2','steelblue','orange','tomato']

for i, grating in enumerate(sp):
    sp[grating]['flux'][~sp[grating].valid] = np.nan
    for ax in axes:
        ax.plot(sp[grating]['wave'], sp[grating]['flux'], alpha=0.5, label=grating, color=colors[i])

leg = axes[3].legend(ncol=4, fontsize=8)
axes[0].set_title(sp[grating].meta['SRCNAME'])

axes[0].set_xlim(0.9, 1.9); axes[0].set_ylim(-1.2, 12)
axes[1].set_xlim(1.85, 3.2); axes[1].set_ylim(-0.5, 5)
axes[2].set_xlim(2.95, 5.3); axes[2].set_ylim(-0.8, 8)
axes[3].set_ylim(-3,30)
axes[3].set_xlabel(r'$\lambda_\mathrm{obs}~[\mu\mathrm{m}]$')

for ax in axes:
    ax.grid()
    ax.set_ylabel(r'$f_\nu~[\mu\mathrm{Jy}]$')

fig.tight_layout(pad=1)