import base64
from io import BytesIO
import cv2
from fastapi import FastAPI, Request
from keras.models import load_model
import numpy as np
from pydantic import BaseModel
from PIL import Image
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
model = load_model("backend/model.h5")
templates = Jinja2Templates(directory="backend/res")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class MathProblem(BaseModel):
    problem: str
    encodedImage: str


def readb64(uri: str) -> Image:
    data = uri.split(",")[1]
    return Image.open(BytesIO(base64.b64decode(data)))


@app.get("/", response_class=HTMLResponse)
def main_page(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/ping/")
def ping():
    return {"ping": "pong"}


@app.post("/predict/")
def predict_problem_solution(request: Request, problem: MathProblem) -> dict:
    prbl: str = problem.problem + "="
    image = readb64(problem.encodedImage)
    image.save("sol.png")
    image = cv2.imread("sol.png", 0)
    _, width = image.shape
    segment_width = width // 2
    segments = []
    for i in range(2):
        start_x = i * segment_width
        end_x = (i + 1) * segment_width
        segment = image[:, start_x:end_x]
        segments.append(segment)
        predictions = []
    for segment in segments:
        processed_segment = cv2.resize(segment, (28, 28))
        processed_segment = processed_segment / 255.0
        processed_segment = np.expand_dims(processed_segment, axis=-1)
        prediction = model.predict(np.array([processed_segment]))
        predictions.append(prediction)
    try:
        res = "".join([str(np.argmax(pr)) for pr in predictions])
        prbl = "a=" + prbl + res
        print(prbl)
        loc = {}
        exec(prbl, globals(), loc)
    except Exception:
        return {"resultCheck": False}
    return {"resultCheck": loc["a"]}