#!/usr/bin/env python
"""Relative spectral-response calibration from a broadband source.

Divides a measured spectrum of a known source by that source's reference
spectrum to get the instrument's relative responsivity R(lambda), then smooths
the continuum:

    R(lambda)  ~  (counts - dark) / reference(lambda)

Apply it later as  irradiance_rel = (counts - dark) / R  (grab/live --response).

The reference is chosen with --reference:
  planck:<T>    ideal blackbody at <T> kelvin, e.g. planck:3300 for a halogen lamp
  <file.csv>    a lamp calibration certificate: wavelength_nm,irradiance

The calibration is relative and non-traceable unless the certificate is in
absolute units with its geometry reproduced. Calibrate through the same optical
path you measure with, since the response depends on the coupling. See README.

Usage: python calibrate_response.py source.csv [--reference REF] [--dark D.csv]
                                    [--integration US] [--band LO HI]
                                    [--out R.csv] [--png fig.png] [--show]
"""
import argparse
import os

import numpy as np
from scipy.signal import medfilt, savgol_filter

from response import dark_at, load_csv, planck

EDGE_PIXELS = 5  # drop this many pixels at each array end (fixed edge artifact)


def reference_spectrum(kind, wl):
    """Return (S on grid `wl`, label) for a --reference value."""
    if kind.startswith("planck:"):
        temperature = float(kind.split(":", 1)[1])
        return planck(wl, temperature), f"Planck {temperature:.0f} K"
    rwl, rval = load_csv(kind)  # otherwise a certificate CSV
    return np.interp(wl, rwl, rval), os.path.basename(kind)


def _odd(n):
    n = int(n)
    return n if n % 2 == 1 else n + 1


def build_responsivity(wl, counts, s, dark=None, band=(380.0, 950.0)):
    """Return (r_raw, r_smooth, m, valid): responsivity vs the source reference.

    `s` is the reference on grid `wl`. `dark` (same grid) is subtracted if
    given; never fabricate one from the spectrum itself.
    """
    lo, hi = band
    m = counts.astype(float)
    if dark is not None:
        m = m - dark

    # Validity mask: fit band, positive signal, away from the edges.
    valid = np.ones_like(wl, dtype=bool)
    valid[:EDGE_PIXELS] = False
    valid[-EDGE_PIXELS:] = False
    valid &= (wl >= lo) & (wl <= hi)
    valid &= m > 0

    # Raw ratio (responsivity) where valid.
    r_raw = np.full_like(wl, np.nan)
    r_raw[valid] = m[valid] / s[valid]

    # Smooth the continuum: interpolate across gaps, then median + Savitzky-Golay.
    r_interp = np.interp(wl, wl[valid], r_raw[valid])
    win = _odd(max(11, int(round(0.03 * wl.size))))  # ~3% of pixels, odd
    r_smooth = savgol_filter(medfilt(r_interp, kernel_size=_odd(9)),
                             window_length=win, polyorder=3)

    # Normalise to a dimensionless relative response (median over valid = 1),
    # scaling the raw ratio by the same factor so both share one axis. Normalise
    # before clipping: the reference scale is arbitrary (a Planck curve is ~1e27),
    # so an absolute floor here would wipe out the shape. After /norm R is order 1.
    norm = np.median(r_smooth[valid])
    r_smooth = np.clip(r_smooth / norm, 1e-6, None)
    r_raw /= norm

    # Outside the fit band we have no calibration; NaN rather than extrapolate.
    r_smooth[(wl < lo) | (wl > hi)] = np.nan
    return r_raw, r_smooth, m, valid


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("source_csv", help="measured spectrum of the known source")
    p.add_argument("--reference", default="planck:3300",
                   help="planck:<T> | a certificate CSV (default: planck:3300)")
    p.add_argument("--dark", metavar="CSV", default=None,
                   help="dark_model.csv from capture_dark.py, subtracted at "
                        "--integration. Apply the calibration with the same --dark.")
    p.add_argument("--integration", type=int, metavar="US", default=None,
                   help="integration time (us) the source was captured at; "
                        "required with --dark")
    p.add_argument("--band", nargs=2, type=float, metavar=("LO", "HI"),
                   default=[380.0, 950.0], help="fit band in nm (default: 380 950)")
    p.add_argument("--out", default="responsivity.csv")
    p.add_argument("--png", default="calibration.png")
    p.add_argument("--show", action="store_true", help="open an interactive window")
    args = p.parse_args()

    import matplotlib
    if not args.show:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    wl, counts = load_csv(args.source_csv)
    s, ref_label = reference_spectrum(args.reference, wl)

    dark = None
    if args.dark:
        if args.integration is None:
            p.error("--integration US is required with --dark")
        dwl, dvals = dark_at(args.dark, args.integration)
        dark = np.interp(wl, dwl, dvals)
        print(f"subtracting dark: {args.dark} at {args.integration} us")

    band = tuple(args.band)
    r_raw, r, m, valid = build_responsivity(wl, counts, s, dark=dark, band=band)

    np.savetxt(args.out, np.column_stack([wl, r]), delimiter=",",
               header="wavelength_nm,relative_responsivity", comments="")
    print(f"reference: {ref_label}   band: {band[0]:.0f}-{band[1]:.0f} nm")
    print(f"wrote {args.out} (relative response, median=1 over fit band)")

    corrected = m / r  # response-corrected source, for the sanity plot

    fig, ax = plt.subplots(3, 1, figsize=(11, 10), sharex=True)
    ax[0].plot(wl, m, lw=0.7, color="C0")
    ax[0].set_ylabel("measured (counts,\ndark-subtracted)")
    ax[0].set_title(f"{ref_label} source -> relative spectral response")

    ax[1].plot(wl, r_raw, ".", ms=2, color="0.6", label="raw ratio")
    ax[1].plot(wl, r, lw=1.6, color="C3", label="smoothed R(λ)")
    ax[1].set_ylabel("relative\nresponsivity")
    ax[1].legend(loc="upper right", fontsize=8)

    def norm(y):
        return y / np.median(y[valid])
    ax[2].plot(wl, norm(corrected), lw=0.7, color="C2", label="corrected source (norm.)")
    ax[2].plot(wl, norm(s), lw=1.2, color="k", alpha=0.6, label=f"{ref_label} reference (norm.)")
    ax[2].set_ylabel("normalised\nirradiance")
    ax[2].set_xlabel("wavelength (nm)")
    ax[2].set_ylim(-0.3, 2.0)
    ax[2].legend(loc="upper right", fontsize=8)

    for a in ax:
        a.grid(True, alpha=0.3)
        a.axvspan(wl[0], band[0], color="0.85", alpha=0.5, lw=0)
        a.axvspan(band[1], wl[-1], color="0.85", alpha=0.5, lw=0)

    fig.tight_layout()
    fig.savefig(args.png, dpi=110)
    print(f"wrote {args.png}  (grey = out of band)")
    if args.show:
        plt.show()


if __name__ == "__main__":
    main()