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


# 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")
model_path = hf_hub_download(repo_id="DIGMMUNI/causal_model", filename="trained_causal_model_whole_city.pkl")
#model_path = hf_hub_download(repo_id="DIGMMUNI/causal_model", filename="trained_causal_model_v1.pkl")

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

# import joblib
# import os
# MODEL_DIR = r"./model"
# MODEL_FILE = "trained_causal_model_whole_city.pkl"
# #MODEL_FILE = "trained_causal_model_v1.pkl"
# model_path = os.path.join(MODEL_DIR, MODEL_FILE)

print("Loading model from {}".format(model_path))
app = FastAPI()

#model = joblib.load(model_path)

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}