import math

import ee


def radians(img):
    """Converts image from degrees to radians"""
    return img.toFloat().multiply(math.pi).divide(180)


def hillshade(az, ze, slope, aspect):
    """Computes hillshade"""
    azimuth = radians(ee.Image.constant(az))
    zenith = radians(ee.Image.constant(90).subtract(ee.Image.constant(ze)))

    return (
        azimuth.subtract(aspect)
        .cos()
        .multiply(slope.sin())
        .multiply(zenith.sin())
        .add(zenith.cos().multiply(slope.cos()))
    )


def hillshadeRGB(
    image,
    elevation,
    weight=1,
    height_multiplier=5,
    azimuth=0,
    zenith=45,
    contrast=0,
    brightness=0,
    saturation=1,
    castShadows=False,
    customTerrain=False,
):
    """Styles RGB image using hillshading, mixes RGB and hillshade using HSV<->RGB transform"""

    hsv = image.visualize().unitScale(0, 255).rgbToHsv()

    z = elevation.multiply(ee.Image.constant(height_multiplier))

    terrain = ee.Algorithms.Terrain(z)
    slope = radians(terrain.select(["slope"])).resample("bicubic")
    aspect = radians(terrain.select(["aspect"])).resample("bicubic")

    if customTerrain:
        raise NotImplementedError("customTerrain argument is not implemented yet")

    hs = hillshade(azimuth, zenith, slope, aspect).resample("bicubic")

    if castShadows:
        hysteresis = True
        neighborhoodSize = 256

        hillShadow = ee.Algorithms.HillShadow(
            elevation,
            azimuth,
            ee.Number(90).subtract(zenith),
            neighborhoodSize,
            hysteresis,
        ).float()

        hillShadow = ee.Image(1).float().subtract(hillShadow)

        # opening
        # hillShadow = hillShadow.multiply(hillShadow.focal_min(3).focal_max(6))

        # cleaning
        hillShadow = hillShadow.focal_mode(3)

        # smoothing
        hillShadow = hillShadow.convolve(ee.Kernel.gaussian(5, 3))

        # transparent
        hillShadow = hillShadow.multiply(0.7)

        hs = hs.subtract(hillShadow).rename("shadow")

    intensity = hs.multiply(ee.Image.constant(weight)).add(
        hsv.select("value").multiply(ee.Image.constant(1).subtract(weight))
    )

    sat = hsv.select("saturation").multiply(saturation)

    hue = hsv.select("hue")

    result = (
        ee.Image.cat(hue, sat, intensity)
        .hsvToRgb()
        .multiply(ee.Image.constant(1).float().add(contrast))
        .add(ee.Image.constant(brightness).float())
    )

    if customTerrain:
        mask = elevation.mask().focal_min(2)

        result = result.updateMask(mask)

    return result
