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


def prepare_image(path):
    pil = Image.open(path).convert("RGB").resize((800, 600))
    return pil
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):
    fix_image_resolution(src_path)
    fix_image_resolution(goal_path)
    src = prepare_image(src_path)
    goal = prepare_image(goal_path)
    frames = tile_reposition_transition(src, goal, grid=(40, 30), n_frames=24)
    fig, ax = plt.subplots(figsize=(8, 6))
    im = ax.imshow(frames[0])
    ax.axis("off")
    for frame in frames:
        im.set_data(frame)
        fig.canvas.draw_idle()
        plt.pause(0.05)
    plt.show()
def tile_reposition_transition(src_pil, goal_pil, grid=(40, 30), n_frames=24):
    
    src = src_pil.convert("RGBA")
    goal = goal_pil.convert("RGBA")
    w, h = src.size
    nx, ny = grid

    def make_tiles(pil):
        tiles = []
        for iy in range(ny):
            y0 = int(round(iy * h / ny))
            y1 = int(round((iy + 1) * h / ny))
            for ix in range(nx):
                x0 = int(round(ix * w / nx))
                x1 = int(round((ix + 1) * w / nx))
                box = (x0, y0, x1, y1)
                crop = pil.crop(box).convert("RGB")
                arr = np.array(crop, dtype=float)
                avg = arr.reshape(-1, 3).mean(axis=0)
                center = ((x0 + x1) / 2.0, (y0 + y1) / 2.0)
                tiles.append({"box": box, "img": crop, "avg": avg, "center": center})
        return tiles

    src_tiles = make_tiles(src)
    goal_tiles = make_tiles(goal)

    max_spatial = np.hypot(w, h)
    src_available = set(range(len(src_tiles)))
    mapping = {} 
    for g_idx, g in enumerate(goal_tiles):
        best = None
        best_cost = None
        for s_idx in list(src_available):
            s = src_tiles[s_idx]
            color_dist = np.linalg.norm(g["avg"] - s["avg"]) / (255.0 * np.sqrt(3))
            spatial_dist = np.linalg.norm(np.array(g["center"]) - np.array(s["center"])) / max_spatial
            cost = color_dist + 0.35 * spatial_dist
            if best_cost is None or cost < best_cost:
                best_cost = cost
                best = s_idx
        if best is None:
            best = src_available.pop()
        else:
            src_available.remove(best)
        mapping[best] = g_idx

    targets = {}
    for s_idx, g_idx in mapping.items():
        s = src_tiles[s_idx]
        g = goal_tiles[g_idx]
        targets[s_idx] = {"start": (s["box"][0], s["box"][1]), "end": (g["box"][0], g["box"][1]), "img": s["img"], "size": (s["box"][2]-s["box"][0], s["box"][3]-s["box"][1])}

    frames = []
    for k in range(n_frames + 1):
        t = k / float(n_frames)
        frame = Image.new("RGBA", (w, h), (0, 0, 0, 255))
        for s_idx, info in targets.items():
            sx, sy = info["start"]
            ex, ey = info["end"]
            cx = int(round(sx + (ex - sx) * t))
            cy = int(round(sy + (ey - sy) * t))
            frame.paste(info["img"], (cx, cy))
        frames.append(frame.convert("RGB"))

    final = Image.new("RGB", (w, h))
    for s_idx, g_idx in mapping.items():
        src_img = src_tiles[s_idx]["img"].convert("RGB")
        x0, y0, x1, y1 = goal_tiles[g_idx]["box"]
        final.paste(src_img, (x0, y0))
    frames[-1] = final
    return frames