import torch
from . import spipf
from torch import Tensor
from typing import TypeAlias, Callable
from math import sqrt
from copy import copy


Resampler: TypeAlias = Callable[[Tensor, tuple[int, int], tuple[int, int]], Tensor]
Scaler: TypeAlias = Callable[[tuple[int, int]], tuple[int, int]]


class ResampleImage:
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "image": ("IMAGE", {}),
            "scaler": ("SCALER", {}),
            "resampler": ("RESAMPLER", {}),
        }}

    RETURN_TYPES = ("IMAGE", )
    RETURN_NAMES = ("image", )
    CATEGORY = "image"
    FUNCTION = "f" 

    @staticmethod
    def f(image: Tensor, scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
        _, h, w, _ = image.shape
        new_res = scaler((h, w))
        linear = spipf.colorspace_srgb_linear_from_gamma(image)
        resampled_linear = resampler(linear, new_res, (1, 2))
        resampled = spipf.colorspace_srgb_gamma_from_linear(resampled_linear)
        return (resampled, )


class ResampleMask:
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "mask": ("MASK", {}),
            "scaler": ("SCALER", {}),
            "resampler": ("RESAMPLER", {}),
        }}

    RETURN_TYPES = ("MASK", )
    RETURN_NAMES = ("mask", )
    CATEGORY = "mask"
    FUNCTION = "f" 

    @staticmethod
    def f(mask: Tensor, scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
        _, h, w = mask.shape
        new_res = scaler((h, w))
        resampled = resampler(mask, new_res, (1, 2))
        return (resampled, )


class ResampleLatent:
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "latent": ("LATENT", {}),
            "scaler": ("SCALER", {}),
            "resampler": ("RESAMPLER", {}),
        }}

    RETURN_TYPES = ("LATENT", )
    RETURN_NAMES = ("latent", )
    CATEGORY = "latent"
    FUNCTION = "f" 

    @staticmethod
    def f(latent: dict[str, Tensor], scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
        samples = latent["samples"]
        _, _, h, w = samples.shape
        new_res = scaler((h, w))
        resampled = resampler(samples, new_res, (2, 3))
        new_latent = copy(latent)
        new_latent["samples"] = resampled
        return (resampled, )


class ResamplerBase:
    RETURN_TYPES = ("RESAMPLER", )
    RETURN_NAMES = ("resampler", )
    CATEGORY = "resampling"
    FUNCTION = "f" 


class ResamplerNearestNeighbor(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {}}

    @staticmethod
    def f() -> tuple[Resampler]:
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            resampled = spipf.resample_nearest_neighbor_2d(tensor, resolution, dims)
            return resampled
        return (resampler, )


class ResamplerTriangle(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "radius": ("INT", {"min": 1, "step": 1, "default": 1})
        }}

    @staticmethod
    def f(radius: int) -> tuple[Resampler]:
        def filter(x: Tensor) -> Tensor:
            return spipf.window_triangle(x, radius)
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            resampled = spipf.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
            return resampled
        return (resampler, )


class ResamplerLanczos(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "radius": ("INT", {"min": 1, "step": 1, "default": 3})
        }}

    @staticmethod
    def f(radius: int) -> tuple[Resampler]:
        def filter(x: Tensor) -> Tensor:
            return torch.sinc(x) * spipf.window_lanczos(x, radius)
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            resampled = spipf.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
            return resampled
        return (resampler, )


class ResamplerMitchellNetravali(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "b": ("FLOAT", {"step": 0.01, "default": 0.33}),
            "c": ("FLOAT", {"step": 0.01, "default": 0.33}),
        }}

    @staticmethod
    def f(b: float, c: float) -> tuple[Resampler]:
        radius = spipf.window_mitchell_netravali_radius()
        def filter(x: Tensor) -> Tensor:
            return spipf.window_mitchell_netravali(x, b, c)
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            resampled = spipf.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
            return resampled
        return (resampler, )


class ResamplerArea(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {}}

    @staticmethod
    def f() -> tuple[Resampler]:
        radius = spipf.window_area_radius()
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            _, h, w, _ = tensor.shape
            H, W = resolution
            def filter_h(x: Tensor) -> Tensor:
                return spipf.window_area(x, h, H)
            def filter_w(x: Tensor) -> Tensor:
                return spipf.window_area(x, w, W)
            resampled = spipf.resample_filter_2d_separable(tensor, resolution, radius, (filter_h, filter_w), dims)
            return resampled
        return (resampler, )


