# -*- coding: utf-8 -*-
"""
SpecPipe - Vegetation Indices - MCARI (Modified Chlorophyll Absorption Reflectance Index)

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

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 mcari(
    spec_array: Annotated[Any, arraylike_validator(ndim=2)],
    wavelength: Annotated[Any, arraylike_validator(ndim=1)],
    axis: int = 0,
) -> pd.DataFrame:
    """
    Compute MCARI (Modified Chlorophyll Absorption Reflectance Index) of 2D arraylike of 1D spectra data series.
    The calculation is based on:
        Si Gao, Kai Yan, Jinxiu Liu, Jiabin Pu, Dongxiao Zou, Jianbo Qi, Xihan Mu, Guangjian Yan,
        Assessment of remote-sensed vegetation indices for estimating forest chlorophyll concentration,
        Ecological Indicators,
        Volume 162,
        2024,
        112001,
        ISSN 1470-160X,
        https://doi.org/10.1016/j.ecolind.2024.112001.

    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
    # blue_range = (458, 523)
    green_range = (543, 578)
    red_range = (650, 680)
    red_edge1 = (698, 713)
    # red_edge2 = (733, 748)
    # red_edge3 = (773, 793)
    # nir_range = (785, 900)
    if (min(wavelength) > green_range[0]) | (max(wavelength) < red_edge1[1]):  # Adjust according to applied VI
        raise ValueError(
            f'Index requires wavelength range ({green_range[0]}, {red_edge1[1]}) '  # Adjust according to applied VI
            + f'exceeds given wavelength range ({min(wavelength)}, {max(wavelength)}).'
        )

    # Spectra indices within the ranges
    # blue_bands_ind = (wavelength >= blue_range[0]) & (wavelength <= blue_range[1])
    green_bands_ind = (wavelength >= green_range[0]) & (wavelength <= green_range[1])
    red_bands_ind = (wavelength >= red_range[0]) & (wavelength <= red_range[1])
    re1_bands_ind = (wavelength >= red_edge1[0]) & (wavelength <= red_edge1[1])
    # re2_bands_ind = (wavelength >= red_edge2[0]) & (wavelength <= red_edge2[1])
    # re3_bands_ind = (wavelength >= red_edge3[0]) & (wavelength <= red_edge3[1])
    # nir_bands_ind = (wavelength >= nir_range[0]) & (wavelength <= nir_range[1])
    if (
        # (sum(blue_bands_ind) < 1) |
        (sum(green_bands_ind) < 1)
        | (sum(red_bands_ind) < 1)
        | (sum(re1_bands_ind) < 1)
        |
        # (sum(re2_bands_ind) < 1) |
        # (sum(re3_bands_ind) < 1) |
        # (sum(nir_bands_ind) < 1) |
        False
    ):
        raise ValueError('One or more required wavelength bands are missing from the dataset.')

    # Mean radiance within the ranges
    # blue = np.nanmean(spec_array[:, blue_bands_ind], axis=1)
    green = np.nanmean(spec_array[:, green_bands_ind], axis=1)
    red = np.nanmean(spec_array[:, red_bands_ind], axis=1)
    re1 = np.nanmean(spec_array[:, re1_bands_ind], axis=1)
    # re2 = np.nanmean(spec_array[:, re2_bands_ind], axis=1)
    # re3 = np.nanmean(spec_array[:, re3_bands_ind], axis=1)
    # nir = np.nanmean(spec_array[:, nir_bands_ind], axis=1)
    if (
        # np.prod(blue != 0) *
        np.prod(green != 0) * np.prod(red != 0) * np.prod(re1 != 0) *
        # np.prod(re2 != 0) *
        # np.prod(re3 != 0) *
        # np.prod(nir != 0) *
        1
        == 0
    ):
        raise ValueError('Band values are all zero in the required wavelength range.')

    # Vegetation indices
    # # Sentinel-2 Red Edge Position
    # s2rep_d = re2 - re1
    # s2rep_d[s2rep_d == 0] = 1e-10
    # s2rep_values = 705 + 35 * ((re3 + red) / 2 - re1) / s2rep_d

    # # MERIS Terrestrial Chlorophyll Index
    # mtci_d = re1 - red
    # mtci_d[mtci_d == 0] = 1e-10
    # mtci_values = (re2 - re1) / mtci_d

    # Chlorophyll Vegetation Index
    # cvi_values = nir * red / green ** 2

    # # Triangular Greenness Index
    # tgi_values = - 0.5 * ((sum(red_range) / 2 - sum(blue_range) / 2) * (red - green) - (sum(red_range) / 2 - sum(green_range) / 2) * (red - green))  # noqa: E501

    # Modified Chlorophyll Absorption Reflectance Index
    mcari_values = (re1 - red) - 0.2 * (re1 - green) * re1 / red

    vi_arr = np.array(
        [
            # s2rep_values,
            # mtci_values,
            # cvi_values,
            # tgi_values,
            mcari_values,
        ]
    )
    vi_names = [
        # 'S2REP',
        # 'MTCI',
        # 'CVI',
        # 'TGI',
        'MCARI',
    ]
    if axis == 0:
        df_vi = pd.DataFrame(vi_arr.T, columns=vi_names).loc[:, ['MCARI']]
    else:
        df_vi = pd.DataFrame(vi_arr, index=vi_names).loc[['MCARI'], :]

    return df_vi


# %% Local test

# from specpipe.vegeind.demo_data import create_vegeind_demo_data

# specdata = create_vegeind_demo_data()
# vidata = mcari(specdata, wavelength=specdata.columns)
