"""
Custom Metrics for Trading Performance Evaluation.

Ref: revisions.txt Section 6
"""

import numpy as np
from typing import Tuple
import logging
import pandas as pd
from sklearn.metrics import brier_score_loss

logger = logging.getLogger(__name__)

def edge_filtered_accuracy(y_true: np.ndarray, p_cal: np.ndarray, thr: float = 0.1) -> Tuple[float, int]:
    """
    Calculates accuracy only on samples where the calibrated prediction 
    has sufficient edge (confidence).

    Args:
        y_true (np.ndarray): True binary labels (0 or 1, potentially soft).
        p_cal (np.ndarray): Calibrated probabilities P(up) (shape: N,).
        thr (float): Edge threshold. Only samples where |2*p_cal - 1| >= thr are included.
                     Defaults to 0.1 (equivalent to p_cal >= 0.55 or p_cal <= 0.45).

    Returns:
        Tuple[float, int]: 
            - Accuracy on the filtered samples (NaN if no samples meet threshold).
            - Number of samples included in the calculation.
    """
    if len(y_true) != len(p_cal):
        raise ValueError("Length mismatch between y_true and p_cal.")
    if len(y_true) == 0:
        return np.nan, 0
    
    y_true = np.asarray(y_true)
    p_cal = np.asarray(p_cal)
    
    # Calculate edge
    edge = np.abs(2 * p_cal - 1)
    
    # Create mask
    mask = edge >= thr
    n_filtered = int(np.sum(mask))
    
    if n_filtered == 0:
        logger.warning(f"No samples met edge threshold {thr:.3f}. Cannot calculate edge-filtered accuracy.")
        return np.nan, 0
        
    # Filter data
    p_cal_filtered = p_cal[mask]
    y_true_filtered = y_true[mask]
    
    # Predict direction based on calibrated probability > 0.5
    y_pred_filtered = (p_cal_filtered > 0.5).astype(int)
    
    # Handle potentially soft true labels
    if not np.all((y_true_filtered == 0) | (y_true_filtered == 1)):
        logger.debug("Soft labels detected in y_true_filtered. Comparing predictions to > 0.5 threshold.")
        y_true_hard_filtered = (y_true_filtered > 0.5).astype(int)
    else:
        y_true_hard_filtered = y_true_filtered.astype(int)
        
    # Calculate accuracy
    accuracy = np.mean(y_pred_filtered == y_true_hard_filtered)
    
    # logger.debug(f"Edge>={thr:.2f}: Acc={accuracy:.4f}, N={n_filtered}/{len(y_true)}")
    return accuracy, n_filtered

def calculate_brier_score(y_true: np.ndarray, p_cal: np.ndarray) -> float:
    """
    Calculates the Brier score for predicted probabilities.

    Args:
        y_true (np.ndarray): True binary labels (0 or 1, potentially soft).
        p_cal (np.ndarray): Calibrated probabilities P(up) (shape: N,).

    Returns:
        float: The Brier score (lower is better).
    """
    if len(y_true) != len(p_cal):
        raise ValueError("Length mismatch between y_true and p_cal.")
    if len(y_true) == 0:
        return np.nan

    y_true = np.asarray(y_true)
    p_cal = np.asarray(p_cal)

    # Handle soft labels by converting them to binary (0 or 1) for Brier score
    if not np.all((y_true == 0) | (y_true == 1)):
        logger.debug("Soft labels detected in y_true. Converting to hard binary for Brier score calculation.")
        y_true_hard = (y_true > 0.5).astype(int)
    else:
        y_true_hard = y_true.astype(int)
        
    # Ensure probabilities are within [0, 1]
    p_cal = np.clip(p_cal, 0.0, 1.0)

    return brier_score_loss(y_true_hard, p_cal)

# --- TODO: Add other metrics from Section 6 --- #
# - CI lower bound calculation helper? (Done implicitly in pipeline check)
# - Re-centred Sharpe calculation?

def calculate_sharpe_ratio(returns: pd.Series, benchmark_return: float = 0.0, annualization_factor: int = 252*24*60) -> float:
    """
    Calculates the annualized Sharpe ratio relative to a benchmark.

    Args:
        returns (pd.Series): Series of portfolio returns (e.g., daily or per-period).
                               Assumes returns are fractional (e.g., 0.01 for 1%).
        benchmark_return (float): The benchmark return per period (e.g., risk-free rate).
                                    Defaults to 0.0.
        annualization_factor (int): Factor to annualize Sharpe ratio (e.g., 252 for daily,
                                      52 for weekly, 12 for monthly, 252*24*60 for 1-min).

    Returns:
        float: The annualized Sharpe ratio.
    """
    if not isinstance(returns, pd.Series):
        returns = pd.Series(returns)
        
    if returns.empty or returns.isnull().all():
        return np.nan

    # Calculate excess returns over the benchmark
    excess_returns = returns - benchmark_return
    
    # Calculate mean and standard deviation of excess returns
    mean_excess_return = excess_returns.mean()
    std_excess_return = excess_returns.std()
    
    if std_excess_return == 0 or np.isnan(std_excess_return):
        # Handle cases with zero or undefined volatility
        return 0.0 if mean_excess_return > 0 else (-np.inf if mean_excess_return < 0 else 0.0)
        
    # Calculate per-period Sharpe ratio
    sharpe_ratio_period = mean_excess_return / std_excess_return
    
    # Annualize Sharpe ratio
    annualized_sharpe_ratio = sharpe_ratio_period * np.sqrt(annualization_factor)
    
    return annualized_sharpe_ratio 

# --- Helper for Youden's J Optimization --- #
def _calculate_optimal_edge_threshold(y_true: np.ndarray, p_cal: np.ndarray) -> float:
    """
    Finds the optimal edge threshold by maximizing Youden's J statistic.
    
    Args:
        y_true: True binary labels (0 or 1).
        p_cal: Calibrated probabilities (P(1)).

    Returns:
        Optimal threshold (float).
    """
    try:
        from sklearn.metrics import roc_curve
    except ImportError:
        logger.error("Scikit-learn not installed. Cannot optimize edge threshold. Returning default 0.1")
        return 0.1 # Default fallback

    if len(np.unique(y_true)) < 2:
         logger.warning("Only one class present in y_true for threshold optimization. Returning default 0.1")
         return 0.1

    fpr, tpr, thresholds = roc_curve(y_true, p_cal)
    youden_j = tpr - fpr
    # Filter out invalid thresholds (often includes inf/-inf)
    valid_thresholds = thresholds[np.isfinite(thresholds)]
    if len(valid_thresholds) == 0:
        logger.warning("No valid thresholds found during ROC curve calculation. Returning default 0.1")
        return 0.1
        
    # Ensure we handle the case where thresholds might not align perfectly with youden_j after filtering
    optimal_idx = np.argmax(youden_j[np.isfinite(thresholds)])
    optimal_threshold = valid_thresholds[optimal_idx]

    # Ensure threshold is within a reasonable range (e.g., 0.01 to 0.49) to avoid extremes
    optimal_threshold = np.clip(optimal_threshold, 0.01, 0.49)
    
    logger.info(f"Optimized edge threshold via Youden's J: {optimal_threshold:.4f}")
    return optimal_threshold
# --- End Helper --- # 