"""
Numba-accelerated functions for PASI algorithm.

This module contains JIT-compiled versions of performance-critical functions
to significantly speed up the tree building process.

Author: PASI Package
"""

import numpy as np
from numba import jit, prange
from typing import Tuple


# =============================================================================
# Individual Accuracy Estimation (measure='indiv')
# =============================================================================

@jit(nopython=True, cache=True)
def indiv_acc_estimation_numba(mu: np.ndarray) -> Tuple[float, float, float, float, float]:
    """
    Numba-accelerated individual accuracy estimation.
    
    Parameters
    ----------
    mu : np.ndarray
        Array of individual accuracy values.
        
    Returns
    -------
    Tuple[float, float, float, float, float]
        (mu_hat, var_hat, sd_hat, ci_lower, ci_upper)
    """
    N = mu.shape[0]
    if N == 0:
        return np.nan, np.nan, np.nan, np.nan, np.nan
    
    # Calculate mean
    mu_hat = 0.0
    for i in range(N):
        mu_hat += mu[i]
    mu_hat /= N
    
    # Calculate variance
    var_sum = 0.0
    for i in range(N):
        diff = mu[i] - mu_hat
        var_sum += diff * diff
    var_hat = var_sum / (N * N)  # var / N for standard error
    
    # Standard deviation and CI
    sd_hat = np.sqrt(var_hat)
    c = 1.96
    ci_l = mu_hat - c * sd_hat
    ci_u = mu_hat + c * sd_hat
    
    return mu_hat, var_hat, sd_hat, ci_l, ci_u


# =============================================================================
# Split Point Search for Individual Accuracy (measure='indiv')
# =============================================================================

@jit(nopython=True, cache=True)
def find_best_split_indiv_numba(
    x_sorted: np.ndarray,
    mu_sorted: np.ndarray,
    cut_pts: np.ndarray,
    min_samples_leaf: int
) -> Tuple[int, float]:
    """
    Numba-accelerated best split search for individual accuracy measure.
    
    This function finds the best split point by evaluating each candidate
    cut point using the split criterion: (mu_left - mu_right)^2 / (var_left + var_right)
    
    Parameters
    ----------
    x_sorted : np.ndarray
        Feature values sorted in ascending order.
    mu_sorted : np.ndarray  
        Individual accuracy values sorted by x.
    cut_pts : np.ndarray
        Candidate split points (midpoints between unique x values).
    min_samples_leaf : int
        Minimum samples required in each leaf.
        
    Returns
    -------
    Tuple[int, float]
        (best_cut_idx, best_score) where best_cut_idx is the index into cut_pts
        and best_score is the split quality score. Returns (-1, 0.0) if no valid split.
    """
    n = x_sorted.shape[0]
    n_cuts = cut_pts.shape[0]
    
    if n_cuts == 0:
        return -1, 0.0
    
    best_cut_idx = -1
    best_score = 0.0
    
    # Use cumulative sums for O(1) mean/variance computation per split
    # Precompute cumulative sums
    cum_sum = np.zeros(n + 1, dtype=np.float64)
    cum_sum_sq = np.zeros(n + 1, dtype=np.float64)
    
    for i in range(n):
        cum_sum[i + 1] = cum_sum[i] + mu_sorted[i]
        cum_sum_sq[i + 1] = cum_sum_sq[i] + mu_sorted[i] * mu_sorted[i]
    
    total_sum = cum_sum[n]
    total_sum_sq = cum_sum_sq[n]
    
    # Find split position for each cut point
    i = 0  # Current position in sorted array
    
    for cut_idx in range(n_cuts):
        cut_pt = cut_pts[cut_idx]
        
        # Move i to the first position where x > cut_pt
        while i < n and x_sorted[i] <= cut_pt:
            i += 1
        
        n_left = i
        n_right = n - n_left
        
        # Check minimum samples constraint
        if n_left < min_samples_leaf or n_right < min_samples_leaf:
            continue
        
        # Compute left statistics
        left_sum = cum_sum[n_left]
        left_sum_sq = cum_sum_sq[n_left]
        left_mean = left_sum / n_left
        left_var = (left_sum_sq / n_left - left_mean * left_mean) / n_left
        
        # Compute right statistics
        right_sum = total_sum - left_sum
        right_sum_sq = total_sum_sq - left_sum_sq
        right_mean = right_sum / n_right
        right_var = (right_sum_sq / n_right - right_mean * right_mean) / n_right
        
        # Compute split score
        de = left_var + right_var
        if de > 0:
            nu = (left_mean - right_mean) ** 2
            score = nu / de
            
            if score > best_score:
                best_score = score
                best_cut_idx = cut_idx
    
    return best_cut_idx, best_score


