#!/usr/bin/env python3
"""Linearized same-input Monte Carlo comparison of legacy WLS and new FDE policy."""

from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path

import numpy as np


CHI2_99 = {1: 6.634896601, 2: 9.210340372, 3: 11.34486673,
           4: 13.27670414, 5: 15.08627247}


def design_matrix(count: int) -> tuple[np.ndarray, np.ndarray]:
    # Fixed, moderately realistic sky: azimuth degrees, elevation degrees.
    sky = np.array([
        [10.0, 72.0], [55.0, 28.0], [102.0, 46.0], [151.0, 19.0],
        [205.0, 61.0], [247.0, 34.0], [292.0, 23.0], [330.0, 52.0],
        [178.0, 38.0],
    ])[:count]
    azimuth = np.radians(sky[:, 0])
    elevation = np.radians(sky[:, 1])
    los = np.column_stack((
        np.cos(elevation) * np.sin(azimuth),
        np.cos(elevation) * np.cos(azimuth),
        np.sin(elevation),
    ))
    matrix = np.column_stack((-los, np.ones(count)))
    # Elevation-dependent standard deviation; low satellites carry less weight.
    sigma = 3.0 / np.maximum(np.sin(elevation), 0.20)
    return matrix, sigma


def solve(matrix: np.ndarray, sigma: np.ndarray, observation: np.ndarray) -> tuple[np.ndarray, float, float, bool]:
    weight = 1.0 / np.square(sigma)
    normal = matrix.T @ (weight[:, None] * matrix)
    try:
        state = np.linalg.solve(normal, matrix.T @ (weight * observation))
    except np.linalg.LinAlgError:
        return np.full(4, np.nan), math.inf, math.inf, False
    residual = observation - matrix @ state
    chi2 = float(np.sum(weight * np.square(residual)))
    dof = len(observation) - 4
    threshold = CHI2_99.get(dof, 0.0)
    ratio = chi2 / threshold if threshold > 0.0 else math.inf
    return state, chi2, ratio, threshold > 0.0 and chi2 <= threshold


def new_policy(matrix: np.ndarray, sigma: np.ndarray, observation: np.ndarray) -> tuple[np.ndarray, bool, bool]:
    full_state, _, full_ratio, full_pass = solve(matrix, sigma, observation)
    if full_pass or len(observation) < 7:
        return full_state, full_pass, False
    selected_state = full_state
    selected_ratio = full_ratio
    selected_pass = False
    fde_applied = False
    for excluded in range(len(observation)):
        keep = np.arange(len(observation)) != excluded
        state, _, ratio, passed = solve(matrix[keep], sigma[keep], observation[keep])
        if (passed and not selected_pass) or (passed == selected_pass and ratio < selected_ratio):
            selected_state = state
            selected_ratio = ratio
            selected_pass = passed
            fde_applied = True
    return selected_state, selected_pass, fde_applied


def percentile(values: np.ndarray, percentage: float) -> float:
    return float(np.percentile(values, percentage))


def scenario(rng: np.random.Generator, count: int, epochs: int, fault_m: float) -> dict[str, float | int]:
    matrix, sigma = design_matrix(count)
    old_horizontal = np.empty(epochs)
    new_horizontal = np.empty(epochs)
    delta = np.empty(epochs)
    valid = 0
    fde = 0
    improved = 0
    for epoch in range(epochs):
        observation = rng.normal(0.0, sigma)
        if fault_m != 0.0:
            observation[epoch % count] += fault_m
        old_state, _, _, _ = solve(matrix, sigma, observation)
        new_state, new_valid, fde_applied = new_policy(matrix, sigma, observation)
        old_horizontal[epoch] = np.linalg.norm(old_state[:2])
        new_horizontal[epoch] = np.linalg.norm(new_state[:2])
        delta[epoch] = np.linalg.norm(new_state[:3] - old_state[:3])
        valid += int(new_valid)
        fde += int(fde_applied)
        improved += int(new_horizontal[epoch] < old_horizontal[epoch])
    return {
        "satellites": count,
        "fault_m": fault_m,
        "epochs": epochs,
        "old_horizontal_mean_m": float(np.mean(old_horizontal)),
        "old_horizontal_p95_m": percentile(old_horizontal, 95.0),
        "new_horizontal_mean_m": float(np.mean(new_horizontal)),
        "new_horizontal_p95_m": percentile(new_horizontal, 95.0),
        "new_minus_old_mean_m": float(np.mean(new_horizontal - old_horizontal)),
        "solver_position_delta_mean_m": float(np.mean(delta)),
        "new_improved_rate_pct": 100.0 * improved / epochs,
        "new_valid_rate_pct": 100.0 * valid / epochs,
        "fde_rate_pct": 100.0 * fde / epochs,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=20000)
    parser.add_argument("--seed", type=int, default=20260712)
    parser.add_argument("--fault-m", type=float, default=75.0)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    rng = np.random.default_rng(args.seed)
    rows = [
        scenario(rng, count, args.epochs, fault)
        for fault in (0.0, args.fault_m)
        for count in range(4, 10)
    ]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    for row in rows:
        print(
            f"fault={row['fault_m']:>4.0f}m sats={row['satellites']} "
            f"old_mean={row['old_horizontal_mean_m']:.2f} "
            f"new_mean={row['new_horizontal_mean_m']:.2f} "
            f"delta={row['new_minus_old_mean_m']:+.2f} "
            f"improved={row['new_improved_rate_pct']:.1f}% "
            f"valid={row['new_valid_rate_pct']:.1f}% fde={row['fde_rate_pct']:.1f}%"
        )
    print(f"wrote {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())