"""
Analysis module for GNDC.

This module provides functions for analyzing compressed data,
including derivatives and vegetation indices.
"""

import numpy as np


def compute_ndvi(red, nir):
    """
    Compute Normalized Difference Vegetation Index.
    
    Args:
        red: Red band data
        nir: Near-infrared band data
    
    Returns:
        NDVI values in range [-1, 1]
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        ndvi = (nir - red) / (nir + red + 1e-8)
        ndvi = np.where(np.isfinite(ndvi), ndvi, np.nan)
    return ndvi


def compute_evi(blue, green, nir, g=2.5, c1=6, c2=7.5, l=1):
    """
    Compute Enhanced Vegetation Index.
    
    Args:
        blue: Blue band data
        green: Green band data  
        nir: Near-infrared band data
        g: Gain factor
        c1, c2: Coefficient parameters
        l: Canopy background adjustment
    
    Returns:
        EVI values
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        evi = g * (nir - red) / (nir + c1 * red - c2 * blue + l + 1e-8)
        evi = np.where(np.isfinite(evi), evi, np.nan)
    return evi


def compute_ndwi(green, nir):
    """
    Compute Normalized Difference Water Index.
    
    Args:
        green: Green band data
        nir: Near-infrared band data
    
    Returns:
        NDWI values in range [-1, 1]
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        ndwi = (green - nir) / (green + nir + 1e-8)
        ndwi = np.where(np.isfinite(ndwi), ndwi, np.nan)
    return ndwi


def compute_ndmi(nir, swir):
    """
    Compute Normalized Difference Moisture Index.
    
    Args:
        nir: Near-infrared band data
        swir: Short-wave infrared band data
    
    Returns:
        NDMI values in range [-1, 1]
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        ndmi = (nir - swir) / (nir + swir + 1e-8)
        ndmi = np.where(np.isfinite(ndmi), ndmi, np.nan)
    return ndmi


def compute_temporal_gradient(time_series_data):
    """
    Compute temporal gradient (first derivative) of time series.
    
    Args:
        time_series_data: 3D array [T, H, W] or 2D array [T, pixels]
    
    Returns:
        Temporal gradient array of same shape
    """
    if time_series_data.ndim == 3:
        T, H, W = time_series_data.shape
        gradient = np.zeros_like(time_series_data)
        gradient[1:-1] = (time_series_data[2:] - time_series_data[:-2]) / 2
        gradient[0] = time_series_data[1] - time_series_data[0]
        gradient[-1] = time_series_data[-1] - time_series_data[-2]
    else:
        T = time_series_data.shape[0]
        gradient = np.zeros_like(time_series_data)
        gradient[1:-1] = (time_series_data[2:] - time_series_data[:-2]) / 2
        gradient[0] = time_series_data[1] - time_series_data[0]
        gradient[-1] = time_series_data[-1] - time_series_data[-2]
    
    return gradient


def compute_spatial_gradient(image_data):
    """
    Compute spatial gradient using Sobel operator.
    
    Args:
        image_data: 2D array [H, W] or 3D array [H, W, C]
    
    Returns:
        Tuple of (grad_x, grad_y) gradient arrays
    """
    if image_data.ndim == 3:
        grad_x = np.zeros_like(image_data)
        grad_y = np.zeros_like(image_data)
        for c in range(image_data.shape[2]):
            gx, gy = _sobel_2d(image_data[:, :, c])
            grad_x[:, :, c] = gx
            grad_y[:, :, c] = gy
    else:
        grad_x, grad_y = _sobel_2d(image_data)
    
    return grad_x, grad_y


def _sobel_2d(img):
    """
    Apply Sobel operator to 2D image.
    """
    sx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    sy = sx.T
    
    pad_img = np.pad(img, 1, mode='edge')
    
    gx = np.zeros_like(img)
    gy = np.zeros_like(img)
    
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            window = pad_img[i:i+3, j:j+3]
            gx[i, j] = np.sum(window * sx)
            gy[i, j] = np.sum(window * sy)
    
    return gx, gy


__all__ = [
    'compute_ndvi',
    'compute_evi', 
    'compute_ndwi',
    'compute_ndmi',
    'compute_temporal_gradient',
    'compute_spatial_gradient'
]
