#! /usr/bin/python3
# Last edited on 2026-03-24 19:40:01 by stolfi

import os, sys, re
import cv2
import numpy as np
from process_funcs import bash

def main(set_name):
  for i in range(2):
    # Read images:
    img_file_a = f"keyframes/{set_name}-{i}.jpg"
    img_file_b = f"keyframes/{set_name}-{i+1}.jpg"
    img_a = cv2.imread(img_file_a)
    img_b = cv2.imread(img_file_b)
    # Read point files:
    pts_file_a = f"keyframes/{set_name}-{i}.txt"
    pts_file_b = f"keyframes/{set_name}-{i+1}.txt"
    pts_a = np.loadtxt(pts_file_a)
    pts_b = np.loadtxt(pts_file_b)

    # Generate 30 frames for the a -> b transition
    nf = 30
    frames = []
    for kf in range(nf):
      t = kf / (nf-1)
      # Warp a toward b, and _b toward _a, then blend
      # This keeps the "ink" looking solid
      w_a = morph_images(img_a, img_b, pts_a, pts_b, t)
      w_b = morph_images(img_b, img_a, pts_b, pts_a, 1 - t)
      merged = cv2.addWeighted(w_a, 1 - t, w_b, t, 0)
      frames.append(merged)
      out_file = f"frames/{set_name}-{i}-{kf:03d}.png"
      cv2.imwrite(out_file, merged)

def morph_images(img1, img2, pts1, pts2, t):
  # Interpolate point positions
  pts_interp = (1 - t) * pts1 + t * pts2

  # Create TPS transformer
  tps = cv2.createThinPlateSplineShapeTransformer()

  # Reshape points for OpenCV (1, N, 2)
  p1 = pts1.reshape(1, -1, 2).astype(np.float32)
  p2 = pts_interp.reshape(1, -1, 2).astype(np.float32)

  # Find matches (1 to 1)
  matches = [cv2.DMatch(i, i, 0) for i in range(len(pts1))]

  # Warp image
  tps.estimateTransformation(p2, p1, matches)
  warped = tps.warpImage(img1)
  return warped

main(sys.argv[1])