import fnmatch
import heapq
import numpy as np
import os


def find_file_name(path, name_pattern):
    names = os.listdir(path)
    for _ in names:
        if fnmatch.fnmatch(_, name_pattern):
            return _
    return None


def backtest_oneday(y_pred, y_true, nmax=2, principal=100000):
    '''
    :return:
        [top 5 stocks with the highest yields]
        accuracy
    '''
    ids_preds = list(map(list(y_pred).index, heapq.nlargest(nmax, y_pred)))
    y_pred_max = y_pred[ids_preds]
    y_true_max = y_true[ids_preds]

    acc_arr = [1 if (y_pred_max[i] > 1 and y_true_max[i] > 1) or (y_pred_max[i] < 1 and y_true_max[i] < 1) else 0
               for i in range(nmax)]
    acc = sum(acc_arr) / nmax

    principals = [principal / nmax for _ in range(nmax)]
    money_add = sum(np.array(principals) * np.array(y_true_max)) - principal

    return ids_preds, acc, money_add
