# coding=utf-8
# Copyright 2025 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from typing import Tuple, Union

import numpy as np
import PIL.Image
import torch

from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.constants import (
    DECODE_ENDPOINT_FLUX,
    DECODE_ENDPOINT_HUNYUAN_VIDEO,
    DECODE_ENDPOINT_SD_V1,
    DECODE_ENDPOINT_SD_XL,
)
from diffusers.utils.remote_utils import (
    remote_decode,
)
from diffusers.video_processor import VideoProcessor

from ..testing_utils import (
    enable_full_determinism,
    slow,
    torch_all_close,
    torch_device,
)


enable_full_determinism()


class RemoteAutoencoderKLMixin:
    shape: Tuple[int, ...] = None
    out_hw: Tuple[int, int] = None
    endpoint: str = None
    dtype: torch.dtype = None
    scaling_factor: float = None
    shift_factor: float = None
    processor_cls: Union[VaeImageProcessor, VideoProcessor] = None
    output_pil_slice: torch.Tensor = None
    output_pt_slice: torch.Tensor = None
    partial_postprocess_return_pt_slice: torch.Tensor = None
    return_pt_slice: torch.Tensor = None
    width: int = None
    height: int = None

    def get_dummy_inputs(self):
        inputs = {
            "endpoint": self.endpoint,
            "tensor": torch.randn(
                self.shape,
                device=torch_device,
                dtype=self.dtype,
                generator=torch.Generator(torch_device).manual_seed(13),
            ),
            "scaling_factor": self.scaling_factor,
            "shift_factor": self.shift_factor,
            "height": self.height,
            "width": self.width,
        }
        return inputs

    def test_no_scaling(self):
        inputs = self.get_dummy_inputs()
        if inputs["scaling_factor"] is not None:
            inputs["tensor"] = inputs["tensor"] / inputs["scaling_factor"]
            inputs["scaling_factor"] = None
        if inputs["shift_factor"] is not None:
            inputs["tensor"] = inputs["tensor"] + inputs["shift_factor"]
            inputs["shift_factor"] = None
        processor = self.processor_cls()
        output = remote_decode(
            output_type="pt",
            # required for now, will be removed in next update
            do_scaling=False,
            processor=processor,
            **inputs,
        )
        assert isinstance(output, PIL.Image.Image)
        self.assertTrue(isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}")
        self.assertEqual(output.height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}")
        self.assertEqual(output.width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}")
        output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
        # Increased tolerance for Flux Packed diff [1, 0, 1, 0, 0, 0, 0, 0, 0]
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1),
            f"{output_slice}",
        )

    def test_output_type_pt(self):
        inputs = self.get_dummy_inputs()
        processor = self.processor_cls()
        output = remote_decode(output_type="pt", processor=processor, **inputs)
        assert isinstance(output, PIL.Image.Image)
        self.assertTrue(isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}")
        self.assertEqual(output.height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}")
        self.assertEqual(output.width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}")
        output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2), f"{output_slice}"
        )

    # output is visually the same, slice is flaky?
    def test_output_type_pil(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pil", **inputs)
        self.assertTrue(isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}")
        self.assertEqual(output.height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}")
        self.assertEqual(output.width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}")

    def test_output_type_pil_image_format(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pil", image_format="png", **inputs)
        self.assertTrue(isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}")
        self.assertEqual(output.height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}")
        self.assertEqual(output.width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}")
        self.assertEqual(output.format, "png", f"Expected image format `png`, got {output.format}")
        output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2), f"{output_slice}"
        )

    def test_output_type_pt_partial_postprocess(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
        self.assertTrue(isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}")
        self.assertEqual(output.height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}")
        self.assertEqual(output.width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}")
        output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2), f"{output_slice}"
        )

    def test_output_type_pt_return_type_pt(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pt", return_type="pt", **inputs)
        self.assertTrue(isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}")
        self.assertEqual(
            output.shape[2], self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[2]}"
        )
        self.assertEqual(
            output.shape[3], self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[3]}"
        )
        output_slice = output[0, 0, -3:, -3:].flatten()
        self.assertTrue(
            torch_all_close(output_slice, self.return_pt_slice.to(output_slice.dtype), rtol=1e-3, atol=1e-3),
            f"{output_slice}",
        )

    def test_output_type_pt_partial_postprocess_return_type_pt(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pt", partial_postprocess=True, return_type="pt", **inputs)
        self.assertTrue(isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}")
        self.assertEqual(
            output.shape[1], self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[1]}"
        )
        self.assertEqual(
            output.shape[2], self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[2]}"
        )
        output_slice = output[0, -3:, -3:, 0].flatten().cpu()
        self.assertTrue(
            torch_all_close(output_slice, self.partial_postprocess_return_pt_slice.to(output_slice.dtype), rtol=1e-2),
            f"{output_slice}",
        )

    def test_do_scaling_deprecation(self):
        inputs = self.get_dummy_inputs()
        inputs.pop("scaling_factor", None)
        inputs.pop("shift_factor", None)
        with self.assertWarns(FutureWarning) as warning:
            _ = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
            self.assertEqual(
                str(warning.warnings[0].message),
                "`do_scaling` is deprecated, pass `scaling_factor` and `shift_factor` if required.",
                str(warning.warnings[0].message),
            )

    def test_input_tensor_type_base64_deprecation(self):
        inputs = self.get_dummy_inputs()
        with self.assertWarns(FutureWarning) as warning:
            _ = remote_decode(output_type="pt", input_tensor_type="base64", partial_postprocess=True, **inputs)
            self.assertEqual(
                str(warning.warnings[0].message),
                "input_tensor_type='base64' is deprecated. Using `binary`.",
                str(warning.warnings[0].message),
            )

    def test_output_tensor_type_base64_deprecation(self):
        inputs = self.get_dummy_inputs()
        with self.assertWarns(FutureWarning) as warning:
            _ = remote_decode(output_type="pt", output_tensor_type="base64", partial_postprocess=True, **inputs)
            self.assertEqual(
                str(warning.warnings[0].message),
                "output_tensor_type='base64' is deprecated. Using `binary`.",
                str(warning.warnings[0].message),
            )


