import uvicorn
import pickle
import os
import json
import logging
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from typing import List, Literal, Optional
from pydantic import BaseModel
import pandas as pd


# logger
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.DEBUG)


# Util Functions & Classes
def loading(fp):
    with open(fp, "rb") as f:
        data = pickle.load(f)

    print(f"INFO: Loaded data : {data}")
    return data


def predict(df, endpoint="simple"):
    """Take a dataframe as input and use it to make predictions"""

    print(
        f"[Info] 'predict' function has been called through the endpoint '{endpoint}'.\n"
    )

    logging.info(f" \n{df.to_markdown()}")

    # scaling
    scaled_df = scaler.transform(df)
    logging.info(f"     Scaler output is of type {type(scaled_df)}")

    # prediction
    prediction = model.predict_proba(scaled_df)
    print(f"INFO: Prediction output: {prediction}")

    # Formatting of the prediction
    ## extract highest proba
    highest_proba = prediction.max(axis=1)
    print(f"INFO: Highest probabilities : {highest_proba}")

    ## extract indexes of the highest proba
    highest_proba_idx = prediction.argmax(axis=1)
    print(f"INFO: Highest probability indexes : {highest_proba_idx}")

    ## Maching prediction with classes
    predicted_classes = [labels[i] for i in highest_proba_idx]
    print(f"INFO: Predicted classes : {predicted_classes}")
    # prediction[:, highest_proba_idx]

    # save in df
    df["predicted proba"] = highest_proba
    df["predicted label"] = predicted_classes

    print(f"INFO: dataframe filled with prediction\n{df.to_markdown()}\n")

    # parsing prediction
    # parsed = json.loads(df.to_json(orient="index")) # or
    parsed = df.to_dict("records")

    return parsed


## INPUT MODELING
class Land(BaseModel):
    """Modeling of one input data in a type-restricted dictionary-like format

    column_name : variable type # strictly respect the name in the dataframe header.

    eg.:
    =========
    customer_age : int
    gender : Literal['male', 'female', 'other']
    """

    N: float
    P: float
    K: float
    temperature: float
    humidity: float
    ph: float
    rainfall: float


class Lands(BaseModel):
    inputs: List[Land]

    def return_list_of_dict(
        cls,
    ):
        # return [land.dict() for land in cls.inputs]
        return [i.dict() for i in cls.inputs]


# API Config
app = FastAPI(
    title="Agri-Tech API",
    description="This is a ML API for classification of crop to plant on a land regarding some features",
)

## Configure static and template files
app.mount(
    "/static", StaticFiles(directory="assets/static"), name="static"
)  # Mount static files
templates = Jinja2Templates(directory="assets/templates")  # Mount templates for HTML


# ML Config
ml_objects = loading(fp=os.path.join("assets", "ml", "crop_recommandation2.pkl"))
## Extract the ml components
model = ml_objects["model"]
scaler = ml_objects["scaler"].set_output(transform="pandas")
labels = ml_objects["labels"]


# Endpoints
# @app.get("/")
# def root():
#     return {
#         "Description": " This is a ML API for classification of crop to plant on a land regarding some features.",
#         "Documentation": "Go to the docs: https://eaedk-agri-tech-fastapi.hf.space/docs",
#     }


# Root endpoint to serve index.html template
@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
    return templates.TemplateResponse("index.html", {'request': request})


@app.get("/checkup")
def test(a: Optional[int], b: int):
    return {"a": a, "b": b}


## ML endpoint
@app.post("/predict")
def make_prediction(
    N: float,
    P: float,
    K: float,
    temperature: float,
    humidity: float,
    ph: float,
    rainfall: float,
):
    """Make prediction with the passed data"""

    df = pd.DataFrame(
        {
            "N": [N],
            "P": [P],
            "K": [K],
            "temperature": [temperature],
            "humidity": [humidity],
            "ph": [ph],
            "rainfall": [rainfall],
        }
    )

    parsed = predict(df=df)  # df.to_dict('records')

    return {
        "output": parsed,
    }


@app.post("/predict_multi")
def make_multi_prediction(multi_lands: Lands):
    """Make prediction with the passed data"""
    print(f"Mutiple inputs passed: {multi_lands}\n")
    df = pd.DataFrame(multi_lands.return_list_of_dict())

    parsed = predict(df=df, endpoint="multi inputs")  # df.to_dict('records')

    return {
        "output": parsed,
        "author": "Stella Archar",
        "api_version": ";)",
    }


if __name__ == "__main__":
    uvicorn.run("main:app", reload=True)