import pandas as pd

train_in = pd.read_csv('../train/in.tsv', header = None, names = ['id', 'nazwa','typ_zbioru', 'rok', 'miesiąc'], sep='\t')
train_exp = pd.read_csv('../train/expected.tsv', header = None,  names = ['opad'], sep='\t')

train  = train_in.copy()
train['opad'] = train_exp


by_miesiac = train.groupby(['rok', 'miesiąc']).mean()


for dataset in ('dev-0', 'test-A'):
    with open(f'../{dataset}/out.tsv','w') as f_out, open(f'../{dataset}/in.tsv', 'r') as f_in:
        for line in f_in:
            _, _, _, rok, miesiac = line.rstrip().split('\t')
            srednia = by_miesiac.loc[int(rok), int(miesiac)]['opad']
            f_out.write(str(srednia) + '\n')