"""Chest injury criteria: deflection and viscous criterion.

Chest Deflection: Peak sternal displacement (mm).
Viscous Criterion (VC): V(t) * C(t) where V = deflection velocity, C = compression ratio.
"""

from __future__ import annotations

from typing import Any

import numpy as np

from impakt.channel.model import Channel, DummyInfo
from impakt.criteria.base import CriterionResult


# Initial chest depth by dummy type (mm)
CHEST_DEPTH: dict[str, float] = {
    "H3-50M": 229.0,
    "H3-05F": 187.0,
    "H3-95M": 254.0,
    "THOR-50M": 229.0,
    "ES-2re": 140.0,  # Approximate rib depth
    "WSID-50": 140.0,
}

DEFAULT_CHEST_DEPTH = 229.0  # mm, Hybrid III 50th Male


def chest_deflection(
    channel: Channel | None = None,
    channels: dict[str, Channel] | None = None,
    dummy: DummyInfo | None = None,
) -> CriterionResult:
    """Compute peak chest deflection.

    For Hybrid III: single-point sternal deflection.
    For THOR: may involve multiple IR-TRACC channels (returns max across all).

    Args:
        channel: Chest deflection Channel (mm).
        channels: Alternative dict (uses first deflection channel found).
        dummy: Dummy info (for future multi-point THOR handling).

    Returns:
        CriterionResult with peak deflection value.
    """
    if channels is not None and channel is None:
        # Use the first channel or the one with the highest peak
        ch_list = list(channels.values())
        channel = max(ch_list, key=lambda c: np.max(np.abs(c.data)))

    if channel is None:
        raise ValueError("A chest deflection channel is required")

    deflection = channel.data.copy()

    # Convert units if needed
    if channel.unit in ("m",):
        deflection = deflection * 1000.0  # m -> mm

    peak_value = float(np.max(np.abs(deflection)))
    peak_idx = int(np.argmax(np.abs(deflection)))
    peak_time = float(channel.time[peak_idx])

    return CriterionResult(
        criterion="Chest Deflection",
        value=peak_value,
        unit="mm",
        time_of_peak=peak_time,
        body_region="Chest",
        details={
            "signed_peak": float(deflection[peak_idx]),
            "input_channel": channel.name,
        },
    )


def viscous_criterion(
    channel: Channel | None = None,
    channels: dict[str, Channel] | None = None,
    dummy: DummyInfo | None = None,
    chest_depth_mm: float | None = None,
) -> CriterionResult:
    """Compute the Viscous Criterion (VC).

    VC = V(t) * C(t)
    where:
        V(t) = d[D(t)]/dt  (velocity of chest deflection, m/s)
        C(t) = D(t) / D0    (instantaneous compression ratio)
        D0 = initial chest depth (mm)

    Args:
        channel: Chest deflection time-history (mm).
        channels: Alternative dict.
        dummy: Dummy info for chest depth lookup.
        chest_depth_mm: Override initial chest depth.

    Returns:
        CriterionResult with max VC value.
    """
    if channels is not None and channel is None:
        ch_list = list(channels.values())
        channel = max(ch_list, key=lambda c: np.max(np.abs(c.data)))

    if channel is None:
        raise ValueError("A chest deflection channel is required")

    # Get chest depth
    d0 = chest_depth_mm
    if d0 is None and dummy is not None:
        for key, depth in CHEST_DEPTH.items():
            if key.lower() in dummy.dummy_type.lower():
                d0 = depth
                break
    if d0 is None:
        d0 = DEFAULT_CHEST_DEPTH

    deflection_mm = channel.data.copy()
    if channel.unit in ("m",):
        deflection_mm = deflection_mm * 1000.0

    dt = channel.dt
    if dt <= 0:
        raise ValueError("Channel has no timing information (dt=0)")

    # Velocity in m/s (deflection is in mm, so divide by 1000)
    velocity = np.gradient(deflection_mm / 1000.0, dt)

    # Compression ratio (dimensionless)
    compression = deflection_mm / d0

    # VC in m/s
    vc = velocity * compression

    peak_vc = float(np.max(np.abs(vc)))
    peak_idx = int(np.argmax(np.abs(vc)))
    peak_time = float(channel.time[peak_idx])

    return CriterionResult(
        criterion="Viscous Criterion",
        value=peak_vc,
        unit="m/s",
        time_of_peak=peak_time,
        body_region="Chest",
        details={
            "chest_depth_mm": d0,
            "max_velocity_m_s": float(np.max(np.abs(velocity))),
            "max_compression": float(np.max(np.abs(compression))),
            "signed_peak_vc": float(vc[peak_idx]),
            "input_channel": channel.name,
        },
    )
