#!/usr/bin/env python3

'''Process and plot results.

Usage: {0} <model name>'''

from   glob import glob
import os
import re
import sys

import matplotlib.pyplot as plt
from   matplotlib import rcParams

def process_file(filename, filename_reg, param_reg, result_reg):
    '''
    Process a log file.
    '''
    # Collect results from the file.
    title = filename_reg.fullmatch(filename)
    assert title is not None
    title = title.expand(r'\1')
    print(title)
    results = []
    with open(filename) as file:
        param, acc, loss = None, None, None
        for line in file:
            # Match parameters pattern. If we have already found param, acc and loss, append them to the results as
            # a tuple. If we already found param but not acc or loss, it means there were two params in a row, so we
            # warn the user.
            match = param_reg.match(line)
            if match is not None:
                param = match.expand(r'\1')
                # pylint: disable=W0123
                param = list(eval(param).values()) # Don't warn about use of eval.
                for i in range(len(param)):
                    if isinstance(param[i], dict):
                        param[i] = ', '.join(map(str, param[i].values()))
                    else:
                        param[i] = str(param[i])
                param = ', '.join(reversed(param)) # Reverse order of opt_params/optimizer.
                continue
            # Match accuracy and loss pattern.
            match = result_reg.match(line)
            if match is not None:
                acc, loss = match.expand(r'\1 \2').split()
                if param is not None:
                    results.append((param,acc,loss))
                param, acc, loss = None, None, None
        if None not in (param,acc,loss):
            results.append((param,acc,loss))
    # Parse the results.
    params, accs, losses = zip(*results)
    assert len(params) == len(accs)
    assert len(params) == len(losses)
    accs   = list(map(float, accs))
    losses = list(map(float, losses))
    return title, params, accs, losses

def process_results(title, params, accs, losses):
    '''
    Process log results to .list file and create plots.
    '''
    rcParams.update({'figure.autolayout': True})
    RATIO  = 16/9
    WIDTH  = 35
    HEIGHT = WIDTH/RATIO
    # Save results.
    with open(os.path.join(LOG_DIR, '{}.list'.format(title)), 'w') as file:
        for param, acc, loss in zip(params, accs, losses):
            file.write('{}\t{}\t{}\n'.format(param, acc, loss))
    # Plot and save accuracies.
    xs = range(len(accs))
    plt.clf()
    fig = plt.gcf()
    plt.tight_layout()
    #fig.set_size_inches(WIDTH, HEIGHT)
    plt.title('Comparison of accuracy for different {}{}'.format(title, 's' if title[-1] != 's' else ''))
    plt.xlabel('Parameters')
    plt.ylabel('Accuracy (%)')
    plt.xticks(xs, params, rotation='vertical')
    plt.bar(xs, accs)
    plt.savefig(os.path.join(LOG_DIR, '{}.acc.png'.format(title)))
    # Plot and save losses.
    plt.clf()
    fig = plt.gcf()
    plt.tight_layout()
    #fig.set_size_inches(WIDTH, HEIGHT)
    plt.title('Comparison of loss for different {}{}'.format(title, 's' if title[-1] != 's' else ''))
    plt.xlabel('Parameters')
    plt.ylabel('Loss')
    plt.xticks(xs, params, rotation='vertical')
    plt.bar(xs, losses)
    plt.savefig(os.path.join(LOG_DIR, '{}.loss.png'.format(title)))

if __name__ == '__main__':
    # Directory containing log files.
    LOG_DIR = os.path.join(os.getcwd(), 'data', '{0}', 'logs')
    # Log filename glob.
    LOG_GLOB = '*.log'
    # Regex patterns.
    PATTERNS = {
        'log':    r'.*[0-9]{2}_(.*)\.log',
        'train':  r'training.log',
        'param':  r'.*Grid [0-9]+\/[0-9]+\: (\{.*\})',
        'result': r'.*acc=([0-9]+\.[0-9]+)%, loss=([0-9]+\.[0-9]+)'
    }
    # Setup.
    if len(sys.argv) != 2:
        print(__doc__.format(sys.argv[0]), file=sys.stderr)
        exit(1)
    model_name = sys.argv[1]
    LOG_DIR    = LOG_DIR.format(model_name)
    log_reg    = re.compile(PATTERNS['log'])
    param_reg  = re.compile(PATTERNS['param'])
    result_reg = re.compile(PATTERNS['result'])
    # Handle tuner logs.
    for log_file in glob(os.path.join(LOG_DIR, LOG_GLOB)):
        title, params, accs, losses = process_file(log_file, log_reg, param_reg, result_reg)
        process_results(title, params, accs, losses)
