import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import mediapipe as mp
import traceback
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from torch_geometric.data import Data
from torch_geometric.nn import GATv2Conv, BatchNorm
from insightface.app import FaceAnalysis

# ==========================================
# 1. ANATOMY & REGION MAPS (GOD MODE)
# ==========================================
REGION_DATA = {
    0: [1, 2, 4, 5, 6, 19, 45, 48, 49, 51, 59, 60, 64, 75, 94, 97, 98, 102, 115, 122, 125, 129, 131, 134, 141, 168, 174, 195, 196, 197, 198, 203, 204, 209, 217, 218, 219, 220, 235, 236, 237, 238, 239, 240, 241, 242, 245, 248, 250, 275, 277, 278, 279, 281, 289, 294, 305, 309, 326, 327, 328, 344, 351, 354, 358, 360, 363, 379, 399, 419, 420, 429, 437, 438, 439, 440, 455, 456, 457, 458, 459, 460, 461, 462, 465, 128, 114, 218, 437],
    1: [7, 23, 27, 33, 133, 144, 145, 153, 154, 155, 157, 158, 159, 160, 161, 163, 173, 246, 46, 52, 53, 55, 65, 56, 70, 63, 105, 66, 107, 25, 28, 29, 30, 34, 35, 110, 112, 130, 221, 222, 223, 224, 225, 226, 228, 229, 230, 231, 232, 233, 243, 244, 468, 469, 470, 471, 472],
    2: [249, 263, 362, 373, 374, 380, 381, 382, 384, 385, 386, 387, 388, 390, 398, 466, 276, 282, 283, 285, 295, 296, 300, 293, 334, 290, 336, 255, 257, 258, 259, 260, 265, 339, 341, 359, 441, 442, 443, 444, 445, 446, 448, 449, 450, 451, 452, 453, 463, 464, 473, 474, 475, 476, 477],
    3: [0, 13, 14, 17, 37, 39, 40, 61, 62, 76, 77, 78, 80, 81, 82, 84, 85, 87, 88, 90, 91, 95, 96, 146, 178, 179, 180, 181, 183, 184, 185, 191, 267, 269, 270, 271, 272, 291, 292, 306, 307, 308, 310, 311, 312, 314, 315, 317, 318, 320, 321, 324, 325, 375, 402, 403, 404, 405, 407, 408, 409, 415],
    4: [18, 32, 83, 140, 148, 149, 152, 170, 171, 175, 176, 182, 194, 199, 200, 201, 208, 211, 262, 313, 335, 369, 377, 396, 400, 406, 418, 421, 424, 431, 395, 378, 428],
    5: [31, 36, 50, 100, 101, 111, 116, 117, 118, 119, 120, 121, 123, 126, 132, 135, 137, 138, 142, 143, 147, 165, 166, 167, 169, 177, 186, 187, 192, 202, 205, 206, 207, 210, 212, 213, 214, 215, 216, 227, 58, 93, 136, 150, 172, 234, 127, 162, 21, 54, 103, 67, 109],
    6: [261, 264, 266, 280, 323, 329, 330, 331, 340, 342, 343, 345, 346, 347, 348, 349, 350, 352, 353, 355, 356, 357, 361, 364, 365, 366, 367, 368, 371, 372, 391, 393, 394, 397, 401, 410, 411, 412, 413, 414, 416, 417, 422, 423, 425, 426, 427, 430, 432, 433, 434, 435, 436, 447, 288, 323, 361, 365, 379, 397, 454, 332, 284, 251, 389, 356],
    7: [8, 9, 10, 11, 12, 15, 16, 20, 24, 41, 42, 43, 44, 47, 22, 26, 38, 57, 68, 69, 71, 72, 73, 74, 79, 86, 89, 92, 99, 104, 106, 108, 113, 124, 139, 151, 156, 164, 188, 189, 190, 193, 252, 253, 254, 256, 268, 273, 274, 286, 287, 297, 298, 299, 301, 302, 303, 304, 316, 319, 322, 333, 337, 338, 353, 370, 376, 383, 392, 410, 46, 52, 53, 65, 55, 70, 63, 105, 66, 107, 276, 282, 283, 295, 285, 300, 293, 334, 296, 336]
}

def get_region_tensor(device):
    region_map = torch.full((468,), 8, dtype=torch.long, device=device)
    for region_id, points in REGION_DATA.items():
        for p in points:
            if p < 468: region_map[p] = region_id
    return region_map

