from pathlib import Path

import numpy as np
import rasterio
from rasterio.transform import Affine
from PIL import Image


# ==========================================================
# Read Sentinel-2 RGB bands
# ==========================================================

def read_sentinel_rgb(
    b02_path,
    b03_path,
    b04_path
):
    """
    Read Sentinel-2 B02, B03 and B04 GeoTIFF bands.

    Band mapping:

        B04 -> Red
        B03 -> Green
        B02 -> Blue

    Returns
    -------
    rgb : numpy.ndarray
        RGB array with shape (H, W, 3)

    profile : rasterio profile
        Original geospatial metadata.
    """

    b02_path = Path(b02_path)
    b03_path = Path(b03_path)
    b04_path = Path(b04_path)

    with rasterio.open(b02_path) as src_b02:
        b02 = src_b02.read(1)
        profile = src_b02.profile.copy()
        transform = src_b02.transform
        crs = src_b02.crs
        width = src_b02.width
        height = src_b02.height

    with rasterio.open(b03_path) as src_b03:
        b03 = src_b03.read(1)

    with rasterio.open(b04_path) as src_b04:
        b04 = src_b04.read(1)

    # Make sure all bands have identical dimensions.
    if b02.shape != b03.shape or b02.shape != b04.shape:
        raise ValueError(
            "B02, B03 and B04 dimensions do not match."
        )

    # ------------------------------------------------------
    # Sentinel-2 RGB order:
    # Red   = B04
    # Green = B03
    # Blue  = B02
    # ------------------------------------------------------

    rgb = np.stack(
        [
            b04,
            b03,
            b02
        ],
        axis=-1
    )

    return rgb, profile


# ==========================================================
# Convert raw Sentinel data to visualization RGB
# ==========================================================

def normalize_rgb(
    rgb,
    lower_percentile=2,
    upper_percentile=98
):
    """
    Convert Sentinel-2 reflectance/DN values to uint8 RGB.

    Percentile normalization is performed independently
    for each channel.
    """

    rgb = rgb.astype(np.float32)

    result = np.zeros_like(
        rgb,
        dtype=np.uint8
    )

    for channel in range(3):

        band = rgb[:, :, channel]

        low = np.percentile(
            band,
            lower_percentile
        )

        high = np.percentile(
            band,
            upper_percentile
        )

        if high <= low:
            result[:, :, channel] = np.clip(
                band,
                0,
                255
            ).astype(np.uint8)

        else:

            normalized = (
                (band - low) /
                (high - low)
            )

            normalized = np.clip(
                normalized,
                0,
                1
            )

            result[:, :, channel] = (
                normalized * 255
            ).astype(np.uint8)

    return result


# ==========================================================
# GeoTIFF writer
# ==========================================================

def save_sr_geotiff(
    sr_image,
    source_profile,
    output_path,
    scale=4
):
    """
    Save a super-resolved RGB image as a GeoTIFF.

    The geographic extent remains unchanged while
    the pixel dimensions increase by `scale`.

    Parameters
    ----------
    sr_image : PIL.Image.Image or numpy.ndarray

    source_profile : dict
        Rasterio profile from the original Sentinel image.

    output_path : str or Path

    scale : int
        Super-resolution scale factor.
    """

    output_path = Path(output_path)
    output_path.parent.mkdir(
        parents=True,
        exist_ok=True
    )

    # ------------------------------------------------------
    # Convert image to numpy
    # ------------------------------------------------------

    if isinstance(
        sr_image,
        Image.Image
    ):
        rgb = np.asarray(
            sr_image.convert("RGB")
        )

    else:
        rgb = np.asarray(
            sr_image
        )

    if rgb.ndim != 3 or rgb.shape[2] != 3:
        raise ValueError(
            "sr_image must have shape (H, W, 3)."
        )

    rgb = np.clip(
        rgb,
        0,
        255
    ).astype(np.uint8)

    height, width, _ = rgb.shape

    # ------------------------------------------------------
    # Original transform
    # ------------------------------------------------------

    original_transform = source_profile[
        "transform"
    ]

    # ------------------------------------------------------
    # Scale pixel size by 1/scale
    #
    # Geographic bounds stay unchanged.
    # ------------------------------------------------------

    new_transform = (
        original_transform *
        Affine.scale(
            1.0 / scale,
            1.0 / scale
        )
    )

    # ------------------------------------------------------
    # Build output profile
    # ------------------------------------------------------

    profile = source_profile.copy()

    profile.update(
        {
            "driver": "GTiff",
            "height": height,
            "width": width,
            "count": 3,
            "dtype": "uint8",
            "transform": new_transform,
            "compress": "deflate",
            "photometric": "RGB",
            "tiled": True,
            "blockxsize": 256,
            "blockysize": 256,
        }
    )

    # Remove incompatible fields where necessary.
    profile.pop(
        "nodata",
        None
    )

    # ------------------------------------------------------
    # Write GeoTIFF
    # ------------------------------------------------------
    # Remove block settings when not using tiled GeoTIFF
    profile.pop("blockxsize", None)
    profile.pop("blockysize", None)
    profile.pop("tiled", None)

    with rasterio.open(
        output_path,
        "w",
        **profile
    ) as dst:

        dst.write(
            rgb[:, :, 0],
            1
        )

        dst.write(
            rgb[:, :, 1],
            2
        )

        dst.write(
            rgb[:, :, 2],
            3
        )

    return output_path


# ==========================================================
# Convenience function
# ==========================================================

def sentinel_rgb_from_bands(
    b02_path,
    b03_path,
    b04_path
):
    """
    Read B02/B03/B04 and return a normalized PIL RGB image
    together with the source GeoTIFF profile.
    """

    rgb, profile = read_sentinel_rgb(
        b02_path,
        b03_path,
        b04_path
    )

    rgb_uint8 = normalize_rgb(rgb)

    image = Image.fromarray(
        rgb_uint8,
        mode="RGB"
    )

    return image, profile