import io
import torch
import numpy as np
from PIL import Image
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import Response
from torchvision import transforms

import matplotlib
matplotlib.use('Agg') # Use non-interactive backend for server
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

# Import architecture, colors, and utils from your model.py
from model import (
    LSNN, 
    IMAGENET_MEAN, 
    IMAGENET_STD, 
    DISPLAY_COLORS, 
    LEGEND_ENTRIES,
    predict_tta
)

app = FastAPI(title="Semantic Change Detection API")

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = None

preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

def class_map_to_rgb(cls_map):
    rgb = np.zeros((*cls_map.shape, 3), dtype=np.uint8)
    for cls, col in DISPLAY_COLORS.items():
        rgb[cls_map == cls] = col
    return rgb

def make_legend_patches():
    return [Patch(facecolor=np.array(c) / 255, label=l) for c, l in LEGEND_ENTRIES]

@app.get("/health")
async def health_check():
    return {
        "status": "ok",
        "model_loaded": model is not None,
        "device": str(device),
    }

@app.on_event("startup")
async def load_model():
    global model
    print("Loading LSNN model...")
    model = LSNN(hidden=96).to(device)
    
    ckpt = torch.load("best_lsnn.pth", map_location=device)
    if "model" in ckpt:
        model.load_state_dict(ckpt["model"])
    else:
        model.load_state_dict(ckpt)
        
    model.eval()
    print("Model loaded!")

@app.post("/predict")
async def predict_change(img1: UploadFile = File(...), img2: UploadFile = File(...)):
    # 1. Read and prep images
    raw_img1 = Image.open(io.BytesIO(await img1.read())).convert("RGB")
    raw_img2 = Image.open(io.BytesIO(await img2.read())).convert("RGB")

    # Resize raw images for plotting so they match the 256x256 prediction map
    disp_img1 = raw_img1.resize((256, 256))
    disp_img2 = raw_img2.resize((256, 256))

    # 2. Preprocess for the model
    t1 = preprocess(raw_img1).unsqueeze(0).to(device)
    t2 = preprocess(raw_img2).unsqueeze(0).to(device)

    # 3. Inference
    with torch.no_grad():
        sem_logits = predict_tta(model, t1, t2)
        pred_mask = torch.argmax(sem_logits, dim=1).squeeze(0).cpu().numpy()

    # 4. Process predictions
    change_pct = 100.0 * (pred_mask != 0).sum() / pred_mask.size
    pred_rgb = class_map_to_rgb(pred_mask)

    # 5. Generate Matplotlib Figure
    fig = plt.figure(figsize=(15, 6))
    gs = fig.add_gridspec(2, 3, height_ratios=[5, 1], hspace=0.1, wspace=0.1)

    # Old Image
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.imshow(disp_img1)
    ax1.set_title("Old Image", fontweight='bold')
    ax1.axis('off')

    # New Image
    ax2 = fig.add_subplot(gs[0, 1])
    ax2.imshow(disp_img2)
    ax2.set_title("New Image", fontweight='bold')
    ax2.axis('off')

    # Prediction
    ax3 = fig.add_subplot(gs[0, 2])
    ax3.imshow(pred_rgb)
    ax3.set_title("Prediction", fontweight='bold')
    ax3.axis('off')

    # Legend
    ax_lg = fig.add_subplot(gs[1, :])
    ax_lg.axis('off')
    ax_lg.legend(handles=make_legend_patches(), loc='center', ncol=7, fontsize=10, frameon=False)

    # Percentage Text
    fig.text(0.5, 0.05, f"Change Detected: {change_pct:.2f}%", 
             ha='center', fontsize=14, fontweight='bold', color='red')

    # 6. Save plot to buffer and return
    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches='tight', dpi=150)
    plt.close(fig)
    buf.seek(0)

    return Response(content=buf.getvalue(), media_type="image/png")