class RemoteAutoencoderKLHunyuanVideoMixin(RemoteAutoencoderKLMixin):
    def test_no_scaling(self):
        inputs = self.get_dummy_inputs()
        if inputs["scaling_factor"] is not None:
            inputs["tensor"] = inputs["tensor"] / inputs["scaling_factor"]
            inputs["scaling_factor"] = None
        if inputs["shift_factor"] is not None:
            inputs["tensor"] = inputs["tensor"] + inputs["shift_factor"]
            inputs["shift_factor"] = None
        processor = self.processor_cls()
        output = remote_decode(
            output_type="pt",
            # required for now, will be removed in next update
            do_scaling=False,
            processor=processor,
            **inputs,
        )
        self.assertTrue(
            isinstance(output, list) and isinstance(output[0], PIL.Image.Image),
            f"Expected `List[PIL.Image.Image]` output, got {type(output)}",
        )
        self.assertEqual(
            output[0].height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
        )
        self.assertEqual(
            output[0].width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
        )
        output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1),
            f"{output_slice}",
        )

    def test_output_type_pt(self):
        inputs = self.get_dummy_inputs()
        processor = self.processor_cls()
        output = remote_decode(output_type="pt", processor=processor, **inputs)
        self.assertTrue(
            isinstance(output, list) and isinstance(output[0], PIL.Image.Image),
            f"Expected `List[PIL.Image.Image]` output, got {type(output)}",
        )
        self.assertEqual(
            output[0].height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
        )
        self.assertEqual(
            output[0].width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
        )
        output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1),
            f"{output_slice}",
        )

    # output is visually the same, slice is flaky?
    def test_output_type_pil(self):
        inputs = self.get_dummy_inputs()
        processor = self.processor_cls()
        output = remote_decode(output_type="pil", processor=processor, **inputs)
        self.assertTrue(
            isinstance(output, list) and isinstance(output[0], PIL.Image.Image),
            f"Expected `List[PIL.Image.Image]` output, got {type(output)}",
        )
        self.assertEqual(
            output[0].height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
        )
        self.assertEqual(
            output[0].width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
        )

    def test_output_type_pil_image_format(self):
        inputs = self.get_dummy_inputs()
        processor = self.processor_cls()
        output = remote_decode(output_type="pil", processor=processor, image_format="png", **inputs)
        self.assertTrue(
            isinstance(output, list) and isinstance(output[0], PIL.Image.Image),
            f"Expected `List[PIL.Image.Image]` output, got {type(output)}",
        )
        self.assertEqual(
            output[0].height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
        )
        self.assertEqual(
            output[0].width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
        )
        output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1),
            f"{output_slice}",
        )

    def test_output_type_pt_partial_postprocess(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
        self.assertTrue(
            isinstance(output, list) and isinstance(output[0], PIL.Image.Image),
            f"Expected `List[PIL.Image.Image]` output, got {type(output)}",
        )
        self.assertEqual(
            output[0].height, self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
        )
        self.assertEqual(
            output[0].width, self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
        )
        output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
        self.assertTrue(
            torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1),
            f"{output_slice}",
        )

    def test_output_type_pt_return_type_pt(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="pt", return_type="pt", **inputs)
        self.assertTrue(isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}")
        self.assertEqual(
            output.shape[3], self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[2]}"
        )
        self.assertEqual(
            output.shape[4], self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[3]}"
        )
        output_slice = output[0, 0, 0, -3:, -3:].flatten()
        self.assertTrue(
            torch_all_close(output_slice, self.return_pt_slice.to(output_slice.dtype), rtol=1e-3, atol=1e-3),
            f"{output_slice}",
        )

    def test_output_type_mp4(self):
        inputs = self.get_dummy_inputs()
        output = remote_decode(output_type="mp4", return_type="mp4", **inputs)
        self.assertTrue(isinstance(output, bytes), f"Expected `bytes` output, got {type(output)}")


