#!/usr/bin/env python3
"""
test_stitcher.py — Offline test: stitch 4 test images into equirectangular panorama.

Usage:
  python3 test_stitcher.py \\
    --images front.jpg right.jpg rear.jpg left.jpg \\
    --calibration-dir /mnt/nvme/saltybot/calibration \\
    --output panorama.jpg
    [--width 1920] [--height 960]

No ROS2 installation required — uses saltybot_panoramic Python modules directly.
"""

import argparse
import sys
import os
import time

import cv2
import numpy as np

# Allow running without installing the package
_script_dir = os.path.dirname(os.path.abspath(__file__))
_pkg_dir = os.path.join(_script_dir, '..', '..')
if _pkg_dir not in sys.path:
    sys.path.insert(0, os.path.abspath(_pkg_dir))

from saltybot_panoramic.calibration_loader import load_all_camera_infos
from saltybot_panoramic.equirect_projector import CameraConfig, EquirectProjector

_CAM_NAMES = ['front', 'right', 'rear', 'left']
_CAM_YAWS  = [0.0, -90.0, 180.0, 90.0]


def parse_args():
    p = argparse.ArgumentParser(
        description='Offline equirectangular stitching test (no ROS2 needed)'
    )
    p.add_argument(
        '--images', nargs=4, metavar=('FRONT', 'RIGHT', 'REAR', 'LEFT'),
        default=None,
        help='Four input images: front right rear left order',
    )
    p.add_argument(
        '--calibration-dir',
        default='/mnt/nvme/saltybot/calibration',
        help='Calibration YAML base directory',
    )
    p.add_argument('--output', default='panorama.jpg', help='Output panorama path')
    p.add_argument('--width',  type=int, default=1920, help='Output width (default 1920)')
    p.add_argument('--height', type=int, default=960,  help='Output height (default 960)')
    p.add_argument('--pitch-deg', type=float, default=0.0, help='Camera pitch in degrees')
    p.add_argument(
        '--blend-margin', type=int, default=30,
        help='Alpha blending margin in pixels',
    )
    p.add_argument(
        '--show', action='store_true',
        help='Display result with cv2.imshow (requires display)',
    )
    p.add_argument(
        '--benchmark', type=int, default=0, metavar='N',
        help='Run N stitch iterations and report mean FPS',
    )
    return p.parse_args()


def load_test_images(image_paths, calibration_dir):
    """Load images from paths or synthesise coloured test patterns."""
    images = {}
    for name, path in zip(_CAM_NAMES, image_paths):
        img = cv2.imread(path)
        if img is None:
            print(f'  WARNING: could not read {path}, using coloured test pattern')
            img = _make_test_pattern(name, 640, 480)
        else:
            print(f'  Loaded {name}: {path} ({img.shape[1]}x{img.shape[0]})')
        images[name] = img
    return images


def make_synthetic_images():
    """Generate coloured + labelled test patterns for each camera."""
    colours = {
        'front': (0,   200, 200),   # yellow-ish
        'right': (200, 0,   200),   # magenta-ish
        'rear':  (200, 200,   0),   # cyan-ish
        'left':  (0,   200,   0),   # green
    }
    images = {}
    for name in _CAM_NAMES:
        images[name] = _make_test_pattern(name, 640, 480, colours[name])
    return images


def _make_test_pattern(name, w=640, h=480, colour=(180, 180, 180)):
    """Create a labelled test image with grid lines."""
    img = np.full((h, w, 3), colour, dtype=np.uint8)
    # Draw grid
    for x in range(0, w, 80):
        cv2.line(img, (x, 0), (x, h), (0, 0, 0), 1)
    for y in range(0, h, 60):
        cv2.line(img, (0, y), (w, y), (0, 0, 0), 1)
    # Draw label
    cv2.putText(img, name.upper(), (w // 2 - 60, h // 2),
                cv2.FONT_HERSHEY_SIMPLEX, 2.0, (0, 0, 0), 3)
    # Draw circle at optical centre
    cx, cy = w // 2, h // 2
    cv2.circle(img, (cx, cy), 10, (255, 255, 255), -1)
    return img


def main():
    args = parse_args()

    print(f'Output: {args.width}x{args.height}')
    print(f'Calibration dir: {args.calibration_dir}')

    # Load calibration
    print('\nLoading calibration...')
    cam_infos = load_all_camera_infos(args.calibration_dir)

    # Build camera configs
    cam_configs = []
    for name, yaw in zip(_CAM_NAMES, _CAM_YAWS):
        info = cam_infos[name]
        cfg = CameraConfig(
            name=name,
            yaw_deg=yaw,
            pitch_deg=args.pitch_deg,
            K=info['K'],
            D=info['D'],
            image_size=info['image_size'],
            fov_deg=160.0,
        )
        cam_configs.append(cfg)

    # Build projector (precomputes remap maps)
    print('\nBuilding EquirectProjector...')
    t_build = time.monotonic()
    projector = EquirectProjector(
        args.width, args.height, cam_configs, args.blend_margin
    )
    print(f'Build time: {(time.monotonic() - t_build)*1000:.0f}ms')

    # Load or synthesise images
    print('\nLoading input images...')
    if args.images:
        images = load_test_images(args.images, args.calibration_dir)
    else:
        print('  No --images provided, using synthetic test patterns')
        images = make_synthetic_images()

    # Single stitch
    print('\nStitching...')
    t0 = time.monotonic()
    panorama = projector.stitch(images)
    dt_ms = (time.monotonic() - t0) * 1000.0
    print(f'Stitch time: {dt_ms:.1f}ms  ({1000/dt_ms:.1f} FPS potential)')

    # Save output
    cv2.imwrite(args.output, panorama)
    print(f'\nSaved: {args.output} ({panorama.shape[1]}x{panorama.shape[0]})')

    # Benchmark
    if args.benchmark > 0:
        print(f'\nBenchmarking ({args.benchmark} iterations)...')
        times = []
        for _ in range(args.benchmark):
            t = time.monotonic()
            projector.stitch(images)
            times.append(time.monotonic() - t)
        times_ms = [t * 1000 for t in times]
        print(f'  Mean: {sum(times_ms)/len(times_ms):.1f}ms  '
              f'Min: {min(times_ms):.1f}ms  Max: {max(times_ms):.1f}ms  '
              f'FPS: {1000/(sum(times_ms)/len(times_ms)):.1f}')

    # Display
    if args.show:
        cv2.imshow('Equirectangular Panorama', panorama)
        print('\nPress any key to close.')
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    return 0


if __name__ == '__main__':
    sys.exit(main())
