#!/usr/bin/env python
import argparse
import json
from pathlib import Path

import numpy as np

from scoobdoob.validation import robustness as rob

HERE = Path(__file__).resolve().parent
DEFAULT_DIR = HERE / "data" / "validation_aib9"
K_DEFAULT, LAG_DEFAULT, DT_PS = 100, 5, 1.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--outdir", type=Path, default=HERE / "figures" / "validation_aib9")
    ap.add_argument(
        "--also",
        type=Path,
        default=None,
        help="second location to write the pair, e.g. the path paper.tex expects",
    )
    a = ap.parse_args()

    src = DEFAULT_DIR / "robustness.json"
    if not src.exists():
        raise SystemExit(f"missing {src}\nrun: python validate_aib9.py --robustness")
    d = json.loads(src.read_text())
    its = np.array(
        [[np.nan if v is None else v for v in row] for row in d["its"]], dtype=float
    )

    a.outdir.mkdir(parents=True, exist_ok=True)
    rob.make_figure(
        d["grid"],
        d["its_lags"],
        its,
        K_default=K_DEFAULT,
        lag_default=LAG_DEFAULT,
        dt_ps=DT_PS,
        outpath=a.outdir / "aib9_robustness.png",
    )
    print(f"-> {a.outdir / 'aib9_robustness.png'}")
    print(f"-> {a.outdir / 'aib9_robustness.pdf'}")

    if a.also:
        a.also.parent.mkdir(parents=True, exist_ok=True)
        for ext in ("png", "pdf"):
            (a.also.with_suffix("." + ext)).write_bytes(
                (a.outdir / f"aib9_robustness.{ext}").read_bytes()
            )
            print(f"-> {a.also.with_suffix('.' + ext)}")


if __name__ == "__main__":
    main()