"""JAX support-estimate corruption model (Step 3).

Reproduces the *error behaviour* of hackathon-everest's sklearn ExtraTrees estimator so an
MJX/Brax training loop can consume a realistic support signal without ever calling sklearn.

Error structure is measured, not assumed (see scripts/calibrate_from_everest.py):

1. Regression-to-the-mean shrinkage. The tree ensemble compresses its output range, so
   `pred ~= intercept + slope * true`. This is the dominant error term and it is
   *systematic* -- a policy cannot average it away the way it can average away noise.
   Modelling it as zero-mean Gaussian would make the task substantially easier than reality.
2. Conditional Gaussian noise around that shrunk value (smaller than the marginal residual
   spread, because most of the marginal spread IS the shrinkage).
3. False-safe events, latched per stance rather than per timestep -- see FALSE_SAFE_NOTE.
4. Dropout that holds the last valid value and drops confidence.
5. Per-environment noise scaling so no single calibration is memorised.

All functions are pure and vmap/jit-safe: no data-dependent Python control flow.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import NamedTuple

import jax
import jax.numpy as jnp

CALIBRATION_PATH = Path(__file__).with_name("calibration.json")

DEPTH_NORM_M = 0.25          # RL terrain range 0..0.25 m
SUPPORT_NORM_N = 900.0       # material_priors bearing_capacity_n upper bound
N_FEET = 2
N_CHANNELS = 4               # depth_est, support_est, confidence, void_prob
OBS_SIZE = N_FEET * N_CHANNELS

# Fitted over the terrain generator's own distribution (see scripts/fit_support_to_depth.py).
# Valid only while bearing is derived from the sinkage law; regenerate if that changes.
SUPPORT_TO_DEPTH_SLOPE = -3.244e-4
SUPPORT_TO_DEPTH_INTERCEPT = 0.1833
# Chosen by measurement, and RE-measured whenever the calibration contract changes.
# Under the real 10-channel rig: corr(implied, true depth) peaks at 0.691 with w=0.50
# (0.683 at 0.75, 0.664 depth-only, 0.327 support-only). Under the optimistic 19-channel
# contract the optimum was 0.75 at corr 0.776 -- the real rig costs ~0.09 of correlation.
# Support contributes depth information mainly in shallow terrain: past the failure sinkage
# the substrate never carries load, so bearing stops distinguishing 12 cm from 25 cm.
DEPTH_CHANNEL_WEIGHT = 0.50

FALSE_SAFE_NOTE = """\
everest measures two very different false-safe rates on the same held-out split:
  0.005 when the consumer sees the CONSERVATIVE lower bound (mean - 2*sigma)
  0.303 when the consumer sees the POINT estimate
The brief's observation block hands the policy a point estimate plus a separate confidence,
which puts the honest rate near 0.303, not 0.005. Applying 0.005 would make the environment
~60x more forgiving than the estimator it claims to model. Default here is the point-estimate
rate; set `false_safe_rate` explicitly to override."""


class Channel(NamedTuple):
    slope: float
    intercept: float
    sigma: float


class EstimatorParams(NamedTuple):
    depth: Channel
    support: Channel
    false_safe_rate: float
    false_safe_support_n: float   # what a false-safe estimate reports
    dropout_prob: float
    noise_scale_range: tuple[float, float]
    confidence_resid_corr: float
    void_false_positive_rate: float
    void_false_negative_rate: float


class EstimatorState(NamedTuple):
    """Per-environment carry. Leading axis is the env batch under vmap."""
    last_obs: jax.Array        # (N_FEET, N_CHANNELS) held values for dropout
    false_safe: jax.Array      # (N_FEET,) bool, latched for the current stance
    noise_scale: jax.Array     # () per-env multiplier on all sigmas


def load_params(
    path: Path | str = CALIBRATION_PATH,
    *,
    false_safe_rate: float | None = None,
    dropout_prob: float = 0.02,
    void_false_positive_rate: float = 0.05,
    void_false_negative_rate: float = 0.10,
) -> EstimatorParams:
    """Build params from measured calibration.

    void_* rates default to non-zero on purpose. everest's void classifier scores a perfect
    1.0/1.0/0.0 on held-out data, which would hand the policy an oracle for voids and quietly
    violate the "rewards never key off ground truth" constraint. Injecting realistic error
    keeps void_prob an estimate rather than a truth channel.
    """
    cal = json.loads(Path(path).read_text())
    ch = cal["channels"]
    fs = cal["false_safe"]
    rate = fs["rate_point_estimate"] if false_safe_rate is None else false_safe_rate
    return EstimatorParams(
        depth=Channel(ch["depth"]["shrink_slope"], ch["depth"]["shrink_intercept"],
                      ch["depth"]["conditional_sigma"]),
        support=Channel(ch["support"]["shrink_slope"], ch["support"]["shrink_intercept"],
                        ch["support"]["conditional_sigma"]),
        false_safe_rate=float(rate),
        false_safe_support_n=float(fs["unsafe_threshold_n"] * 1.25),
        dropout_prob=float(dropout_prob),
        noise_scale_range=(0.7, 1.4),
        confidence_resid_corr=float(ch["support"]["uncertainty_resid_corr"]),
        void_false_positive_rate=float(void_false_positive_rate),
        void_false_negative_rate=float(void_false_negative_rate),
    )


def init_state(key: jax.Array, params: EstimatorParams) -> EstimatorState:
    """Initialise one environment's carry. vmap over keys for a batch."""
    lo, hi = params.noise_scale_range
    return EstimatorState(
        last_obs=jnp.zeros((N_FEET, N_CHANNELS)),
        false_safe=jnp.zeros((N_FEET,), dtype=bool),
        noise_scale=jax.random.uniform(key, (), minval=lo, maxval=hi),
    )


