# ==============================
# IMPORTS
# ==============================
import torch
import re
import time
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import zipfile
import os

# ==============================
# UNZIP MODEL
# ==============================
zip_path = "./1epoch-cds-mistral4bit-training.zip"
extract_path = "./1epoch-cds-mistral4bit-training"

if not os.path.exists(extract_path):
    print("📦 Extracting model...")
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_path)
    print("✅ Done")
else:
    print("✅ Model already extracted")

# ==============================
# GPU STATUS FUNCTION (SAFE)
# ==============================
def print_gpu():
    try:
        os.system("nvidia-smi")
    except:
        print("No GPU (running on CPU)")

# ==============================
# LOAD MODEL
# ==============================
BASE_MODEL = "mistralai/Mistral-7B-v0.1"
MODEL_PATH = "./1epoch-cds-mistral4bit-training/mistral7b_fast/checkpoint-7125"

print("🚀 Loading model...")

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# CPU SAFE LOADING
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map={"": "cpu"}
)

model = PeftModel.from_pretrained(base_model, MODEL_PATH)
model.eval()

# ==============================
# CLEAN OUTPUT
# ==============================
def clean_output(text):
    text = text.strip()
    text = re.sub(r"\n+", "\n", text)
    if "Final Decision:" in text:
        before, after = text.split("Final Decision:", 1)
        decision_line = after.split("\n")[0]
        return before.strip() + "\n\nFinal Decision: " + decision_line.strip()
    return text

# ==============================
# LABEL EXTRACTION
# ==============================
def extract_label(text):
    text = text.lower()
    if "final decision" in text:
        decision_part = text.split("final decision")[-1]
        if "not hate" in decision_part:
            return "non_hate"
        elif "hate" in decision_part:
            return "hate"
    return "unknown"

# ==============================
# METRICS
# ==============================
def compute_metrics(output, post):
    out = output.lower()
    post = post.lower()

    htc = 1 if "final decision" in out else 0

    post_words = post.split()
    qf = 1 if any(word in out for word in post_words[:5]) else 0

    tgi_keywords = ["muslim","black","white","asian","women","men","jews","christian","pakistan","indian"]
    tgi = 1 if any(word in out for word in tgi_keywords) else 0

    pred = extract_label(out)

    if pred == "hate" and "hate" in out:
        cons = 1
    elif pred == "non_hate" and "not hate" in out:
        cons = 1
    else:
        cons = 0

    return htc, qf, tgi, cons

# ==============================
# INFERENCE FUNCTION
# ==============================
def infer(post):
    print("\n📊 STATUS BEFORE:")
    print_gpu()

    inputs = tokenizer(post, return_tensors="pt")
    token_count = inputs.input_ids.shape[1]

    start = time.time()

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,   # SAME as your original
            min_new_tokens=150,
            do_sample=False,
            repetition_penalty=1.15,
            pad_token_id=tokenizer.eos_token_id
        )

    end = time.time()

    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = clean_output(output.replace(post, "").strip())

    pred = extract_label(response)
    htc, qf, tgi, cons = compute_metrics(response, post)

    metrics = f"""Prediction: {pred}
HTC={htc}, QF={qf}, TGI={tgi}, Consistency={cons}
Tokens processed: {token_count}
Time: {end-start:.2f} sec"""

    print("\n📊 STATUS AFTER:")
    print_gpu()

    return response, metrics

# ==============================
# GRADIO UI
# ==============================
iface = gr.Interface(
    fn=infer,
    inputs=gr.Textbox(lines=3, placeholder="Enter a post here..."),
    outputs=[
        gr.Textbox(label="Model Output (Rationales + Decision)", lines=15),
        gr.Textbox(label="Metrics", lines=5)
    ],
    title="Hate Speech Rationales + Decision",
    description="Enter a post to get step-by-step reasoning + final decision"
)

# ==============================
# LAUNCH
# ==============================
iface.launch()