#!/usr/bin/env python3
"""Selected maximum over sky positions AND subtraction choices on independent null cubes."""
import argparse
import json
import math
from pathlib import Path
import numpy as np
from imaging_reduction import reduce_cube, aperture_weights
from synthetic_optics import stellar_expectation


def make_grid(size=80):
    points = []
    for separation in (3, 4, 5, 6, 7):
        count = int(2 * math.pi * separation / 1.1)
        for index in range(count):
            angle = 2 * math.pi * index / count
            points.append((separation, angle))
    weights = np.asarray(
        [
            aperture_weights(
                size, radius * math.cos(angle), radius * math.sin(angle)
            ).ravel()
            for radius, angle in points
        ]
    )
    comparisons = []
    for radius, angle in points:
        comparisons.append(
            [
                index
                for index, (other_radius, other_angle) in enumerate(points)
                if radius == other_radius
                and abs(
                    math.atan2(
                        math.sin(angle - other_angle), math.cos(angle - other_angle)
                    )
                )
                >= math.radians(50) + 2 / radius
            ]
        )
    return points, weights, comparisons


def selected_maximum(data, angles, wavelengths, grid):
    points, weights, comparisons = grid
    best = None
    for algorithm, rank in [
        ("adi", 1),
        ("asdi", 1),
        ("klip", 1),
        ("klip", 2),
        ("klip", 3),
    ]:
        image = reduce_cube(data, angles, wavelengths, algorithm, rank)[2]
        values = weights @ image.ravel()
        for index, refs in enumerate(comparisons):
            sample = values[refs]
            if len(sample) < 3 or np.std(sample, ddof=1) <= 0:
                continue
            statistic = float(
                (values[index] - sample.mean())
                / (np.std(sample, ddof=1) * math.sqrt(1 + 1 / len(sample)))
            )
            if best is None or statistic > best["statistic"]:
                best = {
                    "statistic": statistic,
                    "algorithm": algorithm,
                    "rank": rank,
                    "separation": points[index][0],
                    "angle_degrees": math.degrees(points[index][1]),
                    "comparison_count": len(sample),
                }
    return best


def validate(materials, output, draws=64):
    metadata = json.loads((materials / "synthetic-input.json").read_text())
    arrays = np.load(materials / "synthetic-observation.npz")
    settings, angles = metadata["settings"], metadata["angles"]
    wavelengths = metadata["camera"]["wavelengthsMicrons"]
    null_expected = stellar_expectation(settings, angles, wavelengths)
    grid = make_grid()
    observed = selected_maximum(arrays["data"], angles, wavelengths, grid)
    trials = []
    for trial in range(draws):
        seed = 22024000 + trial
        rng = np.random.default_rng(seed)
        data = rng.poisson(null_expected) + rng.normal(0, 2, null_expected.shape)
        trials.append(
            {"seed": seed, **selected_maximum(data, angles, wavelengths, grid)}
        )
        print("null complete acquisition", trial + 1, "/", draws, flush=True)
    exceedances = sum(row["statistic"] >= observed["statistic"] for row in trials)
    result = {
        "draws": draws,
        "search_positions": len(grid[0]),
        "processing_choices": 5,
        "observed_selected_peak": observed,
        "null_selected_statistic_quantiles": np.quantile(
            [row["statistic"] for row in trials], [0.5, 0.95, 0.99]
        ).tolist(),
        "exceedances": exceedances,
        "add_one_fraction": (exceedances + 1) / (draws + 1),
        "zero_exceedance_95pct_upper_bound": (
            1 - 0.05 ** (1 / draws) if exceedances == 0 else None
        ),
        "trials": trials,
        "limits": "New independent Poisson/read noise and complete processing/search selection, with fixed prescribed OPD. No real atmospheric/survey false-alarm claim. Middle band only; source truth is never read by the candidate search. Zero exceedances have a finite binomial upper bound.",
    }
    output.write_text(json.dumps(result, indent=2) + "\n")
    print(
        "COMPLETE blind null experiment",
        observed,
        "exceedances",
        exceedances,
        flush=True,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--materials", type=Path, default=Path("."))
    parser.add_argument("--output", type=Path, default=Path("blind-null.json"))
    parser.add_argument("--draws", type=int, default=64)
    args = parser.parse_args()
    validate(args.materials, args.output, args.draws)