"""Marsh 3-compartment pharmacokinetic model for propofol.

The Marsh model describes propofol distribution using a three-compartment
mammillary model with first-order inter-compartmental transfer and elimination.

References:
    Marsh, B., White, M., Morton, N., & Kenny, G. N. (1991).
    Pharmacokinetic model driven infusion of propofol in children.
    British Journal of Anaesthesia, 67(1), 41-48.

    Absalom, A. R., Mani, V., De Smet, T., & Struys, M. M. (2009).
    Pharmacokinetic models for propofol: defining and illuminating the
    devil in the detail. British Journal of Anaesthesia, 103(1), 26-37.
"""

from __future__ import annotations

import numpy as np


# SIMPLIFICATION: Using fixed Marsh parameters instead of covariate-adjusted
# Schnider model. Marsh model uses weight-proportional volumes only.
# Clinical TCI systems offer both Marsh and Schnider; we use Marsh for
# its simpler parameterization as a benchmark standard.

PARAMETER_RANGES = {
    "V1_per_kg": {
        "min": 0.15, "max": 0.35, "unit": "L/kg",
        "source": "Marsh et al., 1991", "default": 0.228,
    },
    "k10": {
        "min": 0.05, "max": 0.20, "unit": "1/min",
        "source": "Marsh et al., 1991", "default": 0.119,
    },
    "k12": {
        "min": 0.05, "max": 0.20, "unit": "1/min",
        "source": "Marsh et al., 1991", "default": 0.112,
    },
    "k21": {
        "min": 0.02, "max": 0.10, "unit": "1/min",
        "source": "Marsh et al., 1991", "default": 0.055,
    },
    "k13": {
        "min": 0.02, "max": 0.08, "unit": "1/min",
        "source": "Marsh et al., 1991", "default": 0.042,
    },
    "k31": {
        "min": 0.01, "max": 0.06, "unit": "1/min",
        "source": "Marsh et al., 1991", "default": 0.033,
    },
    "ke0": {
        "min": 0.15, "max": 0.50, "unit": "1/min",
        "source": "Absalom et al., 2009", "default": 0.26,
    },
}


class MarshPKModel:
    """Three-compartment PK model for propofol (Marsh parameterization).

    State variables (amounts in mg):
        A1: central compartment (plasma)
        A2: rapid peripheral compartment
        A3: slow peripheral compartment
        Ce: effect-site concentration (mcg/mL)

    Equations:
        dA1/dt = -(k10 + k12 + k13)*A1 + k21*A2 + k31*A3 + R
        dA2/dt = k12*A1 - k21*A2
        dA3/dt = k13*A1 - k31*A3
        dCe/dt = ke0*(Cp - Ce)

    where R = infusion rate (mg/min), Cp = A1/V1 (mcg/mL).
    """

    def __init__(
        self,
        weight_kg: float = 70.0,
        k10: float = 0.119,
        k12: float = 0.112,
        k21: float = 0.055,
        k13: float = 0.042,
        k31: float = 0.033,
        ke0: float = 0.26,
        V1_per_kg: float = 0.228,
        dt: float = 0.5,
    ) -> None:
        self.weight_kg = weight_kg
        self.V1 = V1_per_kg * weight_kg  # L
        self.k10 = k10
        self.k12 = k12
        self.k21 = k21
        self.k13 = k13
        self.k31 = k31
        self.ke0 = ke0
        self.dt = dt

        self.A1 = 0.0  # mg in central compartment
        self.A2 = 0.0  # mg in rapid peripheral
        self.A3 = 0.0  # mg in slow peripheral
        self.Ce = 0.0  # effect-site concentration (mcg/mL)
        self.total_drug = 0.0  # cumulative drug administered (mg)

    def reset(self) -> None:
        """Reset all compartments to zero."""
        self.A1 = 0.0
        self.A2 = 0.0
        self.A3 = 0.0
        self.Ce = 0.0
        self.total_drug = 0.0

    @property
    def plasma_concentration(self) -> float:
        """Plasma (central compartment) concentration in mcg/mL."""
        return self.A1 / self.V1 if self.V1 > 0 else 0.0

    def step(self, infusion_rate_mcg_kg_min: float = 0.0) -> dict:
        """Advance one timestep with given infusion rate.

        Args:
            infusion_rate_mcg_kg_min: Propofol infusion rate in mcg/kg/min.

        Returns:
            Dict with plasma_concentration, effect_site_concentration.
        """
        # Convert mcg/kg/min to mg/min
        R = infusion_rate_mcg_kg_min * self.weight_kg / 1000.0
        self.total_drug += R * self.dt

        # Euler integration of 3-compartment model
        dA1 = (-(self.k10 + self.k12 + self.k13) * self.A1
               + self.k21 * self.A2 + self.k31 * self.A3 + R) * self.dt
        dA2 = (self.k12 * self.A1 - self.k21 * self.A2) * self.dt
        dA3 = (self.k13 * self.A1 - self.k31 * self.A3) * self.dt

        self.A1 = max(0.0, self.A1 + dA1)
        self.A2 = max(0.0, self.A2 + dA2)
        self.A3 = max(0.0, self.A3 + dA3)

        # Effect-site compartment
        Cp = self.plasma_concentration
        dCe = self.ke0 * (Cp - self.Ce) * self.dt
        self.Ce = max(0.0, self.Ce + dCe)

        return {
            "plasma_concentration": Cp,
            "effect_site_concentration": self.Ce,
        }

    def apply_variability(self, rng: np.random.Generator, scale: float = 0.2) -> None:
        """Apply inter-individual variability to PK parameters.

        Args:
            rng: Random number generator.
            scale: Coefficient of variation (0.2 = 20%).
        """
        for attr in ("k10", "k12", "k21", "k13", "k31", "ke0"):
            base = getattr(self, attr)
            multiplier = 1.0 + rng.uniform(-scale, scale)
            setattr(self, attr, max(0.001, base * multiplier))
        self.V1 = max(1.0, self.V1 * (1.0 + rng.uniform(-scale, scale)))
