# © 2025 rndnanthu – Licensed under CC BY-NC 4.0
import os, torch, numpy as np, cv2, json
import imageio

# Optional EXR support
try:
    import OpenEXR, Imath
    exr_avail = True
except ImportError:
    exr_avail = False

from .log_transforms import (
    srgb_to_linear,
    rec709_to_linear,
    linear_to_logc3,
    linear_to_clog2,
    linear_to_slog3,
    linear_to_proreslog,
)

BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
OCIO_FILE = os.path.join(BASE_PATH, "assets", "presets", "ocio_sim_config.json")


class LogReconstructionNode:
    @classmethod
    def INPUT_TYPES(cls):
        with open(OCIO_FILE, "r") as f:
            cfg = json.load(f)
        return {
            "required": {
                "image": ("IMAGE",),
                "input_space": (cfg["input_spaces"], {"default": cfg["default_input_space"]}),
                "log_transform": (list(cfg["log_transforms"].keys()), {"default": cfg["default_log_transform"]}),
                "use_inverse_tonemap": ("BOOLEAN", {"default": True}),
                "exposure_index": ("FLOAT", {"default": 800.0, "min": 100.0, "max": 6400.0}),
                "auto_analyze": ("BOOLEAN", {"default": True}),
                "local_gain_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
            },
            "optional": {
                "depth_map": ("IMAGE",),
                "normal_map": ("IMAGE",),
                "save_exr": ("BOOLEAN", {"default": False}),
                "exr_filename": ("STRING", {"default": "output.exr"}),
                "save_dpx": ("BOOLEAN", {"default": False}),
                "dpx_filename": ("STRING", {"default": "output.dpx"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "METADATA")
    RETURN_NAMES = ("log_image", "metadata")
    FUNCTION = "run"
    CATEGORY = "rndnanthu/🎨Color Tools"

    def analyze_image(self, img):
        return {
            "mean": float(np.mean(img)),
            "std": float(np.std(img)),
            "histogram": np.histogram(img, bins=64, range=(0, 1))[0].tolist()
        }

    def apply_inverse_tonemap(self, img):
        img = np.clip(img, 0.0, 1.0)
        return np.power(img, 2.0) * 1.2  # Soft shoulder boost

    def inject_local_gain(self, img, depth_map, normal_map, strength):
        H, W, _ = img.shape
        gain = np.ones((H, W), dtype=np.float32)

        if depth_map is not None:
            d = depth_map[0].cpu().numpy()
            if d.ndim == 3 and d.shape[2] == 3:
                d = cv2.cvtColor(d, cv2.COLOR_RGB2GRAY)
            d = cv2.resize(d, (W, H), interpolation=cv2.INTER_LINEAR)
            gain += (1.0 - d) * strength  # Near = more boost

        if normal_map is not None:
            n = normal_map[0].cpu().numpy()
            if n.ndim == 3 and n.shape[2] == 3:
                g = cv2.cvtColor(n, cv2.COLOR_RGB2GRAY)
            else:
                g = n
            g = cv2.resize(g, (W, H), interpolation=cv2.INTER_LINEAR)
            gain += (1.0 - g) * (strength * 0.5)

        # Apply gain to all 3 channels
        return np.clip(img * gain[..., None], 0, 1)

    def export_exr(self, img, path):
        if not exr_avail:
            return
        h, w, _ = img.shape
        hdr = OpenEXR.Header(w, h)
        pt = Imath.PixelType(Imath.PixelType.FLOAT)
        hdr['channels'] = {'R': Imath.Channel(pt), 'G': Imath.Channel(pt), 'B': Imath.Channel(pt)}
        ex = OpenEXR.OutputFile(path, hdr)
        r, g, b = [img[..., i].astype(np.float32).tobytes() for i in range(3)]
        ex.writePixels({'R': r, 'G': g, 'B': b})
        ex.close()

    def run(self, image, input_space, log_transform, use_inverse_tonemap,
            exposure_index, auto_analyze, local_gain_strength,
            depth_map=None, normal_map=None,
            save_exr=False, exr_filename=None, save_dpx=False, dpx_filename=None):

        cfg = json.load(open(OCIO_FILE))

        linearize_func_map = {
            "sRGB": srgb_to_linear,
            "Rec.709": rec709_to_linear,
            "ACEScg": lambda x: x
        }
        linearize = linearize_func_map.get(input_space, lambda x: x)

        log_func_map = {
            "Alexa_LogC3": linear_to_logc3,
            "Canon_Log2": linear_to_clog2,
            "Sony_SLog3": linear_to_slog3,
            "ProRes_Log": linear_to_proreslog,
        }
        logf = log_func_map.get(log_transform, lambda x: x)

        # Convert tensor to numpy
        img = image[0].cpu().numpy().astype(np.float32)
        img = np.clip(img, 0, 1)

        # 1. Linearize color space
        lin = linearize(img)
        metadata = {}

        # 2. Auto exposure using mean luminance
        if auto_analyze:
            stats = self.analyze_image(lin)
            mean = stats["mean"]
            if mean > 0:
                gain = 0.18 / mean
                lin *= gain
                metadata["auto_gain"] = float(gain)

        # 3. Apply EI scaling
        lin *= (exposure_index / 800.0)

        # 4. Apply local gain if maps present
        has_depth = depth_map is not None and isinstance(depth_map, (list, tuple, torch.Tensor))
        has_normal = normal_map is not None and isinstance(normal_map, (list, tuple, torch.Tensor))
        if has_depth or has_normal:
            lin = self.inject_local_gain(lin, depth_map, normal_map, local_gain_strength)

        # 5. Optional shoulder boost (fake inverse tone map)
        if use_inverse_tonemap:
            lin = self.apply_inverse_tonemap(lin)

        # 6. Apply log encoding (camera simulation)
        log_img = logf(lin)

        # 7. Clamp result to [0, 1] for output
        log_img = np.clip(log_img, 0.0, 1.0)

        # 8. Metadata again after log
        if auto_analyze:
            metadata.update(self.analyze_image(log_img))

        # 9. Optional EXR/DPX export
        if save_exr:
            os.makedirs("output", exist_ok=True)
            self.export_exr(log_img, os.path.join("output", exr_filename))
        if save_dpx:
            os.makedirs("output", exist_ok=True)
            imageio.imwrite(os.path.join("output", dpx_filename or "output.dpx"),
                            (log_img * 65535).astype(np.uint16))

        # 10. Return tensor + metadata
        tensor = torch.from_numpy(log_img).unsqueeze(0).float()
        metadata.update({
            "input_space": input_space,
            "log_transform": log_transform,
            "exposure_index": exposure_index,
        })
        return (tensor, metadata)


# 🔌 ComfyUI Mappings
NODE_CLASS_MAPPINGS = {
    "LogReconstructionNode": LogReconstructionNode
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "LogReconstructionNode": "🎯 AI → Camera‑Accurate Log Simulator"
}