import numpy as np
import pandas as pd
from tqdm import tqdm
from scipy import stats
from scipy.stats import mode
from joblib import Parallel, delayed
from numpy.lib.stride_tricks import sliding_window_view
from .api_resample import resample_2D_mean, resample_3D_mean

def aggregate2D(x, exc_shape, ignore_nodata=None):
    """
    x: 2D array
    exc_shape: tuple of (rows, cols)
    return: 2D array
    """
    cur_rows, cur_cols = x.shape
    exc_rows, exc_cols = exc_shape
    rw = cur_rows / exc_rows
    cw = cur_cols / exc_cols
    
    x = np.pad(x, ((0, int(rw+1)), (0, int(cw+1))), 'constant', constant_values=np.nan)
    x[np.isinf(x)] = np.nan
    
    if ignore_nodata is not None:
        x[x==ignore_nodata] = np.nan
    
    x_min = np.nanmin(x)-1
    x = x - x_min # x >= 1

    x[np.isnan(x)] = 0 # 0 is the missing value
    x = x.astype(np.float64)
    y = np.full(exc_shape, np.nan).astype(np.float64)
    y = resample_2D_mean(x, y, rw, cw, exc_rows, exc_cols, 0) # 0 is the missing value
    y = np.asarray(y) + x_min
    
    return y



def aggregate3D_v1(data, exc_shape, n_jobs=-1):
    res = Parallel(n_jobs)(delayed(aggregate2D)(data_i, exc_shape) for data_i in data)
    res = np.array(res)
    return res



def aggregate3D_v2(x, exc_shape, ignore_nodata=None):
    """
    x: 2D array
    exc_shape: tuple of (rows, cols)
    return: 2D array
    """
    cur_times, cur_rows, cur_cols = x.shape
    exc_times, exc_rows, exc_cols = exc_shape
    rw = cur_rows / exc_rows
    cw = cur_cols / exc_cols
    
    x = np.pad(x, (0, (0, int(rw+1)), (0, int(cw+1))), 'constant', constant_values=np.nan)
    x[np.isinf(x)] = np.nan
    
    if ignore_nodata is not None:
        x[x==ignore_nodata] = np.nan
    
    x_min = np.nanmin(x)-1
    x = x - x_min # x >= 1

    x[np.isnan(x)] = 0 # 0 is the missing value
    x = x.astype(np.float64)
    y = np.full(exc_shape, np.nan).astype(np.float64)
    y = resample_3D_mean(x, y, rw, cw, exc_times, exc_rows, exc_cols, 0) # 0 is the missing value
    y = np.asarray(y) + x_min
    
    return y



