import numpy as np

from astropy.io import fits
from astropy.table import Table, Column

from desispec.io.spectra import read_spectra

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rc('font', size=14)

"""
Explanation: Time Domain Spectral Simulations
Demonstrate how to inspect simulated spectra produced using quicktransients.
End of explanation
"""


truth_file = '../../bgs_2020-03-08_0300s_001_truth.fits'
spect_file = '../../bgs_2020-03-08_0300s_001_spect.fits'
coadd_file = '../../bgs_2020-03-08_0300s_001_coadd.fits'

"""
Explanation: Input Files
The quicktransients program in the desisim transients branch will produce two FITS outputs:
1. A truth file with information about the templates used for each object.
2. A spect file with the templates "observed" under conditions specified by the user.
The spectra can then be coadded using the desi_coadd_spectra program available in desispec.
End of explanation
"""


hdus = fits.open(truth_file)
hdus.info()

wave = hdus['WAVE'].data
flux = hdus['FLUX'].data
targets = Table.read(truth_file, 'TARGETS')
truth = Table.read(truth_file, 'TRUTH')
objtr = Table.read(truth_file, 'OBJTRUTH')

truth

r = 22.5 - 2.5*np.log10(truth['FLUX_R'])
z = truth['TRUEZ']

fig, ax = plt.subplots(1,1, figsize=(6,4))
ax.scatter(z, r)
ax.set(xlabel='$z$', ylabel='$r$')
fig.tight_layout();

targets

objtr

"""
Explanation: Contents of the Truth File
The truth file has the following tables:
1. A wavelength table called WAVE.
1. A flux table called FLUX.
1. A TARGETS table simulating a target list available in data.
1. A simulation TRUTH table with information about the object (redshift, flux, etc.).
1. A simulation OBJTRUTH table with line fluxes and other data generated for each object.
End of explanation
"""


fig, axes = plt.subplots(1,2, figsize=(9,4), sharey=True)
ax = axes[0]
n, bins, patches = ax.hist(objtr['TRANSIENT_EPOCH'], bins=10, color='k', alpha=0.3)
x = 0.5*(bins[1:] + bins[:-1])
dx = 0.5*np.diff(bins)
ax.errorbar(x, n, xerr=dx, yerr=np.sqrt(n), fmt=',', color='k')
ax.set(xlabel='epoch [day]', ylabel='count')

ax = axes[1]
n, bins, patches = ax.hist(objtr['TRANSIENT_RFLUXRATIO'], bins=10, color='k', alpha=0.3)
x = 0.5*(bins[1:] + bins[:-1])
dx = 0.5*np.diff(bins)
ax.errorbar(x, n, xerr=dx, yerr=np.sqrt(n), fmt=',', color='k')
ax.set(xlabel=r'$r_\mathrm{trans}/r_\mathrm{host}$')

fig.tight_layout();

"""
Explanation: Object Truth Data
Plot transient data from the objtruth table:
1. Distribution of "epoch," time in days w.r.t. $t_0$.
1. Distribution of transient/host flux ratio in $r$ band.
End of explanation
"""


for i in range(10):
    plt.plot(wave, flux[i])

"""
Explanation: Plot Templates
Draw the first 10 templates generated from the mock catalog.
End of explanation
"""


spectra = read_spectra(spect_file)
hdus = fits.open(spect_file)
hdus.info()

spectra.fibermap

"""
Explanation: Observed Spectra
The spectra generated under observing conditions are stored in a single file. A FIBERMAP table provides a summary of target data for each object, and individual wavelength, flux, variance, resolution, and mask tables are present for each camera.
The data are best accessed using the read_spectra function from desispec, which packs everything into a single object.
End of explanation
"""


coadds = read_spectra(coadd_file)
hdus = fits.open(coadd_file)
hdus.info()

coadds.fibermap

"""
Explanation: Coadd Files
The coadds are generated using the desi_coadd_spectra program available in desispec. For example, to add data across the cameras run
desi_coadd_spectra -i bgs_2020-03-08_0150s_001_spect.fits -o bgs_2020-03-08_0150s_001_coadd.fits --coadd-cameras
The data can then be accessed using the read_spectra function.
End of explanation
"""


if 'MEDIAN_COADD_SNR' not in coadds.scores.columns:
    totsnr = None
    for cam in 'BRZ':
        camsnr = coadds.scores['MEDIAN_COADD_SNR_{}'.format(cam)]
        if totsnr is None:
            totsnr = camsnr**2
        else:
            totsnr += camsnr**2
    totsnr = np.sqrt(totsnr)
    coadds.scores.add_column(Column(totsnr, name='MEDIAN_COADD_SNR'))

coadds.scores

"""
Explanation: Spectral Scores
Per-camera median coadded fluxes and SNRs are available in a scores table.
The code below computes a total SNR and adds it to the table.
End of explanation
"""


from scipy.signal import medfilt

fig, axes = plt.subplots(10,3, figsize=(14,30), sharex=True)
for i in range(10):
    axes[i,0].plot(wave, flux[i], alpha=0.7, label='template')
    axes[i,0].legend()
    
    for _filt, _col in zip(spectra.bands, ['b', 'r', 'saddlebrown']):
        wl = spectra.wave[_filt]
        fl = spectra.flux[_filt][i]
        axes[i,1].plot(wl, fl, color=_col, alpha=0.7, label=_filt)
    axes[i,1].legend()
    
    axes[i,2].plot(coadds.wave['brz'], coadds.flux['brz'][i], alpha=0.7, label='coadd')
    axes[i,2].plot(coadds.wave['brz'], medfilt(coadds.flux['brz'][i], 149), color='yellow')
    axes[i,2].legend()

fig.tight_layout();

"""
Explanation: Plot Results
To visualize the results, the templates, camera data, and coadds for the first 10 spectra in the input files listed at the top of this notebook are plotted.
End of explanation
"""


