#!/usr/bin/env python3
"""Stitch a directory recorded by tools/record.py (or scan.py --save-frames) into a page image."""
import argparse
import csv
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from myscan import calib  # noqa: E402
from myscan.stitching import ROT_GAIN, Corrector, LaserOdometry, PageGraph, Stitcher  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dir")
    ap.add_argument("-o", "--out", help="output image (default <dir>/page.png)")
    ap.add_argument("--live", action="store_true", help="use the incremental frame-by-frame stitcher")
    ap.add_argument("--no-refine", action="store_true", help="with --live: lasers only, no image registration")
    ap.add_argument("--rot-gain", type=float, default=ROT_GAIN)
    ap.add_argument("--no-flat", action="store_true", help="skip the EEPROM flat-field correction")
    ap.add_argument("--no-undistort", action="store_true", help="skip the EEPROM lens correction")
    ap.add_argument("--no-wb", action="store_true", help="skip white-patch balancing")
    ap.add_argument("--start", type=int, default=0)
    ap.add_argument("--end", type=int)
    a = ap.parse_args()

    rows = list(csv.DictReader(open(os.path.join(a.dir, "records.csv"))))
    v = np.array([[int(r[f"v{k}"]) for k in range(6)] for r in rows])
    frames = [int(r["frame"]) for r in rows]
    rng = range(a.start, a.end if a.end is not None else len(rows))

    cal = calib.load()
    corrector = Corrector(None if a.no_flat else calib.FlatField(cal),
                          None if a.no_undistort else calib.Undistort(cal))
    st = None
    for i in rng:
        im = cv2.imread(os.path.join(a.dir, f"f{frames[i]:05d}.png"), cv2.IMREAD_UNCHANGED)
        if st is None:
            st = (Stitcher(im.shape, corrector, LaserOdometry(a.rot_gain, max_step_rot=None), refine=not a.no_refine)
                  if a.live else PageGraph(im.shape, corrector, LaserOdometry(a.rot_gain)))
        st.add(im, v[i])
    out = a.out or os.path.join(a.dir, "page.png")
    cv2.imwrite(out, st.image(white_balance=not a.no_wb))
    print(f"{st.n} frames -> {out}  {st.stats}")


if __name__ == "__main__":
    main()