class RemoteAutoencoderKLSDv1Tests(
    RemoteAutoencoderKLMixin,
    unittest.TestCase,
):
    shape = (
        1,
        4,
        64,
        64,
    )
    out_hw = (
        512,
        512,
    )
    endpoint = DECODE_ENDPOINT_SD_V1
    dtype = torch.float16
    scaling_factor = 0.18215
    shift_factor = None
    processor_cls = VaeImageProcessor
    output_pt_slice = torch.tensor([31, 15, 11, 55, 30, 21, 66, 42, 30], dtype=torch.uint8)
    partial_postprocess_return_pt_slice = torch.tensor([100, 130, 99, 133, 106, 112, 97, 100, 121], dtype=torch.uint8)
    return_pt_slice = torch.tensor([-0.2177, 0.0217, -0.2258, 0.0412, -0.1687, -0.1232, -0.2416, -0.2130, -0.0543])


class RemoteAutoencoderKLSDXLTests(
    RemoteAutoencoderKLMixin,
    unittest.TestCase,
):
    shape = (
        1,
        4,
        128,
        128,
    )
    out_hw = (
        1024,
        1024,
    )
    endpoint = DECODE_ENDPOINT_SD_XL
    dtype = torch.float16
    scaling_factor = 0.13025
    shift_factor = None
    processor_cls = VaeImageProcessor
    output_pt_slice = torch.tensor([104, 52, 23, 114, 61, 35, 108, 87, 38], dtype=torch.uint8)
    partial_postprocess_return_pt_slice = torch.tensor([77, 86, 89, 49, 60, 75, 52, 65, 78], dtype=torch.uint8)
    return_pt_slice = torch.tensor([-0.3945, -0.3289, -0.2993, -0.6177, -0.5259, -0.4119, -0.5898, -0.4863, -0.3845])


class RemoteAutoencoderKLFluxTests(
    RemoteAutoencoderKLMixin,
    unittest.TestCase,
):
    shape = (
        1,
        16,
        128,
        128,
    )
    out_hw = (
        1024,
        1024,
    )
    endpoint = DECODE_ENDPOINT_FLUX
    dtype = torch.bfloat16
    scaling_factor = 0.3611
    shift_factor = 0.1159
    processor_cls = VaeImageProcessor
    output_pt_slice = torch.tensor([110, 72, 91, 62, 35, 52, 69, 55, 69], dtype=torch.uint8)
    partial_postprocess_return_pt_slice = torch.tensor(
        [202, 203, 203, 197, 195, 193, 189, 188, 178], dtype=torch.uint8
    )
    return_pt_slice = torch.tensor([0.5820, 0.5962, 0.5898, 0.5439, 0.5327, 0.5112, 0.4797, 0.4773, 0.3984])


