"""Local FP32 Gradio demo; browser quantization is a separate inference branch."""
import argparse
import json
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
THRESHOLD = 0.5
MODEL_PATH = None
tokenizer = None
model = None
_settings = None


def resolve_model_settings(model_choice="selected", checkpoint=None):
    """Resolve paths without importing torch or loading weights."""
    if model_choice not in ("selected", "original"):
        raise ValueError("model_choice must be selected or original")
    if model_choice == "original" and checkpoint is not None:
        raise ValueError("Use either --model original or --checkpoint")
    explicit = checkpoint if checkpoint is not None else (os.environ.get("MODEL_PATH") if model_choice == "selected" else None)
    if explicit:
        path = Path(explicit)
        settings = dict(display_name="Custom checkpoint", model_id="explicit local checkpoint", mode="custom_fp32")
    elif model_choice == "original":
        path = ROOT / "models/disaster-bert"
        settings = dict(display_name="BERT Mini (original seed 42)", model_id="google/bert_uncased_L-4_H-256_A-4", mode="original_fp32")
    else:
        metadata = ROOT / "results/deployment/selected_model.json"
        if not metadata.is_file():
            raise FileNotFoundError("Missing selected_model.json; use --model original for the original trained Mini")
        selection = json.loads(metadata.read_text(encoding="utf-8"))
        path = Path(selection["checkpoint"])
        settings = dict(display_name=selection.get("display_name", "Selected BERT"), mode="selected_fp32",
                        model_id=selection.get("decision", {}).get("selected_model_id", selection.get("model_id", "selected local checkpoint")))
    path = (path if path.is_absolute() else ROOT / path).resolve()
    if not (path / "model.safetensors").is_file() and not (path / "pytorch_model.bin").is_file():
        raise FileNotFoundError(f"Trained local checkpoint not found: {path}")
    return {**settings, "path": path, "threshold": THRESHOLD, "precision": "FP32", "device": "cpu"}


def load_classifier(model_choice="selected", checkpoint=None):
    global MODEL_PATH, tokenizer, model, _settings
    settings = resolve_model_settings(model_choice, checkpoint)
    if model is None or MODEL_PATH != settings["path"]:
        import torch
        from transformers import AutoTokenizer, AutoModelForSequenceClassification
        torch.set_num_threads(4)
        tokenizer = AutoTokenizer.from_pretrained(settings["path"], local_files_only=True)
        model = AutoModelForSequenceClassification.from_pretrained(settings["path"], local_files_only=True,
            torch_dtype=torch.float32, attn_implementation="eager").cpu().eval()
        MODEL_PATH = settings["path"]
    _settings = settings
    return settings


def classify(text):
    if not text.strip():
        raise ValueError("请输入英文推文。")
    if model is None:
        load_classifier()
    import torch
    with torch.inference_mode():
        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="max_length", max_length=96)
        p = model(**inputs).logits.softmax(-1)[0].tolist()
    # Disaster first gives the positive class when both probabilities equal 0.5.
    return {"灾害 / Disaster": p[1], "非灾害 / Not disaster": p[0]}


def build_demo(settings=None):
    import gradio as gr
    settings = settings or _settings or resolve_model_settings()
    return gr.Interface(fn=classify, inputs=gr.Textbox(lines=4, label="英文推文"), outputs=gr.Label(label="分类及概率"),
        title=f"灾害推文识别 · {settings['display_name']} · 本地 FP32",
        description=("本地使用 FP32 权重，灾害概率 ≥ 0.50 判为灾害。浏览器页面使用另一条量化模型 / WASM 推理分支，"
                     "并在真实浏览器验证集上单独选择阈值；其概率和指标与此页不同。此页固定使用 0.50。"
                     "概率未经校准，结果仅用于课程实验。"),
        examples=[["A strong earthquake damaged several buildings downtown."], ["This movie is on fire! I loved it."]],
        analytics_enabled=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", choices=["selected", "original"], default="selected")
    parser.add_argument("--checkpoint", help="Explicit local FP32 checkpoint, relative to the assignment root")
    parser.add_argument("--show-config", action="store_true", help="Resolve selection without loading weights or starting a server")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()
    settings = resolve_model_settings(args.model, args.checkpoint)
    if args.show_config:
        print(json.dumps({**settings, "path": str(settings["path"])}, ensure_ascii=False, indent=2))
    else:
        load_classifier(args.model, args.checkpoint)
        build_demo(settings).launch(server_name="127.0.0.1",server_port=args.port,share=False,inbrowser=False)