"""
veg_indices.py

A reusable Python module for loading spectral data (from CSV or an existing DataFrame),
calculating vegetation indices, and optionally saving results back to CSV.
"""

import pandas as pd
import numpy as np

###############################################################################
# 1. Vegetation Index Calculation Functions
###############################################################################

def calculate_ndvi(r680, r800):
    """NDVI = (R800 - R680) / (R800 + R680)"""
    return (r800 - r680) / (r800 + r680)

def calculate_gndvi(r550, r750):
    """GNDVI = (R750 - R550) / (R750 + R550)"""
    return (r750 - r550) / (r750 + r550)

def calculate_rndvi(r750, r705):
    """Red-Edge NDVI = (R750 - R705) / (R750 + R705)"""
    return (r750 - r705) / (r750 + r705)

def calculate_wbi(r900, r970):
    """WBI = R900 / R970"""
    return r900 / r970

def calculate_wsct(r970, r850):
    """WSCT = (R970 - R850) / (R970 + R850)"""
    return (r970 - r850) / (r970 + r850)

def calculate_evi(r450, r670, r800):
    """
    EVI = 2.5 * ( (R800 - R670) / (R800 + 6*R670 - 7.5*R450 + 1) )
    """
    return 2.5 * (r800 - r670) / (r800 + 6*r670 - 7.5*r450 + 1.0)

def calculate_mcari(r550, r670, r700):
    """MCARI = ((R700 - R670) - 0.2*(R700 - R550)) * (R700 / R670)"""
    return ((r700 - r670) - 0.2 * (r700 - r550)) * (r700 / r670)

def calculate_rdvi(r670, r800):
    """RDVI = (R800 - R670) / sqrt(R800 + R670)"""
    return (r800 - r670) / np.sqrt(r800 + r670)

def calculate_savi(r670, r800, L=0.5):
    """SAVI = (1 + L) * (R800 - R670) / (R800 + R670 + L)"""
    return (1.0 + L) * (r800 - r670) / (r800 + r670 + L)

def calculate_tvi(r550, r670, r750):
    """TVI = 0.5 * (120*(R750 - R550) - 200*(R670 - R550))"""
    return 0.5 * (120*(r750 - r550) - 200*(r670 - r550))

def calculate_psri(r678, r500, r750):
    """PSRI = (R678 - R500) / R750"""
    return (r678 - r500) / r750

def calculate_pri(r570, r531):
    """PRI = (R570 - R531) / (R570 + R531)"""
    return (r570 - r531) / (r570 + r531)

def calculate_ndre(r842, r783):
    """NDRE = (R842 - R783) / (R842 + R783)"""
    return (r842 - r783) / (r842 + r783)

def calculate_npci(r680, r430):
    """NPCI = (R680 - R430) / (R680 + R430)"""
    return (r680 - r430) / (r680 + r430)

def calculate_lci(r850, r710, r680):
    """LCI = (R850 - R710) / (R850 + R680)"""
    return (r850 - r710) / (r850 + r680)

def calculate_rvi1(r810, r560):
    """RVI1 = R810 / R560"""
    return r810 / r560

def calculate_rvi2(r950, r660):
    """RVI2 = R950 / R660"""
    return r950 / r660

def calculate_gci(r800, r550):
    """GCI = (R800 / R550) - 1.0"""
    return (r800 / r550) - 1.0

def calculate_rgvi(r670, r550):
    """RGVI = R670 / R550"""
    return r670 / r550

def calculate_dvi(r800, r670):
    """DVI = R800 - R670"""
    return r800 - r670

def calculate_msavi(r800, r670):
    """
    MSAVI = 0.5 * ( (2*R800 + 1) - sqrt( (2*R800 + 1)^2 - 8*(R800 - R670) ) )
    """
    return 0.5 * (
        (2.0*r800 + 1.0) -
        np.sqrt((2.0*r800 + 1.0)**2 - 8.0*(r800 - r670))
    )

def calculate_osavi(r800, r670):
    """OSAVI = (R800 - R670) / (R800 + R670 + 0.16)"""
    return (r800 - r670) / (r800 + r670 + 0.16)

def calculate_tvi_alt(r800, r670, r550):
    """
    Alternate "TVI" formula from your code snippet:
    60*(R800 - R550) - 100*(R670 - R550)
    """
    return 60.0 * (r800 - r550) - 100.0 * (r670 - r550)

def calculate_tsavi(r800, r670, a=0.96916, b=0.084726):
    """
    TSAVI = [a * (R800 - a*R670 - b)] / [a*R800 + R670 - a*b]
    """
    return (a*(r800 - a*r670 - b)) / (a*r800 + r670 - a*b)

def calculate_pvi(r800, r670, a=0.96916, b=0.084726):
    """
    PVI = (R800 - a*R670 - b) / sqrt(1 + a^2)
    """
    return (r800 - a*r670 - b) / np.sqrt(1.0 + a**2)

