"""Fit a toy trapezoid model to a synthetic TESS-like light curve.

Run from the repository root:

    python research/tess-timing/scripts/toy_transit_fit.py

The committed TESS light-curve output in the repository is a twelve-row loader
fixture with eight retained cadences and no transit. It cannot support a fit,
so this script generates a clearly labelled synthetic multi-transit light curve
at a two-minute cadence over 27 days, writes it beside the fit products, runs a
box matched-filter period search followed by a least-squares trapezoid fit, and
writes a JSON result and an SVG figure. No network request is made.

Every number in the output is synthetic. The result is not a planet detection,
not a measurement of a real star, and not a claim about any TESS target.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
from pathlib import Path

import numpy as np

PACKAGE_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = PACKAGE_ROOT.parents[1]
SRC = PACKAGE_ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

# isort: off
from tess_timing.toy_transit import (
    DEFAULT_INGRESS_FRACTION,
    fit_toy_transit,
    fold_phase,
    software_versions,
    trapezoid_flux,
)
# isort: on

DEFAULT_OUT_DIR = PACKAGE_ROOT / "data" / "transit-fit-run"
DEFAULT_FIGURE = REPO_ROOT / "docs" / "figures" / "transit-fit.svg"

SEED = 20261022
SECONDS_PER_DAY = 86400.0
CADENCE_DAYS = 2.0 / 1440.0
BASELINE_DAYS = 27.0
T_START_BTJD = 3000.0
TRUE_T0_BTJD = 3000.6
TRUE_PERIOD_DAYS = 3.5
TRUE_DEPTH = 0.012
TRUE_DURATION_DAYS = 0.16
TRUE_BASELINE = 0.0002
NOISE_SIGMA = 2.0e-4
INGRESS_FRACTION = DEFAULT_INGRESS_FRACTION

MIN_PERIOD_DAYS = 1.0
MAX_PERIOD_DAYS = 8.0
PERIOD_STEP_DAYS = 0.002
SEARCH_DURATIONS_DAYS = (0.08, 0.12, 0.16, 0.20, 0.24)

SVG_NAMESPACE = (
    "http://www.w3.org/2000/svg"  # sec-ok: http-url - XML namespace identifier
)
INK = "#1b1b1b"
MUTED = "#67645c"
GRID = "#dedbd2"
AXIS = "#8c887e"
DATA = "#3d5a80"
FIT = "#a24b2a"
TRUTH = "#245c4a"
BAND = "#f0ece2"


def _display(path: Path) -> str:
    try:
        return path.resolve().relative_to(REPO_ROOT.resolve()).as_posix()
    except ValueError:
        return str(path)


def _escape(value: object) -> str:
    return str(value).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def _text(
    x: float,
    y: float,
    value: object,
    size: float = 12.0,
    fill: str = INK,
    anchor: str = "start",
    weight: str = "normal",
) -> str:
    return (
        f'<text x="{x:.1f}" y="{y:.1f}" font-family="Georgia, serif" '
        f'font-size="{size:.1f}" fill="{fill}" text-anchor="{anchor}" '
        f'font-weight="{weight}">{_escape(value)}</text>'
    )


def _line(
    x1: float,
    y1: float,
    x2: float,
    y2: float,
    stroke: str,
    width: float = 1.0,
    dash: str | None = None,
    opacity: float | None = None,
) -> str:
    extra = f' stroke-dasharray="{dash}"' if dash else ""
    if opacity is not None:
        extra += f' stroke-opacity="{opacity:.2f}"'
    return (
        f'<line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" '
        f'stroke="{stroke}" stroke-width="{width:.2f}"{extra} />'
    )


def _rect(
    x: float,
    y: float,
    width: float,
    height: float,
    fill: str,
    stroke: str = "none",
    stroke_width: float = 0.0,
    opacity: float | None = None,
) -> str:
    extra = f' fill-opacity="{opacity:.2f}"' if opacity is not None else ""
    return (
        f'<rect x="{x:.1f}" y="{y:.1f}" width="{width:.1f}" '
        f'height="{height:.1f}" fill="{fill}" stroke="{stroke}" '
        f'stroke-width="{stroke_width:.2f}"{extra} />'
    )


def _circle(x: float, y: float, radius: float, fill: str) -> str:
    return f'<circle cx="{x:.1f}" cy="{y:.1f}" r="{radius:.1f}" fill="{fill}" />'


def _polyline(
    points: list[tuple[float, float]],
    stroke: str,
    width: float,
    dash: str | None = None,
) -> str:
    encoded = " ".join(f"{x:.1f},{y:.1f}" for x, y in points)
    extra = f' stroke-dasharray="{dash}"' if dash else ""
    return (
        f'<polyline points="{encoded}" fill="none" stroke="{stroke}" '
        f'stroke-width="{width:.2f}" stroke-linejoin="round" '
        f'stroke-linecap="round"{extra} />'
    )


def _scale(
    value: float, domain: tuple[float, float], span: tuple[float, float]
) -> float:
    low, high = domain
    return span[0] + (value - low) / (high - low) * (span[1] - span[0])


def _y_scale(
    value: float, domain: tuple[float, float], top: float, bottom: float
) -> float:
    return bottom - (value - domain[0]) / (domain[1] - domain[0]) * (bottom - top)


def _nice_ticks(low: float, high: float, count: int = 5) -> list[float]:
    span = high - low
    if span <= 0.0:
        return [low]
    raw_step = span / max(count - 1, 1)
    magnitude = 10.0 ** math.floor(math.log10(raw_step))
    residual = raw_step / magnitude
    if residual <= 1.0:
        step = 1.0
    elif residual <= 2.0:
        step = 2.0
    elif residual <= 5.0:
        step = 5.0
    else:
        step = 10.0
    step *= magnitude
    start = math.ceil(low / step) * step
    ticks: list[float] = []
    value = start
    while value <= high + 0.5 * step:
        ticks.append(value)
        value += step
    return ticks


def _panel(
    parts: list[str],
    x0: float,
    y0: float,
    width: float,
    height: float,
    x_domain: tuple[float, float],
    y_domain: tuple[float, float],
    x_ticks: list[float],
    y_ticks: list[float],
    x_label: str,
    y_label: str,
    x_format: str,
    y_format: str,
) -> tuple:
    def xmap(value: float) -> float:
        return _scale(value, x_domain, (x0, x0 + width))

    def ymap(value: float) -> float:
        return _y_scale(value, y_domain, y0, y0 + height)

    parts.append(_rect(x0, y0, width, height, "#ffffff", AXIS, 0.8))
    for value in y_ticks:
        y = ymap(value)
        parts.append(_line(x0, y, x0 + width, y, GRID, 0.6))
        parts.append(_text(x0 - 9, y + 4, y_format.format(value), 10.5, MUTED, "end"))
    for value in x_ticks:
        x = xmap(value)
        parts.append(_line(x, y0, x, y0 + height, GRID, 0.6))
        parts.append(
            _text(x, y0 + height + 17, x_format.format(value), 10.5, MUTED, "middle")
        )
    parts.append(
        _text(x0 + width / 2, y0 + height + 38, x_label, 11.5, MUTED, "middle")
    )
    parts.append(
        f'<text x="{x0 - 54:.1f}" y="{y0 + height / 2:.1f}" '
        f'font-family="Georgia, serif" font-size="11.5" fill="{MUTED}" '
        f'text-anchor="middle" transform="rotate(-90 {x0 - 54:.1f} '
        f'{y0 + height / 2:.1f})">{_escape(y_label)}</text>'
    )
    return xmap, ymap


def _bin_series(
    time: np.ndarray, flux: np.ndarray, flux_err: np.ndarray, width: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    edges = np.arange(float(time.min()), float(time.max()) + width, width)
    index = np.digitize(time, edges) - 1
    centres: list[float] = []
    values: list[float] = []
    errors: list[float] = []
    for i in range(edges.size - 1):
        selected = index == i
        if not bool(np.any(selected)):
            continue
        weights = 1.0 / (flux_err[selected] ** 2)
        centres.append(float(np.mean(time[selected])))
        values.append(float(np.sum(weights * flux[selected]) / np.sum(weights)))
        errors.append(float(math.sqrt(1.0 / np.sum(weights))))
    return np.array(centres), np.array(values), np.array(errors)


def _bin_phase(
    phase: np.ndarray, flux: np.ndarray, flux_err: np.ndarray, width: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    edges = np.arange(-0.5, 0.5 + width, width)
    index = np.digitize(phase, edges) - 1
    centres: list[float] = []
    values: list[float] = []
    errors: list[float] = []
    for i in range(edges.size - 1):
        selected = index == i
        if not bool(np.any(selected)):
            continue
        weights = 1.0 / (flux_err[selected] ** 2)
        centres.append(float(0.5 * (edges[i] + edges[i + 1])))
        values.append(float(np.sum(weights * flux[selected]) / np.sum(weights)))
        errors.append(float(math.sqrt(1.0 / np.sum(weights))))
    return np.array(centres), np.array(values), np.array(errors)


def _make_figure(
    path: Path,
    time: np.ndarray,
    flux: np.ndarray,
    flux_err: np.ndarray,
    parameters: dict[str, float],
    uncertainties: dict[str, float],
    truth: dict[str, float],
    ingress_fraction: float,
) -> None:
    width, height = 1080, 830
    parts = [
        (
            f'<svg xmlns="{SVG_NAMESPACE}" viewBox="0 0 {width} {height}" '
            f'width="{width}" height="{height}">'
        ),
        _rect(0, 0, width, height, "#fbfaf6"),
        _text(54, 44, "Toy trapezoid fit to a synthetic multi-transit light curve", 21),
        _text(
            54,
            66,
            "Synthetic two-minute cadence data. This is a model-recovery test, "
            "not a planet detection.",
            12,
            MUTED,
        ),
    ]

    # Top panel: full light curve, binned to 30 minutes for display.
    x0, y0, panel_width, panel_height = 92, 98, 900, 280
    time_centres, flux_bins, _ = _bin_series(time, flux, flux_err, 10.0 / 1440.0)
    model_time = np.arange(
        float(time.min()), float(time.max()) + 5.0 / 1440.0, 5.0 / 1440.0
    )
    model_dense = trapezoid_flux(
        model_time,
        parameters["t0"],
        parameters["period"],
        parameters["depth"],
        parameters["duration"],
        parameters["baseline_offset"],
        ingress_fraction,
    )
    y_low = min(float(np.min(flux_bins)), float(np.min(model_dense)))
    y_high = max(float(np.max(flux_bins)), float(np.max(model_dense)))
    margin = 0.12 * (y_high - y_low)
    x_domain = (float(time.min()), float(time.max()))
    y_domain = (y_low - margin, y_high + margin)
    x_map, y_map = _panel(
        parts,
        x0,
        y0,
        panel_width,
        panel_height,
        x_domain,
        y_domain,
        [3000, 3005, 3010, 3015, 3020, 3025],
        _nice_ticks(y_domain[0], y_domain[1]),
        "BTJD (day)",
        "Relative flux",
        "{:.0f}",
        "{:.3f}",
    )
    parts.append(
        _polyline(
            [
                (x_map(float(t)), y_map(float(v)))
                for t, v in zip(model_time, model_dense)
            ],
            FIT,
            1.5,
        )
    )
    for t, v in zip(time_centres, flux_bins):
        parts.append(_circle(x_map(float(t)), y_map(float(v)), 1.5, DATA))

    epoch = truth["t0"]
    epochs = []
    while epoch < float(time.max()) + truth["period"]:
        if epoch > float(time.min()):
            epochs.append(epoch)
        epoch += truth["period"]
    for epoch in epochs:
        x = x_map(epoch)
        parts.append(_line(x, y0 + panel_height - 9, x, y0 + panel_height, MUTED, 1.0))

    parts.append(
        _text(x0, y0 - 12, "Full light curve, 10-minute display bins", 12, INK)
    )
    parts.append(
        _line(
            x0 + panel_width - 210, y0 - 17, x0 + panel_width - 180, y0 - 17, DATA, 2.0
        )
    )
    parts.append(_text(x0 + panel_width - 172, y0 - 13, "binned data", 10.5, MUTED))
    parts.append(
        _line(x0 + panel_width - 90, y0 - 17, x0 + panel_width - 60, y0 - 17, FIT, 2.0)
    )
    parts.append(_text(x0 + panel_width - 52, y0 - 13, "fit", 10.5, MUTED))

    # Bottom panel: phase-folded curve.
    bx0, by0, bwidth, bheight = 92, 482, 900, 260
    phase = fold_phase(time, parameters["t0"], parameters["period"])
    phase_centres, phase_bins, phase_errors = _bin_phase(phase, flux, flux_err, 0.01)
    model_phase = trapezoid_flux(
        phase_centres,
        0.0,
        parameters["period"],
        parameters["depth"],
        parameters["duration"],
        parameters["baseline_offset"],
        ingress_fraction,
    )
    truth_phase = trapezoid_flux(
        phase_centres,
        0.0,
        parameters["period"],
        truth["depth"],
        truth["duration"],
        truth["baseline"],
        ingress_fraction,
    )
    phase_y_low = min(float(np.min(phase_bins)), float(np.min(truth_phase)))
    phase_y_high = max(float(np.max(phase_bins)), float(np.max(truth_phase)))
    phase_margin = 0.15 * (phase_y_high - phase_y_low)
    phase_y_domain = (phase_y_low - phase_margin, phase_y_high + phase_margin)
    bx_map, by_map = _panel(
        parts,
        bx0,
        by0,
        bwidth,
        bheight,
        (-0.5, 0.5),
        phase_y_domain,
        [-0.5, -0.25, 0.0, 0.25, 0.5],
        _nice_ticks(phase_y_domain[0], phase_y_domain[1]),
        "Phase (periods from mid-transit)",
        "Relative flux",
        "{:.2f}",
        "{:.3f}",
    )
    for centre, value, error in zip(phase_centres, phase_bins, phase_errors):
        x = bx_map(float(centre))
        y = by_map(float(value))
        y_top = by_map(float(value + error))
        y_bottom = by_map(float(value - error))
        parts.append(_line(x, y_top, x, y_bottom, DATA, 0.8, opacity=0.65))
        parts.append(_circle(x, y, 2.2, DATA))
    parts.append(
        _polyline(
            [
                (bx_map(float(t)), by_map(float(v)))
                for t, v in zip(phase_centres, model_phase)
            ],
            FIT,
            2.0,
        )
    )
    parts.append(
        _polyline(
            [
                (bx_map(float(t)), by_map(float(v)))
                for t, v in zip(phase_centres, truth_phase)
            ],
            TRUTH,
            1.2,
            "5 4",
        )
    )

    parts.append(_text(bx0, by0 - 12, "Phase-folded curve, 0.01-period bins", 12, INK))
    parts.append(
        _line(bx0 + bwidth - 300, by0 - 17, bx0 + bwidth - 270, by0 - 17, DATA, 2.0)
    )
    parts.append(_text(bx0 + bwidth - 262, by0 - 13, "binned data", 10.5, MUTED))
    parts.append(
        _line(bx0 + bwidth - 190, by0 - 17, bx0 + bwidth - 160, by0 - 17, FIT, 2.0)
    )
    parts.append(_text(bx0 + bwidth - 152, by0 - 13, "fit", 10.5, MUTED))
    parts.append(
        _line(
            bx0 + bwidth - 110, by0 - 17, bx0 + bwidth - 80, by0 - 17, TRUTH, 1.2, "5 4"
        )
    )
    parts.append(_text(bx0 + bwidth - 72, by0 - 13, "injected", 10.5, MUTED))

    summary = (
        "Recovered: period {period:.6f} +/- {period_err:.6f} d, "
        "depth {depth:.5f} +/- {depth_err:.5f}, "
        "duration {duration:.5f} +/- {duration_err:.5f} d"
    ).format(
        period=parameters["period"],
        period_err=uncertainties["period"],
        depth=parameters["depth"],
        depth_err=uncertainties["depth"],
        duration=parameters["duration"],
        duration_err=uncertainties["duration"],
    )
    parts.append(_text(54, 790, summary, 11.5, MUTED))
    parts.append(
        _text(
            54,
            809,
            "Generated by research/tess-timing/scripts/toy_transit_fit.py. "
            "No archive data, real target, detection, or discovery is involved.",
            11,
            MUTED,
        )
    )
    parts.append("</svg>")
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(parts) + "\n", encoding="utf-8")


def _generate_curve() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    count = math.floor(BASELINE_DAYS / CADENCE_DAYS) + 1
    time = T_START_BTJD + np.arange(count, dtype=float) * CADENCE_DAYS
    model = trapezoid_flux(
        time,
        TRUE_T0_BTJD,
        TRUE_PERIOD_DAYS,
        TRUE_DEPTH,
        TRUE_DURATION_DAYS,
        TRUE_BASELINE,
        INGRESS_FRACTION,
    )
    rng = np.random.default_rng(SEED)
    flux = model + rng.normal(0.0, NOISE_SIGMA, time.size)
    flux_err = np.full(time.size, NOISE_SIGMA)
    return time, flux, flux_err


def _write_csv(
    path: Path, time: np.ndarray, flux: np.ndarray, flux_err: np.ndarray
) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    data = np.column_stack([time, flux, flux_err])
    np.savetxt(
        path,
        data,
        delimiter=",",
        header="time_btjd,flux,flux_err",
        comments="",
        fmt=["%.9f", "%.10f", "%.10f"],
    )
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Fit a toy trapezoid model to a synthetic multi-transit light "
            "curve and write a CSV, a JSON record, and an SVG figure. No "
            "network request is made."
        )
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=None,
        help=(
            "directory for the CSV, JSON, and SVG outputs (default: "
            f"{_display(DEFAULT_OUT_DIR)} and {_display(DEFAULT_FIGURE)}); "
            "when set, the figure is written beside the CSV and JSON"
        ),
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    arguments = _parser().parse_args(argv)
    if arguments.out_dir is None:
        out_dir = DEFAULT_OUT_DIR
        figure_path = DEFAULT_FIGURE
    else:
        out_dir = arguments.out_dir
        figure_path = out_dir / "transit-fit.svg"
    out_dir.mkdir(parents=True, exist_ok=True)
    csv_path = out_dir / "synthetic_transit.csv"
    json_path = out_dir / "transit_fit.json"

    time, flux, flux_err = _generate_curve()
    digest = _write_csv(csv_path, time, flux, flux_err)

    truth = {
        "t0_btjd": TRUE_T0_BTJD,
        "period_days": TRUE_PERIOD_DAYS,
        "depth": TRUE_DEPTH,
        "duration_days": TRUE_DURATION_DAYS,
        "baseline_offset": TRUE_BASELINE,
        "ingress_fraction": INGRESS_FRACTION,
        "noise_sigma": NOISE_SIGMA,
        "seed": SEED,
    }
    fit, search = fit_toy_transit(
        time,
        flux,
        flux_err,
        min_period=MIN_PERIOD_DAYS,
        max_period=MAX_PERIOD_DAYS,
        period_step=PERIOD_STEP_DAYS,
        durations=SEARCH_DURATIONS_DAYS,
        ingress_fraction=INGRESS_FRACTION,
    )

    _make_figure(
        figure_path,
        time,
        flux,
        flux_err,
        fit.parameters,
        fit.uncertainties,
        {
            "t0": TRUE_T0_BTJD,
            "period": TRUE_PERIOD_DAYS,
            "depth": TRUE_DEPTH,
            "duration": TRUE_DURATION_DAYS,
            "baseline": TRUE_BASELINE,
        },
        INGRESS_FRACTION,
    )

    record = {
        "task": "ph-sim-toy",
        "schema_version": 1,
        "synthetic": True,
        "label": (
            "Synthetic TESS-like light curve generated from a trapezoid model. "
            "Not an observation and not a detection."
        ),
        "data": {
            "path": _display(csv_path),
            "sha256": digest,
            "cadence_seconds": CADENCE_DAYS * SECONDS_PER_DAY,
            "cadence_count": int(time.size),
            "baseline_days": BASELINE_DAYS,
            "truth": truth,
        },
        "search": search,
        "fit": fit.to_dict(),
        "figure": _display(figure_path),
        "software": software_versions(),
        "limits": [
            "The model is a trapezoid, not a physical transit model.",
            "The light curve is synthetic and was generated by this script.",
            "The fit result is not a planet detection.",
            "Period, depth, duration, and timing are not measurements of a real object.",
        ],
    }
    json_path.write_text(
        json.dumps(record, indent=2, sort_keys=False) + "\n", encoding="utf-8"
    )

    print(
        "status=OK synthetic=yes "
        f"cadences={fit.cadence_count} "
        f"period={fit.period:.6f}+/-{fit.uncertainties['period']:.6f} "
        f"depth={fit.depth:.6f}+/-{fit.uncertainties['depth']:.6f} "
        f"duration={fit.duration:.6f}+/-{fit.uncertainties['duration']:.6f} "
        f"rms_ppm={fit.rms_ppm:.2f} chi2_dof={fit.chi_square_per_dof:.4f} "
        f"converged={str(fit.converged).lower()} "
        f"wrote={_display(csv_path)},{_display(json_path)},{_display(figure_path)}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())