"""Virtual patient generation for anesthesia simulation.

Generates cohorts of virtual patients with demographically varied
pharmacokinetic parameters for robust environment evaluation.

References:
    Eleveld, D. J., Proost, J. H., Cortinez, L. I., et al. (2014).
    A general purpose pharmacokinetic model for propofol. Anesthesia
    and Analgesia, 118(6), 1221-1237.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass
class VirtualPatient:
    """Virtual patient with demographic and PK parameters."""

    patient_id: int
    age: float  # years
    weight_kg: float
    height_cm: float
    lean_body_mass_kg: float
    sex: str  # "M" or "F"
    group: str  # "young", "middle", "elderly"

    @property
    def bmi(self) -> float:
        """Body mass index in kg/m^2."""
        h_m = self.height_cm / 100.0
        return self.weight_kg / (h_m * h_m) if h_m > 0 else 0.0


def _lean_body_mass(weight_kg: float, height_cm: float, sex: str) -> float:
    """Estimate lean body mass using the James formula.

    Reference: James, W.P.T. (1976). Research on obesity. HMSO, London.
    """
    if sex == "M":
        return 1.1 * weight_kg - 128 * (weight_kg / height_cm) ** 2
    else:
        return 1.07 * weight_kg - 148 * (weight_kg / height_cm) ** 2


def generate_patient_cohort(
    n_patients: int = 30,
    seed: int = 42,
) -> list[VirtualPatient]:
    """Generate a cohort of virtual patients across 3 age groups.

    Creates n_patients spread evenly across young (25-40), middle (40-60),
    and elderly (60-80) groups with realistic demographic distributions.

    Args:
        n_patients: Total number of patients to generate.
        seed: Random seed for reproducibility.

    Returns:
        List of VirtualPatient instances.
    """
    rng = np.random.default_rng(seed)
    per_group = n_patients // 3
    remainder = n_patients - 3 * per_group

    groups = {
        "young": {"age_range": (25, 40), "count": per_group},
        "middle": {"age_range": (40, 60), "count": per_group},
        "elderly": {"age_range": (60, 80), "count": per_group + remainder},
    }

    patients: list[VirtualPatient] = []
    pid = 0

    for group_name, params in groups.items():
        for _ in range(params["count"]):
            sex = rng.choice(["M", "F"])
            age = rng.uniform(*params["age_range"])

            if sex == "M":
                weight = rng.normal(80.0, 12.0)
                height = rng.normal(175.0, 7.0)
            else:
                weight = rng.normal(65.0, 10.0)
                height = rng.normal(163.0, 6.0)

            weight = max(45.0, min(130.0, weight))
            height = max(140.0, min(200.0, height))
            lbm = max(30.0, _lean_body_mass(weight, height, sex))

            patients.append(VirtualPatient(
                patient_id=pid,
                age=round(age, 1),
                weight_kg=round(weight, 1),
                height_cm=round(height, 1),
                lean_body_mass_kg=round(lbm, 1),
                sex=sex,
                group=group_name,
            ))
            pid += 1

    return patients
