"""Functions for handling, checking, and validating rasters."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_raster.ipynb.

# %% auto 0
__all__ = ['warp_raster', 'clip_raster_with_shape', 'set_nodata', 'fill_nodata_with_gdal', 'suggest_nodata',
           'check_raster_nodata', 'check_nodata_in_extent', 'check_projection', 'check_raster', 'print_raster_info',
           'get_unique_classes', 'load_lulc', 'lulc_pixel_stats', 'get_coarsest_resolution']

# %% ../nbs/01_raster.ipynb 4
import rasterio
import numpy as np
import pandas as pd
import geopandas as gpd
import yaml
import os

# %% ../nbs/01_raster.ipynb 5
from muir.core import *
from muir.visualization import *

# %% ../nbs/01_raster.ipynb 8
import rasterio
from rasterio.warp import calculate_default_transform, reproject, Resampling

# adapted from https://rasterio.readthedocs.io/en/latest/topics/reproject.html
def warp_raster(input_path, output_path, dst_crs, resampling=Resampling.nearest):
    with rasterio.open(input_path) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds)
        kwargs = src.meta.copy()
        kwargs.update({
            'crs': dst_crs,
            'transform': transform,
            'width': width,
            'height': height
        })
        with rasterio.open(output_path, 'w', **kwargs) as dst:
            for i in range(1, src.count + 1):
                reproject(
                    source=rasterio.band(src, i),
                    destination=rasterio.band(dst, i),
                    src_transform=src.transform,
                    src_crs=src.crs,
                    dst_transform=transform,
                    dst_crs=dst_crs,
                    resampling=resampling
                )

# %% ../nbs/01_raster.ipynb 9
import fiona
import rasterio
import rasterio.mask

# adapted from https://rasterio.readthedocs.io/en/latest/topics/masking-by-shapefile.html
def clip_raster_with_shape(
    input_path: str,         # Path to the input raster
    shapefile_path: str,     # Path to the shapefile
    output_path: str         # Path to save the clipped raster
):
    with fiona.open(shapefile_path, "r") as shapefile:
        shapes = [feature["geometry"] for feature in shapefile]

    with rasterio.open(input_path) as src:
        out_image, out_transform = rasterio.mask.mask(src, shapes, crop=True)
        out_meta = src.meta.copy()

    out_meta.update({
        "driver": "GTiff",
        "height": out_image.shape[1],
        "width": out_image.shape[2],
        "transform": out_transform
    })

    with rasterio.open(output_path, "w", **out_meta) as dest:
        dest.write(out_image)
    
    return output_path

# %% ../nbs/01_raster.ipynb 10
from muir import SKIP_SYMBOL

def set_nodata(
        input_path:str,  # input raster
        output_path:str, # save new raster to this file
        nodata # nodata value or none
):
    if nodata is None: # skip 
        return SKIP_SYMBOL
    with rasterio.open(input_path) as src:
        meta = src.meta.copy()
        meta.update(nodata=nodata)
        data = src.read()
    with rasterio.open(output_path, "w", **meta) as dst:
        dst.write(data)

# %% ../nbs/01_raster.ipynb 11
from osgeo import gdal

def fill_nodata_with_gdal(
        input_path, 
        output_path,
        max_search_dist=100,
        smoothing_iterations=0
):
    """
    Fill NoData values in a raster using GDAL's nearest neighbor algorithm.
    """
    # Open input raster
    src_ds = gdal.Open(input_path, gdal.GA_ReadOnly)
    band = src_ds.GetRasterBand(1)
    nodata = band.GetNoDataValue()

    # Create a copy for output
    driver = gdal.GetDriverByName('GTiff')
    dst_ds = driver.CreateCopy(output_path, src_ds, 0)
    dst_band = dst_ds.GetRasterBand(1)

    # Fill NoData using nearest neighbor
    gdal.FillNodata(
        targetBand=dst_band, 
        maskBand=None, 
        maxSearchDist=max_search_dist, 
        smoothingIterations=smoothing_iterations
    )

    # Set NoData value
    if nodata is not None:
        dst_band.SetNoDataValue(nodata)

    dst_band.FlushCache()
    dst_ds.FlushCache()
    dst_ds = None
    src_ds = None

# %% ../nbs/01_raster.ipynb 14
import pygeoprocessing as pyg

def suggest_nodata(
        input_path: str # raster path
):
    dtype = pyg.geoprocessing.choose_dtype(input_path)
    suggested_nodata = pyg.geoprocessing.choose_nodata(dtype)
    return suggested_nodata

# %% ../nbs/01_raster.ipynb 15
from muir.core import pathname

def check_raster_nodata(
        rasters:str|list # raster paths to check
):
    """
    Check if raster(s) have a nodata value set.
    Accepts a single path (str), a list of paths, or a dict of paths.
    Prints a status message for each raster.
    """
    def check_one(path):
        with rasterio.open(path) as src:
            if src.nodata is None:
                print(f"❌ No Nodata set for: {pathname(path)}")
            else:
                print(f"✅ Nodata set for: {pathname(path)}")

    if isinstance(rasters, str):
        check_one(rasters)
    elif isinstance(rasters, list):
        for path in rasters:
            check_one(path)

# %% ../nbs/01_raster.ipynb 16
import rasterio
from rasterio.mask import mask
import geopandas as gpd
import numpy as np

# TODO: not working as properly, also considers stuff outside the shapefile
def check_nodata_in_extent(rasters, shapefile_path):
    # Read the shapefile geometry
    shapes = gpd.read_file(shapefile_path)
    geoms = shapes.geometry.values

    def _check_one(input_path):
        # Open the raster and mask it with the shape
        with rasterio.open(input_path) as src:
            out_image, out_transform = mask(src, geoms, crop=True)
            nodata = src.nodata

        # Check for nodata in the masked array
        nodata_count = np.sum(out_image == nodata)
        total_count = out_image.size

        # 
        if nodata_count == 0:
            print(f"✅ No `NoData` values in extent: {pathname(input_path)}")
        else:
            print(f"❌ NoData in extent: {pathname(input_path)}")
            print(f"\tNoData Count: {nodata_count}")
            print(f"\tTotal  Count: {total_count}")
            print(f"\tFrctn NoData: {(nodata_count / total_count):.2f}")

    if isinstance(rasters, str):
        _check_one(rasters)
    elif isinstance(rasters, list):
        for input_path in rasters:
            _check_one(rasters)

# %% ../nbs/01_raster.ipynb 17
from pathlib import Path
def check_projection(projection, paths):
    """
    Checks if each raster in paths has the specified projection (CRS).
    Prints a status message with emojis and the pathname for each raster.
    """
    def check_one(path):
        with rasterio.open(path) as src:
            crs_str = src.crs.to_string() if src.crs else None
            if crs_str == projection:
                print(f"✅ Projection OK for: {pathname(path)}")
            else:
                print(f"❌ Projection mismatch for: {pathname(path)} (found: {crs_str})")

    if isinstance(paths, str) or isinstance(paths, Path):
        check_one(paths)
    elif isinstance(paths, list):
        for path in paths:
            check_one(path)
    elif isinstance(paths, dict):
        for key, path in paths.items():
            check_one(path)


# %% ../nbs/01_raster.ipynb 21
def check_raster(
    raster_path,           # Path to raster file
    projection=None,       # Projection string to check (optional)
    aoi_path=None          # Path to AOI shapefile (optional)
):
    """
    Runs a suite of checks on a raster:
      - Checks if nodata is set
      - Checks for nodata values inside AOI (if provided)
      - Checks projection (if provided)
    """
    print(f"--- Checking raster: {pathname(raster_path)} ---")
    check_raster_nodata(raster_path)
    if projection is not None:
        check_projection(projection, raster_path)
    if aoi_path is not None:
        check_nodata_in_extent(raster_path, aoi_path)

# %% ../nbs/01_raster.ipynb 23
def print_raster_info(input_path):
    with rasterio.open(input_path) as src:
        print("CRS:", src.crs)
        print("Width, Height:", src.width, src.height)
        print("Bounds:", src.bounds)
        print("Transform:", src.transform)
        print("Number of bands:", src.count)
        print("Data type:", src.dtypes)
        print("NoData value:", src.nodata)
        print("Resolution:", src.res)

# %% ../nbs/01_raster.ipynb 25
def get_unique_classes(raster_path):
    with rasterio.open(raster_path) as src:
        raster_data = src.read(1)
        unique_classes = np.unique(raster_data)
    return unique_classes

# %% ../nbs/01_raster.ipynb 30
def load_lulc(lulc_tif_path, lulc_clr_path):
    unique_classes = get_unique_classes(lulc_tif_path)
    class_mapping = parse_clr(lulc_clr_path)
    relevant_classes = class_mapping[class_mapping['class'].isin(unique_classes)].reset_index(drop=True)
    return relevant_classes

# %% ../nbs/01_raster.ipynb 33
def lulc_pixel_stats(
    raster_path,
    include_nodata = False
):
    """
    Returns a DataFrame with the number of pixels and area (in square meters) for each unique class in the raster.
    Excludes nodata pixels unless include_nodata=True.
    """
    with rasterio.open(raster_path) as src:
        data = src.read(1)
        pixel_area = abs(src.transform.a) * abs(src.transform.e)  # pixel width * pixel height
        nodata = src.nodata

    if not include_nodata and nodata is not None:
        data = data[data != nodata]

    classes, counts = np.unique(data, return_counts=True)
    stats = []
    for cls, count in zip(classes, counts):
        area = count * pixel_area
        stats.append({
            'class': cls,
            'pixel_count': count,
            'area_m2': area
        })
    return pd.DataFrame(stats)


# %% ../nbs/01_raster.ipynb 38
def get_coarsest_resolution(
    raster_paths: list[str]  # List of raster file paths
):
    """
    Find the coarsest (largest) pixel size among all rasters.
    Ensures all rasters have the same CRS. Returns (resolution, unit).
    """
    crs_set = set()
    resolutions = []
    for rpath in raster_paths:
        with rasterio.open(rpath) as src:
            crs_set.add(src.crs.to_string())
            res = max(src.res)
            resolutions.append(res)
    if len(crs_set) > 1:
        raise ValueError(f"Rasters have different CRS: {crs_set}")
    crs = list(crs_set)[0]
    # Determine unit
    if "degree" in crs.lower() or "longlat" in crs.lower() or crs.startswith("EPSG:4326"):
        unit = "degree"
    else:
        unit = "meter"
    return max(resolutions), unit
