import os

WORK = "/Users/gbrammer/Research/JWST/Projects/MIRI-LRS/Workdir"
if os.path.exists(WORK):
    os.chdir(WORK)

import yaml
import logging

import matplotlib.pyplot as plt
import numpy as np
import astropy.io.fits as pyfits
import astropy.wcs as pywcs
import gwcs

import scipy.ndimage as nd

import os
import glob
import jwst.datamodels
from jwst.pipeline import Spec2Pipeline, Spec3Pipeline, Image2Pipeline

from grizli import jwst_utils

jwst_utils.QUIET_LEVEL = logging.CRITICAL
jwst_utils.set_quiet_logging(level=logging.CRITICAL)

from grizli import utils

import mastquery.utils

import msaexp.spectrum
from msaexp.cloud import lrs

utils.LOGFILE = "lrs_processing.log.txt"

_ = utils.log_comment(utils.LOGFILE, "Processing started", verbose=True, show_date=True)

from importlib import import_module

for module in ["grizli", "msaexp", "jwst", "astropy"]:
    mod = import_module(module)
    print(f"{module}=={mod.__version__}")

# GOODS-South source with target in both LRS slit and slitless

config = yaml.load("""
root: gds-jw04533
groups:
    jw04533001001a: 
     - jw04533001001_02102_00001_mirimage_rate.fits
     - jw04533001001_02102_00002_mirimage_rate.fits
    jw04533001001b: 
     - jw04533001001_02104_00001_mirimage_rate.fits
     - jw04533001001_02104_00002_mirimage_rate.fits
    jw04533001001c: 
     - jw04533001001_02106_00001_mirimage_rate.fits
     - jw04533001001_02106_00002_mirimage_rate.fits
    jw04533002001a: 
     - jw04533002001_02102_00001_mirimage_rate.fits
     - jw04533002001_02102_00002_mirimage_rate.fits
    jw04533002001b: 
     - jw04533002001_02102_00003_mirimage_rate.fits
     - jw04533002001_02102_00004_mirimage_rate.fits
""", Loader=yaml.Loader)

rate_files = []
for group in config["groups"]:
    rate_files.extend(config["groups"][group])

files_exist = np.array([os.path.exists(file) for file in rate_files])

if np.sum(files_exist) < len(files_exist):
    print(f"Download {(~files_exist).sum()} files:")
    result = mastquery.utils.download_from_mast(rate_files)
else:
    print(f"All {len(rate_files)} rate files found.")

lrs.load_rate_files?

rate_groups = {}

for i, group in enumerate(config["groups"]):
    
    utils.log_comment(
        utils.LOGFILE,
        f"\nGroup {i+1}/{len(config['groups'])}: {group}  N={len(config['groups'][group])} files\n",
        verbose=True,
        show_date=True
    )
    
    rate_groups[group] = lrs.load_rate_files(config["groups"][group])

# Region files with illuminated exposure footprints
lrs.rate_group_regions?

smiles_nirspec = utils.read_catalog(
    "https://archive.stsci.edu/hlsps/smiles/dr2/hlsp_smiles_jwst_nirspec_goodss_multi_v1.0_catalog.fits"
)

with pyfits.open(
    "https://archive.stsci.edu/hlsps/smiles/dr1/hlsp_smiles_jwst_miri_goodss_multi_v1.0_catalog.fits",
) as hdul:
    smiles_miri = utils.read_catalog(hdul["CIRC"])

so = np.argsort(smiles_miri["F560W_CIRC4"])[::-1]
smiles_miri = smiles_miri[so]

idx, dr = smiles_miri.match_to_catalog_sky(smiles_nirspec)
hasm = dr.value < 0.2

smiles_nirspec = smiles_nirspec[hasm]
for c in smiles_miri.colnames:
    if c not in smiles_nirspec.colnames:
        smiles_nirspec[c] = smiles_miri[c][idx[hasm]]

so = np.argsort(smiles_nirspec["F560W_CIRC4"])[::-1]
smiles_nirspec = smiles_nirspec[so]

sregions, fig = lrs.rate_group_regions(rate_groups, root=config["root"], separate_region_files=True)

## Overplot SMILES sources
smiles_nirspec["in_lrs"] = False
coo_nirspec = np.array([smiles_nirspec["RA"], smiles_nirspec["DEC"]]).T

for sr in sregions:
    for path in sr.path:
        smiles_nirspec["in_lrs"] |= path.contains_points(coo_nirspec)

smiles_miri["in_lrs"] = False
coo_miri = np.array([smiles_miri["RA"], smiles_miri["DEC"]]).T

for sr in sregions:
    for path in sr.path:
        smiles_miri["in_lrs"] |= path.contains_points(coo_miri)

