import torch
from transformers import BertForSequenceClassification, BertTokenizer
from huggingface_hub import hf_hub_download
import numpy as np
from sklearn.preprocessing import LabelEncoder
import re
from fastapi import FastAPI
from pydantic import BaseModel

REPO_ID = "waseem11/severity_model"
device = torch.device("cpu")

app = FastAPI()

# ------------------------
# Preprocessing
# ------------------------
def clean_text(text):
    text = str(text).replace('\n', ' ').replace('\r', ' ')
    return re.sub(r'\s+', ' ', text).strip()

def normalize_urdu(text):
    text = re.sub(r'[آأإا]', 'ا', text)
    text = re.sub(r'[يى]', 'ی', text)
    text = re.sub(r'[ةۀ]', 'ہ', text)
    text = re.sub(r'[\u064B-\u0652]', '', text)
    return text

def preprocess(text):
    return normalize_urdu(clean_text(text))

def map_severity(label, confidence):
    if label == 'None':
        return 'Mild'
    elif label == 'Mild':
        return 'Moderate' if confidence >= 0.5 else 'Mild'
    elif label == 'Severe':
        return 'Severe'
    return 'Mild'

# ------------------------
# Load Models (startup par sirf ek dafa load hon ge)
# ------------------------
print("Loading tokenizer...")
tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")

sev_encoder = LabelEncoder()
sev_encoder.classes_ = np.array(['None', 'Mild', 'Severe'])

bin_encoder = LabelEncoder()
bin_encoder.classes_ = np.array(['Not Depressed', 'Depressed'])

print("Loading binary model...")
binary_path = hf_hub_download(REPO_ID, "binary_model.pt")
binary_model = BertForSequenceClassification.from_pretrained(
    "bert-base-multilingual-cased", num_labels=2)
binary_model.load_state_dict(torch.load(binary_path, map_location=device), strict=False)
binary_model.eval()

print("Loading severity model...")
severity_path = hf_hub_download(REPO_ID, "severity_model.pt")
severity_model = BertForSequenceClassification.from_pretrained(
    "bert-base-multilingual-cased", num_labels=3)
severity_model.load_state_dict(torch.load(severity_path, map_location=device), strict=False)
severity_model.eval()

print("All models ready ✓")

# ------------------------
# Request Body
# ------------------------
class TextInput(BaseModel):
    text: str

# ------------------------
# API Endpoint
# ------------------------
@app.post("/predict")
def predict_api(data: TextInput):

    text = data.text

    if not text.strip():
        return {"error": "Text khali hai"}

    processed = preprocess(text)
    inputs = tokenizer(processed, return_tensors="pt",
                       truncation=True, padding=True, max_length=128)

    with torch.no_grad():
        bin_out = binary_model(**inputs)
        bin_pred = torch.argmax(bin_out.logits, dim=1).item()
        bin_conf = torch.softmax(bin_out.logits, dim=1).max().item()
        bin_label = bin_encoder.inverse_transform([bin_pred])[0]

    if bin_label == 'Not Depressed':
        return {
            "status": "Not Depressed",
            "confidence": f"{bin_conf:.1%}",
            "severity": "None"
        }

    with torch.no_grad():
        sev_out = severity_model(**inputs)
        sev_probs = torch.softmax(sev_out.logits, dim=1)
        sev_pred = torch.argmax(sev_probs, dim=1).item()
        sev_conf = sev_probs[0][sev_pred].item()
        sev_label = sev_encoder.inverse_transform([sev_pred])[0]
        sev_mapped = map_severity(sev_label, sev_conf)

    return {
        "status": "Depressed",
        "confidence": f"{bin_conf:.1%}",
        "severity": sev_mapped,
        "severity_confidence": f"{sev_conf:.1%}"
    }