# -*- coding: utf-8 -*-
"""
SpecPipe - Vegetation Indices - SAVI (Soil Adjusted Vegetation Index)

Copyright (c) 2025 Siwei Luo. MIT License.
"""
import numpy as np
import pandas as pd
from typing import Annotated, Any, Union

from ..specio import simple_type_validator, arraylike_validator

# # For local test
# from specpipe.specio import simple_type_validator, arraylike_validator


# %% Vegetation index function


@simple_type_validator
def savi(
    spec_array: Annotated[Any, arraylike_validator(ndim=2)],
    wavelength: Annotated[Any, arraylike_validator(ndim=1)],
    axis: int = 0,
    L: Union[float, int] = 0.5,  # noqa: N803
) -> pd.DataFrame:
    """
    Compute SAVI (Soil Adjusted Vegetation Index) of 2D arraylike of 1D spectra data series at given L.
    The calculation is based on:
        A.R Huete,
        A soil-adjusted vegetation index (SAVI),
        Remote Sensing of Environment,
        Volume 25, Issue 3,
        1988,
        Pages 295-309,
        ISSN 0034-4257,
        https://doi.org/10.1016/0034-4257(88)90106-X.

    Parameters
    ----------
    spec_array : 2D arraylike
        2D arraylike of 1D spectral data series.
    wavelength : 1D arraylike
        Wavelengths for the spectra data.
    axis : int, optional
        Axis of spectral data series index, must be 0 or 1. If 0, each row represents a sample spectrum.
    L : float or int
        Soil influence adjustment parameter. The default is 0.5.

    Returns
    -------
    pd.DataFrame
        Dataframe containing spectral indices of the samples.
    """
    # Validate input spectral data array
    spec_array = np.array(spec_array).astype('float32')
    if np.isnan(spec_array).any():
        raise ValueError('spec_array contains NaN values.')

    # Validate axis
    if axis == 1:
        spec_array = spec_array.T
    elif axis != 0:
        raise ValueError(f'axis must be 0 or 1, got: {axis}.')
    wavelength = np.array(wavelength).astype('float32')

    # Validate data and wavelengths
    if len(wavelength) != spec_array.shape[1]:
        raise ValueError(
            f'Length of given wavelength does not match the band number of spec_array, \
                got wavelength length: {len(wavelength)}, \
                got band number of given spec_array: {spec_array.shape[1]}'
        )

    # Validate L
    if L < 0:
        raise ValueError(f"L must be positive, got: {L}")

    # Sentinel-2 wavelength ranges
    red_range = (650, 680)
    nir_range = (785, 900)
    if (min(wavelength) > red_range[0]) | (max(wavelength) < nir_range[1]):
        raise ValueError(
            f'Index requires wavelength range ({red_range[0]}, {nir_range[1]}) '
            + f'exceeds given wavelength range ({min(wavelength)}, {max(wavelength)}).'
        )
    red_bands_ind = (wavelength >= red_range[0]) & (wavelength <= red_range[1])
    nir_bands_ind = (wavelength >= nir_range[0]) & (wavelength <= nir_range[1])
    if (sum(red_bands_ind) < 1) | (sum(nir_bands_ind) < 1):
        raise ValueError('One or more required wavelength bands are missing from the dataset.')

    # Related band radiance
    radiance_red = np.nanmean(spec_array[:, red_bands_ind], axis=1)
    radiance_nir = np.nanmean(spec_array[:, nir_bands_ind], axis=1)
    if np.prod(radiance_red != 0) * np.prod(radiance_nir != 0) == 0:
        raise ValueError('Band values are all zero in the required wavelength range.')

    # Compute vegetation indices
    savi_values = ((radiance_nir - radiance_red) / (radiance_nir + radiance_red + L)) * (1 + L)

    # Vegetation index array
    vi_arr = np.array(
        [
            savi_values,
        ]
    )
    L_str = str(L).replace('.', 'd')  # noqa: N806
    vi_names = [
        f"SAVI_L{L_str}",
    ]
    if axis == 0:
        df_vi = pd.DataFrame(vi_arr.T, columns=vi_names)
    else:
        df_vi = pd.DataFrame(vi_arr, index=vi_names)

    return df_vi


