import json
import numpy as np
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer

# ============================================================
# CONFIG — ONLY CATEGORY MODELS
# ============================================================

HF_USER = "ClergeF"

CATEGORY_REPOS = {
    "family": "family-model",
    "community": "community-model",
    "education": "education-model",
    "health": "health-model",
    "environment": "environment-model",
    "business": "business-model",
    "finance": "finance-model",
    "history": "history-model",
    "spirituality": "spirituality-model",
    "innovation": "innovation-model",
}

CATEGORY_FILES = {
    "family": "family_level.json",
    "community": "community_level.json",
    "education": "education_level.json",
    "health": "health_level.json",
    "environment": "environment_level.json",
    "business": "business_level.json",
    "finance": "finance_level.json",
    "history": "history_level.json",
    "spirituality": "spirituality_level.json",
    "innovation": "innovation_level.json",
}

# ============================================================
# EMBEDDER
# ============================================================

print("Loading embedder: all-MiniLM-L6-v2 …")
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# ============================================================
# HELPERS
# ============================================================

def embed(text: str):
    return embedder.encode([text])[0]


def load_model(repo, filename):
    print(f"Loading: {repo}/{filename}")
    path = hf_hub_download(
        repo_id=f"{HF_USER}/{repo}",
        filename=filename
    )
    with open(path, "r") as f:
        return json.load(f)


def linear_predict(model_json, vec):
    coef = np.array(model_json["coef"])
    intercept = np.array(model_json["intercept"])
    return float(np.dot(coef, vec) + intercept)

# ============================================================
# LOAD 10 CATEGORY MODELS
# ============================================================

print("Loading category models…")
models = {}

for key in CATEGORY_REPOS:
    repo = CATEGORY_REPOS[key]
    file = CATEGORY_FILES[key]
    models[key] = load_model(repo, file)

print("✔ Category models loaded!")

# ============================================================
# API
# ============================================================

app = FastAPI(title="Category Classification API")


class InputText(BaseModel):
    text: str


@app.get("/")
def home():
    return {"status": "ok", "message": "Category API running"}


@app.post("/predict")
def predict(payload: InputText):
    text = payload.text
    vec = embed(text)

    out = {}
    for cat in CATEGORY_REPOS:
        out[f"{cat}_score"] = linear_predict(models[cat], vec)

    return {
        "input": text,
        "categories": out
    }