def sliding_grid_partition(data: np.ndarray, target_row_blocks: int, target_col_blocks: int) -> np.ndarray:
    """
    Partition a 2D or 3D array into a grid of fixed-size blocks along its last two dimensions using a sliding window approach.
    
    For a 2D array (shape: (rows, cols)), the array is first expanded to (1, rows, cols) and then partitioned into
    target_row_blocks x target_col_blocks blocks. The output shape is:
        - For 2D input: (target_row_blocks, target_col_blocks, block_height, block_width)
        - For 3D input: (D, target_row_blocks, target_col_blocks, block_height, block_width)
    
    If the number of rows and columns can be evenly divided by the target block counts, the function uses reshape
    and transpose for efficient extraction. Otherwise, it computes a fixed block size (using ceiling division) and
    extracts blocks using sliding_window_view based on computed centers, ensuring blocks do not exceed array bounds.
    
    Parameters:
        data (np.ndarray): Input array, either a 2D array (rows, cols) or a 3D array (D, rows, cols), where D typically 
                           represents the number of slices (e.g., channels, time frames, etc.).
        target_row_blocks (int): Desired number of blocks along the row dimension.
        target_col_blocks (int): Desired number of blocks along the column dimension.
        
    Returns:
        np.ndarray: The grid-partitioned array using a sliding window approach. For 2D input, the shape is 
                    (target_row_blocks, target_col_blocks, block_height, block_width), and for 3D input, 
                    the shape is (D, target_row_blocks, target_col_blocks, block_height, block_width).
    """
    # If the input is a 2D array, expand it to 3D for uniform processing; mark to squeeze output later.
    squeeze_output = False
    if data.ndim == 2:
        data = data[np.newaxis, ...]
        squeeze_output = True
    elif data.ndim != 3:
        raise ValueError("Input data must be a 2D array (rows, cols) or a 3D array (D, rows, cols).")
    
    num_slices, orig_rows, orig_cols = data.shape
    row_scale = orig_rows / target_row_blocks
    col_scale = orig_cols / target_col_blocks

    # If the dimensions are evenly divisible, use reshape and transpose for fast extraction.
    if row_scale.is_integer() and col_scale.is_integer():
        block_row_size, block_col_size = int(row_scale), int(col_scale)
        partitioned = data.reshape(num_slices, target_row_blocks, block_row_size, target_col_blocks, block_col_size)
        partitioned = partitioned.transpose(0, 1, 3, 2, 4)
    else:
        # Compute block size with ceiling division.
        block_row_size = int(np.ceil(row_scale))
        block_col_size = int(np.ceil(col_scale))
        
        # Calculate offsets to ensure blocks do not exceed array bounds.
        offset_row_left = block_row_size // 2
        offset_row_right = block_row_size - offset_row_left
        offset_col_top = block_col_size // 2
        offset_col_bottom = block_col_size - offset_col_top
        
        # Compute the center positions for blocks along rows and columns.
        row_centers = np.linspace(offset_row_left, orig_rows - offset_row_right, target_row_blocks).round().astype(int)
        col_centers = np.linspace(offset_col_top, orig_cols - offset_col_bottom, target_col_blocks).round().astype(int)
        
        # Determine the starting indices of each block.
        row_starts = row_centers - offset_row_left
        col_starts = col_centers - offset_col_top
        
        # Use sliding_window_view to generate all possible windows of fixed block size.
        windows = sliding_window_view(data, (block_row_size, block_col_size), axis=(1, 2))
        # Create a grid to select the desired windows based on the computed starting indices.
        grid_row, grid_col = np.meshgrid(row_starts, col_starts, indexing='ij')
        partitioned = windows[:, grid_row, grid_col, :, :]
    
    if squeeze_output:
        partitioned = partitioned[0]
    return partitioned



