import os
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image
from typing import Dict, Any

###########################################
# ComfyUI Node Registration
###########################################
NODE_CLASS_MAPPINGS = {}

###########################################
# Imports from your local code
###########################################
from .models.NonlocalNet import WarpNet, VGG19_pytorch
from .models.ColorVidNet import ColorVidNet
from .models.FrameColor import frame_colorization
# Add any other imports you need from .utils or .models

###########################################
# Paths to checkpoints
###########################################
ROOT_DIR = os.path.dirname(__file__)
NONLOCAL_TEST_PATH = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth")
COLOR_TEST_PATH    = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth")
VGG_PATH           = os.path.join(ROOT_DIR, "data", "vgg19_conv.pth")
VGG_GRAY_PATH      = os.path.join(ROOT_DIR, "data", "vgg19_gray.pth")  # if you use it

###########################################
# Global Single Load (lazy load approach)
###########################################
MODELS_LOADED = False
NONLOCAL_NET  = None
COLOR_NET     = None
VGG_NET       = None

def load_models_if_needed():
    """Load colorization models only once. 
       If they are already loaded, do nothing."""
    global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET
    if MODELS_LOADED:
        return

    # Initialize models
    print("[DeepExemplar] Loading nonlocal, color, and VGG19 models...")
    nonlocal_net = WarpNet(1)
    colornet     = ColorVidNet(7)
    vggnet       = VGG19_pytorch()

    # Load checkpoints
    nonlocal_net.load_state_dict(torch.load(NONLOCAL_TEST_PATH, map_location="cuda"))
    colornet.load_state_dict(torch.load(COLOR_TEST_PATH, map_location="cuda"))
    vggnet.load_state_dict(torch.load(VGG_PATH, map_location="cuda"))

    nonlocal_net.cuda().eval()
    colornet.cuda().eval()
    vggnet.cuda().eval()

    NONLOCAL_NET = nonlocal_net
    COLOR_NET    = colornet
    VGG_NET      = vggnet
    MODELS_LOADED = True
    print("[DeepExemplar] Models loaded successfully.")

###########################################
# Utility: Convert PIL->Lab Tensor (Simplified)
###########################################
def pil_to_lab_tensor(img: Image.Image, size=(432, 768)):
    """
    Resizes to `size`, then does BGR->Lab via OpenCV.
    Normalization is a simplified approach. 
    Replicate your exact 'TestTransforms' from the original code if needed.
    """
    img = img.convert("RGB").resize((size[1], size[0]), Image.BICUBIC)
    arr = np.array(img)[:, :, ::-1]  # PIL to BGR
    arr_lab = cv2.cvtColor(arr, cv2.COLOR_BGR2LAB).astype(np.float32)

    # Convert to tensor shape [1,3,H,W]
    lab_t = torch.from_numpy(arr_lab).permute(2, 0, 1).unsqueeze(0)
    # Simple scaling to ~[-1..+1]
    L = (lab_t[:, 0:1, :, :] - 50.) / 50.
    A = lab_t[:, 1:2, :, :] / 128.
    B = lab_t[:, 2:3, :, :] / 128.
    lab_t = torch.cat([L, A, B], dim=1).cuda()
    return lab_t

###########################################
# Utility: Convert Lab->PIL
###########################################
def lab_to_pil(lab_t: torch.Tensor):
    """
    Inverse of pil_to_lab_tensor. shape [1,3,H,W], roughly in [-1..+1].
    Output: PIL.Image in RGB
    """
    L = (lab_t[:, 0:1, :, :] * 50.) + 50.
    A = lab_t[:, 1:2, :, :] * 128.
    B = lab_t[:, 2:3, :, :] * 128.

    lab_255 = torch.cat([L, A, B], dim=1).clamp(0, 255)
    np_lab = lab_255[0].permute(1, 2, 0).cpu().numpy().astype(np.uint8)

    bgr = cv2.cvtColor(np_lab, cv2.COLOR_Lab2BGR)
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)

