import base64
import io
import math
import os
import time

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def decorator(func):
                return func

            return decorator

    spaces = _SpacesFallback()

import gradio as gr
import numpy as np
import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
from transformers import ViTForImageClassification, ViTImageProcessor

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_model_safely():
    """Load model with fallback options and proper error handling.

    Returns (processor, model, model_id) — model_id names the source that
    actually loaded, so the analysis panel reports the real weights in use.
    """
    try:
        if os.path.exists("./waste-classification"):
            print("Loading model from local clone...")
            processor = ViTImageProcessor.from_pretrained("./waste-classification")
            model = ViTForImageClassification.from_pretrained(
                "./waste-classification", attn_implementation="eager"
            )
            print("Successfully loaded model from local clone")
            return processor, model, "local: ./waste-classification"
    except Exception as e:
        print(f"Failed to load from local clone: {e}")

    try:
        print("Loading model from HuggingFace...")
        processor = ViTImageProcessor.from_pretrained(
            "watersplash/waste-classification", cache_dir="./cache"
        )
        # eager attention so /api/explain can read attention weights (sdpa cannot)
        model = ViTForImageClassification.from_pretrained(
            "watersplash/waste-classification", cache_dir="./cache", attn_implementation="eager"
        )
        print("Successfully loaded model from HuggingFace")
        return processor, model, "watersplash/waste-classification"
    except Exception as e:
        print(f"Failed to load from HuggingFace: {e}")

    try:
        print("Loading base ViT model as fallback...")
        processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
        model = ViTForImageClassification.from_pretrained(
            "google/vit-base-patch16-224-in21k",
            attn_implementation="eager",
            num_labels=12,
            id2label={
                "0": "battery", "1": "biological", "2": "brown-glass", "3": "cardboard",
                "4": "clothes", "5": "green-glass", "6": "metal", "7": "paper",
                "8": "plastic", "9": "shoes", "10": "trash", "11": "white-glass",
            },
            label2id={
                "battery": "0", "biological": "1", "brown-glass": "2", "cardboard": "3",
                "clothes": "4", "green-glass": "5", "metal": "6", "paper": "7",
                "plastic": "8", "shoes": "9", "trash": "10", "white-glass": "11",
            },
        )
        print("Loaded base ViT model as fallback (untrained)")
        return processor, model, "google/vit-base-patch16-224-in21k (untrained fallback)"
    except Exception as e:
        print(f"Failed to load fallback model: {e}")
        return None, None, "none"


print("Initializing model...")
processor, model, MODEL_ID = load_model_safely()
if model is not None:
    model.to(DEVICE)
    model.eval()

CLASS_NAMES = [
    "Battery", "Biological", "Brown-glass", "Cardboard", "Clothes",
    "Green-glass", "Metal", "Paper", "Plastic", "Shoes", "Trash", "White-glass",
]

# Waste segregation mapping: which bin each class belongs to
BIN_MAP = {
    "Biological": "green",
    "Cardboard": "blue", "Paper": "blue", "Plastic": "blue", "Metal": "blue",
    "Brown-glass": "blue", "Green-glass": "blue", "White-glass": "blue",
    "Clothes": "blue", "Shoes": "blue",
    "Battery": "red", "Trash": "red",
}

BIN_INFO = {
    "green": {"name": "Green Bin", "desc": "Wet / Biodegradable Waste"},
    "blue": {"name": "Blue Bin", "desc": "Dry / Recyclable Waste"},
    "red": {"name": "Red Bin", "desc": "Hazardous / Reject Waste"},
}


