#!/usr/bin/env python3
"""Rectify a photo of an ECG sheet or strip: a flat sheet with a square grid at 10 px/mm.

    python3 rectify.py photo.jpg flat.jpg [--illum] [--illum-paper {white,keep}] [--redshift]

Outputs: flat.jpg (the flat sheet), flat.json (status, self-check, scale, timings),
flat-check.jpg (the photo with the mask, working area, crop and grid nodes drawn).

Statuses: ok, warn, fail, no_grid (no 5 mm grid: masked crop without rectification),
too_small (sheet < 50 x 30 mm), not_applicable (the grid falls into several separate pieces).

Pipeline: [redshift] -> mask -> crop -> photo scale -> resize to 5.5 px/mm -> grid nodes ->
scope check -> offsets network on windows -> fusion into one map -> node indices (GNN) ->
global correction (cubic polynomial or homography) -> local correction by nodes -> flat sheet.
"""

import argparse
import json
import sys
import time
import traceback
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

from ecglib import correction, fusion, illum, lattice, nets, nodemap, redshift, render, scale, windows

THREADS = 2
OUTPUT_PX_PER_MM = 10.0
TARGET = scale.TARGET_PX_PER_MM
MASK_THRESHOLD = 0.5
CROP_PAD = 0.03
HULL_MIN_ISLAND = 0.02
EDGE_KERNEL = 15
TOO_SMALL_MM = (50.0, 30.0)
MIN_PIECE = 0.05
WINDOW_SCALES = (0.9, 1.0, 1.1)
WINDOW_OVERLAP = 0.5
WINDOW_KEEP = 0.5
MESH_STEP = 16.0
SMOOTHNESS = 1.0
MIN_NODES = 40
MAX_SHEET_MM = 450.0
MAX_POINTS = 5000
OK_MM, WARN_MM = 0.5, 1.0
NODE_MESH_STEP = 8.0
NODE_OUTLIER_MAX = 0.03
NODE_OK_MM, NODE_WARN_MM = 0.2, 0.4


class Stop(Exception):
    """Early end of the pipeline with a final status; flat / check: images to output instead (or None)."""

    def __init__(self, status, reason, flat=None, check=None):
        super().__init__(reason)
        self.status, self.reason, self.flat, self.check = status, reason, flat, check


@dataclass
class Crop:
    photo: np.ndarray
    mask: np.ndarray
    x0: int
    y0: int
    x1: int
    y1: int


@dataclass
class Prepared:
    image: np.ndarray             # crop at 5.5 px/mm
    mask: np.ndarray
    hull: np.ndarray              # working area: convex hull of the mask
    factor: float                 # px_5.5 = (px_original + 0.5) * factor - 0.5
    px_per_mm: float
    verified: bool                # scale confirmed by the 5 mm cell check


@contextmanager
def timer(report, name):
    t = time.perf_counter()
    try:
        yield
    finally:
        report.setdefault("times", {})[name] = round(time.perf_counter() - t, 2)


def rounded(d):
    return {k: round(v, 3) if isinstance(v, float) else v for k, v in d.items()}


def to_original(crop, prep, xy):
    return np.stack([crop.x0 + (xy[:, 0] + 0.5) / prep.factor - 0.5,
                     crop.y0 + (xy[:, 1] + 0.5) / prep.factor - 0.5], 1)


def draw_contours(image, mask, color, width, shift=(0, 0)):
    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(image, [c + np.array(shift) for c in contours], -1, color, width)


def hull_of(mask):
    """Convex hull of the mask without small islands: blank paper next to the signal is bent too."""
    n, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8))
    areas = stats[1:, cv2.CC_STAT_AREA]
    keep = np.isin(labels, 1 + np.nonzero(areas >= HULL_MIN_ISLAND * max(1, areas.sum()))[0])
    points = cv2.findNonZero(keep.astype(np.uint8))
    out = np.zeros(mask.shape, np.uint8)
    if points is not None:
        cv2.fillConvexPoly(out, cv2.convexHull(points), 1)
    return out > 0


def scale_note(status, reason, verified):
    if not verified:
        if status == "ok":
            status = "warn"
        reason += "; scale not verified (cell = 5 mm not confirmed): mm may be wrong, e.g. x5"
    return status, reason