###########################################
# The Node
###########################################
class DeepExemplarVideoColorNode:
    """
    A ComfyUI node that colorizes a list of frames (VIDEO) using a single reference image.
    The result is a new dictionary with colorized frames.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "video_input": ("VIDEO", ),
                "reference_image": ("IMAGE", ),
                "frame_propagate": ("BOOL", {"default": False}),
            },
        }

    RETURN_TYPES = ("VIDEO",)
    FUNCTION = "colorize_video"
    CATEGORY = "Video/DeepExemplar"

    def colorize_video(self, video_input: Dict[str, Any], reference_image: Image.Image, frame_propagate: bool):
        """
        ComfyUI calls this. 'video_input' is expected to have:
          {
            "frames": [PIL.Image, PIL.Image, ...],
            "fps": int,
            ...
          }
        'reference_image' is a single PIL image.
        'frame_propagate' toggles using the last colorized frame as the next reference.

        Returns a dictionary with colorized frames, e.g.:
          {
            "frames": [...],
            "fps": <int>,
            ...
          }
        """
        load_models_if_needed()

        frames = video_input.get("frames", [])
        if not frames:
            print("[DeepExemplar] No frames found in the input video.")
            return (video_input,)  # just return the input unmodified

        # If frame_propagate is True, we use the first frame as reference
        if frame_propagate:
            reference_image = frames[0]

        # Convert reference to Lab (half-res)
        ref_lab_half = F.interpolate(pil_to_lab_tensor(reference_image), scale_factor=0.5, mode="bilinear")
        # We should compute 'features_B' from the reference if fully replicating test.py logic:
        # But for brevity, let's skip. In a real implementation, you'd replicate test.py's approach 
        #   (converting reference to approximate RGB, passing through vgg, etc.)
        features_B = None  # or your real feature extraction code

        last_lab = None  # for frame_propagation
        out_frames = []

        for idx, frame_pil in enumerate(frames):
            # Convert frame to Lab half-res
            frame_lab_half = F.interpolate(pil_to_lab_tensor(frame_pil), scale_factor=0.5, mode="bilinear")

            if last_lab is None:
                if frame_propagate:
                    last_lab = ref_lab_half.clone()
                else:
                    last_lab = torch.zeros_like(frame_lab_half)

            # Actual colorization logic
            # We call frame_colorization(...). 
            # test.py calls: 
            #   frame_colorization(I_current_lab, I_reference_lab, I_last_lab, features_B, vggnet, nonlocal_net, color_net, ...)
            # We'll do a simplified approach:
            I_current_ab_predict, *_ = frame_colorization(
                frame_lab_half,    # I_current_lab
                ref_lab_half,      # I_reference_lab
                last_lab,          # I_last_lab
                features_B,        # features_B
                VGG_NET,
                NONLOCAL_NET,
                COLOR_NET,
                feature_noise=0,
                temperature=1e-10,
            )

            # Update last_lab
            last_lab = torch.cat([frame_lab_half[:, 0:1, :, :], I_current_ab_predict], dim=1)

            # Upsample predicted ab back to full resolution
            ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear") * 1.25

            # Combine with full-res L
            full_lab = pil_to_lab_tensor(frame_pil)  # Full-res L
            L_full   = full_lab[:, 0:1, :, :]
            colorized_lab = torch.cat([L_full, ab_up], dim=1)
            
            # Convert Lab->PIL
            colorized_pil = lab_to_pil(colorized_lab)
            out_frames.append(colorized_pil)

        # Return new dictionary
        new_video = {
            "frames": out_frames,
            "fps": video_input.get("fps", 30),
            # Copy other fields if needed
        }
        return (new_video,)

# Register the node
NODE_CLASS_MAPPINGS["DeepExemplarVideoColorNode"] = DeepExemplarVideoColorNode