import pandas as pd
import numpy as np
import trueskill
from tqdm import tqdm

df_train = pd.read_csv('./data/train_data.csv')

def calc_weapon_rate(mode):
    
    env = trueskill.TrueSkill()

    weapon_dict = {}
    
    if mode != 'all':
        _df = df_train[df_train['mode'] == mode]
    else:
        _df = df_train.copy()

    # 両チーム4人そろっているものを対象
    _df = _df[~(_df[['A1-weapon', 'A2-weapon', 'A3-weapon', 'A4-weapon', 'B1-weapon', 'B2-weapon', 'B3-weapon', 'B4-weapon']].isnull().sum(axis=1) > 0)]
    
    # 全員のランクが同一のバトルを対象
    if  mode not in ['nawabari' , 'all']:
        _df = _df[(_df['A1-rank'] == _df['A2-rank']) & (_df['A1-rank'] == _df['A3-rank']) & (_df['A1-rank'] == _df['A4-rank']) & (_df['A1-rank'] == _df['B1-rank']) & (_df['A1-rank'] == _df['B2-rank']) & (_df['A1-rank'] == _df['B3-rank']) & (_df['A1-rank'] == _df['B4-rank'])]

    for idx, row in tqdm(_df.sort_values('period').iterrows(), total=len(_df)):
        team_a = {}
        team_b = {}

        for weapon_column in ['A1-weapon', 'A2-weapon', 'A3-weapon', 'A4-weapon']:
            weapon = row[weapon_column]
            team_a[weapon] = weapon_dict.get(weapon, env.create_rating())
        
        for weapon_column in ['B1-weapon', 'B2-weapon', 'B3-weapon', 'B4-weapon']:
            weapon = row[weapon_column]
            team_b[weapon] = weapon_dict.get(weapon, env.create_rating())

        team_a, team_b, = env.rate((team_a, team_b,), ranks=(abs(row['y'] - 1), row['y'],))

        weapon_dict.update(team_a)
        weapon_dict.update(team_b)

    rate_dict = {k:float(v) for k, v in weapon_dict.items()}
    
    return rate_dict

# allは全modeが対象
for mode in ['all', 'nawabari', 'area', 'asari', 'hoko', 'yagura']:
    mode_rate = calc_weapon_rate(mode)
    
    display(f'{mode} レーティング結果')
    
    df_rate = pd.DataFrame(sorted(mode_rate.items(), key=lambda x:x[1]), columns=['weapon', 'rating'])

    display('上位10件')
    display(df_rate.sort_values('rating', ascending=False).head(10))
    
    display('下位10件')
    display(df_rate.head(10))