"""
Analyse the output files from `../calibration/spectral_response_monochromator.py`.
This plots all the intermediate spectra generated by that script.

Command line arguments:
    * `folder`: folder containing monochromator intemediaries. These should be
    NPY files generated using ../calibration/spectral_response_monochromator.py
    By default, these will have been saved to
    root/intermediaries/spectral_response/
"""

import numpy as np
from sys import argv
from spectacle import io, spectral

# Get the data folder from the command line
folder = io.path_from_input(argv)
root = io.find_root_folder(folder)

# Load Camera object
camera = io.load_camera(root)
print(f"Loaded Camera object: {camera}")

# Save locations
savefolder = camera.filename_analysis("spectral_response", makefolders=True)
save_to_raw = savefolder/"monochromator_raw_spectra.pdf"
save_to_calibrated = savefolder/"monochromator_calibrated_spectra.pdf"
save_to_normalised = savefolder/"monochromator_normalised_spectra.pdf"
save_to_final = savefolder/"monochromator_final_spectrum.pdf"
save_to_SNR = savefolder/"monochromator_SNR.pdf"

# Load the wavelength data
wavelengths = np.load(folder/"monochromator_wavelengths.npy")

# Load and plot the raw response curves
raw_mean = np.load(folder/"monochromator_raw_means.npy")
raw_stds = np.load(folder/"monochromator_raw_stds.npy")

spectral.plot_monochromator_curves(wavelengths, raw_mean, raw_stds, title=f"{camera.name}: Raw spectral curves", saveto=save_to_raw)
print("Saved raw spectrum plot")

# Load and plot the calibrated response curves
calibrated_mean = np.load(folder/"monochromator_calibrated_means.npy")
calibrated_stds = np.load(folder/"monochromator_calibrated_stds.npy")

spectral.plot_monochromator_curves(wavelengths, calibrated_mean, calibrated_stds, title=f"{camera.name}: Calibrated spectral curves", saveto=save_to_calibrated)
print("Saved calibrated spectrum plot")

# Load and plot the normalised response curves
normalised_mean = np.load(folder/"monochromator_normalised_means.npy")
normalised_stds = np.load(folder/"monochromator_normalised_stds.npy")

spectral.plot_monochromator_curves(wavelengths, normalised_mean, normalised_stds, title=f"{camera.name}: Normalised spectral curves", saveto=save_to_normalised)
print("Saved normalised spectrum plot")

# Load and plot the final resulting spectral response curves
final_curves = np.load(folder/"monochromator_curve.npy")
# Arrays in [] so they can be looped over (a bit of a hacky solution)
final_mean = [final_curves[1:5].T]
final_stds = [final_curves[5:].T]

spectral.plot_monochromator_curves(wavelengths, final_mean, final_stds, title=f"{camera.name}: Combined spectral curves", unit="normalised", saveto=save_to_final)
print("Saved final spectrum plot")

# Calculate the signal-to-noise ratio (SNR) and plot it
SNR = final_mean[0] / final_stds[0]
SNR_err = np.zeros_like(SNR)  # don't plot errors on the SNR

spectral.plot_monochromator_curves(wavelengths, [SNR], [SNR_err], title=f"{camera.name}: Signal-to-noise ratio", unit="SNR", saveto=save_to_SNR)
print("Saved SNR plot")
