from pathlib import Path
import numpy as np
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
from matplotlib.tri import Triangulation

from image_preparing import fix_image_resolution

import time


transformation_resolution = (1600, 1200)

def auto_correspondences_grid(pil, nx=8, ny=6):
    w, h = pil.size
    xs = np.linspace(0, w - 1, nx)
    ys = np.linspace(0, h - 1, ny)
    pts = []
    for y in ys:
        for x in xs:
            pts.append([float(x), float(y)])
    corners = [[0.0, 0.0], [w - 1.0, 0.0], [w - 1.0, h - 1.0], [0.0, h - 1.0]]
    for c in corners:
        if c not in pts:
            pts.insert(0, c)
    return np.array(pts, dtype=float), np.array(pts, dtype=float).copy()
def compute_delaunay_on_points(pts):
    pts = np.asarray(pts)
    tri = Triangulation(pts[:, 0], pts[:, 1])
    return tri.triangles
def _affine_coeffs(from_tri, to_tri):
    A = []
    B = []
    for (x_dst, y_dst), (x_src, y_src) in zip(from_tri, to_tri):
        A.append([x_dst, y_dst, 1, 0, 0, 0])
        A.append([0, 0, 0, x_dst, y_dst, 1])
        B.append(x_src)
        B.append(y_src)
    sol, *_ = np.linalg.lstsq(np.array(A), np.array(B), rcond=None)
    return tuple(sol.tolist())
def _warp_triangle(src_pil, src_tri, dst_tri, out_size, offset):
    dx, dy = offset
    dst_local = np.array(dst_tri) - np.array([dx, dy])
    coeffs = _affine_coeffs(dst_local, src_tri)
    return src_pil.transform(out_size, Image.AFFINE, coeffs, resample=Image.BILINEAR)
def reconstruct_goal_from_source(src_pil, pts1, pts2):
    w, h = src_pil.size
    pts1 = np.asarray(pts1)
    pts2 = np.asarray(pts2)
    triangles = compute_delaunay_on_points(pts2)
    out = Image.new("RGB", (w, h))
    for tri in triangles:
        tri = np.asarray(tri, dtype=int)
        src_tri = pts1[tri]
        dst_tri = pts2[tri]
        min_x = max(int(np.floor(dst_tri[:, 0].min())), 0)
        max_x = min(int(np.ceil(dst_tri[:, 0].max())), w)
        min_y = max(int(np.floor(dst_tri[:, 1].min())), 0)
        max_y = min(int(np.ceil(dst_tri[:, 1].max())), h)
        if max_x <= min_x or max_y <= min_y:
            continue
        out_w = max_x - min_x
        out_h = max_y - min_y
        warped = _warp_triangle(src_pil, src_tri, dst_tri, (out_w, out_h), (min_x, min_y))
        mask = Image.new("L", (out_w, out_h), 0)
        draw = ImageDraw.Draw(mask)
        tri_local = [(x - min_x, y - min_y) for x, y in dst_tri.tolist()]
        draw.polygon(tri_local, fill=255)
        out.paste(warped, (min_x, min_y), mask)
    return out
def transform_images(src_path, goal_path):
    src = fix_image_resolution(src_path, size=transformation_resolution)
    goal = fix_image_resolution(goal_path, size=transformation_resolution)
    start = time.time()
    image = tile_reposition_transition(src, goal, grid=transformation_resolution)
    end = time.time()
    print(f"Transition took {end - start:.2f} seconds")
    fig, ax = plt.subplots(figsize=(8, 6), facecolor="black")
    fig.patch.set_facecolor("black")
    ax.set_facecolor("black")
    ax.imshow(image, cmap=None)
    ax.axis("off")

    def _close_on_key(event):
        if event.key in ("q", "escape"):
            plt.close(fig)

    fig.canvas.mpl_connect("key_press_event", _close_on_key)

    plt.show()
def _grid_edges(size, divisions):
    return np.rint(np.linspace(0, size, divisions + 1)).astype(np.int32)
