from typing import Union
from .predict import predict
from fastapi import FastAPI
import numpy as np
app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.get("/predict/{data}")
def read_item(data: str):
    final_json = {'probabilities':{}}

    try:
        res = predict(data)
        for i in range(5):
            final_json['probabilities'][f"india_{i}"] = round(float(res[i]),2)
        final_json['success'] = True
    except Exception as e:
        final_json['success'] = False
        final_json['error'] = repr(e)


    return final_json