#!/usr/bin/env python3
"""Optional standalone scientific figures. Requires Matplotlib 3.10.7 and NumPy."""
from pathlib import Path
import argparse
import csv
import gzip
import json
import sys

parser = argparse.ArgumentParser()
parser.add_argument("--directory", type=Path, default=Path("."))
parser.add_argument("--out", type=Path, default=Path("figures"))
parser.add_argument("--library-path", type=Path)
arguments = parser.parse_args()
if arguments.library_path:
    sys.path.insert(0, str(arguments.library_path))
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from timing_data import read_observations, ephemeris
from chopping_model import CoefficientGrid, predict

assert matplotlib.__version__ == "3.10.7"
plt.rcParams.update(
    {
        "font.size": 10,
        "font.family": "DejaVu Sans",
        "svg.hashsalt": "AC19-2026-10-28",
        "axes.spines.top": False,
        "axes.spines.right": False,
    }
)
arguments.out.mkdir(parents=True, exist_ok=True)


def samples(path):
    actual = path if path.exists() else path.with_suffix(path.suffix + ".gz")
    opener = gzip.open if actual.suffix == ".gz" else open
    with opener(actual, "rt") as file:
        rows = list(csv.DictReader(file))
    names = [name for name in rows[0] if name not in ["ensemble", "draw"]]
    return np.array([[float(row[name]) for name in names] for row in rows])


def save(figure, name):
    figure.savefig(
        arguments.out / (name + ".svg"), metadata={"Date": None}, bbox_inches="tight"
    )
    figure.savefig(arguments.out / (name + ".png"), dpi=170, bbox_inches="tight")
    plt.close(figure)


slow = samples(arguments.directory / "posterior/slow-samples.csv")
chopping = samples(arguments.directory / "posterior/chopping-samples.csv")
_, records = read_observations(arguments.directory)
grid = CoefficientGrid()
prediction_samples = np.array(
    [
        predict(row, records, grid, True)
        for row in chopping[:: max(1, len(chopping) // 600)]
    ]
)
median = np.median(prediction_samples, axis=0)
figure, axes = plt.subplots(2, 1, figsize=(3.8, 5.1), layout="constrained")
for axis, planet, title in zip(
    axes,
    ["Kepler-289d", "Kepler-289c"],
    ["Middle planet: 15 measured times", "Outer planet: 10 measured times"],
):
    selected = [
        (index, row) for index, row in enumerate(records) if row["planet"] == planet
    ]
    rows = [row for _, row in selected]
    epoch, period = ephemeris(rows)
    times = [row["observed_reduced_bjd"] for row in rows]
    residuals = [
        (row["observed_reduced_bjd"] - epoch - row["planet_epoch"] * period) * 1440
        for row in rows
    ]
    axis.errorbar(
        times,
        residuals,
        yerr=[row["sigma_days"] * 1440 for row in rows],
        fmt="o",
        markersize=3.8,
        color="#17293c",
        label="Measured + quoted error",
    )
    model = [
        (median[index] - epoch - row["planet_epoch"] * period) * 1440
        for index, row in selected
    ]
    axis.scatter(
        times, model, s=27, marker="x", color="#147eab", label="Model posterior median"
    )
    axis.axhline(0, color="#a8b8c7", linewidth=0.7)
    axis.set_title(title, fontsize=10)
    axis.set_ylabel("Timing residual / min")
    axis.grid(alpha=0.18)
handles, labels = axes[0].get_legend_handles_labels()
figure.legend(
    handles,
    labels,
    fontsize=8,
    loc="upper center",
    bbox_to_anchor=(0.5, 1.07),
    frameon=False,
)
axes[1].set_xlabel("Original reduced BJD / day")
save(figure, "kepler289-measured-timing")

figure, axes = plt.subplots(2, 1, figsize=(3.8, 5.1), layout="constrained")
for axis, column, maximum, label in zip(
    axes, [0, 1], [60, 600], ["Middle mass / Earth masses", "Outer mass / Earth masses"]
):
    bins = np.linspace(0, maximum, 61)
    axis.hist(
        slow[:, column],
        bins=bins,
        density=True,
        histtype="step",
        linewidth=1.6,
        color="#68737c",
        label="Slow only",
    )
    axis.hist(
        chopping[:, column],
        bins=bins,
        density=True,
        histtype="step",
        linewidth=1.8,
        color="#147eab",
        label="With chopping",
    )
    axis.set_xlim(0, maximum)
    axis.set_xlabel(label)
    axis.set_ylabel("Conditional density")
    axis.grid(alpha=0.18)
axes[0].legend(fontsize=9)
save(figure, "kepler289-conditional-masses")

figure, axes = plt.subplots(2, 1, figsize=(3.8, 5.3), layout="constrained")
for data, color, label in [
    (slow, "#68737c", "Slow only"),
    (chopping, "#147eab", "With chopping"),
]:
    points = data[:: max(1, len(data) // 1800)]
    axes[0].scatter(
        points[:, 0],
        points[:, 1],
        s=3,
        alpha=0.18,
        c=color,
        label=label,
        rasterized=False,
    )
    axes[1].scatter(
        points[:, 2], points[:, 1], s=3, alpha=0.18, c=color, rasterized=False
    )
axes[0].set_xlim(0, 60)
axes[0].set_xlabel("Middle mass / Earth masses")
axes[1].set_xlim(-0.15, 0.15)
axes[1].set_xlabel("Real free-eccentricity combination")
for axis in axes:
    axis.set_ylim(0, 600)
    axis.set_ylabel("Outer mass / Earth masses")
    axis.grid(alpha=0.18)
axes[0].legend(fontsize=8, markerscale=3)
save(figure, "kepler289-mass-degeneracy")
print(
    "Saved three independently exportable scientific figures with source-data and conditional-model labels"
)