#!/usr/bin/env python3
"""Generate realistic synthetic MME crash test fixture data.

Creates a proper ISO 13499 MME directory structure simulating a full
frontal crash test (56 km/h rigid barrier, Hybrid III 50th percentile
male driver). Produces ~30 channels covering all body regions needed
for injury criteria computation.

The signals use physically realistic waveforms:
- Half-sine acceleration pulses with appropriate timing and magnitude
- Neck force/moment with realistic tension-flexion patterns
- Chest deflection with viscous loading character
- Femur compressive force profiles
- Tibia axial and bending loads

All signals include pre-trigger baseline, realistic noise, and sensor
characteristics (sample rate, CFC class).

Usage:
    python generate_mme.py [output_dir]

If no output_dir is specified, generates into tests/fixtures/sample_mme/
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np


# ---------------------------------------------------------------------------
# Parameters
# ---------------------------------------------------------------------------

SAMPLE_RATE = 20000.0  # Hz
DT = 1.0 / SAMPLE_RATE
PRE_TRIGGER_S = 0.010  # 10 ms pre-trigger
EVENT_DURATION_S = 0.200  # 200 ms event
TOTAL_DURATION_S = PRE_TRIGGER_S + EVENT_DURATION_S
NUM_SAMPLES = int(TOTAL_DURATION_S * SAMPLE_RATE)
PRE_TRIGGER_SAMPLES = int(PRE_TRIGGER_S * SAMPLE_RATE)

RNG = np.random.default_rng(2024)

# Time array with t=0 at trigger
TIME = np.arange(NUM_SAMPLES, dtype=np.float64) * DT - PRE_TRIGGER_S


def noise(scale: float = 0.3) -> np.ndarray:
    """Add realistic sensor noise."""
    return RNG.normal(0, scale, NUM_SAMPLES)


def half_sine(
    t_start: float,
    t_duration: float,
    amplitude: float,
    phase_shift: float = 0.0,
) -> np.ndarray:
    """Generate a half-sine pulse."""
    data = np.zeros(NUM_SAMPLES)
    mask = (TIME >= t_start) & (TIME <= t_start + t_duration)
    t_local = TIME[mask] - t_start
    data[mask] = amplitude * np.sin(np.pi * t_local / t_duration + phase_shift)
    return data


def crash_pulse(t_start: float, t_rise: float, t_fall: float, amplitude: float) -> np.ndarray:
    """Generate a more realistic crash pulse with fast rise and slower decay."""
    data = np.zeros(NUM_SAMPLES)
    # Rise phase
    mask_rise = (TIME >= t_start) & (TIME < t_start + t_rise)
    t_local = TIME[mask_rise] - t_start
    data[mask_rise] = amplitude * np.sin(np.pi * t_local / (2 * t_rise))
    # Fall phase
    mask_fall = (TIME >= t_start + t_rise) & (TIME <= t_start + t_rise + t_fall)
    t_local = TIME[mask_fall] - (t_start + t_rise)
    data[mask_fall] = amplitude * np.cos(np.pi * t_local / (2 * t_fall))
    return data


# ---------------------------------------------------------------------------
# Channel definitions
# ---------------------------------------------------------------------------

# Each channel: (code, description, unit, cfc_class, data_generator)


def gen_channels() -> list[tuple[str, str, str, int, np.ndarray]]:
    channels = []

    # === HEAD ===
    # Head X acceleration: main frontal pulse, ~45g peak
    head_ax = crash_pulse(0.0, 0.020, 0.060, 45.0) + noise(0.8)
    channels.append(("11HEAD0000ACXA", "Head CG X Acceleration", "g", 1000, head_ax))

    # Head Y acceleration: small lateral component
    head_ay = half_sine(0.005, 0.070, 8.0) * np.cos(5 * np.pi * (TIME - 0.005) / 0.070) + noise(0.4)
    head_ay[TIME < 0.005] = noise(0.4)[TIME < 0.005]
    channels.append(("11HEAD0000ACYA", "Head CG Y Acceleration", "g", 1000, head_ay))

    # Head Z acceleration: downward component
    head_az = crash_pulse(0.003, 0.018, 0.055, 22.0) + noise(0.5)
    channels.append(("11HEAD0000ACZA", "Head CG Z Acceleration", "g", 1000, head_az))

    # Head angular velocity (for future BrIC)
    head_avx = half_sine(0.005, 0.060, 15.0) + noise(0.3)
    channels.append(("11HEAD0000AVXA", "Head Angular Velocity X", "rad/s", 600, head_avx))
    head_avy = half_sine(0.008, 0.055, 25.0) + noise(0.3)
    channels.append(("11HEAD0000AVYA", "Head Angular Velocity Y", "rad/s", 600, head_avy))
    head_avz = half_sine(0.006, 0.050, 10.0) + noise(0.2)
    channels.append(("11HEAD0000AVZA", "Head Angular Velocity Z", "rad/s", 600, head_avz))

    # === NECK (upper) ===
    # Neck Fz: tension pulse ~3200 N
    neck_fz = crash_pulse(0.005, 0.015, 0.045, 3200.0) + noise(15.0)
    channels.append(("11NECKUP00FOZA", "Upper Neck Axial Force", "N", 600, neck_fz))

    # Neck Fx: shear force ~800 N
    neck_fx = crash_pulse(0.008, 0.012, 0.040, 800.0) + noise(8.0)
    channels.append(("11NECKUP00FOXA", "Upper Neck Shear Force X", "N", 600, neck_fx))

    # Neck My: flexion moment ~85 Nm
    neck_my = crash_pulse(0.010, 0.015, 0.050, 85.0) + noise(1.5)
    channels.append(("11NECKUP00MOYA", "Upper Neck Moment Y (Flexion)", "N·m", 600, neck_my))

    # Neck Mx: lateral moment ~25 Nm
    neck_mx = half_sine(0.012, 0.045, 25.0) + noise(0.8)
    channels.append(("11NECKUP00MOXA", "Upper Neck Moment X", "N·m", 600, neck_mx))

    # === CHEST ===
    # Chest X acceleration
    chest_ax = crash_pulse(0.008, 0.020, 0.055, 38.0) + noise(0.6)
    channels.append(("11CHST0000ACXA", "Chest T12 X Acceleration", "g", 180, chest_ax))

    # Chest Y acceleration
    chest_ay = half_sine(0.010, 0.060, 6.0) + noise(0.3)
    channels.append(("11CHST0000ACYA", "Chest T12 Y Acceleration", "g", 180, chest_ay))

    # Chest Z acceleration
    chest_az = crash_pulse(0.009, 0.018, 0.050, 18.0) + noise(0.4)
    channels.append(("11CHST0000ACZA", "Chest T12 Z Acceleration", "g", 180, chest_az))

    # Chest deflection: ramps up to ~36mm, then recovers
    chest_dc = np.zeros(NUM_SAMPLES)
    mask_load = (TIME >= 0.010) & (TIME < 0.055)
    t_load = TIME[mask_load] - 0.010
    chest_dc[mask_load] = 36.0 * np.sin(np.pi * t_load / 0.090)
    mask_unload = (TIME >= 0.055) & (TIME < 0.130)
    t_unload = TIME[mask_unload] - 0.055
    chest_dc[mask_unload] = 36.0 * np.sin(np.pi * 0.045 / 0.090) * np.exp(-t_unload / 0.030)
    chest_dc += noise(0.2)
    channels.append(("11CHST0000DCXA", "Chest Deflection", "mm", 180, chest_dc))

    # === PELVIS ===
    pelv_ax = crash_pulse(0.012, 0.022, 0.060, 35.0) + noise(0.5)
    channels.append(("11PELV0000ACXA", "Pelvis X Acceleration", "g", 180, pelv_ax))
    pelv_ay = half_sine(0.015, 0.055, 5.0) + noise(0.3)
    channels.append(("11PELV0000ACYA", "Pelvis Y Acceleration", "g", 180, pelv_ay))
    pelv_az = crash_pulse(0.014, 0.018, 0.050, 15.0) + noise(0.4)
    channels.append(("11PELV0000ACZA", "Pelvis Z Acceleration", "g", 180, pelv_az))

    # === FEMUR ===
    # Left femur: compressive pulse ~4800 N (negative = compression in SAE)
    femur_l = -crash_pulse(0.020, 0.015, 0.045, 4800.0) + noise(20.0)
    channels.append(("11FEMRLE00FOZA", "Left Femur Axial Force", "N", 600, femur_l))

    # Right femur: slightly lower
    femur_r = -crash_pulse(0.022, 0.014, 0.042, 4200.0) + noise(18.0)
    channels.append(("11FEMRRI00FOZA", "Right Femur Axial Force", "N", 600, femur_r))

    # === TIBIA (left upper) ===
    tibia_fz = -crash_pulse(0.025, 0.012, 0.040, 5500.0) + noise(25.0)
    channels.append(("11TIBILEUOFOZA", "Left Upper Tibia Axial Force", "N", 600, tibia_fz))

    tibia_mx = half_sine(0.028, 0.040, 80.0) + noise(2.0)
    channels.append(("11TIBILEUOMOXA", "Left Upper Tibia Moment X", "N·m", 600, tibia_mx))

    tibia_my = half_sine(0.027, 0.038, 120.0) + noise(2.5)
    channels.append(("11TIBILEUOMOYA", "Left Upper Tibia Moment Y", "N·m", 600, tibia_my))

    # === SEAT BELT ===
    belt_fo = crash_pulse(0.005, 0.025, 0.070, 6500.0) + noise(30.0)
    channels.append(("11BELT0000FOXA", "Lap Belt Force", "N", 60, belt_fo))

    # Shoulder belt force
    sbelt_fo = crash_pulse(0.003, 0.020, 0.065, 5800.0) + noise(25.0)
    channels.append(("11BELTUPOOFOZA", "Shoulder Belt Force", "N", 60, sbelt_fo))

    # === VEHICLE / STRUCTURAL ===
    # B-pillar acceleration
    bpil_ax = crash_pulse(0.001, 0.015, 0.040, 55.0) + noise(1.0)
    channels.append(("10BPILLE00ACXA", "Left B-Pillar X Acceleration", "g", 60, bpil_ax))

    # Steering column displacement
    stcl_ds = np.zeros(NUM_SAMPLES)
    mask = (TIME >= 0.015) & (TIME < 0.100)
    t_local = TIME[mask] - 0.015
    stcl_ds[mask] = 45.0 * (1 - np.exp(-t_local / 0.025))
    stcl_ds[TIME >= 0.100] = stcl_ds[mask][-1] if np.any(mask) else 0
    stcl_ds += noise(0.3)
    channels.append(("10STCL0000DSXA", "Steering Column Displacement", "mm", 60, stcl_ds))

    return channels


# ---------------------------------------------------------------------------
# MME writer
# ---------------------------------------------------------------------------


def write_mme(output_dir: Path) -> None:
    """Write a complete MME directory structure."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # --- MME.ini ---
    ini_content = """\
[Test]
test_number = IMPAKT_SYNTH_001
test_date = 2024-06-15
test_type = Full Frontal Rigid Barrier
test_speed = 56.3
test_facility = Impakt Synthetic Data Generator
test_standard = FMVSS 208
description = Synthetic frontal crash test for Impakt development and testing

[Vehicle]
vehicle_make = Synthetic
vehicle_model = TestCar
vehicle_year = 2024
vehicle_vin = IMPAKT0000000001
vehicle_mass = 1523.0
vehicle_type = Passenger Car

[Barrier]
barrier_type = Rigid
impact_angle = 0
overlap_percent = 100
impact_side = Front

[Dummy]
dummy_type = Hybrid III 50th Percentile Male
dummy_serial = H3-50M-SYNTH-001
dummy_position = Driver
dummy_mass = 78.0
restraint_type = 3-point belt + frontal airbag
seat_position = Mid-track, mid-height

[DataAcquisition]
sample_rate = 20000
pre_trigger_ms = 10
total_duration_ms = 210
num_channels = 28
"""

    (output_dir / "MME.ini").write_text(ini_content, encoding="utf-8")

    # --- Channel files ---
    ch_dir = output_dir / "channels"
    ch_dir.mkdir(exist_ok=True)

    channels = gen_channels()

    for code, description, unit, cfc, data in channels:
        # Write .chn header
        chn_content = f"""\
[Channel]
channel_code = {code}
description = {description}
unit = {unit}
sample_rate = {SAMPLE_RATE:.0f}
num_samples = {NUM_SAMPLES}
dt = {DT:.10f}
pre_trigger = {PRE_TRIGGER_SAMPLES}
cfc = {cfc}
data_format = ascii
"""
        (ch_dir / f"{code}.chn").write_text(chn_content, encoding="utf-8")

        # Write .dat data (ASCII, one value per line)
        np.savetxt(str(ch_dir / f"{code}.dat"), data, fmt="%.8f")

    print(f"Generated MME fixture: {output_dir}")
    print(f"  {len(channels)} channels, {NUM_SAMPLES} samples each")
    print(f"  Sample rate: {SAMPLE_RATE:.0f} Hz")
    print(
        f"  Duration: {PRE_TRIGGER_S * 1000:.0f} ms pre-trigger + {EVENT_DURATION_S * 1000:.0f} ms event"
    )
    print(
        f"  Total size: ~{sum(len(data) for _, _, _, _, data in channels) * 12 / 1024 / 1024:.1f} MB"
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) > 1:
        out = Path(sys.argv[1])
    else:
        out = Path(__file__).parent / "sample_mme"

    write_mme(out)