import numpy as np
from scipy.stats import spearmanr


def compute_rank_ic(scores, returns):
    """
    Compute Rank Information Coefficient (Spearman correlation) between factor scores and returns.

    Parameters:
    - scores: 1D numpy array of factor scores
    - returns: 1D numpy array of next-day returns

    Returns:
    - Single float value of Spearman Rank IC, or np.nan if not enough valid values
    """
    scores = np.asarray(scores)
    returns = np.asarray(returns)
    valid_mask = ~np.isnan(scores) & ~np.isnan(returns)
    if np.sum(valid_mask) >= 2:
        ic, _ = spearmanr(scores[valid_mask], returns[valid_mask])
        return ic
    else:
        return np.nan


def recall_at_n(scores, returns, n):
    """
    Compute Recall@N between predicted top-N scores and actual top-N returns.

    Parameters:
    - scores: 1D numpy array of factor scores
    - returns: 1D numpy array of next-day returns
    - n: Integer, number of top assets to consider

    Returns:
    - Recall value (float)
    """
    scores = np.asarray(scores)
    returns = np.asarray(returns)
    topn_pred = set(np.argsort(scores)[-n:])
    topn_true = set(np.argsort(returns)[-n:])
    return len(topn_pred & topn_true) / n


def movement_precision_at_n(scores, returns, n):
    """
    Compute Movement Precision@N: proportion of assets in top-N scores with positive returns.

    Parameters:
    - scores: 1D numpy array of factor scores
    - returns: 1D numpy array of next-day returns
    - n: Integer, number of top assets to consider

    Returns:
    - Precision value (float)
    """
    scores = np.asarray(scores)
    returns = np.asarray(returns)
    topn_idx = np.argsort(scores)[-n:]
    return np.mean(returns[topn_idx] > 0)


# Example usage
if __name__ == "__main__":
    scores = np.array([0.9, 0.7, 0.2, 0.5])
    returns = np.array([0.01, 0.02, -0.01, 0.005])

    print("Rank IC:", compute_rank_ic(scores, returns))
    print("Recall@2:", recall_at_n(scores, returns, n=2))
    print("Movement Precision@2:", movement_precision_at_n(scores, returns, n=2))
