#!/usr/bin/env python3
"""Reconstruct the published near-circular quarterly angle diagnostic.

Python standard library only. This is not the paper's photodynamical posterior.
Inputs keep the original table ephemerides, asymmetric intervals and BJD label.
"""
from pathlib import Path
import argparse
import csv
import json
import math


def angle_cycles(time, shifts, ephemerides, coefficients):
    return sum(
        coefficient
        * (
            0.25
            + (time - ephemerides[planet]["zeroTransitBjdOffset"] - shifts[planet])
            / ephemerides[planet]["periodDays"]
        )
        for planet, coefficient in zip("bcde", coefficients)
    )


def rectangle_interval(time, measurements, ephemerides, coefficients, level):
    centre_shifts = {planet: measurements[planet]["best_ttv_days"] for planet in "bcde"}
    centre = angle_cycles(time, centre_shifts, ephemerides, coefficients) * 360
    low, high = centre, centre
    for planet, coefficient in zip("bcde", coefficients):
        derivative = -360 * coefficient / ephemerides[planet]["periodDays"]
        bounds = [
            derivative * measurements[planet]["lower_" + level + "_delta_days"],
            derivative * measurements[planet]["upper_" + level + "_delta_days"],
        ]
        low += min(bounds)
        high += max(bounds)
    return [low - centre, high - centre]


def reconstruct(directory):
    metadata = json.loads((directory / "kepler223-source.json").read_text())
    with (directory / "kepler223-quarterly-ttv.csv").open() as source:
        rows = [
            {
                key: value if key == "planet" else float(value)
                for key, value in row.items()
            }
            for row in csv.DictReader(source)
        ]
    assert len(rows) == 44
    assert all(sum(row["planet"] == planet for row in rows) == 11 for planet in "bcde")
    for row in rows:
        assert (
            row["lower_three_delta_days"]
            <= row["lower_one_delta_days"]
            <= 0
            <= row["upper_one_delta_days"]
            <= row["upper_three_delta_days"]
        )
    ephemerides = metadata["ephemerides"]
    coefficients = {
        "inner": [-1, 2, -1, 0],
        "outer": [0, 1, -3, 2],
        "fourBody": [3, -4, -3, 4],
    }
    bins = []
    for index in range(1, 12):
        selected = {
            row["planet"]: row for row in rows if row["common_quarter_index"] == index
        }
        assert len(selected) == 4
        times = [row["mean_bjd_minus_2454900"] for row in selected.values()]
        time = sum(times) / 4
        shifts = {planet: selected[planet]["best_ttv_days"] for planet in "bcde"}
        angles = {}
        for name, weights in coefficients.items():
            angle = angle_cycles(time, shifts, ephemerides, weights) * 360 % 360
            alignment_span = (
                abs(
                    angle_cycles(max(times), shifts, ephemerides, weights)
                    - angle_cycles(min(times), shifts, ephemerides, weights)
                )
                * 360
            )
            angles[name] = {
                "degrees": angle,
                "deltaChiSquareOneRectangleOffsetsDegrees": rectangle_interval(
                    time, selected, ephemerides, weights, "one"
                ),
                "deltaChiSquareNineRectangleOffsetsDegrees": rectangle_interval(
                    time, selected, ephemerides, weights, "three"
                ),
                "quarterTimeChoiceSpanDegrees": alignment_span,
                "zeroTtvLinearEphemerisDegrees": angle_cycles(
                    time, {planet: 0 for planet in "bcde"}, ephemerides, weights
                )
                * 360
                % 360,
            }
        bins.append(
            {
                "commonQuarterIndex": index,
                "representativeBjdOffset": time,
                "planetMeanTimesBjdOffset": times,
                "angles": angles,
            }
        )
    return {
        "kind": "Near-circular finite-window diagnostic from44 published measured quarterly TTVs",
        "timeLabel": metadata["timeConvention"],
        "quarterTimeChoice": "Arithmetic mean of the four actual planet mean times; within-quarter TTV shifts held constant",
        "coefficients": coefficients,
        "bins": bins,
        "rangesDegrees": {
            name: [
                min(row["angles"][name]["degrees"] for row in bins),
                max(row["angles"][name]["degrees"] for row in bins),
            ]
            for name in coefficients
        },
        "maximumQuarterAlignmentSpanDegrees": {
            name: max(
                row["angles"][name]["quarterTimeChoiceSpanDegrees"] for row in bins
            )
            for name in coefficients
        },
        "intervalMeaning": "Extrema of a rectangle made from the four individual asymmetric profile intervals, not a joint confidence/credible interval and not an independence assumption.",
        "limitations": "Quarter binning, near-circular longitude approximation and unknown joint covariance. A short bounded arc alone is not proof of long-term resonance. No gas-disk history or full photodynamical posterior is inferred.",
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--directory", type=Path, default=Path(__file__).parent)
    parser.add_argument("--output", type=Path, default=Path("kepler223-angles.json"))
    arguments = parser.parse_args()
    result = reconstruct(arguments.directory)
    arguments.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n")
    print(
        json.dumps(
            {
                "rangesDegrees": result["rangesDegrees"],
                "maximumQuarterAlignmentSpanDegrees": result[
                    "maximumQuarterAlignmentSpanDegrees"
                ],
            }
        )
    )


if __name__ == "__main__":
    main()