#!/usr/bin/env python3
"""Export shareable synthetic-only scientific figures; no GPI image pixels."""
import argparse
import json
from pathlib import Path
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm, LinearSegmentedColormap
from imaging_reduction import reduce_cube


def main(materials, output):
    output.mkdir(parents=True, exist_ok=True)
    metadata = json.loads((materials / "synthetic-input.json").read_text())
    arrays = np.load(materials / "synthetic-observation.npz")
    images = reduce_cube(
        arrays["data"], metadata["angles"], metadata["camera"]["wavelengthsMicrons"]
    )
    plt.rcParams.update(
        {
            "font.size": 12,
            "svg.fonttype": "path",
            "svg.hashsalt": "AC22-cell-average",
            "axes.spines.top": False,
            "axes.spines.right": False,
            "figure.constrained_layout.h_pad": .16,
            "figure.constrained_layout.w_pad": .06,
        }
    )
    fig, axes = plt.subplots(1, 3, figsize=(10.8, 4.4), constrained_layout=True)
    raw_artist = None
    for axis, index in zip(axes[:2], [0, 17]):
        raw_artist = axis.imshow(
            arrays["data"][2, index],
            origin="upper",
            extent=(-10.125, 9.875, -9.875, 10.125),
            cmap="inferno",
            norm=LogNorm(vmin=1, vmax=100000),
        )
        axis.set_title("Synthetic exposure %d" % (index + 1))
        axis.set_xlabel("Camera x / reference angular units")
        axis.set_ylabel("Camera y / reference angular units")
    limit = np.quantile(np.abs(images[2]), 0.995)
    residual_artist = axes[2].imshow(
        images[2],
        origin="upper",
        extent=(-10.125, 9.875, -9.875, 10.125),
        cmap=LinearSegmentedColormap.from_list("signed_counts", ["#266ab3", "#071221", "#f7cb79"]),
        vmin=-limit,
        vmax=limit,
    )
    axes[2].set_title("Mean ADI, sky aligned")
    axes[2].set_xlabel("Sky x / reference angular units")
    axes[2].set_ylabel("Sky y / reference angular units")
    fig.colorbar(raw_artist, ax=list(axes[:2]), location="bottom", label="Camera electrons / pixel", shrink=.8)
    fig.colorbar(residual_artist, ax=axes[2], label="Residual electrons / pixel", shrink=.8)
    fig.savefig(output / "synthetic-processing.svg", metadata={"Date": None})
    plt.close(fig)
    response = json.loads((materials / "asdi-response.json").read_text())[
        "normalized_response"
    ]
    fig, axis = plt.subplots(figsize=(5.5, 4.7), constrained_layout=True)
    image = axis.imshow(response, cmap="RdBu_r", vmin=-0.14, vmax=0.14)
    labels = ["1.14", "1.19", "1.24", "1.29", "1.34"]
    axis.set_xticks(range(5), labels)
    axis.set_yticks(range(5), labels)
    axis.set_xlabel("Injected wavelength / micrometres")
    axis.set_ylabel("Output wavelength / micrometres")
    for row in range(5):
        for column in range(5):
            axis.text(
                column,
                row,
                "%.3f" % response[row][column],
                ha="center",
                va="center",
                color="white" if abs(response[row][column]) > 0.08 else "black",
            )
    fig.colorbar(image, ax=axis, label="Response / stellar entrance count")
    fig.savefig(output / "spectral-response.svg", metadata={"Date": None})
    plt.close(fig)
    summary = json.loads((materials / "turns-summary.json").read_text())["spectrum"]
    wavelengths = metadata["camera"]["wavelengthsMicrons"]
    fig, axis = plt.subplots(figsize=(6, 3.7), constrained_layout=True)
    truth = 0.002 * np.asarray([1, 0.9, 0.62, 0.48, 0.76])
    axis.plot(
        wavelengths,
        truth * 1000,
        "--",
        color="#bf7b21",
        label="Prescribed synthetic input",
    )
    axis.errorbar(
        wavelengths,
        np.asarray(summary["contrast"]) * 1000,
        yerr=np.asarray(summary["errors"]) * 1000,
        fmt="o",
        color="#087e9a",
        label="Conditional aperture estimate",
    )
    axis.set_xlabel("Wavelength / micrometres")
    axis.set_ylabel("Contrast / 0.001")
    axis.legend(frameon=False)
    fig.savefig(output / "synthetic-spectrum.svg", metadata={"Date": None})
    plt.close(fig)
    print("PASS three synthetic-only standalone SVG figures")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--materials", type=Path, default=Path("."))
    parser.add_argument("--output", type=Path, default=Path("figures"))
    args = parser.parse_args()
    main(args.materials, args.output)