from fastapi import FastAPI
from pydantic import BaseModel
from typing import Dict
import joblib
import numpy as np

app = FastAPI()

# Define the input schema
class UsageInput(BaseModel):
    days_until_cycle_end: int
    voice_total_allowance: float
    voice_remaining: float
    data_total_allowance: float
    data_remaining: float
    plan_price: float

# Load the trained model
model = joblib.load("model.pkl")

@app.post("/predict")
def predict(input_data: UsageInput) -> Dict[str, float]:
    # Convert input to model format
    input_array = np.array([[
        input_data.days_until_cycle_end,
        input_data.voice_total_allowance,
        input_data.voice_remaining,
        input_data.data_total_allowance,
        input_data.data_remaining,
        input_data.plan_price
    ]])

    # Predict using the model
    predicted_voice, predicted_data = model.predict(input_array)[0]

    return {
        "predicted_total_voice_usage": round(predicted_voice, 2),
        "predicted_total_data_usage": round(predicted_data, 2)
    }