def calculate_savi_2(r800, r670, a=0.96916, b=0.084726):
    """
    SAVI_2 = R800 / [R670 - (b / a)]
    """
    return r800 / (r670 - (b/a))

def calculate_atsavi(r800, r670, a=0.96916, b=0.084726, X=0.08):
    """
    ATSAVI = [a * (-a*R670 - b)] / [a*R800 + R670 - a*b + X*(1 + a^2)]
    """
    return (a * (-a*r670 - b)) / (a*r800 + r670 - a*b + X*(1 + a*a))

def calculate_ndwi(r550, r800):
    """NDWI = (R550 - R800) / (R550 + R800)"""
    return (r550 - r800) / (r550 + r800)

def calculate_srpi(r430, r670):
    """SRPI = R430 / R670"""
    return r430 / r670

def calculate_mcari_1(r800, r670, r550):
    """MCARI_1 = 1.2 * [2.5*(R800 - R670) - 1.3*(R800 - R550)]"""
    return 1.2 * (2.5*(r800 - r670) - 1.3*(r800 - r550))

def calculate_mcari_2(r800, r670, r550):
    """
    MCARI_2 = 1.5 * [ (2.5*(R800 - R670) - 1.3*(R800 - R550)) * ((2*R800 + 1)^2) - (6*R800 - 5*R670) - 0.5 ]
    """
    return 1.5 * (
        (2.5*(r800 - r670) - 1.3*(r800 - r550)) *
        ((2.0*r800 + 1.0)**2) - (6.0*r800 - 5.0*r670) - 0.5
    )

def calculate_mtvi1(r800, r670, r550):
    """MTVI1 = 1.2 * (1.2*(R800 - R550) - 2.5*(R670 - R550))"""
    return 1.2 * (1.2*(r800 - r550) - 2.5*(r670 - r550))

def calculate_mtvi2(r800, r670, r550):
    """
    MTVI2 = 1.5 * [ (1.2*(R800 - R550) - 2.5*(R670 - R550)) * ((2*R800 + 1)^2) - (6.0*R800 - 5.0*R670) - 0.5 ]
    """
    return 1.5 * (
        (1.2*(r800 - r550) - 2.5*(r670 - r550)) *
        ((2.0*r800 + 1.0)**2) - (6.0*r800 - 5.0*r670) - 0.5
    )

def calculate_r_mcari_mtvi2(r700, r800, r670, r550):
    """Ratio of MCARI to MTVI2"""
    mcari_component = ((r700 - r670) - 0.2*(r700 - r550)) * (r700 / r670)
    mtvi2_component = calculate_mtvi2(r800, r670, r550)
    return mcari_component / mtvi2_component

def calculate_datt(r800, r700, r670):
    """DATT = (R800 - R700) / (R800 - R670)"""
    return (r800 - r700) / (r800 - r670)

def calculate_ndci(r700, r550):
    """NDCI = (R700 - R550) / (R700 + R550)"""
    return (r700 - r550) / (r700 + r550)

def calculate_sipi(r800, r430, r670):
    """SIPI = (R800 - R430) / (R800 + R670)"""
    return (r800 - r430) / (r800 + r670)

def calculate_spvi(r800, r670, r550):
    """SPVI = 0.4 * 3.7 * (R800 - R670) - 1.2 * abs(R550 - R670)"""
    return 0.4 * 3.7 * (r800 - r670) - 1.2 * abs(r550 - r670)

def calculate_tcari(r700, r670, r550):
    """TCARI = 3 * [ (R700 - R670) - 0.2*(R700 - R550)*(R700/R670) ]"""
    return 3.0 * (
        (r700 - r670) - 0.2*(r700 - r550)*(r700 / r670)
    )

def calculate_reri(r700, r670, r800):
    """RERI = (R700 - R670) / R800"""
    return (r700 - r670) / r800

def calculate_mtci(r800, r700, r670):
    """MTCI = (R800 - R700) / (R700 - R670)"""
    return (r800 - r700) / (r700 - r670)

def calculate_evi2(r800, r670):
    """EVI2 = 2.5 * [ (R800 - R670) / (R800 + 2.4*R670 + 1) ]"""
    return 2.5 * ((r800 - r670) / (r800 + 2.4*r670 + 1.0))

def calculate_reci(r800, r700):
    """RECI = (R800 / R700) - 1"""
    return (r800 / r700) - 1.0

def calculate_nexg(r550, r670, r430):
    """NExG = (2*R550 - R670 - R430) / (R550 + R670 + R430)"""
    return (2.0*r550 - r670 - r430) / (r550 + r670 + r430)