def _classify_waste_impl(image: Image.Image):
    """Classify a PIL image, return (label, confidence, bin_color, top, analysis)"""
    if processor is None or model is None:
        raise RuntimeError("Model failed to load")

    if image.mode != "RGB":
        image = image.convert("RGB")

    original_size = list(image.size)

    started = time.perf_counter()
    inputs = processor(images=image, return_tensors="pt")
    inputs = {name: tensor.to(DEVICE) for name, tensor in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    latency_ms = (time.perf_counter() - started) * 1000

    confidence_scores = predictions[0].tolist()
    ranked = sorted(range(len(confidence_scores)), key=lambda i: confidence_scores[i], reverse=True)

    best_idx = ranked[0]
    label = CLASS_NAMES[best_idx]
    confidence = confidence_scores[best_idx]
    bin_color = BIN_MAP.get(label, "red")
    top = [
        {"label": CLASS_NAMES[i], "confidence": round(confidence_scores[i], 4)}
        for i in ranked[:3]
    ]

    # Shannon entropy in bits; normalised against log2(n) so 0 = certain, 1 = uniform.
    entropy = -sum(p * math.log2(p) for p in confidence_scores if p > 0)
    max_entropy = math.log2(len(confidence_scores))
    margin = confidence_scores[ranked[0]] - confidence_scores[ranked[1]]

    analysis = {
        "distribution": [
            {"label": CLASS_NAMES[i], "confidence": round(confidence_scores[i], 6)}
            for i in ranked
        ],
        "entropy_bits": round(entropy, 4),
        "entropy_normalized": round(entropy / max_entropy, 4) if max_entropy else 0.0,
        "margin": round(margin, 4),
        "logit_top": round(outputs.logits[0][best_idx].item(), 4),
        "latency_ms": round(latency_ms, 1),
        "device": str(next(model.parameters()).device),
        "model": {
            "id": MODEL_ID,
            "architecture": model.config.architectures[0] if model.config.architectures else "ViTForImageClassification",
            "parameters": sum(p.numel() for p in model.parameters()),
            "input_size": [model.config.image_size, model.config.image_size],
            "patch_size": model.config.patch_size,
            "hidden_size": model.config.hidden_size,
            "num_layers": model.config.num_hidden_layers,
            "num_heads": model.config.num_attention_heads,
            "num_classes": len(CLASS_NAMES),
            "original_image_size": original_size,
            "normalize_mean": list(processor.image_mean),
            "normalize_std": list(processor.image_std),
        },
    }

    return label, confidence, bin_color, top, analysis


@spaces.GPU(duration=60)
def classify_waste(image: Image.Image):
    return _classify_waste_impl(image)


# --- FastAPI app: serves the vanilla JS frontend + classification API ---
app = FastAPI(title="Waste Classification")


@app.post("/api/classify")
async def api_classify(file: UploadFile = File(...)):
    if processor is None or model is None:
        raise HTTPException(status_code=503, detail="Model failed to load. Please try again later.")

    try:
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes))
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file.")

    try:
        label, confidence, bin_color, top, analysis = classify_waste(image)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Classification failed: {e}")

    info = BIN_INFO[bin_color]
    return JSONResponse({
        "label": label,
        "confidence": round(confidence, 4),
        "bin_color": bin_color,
        "bin_name": info["name"],
        "bin_desc": info["desc"],
        "top": top,
        "analysis": analysis,
    })


def _colorize(heat):
    """Map a 0..1 attention map to a blue→cyan→green→yellow→red ramp."""
    stops = np.array([
        [59, 76, 192], [77, 165, 220], [92, 200, 160],
        [237, 205, 74], [214, 86, 51],
    ], dtype=np.float32)

    pos = np.clip(heat, 0.0, 1.0) * (len(stops) - 1)
    lo = np.floor(pos).astype(int)
    hi = np.clip(lo + 1, 0, len(stops) - 1)
    t = (pos - lo)[..., None]

    return (stops[lo] * (1 - t) + stops[hi] * t).astype(np.uint8)


