"""
Multi-index water extraction from multispectral satellite imagery.

Implements NDWI, MNDWI, and AWEI indices with automated thresholding
(Otsu's method) and ensemble voting for robust water classification.
"""

import numpy as np


EPS = 1e-10


def compute_ndwi(green_band, nir_band):
    """
    Normalized Difference Water Index.
    NDWI = (Green - NIR) / (Green + NIR)
    Values > 0 indicate water (McFeeters, 1996).
    """
    g = green_band.astype('float32')
    n = nir_band.astype('float32')
    return (g - n) / (g + n + EPS)


def compute_mndwi(green_band, swir1_band):
    """
    Modified Normalized Difference Water Index.
    MNDWI = (Green - SWIR1) / (Green + SWIR1)
    Better at suppressing built-up noise (Xu, 2006).
    """
    g = green_band.astype('float32')
    s = swir1_band.astype('float32')
    return (g - s) / (g + s + EPS)


def compute_aweish(green_band, nir_band, swir1_band,
                   swir2_band=None):
    """
    Automated Water Extraction Index (shadow version).
    AWEIsh = Blue + 2.5*Green - 1.5*(NIR + SWIR1) - 0.25*SWIR2
    Optimised for areas with shadows (Feyisa et al., 2014).
    """
    g = green_band.astype('float32')
    n = nir_band.astype('float32')
    s1 = swir1_band.astype('float32')
    result = g - 0.5 * (n + s1)
    return result


def compute_aweinsh(green_band, nir_band, swir1_band):
    """
    AWEI without shadow correction.
    AWEInsh = 4*(Green - SWIR1) - (0.25*NIR + 2.75*SWIR1)
    Better for non-shadow areas (Feyisa et al., 2014).
    """
    g = green_band.astype('float32')
    n = nir_band.astype('float32')
    s1 = swir1_band.astype('float32')
    return 4 * (g - s1) - (0.25 * n + 2.75 * s1)


def otsu_threshold(image_array):
    """
    Compute Otsu threshold for binary classification.

    Uses numpy histogram approach. Falls back to 0.0 if
    intensity range is too narrow.
    """
    data = image_array[~np.isnan(image_array)]
    if data.size == 0:
        return 0.0

    data_min = data.min()
    data_max = data.max()

    if data_max - data_min < EPS:
        return data_min

    hist, bin_edges = np.histogram(
        data, bins=256,
        range=(float(data_min), float(data_max))
    )
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0

    total_pixels = hist.sum().astype('float64')
    sum_total = (bin_centers * hist.astype('float64')).sum()

    sum_b = 0.0
    w_b = 0.0
    var_max = 0.0
    threshold = bin_centers[0]

    for i in range(len(hist)):
        w_b += hist[i]
        if w_b < EPS:
            continue
        w_f = total_pixels - w_b
        if w_f < EPS:
            break

        sum_b += hist[i] * bin_centers[i]
        m_b = sum_b / w_b
        m_f = (sum_total - sum_b) / w_f

        var_between = w_b * w_f * (m_b - m_f) ** 2
        if var_between > var_max:
            var_max = var_between
            threshold = bin_centers[i]

    return threshold


def classify_index(index_array, method='otsu', fixed_threshold=None):
    """
    Binary classification of a water index.

    Parameters
    ----------
    index_array : np.ndarray
        Water index values (e.g., from compute_ndwi).
    method : str
        'otsu', 'fixed', or 'percentile'.
    fixed_threshold : float or None
        Used if method='fixed'.

    Returns
    -------
    binary_mask : np.ndarray (bool)
        True where water is detected.
    threshold : float
        The threshold value used.
    """
    if method == 'fixed':
        thresh = fixed_threshold if fixed_threshold is not None else 0.0
    elif method == 'percentile':
        valid = index_array[~np.isnan(index_array)]
        if valid.size > 0:
            thresh = np.percentile(valid, 85)
        else:
            thresh = 0.0
    else:
        thresh = otsu_threshold(index_array)

    binary = index_array > thresh
    return binary, thresh


def ensemble_water(ndwi_mask, mndwi_mask, awei_mask=None,
                   method='vote'):
    """
    Combine multiple water classifications into a single mask.

    Parameters
    ----------
    ndwi_mask, mndwi_mask, awei_mask : np.ndarray (bool)
    method : str
        'vote' - majority voting (>=2 out of 3)
        'union' - any positive (OR)
        'intersection' - all positive (AND)

    Returns
    -------
    np.ndarray (bool)
    """
    masks = [ndwi_mask, mndwi_mask]
    if awei_mask is not None:
        masks.append(awei_mask)

    if method == 'union':
        return np.logical_or.reduce(masks)
    elif method == 'intersection':
        return np.logical_and.reduce(masks)
    else:
        stacked = np.stack(masks, axis=0).astype('int32')
        return np.sum(stacked, axis=0) >= (len(masks) + 1) // 2


def polygonize_mask(binary_mask, transform, crs):
    """
    Convert binary water mask to vector polygons.

    Parameters
    ----------
    binary_mask : np.ndarray (bool)
    transform : rasterio.Affine
    crs : rasterio.CRS or str

    Returns
    -------
    geopandas.GeoDataFrame or None
    """
    try:
        import rasterio.features
        from shapely.geometry import shape
        import geopandas as gpd

        shapes = rasterio.features.shapes(
            binary_mask.astype('uint8'),
            mask=binary_mask,
            transform=transform,
        )
        polygons = [
            shape(geom)
            for geom, val in shapes
            if val == 1
        ]
        if not polygons:
            return None

        from shapely.ops import unary_union
        dissolved = unary_union(polygons)
        return gpd.GeoDataFrame(
            {'geometry': [dissolved]},
            crs=str(crs)
        )
    except ImportError:
        return None


def extract_water_pipeline(green, nir, swir1,
                           threshold_method='otsu',
                           ensemble_method='vote',
                           fixed_threshold=None,
                           transform=None, crs=None):
    """
    End-to-end water extraction from band arrays.

    Parameters
    ----------
    green, nir, swir1 : np.ndarray
    threshold_method : str
    ensemble_method : str
    fixed_threshold : float or None
    transform, crs : for polygonization

    Returns
    -------
    dict with 'ndwi', 'mndwi', 'mask', 'thresholds', 'gdf'
    """
    ndwi = compute_ndwi(green, nir)
    mndwi = compute_mndwi(green, swir1)
    awei = compute_aweish(green, nir, swir1)

    ndwi_mask, t1 = classify_index(
        ndwi, threshold_method, fixed_threshold
    )
    mndwi_mask, t2 = classify_index(
        mndwi, threshold_method, fixed_threshold
    )
    awei_mask, t3 = classify_index(
        awei, threshold_method, fixed_threshold
    )

    final_mask = ensemble_water(
        ndwi_mask, mndwi_mask, awei_mask,
        method=ensemble_method
    )

    gdf = None
    if transform is not None and crs is not None:
        gdf = polygonize_mask(final_mask, transform, crs)

    return {
        'ndwi': ndwi,
        'mndwi': mndwi,
        'awei': awei,
        'ndwi_threshold': t1,
        'mndwi_threshold': t2,
        'awei_threshold': t3,
        'mask': final_mask,
        'gdf': gdf,
    }