@jit(nopython=True, cache=True)
def evaluate_all_splits_indiv_numba(
    x_sorted: np.ndarray,
    mu_sorted: np.ndarray,
    min_samples_leaf: int
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Evaluate all possible split points and return scores.
    
    This version iterates through the sorted array and evaluates each unique
    value boundary as a potential split point, using incremental statistics.
    
    Parameters
    ----------
    x_sorted : np.ndarray
        Feature values sorted in ascending order.
    mu_sorted : np.ndarray
        Individual accuracy values sorted by x.
    min_samples_leaf : int
        Minimum samples required in each leaf.
        
    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        (split_positions, scores) where split_positions[i] is the number of 
        samples in the left branch and scores[i] is the corresponding split score.
    """
    n = x_sorted.shape[0]
    
    # Maximum possible splits
    max_splits = n - 1
    split_positions = np.zeros(max_splits, dtype=np.int64)
    scores = np.zeros(max_splits, dtype=np.float64)
    
    if n < 2 * min_samples_leaf:
        return split_positions[:0], scores[:0]
    
    # Precompute total statistics
    total_sum = 0.0
    total_sum_sq = 0.0
    for i in range(n):
        total_sum += mu_sorted[i]
        total_sum_sq += mu_sorted[i] * mu_sorted[i]
    
    # Incremental left statistics
    left_sum = 0.0
    left_sum_sq = 0.0
    
    split_count = 0
    i = 0
    
    while i < n - 1:
        # Add current element to left partition
        left_sum += mu_sorted[i]
        left_sum_sq += mu_sorted[i] * mu_sorted[i]
        
        # Skip to the last element with the same x value
        while i < n - 1 and x_sorted[i] == x_sorted[i + 1]:
            i += 1
            left_sum += mu_sorted[i]
            left_sum_sq += mu_sorted[i] * mu_sorted[i]
        
        n_left = i + 1
        n_right = n - n_left
        
        # Check minimum samples constraint
        if n_left >= min_samples_leaf and n_right >= min_samples_leaf:
            # Compute left statistics
            left_mean = left_sum / n_left
            left_var = (left_sum_sq / n_left - left_mean * left_mean) / n_left
            
            # Compute right statistics
            right_sum = total_sum - left_sum
            right_sum_sq = total_sum_sq - left_sum_sq
            right_mean = right_sum / n_right
            right_var = (right_sum_sq / n_right - right_mean * right_mean) / n_right
            
            # Compute split score
            de = left_var + right_var
            if de > 0:
                nu = (left_mean - right_mean) ** 2
                score = nu / de
                
                split_positions[split_count] = n_left
                scores[split_count] = score
                split_count += 1
        
        i += 1
    
    return split_positions[:split_count], scores[:split_count]


# =============================================================================
# DeLong AUC Computation (optimized)
# =============================================================================

@jit(nopython=True, cache=True)
def compute_midrank_numba(x: np.ndarray) -> np.ndarray:
    """
    Numba-accelerated midrank computation.
    
    Parameters
    ----------
    x : np.ndarray
        1D array of values.
        
    Returns
    -------
    np.ndarray
        Array of midranks.
    """
    J = np.argsort(x)
    Z = x[J]
    N = len(x)
    T = np.zeros(N, dtype=np.float64)
    
    i = 0
    while i < N:
        j = i
        while j < N and Z[j] == Z[i]:
            j += 1
        midrank = 0.5 * (i + j - 1)
        for k in range(i, j):
            T[k] = midrank
        i = j
    
    T2 = np.empty(N, dtype=np.float64)
    for i in range(N):
        T2[J[i]] = T[i] + 1
    
    return T2


@jit(nopython=True, cache=True)
def fast_delong_numba(y_pred_sorted: np.ndarray, m: int) -> Tuple[float, float]:
    """
    Numba-accelerated DeLong AUC and variance computation.
    
    Parameters
    ----------
    y_pred_sorted : np.ndarray
        Predictions sorted by descending ground truth (positives first).
    m : int
        Number of positive examples.
        
    Returns
    -------
    Tuple[float, float]
        (auc, variance)
    """
    n_total = y_pred_sorted.shape[0]
    n = n_total - m  # Number of negatives
    
    if m == 0 or n == 0:
        return 0.5, 0.0
    
    positive_examples = y_pred_sorted[:m]
    negative_examples = y_pred_sorted[m:]
    
    # Compute midranks
    tx = compute_midrank_numba(positive_examples)
    ty = compute_midrank_numba(negative_examples)
    tz = compute_midrank_numba(y_pred_sorted)
    
    # Compute AUC
    tz_pos = tz[:m]
    auc = 0.0
    for i in range(m):
        auc += tz_pos[i]
    auc = auc / m / n - (m + 1.0) / 2.0 / n
    
    # Compute variance components
    v01 = np.empty(m, dtype=np.float64)
    v10 = np.empty(n, dtype=np.float64)
    
    for i in range(m):
        v01[i] = (tz[i] - tx[i]) / n
    for i in range(n):
        v10[i] = 1.0 - (tz[m + i] - ty[i]) / m
    
    # Compute covariance (sample variance)
    if m > 1:
        mean_v01 = 0.0
        for i in range(m):
            mean_v01 += v01[i]
        mean_v01 /= m
        
        sx = 0.0
        for i in range(m):
            diff = v01[i] - mean_v01
            sx += diff * diff
        sx /= (m - 1)
    else:
        sx = 0.0
    
    if n > 1:
        mean_v10 = 0.0
        for i in range(n):
            mean_v10 += v10[i]
        mean_v10 /= n
        
        sy = 0.0
        for i in range(n):
            diff = v10[i] - mean_v10
            sy += diff * diff
        sy /= (n - 1)
    else:
        sy = 0.0
    
    delongcov = sx / m + sy / n
    
    return auc, delongcov


@jit(nopython=True, cache=True)
def delong_auc_variance_numba(y: np.ndarray, y_pred: np.ndarray) -> Tuple[float, float]:
    """
    Numba-accelerated DeLong AUC and variance computation.
    
    Parameters
    ----------
    y : np.ndarray
        Binary ground truth labels (0 or 1).
    y_pred : np.ndarray
        Predicted probabilities.
        
    Returns
    -------
    Tuple[float, float]
        (auc, variance)
    """
    n = y.shape[0]
    if n == 0:
        return np.nan, np.nan
    
    # Count positives
    m = 0
    for i in range(n):
        if y[i] == 1:
            m += 1
    
    if m == 0 or m == n:
        return 0.5, 0.0
    
    # Sort by descending ground truth (positives first)
    order = np.argsort(-y)
    y_pred_sorted = np.empty(n, dtype=np.float64)
    for i in range(n):
        y_pred_sorted[i] = y_pred[order[i]]
    
    return fast_delong_numba(y_pred_sorted, m)


# =============================================================================
# Split Search for AUC measure
# =============================================================================

@jit(nopython=True, cache=True)
def evaluate_all_splits_auc_numba(
    x_sorted: np.ndarray,
    y_sorted: np.ndarray,
    y_pred_sorted: np.ndarray,
    min_samples_leaf: int,
    min_pos_neg: int
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Evaluate all possible split points for AUC measure.
    
    Parameters
    ----------
    x_sorted : np.ndarray
        Feature values sorted in ascending order.
    y_sorted : np.ndarray
        Binary labels sorted by x.
    y_pred_sorted : np.ndarray
        Predictions sorted by x.
    min_samples_leaf : int
        Minimum samples in each leaf.
    min_pos_neg : int
        Minimum positive and negative samples in each leaf.
        
    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        (split_positions, scores)
    """
    n = x_sorted.shape[0]
    
    max_splits = n - 1
    split_positions = np.zeros(max_splits, dtype=np.int64)
    scores = np.zeros(max_splits, dtype=np.float64)
    
    if n < 2 * min_samples_leaf:
        return split_positions[:0], scores[:0]
    
    # Count total positives
    total_pos = 0
    for i in range(n):
        if y_sorted[i] == 1:
            total_pos += 1
    total_neg = n - total_pos
    
    if total_pos < 2 * min_pos_neg or total_neg < 2 * min_pos_neg:
        return split_positions[:0], scores[:0]
    
    split_count = 0
    i = 0
    left_pos = 0
    left_neg = 0
    
    while i < n - 1:
        # Add current element to left
        if y_sorted[i] == 1:
            left_pos += 1
        else:
            left_neg += 1
        
        # Skip to last element with same x value
        while i < n - 1 and x_sorted[i] == x_sorted[i + 1]:
            i += 1
            if y_sorted[i] == 1:
                left_pos += 1
            else:
                left_neg += 1
        
        n_left = i + 1
        n_right = n - n_left
        right_pos = total_pos - left_pos
        right_neg = total_neg - left_neg
        
        # Check constraints
        if (n_left >= min_samples_leaf and n_right >= min_samples_leaf and
            left_pos >= min_pos_neg and left_neg >= min_pos_neg and
            right_pos >= min_pos_neg and right_neg >= min_pos_neg):
            
            # Compute AUC and variance for left partition
            left_auc, left_var = delong_auc_variance_numba(
                y_sorted[:n_left], y_pred_sorted[:n_left]
            )
            
            # Compute AUC and variance for right partition
            right_auc, right_var = delong_auc_variance_numba(
                y_sorted[n_left:], y_pred_sorted[n_left:]
            )
            
            # Compute split score
            de = left_var + right_var
            if de > 0 and not np.isnan(left_auc) and not np.isnan(right_auc):
                nu = (left_auc - right_auc) ** 2
                score = nu / de
                
                split_positions[split_count] = n_left
                scores[split_count] = score
                split_count += 1
        
        i += 1
    
    return split_positions[:split_count], scores[:split_count]


# =============================================================================
# Utility functions
# =============================================================================

@jit(nopython=True, cache=True)
def compute_cut_points(x_sorted: np.ndarray) -> np.ndarray:
    """
    Compute cut points (midpoints between unique values) from sorted array.
    
    Parameters
    ----------
    x_sorted : np.ndarray
        Sorted feature values.
        
    Returns
    -------
    np.ndarray
        Cut points.
    """
    n = x_sorted.shape[0]
    if n < 2:
        return np.empty(0, dtype=np.float64)
    
    # Find unique values
    unique_count = 1
    for i in range(1, n):
        if x_sorted[i] != x_sorted[i - 1]:
            unique_count += 1
    
    if unique_count < 2:
        return np.empty(0, dtype=np.float64)
    
    # Compute cut points
    cut_pts = np.empty(unique_count - 1, dtype=np.float64)
    cut_idx = 0
    prev_val = x_sorted[0]
    
    for i in range(1, n):
        if x_sorted[i] != prev_val:
            cut_pts[cut_idx] = (prev_val + x_sorted[i]) / 2.0
            cut_idx += 1
            prev_val = x_sorted[i]
    
    return cut_pts


@jit(nopython=True, cache=True)
def get_threshold_from_position(x_sorted: np.ndarray, split_pos: int) -> float:
    """
    Get the threshold value for a given split position.
    
    Parameters
    ----------
    x_sorted : np.ndarray
        Sorted feature values.
    split_pos : int
        Number of samples in the left partition.
        
    Returns
    -------
    float
        Threshold value (midpoint between last left and first right value).
    """
    if split_pos <= 0 or split_pos >= len(x_sorted):
        return np.nan
    return (x_sorted[split_pos - 1] + x_sorted[split_pos]) / 2.0