@spaces.GPU(duration=60)
def attention_rollout(image: Image.Image):
    """Overlay ViT attention rollout on the image; returns a PNG data URL.

    Rollout follows Abnar & Zuidema (2020): average heads per layer, add the
    residual connection as an identity, renormalise, then multiply across
    layers to trace CLS attention back to the input patches.
    """
    if processor is None or model is None:
        raise RuntimeError("Model failed to load")

    if image.mode != "RGB":
        image = image.convert("RGB")

    started = time.perf_counter()
    inputs = processor(images=image, return_tensors="pt")
    inputs = {name: tensor.to(DEVICE) for name, tensor in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)

    rollout = None
    for layer in outputs.attentions:
        attn = layer[0].mean(dim=0)
        attn = attn + torch.eye(attn.size(0), device=attn.device)
        attn = attn / attn.sum(dim=-1, keepdim=True)
        rollout = attn if rollout is None else attn @ rollout

    grid = model.config.image_size // model.config.patch_size
    heat = rollout[0, 1:].reshape(grid, grid).detach().cpu().numpy()

    # Rollout is diffuse by nature; stretch the informative band so the
    # attended region reads clearly instead of a flat wash.
    lo, hi = np.percentile(heat, 55), np.percentile(heat, 99.5)
    heat = np.clip((heat - lo) / (hi - lo + 1e-8), 0.0, 1.0) ** 1.3

    # Render at most 512px on the long side to keep the payload small.
    preview = image.copy()
    preview.thumbnail((512, 512), Image.LANCZOS)

    heat_img = Image.fromarray((heat * 255).astype(np.uint8)).resize(preview.size, Image.BICUBIC)
    heat_arr = np.asarray(heat_img, dtype=np.float32) / 255.0

    base = np.asarray(preview, dtype=np.float32)
    alpha = (0.72 * heat_arr)[..., None]
    blended = base * (1 - alpha) + _colorize(heat_arr).astype(np.float32) * alpha

    buffer = io.BytesIO()
    Image.fromarray(blended.astype(np.uint8)).save(buffer, format="PNG", optimize=True)
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")

    return {
        "heatmap": f"data:image/png;base64,{encoded}",
        "grid": [grid, grid],
        "patches": grid * grid,
        "layers": len(outputs.attentions),
        "latency_ms": round((time.perf_counter() - started) * 1000, 1),
    }


@app.post("/api/explain")
async def api_explain(file: UploadFile = File(...)):
    if processor is None or model is None:
        raise HTTPException(status_code=503, detail="Model failed to load. Please try again later.")

    try:
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes))
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file.")

    try:
        return JSONResponse(attention_rollout(image))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Explanation failed: {e}")


@app.get("/api/health")
async def health():
    return {"status": "ok" if (model is not None and processor is not None) else "model_not_loaded"}


# --- Minimal Gradio Blocks, mounted at /gradio purely to satisfy the Gradio SDK ---
@spaces.GPU(duration=60)
def _gradio_classify(image):
    if image is None:
        return {"Please upload an image": 1.0}
    try:
        label, confidence, bin_color, _top, _analysis = classify_waste(image)
        info = BIN_INFO[bin_color]
        return {f"{label} → {info['name']}": confidence}
    except Exception as e:
        return {f"Error: {e}": 1.0}


with gr.Blocks(title="Waste Classification (debug UI)") as demo:
    gr.Markdown("### Debug UI — the real app lives at the Space root URL")
    gr.Interface(
        fn=_gradio_classify,
        inputs=gr.Image(type="pil", label="Upload Waste Image"),
        outputs=gr.Textbox(label="Result"),
    )

# Serve the vanilla JS frontend at the root when the FastAPI app is used outside ZeroGPU.
app.mount("/", StaticFiles(directory="static", html=True), name="static")


@spaces.GPU(duration=60)
def _zerogpu_classify(image):
    if image is None:
        return "Please upload an image."

    label, confidence, bin_color, _top, _analysis = _classify_waste_impl(image)
    info = BIN_INFO[bin_color]
    return (
        f"Prediction: {label}\n"
        f"Confidence: {confidence:.1%}\n"
        f"Bin: {info['name']}\n"
        f"Description: {info['desc']}"
    )


with gr.Blocks(title="Waste Classification") as zerogpu_demo:
    gr.Markdown("### Waste Classification")
    zerogpu_image = gr.Image(type="pil", label="Upload Waste Image")
    zerogpu_button = gr.Button("Classify")
    zerogpu_result = gr.Textbox(label="Result")

    zerogpu_button.click(
        fn=_zerogpu_classify,
        inputs=zerogpu_image,
        outputs=zerogpu_result,
        api_name="classify",
    )


if __name__ == "__main__":
    zerogpu_demo.launch(server_name="0.0.0.0", server_port=7860)