#!/usr/bin/env python3
"""
Stitch north (cylindrical warp) + sunrise (unmodified).

Usage:
  python3 stitch-cameras.py north.jpg sunrise.jpg out_prefix [focal_px]

  focal_px: cylindrical focal length in pixels (default: auto-search).
            Approx guide: 1920 = ~90° HFOV for a 3840px-wide sensor.

Pipeline
--------
1. Cylindrical-warp north only; sunrise is never modified.
2. Build a soft validity mask for the warped north so black-corner
   pixels fade out smoothly.
3. Phase-correlate edge-enhanced treeline strips to find tx, ty.
4. Colour-match north's overlap strip to sunrise's (per-channel gain)
   so the two cameras' exposure/WB differences don't leave a visible seam.
5. Feathered alpha blend in the overlap zone, weighted by north's validity.
6. Crop black margins; output full-res panorama + half-size preview.
"""
import sys, cv2, numpy as np

north_p  = sys.argv[1]
sunrise_p= sys.argv[2]
prefix   = sys.argv[3]
forced_f = float(sys.argv[4]) if len(sys.argv) > 4 else None

north   = cv2.imread(north_p)
sunrise = cv2.imread(sunrise_p)
assert north is not None,  f"Cannot read {north_p}"
assert sunrise is not None, f"Cannot read {sunrise_p}"
H, W = north.shape[:2]
cx, cy = W / 2.0, H / 2.0

# ── Helpers ───────────────────────────────────────────────────────────────────