def _shrink(true_value: jax.Array, ch: Channel) -> jax.Array:
    return ch.intercept + ch.slope * true_value


def estimate(
    key: jax.Array,
    state: EstimatorState,
    true_depth_m: jax.Array,      # (N_FEET,)
    true_support_n: jax.Array,    # (N_FEET,)
    true_void: jax.Array,         # (N_FEET,) bool
    new_stance: jax.Array,        # (N_FEET,) bool -- foot just made contact
    params: EstimatorParams,
) -> tuple[jax.Array, EstimatorState]:
    """One 50 Hz estimator tick for one environment.

    Returns (obs, state) where obs is (OBS_SIZE,) laid out per foot as
    [depth_est_norm, support_est_norm, confidence, void_prob].
    """
    k_d, k_s, k_fs, k_drop, k_conf, k_void = jax.random.split(key, 6)
    scale = state.noise_scale

    depth = _shrink(true_depth_m, params.depth) + params.depth.sigma * scale * jax.random.normal(
        k_d, true_depth_m.shape
    )
    support = _shrink(true_support_n, params.support) + params.support.sigma * scale * jax.random.normal(
        k_s, true_support_n.shape
    )

    # False-safe: resampled only on a NEW STANCE, then held for that stance's duration.
    # Sampling per 50 Hz timestep instead would fire ~20x per 0.4 s stance and inflate the
    # effective per-footstep rate by more than an order of magnitude.
    unsafe = true_support_n < params.false_safe_support_n / 1.25
    draw = jax.random.uniform(k_fs, true_support_n.shape) < params.false_safe_rate
    fresh = jnp.where(new_stance, unsafe & draw, state.false_safe)
    support = jnp.where(fresh, params.false_safe_support_n, support)

    # Confidence: weakly informative, matching the measured |resid| vs uncertainty correlation.
    # A false-safe estimate looks CONFIDENT -- that is what makes it dangerous.
    resid_proxy = jnp.abs(support - _shrink(true_support_n, params.support)) / (
        params.support.sigma * 3.0
    )
    # A false-safe is exactly the case where the tree ensemble AGREES and is wrong -- tight
    # spread is why it cleared everest's 2-sigma gate. So it reports near-zero apparent
    # residual, i.e. high confidence. That is what makes it dangerous rather than merely noisy.
    resid_proxy = jnp.where(fresh, 0.0, resid_proxy)
    noise_c = jax.random.uniform(k_conf, true_support_n.shape, minval=-0.3, maxval=0.3)
    rho = params.confidence_resid_corr
    confidence = jnp.clip(1.0 - rho * resid_proxy + noise_c, 0.0, 1.0)

    flip = jax.random.uniform(k_void, true_void.shape)
    void_prob = jnp.where(
        true_void,
        jnp.where(flip < params.void_false_negative_rate, 0.15, 0.85),
        jnp.where(flip < params.void_false_positive_rate, 0.75, 0.10),
    )

    obs = jnp.stack(
        [
            jnp.clip(depth / DEPTH_NORM_M, 0.0, 1.5),
            jnp.clip(support / SUPPORT_NORM_N, 0.0, 1.5),
            confidence,
            void_prob,
        ],
        axis=-1,
    )  # (N_FEET, N_CHANNELS)

    # Dropout: hold the last valid sample and collapse confidence.
    dropped = jax.random.uniform(k_drop, (N_FEET,)) < params.dropout_prob
    held = state.last_obs.at[:, 2].multiply(0.3)
    obs = jnp.where(dropped[:, None], held, obs)

    return obs.reshape(OBS_SIZE), EstimatorState(
        last_obs=obs, false_safe=fresh, noise_scale=scale
    )


def deshrink(value: jax.Array, ch: Channel) -> jax.Array:
    """Invert the measured shrinkage to recover a correctly-SCALED estimate.

    The raw estimate is compressed toward the training mean (depth slope 0.115), so as an
    absolute number it is worse than a constant (R^2 -0.06) even though it still TRACKS truth
    (corr +0.45). A learned policy could rescale such an input itself, but a fixed reward
    formula cannot -- `depth_est + 0.05` would be a near-constant target.

    De-shrinking restores the scale at the cost of amplifying noise by 1/slope. The result is
    noisy but unbiased and correctly scaled, which is what a reward target needs.

    Note this is only legitimate because the shrinkage was MEASURED on held-out data. It uses
    no per-timestep ground truth and is applied to the corrupted estimate, so it does not leak.
    """
    return (value - ch.intercept) / ch.slope


def implied_depth_m(obs: jax.Array, params: EstimatorParams) -> jax.Array:
    """Clearance-relevant snow depth implied by the estimator block, per foot.

    Blends the de-shrunk depth channel with the depth implied by support. In terrain where
    bearing is derived from the sinkage law the two are coupled (corr -0.50), so support
    carries independent information about how deep the snow is.
    """
    depth_raw = obs[..., 0] * DEPTH_NORM_M
    support_raw = obs[..., 1] * SUPPORT_NORM_N
    from_depth = deshrink(depth_raw, params.depth)
    from_support = SUPPORT_TO_DEPTH_SLOPE * deshrink(support_raw, params.support) + SUPPORT_TO_DEPTH_INTERCEPT
    blended = DEPTH_CHANNEL_WEIGHT * from_depth + (1.0 - DEPTH_CHANNEL_WEIGHT) * from_support
    return jnp.clip(blended, 0.0, DEPTH_NORM_M)


batched_estimate = jax.vmap(estimate, in_axes=(0, 0, 0, 0, 0, 0, None))
batched_init = jax.vmap(init_state, in_axes=(0, None))