ax = fig.axes[0]
ax.scatter(*coo_miri[smiles_miri["in_lrs"],:].T, marker='.', alpha=0.5, label="SMILES MIRI", color="0.5")

ax.scatter(*coo_nirspec[smiles_nirspec["in_lrs"],:].T, marker='s', alpha=0.5, label="SMILES NIRSpec", ec="tomato", fc="None")

ax.legend()

smiles_miri = smiles_miri[smiles_miri["in_lrs"]]

# Drizzled mosaics of group exposures
lrs.rate_group_mosaic?

products = lrs.rate_group_mosaic(rate_groups, root=config['root'], verbose=2)

with pyfits.open(f"{config['root']}_miri_lrs_sci.fits") as im:
    mosaic_data = im[0].data

fig, ax = plt.subplots(1,1,figsize=(8,8))

ax.imshow(mosaic_data * np.nan**(mosaic_data == 0), vmin=-0.3, vmax=0.3, cmap="bone")

fig.tight_layout(pad=1)

coords = """

# Source in LRS slit and slitless area
# https://s3.amazonaws.com/msaexp-nirspec/miri-lrs/jw04533001-mds.-deep-survey-alma-10_lrs.spec.fits
53.166943 -27.798828

# Other PAH sources
53.1726177 -27.7810377
53.1623705 -27.7751336

# Very bright with PAH
53.1616034 -27.7923319

""".strip().split("\n")[-1]

ra, dec = np.array(coords.split()[:2]).astype(float)

pos_key = f"{ra:.6f}_{dec:.6f}"

ra, dec, pos_key

if 0:
    # Step through SMILES
    j = -1

if 0:
    j += 1
    
    if 0:
        subset = smiles_nirspec["in_lrs"] & (smiles_nirspec["Q"] == "A")
        ra = smiles_nirspec[subset]["RA"][j]
        dec = smiles_nirspec[subset]["DEC"][j]        
        smiles_nirspec[subset]["NIRCAM_ID", "RA", "DEC", "Z_SPEC", "KEYWORDS", "F560W_CIRC3"][j:j+1].pprint_all()
    
    else:
    
        smiles_miri["ID", "RA", "DEC", "F560W_CIRC4", "F770W_CIRC4"][j:j+1].pprint_all()
        ra = smiles_miri["RA"][j]
        dec = smiles_miri["DEC"][j]
        
    pos_key = f"{ra:.6f}_{dec:.6f}"

# Get the cutouts

fig, data = lrs.cutout_pixel_table(ra, dec, rate_groups, make_figure=True, subtract_pedestal=True, check_window=True)

if fig is not None:
    fig.savefig(f"{config['root']}_{pos_key}_exposures.png")

group_names, ung = lrs.slitless_data_groups(data, config["groups"])

if 0:
    # Manually mask some contamination in the Monster
    une = utils.Unique(data["exposure"], verbose=False)
    bad = (data["wave"] > 12) # & (~np.isin(data["exposure"], [8, 9]))
    bad = (data["wave"] > 13.4) # & (~np.isin(data["exposure"], [8, 9]))
    data["data"][bad] = np.nan

if 0:
    bad = (data["dx"] > 0.4) & (data["wave"] > 9.1)
    bad = (data["dx"] < -1)
    bad = (data["dx"] > 1.2) & (data["wave"] >10.5)
    bad = (data["dx"] < -1.2)
    data["data"][bad] = np.nan

if 0:
    data = lrs.get_data_slice(data, (data["exposure"] == 14))

psf_data = lrs.evaluate_model_psf(
    data, np.array([0.2, 0.25, 0.15]), do_optimal_extraction=True
)

_ = lrs.process_2d_products(data, recenter=False)

# Bkg pedestal
from scipy.stats import binned_statistic

SUBTRACT_EXTRA = False

une = utils.Unique(data["exposure"], verbose=False)

fig, ax = plt.subplots(1,1,figsize=(8, 5))

xwbin = np.linspace(-1,1, len(data["wbin"]))
xwb = np.interp(data["wave"], data["wbin"], xwbin)

for exp in une.values:
    sub = une[exp] & (np.abs(data["dx"]) > 1.4) & np.isfinite(data["data"])
    sub &= data["data"] < 3 * data["err"]
    wsub = data["wave"] > 1.8
    # wsub = data["wave"] < 10
    
    # wsub = data["wave"] < 10
    sub &= wsub
    if sub.sum() == 0:
        continue
        
    stat = binned_statistic(data["wave"][sub], data["data"][sub], statistic="median", bins=data["wbin_edge"])

    pl = ax.plot(data["wbin"], stat.statistic, label=f"Exp {exp}", alpha=0.4)
    ok = np.isfinite(stat.statistic)
    
    c = np.polyfit(xwbin[ok], stat.statistic[ok], 0)
    ax.plot(data["wbin"], np.polyval(c, xwbin), color=pl[0].get_color(), alpha=0.4)

    if SUBTRACT_EXTRA == False:
        data["data"][une[exp]] -= np.polyval(c, xwb[une[exp]])
    