def calculate_ngrdi(r550, r670):
    """NGRDI = (R550 - R670) / (R550 + R670)"""
    return (r550 - r670) / (r550 + r670)

def calculate_endvi(r800, r550, r430):
    """ENDVI = (R800 + R550 - 2*R430) / (R800 + R550 + 2*R430)"""
    return (r800 + r550 - 2.0*r430) / (r800 + r550 + 2.0*r430)

def calculate_ari_2(r800, r550, r700):
    """ARI2 = R800 * [ (1/R550) - (1/R700) ]"""
    return r800 * ((1.0/r550) - (1.0/r700))

def calculate_cri_2(r550, r700):
    """CRI2 = (1/R550) - (1/R700)"""
    return (1.0/r550) - (1.0/r700)

def calculate_r_tcari_osavi(r700, r670, r550, r800):
    """
    Ratio of TCARI to OSAVI
    RTCRI_OSAVI = TCARI / OSAVI
    """
    tcari_val = calculate_tcari(r700, r670, r550)
    osavi_val = calculate_osavi(r800, r670)
    return tcari_val / osavi_val


###############################################################################
# 2. Dictionary of Indices and the Required Wavelengths
###############################################################################

INDEX_DEFINITIONS = {
    "NDVI": ([680, 800], calculate_ndvi),
    "GNDVI": ([550, 750], calculate_gndvi),
    "RNDVI": ([750, 705], calculate_rndvi),
    "WBI": ([900, 970], calculate_wbi),
    "WSCT": ([970, 850], calculate_wsct),
    "EVI": ([450, 670, 800], calculate_evi),
    "MCARI": ([550, 670, 700], calculate_mcari),
    "SAVI": ([670, 800], calculate_savi),
    "TVI": ([550, 670, 750], calculate_tvi),
    "PSRI": ([678, 500, 750], calculate_psri),
    "PRI": ([570, 531], calculate_pri),
    "NDRE": ([842, 783], calculate_ndre),
    "NPCI": ([680, 430], calculate_npci),
    "LCI": ([850, 710, 680], calculate_lci),
    "RVI1": ([810, 560], calculate_rvi1),
    "RVI2": ([950, 660], calculate_rvi2),
    "GCI": ([800, 550], calculate_gci),
    "rgvi": ([670, 550], calculate_rgvi),
    "DVI": ([800, 670], calculate_dvi),
    "MSAVI": ([800, 670], calculate_msavi),
    "OSAVI": ([800, 670], calculate_osavi),
    "TVI_ALT": ([800, 670, 550], calculate_tvi_alt),
    "TSAVI": ([800, 670], calculate_tsavi),
    "PVI": ([800, 670], calculate_pvi),
    "SAVI_2": ([800, 670], calculate_savi_2),
    "ATSAVI": ([800, 670], calculate_atsavi),
    "NDWI": ([550, 800], calculate_ndwi),
    "SRPI": ([430, 670], calculate_srpi),
    "MCARI_1": ([800, 670, 550], calculate_mcari_1),
    "MCARI_2": ([800, 670, 550], calculate_mcari_2),
    "MTVI1": ([800, 670, 550], calculate_mtvi1),
    "MTVI2": ([800, 670, 550], calculate_mtvi2),
    "r_MCARI_MTVI2": ([700, 800, 670, 550], calculate_r_mcari_mtvi2),
    "DATT": ([800, 700, 670], calculate_datt),
    "NDCI": ([700, 550], calculate_ndci),
    "SIPI": ([800, 430, 670], calculate_sipi),
    "SPVI": ([800, 670, 550], calculate_spvi),
    "TCARI": ([700, 670, 550], calculate_tcari),
    "RERI": ([700, 670, 800], calculate_reri),
    "MTCI": ([800, 700, 670], calculate_mtci),
    "EVI2": ([800, 670], calculate_evi2),
    "RECI": ([800, 700], calculate_reci),
    "NExG": ([550, 670, 430], calculate_nexg),
    "NGRDI": ([550, 670], calculate_ngrdi),
    "ENDVI": ([800, 550, 430], calculate_endvi),
    "ARI_2": ([800, 550, 700], calculate_ari_2),
    "CRI_2": ([550, 700], calculate_cri_2),
    "RTCRI_OSAVI": ([700, 670, 550, 800], calculate_r_tcari_osavi),
}


###############################################################################
# 3. Core Helper Functions
###############################################################################

def load_spectral_csv(csv_path):
    """
    Loads the CSV file into a DataFrame.
    By convention, the first row is assumed to have wavelength information,
    and subsequent rows contain sample reflectances.
    """
    return pd.read_csv(csv_path)

