from .comfyui_invsr_trimmed import get_configs, InvSamplerSR, BaseSampler
import torch
from comfy.utils import ProgressBar
from folder_paths import get_full_path, get_folder_paths
import os

def split_tensor_into_batches(tensor, batch_size):
    """
    Split a tensor into smaller batches of specified size
    
    Args:
        tensor (torch.Tensor): Input tensor of shape (N, C, H, W)
        batch_size (int): Desired batch size for splitting
        
    Returns:
        list: List of tensors, each with batch_size (except possibly the last one)
    """
    # Get original batch size
    original_batch_size = tensor.size(0)
    
    # Calculate number of full batches and remaining samples
    num_full_batches = original_batch_size // batch_size
    remaining_samples = original_batch_size % batch_size
    
    # Split tensor into chunks
    batches = []
    
    # Handle full batches
    for i in range(num_full_batches):
        start_idx = i * batch_size
        end_idx = start_idx + batch_size
        batch = tensor[start_idx:end_idx]
        batches.append(batch)
    
    # Handle remaining samples if any
    if remaining_samples > 0:
        last_batch = tensor[-remaining_samples:]
        batches.append(last_batch)
    
    return batches

class Namespace:
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)
    
    def __repr__(self):
        items = [f"{key}={repr(value)}" for key, value in vars(self).items()]
        return f"Namespace({', '.join(items)})"

class LoadInvSRModels:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sd_model": (['stabilityai/sd-turbo'],),
                "invsr_model": (['noise_predictor_sd_turbo_v5.pth'],),
                "dtype": (['fp16', 'fp32', 'bf16'], {"default": "fp16"}),
                "tiled_vae": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("INVSR_PIPE",)
    RETURN_NAMES = ("invsr_pipe",)
    FUNCTION = "loadmodel"
    CATEGORY = "INVSR"

    def loadmodel(self, sd_model, invsr_model, dtype, tiled_vae):
        match dtype:
            case "fp16":
                dtype = "torch.float16"
            case "fp32":
                dtype = "torch.float32"
            case "bf16":
                dtype = "torch.bfloat16"

        cfg_path = os.path.join(
            os.path.dirname(__file__), "configs", "sample-sd-turbo.yaml"
        )
        sd_path = get_folder_paths("diffusers")[0]
        ckpt_path = get_full_path("invsr", "noise_predictor_sd_turbo_v5.pth")

        args = Namespace(
            bs=1,
            chopping_bs=8,
            timesteps=None,
            num_steps=1,
            cfg_path=cfg_path,
            sd_path=sd_path,
            started_ckpt_path=ckpt_path,
            tiled_vae=tiled_vae,
            color_fix="",
            chopping_size=128,
        )
        configs = get_configs(args)
        configs["sd_pipe"]["params"]["torch_dtype"] = dtype
        base_sampler = BaseSampler(configs)

        return (base_sampler,)

class InvSRSampler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "invsr_pipe": ("INVSR_PIPE",),
                "images": ("IMAGE",),
                "num_steps": ("INT",{"default": 1, "min": 1, "max": 5}),
                "batch_size": ("INT",{"default": 1}),
                "chopping_batch_size": ("INT",{"default": 8}),
                "chopping_size": ([128, 256, 512],{"default": 128}),
                "color_fix": (['none', 'wavelet', 'ycbcr'], {"default": "none"}),
                "seed": ("INT", {"default": 123, "min": 0, "max": 2**32 - 1, "step": 1}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "process"
    CATEGORY = "INVSR"

    def process(self, invsr_pipe, images, num_steps, batch_size, chopping_batch_size, chopping_size, color_fix, seed):
        base_sampler = invsr_pipe
        if color_fix == "none":
            color_fix = ""

        cfg_path = os.path.join(
            os.path.dirname(__file__), "configs", "sample-sd-turbo.yaml"
        )
        sd_path = get_folder_paths("diffusers")[0]
        ckpt_path = get_full_path("invsr", "noise_predictor_sd_turbo_v5.pth")

        args = Namespace(
            bs=batch_size,
            chopping_bs=chopping_batch_size,
            timesteps=None,
            num_steps=num_steps,
            cfg_path=cfg_path,
            sd_path=sd_path,
            started_ckpt_path=ckpt_path,
            tiled_vae=base_sampler.configs.tiled_vae,
            color_fix=color_fix,
            chopping_size=chopping_size,
        )
        configs = get_configs(args)
        base_sampler.configs = get_configs(args, log=True)
        base_sampler.setup_seed(seed)
        sampler = InvSamplerSR(base_sampler)

        images_bchw = images.permute(0,3,1,2)
        batches = split_tensor_into_batches(images_bchw, batch_size)

        results = []
        pbar = ProgressBar(len(batches))

        for batch in batches:
            result = sampler.inference(image_bchw=batch)
            results.append(torch.from_numpy(result))
            pbar.update(1)

        return (torch.cat(results, dim=0),)