import inout as io


if __name__ == '__main__':
    target = [x[0].replace('\n', '') for x in io.read('train/expected.tsv.xz')]
    categories = {}
    i = 0
    for x in target:
        if x not in categories.values(): 
            categories[i] = x
            i += 1

    files = ['dev-0', 'test-A', 'test-B']

    for file in files:
        predicted = io.read(file)
        predicted = [categories[round(float(x))] for x in predicted]
        io.write(predicted, file + '/out.tsv')