"""Two-compartment urea kinetics model for hemodialysis.

Based on the Gotch-Sargent model and the double-pool kinetics described
in Smye et al. (1997). Urea distributes between intracellular (ICF) and
extracellular (ECF) compartments separated by a cell membrane.

References:
    Gotch FA, Sargent JA. Kidney Int. 1985;28(3):526-534.
    Smye SW, et al. Kidney Int. 1997;51(5):1541-1547.
    Daugirdas JT. J Am Soc Nephrol. 1993;4(5):1205-1213.
"""

import numpy as np
from scipy.integrate import solve_ivp


class UreaKineticsModel:
    """Two-compartment urea kinetics model.

    SIMPLIFICATION: Uses lumped two-compartment model (ECF + ICF) instead of
    the full regional blood flow model. The Gotch-Sargent NCDS analysis used
    a single-pool model; the two-pool extension captures urea rebound but
    omits regional perfusion heterogeneity. Acceptable for benchmark simulation.
    """

    def __init__(
        self,
        ecf_volume: float,
        icf_volume: float,
        koa: float = 800.0,
        intercompartment_kc: float = 800.0,
        urea_generation_rate: float = 8.0,
    ):
        """Initialize urea kinetics model.

        Args:
            ecf_volume: Extracellular fluid volume in mL.
            icf_volume: Intracellular fluid volume in mL.
            koa: Dialyzer mass transfer-area coefficient in mL/min.
            intercompartment_kc: Inter-compartment clearance in mL/min.
            urea_generation_rate: Endogenous urea generation in mg/min.
        """
        self.v_ecf = ecf_volume
        self.v_icf = icf_volume
        self.koa = koa
        self.kc = intercompartment_kc
        self.g_urea = urea_generation_rate

        # State: [C_ecf (mg/dL), C_icf (mg/dL)]
        self.state = np.array([70.0, 70.0])  # pre-dialysis BUN

    def reset(self, c_ecf: float = 70.0, c_icf: float = 70.0) -> np.ndarray:
        """Reset model to initial concentrations."""
        self.state = np.array([c_ecf, c_icf])
        return self.state.copy()

    def dialyzer_clearance(
        self, blood_flow_rate: float, dialysate_flow_rate: float
    ) -> float:
        """Compute effective dialyzer clearance using the KoA equation.

        K = Qb * Qd * (1 - exp(-KoA*(1/Qb - 1/Qd))) / (Qd - Qb*exp(-KoA*(1/Qb - 1/Qd)))

        SIMPLIFICATION: Uses steady-state mass transfer equation. Clinical
        dialyzer clearance varies with blood and dialysate flow distribution
        within the fiber bundle. Acceptable for benchmark simulation.

        Args:
            blood_flow_rate: Blood flow rate Qb in mL/min.
            dialysate_flow_rate: Dialysate flow rate Qd in mL/min.

        Returns:
            Effective clearance K in mL/min.
        """
        qb = max(blood_flow_rate, 1.0)
        qd = max(dialysate_flow_rate, 1.0)

        if abs(qb - qd) < 1.0:
            # Avoid division by zero when Qb ~ Qd
            return qb * self.koa / (qb + self.koa)

        exponent = -self.koa * (1.0 / qb - 1.0 / qd)
        exp_term = np.exp(np.clip(exponent, -20, 20))
        k = qb * qd * (1.0 - exp_term) / (qd - qb * exp_term)
        return float(np.clip(k, 0, min(qb, qd)))

    def _ode_rhs(self, _t: float, y: np.ndarray, k_dialyzer: float) -> np.ndarray:
        """Right-hand side of the two-compartment ODE system.

        dC_ecf/dt = (G + Kc*(C_icf - C_ecf) - K*C_ecf) / V_ecf
        dC_icf/dt = -Kc*(C_icf - C_ecf) / V_icf

        Where:
            G = urea generation rate (mg/min)
            Kc = inter-compartment clearance (mL/min)
            K = dialyzer clearance (mL/min)
        """
        c_ecf, c_icf = y

        dc_ecf = (
            self.g_urea
            + self.kc * (c_icf - c_ecf)
            - k_dialyzer * c_ecf
        ) / self.v_ecf

        dc_icf = -self.kc * (c_icf - c_ecf) / self.v_icf

        return np.array([dc_ecf, dc_icf])

    def step(
        self,
        blood_flow_rate: float,
        dialysate_flow_rate: float,
        dt: float = 1.0,
    ) -> dict:
        """Advance the model by dt minutes.

        Args:
            blood_flow_rate: Current blood flow rate in mL/min.
            dialysate_flow_rate: Current dialysate flow rate in mL/min.
            dt: Time step in minutes.

        Returns:
            Dict with c_ecf, c_icf, clearance, urea_removed.
        """
        k = self.dialyzer_clearance(blood_flow_rate, dialysate_flow_rate)

        c_ecf_before = self.state[0]

        sol = solve_ivp(
            self._ode_rhs,
            [0, dt],
            self.state,
            args=(k,),
            method="RK45",
            rtol=1e-6,
            atol=1e-8,
        )

        self.state = sol.y[:, -1]
        self.state = np.clip(self.state, 0, 300)  # physiological bounds

        urea_removed = max(0, (c_ecf_before - self.state[0]) * self.v_ecf / 100.0)

        return {
            "c_ecf": float(self.state[0]),
            "c_icf": float(self.state[1]),
            "clearance": float(k),
            "urea_removed_mg": float(urea_removed),
        }

    def compute_kt_v(
        self,
        clearance: float,
        elapsed_minutes: float,
    ) -> float:
        """Compute current Kt/V (dialysis adequacy measure).

        Kt/V = K * t / V where K = clearance, t = time, V = distribution volume.

        Args:
            clearance: Average effective clearance in mL/min.
            elapsed_minutes: Total dialysis time elapsed in minutes.

        Returns:
            Current Kt/V value (target >= 1.2 per KDOQI).
        """
        v_total = self.v_ecf + self.v_icf
        return clearance * elapsed_minutes / v_total
