#!/usr/bin/env python3
"""Detached AC19 timing tutorial. Python >=3.9, standard library only.

Download the four adjacent Python modules plus the attributed original MRT.
Reproduction uses only public science files, never an executable Lab Engine.
"""
from pathlib import Path
import argparse
import csv
import hashlib
import json
import math
import time
from timing_data import read_observations, ephemeris
from chopping_model import CoefficientGrid, coefficients, predict
from timing_posterior import run_ensemble, summarize


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, default=Path("."))
    parser.add_argument("--out", type=Path, default=Path("timing-output"))
    parser.add_argument("--iterations", type=int, default=16000)
    parser.add_argument("--ensembles", type=int, default=4)
    parser.add_argument("--eccentricity-scale", type=float, default=0.03)
    parser.add_argument("--model", choices=["both", "chopping", "slow"], default="both")
    parser.add_argument("--injection", action="store_true")
    arguments = parser.parse_args()
    if arguments.ensembles < 2 or arguments.iterations < 600:
        parser.error(
            "Use at least two independent ensembles and600 iterations; short probes are not posterior convergence evidence"
        )
    arguments.out.mkdir(parents=True, exist_ok=True)
    observed, records = read_observations(arguments.input)
    grid = CoefficientGrid()
    source_records = [dict(row) for row in records]
    injection = [6.0, 150.0, 0.008, -0.004, 1635.87, 1698.99, 66.028, 125.865, 0.0]
    if arguments.injection:
        import random

        generator = random.Random(190191)
        values = predict(injection, records, grid, True)
        for row, value in zip(records, values):
            row["observed_reduced_bjd"] = value + generator.gauss(0, row["sigma_days"])
    results = {
        "status": "Conditional approximate two-planet timing inference; not the published full-system mass posterior",
        "data_kind": (
            "synthetic injection at the actual 25 epochs and errors"
            if arguments.injection
            else "25 actual Kepler transit mid-times"
        ),
        "time_coordinate": "original BJD-2454000 days; source does not declare UTC/TDB",
        "observed_count": len(observed),
        "selected_count": len(records),
        "injection_truth": injection if arguments.injection else None,
        "injection_noise_seed": 190191 if arguments.injection else None,
        "injection_scope": (
            "Analytic-model closure at the same source epochs/errors; not an independent N-body validation"
            if arguments.injection
            else None
        ),
        "source_sha256": hashlib.sha256(
            (arguments.input / "ajac8553t1_mrt.txt").read_bytes()
        ).hexdigest(),
        "priors": {
            "middle_mass_earth": "Uniform from 0.05 to 60 Earth masses",
            "outer_mass_earth": "Uniform from 5 to 600 Earth masses",
            "Z_real_imag": "independent zero-mean Gaussian, sigma="
            + str(arguments.eccentricity_scale)
            + "; each truncated at +/-0.15",
            "middle_center_day": "Uniform from 1634.5 to 1636.8 reduced BJD at linear-clock epoch 10",
            "outer_center_day": "Uniform from 1698.5 to 1699.5 reduced BJD at linear-clock epoch 5",
            "middle_period_day": "Uniform from 65.9 to 66.2 days",
            "outer_period_day": "Uniform from 125.7 to 126.0 days",
            "jitter_day": "Half-normal scale 30 minutes, restricted to 0.00001 through 0.2 days; common white timing scatter",
            "stellar_mass": "Fixed 1.08 solar masses; 332946.0487 Earth masses per solar mass",
        },
        "models": {},
    }
    for chopping in [False, True]:
        key = "chopping" if chopping else "slow"
        if arguments.model not in ["both", key]:
            continue
        began = time.monotonic()
        runs = [
            run_ensemble(
                records,
                grid,
                chopping,
                19019 + 101 * index,
                arguments.iterations,
                arguments.eccentricity_scale,
            )
            for index in range(arguments.ensembles)
        ]
        summary = summarize(runs)
        best = max(runs, key=lambda run: run["best_log_density"])
        model_times = predict(best["best"], records, grid, chopping)
        residuals = [
            (row["observed_reduced_bjd"] - value) * 1440
            for row, value in zip(records, model_times)
        ]
        raw_chi = sum(
            (residual / (row["sigma_days"] * 1440)) ** 2
            for row, residual in zip(records, residuals)
        )
        inflated_chi = sum(
            (residual / math.hypot(row["sigma_days"] * 1440, best["best"][8] * 1440))
            ** 2
            for row, residual in zip(records, residuals)
        )
        results["models"][key] = {
            "summary": summary,
            "best_posterior_sample": best["best"],
            "representative_sample_definition": "maximum retained density in logmass/logjitter coordinates; residual chi-square is not a minimized chi-square",
            "residual_minutes": residuals,
            "raw_error_chi_square": raw_chi,
            "jitter_error_chi_square": inflated_chi,
            "runs": [{k: v for k, v in run.items() if k != "samples"} for run in runs],
        }
        with (arguments.out / (key + "-samples.csv")).open("w") as file:
            writer = csv.writer(file, lineterminator="\n")
            writer.writerow(["ensemble", "draw", *summary.keys()])
            for index, run in enumerate(runs):
                for draw, values in enumerate(run["samples"]):
                    writer.writerow(
                        [index, draw, *[format(value, ".10g") for value in values]]
                    )
        print(
            key,
            "elapsed_seconds",
            round(time.monotonic() - began, 2),
            "masses",
            summary["middle_mass_earth"],
            summary["outer_mass_earth"],
            flush=True,
        )
    with (arguments.out / "selected-times.csv").open("w") as file:
        exported = []
        for original, analyzed in zip(source_records, records):
            exported.append(
                {
                    **original,
                    "analysis_time_reduced_bjd": analyzed["observed_reduced_bjd"],
                    "analysis_kind": (
                        "synthetic analytic injection"
                        if arguments.injection
                        else "actual measured time"
                    ),
                }
            )
        writer = csv.DictWriter(file, fieldnames=list(exported[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(exported)
    (arguments.out / "timing-results.json").write_text(
        json.dumps(results, indent=2, allow_nan=False) + "\n"
    )
    print("Saved", arguments.out, flush=True)


if __name__ == "__main__":
    main()