"""
Hugging Face Space backend for the RadAI WM-811K wafer defect model.

This is the ONLY place torch/torchvision run. Your local machine calls this
Space over HTTP instead of loading the model itself, so the local venv no
longer needs torch/torchvision (the bulk of the disk usage).

Model: https://huggingface.co/radai-agent/radai-wm811k-defect-detection

Endpoints:
    GET  /                -> health check
    POST /predict_map     -> {"wafer_map": <2D list, any size, any grayscale/bin values>}
                              returns {"predictions": [...], "preview_raw": <64x64 list, 0-1>}
"""

import os

import numpy as np
import torch
import torch.nn as nn
from flask import Flask, jsonify, request
from flask_cors import CORS
from huggingface_hub import hf_hub_download
from scipy.ndimage import zoom
from torchvision import models

REPO_ID = "radai-agent/radai-wm811k-defect-detection"
MODEL_PATH = os.path.join(os.path.dirname(__file__), "best_radai_resnet.pt")

CLASSES = [
    "Center", "Donut", "Edge-Loc", "Edge-Ring",
    "Loc", "Random", "Scratch", "Near-full",
]

CLASS_INFO = {
    "Center":    "缺陷集中在晶圓中心區域",
    "Donut":     "缺陷呈環狀分佈，中心與邊緣正常",
    "Edge-Loc":  "缺陷集中在邊緣的局部區域",
    "Edge-Ring": "缺陷沿整個邊緣呈環狀分佈",
    "Loc":       "缺陷集中在某個局部區域（非邊緣）",
    "Random":    "缺陷隨機散佈，無明顯規律",
    "Scratch":   "缺陷呈線狀，類似刮痕",
    "Near-full": "近乎整片晶圓都被判定為缺陷",
}


class RadAI_ResNet(nn.Module):
    """Architecture must match the checkpoint exactly (see model card)."""

    def __init__(self, num_classes=8):
        super().__init__()
        self.base = models.resnet34(weights=None)
        self.base.conv1 = nn.Conv2d(
            1, 64, kernel_size=7, stride=2, padding=3, bias=False
        )
        self.base.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(self.base.fc.in_features, num_classes),
        )

    def forward(self, x):
        return self.base(x)


app = Flask(__name__)
# Allow your local page / local Flask app to call this Space from a
# different origin. Tighten origins= to your own domain(s) if you want.
CORS(app)

_model = None


def ensure_weights():
    if os.path.exists(MODEL_PATH):
        return
    from huggingface_hub import list_repo_files
    files = list_repo_files(REPO_ID)
    candidates = [f for f in files if f.endswith(".pt") or f.endswith(".pth")]
    filename = next((f for f in candidates if "best_radai_resnet" in f), None) \
        or (candidates[0] if candidates else None)
    if not filename:
        raise RuntimeError(f"在 {REPO_ID} 找不到 .pt/.pth 權重檔")
    path = hf_hub_download(repo_id=REPO_ID, filename=filename)
    if os.path.abspath(path) != os.path.abspath(MODEL_PATH):
        import shutil
        shutil.copy(path, MODEL_PATH)


def get_model():
    global _model
    if _model is None:
        ensure_weights()
        model = RadAI_ResNet(num_classes=8)
        checkpoint = torch.load(MODEL_PATH, map_location="cpu")
        state_dict = checkpoint.get("model_state_dict", checkpoint)
        model.load_state_dict(state_dict)
        model.eval()
        _model = model
    return _model


@app.route("/", methods=["GET"])
def health():
    return jsonify({"status": "ok", "model": REPO_ID})


@app.route("/predict_map", methods=["POST"])
def predict_map():
    data = request.get_json(force=True, silent=True) or {}
    raw = data.get("wafer_map")
    if raw is None:
        return jsonify({"error": "缺少 wafer_map"}), 400

    try:
        wafer_map = np.array(raw, dtype=np.float32)
        if wafer_map.ndim != 2 or wafer_map.size == 0:
            raise ValueError("wafer_map 必須是非空的 2D 陣列")
    except Exception as e:
        return jsonify({"error": f"wafer_map 格式錯誤：{e}"}), 400

    h, w = wafer_map.shape
    resized = zoom(wafer_map, (64 / h, 64 / w), order=1)[:64, :64]
    if resized.max() > 0:
        resized = resized / resized.max()

    tensor = torch.FloatTensor(resized).unsqueeze(0).unsqueeze(0)

    try:
        model = get_model()
    except Exception as e:
        return jsonify({"error": f"模型載入失敗：{e}"}), 500

    with torch.no_grad():
        output = model(tensor)
        probs = torch.softmax(output, dim=1)[0].numpy()

    ranked = sorted(
        (
            {"label": CLASSES[i], "desc": CLASS_INFO[CLASSES[i]], "prob": float(probs[i])}
            for i in range(len(CLASSES))
        ),
        key=lambda item: item["prob"],
        reverse=True,
    )

    return jsonify({"predictions": ranked, "preview_raw": resized.tolist()})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)