"""
Tests for custom metric functions.

Ref: revisions.txt Task 6.5
"""

import pytest
import numpy as np
import pandas as pd
import sys, os

# --- Add path for src imports --- #
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
src_path = os.path.join(project_root, 'src')
if src_path not in sys.path:
    sys.path.insert(0, src_path)
# --- End Add path --- #

from metrics import edge_filtered_accuracy, calculate_sharpe_ratio

# --- Tests for edge_filtered_accuracy --- #

def test_edge_filtered_accuracy_basic():
    """Test basic functionality with hard labels and clear edge."""
    y_true = np.array([1, 0, 1, 0, 1, 1, 0, 0])
    p_cal  = np.array([0.9, 0.1, 0.8, 0.2, 0.7, 0.6, 0.3, 0.4]) # Edge > 0.1 for all
    thr = 0.1
    
    accuracy, n_filtered = edge_filtered_accuracy(y_true, p_cal, thr=thr)
    
    assert n_filtered == 8
    # Predictions: 1, 0, 1, 0, 1, 1, 0, 0. All correct.
    assert accuracy == pytest.approx(1.0)

def test_edge_filtered_accuracy_thresholding():
    """Test that the threshold correctly filters samples."""
    y_true = np.array([1, 0, 1, 0, 1, 1, 0, 0]) 
    p_cal  = np.array([0.9, 0.1, 0.8, 0.2, 0.51, 0.49, 0.55, 0.45]) # Edge: 0.8, 0.8, 0.6, 0.6, 0.02, 0.02, 0.1, 0.1
    
    # Test with thr=0.15 (should exclude last 4 samples)
    thr1 = 0.15
    accuracy1, n_filtered1 = edge_filtered_accuracy(y_true, p_cal, thr=thr1)
    assert n_filtered1 == 4
    # Predictions on first 4: 1, 0, 1, 0. All correct.
    assert accuracy1 == pytest.approx(1.0)
    
    # Test with thr=0.05 (should include all but middle 2)
    thr2 = 0.05
    accuracy2, n_filtered2 = edge_filtered_accuracy(y_true, p_cal, thr=thr2)
    assert n_filtered2 == 6
    # Included: 1,0,1,0, 1, 0. Correct: 1,0,1,0, ?, ?. Preds: 1,0,1,0, 1, 0. 6/6 correct.
    assert accuracy2 == pytest.approx(1.0)

def test_edge_filtered_accuracy_soft_labels():
    """Test with soft labels."""
    y_true_soft = np.array([0.9, 0.1, 0.8, 0.2, 0.7, 0.6]) # Soft labels
    p_cal       = np.array([0.8, 0.3, 0.9, 0.1, 0.6, 0.7]) # All edge > 0.1
    thr = 0.1
    
    accuracy, n_filtered = edge_filtered_accuracy(y_true_soft, p_cal, thr=thr)
    
    assert n_filtered == 6
    # y_true_hard: 1, 0, 1, 0, 1, 1
    # y_pred     : 1, 0, 1, 0, 1, 1. All correct.
    assert accuracy == pytest.approx(1.0)

def test_edge_filtered_accuracy_no_samples():
    """Test case where no samples meet the edge threshold."""
    y_true = np.array([1, 0, 1, 0])
    p_cal  = np.array([0.51, 0.49, 0.52, 0.48]) # All edge < 0.1
    thr = 0.1
    
    accuracy, n_filtered = edge_filtered_accuracy(y_true, p_cal, thr=thr)
    assert n_filtered == 0
    assert np.isnan(accuracy)

def test_edge_filtered_accuracy_empty_input():
    """Test with empty input arrays."""
    y_true = np.array([])
    p_cal  = np.array([])
    thr = 0.1
    
    accuracy, n_filtered = edge_filtered_accuracy(y_true, p_cal, thr=thr)
    assert n_filtered == 0
    assert np.isnan(accuracy)

# --- Tests for calculate_sharpe_ratio --- #

def test_calculate_sharpe_ratio_basic():
    """Test basic Sharpe calculation."""
    returns = pd.Series([0.01, -0.005, 0.02, 0.005, -0.01])
    # mean = 0.004, std = 0.01166, Sharpe_period = 0.343
    # Annualized (252) = 0.343 * sqrt(252) = 5.44
    expected_sharpe = 5.44441
    sharpe = calculate_sharpe_ratio(returns, benchmark_return=0.0, annualization_factor=252)
    assert sharpe == pytest.approx(expected_sharpe, abs=1e-4)

def test_calculate_sharpe_ratio_different_annualization():
    """Test Sharpe with different annualization factor."""
    returns = pd.Series([0.01, -0.005, 0.02, 0.005, -0.01])
    # Annualized (52) = 0.343 * sqrt(52) = 2.47
    expected_sharpe = 2.4738
    sharpe = calculate_sharpe_ratio(returns, benchmark_return=0.0, annualization_factor=52)
    assert sharpe == pytest.approx(expected_sharpe, abs=1e-4)

def test_calculate_sharpe_ratio_with_benchmark():
    """Test Sharpe with a non-zero benchmark return."""
    returns = pd.Series([0.01, -0.005, 0.02, 0.005, -0.01]) # mean=0.004
    benchmark = 0.001 # Per period
    # excess mean = 0.003, std = 0.01166, Sharpe_period = 0.257
    # Annualized (252) = 0.257 * sqrt(252) = 4.08
    expected_sharpe = 4.0833
    sharpe = calculate_sharpe_ratio(returns, benchmark_return=benchmark, annualization_factor=252)
    assert sharpe == pytest.approx(expected_sharpe, abs=1e-4)

def test_calculate_sharpe_ratio_zero_std():
    """Test Sharpe when returns have zero standard deviation."""
    returns_positive = pd.Series([0.01, 0.01, 0.01])
    returns_negative = pd.Series([-0.01, -0.01, -0.01])
    returns_zero = pd.Series([0.0, 0.0, 0.0])
    
    assert calculate_sharpe_ratio(returns_positive) == 0.0 # Positive mean, zero std -> 0?
    # assert calculate_sharpe_ratio(returns_negative) == -np.inf # Negative mean, zero std -> -inf?
    assert calculate_sharpe_ratio(returns_zero) == 0.0
    
    # Let's refine zero std handling based on function's logic
    # Function returns 0 if mean>0, -inf if mean<0, 0 if mean=0
    assert calculate_sharpe_ratio(returns_positive) == 0.0
    assert calculate_sharpe_ratio(returns_negative) == -np.inf
    assert calculate_sharpe_ratio(returns_zero) == 0.0

def test_calculate_sharpe_ratio_empty_or_nan():
    """Test Sharpe with empty or all-NaN input."""
    assert np.isnan(calculate_sharpe_ratio(pd.Series([], dtype=float)))
    assert np.isnan(calculate_sharpe_ratio(pd.Series([np.nan, np.nan], dtype=float))) 