import os
import torch
import torch.nn.functional as F
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights
from torchvision.utils import flow_to_image as torchvision_flow_to_image
import numpy as np

_raft_model = None

def get_raft_model(device):
    global _raft_model
    if _raft_model is None:
        # First, try to load a user-provided checkpoint from ComfyUI's configurable
        # "optical_flow" folder (preferred). If that fails, fall back to the
        # torchvision default weights.
        try:
            # Imported here to avoid adding a hard dependency at module import time
            # when this utility may be used outside of a full ComfyUI environment.
            import folder_paths
            import comfy.utils

            names = folder_paths.get_filename_list("optical_flow")
            model_name = None
            if names:
                # Choose the first available file by default (matching other UI lists).
                model_name = names[0]
            else:
                # Try to use torchvision weights but save them to the optical_flow folder
                try:
                    weights = Raft_Large_Weights.DEFAULT

                    # Instantiate the torchvision model (this will download weights into the torchvision cache)
                    try:
                        model_dl = raft_large(weights=weights, progress=True)
                    except Exception:
                        model_dl = raft_large(weights=weights, progress=False)

                    # Move state dict to CPU and detach
                    state_dict = {k: v.cpu() for k, v in model_dl.state_dict().items()}

                    # Try to infer original filename/extension from the weights metadata or URL
                    candidate_name = None
                    try:
                        url = getattr(weights, 'url', None) or getattr(weights, '_url', None)
                        if url:
                            candidate_name = os.path.basename(str(url).split('?')[0])
                    except Exception:
                        candidate_name = None
                    if not candidate_name:
                        try:
                            meta = getattr(weights, 'meta', None) or {}
                            if isinstance(meta, dict):
                                candidate_name = meta.get('filename') or meta.get('name')
                        except Exception:
                            candidate_name = None

                    # Ensure folder exists: use folder_paths mapping when available
                    save_dir = None
                    if hasattr(folder_paths, 'folder_names_and_paths') and 'optical_flow' in folder_paths.folder_names_and_paths:
                        dirs = folder_paths.folder_names_and_paths['optical_flow'][0]
                        if isinstance(dirs, (list, tuple)) and len(dirs) > 0:
                            save_dir = dirs[0]
                    if not save_dir:
                        save_dir = os.path.join(os.path.dirname(__file__), "..", "..", "models", "optical_flow")
                    os.makedirs(save_dir, exist_ok=True)

                    # Decide filename and format
                    if candidate_name:
                        model_name = candidate_name
                    else:
                        model_name = "raft_large_default.pth"

                    model_path = folder_paths.get_full_path("optical_flow", model_name) or os.path.join(save_dir, model_name)

                    # Save in original format if known (.safetensors or .pth/.pt), otherwise default to .pth
                    _, ext = os.path.splitext(model_name)
                    ext = ext.lower()
                    saved_ok = False
                    try:
                        if ext in ('.safetensors', '.sft'):
                            # Save as safetensors
                            try:
                                comfy.utils.save_torch_file(state_dict, model_path)
                                saved_ok = True
                            except Exception:
                                import safetensors.torch
                                safetensors.torch.save_file(state_dict, model_path)
                                saved_ok = True
                        else:
                            # Default: save as a PyTorch checkpoint (.pth)
                            try:
                                import torch as _torch
                                _torch.save(state_dict, model_path)
                                saved_ok = True
                            except Exception:
                                # Fallback to comfy.utils (may write safetensors)
                                try:
                                    comfy.utils.save_torch_file(state_dict, model_path)
                                    saved_ok = True
                                except Exception:
                                    saved_ok = False
                    except Exception:
                        saved_ok = False

                    if not saved_ok:
                        model_name = None
                except Exception:
                    model_name = None

            if model_name:
                try:
                    model_path = folder_paths.get_full_path_or_raise("optical_flow", model_name)
                    sd = comfy.utils.load_torch_file(model_path, safe_load=True)

                    # Basic validation for a RAFT-large state dict
                    has_raft_keys = (
                        any(k.startswith("feature_encoder.") for k in sd)
                        and any(k.startswith("context_encoder.") for k in sd)
                        and any(k.startswith("update_block.") for k in sd)
                    )
                    if has_raft_keys:
                        model = raft_large(weights=None, progress=False)
                        model.load_state_dict(sd)
                        _raft_model = model.to(device).eval()
                        return _raft_model
                except Exception:
                    pass
        except Exception:
            # If anything goes wrong with folder resolution / loading, fall back below.
            pass

        # Fallback to torchvision-provided default weights
        weights = Raft_Large_Weights.DEFAULT
        _raft_model = raft_large(weights=weights, progress=False).to(device).eval()
    return _raft_model

