import numpy as np
from policies.expert_mode_policies import ExpertModeMultiplicativeWeights
from scipy.spatial import distance
from utils import nash_solver, get_losses
from policies.online_oracle_kuhn import OnlineOracleKuhn, KuhnMetagame



def online_learning_kuhn(row_num_pol_init, col_num_pol_init, n_brs, eps, lp_nash_iter=True, refractory=20, whole_time_average=True, ref_func=None):
    game = KuhnMetagame(pop_size1=row_num_pol_init, pop_size2=col_num_pol_init, seed=0)
    matrix_a = -game.metagame  # Numpy array with metagame

    eqs_row, _, v = nash_solver(matrix_a, False)
    row_player_algo = OnlineOracleKuhn(game.pop1, eps, whole_time_average=whole_time_average, ref_func=ref_func)
    col_player_algo = OnlineOracleKuhn(game.pop2, eps, whole_time_average=whole_time_average, ref_func=ref_func)


    n_br = 0
    while n_br<n_brs:
        print(f'ITER {n_br}/{n_brs}')
        row_loss, col_loss, payoff = get_losses(row_player_algo.policy, col_player_algo.policy, matrix_a)

        row_player_algo.update(row_loss, col_player_algo)
        col_player_algo.update(col_loss, row_player_algo)

        if len(row_player_algo.policy)!=matrix_a.shape[0] or len(col_player_algo.policy)!=matrix_a.shape[1]:
            matrix_a = -game.update_metagame(len(row_player_algo.policy), len(col_player_algo.policy)) 

        n_br = len(row_player_algo.exps) + len(col_player_algo.exps)



        if lp_nash_iter:
            eqs_row, _, v = nash_solver(matrix_a, False)

    exploitabilities = [row_player_algo.exps, col_player_algo.exps]
    return exploitabilities