def calculate_density(coords, edge_index):
    row, col = edge_index
    dist = torch.norm(coords[row] - coords[col], dim=1)
    sum_dist = torch.zeros(468, device=coords.device)
    sum_dist.scatter_add_(0, row, dist)
    count = torch.zeros(468, device=coords.device)
    count.scatter_add_(0, row, torch.ones_like(dist))
    mean_dist = sum_dist / (count + 1e-6)
    density = 1.0 / (mean_dist + 1e-6)
    return density.unsqueeze(1)

# ==========================================
# 2. MODEL ARCHITECTURES
# ==========================================
class RegionAwareExpert(nn.Module):
    def __init__(self):
        super().__init__()
        self.visual_proj = nn.Linear(64, 64)
        self.point_id_emb = nn.Embedding(468, 64)
        self.region_emb = nn.Embedding(9, 32)
        self.coord_proj = nn.Linear(3, 32)
        self.neck = nn.Sequential(
            nn.Linear(192, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.3),
            nn.Linear(256, 128), nn.LeakyReLU(0.2)
        )
        self.head_gate = nn.Linear(128, 1)
        self.head_tech = nn.Linear(128, 3)    
        self.head_dosage = nn.Linear(128, 8)  
        self.head_depth = nn.Linear(128, 4)   
        self.head_prod = nn.Linear(128, 8)    

    def forward(self, features, coords, point_ids, region_ids):
        vis = self.visual_proj(features)
        pid = self.point_id_emb(point_ids)
        reg = self.region_emb(region_ids)
        xyz = self.coord_proj(coords)
        combined = torch.cat([vis, pid, reg, xyz], dim=-1)
        x = self.neck(combined.view(-1, 192))
        return {
            "tech": self.head_tech(x), 
            "dosage": self.head_dosage(x),
            "depth": self.head_depth(x), 
            "product": self.head_prod(x)
        }

class AnatomyLocationNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.coord_proj = nn.Linear(3, 32)
        self.point_id_emb = nn.Embedding(468, 64)
        self.region_emb = nn.Embedding(9, 32)
        self.density_proj = nn.Linear(1, 16)
        self.context_proj = nn.Linear(512, 32)
        self.neck = nn.Sequential(
            nn.Linear(176, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.4),
            nn.Linear(256, 128), nn.LeakyReLU(0.2), nn.Linear(128, 1)
        )

    def forward(self, coords, pids, rids, den, ctx):
        B, N, _ = coords.shape
        c_emb = self.context_proj(ctx).unsqueeze(1).repeat(1, N, 1)
        combined = torch.cat([self.coord_proj(coords), self.point_id_emb(pids), 
                              self.region_emb(rids), self.density_proj(den), c_emb], dim=-1)
        return self.neck(combined.view(-1, 176)).view(B, N, 1)

