import os
import sys

from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates

import torch
from torchvision import transforms
from PIL import Image


# ---------------- PATH SETUP ----------------

script_dir = os.path.dirname(
    os.path.abspath(__file__)
)

root_dir = os.path.dirname(script_dir)

if root_dir not in sys.path:
    sys.path.append(root_dir)


from model.alzheimers_model import AlzheimerNet



# ---------------- MODEL LOAD ----------------

model_path = os.path.join(
    root_dir,
    "saved_models",
    "alzheimer_model.pth"
)


if not os.path.exists(model_path):
    raise FileNotFoundError(model_path)



model = AlzheimerNet(
    num_classes=4,
    sophisticated=False
)


model.load_state_dict(
    torch.load(
        model_path,
        map_location="cpu"
    )
)

model.eval()



# ---------------- TRANSFORM ----------------

transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[
            0.485,
            0.456,
            0.406
        ],
        std=[
            0.229,
            0.224,
            0.225
        ]
    )
])


CLASSES = [
    "nondemented",
    "very mild",
    "mild demented",
    "moderate demented"
]



# ---------------- FASTAPI ----------------

app = FastAPI()



templates = Jinja2Templates(
    directory=os.path.join(
        os.path.dirname(__file__),
        "templates"
    )
)



@app.get(
    "/",
    response_class=HTMLResponse
)
async def home(request: Request):

    return templates.TemplateResponse(
        "index.html",
        {
            "request": request
        }
    )



@app.post("/predict")
async def predict(
    file: UploadFile = File(...)
):

    image = Image.open(
        file.file
    ).convert("RGB")


    tensor = transform(image)

    tensor = tensor.unsqueeze(0)


    with torch.no_grad():

        output = model(tensor)

        probs = torch.nn.functional.softmax(
            output,
            dim=1
        )[0]


    prediction = torch.argmax(
        probs
    ).item()



    confidence = {
        CLASSES[i]:
        round(
            float(probs[i])*100,
            2
        )

        for i in range(len(CLASSES))
    }



    return {
        "prediction": CLASSES[prediction],
        "confidence": confidence
    }