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)
        self.survivals = {}
        with open('survivals.json', 'r') as f:
            self.survivals = json.load(f)
        self.neighbors = {}
        self.read_neighbor()

    def read_neighbor(self):
        for i, m in enumerate(self.matches.data):
            n = {}
            for r in m.query('neighbor', 'create').raw_data:
                if r['a'] in n:
                    n[r['a']].append(r['b'])
                else:
                    n[r['a']] = [r['b']]

                if r['b'] in n:
                    n[r['b']].append(r['a'])
                else:
                    n[r['b']] = [r['a']]
            self.neighbors[self.matches.names[i]] = n

    def new_partner_after(self):
        re_after = [0] * 2 #[c,d]
        su_after = [0] * 2
        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):
                    previous_round_partner = []
                    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['a'] == pid:
                            previous_round_partner.append(r['b'])
                            if r['act_a'] == 'C':
                                c += 1
                        else:
                            previous_round_partner.append(r['a'])
                            if r['act_b'] == 'C':
                                c += 1
                        t += 1
                    is_coop = 0 if c * 2 >= t else 1

                    new_partner_request = 0
                    for r in m.query('action', 'request').where(lambda x: x['rno'] == i and (x['from'] == pid or x['to'] == pid)).raw_data:
                        if r['from'] == pid:
                            if r['to'] not in previous_round_partner:
                                new_partner_request += 1
                        else:
                            if r['from'] not in previous_round_partner:
                                new_partner_request += 1
                    if new_partner_request == 0:
                        continue
                    re_after[is_coop] += new_partner_request                        
                    new_partner_succ = 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:
                            if r['b'] not in previous_round_partner:
                                new_partner_succ += 1
                        else:
                            if r['a'] not in previous_round_partner:
                                new_partner_succ += 1
                    su_after[is_coop] += new_partner_succ

        t,v = 0,0
        for i, _ in enumerate(su_after):
            print(su_after[i], re_after[i], su_after[i] / re_after[i])
            t += re_after[i]
            v += su_after[i]
            su_after[i] /= re_after[i]
        print(v, t, v/t)
        return su_after


if __name__ == '__main__':
    sc = Solution('CLASSIC')
    npc_after = sc.new_partner_after()
    ss = Solution('SURVIVE')
    nps_after = ss.new_partner_after()
    index = np.arange(2)
    
    blue = '#0984e3'
    red = '#d63031'

    fig = plt.figure(figsize=(3, 3))
    ax = fig.gca()
    bar_width = 0.35
    opacity = 1
    error_config = {'ecolor': '0.3', 'capsize': 4}
    rects1 = ax.bar(index, nps_after, bar_width,
                    color=red, label='Dissipative')
    rects2 = ax.bar(index + bar_width, npc_after, bar_width,
                    alpha=opacity, color='#00b894',
                    label='Classic')

    ax.set_ylabel('Rate')
    # ax.set_title('Rate of Successfully Making New Partner')
    ax.set_xlabel('Previous Move')
    ax.set_xticks(index + bar_width / 2)
    ax.set_xticklabels(['C', 'D'])
    ax.set_ylim(0,1)
    ax.legend()

    fig.tight_layout()
    # plt.show()
    # plt.savefig("graph/new_partner_bar.eps")