def get_stats(data):
    """
    Compute a comprehensive descriptive statistics summary for the input numerical data.

    This function calculates basic statistics (minimum, maximum, mean, median, standard deviation,
    and variance) as well as advanced metrics (skewness, kurtosis, various percentiles, interquartile range, 
    and range). It also counts the number of valid (non-NaN) data points and the number of NaN values.
    
    Parameters:
        data (np.ndarray): 
            A numerical array (1D or multi-dimensional) that may contain NaN values. The function 
            automatically filters out NaN values and performs all statistical calculations on the valid data.
    
    Returns:
        pd.DataFrame: 
            A single-row DataFrame containing the following statistics:
              - Minimum: The minimum value in the data.
              - Mean: The average value.
              - Maximum: The maximum value in the data.
              - Median: The median value.
              - Standard Deviation: The standard deviation.
              - Variance: The sample variance (with degrees of freedom = 1).
              - Skewness: The skewness of the data distribution (indicating asymmetry).
              - Kurtosis: The kurtosis of the data distribution (indicating tailedness).
              - Range: The range (difference between maximum and minimum).
              - Count / Valid Count: The number of valid (non-NaN) data points.
              - NaN Count: The number of NaN values in the original data.
              - Interquartile Range (IQR): The difference between the 75th and 25th percentiles.
              - 1th Percentile: The 1st percentile of the data.
              - 5th Percentile: The 5th percentile of the data.
              - 10th Percentile: The 10th percentile of the data.
              - 90th Percentile: The 90th percentile of the data.
              - 95th Percentile: The 95th percentile of the data.
              - 99th Percentile: The 99th percentile of the data.
              - Mean - 3*StdDev: The value of the mean minus three times the standard deviation.
              - Mean + 3*StdDev: The value of the mean plus three times the standard deviation.
              - Sum: The sum of all valid data points.
    
    Usage Example:
        >>> np.random.seed(42)
        >>> data = np.random.randn(1000)
        >>> nan_indices = np.random.choice(len(data), size=50, replace=False)
        >>> data[nan_indices] = np.nan
        >>> stats_df = get_stats(data)
        >>> print(stats_df)
    
    The output is a DataFrame with one row containing the computed statistical metrics.
    """
    # Count the number of NaN values
    nan_count = np.isnan(data).sum()
    # Filter out NaN values for calculations
    clean_data = data[~np.isnan(data)]
    valid_count = len(clean_data)
    
    # Calculate basic statistics
    mean = np.mean(clean_data)
    std_dev = np.std(clean_data)  # Standard deviation
    
    # Create a dictionary of statistics
    stats_dict = {
        'Minimum': np.min(clean_data),
        'Mean': mean,
        'Maximum': np.max(clean_data),
        'Median': np.median(clean_data),
        'Standard Deviation': std_dev,
        'Variance': np.var(clean_data, ddof=1),
        'Skewness': stats.skew(clean_data),
        'Kurtosis': stats.kurtosis(clean_data),
        'Range': np.ptp(clean_data),
        'Count': valid_count,
        'Valid Count': valid_count,
        'NaN Count': nan_count,
        'Interquartile Range (IQR)': np.percentile(clean_data, 75) - np.percentile(clean_data, 25),
        '1th Percentile': np.percentile(clean_data, 1),
        '5th Percentile': np.percentile(clean_data, 5),
        '10th Percentile': np.percentile(clean_data, 10),
        '90th Percentile': np.percentile(clean_data, 90),
        '95th Percentile': np.percentile(clean_data, 95),
        '99th Percentile': np.percentile(clean_data, 99),
        'Mean - 3*StdDev': mean - 3 * std_dev,
        'Mean + 3*StdDev': mean + 3 * std_dev,
        'Sum': np.sum(clean_data),
    }
    
    # Convert the dictionary to a DataFrame and return
    stats_df = pd.DataFrame([stats_dict])
    return stats_df



