#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os

import re

import numpy as np
from tqdm import tqdm

REGEX_FILE_NAME = re.compile(r'^out.tsv[0-9]+$')


def avarage(dir_name: str):
    print(f'Processing {dir_name}')
    file_names = [f for f in os.listdir(dir_name)
                  if REGEX_FILE_NAME.match(f)]

    if not file_names:
        print('ERROR! Not found files!')
        return

    print(f'Reading from files: {file_names}')
    files = [open(dir_name + '/' + f) for f in file_names]
    f_out = open(dir_name + '/out-model=best_sum.tsv', 'w')
    progress = tqdm(desc=dir_name)

    while True:
        try:
            hyps = [float(next(x).rstrip()) for x in files]
        except StopIteration:
            break
        avg_v = np.mean(hyps)
        f_out.write(str(avg_v) + '\n')
        progress.update(1)

    f_out.close()


if __name__ == '__main__':
    avarage('data/dev-0')
    avarage('data/dev-1')
    avarage('data/test-A')