import os
import gc
import torch
import torch.nn as nn
import joblib
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel
from torch.utils.data import Dataset, DataLoader
from typing import Optional
import uvicorn
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# =========================
# CONFIG
# =========================
#MODEL_NAME = "BAAI/bge-small-en-v1.5"
MODEL_NAME = os.path.join(
    BASE_DIR,
    "BAAI_bge-small-en-v1.5_best_progressive"
)
MAX_LENGTH = 256
MAX_CHUNKS = 10
BATCH_SIZE = 8

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

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
LR_PIPELINE_PATH = os.path.join(BASE_DIR, "lr_assets", "tfidf_lr_pipeline.joblib")
CHECKPOINT_PATH = os.path.join(BASE_DIR, "models", "BAAI_bge-small-en-v1.5_progressive_checkpoint.pt")


# =========================
# MODEL (minimal safe load)
# =========================
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = AutoModel.from_pretrained(
            MODEL_NAME,
            local_files_only=True
        )

        H = self.encoder.config.hidden_size
        self.classifier = nn.Linear(H, 1)  

        
print("🚀 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    local_files_only=True
)

print("🚀 Loading model...")
model = SimpleModel().to(DEVICE)

print("🚀 Loading checkpoint...")
ckpt = torch.load(CHECKPOINT_PATH, map_location=DEVICE, weights_only=False)
state = ckpt.get("model_state_dict") or ckpt.get("best_overall_state")
model.load_state_dict(state, strict=False)

model.eval()

print("🚀 Loading LR...")
lr_model = joblib.load(LR_PIPELINE_PATH)


# =========================
# API
# =========================
app = FastAPI(title="Readmission API")


class Request(BaseModel):
    text: str


@app.get("/")
def home():
    return {"status": "running"}


@app.post("/predict")
def predict(req: Request):
    if not req.text:
        raise HTTPException(400, "empty text")

    inputs = tokenizer(
        req.text,
        truncation=True,
        padding=True,
        return_tensors="pt",
        max_length=MAX_LENGTH
    ).to(DEVICE)

    with torch.no_grad():
        logits = model(inputs["input_ids"], inputs["attention_mask"])
        bert_prob = torch.sigmoid(logits).item()

    lr_prob = lr_model.predict_proba([req.text])[0][1]

    final = 0.7 * bert_prob + 0.3 * lr_prob
    pred = int(final > 0.5)

    return {
        "probability": final,
        "prediction": pred,
        "bert_prob": bert_prob,
        "lr_prob": lr_prob
    }


# =========================
# RUN (IMPORTANT for HF Spaces)
# =========================
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)