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")