ax.grid()
ax.legend()
fig.tight_layout(pad=0.5)

# Initialize trace parameters
group_theta = {}

ung = utils.Unique(np.array(data["meta"]["group_names"])[data["group"]], verbose=True)

for ig in ung.values:
    group_theta[ig] = np.array([0.12, 0.21, 0.1])

if 0:
    # check delta chi2
    shifts = np.arange(-0.2, 0.41, 0.1)
    sizes = np.arange(0, 0.21, 0.05)
    
    theta = []
    chi2 = []
    
    for sh in shifts:
        for si in sizes:
            theta_i = [sh, si]
            psf_data = lrs.evaluate_model_psf(
                data, theta_i, do_optimal_extraction=True
            )
            theta.append(theta_i)
            chi2.append(data["chi2"])
    
            print(f"{theta_i[0]:5.1f} {theta_i[1]:5.1f}  {chi2[-1]:>10.3f}")
            
    t = utils.GTable()
    t["theta"] = theta
    t["chi2"] = chi2
    # t.pprint_all()
    
    print(theta[np.argmin(chi2)])

# Do the fit
ung, group_theta = lrs.slitless_trace_by_group(
    data, verbose=1,
    method="bfgs", tol=1.e-4,
    reset_initial=True,
    initial_size=0.08,
    fit_size=True,
    show_each_2d=True,
    fix_first_size=True,  # Force sizes to be the same as the first
    fix_first_shift=False,
)

##########
# Full wavelength range for plots
xlim = None

if 0:
    xlim = (3.9, 12.1)

##########
# Redshift label
z = None

if 0:
    z = 0.3029
    
######

# Make sure separate thetas are set
lrs.set_full_profile(data, group_theta)

lrs.optimal_extraction(data, recenter=True)
hdu, fig = lrs.process_2d_products(data, prefix=f"{config['root']}_{pos_key}", xlim=xlim)

lspec = msaexp.spectrum.SpectrumSampler(hdu)
unit = "fnu"
summary_figs = {}
for unit in ["flam", "fnu"]:
    summary_figs[unit] = lspec.drizzled_hdu_figure(hdu, unit=unit, ny=12, xlim=xlim, z=z)

hdus = {}
lspecs = {}

# Show individual group extractions
if (ung.N > 1) & (1):
    for ig in group_theta:
        data_i = lrs.get_data_slice(data, ung[ig])
        
        lrs.optimal_extraction(data_i, recenter=True)
        hdu_i, fig = lrs.process_2d_products(data_i, prefix=f"{ig}_{pos_key}", xlim=xlim, make_figure=True)

        lspec_i = msaexp.spectrum.SpectrumSampler(hdu_i)
        for unit in summary_figs:
            msk = np.nan**(~lspec_i['valid'])
            summary_figs[unit].axes[2].step(np.arange(len(lspec_i["wave"])), lspec_i["flux"] * lspec_i["to_flam"]**(unit == "flam") * msk, alpha=0.5, zorder=-100)

        hdus[ig] = hdu_i
        lspecs[ig] = lspec_i

prefix = f"{config['root']}_{pos_key}_lrs.spec"

for unit in summary_figs:
    if ung.N > 1:
        for unit in summary_figs:
            msk = np.nan**(~lspec['valid'])
            summary_figs[unit].axes[2].step(np.arange(len(lspec["wave"])), lspec["flux"] * lspec["to_flam"]**(unit == "flam") * msk, alpha=0.3, zorder=100, color='k')

            summary_figs[unit].savefig(f"{prefix}.{unit}.png")

# B-Spline continuum templates
bspl = lspec.bspline_array(nspline=9, by_wavelength=False, orders=[1])

# seems like wavelength grid is often/always off by one pixel??
lspec.spec_wobs = lspec["wave"] + np.gradient(lspec["wave"])

fit_kwargs = dict(
    continuum_models=bspl, fitter="nnls",
    scale_disp=0.8
)

result = lrs.fit_lrs_redshift(lspec, zlim=(0.02, 2.5), dz=0.005, **fit_kwargs, make_plots=True)

next_zlim = result["zbest"] + np.array([-0.025, 0.025]) * (1 + result["zbest"])

result = lrs.fit_lrs_redshift(lspec, zlim=next_zlim, dz=0.0003, **fit_kwargs, make_plots=2)

# Break run all
print(xbreak)

data_slit = lrs.get_data_slice(data, data["in_slit"])

psf_data = lrs.evaluate_model_psf(
    data_slit, [0.0, 0.0], do_optimal_extraction=True
)

psf_dict = (psf_data["psf_mu"], psf_data["psf_width"])

