from typing import Tuple, List
import pandas as pd
import matplotlib.pyplot as plt

columns = ['map name', 'TvZ', 'ZvP', 'PvT']
data = [
    ('Alcyone', (149, 139), (157, 120), (162, 197)),
    ('Amphion', (12, 9), (10, 6), (7, 13)),
    ('Crimson Court', (20, 8), (4, 11), (22, 25)),
    ('Dynasty', (15, 3), (9, 14), (21, 17)),
    ('Ghost River', (11, 10), (12, 17), (26, 22)),
    ('Goldenaura', (88, 95), (80, 107), (135, 143)),
    ('Oceanborn', (126, 103), (205, 184), (251, 267)),
    ('Post Youth', (9, 11), (11, 13), (8, 21)),
    ('Site Delta', (113, 137), (142, 139), (138, 144)),
    ('Hard Lead', (99, 101), (159, 159), (112, 173)),
    ('Hecate', (68, 55), (102, 114), (170, 178)),
    ('Radhuset', (33, 37), (32, 40), (62, 51)),
    ('Solaris', (63, 92), (85, 84), (104, 93)),
    ('Equilibrium', (50, 60), (68, 40), (70, 53))
]

s1_maps = ('Alcyone', 'Equilibrium', 'Goldenaura', 'Hard Lead', 'Hecate', 'Oceanborn', 'Radhuset', 'Site Delta', 'Solaris')
s2_maps = ('Alcyone', 'Amphion', 'Crimson Court', 'Dynasty', 'Ghost River', 'Goldenaura', 'Oceanborn', 'Post Youth', 'Site Delta')
new_maps = ('Amphion', 'Crimson Court', 'Dynasty', 'Ghost River', 'Post Youth')

def print_winrates(sample: Tuple) -> None:
    name, TvZ, ZvP, PvT = sample
    TvZ_w = round(TvZ[0] / sum(TvZ) * 100, 1)
    ZvP_w = round(ZvP[0] / sum(ZvP) * 100, 1)
    PvT_w = round(PvT[0] / sum(PvT) * 100, 1)
    print(f'Winrate for {name}: TvZ: {TvZ_w} ZvP: {ZvP_w}, PvT: {PvT_w}')

def merge_samples(samples: List[Tuple], name: str) -> Tuple:
    TvZ, ZvT = 0, 0
    ZvP, PvZ = 0, 0
    TvP, PvT = 0, 0
    for _, (s_TvZ, s_ZvT), (s_ZvP, s_PvZ), (s_PvT, s_TvP) in samples:
        TvZ += s_TvZ
        ZvT += s_ZvT
        ZvP += s_ZvP
        PvZ += s_PvZ
        PvT += s_PvT
        TvP += s_TvP
    return (name, (TvZ, ZvT), (ZvP, PvZ), (PvT, TvP))

def filter_samples(samples: List[Tuple], targets: Tuple[str]) -> List[Tuple]:
    return list(filter(lambda sample: sample[0] in targets, samples))

def matchup_winrates(samples: List[Tuple], matchup: str) -> List[Tuple]:
    matchup, side = get_matchup_indices(matchup)
    res = []
    for sample in samples:
        name = sample[0]
        sample_winrate = round(sample[matchup][side] / sum(sample[matchup]) * 100, 1)
        res.append((name, sample_winrate))
    return sorted(res, key=lambda sample: sample[1])

def race_winrates(samples: List[Tuple], race: str):
    mu1, side1, mu2, side2 = get_race_indices(race)
    res = []
    for sample in samples:
        name = sample[0]
        mu1_wins = sample[mu1][side1]
        mu1_total = sum(sample[mu1])
        mu2_wins = sample[mu2][side2]
        mu2_total = sum(sample[mu2])
        sample_winrate = round((mu1_wins + mu2_wins) / (mu1_total + mu2_total) * 100, 1)
        res.append((name, sample_winrate))
    # return res
    return sorted(res, key=lambda w: w[1])


def get_matchup_indices(matchup: str) -> Tuple[int, int]:
    if matchup == 'TvZ' or matchup == 'ZvT':
        return 1, int(matchup == 'ZvT')
    elif matchup == 'ZvP' or matchup == 'PvZ':
        return 2, int(matchup == 'PvZ')
    else:
        return 3, int(matchup == 'TvP')

def get_race_indices(race: str) -> Tuple[int, int, int, int]:
    if race == 'Z':
        return *get_matchup_indices('ZvT'), *get_matchup_indices('ZvP')
    elif race == 'T':
        return *get_matchup_indices('TvZ'), *get_matchup_indices('TvP')
    else:
        return *get_matchup_indices('PvT'), *get_matchup_indices('PvZ')

def plot_matchup_winrates(winrates: List[Tuple[str, float]], matchup: str):
    fig, ax = plt.subplots()
    samples, rates = zip(*winrates)
    p = ax.bar(samples, rates)
    ax.bar_label(p, label_type='center')
    ax.set_title(f'{matchup} winrates')
    ax.set_ylabel('winrate')

    plt.show()

# Matchup winrates per map pools
all_maps_winrates = merge_samples(data, 'all_maps')
s1_winrates = merge_samples(filter_samples(data, s1_maps), 's1_maps')
s2_winrates = merge_samples(filter_samples(data, s2_maps), 's2_maps')
new_maps_winrates = merge_samples(filter_samples(data, new_maps), 'new_maps')
aggregated_data = [all_maps_winrates, s1_winrates, s2_winrates, new_maps_winrates]
for sample in aggregated_data:
    print_winrates(sample)

# Race winrates
## Zerg
### Map pools
zerg_map_pools = race_winrates(aggregated_data, 'Z')
plot_matchup_winrates(zerg_map_pools, 'Map pools Zerg')

## Protoss
protoss_map_pool = race_winrates(aggregated_data, 'P')
plot_matchup_winrates(protoss_map_pool, 'Map pools Protoss')

## Terran
terran_map_pool = race_winrates(aggregated_data, 'T')
plot_matchup_winrates(terran_map_pool, 'Map pools Terran')

### S2 results
## Zerg
zerg_s2 = race_winrates(filter_samples(data, s2_maps), 'Z')
plot_matchup_winrates(zerg_s2, 'Zerg S2')

## Protoss
protoss_s2 = race_winrates(filter_samples(data, s2_maps), 'P')
plot_matchup_winrates(protoss_s2, 'Protoss S2')

## Terran
terran_s2 = race_winrates(filter_samples(data, s2_maps), 'T')
plot_matchup_winrates(terran_s2, 'Terran S2')

# Matchup winrates
## Zerg
ZvT_s2 = matchup_winrates(filter_samples(data, s2_maps), 'ZvT')
TvP_s2 = matchup_winrates(filter_samples(data, s2_maps), 'TvP')
ZvP_s2 = matchup_winrates(filter_samples(data, s2_maps), 'ZvP')
plot_matchup_winrates(ZvT_s2, 'ZvT S2')
plot_matchup_winrates(TvP_s2, 'TvP S2')
plot_matchup_winrates(ZvP_s2, 'ZvP S2')