def extract_band_values(df, target_wavelengths):
    """
    Extract band reflectance arrays for the requested target wavelengths.

    - The first row of df is assumed to have numeric wavelength values 
      (for columns 1..end), i.e. df.iloc[0,1:].
    - The actual reflectance data start from row 1 onward.
    """
    wavelengths = df.iloc[0, 1:].values  # skip first column if it's an ID
    band_values = {}

    for target_wv in target_wavelengths:
        # Find the closest column in terms of absolute difference
        closest_idx = min(
            range(len(wavelengths)),
            key=lambda i: abs(wavelengths[i] - target_wv)
        )
        # Reflectances: from row 1 onward, same column offset
        band_values[target_wv] = df.iloc[1:, closest_idx + 1].values
    return band_values

def calculate_index_for_all_samples(df, index_name, custom_wavelengths=None):
    """
    Calculate a single vegetation index for all samples in the provided DataFrame.
    
    Args:
        df (DataFrame): The spectral DataFrame with the first row as wavelengths.
        index_name (str): One of the keys in INDEX_DEFINITIONS.
        custom_wavelengths (list, optional): Override default wavelengths for this index.
    
    Returns:
        A 1D NumPy array of index values (one per sample row).
    """
    if index_name not in INDEX_DEFINITIONS:
        raise ValueError(f"Index '{index_name}' is not defined in INDEX_DEFINITIONS.")

    default_wv, index_func = INDEX_DEFINITIONS[index_name]

    if custom_wavelengths is not None:
        if len(custom_wavelengths) != len(default_wv):
            raise ValueError(
                f"Number of custom wavelengths provided ({len(custom_wavelengths)}) "
                f"does not match what '{index_name}' requires ({len(default_wv)})."
            )
        wv_list = custom_wavelengths
    else:
        wv_list = default_wv

    # Extract reflectance arrays for each required wavelength
    band_vals = extract_band_values(df, wv_list)
    
    # Put them in the correct order to pass into the index function
    reflectance_arrays = [band_vals[wv] for wv in wv_list]

    # Perform the index calculation (vectorized) over all samples
    result = index_func(*reflectance_arrays)
    return result

def calculate_multiple_indices(df, index_list=None, custom_wv_dict=None):
    """
    Calculate multiple indices for all samples in the DataFrame.
    
    Args:
        df (DataFrame): The input spectral DataFrame (first row = wavelengths).
        index_list (list of str, optional): Names of the indices to calculate. 
                                            If None, calculates all.
        custom_wv_dict (dict, optional): Override wavelengths for some indices.
                                         e.g. { "NDVI": [681, 801], "EVI": [451, 671, 801], ... }
    Returns:
        result_df (pd.DataFrame): One column per index, one row per sample.
    """
    if index_list is None:
        # If none provided, compute all known indices
        index_list = list(INDEX_DEFINITIONS.keys())

    output_data = {}

    for idx_name in index_list:
        # If user provided custom wavelengths for this index, use them
        cw = None
        if custom_wv_dict and (idx_name in custom_wv_dict):
            cw = custom_wv_dict[idx_name]

        index_values = calculate_index_for_all_samples(df, idx_name, cw)
        output_data[idx_name] = index_values

    # Convert to DataFrame. Each column = an index, each row = a sample
    result_df = pd.DataFrame(output_data)
    return result_df

def save_indices_to_csv(result_df, out_path):
    """
    Saves the result DataFrame (indices) to a CSV at 'out_path'.
    """
    result_df.to_csv(out_path, index=False)


###############################################################################
# 4. A High-Level Wrapper Function
###############################################################################

def compute_indices(
    input_data,
    index_list=None,
    custom_wv_dict=None,
    out_csv=None
):
    """
    High-level function to compute vegetation indices from either:
      - a CSV file path (str)
      - an existing DataFrame
    
    Args:
        input_data (str or pd.DataFrame):
            If it's a string, we treat it as a path to a CSV file.
            If it's already a DataFrame, we use it directly.
        index_list (list of str, optional):
            Names of indices to calculate. If None, calculates all available.
        custom_wv_dict (dict, optional):
            Override default wavelengths for specific indices, e.g.:
            { "NDVI": [681, 801], "EVI": [451, 671, 801] }
        out_csv (str, optional):
            If provided, the resulting DataFrame of indices will be saved to this CSV path.

    Returns:
        pd.DataFrame: The calculated indices, one column per index, one row per sample.
    """
    # 1. If input_data is a CSV path, load it. Else assume it's already a DataFrame.
    if isinstance(input_data, str):
        df = load_spectral_csv(input_data)
    elif isinstance(input_data, pd.DataFrame):
        df = input_data
    else:
        raise TypeError("`input_data` must be either a CSV filepath (str) or a pandas DataFrame.")

    # 2. Calculate the requested indices
    result_df = calculate_multiple_indices(df, index_list, custom_wv_dict)

    # 3. Optionally save to CSV
    if out_csv is not None:
        save_indices_to_csv(result_df, out_csv)

    return result_df