ext_psf_dict = [
    psf_data[k]
    for k in ["sigma_scales", "sigma_widths", "psf_shifts", "sigma_coeffs"]
]

print(group_theta[ig])

path = lrs.profile_pathloss(
    data_slit, group_theta[ig],
    psf_data=psf_dict,
    extended_psf_data=ext_psf_dict,
    do_optimal_extraction=False,
    slit_width=0.51,
)

so = np.argsort(data_slit["wave"])

plt.plot(data_slit["wave"][so], path["pathloss"][so])
plt.xlabel("wavelength")
plt.ylabel("estimated pathloss")

lrs.optimal_extraction(data_slit, recenter=True)
hdu, fig = lrs.process_2d_products(data_slit, prefix=f"slit_{pos_key}", xlim=xlim, make_figure=False)
lspec_slit = msaexp.spectrum.SpectrumSampler(hdu)

# Apply the correction
data_slit["data"] /= path["pathloss"]

lrs.optimal_extraction(data_slit, recenter=True)
hdu, fig = lrs.process_2d_products(data_slit, prefix=f"slitpath_{pos_key}", xlim=xlim)

lspec_slit_pathcorr = msaexp.spectrum.SpectrumSampler(hdu)


# Slitless
data_slitless = lrs.get_data_slice(data, ~data["in_slit"])

lrs.optimal_extraction(data_slitless, recenter=True)

hdu, fig = lrs.process_2d_products(data_slitless, prefix=f"slitless_{pos_key}", xlim=xlim, make_figure=False)
lspec_slitless = msaexp.spectrum.SpectrumSampler(hdu)

fig, ax = plt.subplots(1,1,figsize=(8,5))

ax.step(lspec_slit["wave"], lspec_slitless["flux"] * np.nan**(~lspec_slit["valid"]), label="Slitless", alpha=0.5)

ax.step(lspec_slit["wave"], lspec_slit_pathcorr["flux"] * np.nan**(~lspec_slit["valid"]) * 3, label="Slit, pathloss corrected x 3", alpha=0.5)

ax.step(lspec_slit["wave"], lspec_slit["flux"] * np.nan**(~lspec_slit["valid"]) * 3.5, label="Slit x 3.5", alpha=0.5)

ax.step(lspec_slit["wave"], lspec_slit_pathcorr["flux"] * np.nan**(~lspec_slit["valid"]), label="Slit, pathloss corrected", alpha=0.5)

ax.step(lspec_slit["wave"], lspec_slit["flux"] * np.nan**(~lspec_slit["valid"]), label="Slit", alpha=0.5)

ax.set_ylim(-0.1, 50)
ax.legend()

fig.tight_layout(pad=1)

APPLY_SPLINE_CORRECTION = False

fig, axes = plt.subplots(2,1,figsize=(10,8), height_ratios=[1, 0.5], sharex=True)

axes[0].plot(lspec['wave'], lspec['flux_bin'] * lspec['calsp'] / data["unit_bin"], color='k', alpha=0.5, zorder=1001)

axes[1].plot(lspec['wave'], lspec['flux_bin'] * lspec['calsp'] * 0, color='k', alpha=0.5, zorder=1001)

xwbin = np.linspace(-1,1, len(data["wbin"]))
xwb = np.interp(data["wave"], data["wbin"], xwbin)

df = 13

bsplf = utils.bspline_templates(xwb, df=df, minmax=(-1,1), get_matrix=True)
bspl = utils.bspline_templates(xwbin, df=df, minmax=(-1,1), get_matrix=True).T

for ig in lspecs:
    lspec_i = lspecs[ig]
    pl = axes[0].plot(lspec_i['wave'], lspec_i['flux_bin'] * lspec_i['calsp'] / data["unit_bin"], alpha=0.4, label=ig)

    diff = (lspec_i['flux_bin'] - lspec['flux_bin']) * lspec_i['calsp'] / data["unit_bin"]
    ediff = lspec_i['err'] * lspec_i['calsp'] / data["unit_bin"]
    ok = np.isfinite(diff / ediff)
    
    c = np.linalg.lstsq((bspl / ediff).T[ok,:], (diff / ediff)[ok], rcond=None)
    splm = bspl.T.dot(c[0])

    if APPLY_SPLINE_CORRECTION:
        splmf = bsplf[ung[ig],:].dot(c[0])
        data["data"][ung[ig]] -= splmf
        
    pl = axes[1].plot(lspec_i['wave'], diff, alpha=0.4, label=ig)
    pl = axes[1].plot(lspec_i['wave'], splm, alpha=0.4, color=pl[0].get_color())

ax = axes[0]
ax.legend()
ax.grid()
ax.set_xlabel(r'$\lambda$')
ax.set_ylabel('flux  [DN/s]')

axes[1].grid()

fig.tight_layout(pad=1)