import torch
import numpy as np
import cv2
import math

from .projection import equirect_to_perspective, strip_to_equirect


class EquirectToStrip90:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "face_size": ("INT", {"default": 512, "min": 128, "max": 2048}),
                "order": ("STRING", {"default": "3,2,1,0,4,5", "multiline": False}),
                "flip_sides": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING", "BOOLEAN")
    RETURN_NAMES = ("strip", "order", "flip_sides")
    FUNCTION = "convert"
    CATEGORY = "360"

    def convert(self, image, face_size, order, flip_sides):
        img = (image[0].cpu().numpy() * 255).astype(np.uint8)

        # Parse order string to get indices
        order_list = [int(x.strip()) for x in order.split(",")]

        # All 6 views in standard order: front, left, back, right, top, bottom
        all_views = [
            ("0_front",  0,          0),
            ("1_left",  -math.pi/2, 0),
            ("2_back",   math.pi,    0),
            ("3_right",  math.pi/2,  0),
            ("4_top",    0, -math.pi/2),
            ("5_bottom", 0,  math.pi/2),
        ]

        faces = []
        for i, idx in enumerate(order_list):
            _, yaw, pitch = all_views[idx]
            face = equirect_to_perspective(img, yaw, pitch, 90, face_size)
            # Optionally flip first 4 faces (horizontal views) vertically
            if flip_sides and idx < 4:
                face = cv2.flip(face, 0)
            faces.append(face)

        strip = np.concatenate(faces, axis=1)
        strip = torch.from_numpy(strip.astype(np.float32) / 255.0).unsqueeze(0)

        return (strip, order, flip_sides)


class Strip90ToEquirect:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "strip": ("IMAGE",),
                "width": ("INT", {"default": 2048}),
                "height": ("INT", {"default": 1024}),
                "face_size": ("INT", {"default": 512}),
                "order": ("STRING", {"default": "0,1,2,3,4,5", "multiline": False}),
                "flip_sides": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "convert"
    CATEGORY = "360"

    def convert(self, strip, width, height, face_size, order, flip_sides):
        strip_img = (strip[0].cpu().numpy() * 255).astype(np.uint8)
        order_list = [int(x.strip()) for x in order.split(",")]
        out = strip_to_equirect(strip_img, width, height, face_size, order_list, flip_sides)
        out = torch.from_numpy(out.astype(np.float32) / 255.0).unsqueeze(0)
        return (out,)

# -------------------------
# ComfyUI registration
# -------------------------

NODE_CLASS_MAPPINGS = {
    "EquirectToStrip90": EquirectToStrip90,
    "Strip90ToEquirect": Strip90ToEquirect,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "EquirectToStrip90": "360 → 6×90° Strip",
    "Strip90ToEquirect": "6×90° Strip → 360",
}