def step_mask(models, photo, report):
    mask = nets.mask_net(models, photo) > MASK_THRESHOLD
    report["mask_fraction"] = float(mask.mean())
    if not mask.any():
        raise Stop("fail", "empty mask")
    return mask


def step_crop_to_mask(photo, mask, report):
    ys, xs = np.nonzero(mask)
    pad = int(CROP_PAD * max(photo.shape[:2]))
    x0, y0 = max(0, xs.min() - pad), max(0, ys.min() - pad)
    x1, y1 = min(photo.shape[1], xs.max() + pad + 1), min(photo.shape[0], ys.max() + pad + 1)
    report["crop"] = [int(x0), int(y0), int(x1), int(y1)]
    return Crop(photo[y0:y1, x0:x1], mask[y0:y1, x0:x1], int(x0), int(y0), int(x1), int(y1))


def step_find_scale(models, photo, mask, crop, report):
    ppm, centre, tried, verified = scale.find_scale_in_mask(models, crop.photo, crop.mask)
    report.update(scale_tries=tried, scale_verified=verified)
    if ppm is None:
        flat = render.crop_to(np.where(mask[..., None], photo, 0), mask)
        raise Stop("no_grid", "no 5 mm lattice in the mask", flat=flat)
    report.update(photo_px_per_mm=ppm, scale_centre=centre)
    return ppm, verified


def step_prepare(crop, ppm, verified, report):
    factor = TARGET / ppm
    image = scale.resize(crop.photo, factor)
    h, w = image.shape[:2]
    mask = cv2.resize(crop.mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) > 0
    hull = hull_of(mask)
    report["prepared"] = [w, h]
    report["hull_extra_fraction"] = float((hull & ~mask).sum() / max(1, mask.sum()))
    return Prepared(image, mask, hull, factor, ppm, verified)


def step_check_size(prep, report):
    ys, xs = np.nonzero(prep.mask)
    size_mm = sorted([(xs.max() - xs.min() + 1) / TARGET, (ys.max() - ys.min() + 1) / TARGET], reverse=True)
    report["sheet_size_mm"] = [round(float(v), 1) for v in size_mm]
    if size_mm[0] > MAX_SHEET_MM:
        raise Stop("fail", "sheet {:.0f} mm long (> {:.0f} mm): photo scale is wrong".format(size_mm[0], MAX_SHEET_MM))
    if prep.verified and size_mm[0] < TOO_SMALL_MM[0] and size_mm[1] < TOO_SMALL_MM[1]:
        raise Stop("too_small", "sheet {:.0f} x {:.0f} mm < {:.0f} x {:.0f} mm".format(
            size_mm[0], size_mm[1], *TOO_SMALL_MM))


def step_detect_nodes(models, prep, report):
    """-> (nodes inside the mask, nodes inside the working area); both slightly dilated to keep edge nodes."""
    h, w = prep.image.shape[:2]
    xy = nets.detector_net(models, prep.image)[:, :2]
    iy = np.clip(np.round(xy[:, 1]).astype(int), 0, h - 1)
    ix = np.clip(np.round(xy[:, 0]).astype(int), 0, w - 1)
    kernel = np.ones((EDGE_KERNEL, EDGE_KERNEL), np.uint8)
    in_mask = (cv2.dilate(prep.mask.astype(np.uint8), kernel) > 0)[iy, ix]
    in_hull = (cv2.dilate(prep.hull.astype(np.uint8), kernel) > 0)[iy, ix]
    report["detector_points"] = int(len(xy))
    if max(in_mask.sum(), in_hull.sum()) > MAX_POINTS:
        raise Stop("fail", "{} detector points (> {}): wrong scale?".format(int(max(in_mask.sum(), in_hull.sum())), MAX_POINTS))
    return xy[in_mask], xy[in_hull]


