from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import pickle
from huggingface_hub import hf_hub_download
from typing import List, Optional

app = FastAPI()

# Download your model pickle from the Hub on startup
model_path = hf_hub_download(repo_id="Projects-by-IF/causal-model-Z15-v2", filename="trained_causal_model_v4.pkl")

with open(model_path, "rb") as f:
    model = pickle.load(f)

class InputData(BaseModel):
    X: List
    T0: Optional[List] = None
    T1: Optional[List] = None

@app.post("/model-effect")
def predict(data: InputData):
    X = np.array(data.X)

    if data.T0 is not None and data.T1 is not None:
        T0 = np.array(data.T0)
        T1 = np.array(data.T1)
        effect = model.effect(X=X, T0=T0, T1=T1).tolist()
    else:
        effect = model.effect(X).tolist()

    return {"effect": effect}