@simple_type_validator
def savis(
    spec_array: Annotated[Any, arraylike_validator(ndim=2)],
    wavelength: Annotated[Any, arraylike_validator(ndim=1)],
    axis: int = 0,
) -> pd.DataFrame:
    """
    Compute SAVIs (Soil Adjusted Vegetation Indices) of 2D arraylike of 1D spectra data series.
    This function applies a series of default L values, includes: 0, 0.25, 0.5, 0.75, 1 and 2.
    The calculation is based on:
        A.R Huete,
        A soil-adjusted vegetation index (SAVI),
        Remote Sensing of Environment,
        Volume 25, Issue 3,
        1988,
        Pages 295-309,
        ISSN 0034-4257,
        https://doi.org/10.1016/0034-4257(88)90106-X.

    Parameters
    ----------
    spec_array : 2D arraylike
        2D arraylike of 1D spectral data series.
    wavelength : 1D arraylike
        Wavelengths for the spectra data.
    axis : int, optional
        Axis of spectral data series index, must be 0 or 1. If 0, each row represents a sample spectrum.

    Returns
    -------
    pd.DataFrame
        Dataframe containing spectral indices of the samples.
    """
    # Validate input spectral data array
    spec_array = np.array(spec_array).astype('float32')
    if np.isnan(spec_array).any():
        raise ValueError('spec_array contains NaN values.')

    # Validate axis
    if axis == 1:
        spec_array = spec_array.T
    elif axis != 0:
        raise ValueError(f'axis must be 0 or 1, got: {axis}.')
    wavelength = np.array(wavelength).astype('float32')

    # Validate data and wavelengths
    if len(wavelength) != spec_array.shape[1]:
        raise ValueError(
            f'Length of given wavelength does not match the band number of spec_array, \
                got wavelength length: {len(wavelength)}, \
                got band number of given spec_array: {spec_array.shape[1]}'
        )

    # Sentinel-2 wavelength ranges
    red_range = (650, 680)
    nir_range = (785, 900)
    if (min(wavelength) > red_range[0]) | (max(wavelength) < nir_range[1]):
        raise ValueError(
            f'Index requires wavelength range ({red_range[0]}, {nir_range[1]}) '
            + f'exceeds given wavelength range ({min(wavelength)}, {max(wavelength)}).'
        )
    red_bands_ind = (wavelength >= red_range[0]) & (wavelength <= red_range[1])
    nir_bands_ind = (wavelength >= nir_range[0]) & (wavelength <= nir_range[1])
    if (sum(red_bands_ind) < 1) | (sum(nir_bands_ind) < 1):
        raise ValueError('One or more required wavelength bands are missing from the dataset.')

    # Related band radiance
    radiance_red = np.nanmean(spec_array[:, red_bands_ind], axis=1)
    radiance_nir = np.nanmean(spec_array[:, nir_bands_ind], axis=1)
    if np.prod(radiance_red != 0) * np.prod(radiance_nir != 0) == 0:
        raise ValueError('Band values are all zero in the required wavelength range.')

    # Compute vegetation indices
    savi_list = []
    savi_names = []
    for L in [0, 0.25, 0.5, 0.75, 1, 2]:  # noqa: N806
        savi_values = ((radiance_nir - radiance_red) / (radiance_nir + radiance_red + L)) * (1 + L)
        savi_list.append(savi_values)
        L_str = str(L).replace('.', 'd')  # noqa: N806
        savi_names.append(f"SAVI_L{L_str}")

    # Vegetation index array
    vi_arr = np.array(savi_list)
    vi_names = savi_names
    if axis == 0:
        df_vi = pd.DataFrame(vi_arr.T, columns=vi_names)
    else:
        df_vi = pd.DataFrame(vi_arr, index=vi_names)

    return df_vi


# %% Local test

# from specpipe.vegeind.demo_data import create_vegeind_demo_data

# specdata = create_vegeind_demo_data()
# vidata = savi(specdata, wavelength=specdata.columns, L=1)
