from island.match import Match
from island.matches import Matches
import numpy as np
import csv
import json

NAME = 'wos-data-2022-sd'


def calc_and_save(mode, network_type=None):
    matches = Matches(NAME, network_type=network_type)
    max_round = 0  # = game_end_at
    for m in matches.data:
        max_round = max(max_round, int(m.query('game', 'created').first()['info']['game_end_at']))
    # matches = Matches.from_profile(mode)
    times = np.zeros(max_round)
    cs = np.zeros(max_round)
    crs = []
    for m in matches.data:
        maxr = int(m.query('game', 'created').first()['info']['game_end_at'])
        c = [0.0] * maxr
        for r in range(1, maxr + 1):
            rows = m.query('action', 'done').where(lambda x: x['rno'] == r).raw_data
            actors = dict()  # pid -> [C,D]
            coop_rate = 0.0
            for row in rows:
                for side in ['a', 'b']:
                    if row[side] not in actors:
                        actors[row[side]] = [0, 0]
                    if row[f'act_{side}'] == 'C':
                        actors[row[side]][0] += 1
                    else:
                        actors[row[side]][1] += 1
            if rows:
                for v in actors.values():
                    coop_rate += float(v[0]) / float(v[0] + v[1])

                times[r - 1] += len(actors)
                cs[r - 1] += coop_rate
                c[r - 1] = coop_rate / len(actors)
        with open(f"outputs/CR_{m.name}.csv", 'w') as f:
            csv.writer(f).writerow(c)
        crs.append(c)

    for i in range(max_round):
        if times[i] == 0:
            times[i] = 1
    cs /= times
    with open("outputs/CR_%s_%s.csv" % (mode, network_type), 'w') as f:
        csv.writer(f).writerow(cs)
    with open("outputs/CRAW_%s_%s.json" % (mode, network_type), 'w') as f:
        json.dump(crs, f)


if __name__ == '__main__':
    # calcAndSave('CLASSIC')
    # calcAndSave('SURVIVE')
    calc_and_save('NEW', 'BA')
    calc_and_save('NEW', 'WS')
