"""Neck Injury Criterion (Nij) calculation.

Nij = Fz/Fzc + My/Myc

Four loading modes:
    NTE: Tension + Extension
    NTF: Tension + Flexion
    NCE: Compression + Extension
    NCF: Compression + Flexion

The reported Nij is the maximum across all four modes and all time steps.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np

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


@dataclass(frozen=True)
class NijIntercepts:
    """Critical intercept values for Nij calculation.

    These are dummy-specific and define the denominator terms.
    """

    fzc_tension: float  # N (positive Fz = tension)
    fzc_compression: float  # N (negative Fz = compression)
    myc_flexion: float  # N·m (positive My = flexion)
    myc_extension: float  # N·m (negative My = extension)


# Standard intercepts by dummy type
NIJ_INTERCEPTS: dict[str, NijIntercepts] = {
    "H3-50M": NijIntercepts(
        fzc_tension=6806.0,
        fzc_compression=6160.0,
        myc_flexion=310.0,
        myc_extension=135.0,
    ),
    "Hybrid III 50th Percentile Male": NijIntercepts(
        fzc_tension=6806.0,
        fzc_compression=6160.0,
        myc_flexion=310.0,
        myc_extension=135.0,
    ),
    "H3-05F": NijIntercepts(
        fzc_tension=4287.0,
        fzc_compression=3880.0,
        myc_flexion=155.0,
        myc_extension=67.0,
    ),
    "Hybrid III 5th Percentile Female": NijIntercepts(
        fzc_tension=4287.0,
        fzc_compression=3880.0,
        myc_flexion=155.0,
        myc_extension=67.0,
    ),
    "H3-95M": NijIntercepts(
        fzc_tension=8216.0,
        fzc_compression=7440.0,
        myc_flexion=415.0,
        myc_extension=179.0,
    ),
    "H3-6YO": NijIntercepts(
        fzc_tension=2800.0,
        fzc_compression=2800.0,
        myc_flexion=93.0,
        myc_extension=37.0,
    ),
    "H3-3YO": NijIntercepts(
        fzc_tension=2120.0,
        fzc_compression=2120.0,
        myc_flexion=68.0,
        myc_extension=27.0,
    ),
    "H3-10YO": NijIntercepts(
        fzc_tension=3500.0,
        fzc_compression=3500.0,
        myc_flexion=125.0,
        myc_extension=50.0,
    ),
}

# Default to Hybrid III 50th Male
DEFAULT_INTERCEPTS = NIJ_INTERCEPTS["H3-50M"]


def _get_intercepts(dummy: DummyInfo | None) -> NijIntercepts:
    """Look up Nij intercepts for a given dummy."""
    if dummy is None:
        return DEFAULT_INTERCEPTS

    # Try direct match
    for key, intercepts in NIJ_INTERCEPTS.items():
        if key.lower() in dummy.dummy_type.lower():
            return intercepts

    # Heuristic matching
    dtype = dummy.dummy_type.upper()
    if "50" in dtype and ("MALE" in dtype or "50M" in dtype or "H3" in dtype):
        return NIJ_INTERCEPTS["H3-50M"]
    elif "5" in dtype and ("FEM" in dtype or "05F" in dtype or "5F" in dtype):
        return NIJ_INTERCEPTS["H3-05F"]
    elif "95" in dtype:
        return NIJ_INTERCEPTS["H3-95M"]

    return DEFAULT_INTERCEPTS


def nij(
    fz_channel: Channel | None = None,
    my_channel: Channel | None = None,
    channels: dict[str, Channel] | None = None,
    dummy: DummyInfo | None = None,
    intercepts: NijIntercepts | None = None,
) -> CriterionResult:
    """Compute the Neck Injury Criterion (Nij).

    Args:
        fz_channel: Upper neck axial force channel (N). Positive = tension.
        my_channel: Upper neck sagittal moment channel (N·m). Positive = flexion.
        channels: Alternative: dict with 'fz' and 'my' keys.
        dummy: Dummy info for intercept lookup.
        intercepts: Override intercept values directly.

    Returns:
        CriterionResult with the maximum Nij value.
    """
    # Resolve channels
    if channels is not None:
        fz_channel = channels.get("fz") or channels.get("Fz") or channels.get("FZ")
        my_channel = channels.get("my") or channels.get("My") or channels.get("MY")

    if fz_channel is None or my_channel is None:
        raise ValueError("Both Fz (axial force) and My (moment) channels are required")

    # Get intercepts
    ints = intercepts or _get_intercepts(dummy)

    fz = fz_channel.data.copy()
    my = my_channel.data.copy()

    # Convert units if needed (expect N and N·m)
    if fz_channel.unit in ("kN",):
        fz = fz * 1000.0
    if my_channel.unit in ("kN·m", "kNm"):
        my = my * 1000.0

    n = min(len(fz), len(my))
    fz = fz[:n]
    my = my[:n]

    # Compute all four Nij modes at each time step
    # NTE: tension (Fz > 0) + extension (My < 0)
    # NTF: tension (Fz > 0) + flexion (My > 0)
    # NCE: compression (Fz < 0) + extension (My < 0)
    # NCF: compression (Fz < 0) + flexion (My > 0)

    fz_tension = np.maximum(fz, 0.0) / ints.fzc_tension
    fz_compression = np.maximum(-fz, 0.0) / ints.fzc_compression
    my_flexion = np.maximum(my, 0.0) / ints.myc_flexion
    my_extension = np.maximum(-my, 0.0) / ints.myc_extension

    nte = fz_tension + my_extension
    ntf = fz_tension + my_flexion
    nce = fz_compression + my_extension
    ncf = fz_compression + my_flexion

    # Find maximum across all modes and time
    modes = {"NTE": nte, "NTF": ntf, "NCE": nce, "NCF": ncf}
    max_nij = 0.0
    max_mode = ""
    max_idx = 0

    for mode_name, mode_values in modes.items():
        idx = int(np.argmax(mode_values))
        val = float(mode_values[idx])
        if val > max_nij:
            max_nij = val
            max_mode = mode_name
            max_idx = idx

    time_arr = fz_channel.time[:n]
    peak_time = float(time_arr[max_idx])

    return CriterionResult(
        criterion="Nij",
        value=max_nij,
        unit="",
        time_of_peak=peak_time,
        body_region="Neck",
        details={
            "mode": max_mode,
            "NTE_max": float(np.max(nte)),
            "NTF_max": float(np.max(ntf)),
            "NCE_max": float(np.max(nce)),
            "NCF_max": float(np.max(ncf)),
            "intercepts": {
                "fzc_tension": ints.fzc_tension,
                "fzc_compression": ints.fzc_compression,
                "myc_flexion": ints.myc_flexion,
                "myc_extension": ints.myc_extension,
            },
            "fz_peak_N": float(np.max(np.abs(fz))),
            "my_peak_Nm": float(np.max(np.abs(my))),
        },
    )
