#!/usr/bin/env python3

import argparse
import re

import matplotlib.pyplot as plt
import numpy as np


def extract_data(filename):

    regex_step = re.compile(
        r"STEP [0-9]* ends.*\n.* Avg\. per step = ([0-9]*[.])?[0-9]+ s", re.MULTILINE)

    string_data = []

    print("Processing " + filename + " ...", end='')
    with open(filename) as f:
        text = f.read()

        string_data = [s.group(0) for s in regex_step.finditer(text)]

    regex_real = re.compile(
        r" -?[\d.]+(?:e-?\d+)?", re.MULTILINE)

    time_data = np.zeros([len(string_data), 6])
    for i, ss in enumerate(string_data):
        numbers = regex_real.findall(ss)
        time_data[i,:] = np.array(numbers)

    print("...done!")
    return time_data


def plot_timestep_duration(time_data, name):
    fig_name = name + "_ts_duration.png"
    print("Generating " + fig_name + "...", end='')

    plt.rcParams.update({'font.size': 20})
    plt.rcParams['axes.linewidth'] = 3

    f, ax = plt.subplots(figsize=(12,6))

    ax.set_ylabel("timestep duration [s]")
    ax.set_xlabel("step [#]")

    ax.semilogy(time_data[:,0], time_data[:,4])

    ax.spines['bottom'].set_color('gray')
    ax.spines['top'].set_visible(False)
    ax.spines['left'].set_color('gray')
    ax.spines['right'].set_visible(False)

    plt.tight_layout()

    plt.savefig(fig_name, transparent=False, dpi=300)
    print("...done!")


def plot_cumulative_duration(time_data, name):
    fig_name = name + "_cumulative_duration.png"
    print("Generating " + fig_name + "...", end='')

    plt.rcParams.update({'font.size': 20})
    plt.rcParams['axes.linewidth'] = 3

    f, ax = plt.subplots(figsize=(12,6))

    ax.set_ylabel("cumulative duration [s]")
    ax.set_xlabel("step [#]")

    ax.plot(time_data[:,0], np.cumsum(time_data[:,4]))

    ax.spines['bottom'].set_color('gray')
    ax.spines['top'].set_visible(False)
    ax.spines['left'].set_color('gray')
    ax.spines['right'].set_visible(False)

    plt.tight_layout()

    plt.savefig(name + "_cumulative_duration.png", transparent=False, dpi=300)
    print("...done!")


def do_plot_timestep_duration():
    parser = argparse.ArgumentParser(description='Generates plots of timestep duration from WarpX standard output logs')
    parser.add_argument('file_name', metavar='file_name', type=str, nargs=1,
        help='the name of the WarpX output log to process')

    args = parser.parse_args()
    log_file_name = args.file_name[0]

    time_data = extract_data(log_file_name)

    plot_timestep_duration(time_data, log_file_name)
    plot_cumulative_duration(time_data, log_file_name)

if __name__ == "__main__":
    do_plot_timestep_duration()