class ResamplerJincLanczos(ResamplerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "radius": ("INT", {"min": 1, "step": 1, "default": 3})
        }}

    @staticmethod
    def f(radius: int) -> tuple[Resampler]:
        def filter(x: Tensor) -> Tensor:
            return spipf.special_jinc(x) * spipf.window_lanczos(x, radius)
        def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
            resampled = spipf.resample_filter_2d(tensor, resolution, radius, filter, dims)
            return resampled
        return (resampler, )



class ScalerBase:
    RETURN_TYPES = ("SCALER", )
    RETURN_NAMES = ("scaler", )
    CATEGORY = "resampling/scaler"
    FUNCTION = "f" 


class ScalerSide(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "side_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
        }}

    @staticmethod
    def f(side_multiplier: float) -> tuple[Scaler]:
        def scaler(res: tuple[int, int]) -> tuple[int, int]:
            scaled = tuple((int(round(x * side_multiplier)) for x in res))
            return scaled
        return (scaler, )


class ScalerArea(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "area_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
        }}

    @staticmethod
    def f(area_multiplier: float) -> tuple[Scaler]:
        side_multiplier = sqrt(area_multiplier)
        def scaler(res: tuple[int, int]) -> tuple[int, int]:
            scaled = tuple((int(round(x * side_multiplier)) for x in res))
            return scaled
        return (scaler, )


class ScalerUnlinked(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "width_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
            "height_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
        }}

    @staticmethod
    def f(width_multiplier: float, height_multiplier: float) -> tuple[Scaler]:
        def scaler(res: tuple[int, int]) -> tuple[int, int]:
            h, w = res
            scaled = (height_multiplier * h, width_multiplier * w)
            return scaled
        return (scaler, )


class ScalerPixelDeltas(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "width_delta": ("INT", {"default": 0}),
            "height_delta": ("INT", {"default": 0}),
        }}

    @staticmethod
    def f(width_multiplier: float, height_multiplier: float) -> tuple[Scaler]:
        def scaler(res: tuple[int, int]) -> tuple[int, int]:
            h, w = res
            scaled = (h + height_delta, w + width_delta)
            return scaled
        return (scaler, )


class ScalerFixed(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "width": ("INT", {"default": 1024, "min": 1}),
            "height": ("INT", {"default": 1024, "min": 1}),
        }}

    @staticmethod
    def f(width: int, height: int) -> tuple[Scaler]:
        def scaler(_: tuple[int, int]) -> tuple[int, int]:
            scaled = (height, width)
            return scaled
        return (scaler, )


class ScalerMegapixels(ScalerBase):
    @staticmethod
    def INPUT_TYPES():
        return {"required": {
            "megapixels": ("FLOAT", {"default": 1.0, "min": 0.001, "step": 0.001}),
        }}

    @staticmethod
    def f(megapixels: float) -> tuple[Scaler]:
        def scaler(res: tuple[int, int]) -> tuple[int, int]:
            h, w = res
            cur_megapixels = h * w / 1_000_000
            area_mult = megapixels / cur_megapixels
            side_mult = sqrt(area_mult)
            scaled = tuple((int(round(side_mult * x)) for x in res))
            return scaled
        return (scaler, )


NODE_CLASS_MAPPINGS = {
    "Resample Image"                 : ResampleImage,
    "Resample Mask"                  : ResampleMask,
    "Resample Latent"                : ResampleLatent,
    "Resampler | Nearest-Neighbor"   : ResamplerNearestNeighbor,
    "Resampler | Triangle"           : ResamplerTriangle,
    "Resampler | Lanczos"            : ResamplerLanczos,
    "Resampler | Jinc-Lanczos"       : ResamplerJincLanczos,
    "Resampler | Mitchell-Netravali" : ResamplerMitchellNetravali,
    "Resampler | Area"               : ResamplerArea,
    "Scaler | Side"                  : ScalerSide,
    "Scaler | Area"                  : ScalerArea,
    "Scaler | Fixed"                 : ScalerFixed,
    "Scaler | Sides Unlinked"        : ScalerUnlinked,
    "Scaler | Pixel Deltas"          : ScalerPixelDeltas,
    "Scaler | Megapixels"            : ScalerMegapixels,
}

__all__ = ["NODE_CLASS_MAPPINGS"]