class GatedFusion(nn.Module):
    def __init__(self, geo_dim, context_dim):
        super().__init__()
        self.context_adapter = nn.Linear(context_dim, geo_dim)
        self.gate_net = nn.Sequential(nn.Linear(geo_dim * 2, geo_dim // 2), nn.ReLU(), nn.Linear(geo_dim // 2, geo_dim), nn.Sigmoid())
    def forward(self, x_geo, x_ctx):
        ctx_adapted = self.context_adapter(x_ctx)
        ctx_expanded = ctx_adapted.unsqueeze(1).repeat(1, 468, 1)
        combined = torch.cat([x_geo, ctx_expanded], dim=-1)
        return x_geo + (self.gate_net(combined) * ctx_expanded)

class SmartClinicalNet(nn.Module):
    def __init__(self, hidden=32, heads=2):
        super().__init__()
        self.id_emb = nn.Embedding(468, 16)
        self.geo_proj = nn.Linear(19, hidden)
        self.fusion = GatedFusion(geo_dim=hidden, context_dim=515)
        self.conv1 = GATv2Conv(hidden, hidden, heads=heads, concat=True)
        self.bn1 = BatchNorm(hidden * heads)
        self.conv2 = GATv2Conv(hidden * heads, hidden, heads=heads, concat=True)
        self.bn2 = BatchNorm(hidden * heads)

    def get_features(self, data):
        x, edges = data.x, data.edge_index
        batch_size = data.embedding.shape[0]
        global_ctx = torch.cat([data.embedding, data.virt_prof], dim=1)
        ids = torch.arange(468, device=x.device).repeat(batch_size)
        if len(ids) > x.shape[0]: ids = ids[:x.shape[0]]
        geo_feat = self.geo_proj(torch.cat([x, self.id_emb(ids)], dim=1))
        geo_reshaped = geo_feat.view(batch_size, 468, -1)
        fused = self.fusion(geo_reshaped, global_ctx)
        fused_flat = fused.view(-1, 32)
        h = F.elu(self.bn1(self.conv1(fused_flat, edges)))
        h = F.elu(self.bn2(self.conv2(h, edges)))
        return h 

# ==========================================
# 3. INITIALIZATION
# ==========================================
app = FastAPI()
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print("🔄 Initializing SOTA Clinical System...")

mp_mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
face_app.prepare(ctx_id=0, det_size=(640, 640))

mp_edges = mp.solutions.face_mesh.FACEMESH_TESSELATION
s, t = [], []
for src, dst in mp_edges:
    s.extend([src, dst]); t.extend([dst, src])
GLOBAL_EDGE_INDEX = torch.tensor([s, t], dtype=torch.long).to(DEVICE)
STATIC_REGION_IDS = get_region_tensor(DEVICE).unsqueeze(0)
STATIC_POINT_IDS = torch.arange(468, dtype=torch.long, device=DEVICE).unsqueeze(0)

backbone = SmartClinicalNet(hidden=32, heads=2).to(DEVICE)
expert_model = RegionAwareExpert().to(DEVICE)
location_model = AnatomyLocationNet().to(DEVICE)

# Load weights safely
if os.path.exists("smart_clinical_model.pth"):
    backbone.load_state_dict(torch.load("smart_clinical_model.pth", map_location=DEVICE), strict=False)
    backbone.eval()
if os.path.exists("region_expert.pth"):
    expert_model.load_state_dict(torch.load("region_expert.pth", map_location=DEVICE))
    expert_model.eval()
if os.path.exists("anatomy_location.pth"):
    location_model.load_state_dict(torch.load("anatomy_location.pth", map_location=DEVICE))
    location_model.eval()

print("✅ System Ready.")

# ==========================================
# 4. HELPERS
# ==========================================
def get_virtual_profile_norm(x_tensor):
    NOSE=1; LIP=13; CHIN=152; L_CHEEK=234; L_JAW=172; L_EYE=33; R_EYE=263
    eye_dist = torch.norm(x_tensor[L_EYE] - x_tensor[R_EYE]) + 1e-6
    chin = (x_tensor[CHIN, 2] - x_tensor[LIP, 2]) / eye_dist
    cheek = (x_tensor[L_CHEEK, 2] - x_tensor[NOSE, 2]) / eye_dist
    jaw = torch.norm(x_tensor[CHIN] - x_tensor[L_JAW]) / eye_dist
    return torch.tensor([chin, cheek, jaw], dtype=torch.float)

#spacing is where we cluster points into one default was 12.0
def apply_nms_indices(landmarks_np, probs_np, spacing_mm=12.0, threshold=0.5):
    valid = np.where(probs_np > threshold)[0]
    if len(valid) == 0: return []
    sorted_idx = valid[np.argsort(probs_np[valid])[::-1]]
    keep = []
    eye_l, eye_r = landmarks_np[33], landmarks_np[263]
    eye_dist_px = np.linalg.norm(eye_l - eye_r)
    mm_per_px = 63.0 / (eye_dist_px + 1e-6)
    min_dist_sq = (spacing_mm / mm_per_px) ** 2
    while len(sorted_idx) > 0:
        curr = sorted_idx[0]
        keep.append(int(curr))
        if len(sorted_idx) == 1: break
        rest = sorted_idx[1:]
        dists = np.sum((landmarks_np[rest] - landmarks_np[curr])**2, axis=1)
        survivors = np.where(dists > min_dist_sq)[0]
        sorted_idx = rest[survivors]
    return keep

def get_top_probs(probs_array, class_list):
    sorted_idxs = np.argsort(probs_array)[::-1][:3]
    return {class_list[i]: float(f"{probs_array[i]:.4f}") for i in sorted_idxs}

# ==========================================
# 5. API ENDPOINT (FINAL - WITH ALL COORDINATES)
# ==========================================
@app.get("/")
def home():
    return {"message": "SOTA Clinical AI - Ready"}

@app.post("/predict")
async def predict_injections(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img_bgr is None: return JSONResponse(status_code=400, content={"error": "Invalid image"})
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        
        res = mp_mesh.process(img_rgb)
        if not res.multi_face_landmarks:
            return JSONResponse(status_code=400, content={"error": "No face detected"})
        
        lms = res.multi_face_landmarks[0].landmark
        coords = [[p.x, p.y, p.z] for p in lms[:468]]
        x_geo = torch.tensor(coords, dtype=torch.float).to(DEVICE)
        x_geo_norm = x_geo - x_geo.mean(dim=0)
        
        faces = face_app.get(img_bgr)
        emb = torch.tensor(faces[0].embedding).float().to(DEVICE) if faces else torch.zeros(512).to(DEVICE)
        virt_prof = get_virtual_profile_norm(x_geo_norm.cpu()).to(DEVICE)
        
        data = Data(x=x_geo_norm, edge_index=GLOBAL_EDGE_INDEX, embedding=emb.unsqueeze(0), virt_prof=virt_prof.unsqueeze(0))
        density = calculate_density(x_geo_norm, GLOBAL_EDGE_INDEX).unsqueeze(0).to(DEVICE)
        
        with torch.no_grad():
            logits_loc = location_model(x_geo_norm.unsqueeze(0), STATIC_POINT_IDS, STATIC_REGION_IDS, density, emb.unsqueeze(0))
            probs_loc = torch.sigmoid(logits_loc).squeeze().cpu().numpy()
            
            smart_features = backbone.get_features(data).unsqueeze(0)
            coords_input = x_geo_norm.unsqueeze(0)
            preds = expert_model(smart_features, coords_input, STATIC_POINT_IDS, STATIC_REGION_IDS)
            
            prob_t = torch.softmax(preds['tech'], dim=-1).squeeze().cpu().numpy()
            prob_d = torch.softmax(preds['dosage'], dim=-1).squeeze().cpu().numpy()
            prob_de = torch.softmax(preds['depth'], dim=-1).squeeze().cpu().numpy()
            prob_p = torch.softmax(preds['product'], dim=-1).squeeze().cpu().numpy()

        h, w, _ = img_bgr.shape
        pixel_coords = np.array([[p.x*w, p.y*h, p.z*w] for p in lms[:468]])
        optimal_indices_list = apply_nms_indices(pixel_coords, probs_loc, spacing_mm=10.0, threshold=0.4)
        optimal_set = set(optimal_indices_list)
        
        classes_tech = ["Bolus", "Fanning", "Microbolus"]
        classes_dosage = ["0.01ml", "0.02ml", "0.05ml", "0.1ml", "0.2ml", "0.3ml", "0.5ml", "1.0ml"]
        classes_depth = ["Periosteal", "Subdermal", "Hypodermic", "Dermal"]
        classes_prod = ["XXL", "XL", "L", "M", "S", "Hydro", "Induce", "Lips"]

        all_points_list = []
        for idx in range(468):
            p_t, p_d, p_de, p_p = prob_t[idx], prob_d[idx], prob_de[idx], prob_p[idx]
            pt_info = {
                "point_id": int(idx),
                "confidence": float(f"{probs_loc[idx]:.4f}"),
                "is_optimal": idx in optimal_set,
                "coordinates": {"x": lms[idx].x, "y": lms[idx].y},
                "attributes": {
                    "technique": classes_tech[np.argmax(p_t)],
                    "dosage": classes_dosage[np.argmax(p_d)],
                    "depth": classes_depth[np.argmax(p_de)],
                    "product": classes_prod[np.argmax(p_p)],
                    "technique_probs": get_top_probs(p_t, classes_tech),
                    "dosage_probs": get_top_probs(p_d, classes_dosage),
                    "depth_probs": get_top_probs(p_de, classes_depth),
                    "product_probs": get_top_probs(p_p, classes_prod)
                }
            }
            all_points_list.append(pt_info)
        
        # ✅ BACKEND FIX: Explicitly generating the list of landmarks here
        all_coords_list = [[float(f"{p.x:.4f}"), float(f"{p.y:.4f}")] for p in lms[:468]]

        return {
            "status": "success",
            "message": "SOTA V10: All Data + Landmarks",
            "injection_sites": all_points_list,
            "all_coordinates": all_coords_list, # <--- SENDING SEPARATE LIST AGAIN
            "summary": {"total_optimal": len(optimal_set)}
        }

    except Exception as e:
        print(traceback.format_exc())
        return JSONResponse(status_code=400, content={"error": str(e), "trace": traceback.format_exc()})