import json
from matplotlib import pyplot as plt
from matplotlib import patches as pch
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):
        survivals = {}
        with open('survivals.json', 'r') as f:
            survivals = json.load(f)
        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):
                    if pid not in survivals[m.name][str(i)]: break
                    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 if self.mode == 'SURVIVE' else 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)
        print(len(fl[0]), len(fl[1]))
        print(sum(fl[0]) / len(fl[0]), sum(fl[1]) / len(fl[1]))
        return fl


if __name__ == '__main__':
    sc = Solution('CLASSIC')
    fc = sc.food_loss()
    ss = Solution('SURVIVE')
    fs = ss.food_loss()
    blue = '#0984e3'
    red = '#d63031'
    green = '#00b894'
    c = [blue, red]
    labels = ['Dissipative', 'Classic']
    fig = plt.figure(figsize=(8, 6))
    ax = fig.gca()
    index = np.arange(2)
    widths = 0.3

    vp1 = ax.violinplot(fs, positions=index-widths/2, widths=widths, showmeans=True, showmedians=True)
    vp2 = ax.violinplot(fc, positions=index+widths/2, widths=widths, showmeans=True, showmedians=True)
    hdl1 = pch.Patch(facecolor=vp1['bodies'][0].get_facecolor()[0])
    hdl2 = pch.Patch(facecolor=vp2['bodies'][0].get_facecolor()[0])
    vp1['cmeans'].set_edgecolor(green)
    vp1['cmeans'].set_linestyle(':')
    vp2['cmeans'].set_edgecolor(green)
    vp2['cmeans'].set_linestyle(':')
    ax.set_ylabel('Food Loss')
    ax.yaxis.grid(True, linestyle='--')
    # ax.set_title('Scores by group and gender')
    ax.set_xticks(index)
    ax.set_xticklabels(['C', 'D'])
    ax.set_ylim(-6, 4)
    ax.set_xlabel("Previous Move")
    ax.legend([hdl1, hdl2], labels)

    fig.tight_layout()
    # plt.show()
    plt.savefig("graph/FoodLoss.svg")