import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from scipy.integrate import solve_ivp

app = FastAPI(
    title="Advanced PK/PD Physiological Simulation Engine",
    description="Multi-compartment PBPK engine for vaccine visual tracking across specific organs",
    version="2.0.0"
)

# Robust CORS configuration to allow your local desktop Electron client to talk to the cloud
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- MATHEMATICAL SCHEMAS ---

class OrganParameters(BaseModel):
    name: str = Field(..., description="Name of target organ/tissue profile clicked")
    volume: float = Field(..., description="Tissue compartment volume (V_organ) in Liters")
    blood_flow: float = Field(..., description="Blood flow rate to the tissue (Q_organ) in L/min")
    partition_coefficient: float = Field(1.2, description="Tissue-to-plasma partition ratio (Kp)")

class PKPDRequest(BaseModel):
    initial_dose: float = Field(..., description="Administered vaccine amount/payload concentration")
    organ: OrganParameters = Field(..., description="Anatomical structural metrics fetched from Cloudflare/HRA")
    clearance_systemic: float = Field(0.2, description="Total metabolic clearance (Cl) in L/min")
    e_max: float = Field(100.0, description="Maximum pharmacodynamic inflammatory threshold response")
    ec_50: float = Field(10.0, description="Concentration causing 50% maximal effect parameter")
    simulation_hours: int = Field(72, ge=1, le=336, description="Total simulation window limit")

# --- SIMULATION ENGINE ENGINE ---

def pbpk_ode_system(t, y, Q_org, V_org, Kp, Cl, V_central=5.0):
    """
    Simulates a closed loop multi-compartment system:
    y[0] = Amount of vaccine localized at the primary tissue site
    y[1] = Amount of vaccine flowing through central vascular circulation
    """
    A_org = y[0]
    A_central = y[1]
    
    # Deriving concentration states
    C_org = A_org / V_org if V_org > 0 else 0
    C_central = A_central / V_central
    
    # ODE Mass Balance Equations
    # Rate of change in target organ = Delivery from arterial flow - Venous drainage outward
    dA_org_dt = Q_org * (C_central - (C_org / Kp))
    
    # Rate of change in central system = Return from tissue - Clearance/Elimination
    dA_central_dt = Q_org * ((C_org / Kp) - C_central) - (Cl * C_central)
    
    return [dA_org_dt, dA_central_dt]

@app.post("/api/v1/simulate/pkpd")
async def execute_simulation(payload: PKPDRequest):
    try:
        # 1. Setup Time Matrix
        t_span = (0, payload.simulation_hours)
        # Generate 150 points for extra smooth 3D graph animations inside Three.js
        t_eval = np.linspace(0, payload.simulation_hours, 150)
        
        # 2. Extract Physiological parameters
        Q_org = payload.organ.blood_flow
        V_org = payload.organ.volume
        Kp = payload.organ.partition_coefficient
        Cl = payload.clearance_systemic
        
        # 3. Set Initial Conditions
        # Assume initial dose enters the specified local target tissue compartment immediately
        initial_state = [payload.initial_dose, 0.0]
        
        # 4. Numerical ODE Solver (Runge-Kutta 45 integration)
        solution = solve_ivp(
            pbpk_ode_system,
            t_span,
            initial_state,
            args=(Q_org, V_org, Kp, Cl),
            t_eval=t_eval,
            method='RK45'
        )
        
        if not solution.success:
            raise HTTPException(status_code=500, detail="The differential integration solver failed to converge.")
            
        # 5. Extract Amounts & Calculate Concentrations
        amounts_tissue = solution.y[0]
        amounts_blood = solution.y[1]
        
        concentrations_tissue = amounts_tissue / V_org
        concentrations_blood = amounts_blood / 5.0 # Central Blood Volume constant
        
        # 6. Calculate Pharmacodynamics (PD) / Vitals Impact
        # Emax sigmoid mathematical function evaluating localized physiological shift
        local_vitals_delta = (payload.e_max * concentrations_tissue) / (payload.ec_50 + concentrations_tissue)
        
        # Systemic immune response representation (e.g., simulated temperature spike)
        fever_metric = 37.0 + ((payload.e_max * 1.5 * concentrations_blood) / (payload.ec_50 + concentrations_blood))
        
        return {
            "status": "success",
            "time_series_hours": solution.t.tolist(),
            "kinetics": {
                "organ_amount": amounts_tissue.tolist(),
                "organ_concentration": concentrations_tissue.tolist(),
                "blood_amount": amounts_blood.tolist(),
                "blood_concentration": concentrations_blood.tolist()
            },
            "vitals_metrics": {
                "localized_cellular_activation_percent": local_vitals_delta.tolist(),
                "simulated_core_body_temperature": np.clip(fever_metric, 37.0, 41.5).tolist()
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Simulation parameter error: {str(e)}")

@app.get("/")
@app.head("/")
async def status_ping():
    return {
        "status": "active",
        "engine": "PBPK Math Engine V2",
        "concurrency": "Multi-Worker Optimized"
    }

@app.get("/ping")
@app.head("/ping")
async def uptime_ping():
    """
    Endpoint for UptimeRobot monitoring
    Supports both GET and HEAD requests for health checks
    """
    return {
        "status": "ok",
        "service": "pharma-pkpd",
        "timestamp": "active"
    }

@app.get("/health")
@app.head("/health")
async def health_check():
    """
    Detailed health check endpoint
    Supports both GET and HEAD requests
    """
    return {
        "status": "healthy",
        "service": "AEGIS Pharma PK/PD Engine",
        "version": "2.0.0",
        "endpoints": {
            "simulation": "/api/v1/simulate/pkpd",
            "status": "/",
            "ping": "/ping",
            "health": "/health"
        }
    }