def warp(x, flow, padding_mode='reflection'):
    """
    Warps an image or tensor using the given flow.
    x: [B, C, H, W]
    flow: [B, 2, H, W] (flow in pixels)
    padding_mode: passed to F.grid_sample (e.g. 'zeros' for occlusion detection)
    """
    b, c, h, w = x.shape
    device = flow.device # Use flow device as the master device
    x = x.to(device)

    # Create grid
    grid_y, grid_x = torch.meshgrid(
        torch.linspace(0, h - 1, h, device=device),
        torch.linspace(0, w - 1, w, device=device),
        indexing='ij'
    )
    # [H, W]

    grid = torch.stack((grid_x, grid_y), dim=0).unsqueeze(0).repeat(b, 1, 1, 1) # [B, 2, H, W]

    # Add flow to grid
    v_grid = grid + flow # [B, 2, H, W]

    # Normalize to [-1, 1] for grid_sample
    v_grid_x = 2.0 * v_grid[:, 0, :, :] / max(w - 1, 1) - 1.0
    v_grid_y = 2.0 * v_grid[:, 1, :, :] / max(h - 1, 1) - 1.0

    v_grid = torch.stack((v_grid_x, v_grid_y), dim=3) # [B, H, W, 2] with (x, y)

    return F.grid_sample(x, v_grid, mode='bilinear', padding_mode=padding_mode, align_corners=True)

def preprocess_image(img, device):
    """
    img: [B, H, W, C] in range [0, 1]
    Returns: [B, 3, H, W] normalized for RAFT
    """
    if img.ndim == 3:
        img = img.unsqueeze(0)

    # ComfyUI [B, H, W, C] -> Torch [B, C, H, W]
    img = img.movedim(-1, 1)

    # Ensure RGB
    if img.shape[1] == 1:
        img = img.expand(-1, 3, -1, -1)
    elif img.shape[1] > 3:
        img = img[:, :3, :, :]

    # Scale to [0, 255] as RAFT transforms usually expect this
    img = img * 255.0

    # The transform expects [0, 255] and returns normalized [-1, 1]
    # It takes (img1, img2) but we can use it for one or just follow its logic
    # Actually, let's just follow the logic: 2 * (img / 255.0) - 1.0
    # Wait, if I do that, it's just 2 * img_01 - 1.0.

    return (2.0 * (img / 255.0) - 1.0).to(device)

