"""Rerun every AC85 tutorial number and compare it with expected-outputs.json.

Usage:  python3 reproduce.py            # compare against expected-outputs.json
        python3 reproduce.py --write    # overwrite expected-outputs.json
Requires Python >= 3.9, numpy, scipy and mpmath. Runtime is about one minute on a 2020s laptop.
"""
import json
import platform
import sys
import numpy as np
import scipy
import mpmath as mp
from point_lens_reference import analytic_amplification, geometric_optics_amplification
from independent_fresnel_integral import tabulate, fresnel_amplification
import signal_injection

FRESNEL_IMPACT_PARAMETERS = [0.1, 0.3, 1.0, 3.0]
FRESNEL_FREQUENCIES = [0.1, 1.0, 3.0, 10.0, 30.0]
TOLERANCE = {"relative_amplification": 2e-6, "delta_log_likelihood": 1e-6}


def fresnel_benchmark():
    rows = []
    for y in FRESNEL_IMPACT_PARAMETERS:
        tau, values = tabulate(y)
        for w in FRESNEL_FREQUENCIES:
            numeric, (c1, _) = fresnel_amplification(w, tau, values, y)
            analytic = analytic_amplification(w, y)
            rows.append({
                "w": w, "y": y,
                "fresnel": [numeric.real, numeric.imag],
                "analytic": [analytic.real, analytic.imag],
                "relative_difference": abs(numeric - analytic) / abs(analytic),
                "tail_coefficient": c1,
            })
    return rows


def limit_checks():
    low = []
    for y in [0.1, 1.0, 3.0]:
        for w in [1e-4, 1e-3, 1e-2]:
            value = analytic_amplification(w, y)
            low.append({"w": w, "y": y, "distance_from_unlensed": abs(value - 1)})
    high = []
    for y in [0.3, 1.0, 3.0]:
        for w in [30.0, 100.0, 300.0, 1000.0]:
            exact = analytic_amplification(w, y)
            high.append({"w": w, "y": y, "relative_geometric_optics_error": abs(exact - geometric_optics_amplification(w, y)) / abs(exact)})
    return {"low_frequency": low, "high_frequency": high}


def main():
    inputs = json.load(open("injection-inputs.json"))
    outputs = {
        "environment": {"python": platform.python_version(), "numpy": np.__version__, "scipy": scipy.__version__, "mpmath": mp.__version__},
        "fresnel_versus_analytic": fresnel_benchmark(),
        "limits": limit_checks(),
        "injections": signal_injection.run(inputs),
    }
    if "--write" in sys.argv:
        json.dump(outputs, open("expected-outputs.json", "w"), indent=2)
        print("wrote expected-outputs.json")
        return
    expected = json.load(open("expected-outputs.json"))
    worst = 0.0
    for got, want in zip(outputs["fresnel_versus_analytic"], expected["fresnel_versus_analytic"]):
        for key in ["fresnel", "analytic"]:
            a, b = complex(*got[key]), complex(*want[key])
            worst = max(worst, abs(a - b) / abs(b))
    injection_worst = max(abs(g.get("delta_log_likelihood_median", g.get("search_delta_log_likelihood"))
                              - w.get("delta_log_likelihood_median", w.get("search_delta_log_likelihood")))
                          for g, w in zip(outputs["injections"], expected["injections"]))
    ok = worst <= TOLERANCE["relative_amplification"] and injection_worst <= TOLERANCE["delta_log_likelihood"]
    print(json.dumps({"largest_relative_amplification_change": worst, "largest_injection_median_change": injection_worst, "pass": ok}, indent=2))
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()