import torch
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any

from core.tensor_bridge import TensorBridge
from core.amlm import load_amlm_model
from core.aeg import load_aeg_model
from core.triage import TriageEngineV2

# ==========================================
# 1. SETUP & CONFIG
# ==========================================
app = FastAPI(title="Clario Training Module", version="1.0.0")

# Global Model Containers
MODELS = {}
DEVICE = "cpu" # Force CPU for Hugging Face Spaces (unless GPU upgraded)

# Paths to weights (Uploaded in Phase 5)
AMLM_PATH = "models/amlm_v5.pth"
AEG_PATH = "models/aeg_v5.pth"

# Activity Format Map (Must match AEG Training Data)
FORMAT_MAP = [
    "Auditory_Bombardment", "Phoneme_Isolation", "Minimal_Pairs",
    "Syllable_Chaining", "Carrier_Phrases", "Pacing",
    "Shadowing", "Speed_Drills"
]
STRUCTURE_MAP = ["CV", "VC", "CVC", "VCV", "CVCV", "CCV", "VCC", "Complex"]

# ==========================================
# 2. DATA MODELS (Pydantic)
# ==========================================
class TensorUpdatePayload(BaseModel):
    current_tensor: Optional[List[List[List[float]]]] = None # 3x64x9 List
    new_results: List[Dict[str, Any]] # From Backend or Practice
    
class UserStats(BaseModel):
    severity: float # 0.0 - 1.0
    fatigue: float # 0.0 - 1.0
    age: float # Normalized 0.0 - 1.0
    delta: float # Velocity (-0.5 to +0.5)

class StrategyGeneratePayload(BaseModel):
    current_tensor: List[List[List[float]]] # The raw observed tensor
    user_stats: UserStats
    force_easier: bool = False # The "Frustration Valve"

# ==========================================
# 3. LIFECYCLE EVENTS
# ==========================================
@app.on_event("startup")
async def startup_event():
    print("🚀 Initializing Clario Training Module...")
    try:
        # Load Cognitive Core
        MODELS['amlm'] = load_amlm_model(AMLM_PATH, DEVICE)
        MODELS['aeg'] = load_aeg_model(AEG_PATH, DEVICE)
        MODELS['bridge'] = TensorBridge()
        MODELS['triage'] = TriageEngineV2(DEVICE)
        print("✅ All Systems Operational.")
    except Exception as e:
        print(f"❌ Startup Failed: {e}")

# ==========================================
# 4. HELPER FUNCTIONS
# ==========================================
def _prepare_amlm_input(raw_tensor_list):
    """
    Converts raw 3x64x9 list into Batchx6x64x9 Tensor for AMLM.
    Generates visibility mask based on non-zero entries.
    """
    # 1. Convert to Numpy (3, 64, 9)
    val_map = np.array(raw_tensor_list, dtype=np.float32)
    
    # 2. Generate Mask (Assumption: Val > 0 means "Observed Error")
    # Zeros are treated as "Ambiguous/Untested" for the Inference Engine
    vis_map = (val_map > 0.0).astype(np.float32)
    
    # 3. Concatenate (6, 64, 9)
    input_numpy = np.concatenate([val_map, vis_map], axis=0)
    
    # 4. Add Batch Dim & Convert to Torch
    return torch.FloatTensor(input_numpy).unsqueeze(0).to(DEVICE)

# ==========================================
# 5. ENDPOINTS
# ==========================================

@app.post("/tensor/update")
async def update_tensor(payload: TensorUpdatePayload):
    """
    THE FAST LOOP: Updates the patient's digital twin with new error data.
    """
    bridge = MODELS['bridge']
    
    # Run Bridge Logic
    updated_tensor = bridge.update_tensor(
        payload.current_tensor,
        payload.new_results
    )
    
    return {
        "status": "success",
        "message": "Tensor updated successfully.",
        "updated_tensor": updated_tensor
    }

@app.post("/strategy/generate")
async def generate_strategy(payload: StrategyGeneratePayload):
    """
    THE SLOW LOOP: Generates a full therapy session blueprint.
    """
    amlm = MODELS['amlm']
    aeg = MODELS['aeg']
    triage = MODELS['triage']
    
    # 1. AMLM Inference (Fill Gaps)
    # We use the raw tensor to infer the full clinical image
    amlm_input = _prepare_amlm_input(payload.current_tensor)
    with torch.no_grad():
        diagnosis_map, _ = amlm(amlm_input)
        # Squeeze batch dim -> (3, 64, 9) Numpy
        diag_np = diagnosis_map[0].cpu().numpy()
        
    # 2. Triage (Find Focus)
    all_diagnoses = triage.scan_tensor(diag_np)
    
    if not all_diagnoses:
        return {"status": "healthy", "message": "No significant errors detected."}
        
    focus_diag = all_diagnoses[0] # Highest priority issue
    
    # 3. Prepare Stats (Frustration Valve Logic)
    stats = payload.user_stats
    effective_fatigue = 1.0 if payload.force_easier else stats.fatigue
    
    # 4. AEG Input Preparation
    aeg_inputs = triage.prepare_aeg_input(
        focus_diag, 
        all_diagnoses, 
        [stats.severity, effective_fatigue, stats.age, stats.delta]
    )
    
    # 5. AEG Inference (Generate Blueprint)
    with torch.no_grad():
        p_vol, p_diff, p_fmt, p_targ, p_struct, p_forbid, p_dyn = aeg(
            aeg_inputs['error_idx'], aeg_inputs['major_idx'], aeg_inputs['pos_idx'],
            aeg_inputs['cluster'], aeg_inputs['trap'], aeg_inputs['feat'],
            aeg_inputs['flags'], aeg_inputs['stats']
        )
        
    # 6. Decode Outputs
    vol = int(torch.clamp(p_vol, 3, 10).item())
    
    # Extract Sequences (Slice by volume)
    diff_seq = torch.argmax(p_diff, dim=2)[0].tolist()[:vol]
    # Difficulty is 0-4 internally, map to 1-5 for UI
    diff_seq_ui = [d + 1 for d in diff_seq]
    
    fmt_indices = torch.argmax(p_fmt, dim=2)[0].tolist()[:vol]
    fmt_names = [FORMAT_MAP[i] for i in fmt_indices]
    
    struct_indices = torch.argmax(p_struct, dim=2)[0].tolist()[:vol]
    struct_names = [STRUCTURE_MAP[i] for i in struct_indices]
    
    # Target Phonemes (Multi-hot)
    target_logits = p_targ[0][:vol] # (Vol, 64)
    target_phonemes = []
    
    from core.triage import PHONEMES # Helper to map back
    
    for step_logits in target_logits:
        # Sigmoid > 0.5 implies active target
        active_indices = (torch.sigmoid(step_logits) > 0.5).nonzero(as_tuple=True)[0].tolist()
        targets = [PHONEMES[i] for i in active_indices]
        if not targets: targets = focus_diag.active_phonemes # Fallback to main targets
        target_phonemes.append(targets)

    return {
        "status": "generated",
        "focus_diagnosis": {
            "name": focus_diag.error_name,
            "type": focus_diag.major_type,
            "severity": focus_diag.severity
        },
        "blueprint": {
            "volume": vol,
            "plan": [
                {
                    "step": i+1,
                    "difficulty": diff_seq_ui[i],
                    "format": fmt_names[i],
                    "structure": struct_names[i],
                    "targets": target_phonemes[i]
                }
                for i in range(vol)
            ]
        }
    }