def cylindrical_warp(img, f):
    xc = np.arange(W, dtype=np.float32) - cx
    yc = np.arange(H, dtype=np.float32) - cy
    XC, YC = np.meshgrid(xc, yc)
    theta = XC / f
    map_x = (f * np.tan(theta) + cx).astype(np.float32)
    map_y = (YC / np.cos(theta) + cy).astype(np.float32)
    warped = cv2.remap(img, map_x, map_y, cv2.INTER_LINEAR,
                       borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    valid = ((map_x >= 0) & (map_x < W) &
             (map_y >= 0) & (map_y < H)).astype(np.float32)
    return warped, valid

def phase_corr(n_cyl):
    """Return (peak_score, tx, ty) aligning warped-north to sunrise."""
    bw = int(W * 0.35)
    y0, y1 = int(H * 0.60), int(H * 0.92)
    n_strip = cv2.cvtColor(n_cyl  [:, W-bw:], cv2.COLOR_BGR2GRAY)[y0:y1].astype(np.float64)
    s_strip = cv2.cvtColor(sunrise[:, :bw],   cv2.COLOR_BGR2GRAY)[y0:y1].astype(np.float64)
    ne = cv2.Laplacian(n_strip, cv2.CV_64F, ksize=3)
    se = cv2.Laplacian(s_strip, cv2.CV_64F, ksize=3)
    N = np.fft.fft2(ne); S = np.fft.fft2(se)
    R = N * np.conj(S); nrm = np.abs(R); nrm[nrm == 0] = 1
    corr = np.fft.ifft2(R / nrm).real
    peak_val = corr.max()
    dy_r, dx_r = [int(v) for v in np.unravel_index(np.argmax(corr), corr.shape)]
    rh, rw = corr.shape
    if dy_r > rh // 2: dy_r -= rh
    if dx_r > rw // 2: dx_r -= rw
    return peak_val, -(W - bw + dx_r), -dy_r

def colour_match(src, ref_strip_src, ref_strip_ref):
    """Scale each BGR channel of src so strip means match between the two strips."""
    out = src.astype(np.float32)
    for c in range(3):
        s_mean = ref_strip_src[..., c].astype(np.float32).mean()
        r_mean = ref_strip_ref[..., c].astype(np.float32).mean()
        gain = np.clip(r_mean / s_mean if s_mean > 1e-3 else 1.0, 0.5, 2.0)
        out[..., c] = np.clip(out[..., c] * gain, 0, 255)
    return out

# ── Focal length search (or forced) ──────────────────────────────────────────
focal_candidates = [forced_f] if forced_f else [1400, 1550, 1700, 1800, 1920, 2100, 2300, 2600]
print(f"{'focal':>6}  {'score':>8}  {'tx':>7}  {'ty':>5}")
best_score, best_f, best_tx, best_ty = -1.0, 1920.0, 0, 0
for f in focal_candidates:
    n_cyl, _ = cylindrical_warp(north, f)
    score, tx, ty = phase_corr(n_cyl)
    print(f"{f:6.0f}  {score:8.5f}  {tx:7d}  {ty:5d}")
    if score > best_score:
        best_score, best_f, best_tx, best_ty = score, f, tx, ty

f   = best_f
tx  = best_tx
ty  = best_ty
print(f"\nBest focal={f:.0f}  score={best_score:.5f}  tx={tx}  ty={ty}")

# ── Warp + validity mask ──────────────────────────────────────────────────────
n_cyl, valid_raw = cylindrical_warp(north, f)

kernel      = np.ones((31, 31), np.float32) / (31 * 31)
valid_soft  = cv2.filter2D(valid_raw, -1, kernel).clip(0, 1)

# First column of north with >50% valid pixels in the central band
central_y0, central_y1 = int(H * 0.10), int(H * 0.90)
col_valid = valid_raw[central_y0:central_y1, :].mean(axis=0)
first_valid_col = int(np.argmax(col_valid > 0.5))

# ── Canvas geometry ───────────────────────────────────────────────────────────
xmin = min(0, tx); ymin = min(0, ty)
xmax = max(W, W + tx); ymax = max(H, H + ty)
cW, cH = int(xmax - xmin), int(ymax - ymin)
sx0, sy0 = int(-xmin), int(-ymin)
nx0, ny0 = int(sx0 + tx), int(sy0 + ty)

ov_x0 = max(sx0, nx0);  ov_x1 = min(sx0 + W, nx0 + W)
ov_y0 = max(sy0, ny0);  ov_y1 = min(sy0 + H, ny0 + H)

# ── Colour-match north → sunrise in the overlap strip ────────────────────────
# Use a central vertical slice of the overlap (avoid the edge-feather zones).
ov_w = ov_x1 - ov_x0
# Sample the middle half of the overlap, treeline rows only (avoid sky + timestamp)
y_lo, y_hi = int(H * 0.55), int(H * 0.88)
mid_n_x0 = ov_x0 - nx0 + ov_w // 4
mid_n_x1 = ov_x0 - nx0 + 3 * ov_w // 4
mid_s_x0 = ov_x0 - sx0 + ov_w // 4
mid_s_x1 = ov_x0 - sx0 + 3 * ov_w // 4
n_sample = n_cyl  [y_lo:y_hi, mid_n_x0:mid_n_x1]
s_sample = sunrise[y_lo:y_hi, mid_s_x0:mid_s_x1]
n_cyl_matched = colour_match(n_cyl, n_sample, s_sample)
print(f"Colour-matched north to sunrise (overlap centre rows {y_lo}-{y_hi})")

# ── Composite ────────────────────────────────────────────────────────────────
canvas = np.zeros((cH, cW, 3), dtype=np.float32)

# 1. Paint north (valid content only, black corners excluded)
nv   = valid_soft[..., np.newaxis]
c_s  = nx0 + first_valid_col
c_e  = min(nx0 + W, cW)
canvas[ny0:ny0 + H, c_s:c_e] = (
    n_cyl_matched[:, first_valid_col:c_e - nx0].astype(np.float32)
    * nv[:, first_valid_col:c_e - nx0]
)

# 2. Sunrise overwrites its entire area (completely unmodified)
canvas[sy0:sy0 + H, sx0:sx0 + W] = sunrise.astype(np.float32)

# 3. Feathered blend in overlap: when north is invalid, fall back to sunrise
if ov_x1 > ov_x0 and ov_y1 > ov_y0:
    ov_w = ov_x1 - ov_x0
    alpha    = np.linspace(0.0, 1.0, ov_w, dtype=np.float32)[np.newaxis, :, np.newaxis]
    n_patch  = n_cyl_matched[ov_y0 - ny0:ov_y1 - ny0, ov_x0 - nx0:ov_x1 - nx0].astype(np.float32)
    s_patch  = sunrise      [ov_y0 - sy0:ov_y1 - sy0, ov_x0 - sx0:ov_x1 - sx0].astype(np.float32)
    nv_patch = valid_soft   [ov_y0 - ny0:ov_y1 - ny0, ov_x0 - nx0:ov_x1 - nx0][..., np.newaxis]
    s_weight = np.maximum(alpha, 1.0 - nv_patch)
    canvas[ov_y0:ov_y1, ov_x0:ov_x1] = s_patch * s_weight + n_patch * (1.0 - s_weight)
    print(f"Overlap blend: {ov_w}x{ov_y1 - ov_y0}px")

out = canvas.clip(0, 255).astype(np.uint8)

# Vertical crop to rows where both images exist
top = max(sy0, ny0); bot = min(sy0 + H, ny0 + H)
out = out[top:bot, :]

# Left-crop black corner columns
col_has_content = out.max(axis=(0, 2)) > 0
left_crop = int(np.argmax(col_has_content))
out = out[:, left_crop:]
print(f"Final size: {out.shape[1]}x{out.shape[0]}")

cv2.imwrite(f"{prefix}-panorama.jpg", out, [cv2.IMWRITE_JPEG_QUALITY, 93])
half = cv2.resize(out, (out.shape[1] // 2, out.shape[0] // 2))
cv2.imwrite(f"{prefix}-preview.jpg", half, [cv2.IMWRITE_JPEG_QUALITY, 88])

with open(f"{prefix}-params.txt", "w") as fh:
    fh.write(f"focal={f:.0f}\ntx={tx}\nty={ty}\n")