import os
import torch
from omegaconf import OmegaConf
from torchvision.transforms import v2
from .normal_to_height import normal_to_height

# Modules from ComfyUI
import folder_paths
from comfy_api.latest import ComfyExtension, io

# Modules from ComfyUI-Chord
from chord import ChordModel
from chord.io import read_image, save_maps

def apply_padding(model, mode):
    for layer in [layer for _, layer in model.named_modules() if isinstance(layer, torch.nn.Conv2d)]:
        layer.padding_mode = mode

def apply_circular_padding(model):
    if hasattr(model, 'sd'):
        apply_padding(model.sd.vae, 'circular')
        apply_padding(model.sd.unet, 'circular')
    else:
        apply_padding(model, 'circular')

class ChordLoadModel(io.ComfyNode):
    """Node to load Chord Model"""

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="ChordLoadModel",
            display_name="Chord - Load Model",
            category="Chord",
            description="A node to load a Chord model.",
            inputs=[
                io.Combo.Input(
                    "ckpt_name",
                    options=([x for x in folder_paths.get_filename_list("checkpoints") if x.endswith("ckpt")], {})
                ),
            ],
            outputs=[io.Model.Output("model")],
        )
    
    @classmethod
    def execute(cls, ckpt_name) -> io.NodeOutput:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if type(ckpt_name) is list:
            ckpt_name = ckpt_name[0]
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
        config = OmegaConf.load(os.path.join(os.path.dirname(__file__), "chord/config/chord.yaml"))
        model = ChordModel(config)
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        try:
            model.load_state_dict(ckpt["state_dict"])
        except RuntimeError as e:
            raise RuntimeError('Failed to load model, check if the checkpoint file is correct.\n{}'.format(repr(e)))
        model.to(device)
        model.eval()
        return io.NodeOutput(model)

class ChordMaterialEstimation(io.ComfyNode):
    """Chord Material Estimation Node"""

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="ChordMaterialEstimation",
            display_name="Chord - Material Estimation",
            category="Chord",
            description="A node to estimate material maps from a texture image using the Chord model.",
            inputs=[
                io.Model.Input(
                    "model",
                    tooltip="The Chord model used to estimate material."
                ),
                io.Image.Input(
                    "image",
                ),
            ],
            outputs=[
                io.Image.Output("basecolor"),
                io.Image.Output("normal"),
                io.Image.Output("roughness"),
                io.Image.Output("metalness"),
            ]
        )
    
    @classmethod
    def execute(cls, model, image) -> io.NodeOutput:
        device = next(model.parameters()).device
        apply_circular_padding(model)
        image = image.permute(0,3,1,2).to(device)
        ori_h, ori_w = image.shape[-2:]
        x = v2.Resize(size=(1024, 1024), antialias=True)(image)
        with torch.no_grad() as no_grad, torch.autocast(device_type=device.type) as amp:
            output = model(x)
        for key in output.keys():
            output[key] = v2.Resize(size=(ori_h, ori_w), antialias=True)(output[key])
            if output[key].ndim == 4:
                output[key] = output[key].permute(0,2,3,1)
        return io.NodeOutput(output['basecolor'], output['normal'], output['roughness'], output['metalness'])
    
class ChordNormalToHeight(io.ComfyNode):
    """Integrate normal map to height map using Poisson solver with overlapping subregions."""

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="ChordNormalToHeight",
            display_name="Chord - Normal to Height",
            category="Chord",
            description="A node to convert normal map to height map using Poisson solver with overlapping subregions.",
            inputs=[
                io.Image.Input(
                    "normal",
                ),
            ],
            outputs=[
                io.Image.Output("height"), 
            ],
        )
    
    @classmethod
    def execute(cls, normal) -> io.NodeOutput:
        normal = normal.permute(0,3,1,2)
        height_var_threshold = 5e-4
        ori_h, ori_w = normal.shape[-2:]
        x = v2.Resize(size=(1024, 1024), antialias=True)(normal)
        height = normal_to_height(x)[None, None].squeeze(1)
        if height.var() < height_var_threshold and height.var() > 0:
            height = normal_to_height(x, skip_normalize_normal=True)[None, None].squeeze(1)

        height = v2.Resize(size=(ori_h, ori_w), antialias=True)(height)
        return io.NodeOutput(height)