def get_optical_flow(img1, img2, tiling_size=0, iterations=12, multi_scale=False):
    """
    img1, img2: [H, W, C] or [B, H, W, C]
    Returns: [B, H, W, 2] flow vectors
    """
    device = img1.device
    model = get_raft_model(device)

    img1_proc = preprocess_image(img1, device)
    img2_proc = preprocess_image(img2, device)

    b, _, h, w = img1_proc.shape

    # Auto-tiling for large images to prevent OOM
    # 2MPx is a reasonable limit for 8GB VRAM without tiling
    if tiling_size <= 0 and (h * w) > 2000000:
        tiling_size = 1024

    # Handle fraction or float conversion
    if 0 < tiling_size < 1:
        tiling_size = int(min(h, w) * tiling_size)
    else:
        tiling_size = int(tiling_size)

    # Ensure divisible by 8 for RAFT
    h_new, w_new = (h // 8) * 8, (w // 8) * 8
    if h != h_new or w != w_new:
        img1_proc = F.interpolate(img1_proc, size=(h_new, w_new), mode='bilinear', align_corners=False)
        img2_proc = F.interpolate(img2_proc, size=(h_new, w_new), mode='bilinear', align_corners=False)

    # Process sequentially to save memory
    flows = []
    with torch.no_grad():
        for i in range(b):
            p1 = img1_proc[i:i+1]
            p2 = img2_proc[i:i+1]

            global_flow = None
            if multi_scale:
                # 1. Global Pass
                # Downsample more aggressively to capture massive jumps (max 256)
                g_size = 256
                g_h = (min(h_new, g_size) // 8) * 8
                g_w = (min(w_new, g_size) // 8) * 8

                gp1 = F.interpolate(p1, size=(g_h, g_w), mode='bilinear', align_corners=False)
                gp2 = F.interpolate(p2, size=(g_h, g_w), mode='bilinear', align_corners=False)

                # Global flow estimation
                global_flow = model(gp1, gp2, num_flow_updates=iterations)[-1]

                # Upsample global flow to high-res
                global_flow = F.interpolate(global_flow, size=(h_new, w_new), mode='bilinear', align_corners=False)
                global_flow[:, 0] *= float(w_new) / g_w
                global_flow[:, 1] *= float(h_new) / g_h

                # print(f"DEBUG: Global flow at center: {global_flow[0, :, h_new//2, w_new//2]}")

                # 2. Warp img2 using global flow
                p2_warped = warp(p2, global_flow)
                target_p2 = p2_warped
            else:
                target_p2 = p2

            # 3. Local Refinement (Tiled or Full)
            if tiling_size > 0:
                f_residual = _tiled_flow(model, p1, target_p2, tiling_size, iterations)
            else:
                f_residual = model(p1, target_p2, num_flow_updates=iterations)[-1]

            # 4. Combine
            if global_flow is not None:
                total_flow = global_flow + f_residual
            else:
                total_flow = f_residual

            flows.append(total_flow)

            if i % 4 == 0 and b > 4:
                torch.cuda.empty_cache()

    flow = torch.cat(flows, dim=0)

    # Resize flow back if needed
    if h != h_new or w != w_new:
        flow = F.interpolate(flow, size=(h, w), mode='bilinear', align_corners=False)
        # Rescale flow values
        flow[:, 0] *= float(w) / w_new
        flow[:, 1] *= float(h) / h_new

    # [B, 2, H, W] -> [B, H, W, 2]
    return flow.movedim(1, -1)

def calculate_occlusion_mask(flow):
    """
    Calculates a disocclusion mask: pixels revealed by motion that have no
    valid source data (e.g. background uncovered when an object moves away).

    flow: [B, H, W, 2] or [H, W, 2]
    Returns: [B, H, W] mask (1.0 = disoccluded/revealed, 0.0 = valid)
    """
    was_single = (flow.ndim == 3)
    if was_single:
        flow = flow.unsqueeze(0)

    # sanitize flow
    flow = torch.nan_to_num(flow, nan=0.0, posinf=0.0, neginf=0.0)

    b, h, w, _ = flow.shape
    device = flow.device

    # [B, H, W, 2] -> [B, 2, H, W] for warp()
    flow_nchw = flow.movedim(-1, 1)

    # Warp a solid-ones image with the flow.
    # Use zero padding so samples from outside are 0, making disocclusions visible
    ones = torch.ones((b, 1, h, w), device=device)
    warped_ones = warp(ones, flow_nchw, padding_mode='zeros')  # [B, 1, H, W]

    # Pixels that are < threshold are disoccluded
    mask = 1.0 - warped_ones.squeeze(1).clamp(0.0, 1.0)  # [B, H, W]

    return mask.squeeze(0) if was_single else mask

def apply_flow(image, flow):
    """
    Warps an image using optical flow vectors.
    image: [B, H, W, C] (ComfyUI format)
    flow:  [B, H, W, 2] (from rife())
    Returns: [B, H, W, C] warped image
    """
    if image.ndim == 3:
        image = image.unsqueeze(0)
    if flow.ndim == 3:
        flow = flow.unsqueeze(0)

    # [B, H, W, 2] -> [B, 2, H, W]
    flow_nchw = flow.movedim(-1, 1)

    # Ensure image is on the same device as flow
    img_nchw = image.movedim(-1, 1).to(flow.device)

    warped = warp(img_nchw, flow_nchw)

    # [B, C, H, W] -> [B, H, W, C]
    return warped.movedim(1, -1)

def flow_to_image(flow):
    """
    Converts flow vectors [B, H, W, 2] to RGB images [B, H, W, 3] for visualization.
    """
    if flow.ndim == 3:
        flow = flow.unsqueeze(0)

    # [B, H, W, 2] -> [B, 2, H, W]
    flow_torch = flow.movedim(-1, 1)

    # torchvision_flow_to_image expects (N, 2, H, W)
    img_uint8 = torchvision_flow_to_image(flow_torch)

    # [B, 3, H, W] -> [B, H, W, 3]
    img = img_uint8.movedim(1, -1).float() / 255.0

    return img

def _tiled_flow(model, img1, img2, tiling_size, iterations):
    """
    Robust tiling implementation with weighted blending and edge coverage.
    """
    b, c, h, w = img1.shape
    device = img1.device

    # Enforce minimum tile size for RAFT
    if tiling_size < 256:
        # If user provided a very small value, maybe they meant fraction? 
        # But for ComfyUI, we'll just enforce a safe minimum.
        tiling_size = max(256, tiling_size)

    # Ensure tiling_size is divisible by 8
    tiling_size = (tiling_size // 8) * 8

    # Clip tiling_size to image dimensions
    t_h = int(min(tiling_size, h))
    t_w = int(min(tiling_size, w))

    # If image is smaller than tiling_size after divisibility adjustments
    if t_h < 8 or t_w < 8:
        return model(img1, img2, num_flow_updates=iterations)[-1]

    output_flow = torch.zeros((b, 2, h, w), device=device)
    weights = torch.zeros((b, 1, h, w), device=device)

    # Stride of 50% overlap is usually good
    stride_h = int(t_h // 2)
    stride_w = int(t_w // 2)

    # Generate coordinates for tiles
    y_coords = list(range(0, h - t_h, stride_h)) + [max(0, h - t_h)]
    x_coords = list(range(0, w - t_w, stride_w)) + [max(0, w - t_w)]

    # Remove duplicates if any (when image size matches stride)
    y_coords = sorted(list(set(y_coords)))
    x_coords = sorted(list(set(x_coords)))

    for y in y_coords:
        for x in x_coords:
            patch1 = img1[:, :, y:y+t_h, x:x+t_w]
            patch2 = img2[:, :, y:y+t_h, x:x+t_w]

            # RAFT expects inputs divisible by 8, but we already ensured t_h, t_w are.
            patch_flow = model(patch1, patch2, num_flow_updates=iterations)[-1]

            # Create local mask for this tile
            # Linear tapering only on overlapping edges
            layer_mask = torch.ones((1, 1, t_h, t_w), device=device)
            overlap_h = stride_h
            overlap_w = stride_w

            if y > 0: # Fade in top
                layer_mask[:, :, :overlap_h, :] *= torch.linspace(0, 1, overlap_h, device=device).view(1, 1, overlap_h, 1)
            if y + t_h < h: # Fade out bottom
                layer_mask[:, :, -overlap_h:, :] *= torch.linspace(1, 0, overlap_h, device=device).view(1, 1, overlap_h, 1)
            if x > 0: # Fade in left
                layer_mask[:, :, :, :overlap_w] *= torch.linspace(0, 1, overlap_w, device=device).view(1, 1, 1, overlap_w)
            if x + t_w < w: # Fade out right
                layer_mask[:, :, :, -overlap_w:] *= torch.linspace(1, 0, overlap_w, device=device).view(1, 1, 1, overlap_w)

            output_flow[:, :, y:y+t_h, x:x+t_w] += patch_flow * layer_mask
            weights[:, :, y:y+t_h, x:x+t_w] += layer_mask

    return output_flow / torch.clamp(weights, min=1e-6)