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

Phosphate has slower inter-compartment transfer than urea, leading to
significant post-dialysis rebound. The model captures both dialytic
clearance and dietary phosphate absorption.

References:
    Locatelli F, et al. Nephrol Dial Transplant. 2002;17(Suppl 11):24-27.
    Daugirdas JT, et al. Semin Dial. 2025;38(1):e13252.
"""

import numpy as np
from scipy.integrate import solve_ivp


class PhosphateKineticsModel:
    """Two-compartment phosphate kinetics with dietary intake.

    SIMPLIFICATION: Uses two-compartment model with fixed transfer coefficient.
    Clinical phosphate kinetics involves additional bone-mineral compartment
    and PTH-mediated regulation. The pseudo one-compartment model from
    the Daugirdas 2025 review is adequate for dialysis-session timescale
    simulation. Acceptable for benchmark simulation.
    """

    def __init__(
        self,
        ecf_volume: float = 14000.0,
        icf_volume: float = 28000.0,
        phosphate_koa: float = 400.0,
        intercompartment_kp: float = 150.0,
        dietary_phosphate_mg_per_day: float = 1200.0,
        gi_absorption_fraction: float = 0.65,
    ):
        """Initialize phosphate kinetics model.

        Args:
            ecf_volume: Extracellular fluid volume in mL.
            icf_volume: Intracellular fluid volume in mL.
            phosphate_koa: Dialyzer phosphate clearance coefficient (mL/min).
            intercompartment_kp: Inter-compartment transfer for phosphate (mL/min).
            dietary_phosphate_mg_per_day: Daily dietary phosphate intake (mg).
            gi_absorption_fraction: GI absorption fraction (typically 0.6-0.7).
        """
        self.v_ecf = ecf_volume
        self.v_icf = icf_volume
        self.koa_p = phosphate_koa
        self.kp = intercompartment_kp
        self.dietary_p = dietary_phosphate_mg_per_day
        self.gi_absorption = gi_absorption_fraction

        # State: [P_ecf (mg/dL), P_icf (mg/dL)]
        self.state = np.array([5.5, 5.5])

    def reset(self, p_ecf: float = 5.5, p_icf: float = 5.5) -> np.ndarray:
        """Reset phosphate state."""
        self.state = np.array([p_ecf, p_icf])
        return self.state.copy()

    def _ode_rhs_dialysis(self, _t: float, y: np.ndarray, clearance: float) -> np.ndarray:
        """ODE during dialysis session."""
        p_ecf, p_icf = y

        # Dietary input distributed over 24h, assume meals during dialysis
        dietary_rate = (self.dietary_p * self.gi_absorption) / (24 * 60)  # mg/min

        dp_ecf = (
            dietary_rate
            + self.kp * (p_icf - p_ecf)
            - clearance * p_ecf
        ) / self.v_ecf * 100  # convert to mg/dL

        dp_icf = -self.kp * (p_icf - p_ecf) / self.v_icf * 100

        return np.array([dp_ecf, dp_icf])

    def _ode_rhs_interdialytic(self, _t: float, y: np.ndarray) -> np.ndarray:
        """ODE during inter-dialytic period (no clearance)."""
        return self._ode_rhs_dialysis(_t, y, clearance=0.0)

    def step_dialysis(
        self,
        blood_flow_rate: float,
        dialysate_flow_rate: float,
        binder_dose_mg: float = 0.0,
        dt: float = 1.0,
    ) -> dict:
        """Advance during a dialysis session.

        Args:
            blood_flow_rate: Blood flow rate in mL/min.
            dialysate_flow_rate: Dialysate flow rate in mL/min.
            binder_dose_mg: Phosphate binder dose in mg (reduces GI absorption).
            dt: Time step in minutes.

        Returns:
            Dict with current phosphate levels.
        """
        # Effective clearance for phosphate (lower than urea due to larger molecule)
        qb = max(blood_flow_rate, 1.0)
        qd = max(dialysate_flow_rate, 1.0)
        clearance = min(qb, qd, self.koa_p) * 0.6

        # Binder effect: reduce effective GI absorption
        binder_effect = min(1.0, binder_dose_mg / 2000.0) * 0.5
        original_gi = self.gi_absorption
        self.gi_absorption = max(0.1, original_gi * (1.0 - binder_effect))

        sol = solve_ivp(
            self._ode_rhs_dialysis,
            [0, dt],
            self.state,
            args=(clearance,),
            method="RK45",
            rtol=1e-6,
            atol=1e-8,
        )
        self.state = np.clip(sol.y[:, -1], 0.5, 15.0)
        self.gi_absorption = original_gi

        return {
            "p_ecf": float(self.state[0]),
            "p_icf": float(self.state[1]),
            "clearance": float(clearance),
            "calcium_phosphate_product": float(self.state[0] * 9.5),
        }

    def step_interdialytic(
        self,
        binder_dose_mg: float = 0.0,
        dt_hours: float = 8.0,
    ) -> dict:
        """Advance during inter-dialytic period.

        Args:
            binder_dose_mg: Phosphate binder dose in mg.
            dt_hours: Time step in hours.

        Returns:
            Dict with current phosphate levels.
        """
        binder_effect = min(1.0, binder_dose_mg / 2000.0) * 0.5
        original_gi = self.gi_absorption
        self.gi_absorption = max(0.1, original_gi * (1.0 - binder_effect))

        dt_min = dt_hours * 60.0
        sol = solve_ivp(
            self._ode_rhs_interdialytic,
            [0, dt_min],
            self.state,
            method="RK45",
            rtol=1e-6,
            atol=1e-8,
        )
        self.state = np.clip(sol.y[:, -1], 0.5, 15.0)
        self.gi_absorption = original_gi

        return {
            "p_ecf": float(self.state[0]),
            "p_icf": float(self.state[1]),
            "calcium_phosphate_product": float(self.state[0] * 9.5),
        }