def step_check_scope(photo, mask, crop, prep, xy_mask, report):
    """The sheet must fit a network window and the grid must be one piece (not strips in frames)."""
    h, w = prep.image.shape[:2]
    if min(w, h) < windows.CROP_SIZE:
        raise Stop("fail", "sheet too small for a 224 px window")
    grid = lattice.grid_area(xy_mask, (h, w))
    report["grid_fraction_of_mask"] = float((prep.mask & grid).sum() / max(1, prep.mask.sum()))
    _, _, stats, _ = cv2.connectedComponentsWithStats(grid.astype(np.uint8))
    areas = stats[1:, cv2.CC_STAT_AREA].astype(float)
    report["grid_pieces"] = int((areas >= MIN_PIECE * max(1.0, areas.sum())).sum())
    if report["grid_pieces"] >= 2:
        check = photo.copy()
        line = max(2, photo.shape[1] // 400)
        draw_contours(check, mask, (255, 0, 0), line)
        grid_crop = cv2.resize(grid.astype(np.uint8), (crop.photo.shape[1], crop.photo.shape[0]),
                               interpolation=cv2.INTER_NEAREST)
        draw_contours(check, grid_crop, (0, 0, 255), line, shift=(crop.x0, crop.y0))
        raise Stop("not_applicable", "grid in {} separate pieces, covers {:.0%} of the mask: not one sheet grid "
                   "(strips in frames?)".format(report["grid_pieces"], report["grid_fraction_of_mask"]), check=check)


def step_rectify_windows(models, prep, xy_hull, report):
    """Offsets network on windows inside the working area and on the grid."""
    h, w = prep.image.shape[:2]
    grid = lattice.grid_area(xy_hull, (h, w))
    use = prep.hull & grid if grid.sum() > 0.1 * prep.mask.sum() else prep.hull
    all_windows = windows.windows_for(w, h, WINDOW_SCALES, WINDOW_OVERLAP)
    keep = windows.window_inside(all_windows, use) >= WINDOW_KEEP
    chosen = [win for win, k in zip(all_windows, keep) if k]
    report.update(windows_all=int(len(all_windows)), windows_kept=len(chosen))
    if len(chosen) < 4:
        raise Stop("fail", "too few windows in the mask")
    return chosen, nets.rectify_net(models, prep.image, chosen)


def step_fuse(prep, chosen, offsets, report):
    h, w = prep.image.shape[:2]
    mesh, nodes, report["window_agreement_mm"] = fusion.fuse_windows(w, h, chosen, offsets, MESH_STEP, SMOOTHNESS)
    return mesh, nodes


def step_index_nodes(models, xy, report):
    """Lattice indices by the GNN; the largest connected piece is kept."""
    pts, ij = np.zeros((0, 2)), np.zeros((0, 2), int)
    if len(xy) >= MIN_NODES:
        ij_all, component = nets.lattice_gnn(models, xy)
        keep = component >= 0
        if keep.any():
            main = component == np.bincount(component[keep]).argmax()
            pts, ij = xy[main], ij_all[main]
    report["detector_nodes"] = int(len(pts))
    return pts, ij


def step_correct(mesh, network_nodes, pts, ij, verified, report):
    """Global correction: network + cubic polynomial or a homography, whichever has the lower self-check."""
    nodes, status, reason = network_nodes, "warn", "too few detector nodes: no correction, no self-check"
    if len(pts) >= MIN_NODES:
        check = correction.holdout_network(mesh, network_nodes, pts, ij)
        check_h = correction.holdout_homography(mesh, pts, ij)
        report["selfcheck_network_mm"] = rounded(check)
        if check_h is not None:
            report["selfcheck_homography_mm"] = rounded(check_h)
        apply = correction.fit_correction(mesh.sample(network_nodes, pts), ij)
        h_nodes = correction.homography_nodes(mesh, pts, ij)
        if h_nodes is not None and check_h is not None and check_h["median_mm"] < check["median_mm"]:
            nodes, check, report["method"] = h_nodes, check_h, "homography"
        else:
            nodes, report["method"] = apply(network_nodes), "network+poly3"
        report["selfcheck_holdout_mm"] = rounded(check)
        med = check["median_mm"]
        status = "ok" if med <= OK_MM else "warn" if med <= WARN_MM else "fail"
        reason = "hold-out median {:.2f} mm".format(med)
    status, reason = scale_note(status, reason, verified)
    return nodes, status, reason


def map_from_half(mesh, network_nodes, pts, ij, fit, method):
    """Global map fitted on the nodes `fit` only (for the self-check)."""
    if method == "homography":
        h_nodes = correction.homography_nodes(mesh, pts[fit], ij[fit])
        if h_nodes is not None:
            return lambda p: mesh.sample(h_nodes, p)
    apply = correction.fit_correction(mesh.sample(network_nodes, pts[fit]), ij[fit])
    return lambda p: apply(mesh.sample(network_nodes, p))


def step_node_correction(prep, mesh, nodes, network_nodes, pts, ij, verified, status, reason, report):
    """Local correction by the nodes; skipped when more than NODE_OUTLIER_MAX of the nodes are outliers."""
    report["node_correction"] = "off"
    if len(pts) < MIN_NODES:
        return mesh, nodes, status, reason
    apply, outliers = nodemap.correct(mesh.sample(nodes, pts), ij)
    fraction = outliers / len(pts)
    report["node_outliers"] = int(outliers)
    report["node_outlier_fraction"] = round(fraction, 4)
    if fraction > NODE_OUTLIER_MAX:
        report["node_correction"] = "skipped"
        if status == "ok":
            status = "warn"
        reason += "; {:.1%} of grid nodes are outliers (> {:.0%}): grid found with errors, node correction not applied".format(
            fraction, NODE_OUTLIER_MAX)
        return mesh, nodes, status, reason
    fit, test = correction.halves(ij)
    base = map_from_half(mesh, network_nodes, pts, ij, fit, report.get("method"))
    apply_half, _ = nodemap.correct(base(pts[fit]), ij[fit])
    check = correction.score(apply_half(base(pts[test])), ij[test])
    report["selfcheck_global_mm"] = report.get("selfcheck_holdout_mm")
    report["selfcheck_holdout_mm"] = rounded(check)
    report["method"] = str(report.get("method")) + "+nodes"
    report["node_correction"] = "applied"
    med = check["median_mm"]
    status = "ok" if med <= NODE_OK_MM else "warn" if med <= NODE_WARN_MM else "fail"
    reason = "hold-out median {:.2f} mm (with node correction)".format(med)
    status, reason = scale_note(status, reason, verified)
    h, w = prep.image.shape[:2]
    fine = fusion.Mesh(w, h, NODE_MESH_STEP)
    fine_px = np.stack(np.meshgrid(np.arange(fine.nx) * fine.step, np.arange(fine.ny) * fine.step), -1).reshape(-1, 2)
    return fine, apply(mesh.sample(nodes, fine_px)), status, reason


def step_render(photo, crop, prep, mesh, nodes, report):
    source = photo[crop.y0:crop.y1, crop.x0:crop.x1]
    map_x, map_y = render.inverse_map(source.shape, prep.factor, mesh, nodes, OUTPUT_PX_PER_MM)
    flat = cv2.remap(source, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))
    flat_mask = cv2.remap(crop.mask.astype(np.uint8) * 255, map_x, map_y, cv2.INTER_NEAREST,
                          borderMode=cv2.BORDER_CONSTANT, borderValue=0) > 0
    report["output_px_per_mm"] = OUTPUT_PX_PER_MM
    return render.crop_to(flat, flat_mask)


