```python
#!/usr/bin/env python3
"""
Bayesian Sequential Urinalysis Monitoring for Early Lupus Nephritis Detection
Claw4S Skill — DNAI-ClinicalAI

Hidden Markov Model for lupus nephritis class prediction from serial urinalysis,
with biopsy recommendation and medication adjustment triggers.
"""

import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional
import json, sys

# --- ISN/RPS Lupus Nephritis Classes ---
LN_CLASSES = ["No_Nephritis", "Class_II_Mesangial", "Class_III_Focal",
              "Class_IV_Diffuse", "Class_V_Membranous"]

@dataclass
class UrinalysisResult:
    protein: int          # 0=neg, 1=trace, 2=1+, 3=2+, 4=3+, 5=4+
    blood: int            # 0=neg, 1=trace, 2=small, 3=moderate, 4=large
    rbc_casts: bool       # Red blood cell casts present
    wbc_casts: bool       # White blood cell casts present
    granular_casts: bool  # Granular casts
    specific_gravity: float  # 1.001-1.035
    upcr: Optional[float] = None  # Urine protein:creatinine ratio (mg/mg)

@dataclass
class SerologicData:
    anti_dsdna: Optional[float] = None   # IU/mL
    c3: Optional[float] = None           # mg/dL (normal 90-180)
    c4: Optional[float] = None           # mg/dL (normal 10-40)
    sledai_renal: Optional[int] = None   # 0-16

class LupusNephritisHMM:
    """Hidden Markov Model for lupus nephritis staging from serial urinalysis."""

    def __init__(self):
        self.states = LN_CLASSES
        self.n_states = len(self.states)

        # Prior: most SLE patients start without nephritis
        self.prior = np.array([0.50, 0.20, 0.12, 0.10, 0.08])

        # Transition matrix (per 3-month interval)
        # Based on Hopkins Lupus Cohort and LUNAR trial data
        self.transition = np.array([
            #  NoN    II     III    IV     V
            [0.85, 0.08, 0.03, 0.02, 0.02],  # No Nephritis
            [0.05, 0.70, 0.12, 0.08, 0.05],  # Class II
            [0.02, 0.05, 0.65, 0.20, 0.08],  # Class III
            [0.01, 0.02, 0.07, 0.80, 0.10],  # Class IV
            [0.02, 0.03, 0.05, 0.10, 0.80],  # Class V
        ])

        # Emission parameters: P(urinalysis findings | LN class)
        # protein_mean, blood_mean, cast_prob, low_sg_prob
        self.emission_params = {
            "No_Nephritis":      {"protein_mu": 0.3, "blood_mu": 0.2, "cast_p": 0.02, "upcr_mu": 0.1},
            "Class_II_Mesangial": {"protein_mu": 1.2, "blood_mu": 1.0, "cast_p": 0.10, "upcr_mu": 0.3},
            "Class_III_Focal":   {"protein_mu": 2.5, "blood_mu": 2.5, "cast_p": 0.35, "upcr_mu": 1.2},
            "Class_IV_Diffuse":  {"protein_mu": 3.8, "blood_mu": 3.2, "cast_p": 0.55, "upcr_mu": 3.5},
            "Class_V_Membranous": {"protein_mu": 4.0, "blood_mu": 1.0, "cast_p": 0.15, "upcr_mu": 4.0},
        }

    def emission_probability(self, ua: UrinalysisResult, state: str) -> float:
        """P(urinalysis | LN class) using Poisson-like model."""
        p = self.emission_params[state]
        # Protein likelihood
        protein_ll = np.exp(-0.5 * ((ua.protein - p["protein_mu"]) / 1.2) ** 2)
        # Blood likelihood
        blood_ll = np.exp(-0.5 * ((ua.blood - p["blood_mu"]) / 1.2) ** 2)
        # Cast likelihood
        has_casts = ua.rbc_casts or ua.wbc_casts or ua.granular_casts
        cast_ll = p["cast_p"] if has_casts else (1 - p["cast_p"])
        # UPCR if available
        upcr_ll = 1.0
        if ua.upcr is not None:
            upcr_ll = np.exp(-0.5 * ((ua.upcr - p["upcr_mu"]) / 1.5) ** 2)

        return protein_ll * blood_ll * cast_ll * upcr_ll + 1e-10

    def update_posterior(self, prior: np.ndarray, ua: UrinalysisResult,
                        serology: Optional[SerologicData] = None) -> np.ndarray:
        """Bayesian update: P(state | observation) ∝ P(obs | state) × P(state)."""
        # Predict step (transition)
        predicted = prior @ self.transition

        # Update step (emission)
        likelihoods = np.array([self.emission_probability(ua, s) for s in self.states])

        # Serologic adjustment
        if serology:
            sero_factor = np.ones(self.n_states)
            if serology.c3 is not None and serology.c3 < 90:
                hypoC3 = max(0.1, serology.c3 / 90)
                sero_factor[2:4] *= (1 + (1 - hypoC3) * 2)  # Boost III/IV
            if serology.anti_dsdna is not None and serology.anti_dsdna > 30:
                sero_factor[2:5] *= min(2.0, serology.anti_dsdna / 30)
            if serology.sledai_renal is not None and serology.sledai_renal >= 4:
                sero_factor[2:5] *= 1.5
            likelihoods *= sero_factor

        posterior = predicted * likelihoods
        posterior /= posterior.sum()
        return posterior

    def sequential_monitor(self, observations: List[Tuple[UrinalysisResult, Optional[SerologicData]]]) -> List[dict]:
        """Process serial urinalysis observations and track LN probability over time."""
        posterior = self.prior.copy()
        results = []

        for i, (ua, sero) in enumerate(observations):
            posterior = self.update_posterior(posterior, ua, sero)
            probs = {s: round(float(p), 4) for s, p in zip(self.states, posterior)}
            proliferative_risk = float(posterior[2] + posterior[3])  # III + IV

            # Biopsy recommendation
            biopsy_rec = "RECOMMEND BIOPSY" if proliferative_risk > 0.40 else \
                         "Consider biopsy" if proliferative_risk > 0.25 else "Continue monitoring"

            # Medication triggers
            med_action = self._medication_recommendation(posterior, proliferative_risk)

            results.append({
                "visit": i + 1,
                "posterior": probs,
                "most_likely": self.states[np.argmax(posterior)],
                "proliferative_risk": round(proliferative_risk, 4),
                "biopsy_recommendation": biopsy_rec,
                "medication_action": med_action,
            })
        return results

    def _medication_recommendation(self, posterior: np.ndarray, prolif_risk: float) -> str:
        if prolif_risk > 0.60:
            return "Consider cyclophosphamide or high-dose MMF (2-3g/day) + pulse methylprednisolone"
        elif prolif_risk > 0.40:
            return "Start/increase MMF to 2-3g/day, consider renal biopsy before escalation"
        elif posterior[4] > 0.40:  # Class V
            return "Consider calcineurin inhibitor (tacrolimus) + low-dose MMF for membranous"
        elif posterior[1] > 0.40:  # Class II
            return "Low-dose MMF (1-1.5g/day) or monitoring if mild"
        else:
            return "Continue current regimen, serial urinalysis q3months"

# --- Synthetic data generator ---
def simulate_lupus_patient(trajectory: str = "progressive", n_visits: int = 6, seed: int = 42):
    """Simulate a lupus patient's serial urinalysis."""
    rng = np.random.default_rng(seed)
    observations = []

    trajectories = {
        "progressive": [0, 0, 1, 2, 3, 3, 3],  # No nephritis → Class IV
        "stable": [0, 0, 0, 1, 1, 1, 1],        # Mild, stays mesangial
        "membranous": [0, 1, 1, 4, 4, 4, 4],     # Develops Class V
        "treatment_response": [3, 3, 2, 1, 0, 0], # IV → remission
    }
    class_sequence = trajectories.get(trajectory, trajectories["progressive"])

    for i in range(min(n_visits, len(class_sequence))):
        cls_idx = class_sequence[i]
        cls = LN_CLASSES[cls_idx]
        hmm = LupusNephritisHMM()
        p = hmm.emission_params[cls]

        ua = UrinalysisResult(
            protein=max(0, min(5, int(rng.normal(p["protein_mu"], 0.8)))),
            blood=max(0, min(4, int(rng.normal(p["blood_mu"], 0.8)))),
            rbc_casts=rng.random() < p["cast_p"],
            wbc_casts=rng.random() < p["cast_p"] * 0.5,
            granular_casts=rng.random() < p["cast_p"] * 0.3,
            specific_gravity=round(rng.normal(1.015, 0.005), 3),
            upcr=round(max(0, rng.normal(p["upcr_mu"], 0.5)), 2)
        )
        # Serology correlates
        sero = SerologicData(
            anti_dsdna=round(max(0, rng.normal(15 + cls_idx * 25, 10)), 1),
            c3=round(max(20, rng.normal(130 - cls_idx * 20, 15)), 1),
            c4=round(max(5, rng.normal(25 - cls_idx * 4, 5)), 1),
            sledai_renal=min(16, max(0, cls_idx * 4 + int(rng.normal(0, 1))))
        )
        observations.append((ua, sero))
    return observations

def main():
    print("=" * 70)
    print("BAYESIAN LUPUS NEPHRITIS MONITOR — Claw4S Skill")
    print("=" * 70)

    hmm = LupusNephritisHMM()

    for trajectory in ["progressive", "treatment_response", "membranous"]:
        print(f"\n{'─' * 70}")
        print(f"Patient trajectory: {trajectory.upper()}")
        print(f"{'─' * 70}")

        obs = simulate_lupus_patient(trajectory, n_visits=6, seed=42)
        results = hmm.sequential_monitor(obs)

        for r in results:
            print(f"\n  Visit {r['visit']}: Most likely = {r['most_likely']}")
            print(f"    Proliferative risk (III+IV): {r['proliferative_risk']:.1%}")
            print(f"    → {r['biopsy_recommendation']}")
            print(f"    → Rx: {r['medication_action']}")
            top3 = sorted(r['posterior'].items(), key=lambda x: -x[1])[:3]
            print(f"    Top probabilities: {', '.join(f'{k}={v:.1%}' for k,v in top3)}")

if __name__ == "__main__":
    main()

```