#!/usr/bin/env python
"""Live-streaming spectrum viewer for an Ocean Optics spectrometer.

Opens a matplotlib window and continuously redraws the current spectrum.

Usage: python live_spectrum.py [integration_time_us]

Keys (while the plot window is focused):
    q / Esc   quit
    a         toggle autoscale of the y-axis on/off
    d         capture the current spectrum as a dark reference
    x         clear the dark reference
    space     freeze / unfreeze acquisition
    t         toggle a live Planck temperature fit (use --response for a real T)
"""
import argparse

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from seabreeze.spectrometers import Spectrometer

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("integration_us", nargs="?", type=int, default=50000,
                    help="integration time in microseconds (default: 50000)")
parser.add_argument("--interval", type=int, default=30,
                    help="redraw interval in ms (default: 30)")
parser.add_argument("--fit-band", nargs=2, type=float, default=[420.0, 800.0],
                    metavar=("LO", "HI"),
                    help="band (nm) for the live temperature fit (default: 420 800)")
add_calibration_args(parser)
args = parser.parse_args()

cal = load_csv(args.response) if args.response else None
dark = dark_at(args.dark, args.integration_us)

EDGE_PIXELS = 5  # ignore the fixed edge artifacts (e.g. the ~340 nm pixel-1 spike)

spec = Spectrometer.from_first_available()
spec.integration_time_micros(args.integration_us)
wl = spec.wavelengths()
print(f"{spec.model}  serial={spec.serial_number}  pixels={spec.pixels}")
print(f"range {wl[0]:.1f}-{wl[-1]:.1f} nm  integration {args.integration_us} us")
if cal is not None:
    print(f"applied response calibration: {args.response}")
print("keys: q quit | a autoscale | d dark | x clear dark | space freeze | t fit T")

state = {"autoscale": True, "frozen": False, "dark": None, "fit": False}

fig, ax = plt.subplots(figsize=(10, 5))
(line,) = ax.plot(wl, np.zeros_like(wl), lw=0.8)
# Two lines for the fit: the full-range curve (extrapolated) and, on top, the
# segment over the band actually used to constrain it.
(fit_full,) = ax.plot([], [], lw=1.2, color="C1")
(fit_band,) = ax.plot([], [], lw=2.0, color="C3")
fit_full.set_visible(False)
fit_band.set_visible(False)
ax.set_xlim(wl[0], wl[-1])
ax.set_ylim(0, 4096)
ax.set_xlabel("wavelength (nm)")
ax.set_ylabel("relative irradiance" if cal is not None else "intensity (counts)")
title = ax.set_title("")
ax.grid(True, alpha=0.3)


def on_key(event):
    if event.key in ("q", "escape"):
        plt.close(fig)
    elif event.key == "a":
        state["autoscale"] = not state["autoscale"]
    elif event.key == " ":
        state["frozen"] = not state["frozen"]
    elif event.key == "d":
        state["dark"] = spec.intensities(correct_nonlinearity=True)
        print("dark reference captured")
    elif event.key == "x":
        state["dark"] = None
        print("dark reference cleared")
    elif event.key == "t":
        state["fit"] = not state["fit"]
        if state["fit"] and cal is None:
            print("fitting uncalibrated counts; use --response for a meaningful T")


fig.canvas.mpl_connect("key_press_event", on_key)


def _hide_fit():
    fit_full.set_visible(False)
    fit_band.set_visible(False)
    if ax.get_legend():
        ax.get_legend().remove()


def update(_frame):
    if state["frozen"]:
        return line, fit_full, fit_band, title
    y = spec.intensities(correct_nonlinearity=True)
    if state["dark"] is not None:
        y = y - state["dark"]
    if cal is not None:
        y = apply_response(wl, y, cal, dark)
    line.set_ydata(y)
    if state["autoscale"]:
        ymin, ymax = float(np.nanmin(y)), float(np.nanmax(y))
        pad = max(1e-9, 0.05 * (ymax - ymin))
        ax.set_ylim(ymin - pad, ymax + pad)
    # Peak over the interior only, so the fixed edge spikes never win.
    pk = EDGE_PIXELS + int(np.nanargmax(y[EDGE_PIXELS:-EDGE_PIXELS]))
    if state["fit"]:
        result = fit_planck(wl, y, tuple(args.fit_band))
        if result is None:
            _hide_fit()
        else:
            temperature, dT, predict = result
            lo, hi = args.fit_band
            model = predict(wl)
            in_band = (wl >= lo) & (wl <= hi)
            fit_full.set_data(wl, model)          # extrapolated over the full range
            fit_band.set_data(wl[in_band], model[in_band])  # region used for the fit
            note = "" if cal is not None else " (uncalibrated)"
            fit_full.set_label("Planck fit (extrapolated)")
            fit_band.set_label(f"used {lo:.0f}-{hi:.0f} nm, "
                               f"T = {temperature:.2f} +/- {dT:.2f} K{note}")
            fit_full.set_visible(True)
            fit_band.set_visible(True)
            ax.legend(loc="upper right", fontsize=9)
    else:
        _hide_fit()
    flags = []
    if state["dark"] is not None:
        flags.append("dark-sub")
    if state["frozen"]:
        flags.append("FROZEN")
    if not state["autoscale"]:
        flags.append("fixed-y")
    suffix = ("  [" + " ".join(flags) + "]") if flags else ""
    title.set_text(f"peak {wl[pk]:.2f} nm = {y[pk]:.2f}{suffix}")
    return line, fit_full, fit_band, title


# blit=False because autoscale mutates the axes limits each frame.
ani = FuncAnimation(fig, update, interval=args.interval, blit=False,
                    cache_frame_data=False)

try:
    plt.show()
finally:
    spec.close()