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

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:
            id, _, _, rok, miesiac = line.rstrip().split('\t')
            rok_2010 = train[(train['id'] == int(id) ) & (train['rok'] == 2010) & (train['miesiąc'] == int(miesiac))  ]['opad'].item()
            f_out.write(str(rok_2010) + '\n')