import unittest
import numpy as np
import os
import matplotlib.pyplot as plt

import acpc_python_client as acpc

from tools.constants import Action
from weak_agents.action_tilted_agent import create_agent_strategy_from_trained_strategy, TiltType
from tools.io_util import read_strategy_from_file
from evaluation.exploitability import Exploitability


FIGURES_FOLDER = 'verification/action_tilted_agents'

KUHN_EQUILIBRIUM_STRATEGY_PATH = 'strategies/kuhn.limit.2p-equilibrium.strategy'
LEDUC_EQUILIBRIUM_STRATEGY_PATH = 'strategies/leduc.limit.2p-equilibrium.strategy'

TILT_TYPES = [
    ('fold-add', Action.FOLD, TiltType.ADD),
    ('call-add', Action.CALL, TiltType.ADD),
    ('raise-add', Action.RAISE, TiltType.ADD),
    ('fold-multiply', Action.FOLD, TiltType.MULTIPLY),
    ('call-multiply', Action.CALL, TiltType.MULTIPLY),
    ('raise-multiply', Action.RAISE, TiltType.MULTIPLY),
]


class ActionTiltedAgentsExploitabilityTest(unittest.TestCase):
    def test_plot_kuhn_agent_exploitabilities(self):
        self.create_agents_and_plot_exploitabilities({
            'title': 'Kuhn poker action tilted agents exploitability',
            'figure_filename': 'kuhn_action_tilted_agents',
            'base_strategy_path': KUHN_EQUILIBRIUM_STRATEGY_PATH,
            'game_file_path': 'games/kuhn.limit.2p.game',
            'tilt_probabilities': np.arange(-1, 1, 0.01),
        })

    def test_plot_leduc_agent_exploitabilities(self):
        self.create_agents_and_plot_exploitabilities({
            'title': 'Leduc Hold\'em action tilted agents exploitability',
            'figure_filename': 'leduc_action_tilted_agents',
            'base_strategy_path': LEDUC_EQUILIBRIUM_STRATEGY_PATH,
            'game_file_path': 'games/leduc.limit.2p.game',
            'tilt_probabilities': np.arange(-1, 1, 0.1),
        })

    def create_agents_and_plot_exploitabilities(self, test_spec):
        base_strategy = read_strategy_from_file(
            test_spec['game_file_path'],
            test_spec['base_strategy_path'])

        game = acpc.read_game_file(test_spec['game_file_path'])
        exploitability = Exploitability(game)

        equilibrium_exploitability = exploitability.evaluate(base_strategy)

        tilt_probabilities = test_spec['tilt_probabilities']
        exploitability_values = np.zeros([len(TILT_TYPES), len(tilt_probabilities)])

        for i, tilt_type in enumerate(TILT_TYPES):
            for j, tilt_probability in enumerate(tilt_probabilities):
                tilted_agent = create_agent_strategy_from_trained_strategy(
                    test_spec['game_file_path'],
                    base_strategy,
                    tilt_type[1],
                    tilt_type[2],
                    tilt_probability)
                exploitability_values[i, j] = exploitability.evaluate(tilted_agent)

            plt.figure(dpi=160)
            for j in range(i + 1):
                plt.plot(
                    tilt_probabilities,
                    exploitability_values[j],
                    label=TILT_TYPES[j][0],
                    linewidth=0.8)

            plt.plot(
                tilt_probabilities,
                [equilibrium_exploitability] * len(tilt_probabilities),
                'r--',
                label='Equilibrium',
                linewidth=1.5)

            plt.title(test_spec['title'])
            plt.xlabel('Tilt probability')
            plt.ylabel('Agent exploitability [mbb/g]')
            plt.grid()
            plt.legend()

            figure_output_path = '%s/%s.png' % (FIGURES_FOLDER, test_spec['figure_filename'])

            figures_directory = os.path.dirname(figure_output_path)
            if not os.path.exists(figures_directory):
                os.makedirs(figures_directory)

            plt.savefig(figure_output_path)

Add test that explores exploitability of tilted agents