from fastapi import FastAPI
import joblib
import pandas as pd
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
import onnxruntime as ort
import numpy as np

# Initialize FastAPI app
app = FastAPI()

# Download the preprocessor and model from Hugging Face
preprocessor_path = hf_hub_download(
    repo_id="sravan837/slope-stability-api",
    filename="slope_stability_preprocessor.pkl",
)

onnx_model_path = hf_hub_download(
    repo_id="sravan837/slope-stability-api",
    filename="slope_stability_model.onnx",
)

# Load preprocessor
preprocessor = joblib.load(preprocessor_path)

# Load ONNX model for inference
onnx_session = ort.InferenceSession(onnx_model_path)

# Define input data structure
class SlopeStabilityInput(BaseModel):
    unit_weight: float
    cohesion: float
    friction_angle: float
    slope_angle: float
    slope_height: float
    water_pressure_ratio: float
    reinforcement_type: str

# API Endpoint for prediction
@app.post("/predict")
def predict_slope_stability(data: SlopeStabilityInput):
    # Convert input to DataFrame
    input_data = pd.DataFrame([data.dict()])

    # Transform input using preprocessor
    transformed_input = preprocessor.transform(input_data)

    # Convert to ONNX format
    input_array = np.array(transformed_input, dtype=np.float32)

    # Check ONNX model input name
    input_name = onnx_session.get_inputs()[0].name

    # Make prediction
    prediction = onnx_session.run(None, {input_name: input_array})[0][0]

    return {"Factor of Safety": float(prediction)}