import numpy as np
import torch


def channels_layer(images, channels, function):
    if channels == "rgb":
        img = images[:, :, :, :3]
    elif channels == "rgba":
        img = images[:, :, :, :4]
    elif channels == "rg":
        img = images[:, :, :, [0, 1]]
    elif channels == "rb":
        img = images[:, :, :, [0, 2]]
    elif channels == "ra":
        img = images[:, :, :, [0, 3]]
    elif channels == "gb":
        img = images[:, :, :, [1, 2]]
    elif channels == "ga":
        img = images[:, :, :, [1, 3]]
    elif channels == "ba":
        img = images[:, :, :, [2, 3]]
    elif channels == "r":
        img = images[:, :, :, 0]
    elif channels == "g":
        img = images[:, :, :, 1]
    elif channels == "b":
        img = images[:, :, :, 2]
    elif channels == "a":
        img = images[:, :, :, 3]
    else:
        raise ValueError("Unsupported channels")

    result = torch.from_numpy(function(img.numpy()))

    if channels == "rgb":
        images[:, :, :, :3] = result
    elif channels == "rgba":
        images[:, :, :, :4] = result
    elif channels == "rg":
        images[:, :, :, [0, 1]] = result
    elif channels == "rb":
        images[:, :, :, [0, 2]] = result
    elif channels == "ra":
        images[:, :, :, [0, 3]] = result
    elif channels == "gb":
        images[:, :, :, [1, 2]] = result
    elif channels == "ga":
        images[:, :, :, [1, 3]] = result
    elif channels == "ba":
        images[:, :, :, [2, 3]] = result
    elif channels == "r":
        images[:, :, :, 0] = result
    elif channels == "g":
        images[:, :, :, 1] = result
    elif channels == "b":
        images[:, :, :, 2] = result
    elif channels == "a":
        images[:, :, :, 3] = result

    return images


class ImageNoiseBeta:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "a": ("INT", {
                    "default": 1,
                    "min": 1,
                }),
                "b": ("INT", {
                    "default": 1,
                    "min": 1,
                }),
                "monochromatic": (["false", "true"],),
                "invert": (["false", "true"],),
                "channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "node"
    CATEGORY = "image/noise"

    def noise(self, images, a, b, monochromatic, invert):
        if monochromatic and images.shape[3] > 1:
            noise = np.random.beta(a, b, images.shape[:3])
        else:
            noise = np.random.beta(a, b, images.shape)

        if monochromatic and images.shape[3] > 1:
            noise = noise[..., np.newaxis].repeat(images.shape[3], -1)

        if invert:
            noise = images - noise
        else:
            noise = images + noise

        noise = noise.astype(images.dtype)

        return noise

    def node(self, images, a, b, monochromatic, invert, channels):
        tensor = images.clone().detach()

        monochromatic = True if monochromatic == "true" else False
        invert = True if invert == "true" else False

        return (channels_layer(tensor, channels, lambda x: self.noise(
            x, a, b, monochromatic, invert
        )),)


class ImageNoiseBinomial:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "n": ("INT", {
                    "default": 128,
                    "min": 1,
                    "max": 255,
                }),
                "p": ("FLOAT", {
                    "default": 0.5,
                    "max": 1.0,
                    "step": 0.01
                }),
                "monochromatic": (["false", "true"],),
                "invert": (["false", "true"],),
                "channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "node"
    CATEGORY = "image/noise"

    def noise(self, images, n, p, monochromatic, invert):
        if monochromatic and images.shape[3] > 1:
            noise = np.random.binomial(n, p, images.shape[:3])
        else:
            noise = np.random.binomial(n, p, images.shape)

        noise = noise.astype(images.dtype)
        noise /= 255

        if monochromatic and images.shape[3] > 1:
            noise = noise[..., np.newaxis].repeat(images.shape[3], -1)

        if invert:
            noise = images - noise
        else:
            noise = images + noise

        noise = np.clip(noise, 0.0, 1.0)

        return noise

    def node(self, images, n, p, monochromatic, invert, channels):
        tensor = images.clone().detach()

        monochromatic = True if monochromatic == "true" else False
        invert = True if invert == "true" else False

        return (channels_layer(tensor, channels, lambda x: self.noise(
            x, n, p, monochromatic, invert
        )),)


class ImageNoiseBytes:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "monochromatic": (["false", "true"],),
                "invert": (["false", "true"],),
                "channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "node"
    CATEGORY = "image/noise"

    def noise(self, images, monochromatic, invert):
        if monochromatic and images.shape[3] > 1:
            noise = np.random.bytes(np.prod(images.shape[:3]))
            noise = np.frombuffer(noise, np.uint8)
            noise = np.reshape(noise, images.shape[:3])
        else:
            noise = np.random.bytes(np.prod(images.shape))
            noise = np.frombuffer(noise, np.uint8)
            noise = np.reshape(noise, images.shape)

        noise = noise.astype(images.dtype)
        noise /= 255

        if monochromatic and images.shape[3] > 1:
            noise = noise[..., np.newaxis].repeat(images.shape[3], -1)

        if invert:
            noise = images - noise
        else:
            noise = images + noise

        noise = np.clip(noise, 0.0, 1.0)

        return noise

    def node(self, images, monochromatic, invert, channels):
        tensor = images.clone().detach()

        monochromatic = True if monochromatic == "true" else False
        invert = True if invert == "true" else False

        return (channels_layer(tensor, channels, lambda x: self.noise(
            x, monochromatic, invert
        )),)


class ImageNoiseGaussian:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "strength": ("FLOAT", {
                    "default": 0.5,
                    "step": 0.01
                }),
                "monochromatic": (["false", "true"],),
                "invert": (["false", "true"],),
                "channels": (["rgb", "rgba", "rg", "rb", "ra", "gb", "ga", "ba", "r", "g", "b", "a"],),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "node"
    CATEGORY = "image/noise"

    def noise(self, images, strength, monochromatic, invert):
        if monochromatic and images.shape[3] > 1:
            noise = np.random.normal(0, 1, images.shape[:3])
        else:
            noise = np.random.normal(0, 1, images.shape)

        noise = np.abs(noise)
        noise /= noise.max()

        if monochromatic and images.shape[3] > 1:
            noise = noise[..., np.newaxis].repeat(images.shape[3], -1)

        if invert:
            noise = images - noise * strength
        else:
            noise = images + noise * strength

        noise = np.clip(noise, 0.0, 1.0)
        noise = noise.astype(images.dtype)

        return noise

    def node(self, images, strength, monochromatic, invert, channels):
        tensor = images.clone().detach()

        monochromatic = True if monochromatic == "true" else False
        invert = True if invert == "true" else False

        return (channels_layer(tensor, channels, lambda x: self.noise(
            x, strength, monochromatic, invert
        )),)


NODE_CLASS_MAPPINGS = {
    "ImageNoiseBeta": ImageNoiseBeta,
    "ImageNoiseBinomial": ImageNoiseBinomial,
    "ImageNoiseBytes": ImageNoiseBytes,
    "ImageNoiseGaussian": ImageNoiseGaussian
}