class Aggregate:

    def __init__(self, data: np.ndarray, exc_shape:tuple):
        '''
        Initialize Aggregator object.

        Parameters:
        data (numpy.ndarray): The input data to be aggregated.
        exc_shape (tuple): The expected shape of the output data.
        '''
        self.data = data
        self.exc_shape = exc_shape
        self.exc_rows, self.exc_cols = self.exc_shape
        self.cur_rows, self.cur_cols = self.data.shape
        self.r_res, self.c_res = self.cur_rows/self.exc_rows, self.cur_cols/self.exc_cols

        pass
        
    def mode(self): return self.fit('mode')
    def mean(self): return self.fit('mean')

    
    def fit(self, operation: str = 'mode'):
        '''
        Aggregate the data using mode operation.

        Returns:
        numpy.ndarray: The aggregated data.
        '''
       
        if self.r_res.is_integer() and self.c_res.is_integer():
            self.data = self.data.reshape(self.exc_rows, int(self.r_res), self.exc_cols, int(self.c_res))
            self.data = self.data.transpose(0, 2, 1, 3).reshape(self.exc_rows, self.exc_cols, -1)
            if operation == 'mode':
                return self._mode_exa_div()
            elif operation == 'mean':
                return np.nanmean(self.data, axis=-1)
            else:
                raise ValueError("Invalid operation")
            
        else:
            return self._mode_not_div(operation)




    def _get_boundary(self, r, ratio):
        r_sta = r*ratio
        r_end = (r+1)*ratio

        rmin = int(r_sta)
        rmax = np.ceil(r_end)

        offset_1 = np.ceil(r_sta) - r_sta # 上/左 偏移量
        offset_2 = r_end - int(r_end) # 下/右 偏移量
        return rmin, rmax, offset_1, offset_2
    


    def _get_area(self, win_shape, offsets):

        offset_1, offset_2, offset_3, offset_4 = offsets # 上下左右
        areas = np.full(win_shape, 1, dtype=float)
        if offset_1 > 0: areas[ 0, :] = offset_1 
        if offset_2 > 0: areas[-1, :] = offset_2
        if offset_3 > 0: areas[:,  0] = offset_3 * areas[:,  0]
        if offset_4 > 0: areas[:, -1] = offset_4 * areas[:, -1]

        return areas



    def _mode_not_div(self, operation):
        result = np.full(self.exc_shape, np.nan)
        for r in tqdm(range(self.exc_rows)):
            rmin, rmax, offset_1, offset_2 = self._get_boundary(r, self.r_res)
            
            for c in range(self.exc_cols):
                cmin, cmax, offset_3, offset_4 = self._get_boundary(c, self.c_res)
                data_win = self.data[rmin:rmax, cmin:cmax]
                mask = ~np.isnan(data_win)
                if not mask.any():
                    continue

                area_win = self._get_area(data_win.shape, (offset_1, offset_2, offset_3, offset_4))
                data_ij = data_win[mask].astype(np.int64)
                area_ij = area_win[mask]

                if operation == 'mode':
                    counts = np.bincount(data_ij, weights=area_ij)*1e9
                    result[r, c] = np.argmax(counts.astype(np.int64)) # areas整型或浮点型 对np.argmax(areas)的结果有影响。具体影响在于当多个val并列第一的时候, 该取哪一个? 据简单观测, 整型-->取最小的, 浮点型-->取最大
                elif operation == 'mean':
                    result[r, c] = np.sum(data_ij*area_ij)/np.sum(area_ij)
        
        return result


    def _mode_exa_div(self,):
        mask = np.any(~np.isnan(self.data), axis=-1)
        result = np.full(mask.shape, np.nan)
        result[mask] = np.apply_along_axis(lambda x: np.argmax(np.bincount(np.int64(x[~np.isnan(x)]))), axis=-1, arr=self.data[mask])
        return result
    


class SlidingWindowAnalysis: # sliding window analysis
    def __init__(self, step=10, window_size=10):
        self.step = step
        self.winsize = window_size
        self.l_ws = window_size // 2
        self.r_ws = window_size - self.l_ws
        self.data = None
        self.res = None
        self.concatenate_axis = 1

        
    def _pad(self, data):
        pad_width = [(self.r_ws, self.l_ws)]*2 + [(0, 0)]*(data.ndim - 2)
        if self.concatenate_axis is not None:
            if self.concatenate_axis != 1:
                data = np.swapaxes(data, self.concatenate_axis, 1)
            data = np.concatenate((data[:, -self.r_ws:], data, data[:, :self.l_ws]), axis=1) # 左右填充
            if self.concatenate_axis != 1:
                data = np.swapaxes(data, self.concatenate_axis, 1)
            pad_width[1] = (0, 0)
        data = np.pad(data, pad_width, mode='constant', constant_values=np.nan)
        return data


    def slide(self, data):
        data_list = []
        height, width, bands = data.shape # (500, 400, 3)
        x_range = np.arange(self.l_ws, height + self.l_ws, self.step).astype(int)
        y_range = np.arange(self.l_ws, width + self.l_ws, self.step).astype(int)
        data = self._pad(data)
        for i in x_range:
            i_start = i - self.l_ws
            i_end = i + self.r_ws
            for j in y_range:
                j_start = j - self.l_ws
                j_end = j + self.r_ws
                data_list.append(data[i_start:i_end, j_start:j_end])
        self.slide_data = np.array(data_list).reshape(len(x_range), len(y_range), self.winsize, self.winsize, bands)
        return self.slide_data
    
    def pcorr():
        pass


