"""Bilateral support reserve (ported from hackathon-everest's `control.py`).

everest's central quantity: how much load each foot has left before its foothold fails.

    reserve = lower_confidence_bearing - current_load

Worth porting specifically because it is the part of everest that demonstrably *worked*.
Decomposing everest's benchmark showed all 7 unsafe transfers were stance-capacity failures
-- the bilateral gate caught them, and the belief map contributed nothing measurable to that
number. So this is the load-bearing idea, and the map is the unproven one.

The translation to RL is not one-to-one. everest's controller emits a discrete decision
(COMMIT / HOLD_DOUBLE_SUPPORT / REPLANT) and can refuse to move. A locomotion policy is
handed a velocity command and must produce joint targets 50 times a second -- it cannot
refuse. So the gate becomes two continuous things instead:

* an **observation**: the policy sees each foot's reserve and can act on it;
* a **penalty**: loading a foot beyond its conservatively-estimated support is punished.

Both read the ESTIMATE, never the truth, so the policy cannot succeed by ignoring the sensor.
"""
from __future__ import annotations

from typing import NamedTuple

import jax
import jax.numpy as jnp

G1_WEIGHT_N = 343.0
SAFETY_BUFFER_N = 10.0
RESERVE_NORM_N = 400.0     # observation scaling
UNCERTAINTY_SIGMAS = 2.0   # everest uses mean - 2*sigma


class SupportState(NamedTuple):
    """Per-foot support accounting. All values derived from the estimate."""
    reserve_n: jax.Array        # (2,) lower-confidence support minus current load
    load_n: jax.Array           # (2,) vertical load the snow is currently carrying
    total_margin_n: jax.Array   # () sum of both reserves


def evaluate(
    support_est_n: jax.Array,       # (2,) estimated bearing capacity per foot
    support_std_n: jax.Array,       # (2,) estimator uncertainty
    load_n: jax.Array,              # (2,) vertical snow force per foot
) -> SupportState:
    """Reserve on each foot, using the conservative lower bound rather than the mean."""
    lower = support_est_n - UNCERTAINTY_SIGMAS * support_std_n
    reserve = lower - load_n
    return SupportState(
        reserve_n=reserve,
        load_n=load_n,
        total_margin_n=jnp.sum(reserve),
    )


def observation(state: SupportState) -> jax.Array:
    """(5,) normalised: per-foot reserve, per-foot load, and total margin."""
    return jnp.concatenate([
        state.reserve_n / RESERVE_NORM_N,
        state.load_n / RESERVE_NORM_N,
        jnp.atleast_1d(state.total_margin_n / RESERVE_NORM_N),
    ])


def overload_cost(state: SupportState, contact: jax.Array) -> jax.Array:
    """Penalty for standing on a foothold the estimate says cannot hold you.

    everest's gate as a continuous cost: a negative reserve on a loaded foot means the
    conservative estimate of what this foothold carries is already exceeded. Counted only on
    feet actually bearing load -- a swinging foot with no support underneath is not a fault.
    """
    deficit = jnp.maximum(-state.reserve_n, 0.0) / RESERVE_NORM_N
    return jnp.sum(deficit * contact.astype(float))


def transfer_is_safe(state: SupportState, swing_index: int) -> jax.Array:
    """everest's `transfer_is_safe`, kept for evaluation and diagnostics.

    True when the stance foot's reserve covers the load about to move onto it. Not used as a
    gate in training -- the policy cannot halt -- but reported so a rollout can be scored
    against everest's own criterion.
    """
    stance = 1 - swing_index
    return state.reserve_n[stance] >= state.load_n[swing_index] + SAFETY_BUFFER_N