import numpy as np
import warnings


def rgb_to_hsv(rgb):
    with warnings.catch_warnings():
        # Suppress the specific RuntimeWarning
        warnings.filterwarnings(
            "ignore", category=RuntimeWarning, message="invalid value encountered in divide")

        # Ensure the input is a float numpy array and in the range [0, 1]
        rgb = np.asarray(rgb)

        # Separate the R, G, B channels
        r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]

        # Max and min of RGB values
        max_rgb = np.max(rgb, axis=-1)
        min_rgb = np.min(rgb, axis=-1)
        delta = max_rgb - min_rgb

        # Hue calculation
        hue = np.zeros_like(max_rgb)

        # Avoid division by zero: Only calculate where delta is non-zero
        mask = delta != 0

        # Red is max
        idx = (max_rgb == r) & mask
        hue[idx] = (60 * ((g - b) / delta % 6))[idx]

        # Green is max
        idx = (max_rgb == g) & mask
        hue[idx] = (60 * ((b - r) / delta + 2))[idx]

        # Blue is max
        idx = (max_rgb == b) & mask
        hue[idx] = (60 * ((r - g) / delta + 4))[idx]

        # Saturation calculation
        saturation = np.zeros_like(max_rgb)
        saturation[mask] = delta[mask] / max_rgb[mask]

        # Handle the edge case where max_rgb is 0 (black pixels)
        saturation[max_rgb == 0] = 0

        # Value calculation
        value = max_rgb

        # Stack the HSV components together
        hsv = np.stack((hue, saturation, value), axis=-1)
    return hsv


def hsv_to_rgb(hsv):
    # Ensure the input is a float numpy array and in the range [0, 1]
    hsv = np.asarray(hsv)

    # Separate the H, S, V channels
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]

    # Chromaticity component
    c = v * s
    x = c * (1 - np.abs((h / 60) % 2 - 1))
    m = v - c

    # Initialize RGB
    r, g, b = np.zeros_like(h), np.zeros_like(h), np.zeros_like(h)

    # Compute the RGB components based on the hue range
    h0_60 = (h >= 0) & (h < 60)
    r[h0_60], g[h0_60], b[h0_60] = c[h0_60], x[h0_60], 0

    h60_120 = (h >= 60) & (h < 120)
    r[h60_120], g[h60_120], b[h60_120] = x[h60_120], c[h60_120], 0

    h120_180 = (h >= 120) & (h < 180)
    r[h120_180], g[h120_180], b[h120_180] = 0, c[h120_180], x[h120_180]

    h180_240 = (h >= 180) & (h < 240)
    r[h180_240], g[h180_240], b[h180_240] = 0, x[h180_240], c[h180_240]

    h240_300 = (h >= 240) & (h < 300)
    r[h240_300], g[h240_300], b[h240_300] = x[h240_300], 0, c[h240_300]

    h300_360 = (h >= 300) & (h < 360)
    r[h300_360], g[h300_360], b[h300_360] = c[h300_360], 0, x[h300_360]

    # Add m to each component and stack them together
    rgb = np.stack((r + m, g + m, b + m), axis=-1)

    return rgb


def test_conversion():
    # RGB array example (range [0, 1])
    rgb = np.array([[[0.5, 0.4, 0.7], [0.1, 0.2, 0.3]],
                    [[0.8, 0.7, 0.6], [0.9, 0.1, 0.2]]])
    # rgb = np.random.rand(1000, 100, 3)
    # Convert RGB to HSV
    hsv = rgb_to_hsv(rgb)
    print(hsv.shape)

    # Convert HSV back to RGB
    rgb_back = hsv_to_rgb(hsv)
    assert np.allclose(rgb, rgb_back), "Conversion failed!"


if __name__ == "__main__":
    test_conversion()