from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import json
import pandas as pd

app = FastAPI(title="Game of Thrones House Predictor")

model = joblib.load("model.pkl")

with open("feature_columns.json", "r") as f:
    feature_columns = json.load(f)["columns"]

with open("label_classes.json", "r") as f:
    label_classes = json.load(f)["classes"]


class CharacterInput(BaseModel):
    region: str
    primary_role: str
    alignment: str
    status: str
    species: str
    honour_1to5: int
    ruthlessness_1to5: int
    intelligence_1to5: int
    combat_skill_1to5: int
    diplomacy_1to5: int
    leadership_1to5: int
    trait_loyal: bool
    trait_scheming: bool


class PredictionOutput(BaseModel):
    predicted_house: str


@app.get("/")
def root():
    return {"message": "Game of Thrones House Predictor API", "docs": "/docs"}


@app.post("/predict", response_model=PredictionOutput)
def predict(character: CharacterInput):
    input_dict = {
        "region": character.region,
        "primary_role": character.primary_role,
        "alignment": character.alignment,
        "status": character.status,
        "species": character.species,
        "honour_1to5": character.honour_1to5,
        "ruthlessness_1to5": character.ruthlessness_1to5,
        "intelligence_1to5": character.intelligence_1to5,
        "combat_skill_1to5": character.combat_skill_1to5,
        "diplomacy_1to5": character.diplomacy_1to5,
        "leadership_1to5": character.leadership_1to5,
        "trait_loyal": int(character.trait_loyal),
        "trait_scheming": int(character.trait_scheming),
    }

    df = pd.DataFrame([input_dict])
    df_encoded = pd.get_dummies(df, dummy_na=True)

    for col in feature_columns:
        if col not in df_encoded.columns:
            df_encoded[col] = 0

    df_encoded = df_encoded[feature_columns]
    prediction = model.predict(df_encoded)[0]

    return PredictionOutput(predicted_house=prediction)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)