"""Per-installation depth-bias correction for the Nuwa-HP60C.

The Nuwa depth sensor exhibits an offset on `depth_mm` that doesn't fit a
simple multiplicative scale (see /api/perception/measure_rgb experiments
on a known marker — bias varies with distance). We model the bias as an
affine map:

    depth_raw = scale * depth_true + offset_mm     (sensor)
    depth_true = (depth_raw - offset_mm) / scale   (correction)

and fit (scale, offset_mm) from N≥2 samples where `depth_true` is taken
from a trusted source (typically a PnP-derived Z on a known-size ArUco
marker held at varying distances).

State (params + samples) lives in a single JSON file so it can be
inspected/edited by hand:

    data/depth_correction.json
    {
      "enabled": true,
      "scale": 0.871,
      "offset_mm": 185.0,
      "fit_meta": {"n_samples": 4, "rms_residual_mm": 2.1, ...},
      "samples": [
        {"depth_raw_mm": 1030, "depth_true_mm": 970, "note": "...", "added_at": ...},
        ...
      ]
    }
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

import numpy as np


def default_state() -> dict[str, Any]:
    return {
        "enabled": False,
        "scale": 1.0,
        "offset_mm": 0.0,
        "fit_meta": None,
        "samples": [],
    }


def load(path: Path) -> dict[str, Any]:
    if not path.exists():
        return default_state()
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, dict):
            return default_state()
        # Fill in any missing keys with defaults.
        merged = default_state()
        merged.update(data)
        return merged
    except json.JSONDecodeError:
        return default_state()


def save(path: Path, state: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


def apply_correction(depth_raw_mm: int, state: dict[str, Any] | None) -> int:
    """Convert a raw sensor depth (mm) to the corrected depth (mm). Returns
    the input unchanged if correction is disabled or state is None."""
    if not state or not state.get("enabled") or depth_raw_mm <= 0:
        return depth_raw_mm
    scale = float(state.get("scale") or 1.0)
    offset = float(state.get("offset_mm") or 0.0)
    if scale <= 0:
        return depth_raw_mm
    corrected = (depth_raw_mm - offset) / scale
    return max(0, int(round(corrected)))


def fit_affine(samples: list[dict[str, Any]]) -> dict[str, Any]:
    """Least-squares fit `depth_raw = scale * depth_true + offset_mm`.

    Needs ≥2 samples with non-degenerate range of depth_true. With exactly
    2 samples the fit is exact and rms_residual_mm = 0. With more samples
    the residual is meaningful for assessing fit quality.
    """
    if len(samples) < 2:
        raise ValueError("need at least 2 samples to fit affine model")
    raw = np.array([float(s["depth_raw_mm"]) for s in samples], dtype=np.float64)
    true = np.array([float(s["depth_true_mm"]) for s in samples], dtype=np.float64)
    if true.max() - true.min() < 50:
        raise ValueError("samples cover too narrow a range of depth_true (<50 mm)")
    A = np.column_stack([true, np.ones_like(true)])
    sol, *_ = np.linalg.lstsq(A, raw, rcond=None)
    scale = float(sol[0])
    offset = float(sol[1])
    pred = scale * true + offset
    residuals = pred - raw
    rms = float(np.sqrt(np.mean(residuals ** 2)))
    max_abs = float(np.max(np.abs(residuals)))
    return {
        "scale": scale,
        "offset_mm": offset,
        "n_samples": len(samples),
        "rms_residual_mm": rms,
        "max_abs_residual_mm": max_abs,
        "depth_true_min_mm": float(true.min()),
        "depth_true_max_mm": float(true.max()),
        "fitted_at": time.time(),
    }


def add_sample(state: dict[str, Any], depth_raw_mm: int, depth_true_mm: int,
               note: str = "") -> dict[str, Any]:
    samples = list(state.get("samples") or [])
    samples.append({
        "depth_raw_mm": int(depth_raw_mm),
        "depth_true_mm": int(depth_true_mm),
        "note": note,
        "added_at": time.time(),
    })
    new_state = dict(state)
    new_state["samples"] = samples
    return new_state


def clear_samples(state: dict[str, Any]) -> dict[str, Any]:
    new_state = dict(state)
    new_state["samples"] = []
    return new_state