#!/usr/bin/env python
"""Overlay one or more emission spectra, response-corrected, with Planck fits.

For each spectrum CSV: optionally response-correct it (--response/--dark plus
--integration, as grab_spectrum.py --response would), fit a Planck curve over
--band, and plot the corrected spectrum with its fit and the fitted temperature.
Reuses the correction and the fit from the rest of the project.

Usage: python plot_emission.py f1.csv f2.csv ... [--response R.csv --dark D.csv
                              --integration US] [--band LO HI] [--xlim LO HI]
                              [--png fig.png]
"""
import argparse

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from response import add_calibration_args, apply_response, dark_at, load_csv
from fit_temperature import fit_planck

p = argparse.ArgumentParser(description=__doc__,
                            formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("spectra", nargs="+", help="emission spectrum CSVs")
p.add_argument("--integration", type=int, metavar="US",
               help="integration time (us) the spectra were captured at; required with --dark")
p.add_argument("--band", nargs=2, type=float, default=[650.0, 850.0],
               metavar=("LO", "HI"), help="Planck fit band in nm (default: 650 850)")
p.add_argument("--xlim", nargs=2, type=float, default=[560.0, 930.0],
               metavar=("LO", "HI"), help="plot x-range in nm (default: 560 930)")
p.add_argument("--png", default="emission_fit.png")
add_calibration_args(p)
args = p.parse_args()

cal = load_csv(args.response) if args.response else None
dark = None
if args.dark:
    if args.integration is None:
        p.error("--integration US is required with --dark")
    dark = dark_at(args.dark, args.integration)

lo, hi = args.band
fig, ax = plt.subplots(figsize=(9, 5))
for f in args.spectra:
    wl, y = load_csv(f)
    if cal is not None:
        y = apply_response(wl, y, cal, dark)
    result = fit_planck(wl, y, (lo, hi))
    m = (wl >= args.xlim[0]) & (wl <= args.xlim[1])
    (base,) = ax.plot(wl[m], y[m], lw=0.5, alpha=0.35)
    if result is None:
        print(f"{f}: no fit in {lo:.0f}-{hi:.0f} nm")
        continue
    temperature, dT, predict = result
    ax.plot(wl[m], predict(wl[m]), lw=2.0, color=base.get_color(),
            label=f"{f}:  T = {temperature:.0f} K")
    print(f"{f}: T = {temperature:.2f} +/- {dT:.2f} K")

ax.axvspan(lo, hi, color="0.85", alpha=0.5, lw=0, zorder=0)
ax.set_xlabel("wavelength (nm)")
ax.set_ylabel("relative irradiance (response-corrected)")
ax.set_title(f"Emission with Planck fit over {lo:.0f}-{hi:.0f} nm (shaded)")
ax.set_xlim(*args.xlim)
ax.set_ylim(bottom=0)
ax.grid(alpha=0.3)
ax.legend(fontsize=9, loc="upper left")
fig.tight_layout()
fig.savefig(args.png, dpi=130)
print(f"wrote {args.png}")