class RemoteAutoencoderKLFluxPackedTests(
    RemoteAutoencoderKLMixin,
    unittest.TestCase,
):
    shape = (
        1,
        4096,
        64,
    )
    out_hw = (
        1024,
        1024,
    )
    height = 1024
    width = 1024
    endpoint = DECODE_ENDPOINT_FLUX
    dtype = torch.bfloat16
    scaling_factor = 0.3611
    shift_factor = 0.1159
    processor_cls = VaeImageProcessor
    # slices are different due to randn on different shape. we can pack the latent instead if we want the same
    output_pt_slice = torch.tensor([96, 116, 157, 45, 67, 104, 34, 56, 89], dtype=torch.uint8)
    partial_postprocess_return_pt_slice = torch.tensor(
        [168, 212, 202, 155, 191, 185, 150, 180, 168], dtype=torch.uint8
    )
    return_pt_slice = torch.tensor([0.3198, 0.6631, 0.5864, 0.2131, 0.4944, 0.4482, 0.1776, 0.4153, 0.3176])


class RemoteAutoencoderKLHunyuanVideoTests(
    RemoteAutoencoderKLHunyuanVideoMixin,
    unittest.TestCase,
):
    shape = (
        1,
        16,
        3,
        40,
        64,
    )
    out_hw = (
        320,
        512,
    )
    endpoint = DECODE_ENDPOINT_HUNYUAN_VIDEO
    dtype = torch.float16
    scaling_factor = 0.476986
    processor_cls = VideoProcessor
    output_pt_slice = torch.tensor([112, 92, 85, 112, 93, 85, 112, 94, 85], dtype=torch.uint8)
    partial_postprocess_return_pt_slice = torch.tensor(
        [149, 161, 168, 136, 150, 156, 129, 143, 149], dtype=torch.uint8
    )
    return_pt_slice = torch.tensor([0.1656, 0.2661, 0.3157, 0.0693, 0.1755, 0.2252, 0.0127, 0.1221, 0.1708])


class RemoteAutoencoderKLSlowTestMixin:
    channels: int = 4
    endpoint: str = None
    dtype: torch.dtype = None
    scaling_factor: float = None
    shift_factor: float = None
    width: int = None
    height: int = None

    def get_dummy_inputs(self):
        inputs = {
            "endpoint": self.endpoint,
            "scaling_factor": self.scaling_factor,
            "shift_factor": self.shift_factor,
            "height": self.height,
            "width": self.width,
        }
        return inputs

    def test_multi_res(self):
        inputs = self.get_dummy_inputs()
        for height in {320, 512, 640, 704, 896, 1024, 1208, 1384, 1536, 1608, 1864, 2048}:
            for width in {320, 512, 640, 704, 896, 1024, 1208, 1384, 1536, 1608, 1864, 2048}:
                inputs["tensor"] = torch.randn(
                    (1, self.channels, height // 8, width // 8),
                    device=torch_device,
                    dtype=self.dtype,
                    generator=torch.Generator(torch_device).manual_seed(13),
                )
                inputs["height"] = height
                inputs["width"] = width
                output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
                output.save(f"test_multi_res_{height}_{width}.png")


@slow
class RemoteAutoencoderKLSDv1SlowTests(
    RemoteAutoencoderKLSlowTestMixin,
    unittest.TestCase,
):
    endpoint = DECODE_ENDPOINT_SD_V1
    dtype = torch.float16
    scaling_factor = 0.18215
    shift_factor = None


@slow
class RemoteAutoencoderKLSDXLSlowTests(
    RemoteAutoencoderKLSlowTestMixin,
    unittest.TestCase,
):
    endpoint = DECODE_ENDPOINT_SD_XL
    dtype = torch.float16
    scaling_factor = 0.13025
    shift_factor = None


@slow
class RemoteAutoencoderKLFluxSlowTests(
    RemoteAutoencoderKLSlowTestMixin,
    unittest.TestCase,
):
    channels = 16
    endpoint = DECODE_ENDPOINT_FLUX
    dtype = torch.bfloat16
    scaling_factor = 0.3611
    shift_factor = 0.1159