def _build_tile_boxes_and_centers(w, h, nx, ny):
    x_edges = _grid_edges(w, nx)
    y_edges = _grid_edges(h, ny)

    x0 = np.repeat(x_edges[:-1], ny)
    x1 = np.repeat(x_edges[1:], ny)
    y0 = np.tile(y_edges[:-1], nx)
    y1 = np.tile(y_edges[1:], nx)

    boxes = np.stack([x0, y0, x1, y1], axis=1)
    centers = np.stack([(x0 + x1) * 0.5, (y0 + y1) * 0.5], axis=1).astype(np.float32)
    areas = np.maximum((x1 - x0) * (y1 - y0), 1).astype(np.float32)
    return boxes, centers, areas
def _integral_image_rgb(arr):
    integral = arr.cumsum(axis=0).cumsum(axis=1)
    return np.pad(integral, ((1, 0), (1, 0), (0, 0)), mode="constant", constant_values=0)
def _tile_means_from_integral(integral, boxes, areas):
    x0 = boxes[:, 0]
    y0 = boxes[:, 1]
    x1 = boxes[:, 2]
    y1 = boxes[:, 3]

    sums = (
        integral[y1, x1]
        - integral[y0, x1]
        - integral[y1, x0]
        + integral[y0, x0]
    )
    return sums / areas[:, None]
def _match_tiles_by_feature_rank(src_avgs, src_centers, goal_avgs, goal_centers, w, h):
    color_scale = 255.0 * np.sqrt(3.0)
    spatial_scale = max(np.hypot(w, h), 1e-6)

    src_feat = np.hstack([
        src_avgs / color_scale,
        (src_centers / spatial_scale) * 0.35,
    ])
    goal_feat = np.hstack([
        goal_avgs / color_scale,
        (goal_centers / spatial_scale) * 0.35,
    ])

    # Fixed projection keeps pairing deterministic and avoids O(n^2) matching.
    projection = np.array([0.50, 0.30, 0.20, 0.60, 0.40], dtype=np.float32)
    src_order = np.argsort(src_feat @ projection)
    goal_order = np.argsort(goal_feat @ projection)
    return src_order, goal_order
def tile_reposition_transition(src_pil, goal_pil, grid=transformation_resolution):
    src = src_pil.convert("RGB")
    goal = goal_pil.convert("RGB")
    w, h = src.size
    nx, ny = grid

    src_boxes, src_centers, src_areas = _build_tile_boxes_and_centers(w, h, nx, ny)
    goal_boxes, goal_centers, goal_areas = _build_tile_boxes_and_centers(w, h, nx, ny)

    src_arr = np.asarray(src, dtype=np.float32)
    goal_arr = np.asarray(goal, dtype=np.float32)
    src_integral = _integral_image_rgb(src_arr)
    goal_integral = _integral_image_rgb(goal_arr)

    src_avgs = _tile_means_from_integral(src_integral, src_boxes, src_areas)
    goal_avgs = _tile_means_from_integral(goal_integral, goal_boxes, goal_areas)

    src_order, goal_order = _match_tiles_by_feature_rank(
        src_avgs,
        src_centers,
        goal_avgs,
        goal_centers,
        w,
        h,
    )

    final = Image.new("RGB", (w, h))
    for s_idx, g_idx in zip(src_order.tolist(), goal_order.tolist()):
        sx0, sy0, sx1, sy1 = src_boxes[s_idx].tolist()
        gx0, gy0, gx1, gy1 = goal_boxes[g_idx].tolist()

        tile = src.crop((sx0, sy0, sx1, sy1))
        goal_w = max(gx1 - gx0, 1)
        goal_h = max(gy1 - gy0, 1)
        if tile.size != (goal_w, goal_h):
            tile = tile.resize((goal_w, goal_h), Image.Resampling.BILINEAR)
        final.paste(tile, (gx0, gy0))
    return final