import json
from matplotlib import pyplot as plt
from island.match import Match
from island.matches import Matches
import numpy as np
from scipy.stats import pearsonr


class Solution:
    def __init__(self, mode):
        self.mode = mode
        self.matches = Matches.from_profile_expr(lambda r: mode in r)

    def food_loss(self):
        fl = [[],[]]
        payoff = {
            'C': {
                'C': 2.6,
                'D': 0
            },
            'D': {
                'C': 4.5,
                'D': 0.5
            }
        }
        for m in self.matches.data:
            info = m.query('game', 'created').select('info').first()['info']
            game_end_at = int(info['game_end_at'])

            for p in m.query('player', 'join').raw_data:
                pid = p['pid']
                for i in range(2, game_end_at):
                    c, t = 0, 0
                    for r in m.query('action', 'done').where(lambda x: x['rno'] == i - 1 and (x['a'] == pid or x['b'] == pid)).raw_data:
                        if r['act_a' if r['a'] == pid else 'act_b'] == 'C':
                            c += 1
                        t += 1
                    is_coop = 0 if c * 2 >= t else 1
                    f = -3.0
                    for r in m.query('action', 'done').where(lambda x: x['rno'] == i and (x['a'] == pid or x['b'] == pid)).raw_data:
                        if r['a'] == pid:
                            f += payoff[r['act_a']][r['act_b']] * r['tr'] / 1440.0
                        else:
                            f += payoff[r['act_b']][r['act_a']] * r['tr'] / 1440.0
                    fl[is_coop].append(f)
        return fl


if __name__ == '__main__':
    mode = 'SURVIVE'
    s = Solution(mode)
    f = s.food_loss()
    blue = '#0984e3'
    red = '#d63031'
    c = [blue, red]
    labels = ['C', 'D']
    fig = plt.figure(figsize=(3, 3))
    ax = fig.gca()
    error_config = {'ecolor': '0.3', 'capsize': 4}
    bplot = ax.boxplot(f,
                        vert=True,  # vertical box alignment
                        patch_artist=True,  # fill with color
                        notch=True,  # notch shape
                        labels=labels)  # will be used to label x-ticks

    for patch, color in zip(bplot['boxes'], c):
        patch.set_facecolor(color)
    ax.set_ylabel('Frequency')
    # ax.set_title('Scores by group and gender')

    ax.legend()

    fig.tight_layout()
    plt.show()
    # plt.savefig("graph/new_partner_bar_%s.eps" % mode)