"""Web demo (Gradio): upload a photo, press Rectify, get the flat sheet and the status.

    pip install gradio && python3 app.py
"""

import tempfile
import uuid
from pathlib import Path

import cv2
import gradio as gr

from ecglib import illum, nets
from rectify import OUTPUT_PX_PER_MM, THREADS, rectify_photo

MAX_SIDE = 2048
MODELS = nets.load_models(threads=THREADS)
OUT = Path(tempfile.mkdtemp())


def status_text(report, shrunk):
    lines = ["**Status: {}** — {}".format(report.get("status"), report.get("reason"))]
    holdout = report.get("selfcheck_holdout_mm", {}).get("median_mm")
    if holdout is not None:
        lines.append("Self-check median: {} mm".format(holdout))
    if report.get("photo_px_per_mm"):
        lines.append("Photo scale: {:.2f} px/mm".format(report["photo_px_per_mm"]))
    if report.get("sheet_size_mm"):
        lines.append("Sheet: {:.0f} x {:.0f} mm".format(*report["sheet_size_mm"]))
    if shrunk:
        lines.append("Photo resized to {} px on the longer side".format(MAX_SIDE))
    lines.append("Time: {} s".format(report.get("times", {}).get("total")))
    return "\n\n".join(lines)


def run(path, fix_illumination):
    if path is None:
        return None, None, "Upload a photo first."
    photo = cv2.imread(path, cv2.IMREAD_COLOR)
    if photo is None:
        return None, None, "Cannot read the image."
    scale = MAX_SIDE / max(photo.shape[:2])
    if scale < 1:
        photo = cv2.resize(photo, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    try:
        flat, check, report = rectify_photo(MODELS, photo)
    except Exception as error:
        return None, None, "**Status: fail** — {!r}".format(error)
    stem = OUT / uuid.uuid4().hex
    flat_path = check_path = None
    if flat is not None:
        if fix_illumination:
            flat = illum.flatten(flat, OUTPUT_PX_PER_MM, paper="white")
        flat_path = str(stem) + "-flat.jpg"
        cv2.imwrite(flat_path, flat, [cv2.IMWRITE_JPEG_QUALITY, 92])
    if check is not None:
        check_path = str(stem) + "-check.jpg"
        cv2.imwrite(check_path, check, [cv2.IMWRITE_JPEG_QUALITY, 85])
    return flat_path, check_path, status_text(report, scale < 1)


with gr.Blocks(title="ECG Rectifier") as demo:
    gr.Markdown("# ECG Rectifier\nA phone photo of a paper ECG (bent, crumpled, folded) -> a flat sheet "
                "with a square grid at 10 px/mm. CPU only; one photo takes about a minute.")
    with gr.Row():
        with gr.Column():
            photo = gr.Image(type="filepath", label="Photo")
            fix = gr.Checkbox(value=True, label="Fix illumination")
            button = gr.Button("Rectify", variant="primary")
            gr.Examples([["examples/folded.jpg"], ["examples/crumpled.jpg"]], inputs=[photo])
        with gr.Column():
            flat = gr.Image(type="filepath", label="Flat sheet")
            status = gr.Markdown()
            check = gr.Image(type="filepath", label="Mask (blue), working area (yellow), crop (green), grid nodes used (red)")
    button.click(run, inputs=[photo, fix], outputs=[flat, check, status])

if __name__ == "__main__":
    demo.launch()