import logging
import pandas as pd
import numpy as np
from scipy.signal import find_peaks
from typing import Tuple
from typing import Tuple, Union

# Configure logging at the module level
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def find_peaks_from_chl_data(
    data: pd.DataFrame,
    time_col: str,
    chl_col: str,
    threshold: float
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Detect peaks in chlorophyll data above a specified threshold using scipy.signal.find_peaks.

    Args:
        data (pd.DataFrame): DataFrame containing time series data.
        time_col (str): Column name for the time variable (e.g., 'weekofyear' or 'dayofyear').
        chl_col (str): Column name for chlorophyll values (e.g., 'CHL').
        threshold (float): Minimum chlorophyll value for peak detection.

    Returns:
        Tuple[np.ndarray, np.ndarray]: Arrays of time values and corresponding chlorophyll values at peaks.

    Raises:
        TypeError: If data is not a pandas DataFrame or if chl_col is not numeric.
        ValueError: If required columns are missing or data is empty.
    """
    if not isinstance(data, pd.DataFrame):
        raise TypeError("data must be a pandas DataFrame")
    if time_col not in data.columns or chl_col not in data.columns:
        raise ValueError(f"Columns {time_col} and/or {chl_col} not found in data")
    if not np.issubdtype(data[chl_col].dtype, np.number):
        raise TypeError(f"Column {chl_col} must contain numeric values")
    if data.empty:
        logger.debug("Empty DataFrame provided, returning empty arrays")
        return np.array([]), np.array([])

    data = data.sort_values(time_col)
    peaks, _ = find_peaks(data[chl_col], height=threshold)
    peak_times = data.iloc[peaks][time_col].values
    peak_chl = data.iloc[peaks][chl_col].values
    return peak_times, peak_chl

# 1: The following funciton detects phytoplankton bloom from daily climatology data
def detect_phyto_pheno_1(
    df: pd.DataFrame,
    lat_col: str = 'latitude',
    lon_col: str = 'longitude',
    time_col: str = 'day',
    chl_col: str = 'CHL',
    threshold_percent: int = 10,
    min_duration: int = 0,
    base_year: int = 2000,
    verbose: bool = False
) -> pd.DataFrame:
    """
    Detect phytoplankton bloom phenology from daily chlorophyll climatology data using a group-based median threshold.
    Peaks are detected using find_peaks_from_chl_data. Requires latitude and longitude for grouping.
    Requires columns 'latitude', 'longitude', 'day', and 'CHL'. Peaks are detected using scipy.signal.find_peaks.

    Args:
        df (pd.DataFrame): Input DataFrame with daily climatology data.
        lat_col (str): Column name for latitude. Default: 'latitude'.
        lon_col (str): Column name for longitude. Default: 'longitude'.
        time_col (str): Column name for time variable (e.g., 'dayofyear'). Default: 'dayofyear'.
        chl_col (str): Column name for chlorophyll values (e.g., 'CHL'). Default: 'CHL'.
        threshold_percent (int): Percentage above median for bloom detection (e.g., 10 for 10%). Default: 10.
        min_duration (int): Minimum bloom duration (days) for valid event. Default: 0.
        base_year (int): Base year for day-to-month conversion (must be a valid year). Default: 2000.
        verbose (bool): If True, logs detailed processing information. Default: False.

    Returns:
        pd.DataFrame: DataFrame with phenology indices for each location, including:
            - latitude, longitude
            - bloom_initiation, bloom_initiation_chl
            - bloom_peak, bloom_peak_chl
            - bloom_termination, bloom_termination_chl
            - bloom_duration
            - bloom_initiation_month, bloom_peak_month, bloom_termination_month
    """
    # Set logging level based on verbose
    if verbose:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)

    # Validate inputs
    required_cols = [lat_col, lon_col, time_col, chl_col]
    missing = [col for col in required_cols if col not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")
    if not np.issubdtype(df[chl_col].dtype, np.number):
        raise TypeError(f"Column {chl_col} must contain numeric values")
    if not isinstance(threshold_percent, int) or threshold_percent < 0:
        raise ValueError("threshold_percent must be a non-negative integer")
    if min_duration < 0:
        raise ValueError("min_duration must be non-negative")

    # Base date for month conversion
    base_date = pd.Timestamp(f"{base_year}-01-01")

    # Helper function to convert day of year to month
    def to_month(day):
        if pd.isna(day):
            return np.nan
        target_date = base_date + pd.Timedelta(days=int(day) - 1)
        return target_date.strftime('%B')

    # Group data by latitude and longitude
    grouped = df.groupby([lat_col, lon_col])
    logger.info(f"Grouping by [{lat_col}, {lon_col}]")

    results = []
    for (lat, lon), group in grouped:
        logger.debug(f"Processing group: ({lat}, {lon})")
        group = group.sort_values(time_col).groupby(time_col)[chl_col].mean().reset_index()

        if len(group) < 2:
            logger.warning(f"Insufficient data for group ({lat}, {lon})")
            continue

        # Calculate threshold
        median_chl = group[chl_col].median()
        threshold = median_chl * (1 + threshold_percent / 100)

        # Bloom initiation
        above_threshold = group[chl_col] > threshold
        if not above_threshold.any():
            continue
        init_idx = above_threshold.idxmax()
        bloom_initiation = group.loc[init_idx, time_col]
        bloom_initiation_chl = group.loc[init_idx, chl_col]

        # Bloom peak using the helper function
        peak_days, peak_chl_values = find_peaks_from_chl_data(group, time_col, chl_col, threshold)
        if len(peak_days) == 0:
            continue
        valid_peaks = [(day, chl) for day, chl in zip(peak_days, peak_chl_values) if day >= bloom_initiation]
        if not valid_peaks:
            continue
        bloom_peak, bloom_peak_chl = max(valid_peaks, key=lambda x: x[1])

        # Bloom termination
        post_peak = group[group[time_col] > bloom_peak]
        if post_peak.empty:
            continue
        below_threshold = post_peak[chl_col] < threshold
        if not below_threshold.any():
            continue
        term_idx = below_threshold.idxmax()
        bloom_termination = post_peak.loc[term_idx, time_col]
        bloom_termination_chl = post_peak.loc[term_idx, chl_col]

        # Duration check
        bloom_duration = bloom_termination - bloom_initiation
        if bloom_duration < min_duration:
            continue

        # Store results
        results.append({
            lat_col: lat,
            lon_col: lon,
            'bloom_initiation': bloom_initiation,
            'bloom_initiation_chl': bloom_initiation_chl,
            'bloom_peak': bloom_peak,
            'bloom_peak_chl': bloom_peak_chl,
            'bloom_termination': bloom_termination,
            'bloom_termination_chl': bloom_termination_chl,
            'bloom_duration': bloom_duration,
            'bloom_initiation_month': to_month(bloom_initiation),
            'bloom_peak_month': to_month(bloom_peak),
            'bloom_termination_month': to_month(bloom_termination)
        })

    # Return results as DataFrame
    if not results:
        logger.warning("No valid blooms detected")
        return pd.DataFrame()
    result_df = pd.DataFrame(results)
    cols = [lat_col, lon_col] + [col for col in result_df.columns if col not in [lat_col, lon_col]]
    return result_df[cols]


# 2: The following funciton detects phytoplankton bloom from weekly climatology data
def detect_phyto_pheno_2(
    df: pd.DataFrame,
    lat_col: str = 'latitude',
    lon_col: str = 'longitude',
    time_col: str = 'week',
    chl_col: str = 'CHL',
    threshold_percent: int = 10,
    min_duration: int = 0,
    base_year: int = 2000,
    verbose: bool = False
) -> pd.DataFrame:
    """
    Detect phytoplankton bloom phenology from weekly chlorophyll climatology data using a group-based median threshold.
    Requires columns 'latitude', 'longitude', 'weekofyear', and 'CHL'. Peaks are detected using scipy.signal.find_peaks.

    Args:
        df (pd.DataFrame): Input DataFrame with weekly climatology data.
        lat_col (str): Column name for latitude. Default: 'latitude'.
        lon_col (str): Column name for longitude. Default: 'longitude'.
        time_col (str): Column name for time variable (e.g., 'weekofyear'). Default: 'weekofyear'.
        chl_col (str): Column name for chlorophyll values (e.g., 'CHL'). Default: 'CHL'.
        threshold_percent (int): Percentage above median for bloom detection (e.g., 10 for 10%). Default: 10.
        min_duration (int): Minimum bloom duration (weeks) for valid event. Default: 0.
        base_year (int): Base year for week-to-month conversion (must be a valid year). Default: 2000.
        verbose (bool): If True, logs detailed processing information. Default: False.

    Returns:
        pd.DataFrame: DataFrame with phenology indices for each location, including:
            - latitude, longitude
            - bloom_initiation, bloom_initiation_chl
            - bloom_peak, bloom_peak_chl
            - bloom_termination, bloom_termination_chl
            - bloom_duration
            - bloom_initiation_month, bloom_peak_month, bloom_termination_month
    """
    # Set logging level
    if verbose:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    
    # Validate inputs
    required_cols = [lat_col, lon_col, time_col, chl_col]
    missing = [col for col in required_cols if col not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")
    if df.empty:
        raise ValueError("Input DataFrame is empty")
    if not np.issubdtype(df[chl_col].dtype, np.number):
        raise TypeError(f"Column {chl_col} must contain numeric values")
    if not isinstance(threshold_percent, int):
        raise TypeError("threshold_percent must be an integer")
    if threshold_percent < 0:
        raise ValueError("threshold_percent must be non-negative")
    if min_duration < 0:
        raise ValueError("min_duration must be non-negative")

    # Create a base date for month conversion
    try:
        base_date = pd.Timestamp(f"{base_year}-01-01")
    except ValueError:
        raise ValueError(f"Invalid base_year: {base_year}. Must be a valid year.")

    # Function to convert week of year to month name
    def to_month(week: Union[float, int]) -> str:
        if pd.isna(week):
            return np.nan
        try:
            # Approximate the middle day of the week: (week-1)*7 + 3.5
            approx_day = (int(week) - 1) * 7 + 3.5
            target_date = base_date + pd.Timedelta(days=approx_day)
            return target_date.strftime('%B')
        except ValueError:
            logger.warning(f"Invalid week value: {week}. Returning NaN.")
            return np.nan

    # Group by latitude and longitude
    grouped = df.groupby([lat_col, lon_col])
    logger.info(f"Grouping by [{lat_col}, {lon_col}]")

    # Process each location group
    results = []
    for (lat, lon), group in grouped:
        logger.debug(f"Processing group: ({lat}, {lon})")
        
        # Ensure sorted by time and handle duplicates
        group = group.sort_values(time_col)
        group = group.groupby(time_col)[chl_col].mean().reset_index()
        
        # Check for sufficient data
        if len(group) < 2:
            logger.warning(f"Group ({lat}, {lon}) has insufficient data points ({len(group)})")
            continue
        
        # Calculate the median and threshold
        median_chl = group[chl_col].median()
        threshold = median_chl * (1 + threshold_percent / 100)
        logger.debug(f"Threshold for group ({lat}, {lon}): {threshold}")
        
        # Find bloom initiation
        above_threshold = group[chl_col] > threshold
        if not above_threshold.any():
            logger.debug(f"No values above threshold for group ({lat}, {lon})")
            continue
            
        init_idx = above_threshold.idxmax()
        bloom_initiation = group.loc[init_idx, time_col]
        bloom_initiation_chl = group.loc[init_idx, chl_col]
        
        # Find bloom peak
        peak_times, peak_chl_values = find_peaks_from_chl_data(group, time_col, chl_col, threshold)
        if len(peak_times) == 0:
            logger.debug(f"No peaks found for group ({lat}, {lon})")
            continue
        
        # Select the highest peak after initiation
        valid_peaks = [(time, chl) for time, chl in zip(peak_times, peak_chl_values) if time >= bloom_initiation]
        if not valid_peaks:
            logger.debug(f"No valid peaks after initiation for group ({lat}, {lon})")
            continue
        bloom_peak, bloom_peak_chl = max(valid_peaks, key=lambda x: x[1])
        
        # Find termination
        post_peak = group[group[time_col] > bloom_peak]
        if post_peak.empty:
            logger.debug(f"No data after peak for group ({lat}, {lon})")
            continue
            
        below_threshold = post_peak[chl_col] < threshold
        if not below_threshold.any():
            logger.debug(f"No values below threshold after peak for group ({lat}, {lon})")
            continue
            
        term_idx = below_threshold.idxmax()
        bloom_termination = post_peak.loc[term_idx, time_col]
        bloom_termination_chl = post_peak.loc[term_idx, chl_col]
        
        # Calculate duration and validate
        bloom_duration = bloom_termination - bloom_initiation
        if bloom_duration < min_duration:
            logger.debug(f"Bloom duration {bloom_duration} too short for group ({lat}, {lon})")
            continue
        
        # Add to results
        results.append({
            lat_col: lat,
            lon_col: lon,
            'bloom_initiation': bloom_initiation,
            'bloom_initiation_chl': bloom_initiation_chl,
            'bloom_peak': bloom_peak,
            'bloom_peak_chl': bloom_peak_chl,
            'bloom_termination': bloom_termination,
            'bloom_termination_chl': bloom_termination_chl,
            'bloom_duration': bloom_duration,
            'bloom_initiation_month': to_month(bloom_initiation),
            'bloom_peak_month': to_month(bloom_peak),
            'bloom_termination_month': to_month(bloom_termination)
        })
    
    # Create and return results DataFrame
    if not results:
        logger.warning("No valid blooms detected")
        return pd.DataFrame()
    
    result_df = pd.DataFrame(results)
    # Reorder columns to put lat_col, lon_col first
    cols = [lat_col, lon_col] + [col for col in result_df.columns if col not in [lat_col, lon_col]]
    result_df = result_df[cols]
    
    logger.info(f"Detected {len(result_df)} blooms")
    return result_df



# 3: The following funciton detects phytoplankton bloom from daily climatology data. The data can't contain latitude and longitude
def detect_phyto_pheno_3(
    df: pd.DataFrame,
    time_col: str = 'day_of_year',
    chl_col: str = 'CHL',
    threshold_percent: int = 10,
    min_gap: int = 1,
    min_duration: int = 1,
    min_peak_height: float = None
) -> dict:
    """
    Detect phytoplankton bloom events using percentage-based threshold.
    
    Threshold condition: threshold = median_chl * (1 + threshold_percent/100)
    
    Args:
        df: DataFrame containing chlorophyll time series
        time_col: Column name for temporal index (day/week of year)
        chl_col: Column name for chlorophyll values
        threshold_percent: Percentage above median for bloom threshold
        min_gap: Maximum gap between measurements to be considered same bloom
        min_duration: Minimum event duration to be considered a bloom
        min_peak_height: Minimum peak height to qualify as bloom (optional)
        
    Returns:
        Dictionary containing:
            - median_chl: Median chlorophyll value
            - threshold: Bloom detection threshold used
            - threshold_percent: Percentage above median used
            - bloom_events: List of bloom dictionaries
            
    Example:
        bloom_data = detect_bloom_events(df, threshold_percent=15.0,
                                        time_col='weekofyear')
    """
    # Validate inputs
    if not isinstance(df, pd.DataFrame):
        raise TypeError("Input must be a pandas DataFrame")
    if time_col not in df.columns or chl_col not in df.columns:
        missing = [c for c in [time_col, chl_col] if c not in df.columns]
        raise ValueError(f"Columns missing: {missing}")
    if threshold_percent < 0:
        raise ValueError("threshold_percent must be non-negative")
    
    # Calculate threshold using specified condition
    median_chl = np.median(df[chl_col])
    threshold = median_chl * (1 + threshold_percent / 100)
    
    # Identify values above threshold
    above_threshold = df[df[chl_col] > threshold].copy()
    
    if above_threshold.empty:
        return {
            'median_chl': median_chl,
            'threshold': threshold,
            'threshold_percent': threshold_percent,
            'bloom_events': []
        }
    
    # Sort by time and group consecutive events
    above_threshold = above_threshold.sort_values(time_col)
    time_diffs = above_threshold[time_col].diff()
    above_threshold['bloom_group'] = (time_diffs > min_gap).cumsum()
    
    # Extract bloom events
    bloom_events = []
    for _, group in above_threshold.groupby('bloom_group'):
        duration = group[time_col].iloc[-1] - group[time_col].iloc[0] + 1
        
        if duration < min_duration:
            continue
            
        # Find peak within bloom
        peak_idx = group[chl_col].idxmax()
        peak_value = group.loc[peak_idx, chl_col]
        
        # Apply peak height filter if specified
        if min_peak_height is not None and peak_value < min_peak_height:
            continue
            
        event = {
            'start_day': group.iloc[0][time_col],
            'start_value': group.iloc[0][chl_col],
            'end_day': group.iloc[-1][time_col],
            'end_value': group.iloc[-1][chl_col],
            'peak_day': group.loc[peak_idx][time_col],
            'peak_value': peak_value,
            'duration': duration
        }
        bloom_events.append(event)
    
    return {
        'median_chl': median_chl,
        'threshold': threshold,
        'threshold_percent': threshold_percent,
        'bloom_events': bloom_events
    }