def draw_check(photo, mask, crop, prep, pts):
    """Mask (blue), working area (yellow), crop (green), grid nodes (red crosses)."""
    check = photo.copy()
    w = photo.shape[1]
    draw_contours(check, mask, (255, 0, 0), max(2, w // 400))
    cv2.rectangle(check, (crop.x0, crop.y0), (crop.x1 - 1, crop.y1 - 1), (0, 200, 0), max(1, w // 600))
    hull = cv2.resize(prep.hull.astype(np.uint8), (crop.photo.shape[1], crop.photo.shape[0]),
                      interpolation=cv2.INTER_NEAREST)
    draw_contours(check, hull, (0, 220, 255), max(2, w // 500), shift=(crop.x0, crop.y0))
    for x, y in to_original(crop, prep, pts):
        cv2.drawMarker(check, (int(x), int(y)), (0, 0, 255), cv2.MARKER_CROSS, max(6, w // 150), max(1, w // 800))
    return check


def rectify_photo(models, photo, use_redshift=False):
    """BGR photo -> (flat sheet or None, check image or None, report).
    With use_redshift the networks see the recoloured photo; the outputs are built from the original."""
    report = {"redshift": False}
    t0 = time.perf_counter()
    flat = check = None
    try:
        seen = photo
        if use_redshift:
            with timer(report, "redshift"):
                report["redshift"] = True
                seen = redshift.redshift(photo)
        with timer(report, "mask"):
            mask = step_mask(models, seen, report)
        crop = step_crop_to_mask(seen, mask, report)
        with timer(report, "scale"):
            ppm, verified = step_find_scale(models, photo, mask, crop, report)
        prep = step_prepare(crop, ppm, verified, report)
        step_check_size(prep, report)
        with timer(report, "detector"):
            xy_mask, xy_hull = step_detect_nodes(models, prep, report)
        step_check_scope(photo, mask, crop, prep, xy_mask, report)
        with timer(report, "windows"):
            chosen, offsets = step_rectify_windows(models, prep, xy_hull, report)
        with timer(report, "fusion"):
            mesh, network_nodes = step_fuse(prep, chosen, offsets, report)
        with timer(report, "correction"):
            pts, ij = step_index_nodes(models, xy_hull, report)
            nodes, status, reason = step_correct(mesh, network_nodes, pts, ij, verified, report)
            mesh, nodes, status, reason = step_node_correction(prep, mesh, nodes, network_nodes, pts, ij,
                                                               verified, status, reason, report)
        with timer(report, "render"):
            flat = step_render(photo, crop, prep, mesh, nodes, report)
        check = draw_check(photo, mask, crop, prep, pts)
    except Stop as stop:
        status, reason, flat, check = stop.status, stop.reason, stop.flat, stop.check
    report.update(status=status, reason=reason)
    report["times"]["total"] = round(time.perf_counter() - t0, 2)
    return flat, check, report


def main():
    parser = argparse.ArgumentParser(description="Rectify a photo of an ECG sheet.")
    parser.add_argument("photo", help="input photo")
    parser.add_argument("output", help="output image (flat sheet)")
    parser.add_argument("--illum", action="store_true", help="also write <output>-illum.jpg with flattened illumination")
    parser.add_argument("--illum-paper", choices=["white", "keep"], default=None,
                        help="paper colour for --illum (default white); implies --illum")
    parser.add_argument("--redshift", action="store_true", help="recolour dark lines red before the networks (black or grey grids)")
    args = parser.parse_args()
    if args.illum_paper is not None:
        args.illum = True
    in_path, out_path = Path(args.photo), Path(args.output)
    photo = cv2.imread(str(in_path), cv2.IMREAD_COLOR)
    if photo is None:
        sys.exit("cannot read {}".format(in_path))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    stem = out_path.with_suffix("")
    models = nets.load_models(threads=THREADS)
    try:
        flat, check, report = rectify_photo(models, photo, use_redshift=args.redshift)
    except Exception as error:
        traceback.print_exc()
        flat, check, report = None, None, {"status": "fail", "reason": repr(error)}
    report["photo"] = str(in_path)
    if flat is not None:
        cv2.imwrite(str(out_path), flat, [cv2.IMWRITE_JPEG_QUALITY, 90])
        if args.illum:
            t = time.perf_counter()
            cv2.imwrite(str(stem) + "-illum.jpg", illum.flatten(flat, OUTPUT_PX_PER_MM, paper=args.illum_paper or "white"),
                        [cv2.IMWRITE_JPEG_QUALITY, 92])
            report["times"]["illum"] = round(time.perf_counter() - t, 2)
    if check is not None:
        cv2.imwrite(str(stem) + "-check.jpg", check, [cv2.IMWRITE_JPEG_QUALITY, 85])
    Path(str(stem) + ".json").write_text(json.dumps(report, indent=1, default=float, ensure_ascii=False))
    selfcheck = report.get("selfcheck_holdout_mm", {}).get("median_mm")
    ppm = report.get("photo_px_per_mm")
    print("{}: {} ({})".format(in_path.name, report.get("status"), report.get("reason")))
    print("  method: {}, photo scale: {} px/mm, self-check median: {} mm, time: {} s".format(
        report.get("method"), None if ppm is None else round(ppm, 2), selfcheck, report.get("times", {}).get("total")))
    print("  flat sheet: {}".format(out_path if flat is not None else "not produced"))
    sys.exit(0 if flat is not None else 1)


if __name__ == "__main__":
    main()