"""Snow sinkage force (Step 2, part 2).

Follows the resistance law in hackathon-everest's `physics.py::_resistance_n`:

    F = k_eff * d^1.12  +  c * max(v_down, 0)  +  3 * k_eff * max(d - d_support, 0)

with a void band that collapses resistance to 8%. Term by term:

* `k_eff * d^1.12`   progressive compaction -- snow stiffens as it packs.
* `c * v_down`       rate dependence -- snow resists a fast foot harder than a slow one.
                     Only the downward half; snow does not pull a lifting foot back.
* `3 * k_eff * (...)`  bottoming out on the hard substrate once the foot passes through
                     the snow layer. The factor 3 makes that transition sharp.
* void               a hollow under a crust: the foot meets almost nothing, then the
                     substrate term catches it. This is the breakthrough the policy must survive.

Applied to the feet as `xfrc_applied`, not as contact geometry -- see terrain.py for why.
The brief calls for depth-and-velocity dependence and explicitly not granular physics.
"""
from __future__ import annotations

import jax
import jax.numpy as jnp

COMPACTION_EXPONENT = 1.12
SUBSTRATE_STIFFNESS_GAIN = 3.0
VOID_RESISTANCE_FACTOR = 0.08
MAX_FORCE_N = 2_000.0  # numerical guard; ~6x G1 body weight

# Ploughing. Vertical support alone is NOT enough to make deep snow hard: under uniform
# stiffness a foot reaches the same equilibrium sinkage in 0.07 m and 0.19 m of snow, so
# the two regimes would feel identical and the policy would have nothing to learn.
# What actually punishes deep snow is dragging the swing foot through it. Resistance scales
# with submerged frontal area (~ depth x foot width) and snow strength, saturating in speed.
# Sized so a foot swung through deep snow at walking speed meets ~30 N -- roughly 9% of body
# weight, enough to alter gait without immobilising the robot. The learning signal comes from
# clearance: a foot lifted clear of the surface ploughs nothing, so deeper snow demands a
# higher swing to escape the same drag.
PLOUGH_COEFF = 0.25
FOOT_WIDTH_M = 0.09           # G1 foot, approximate
PLOUGH_REFERENCE_SPEED_MPS = 0.4


def resistance_n(
    penetration_m: jax.Array,
    vertical_velocity_mps: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    damping_ns_per_m: jax.Array,
    void: jax.Array,
    include_substrate: bool = True,
) -> jax.Array:
    """Upward resisting force on a foot. Zero when the foot is above the snow surface.

    penetration_m: depth below the snow surface, negative when clear of it.
    vertical_velocity_mps: foot velocity, negative downward.

    include_substrate: whether to add the bottoming-out term. TRUE when deriving bearing
    capacity as a label (everest has no geometry, so the term is the only hard ground there).
    FALSE when computing a force to APPLY in MuJoCo, where the substrate is a real plane with
    a real contact solver -- adding it there double-counts the floor and fights the solver,
    which produced spikes to 1530 N (~4.5x body weight) against the 2000 N clamp.
    """
    d = jnp.maximum(penetration_m, 0.0)
    v_down = jnp.maximum(-vertical_velocity_mps, 0.0)

    force = stiffness_n_per_m * d**COMPACTION_EXPONENT + damping_ns_per_m * v_down
    # Void collapses the snow-column term only; the substrate below still stops the foot.
    force = jnp.where(void, force * VOID_RESISTANCE_FACTOR, force)
    if include_substrate:
        force = force + SUBSTRATE_STIFFNESS_GAIN * stiffness_n_per_m * jnp.maximum(
            d - snow_depth_m, 0.0
        )
    return jnp.clip(jnp.where(d > 0.0, force, 0.0), 0.0, MAX_FORCE_N)


def plough_force_n(
    penetration_m: jax.Array,
    horizontal_speed_mps: jax.Array,
    stiffness_n_per_m: jax.Array,
    void: jax.Array,
) -> jax.Array:
    """Magnitude of the horizontal drag on a foot moving through snow.

    Opposes horizontal motion. Zero above the surface; a void offers almost nothing to plough.
    """
    submerged = jnp.maximum(penetration_m, 0.0)
    strength = jnp.where(void, stiffness_n_per_m * VOID_RESISTANCE_FACTOR, stiffness_n_per_m)
    force = (
        PLOUGH_COEFF * strength * submerged * FOOT_WIDTH_M
        * jnp.tanh(horizontal_speed_mps / PLOUGH_REFERENCE_SPEED_MPS)
    )
    return jnp.clip(force, 0.0, MAX_FORCE_N)


def foot_force_world_n(
    foot_pos_z_m: jax.Array,
    foot_vel_xyz_mps: jax.Array,
    ground_z_m: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    damping_ns_per_m: jax.Array,
    void: jax.Array,
) -> jax.Array:
    """World-frame (fx, fy, fz) per foot, ready for `data.xfrc_applied[body, :3]`.

    The snow surface sits at `ground_z_m + snow_depth_m`; the hard substrate at `ground_z_m`.
    Vertical term supports the foot; horizontal term ploughs against its motion.
    """
    penetration = ground_z_m + snow_depth_m - foot_pos_z_m
    # include_substrate=False: MuJoCo's floor plane already stops the foot.
    fz = resistance_n(penetration, foot_vel_xyz_mps[..., 2], snow_depth_m,
                      stiffness_n_per_m, damping_ns_per_m, void, include_substrate=False)
    v_h = foot_vel_xyz_mps[..., :2]
    speed = jnp.linalg.norm(v_h, axis=-1)
    drag = plough_force_n(penetration, speed, stiffness_n_per_m, void)
    direction = v_h / jnp.maximum(speed, 1e-6)[..., None]
    return jnp.concatenate([-drag[..., None] * direction, fz[..., None]], axis=-1)


def equilibrium_sinkage_m(
    load_n: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    void: jax.Array,
    *,
    iterations: int = 40,
) -> jax.Array:
    """Static sinkage where resistance balances `load_n`, by bisection.

    Diagnostic and test helper -- the training loop integrates the force instead.
    """
    def body(bounds, _):
        lo, hi = bounds
        mid = 0.5 * (lo + hi)
        f = resistance_n(mid, jnp.zeros_like(mid), snow_depth_m, stiffness_n_per_m,
                         jnp.zeros_like(mid), void)
        too_soft = f < load_n
        return (jnp.where(too_soft, mid, lo), jnp.where(too_soft, hi, mid)), None

    lo = jnp.zeros_like(snow_depth_m)
    hi = jnp.full_like(snow_depth_m, 1.0)
    (lo, hi), _ = jax.lax.scan(body, (lo, hi), None, length=iterations)
    return 0.5 * (lo + hi)