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

from __future__ import division, print_function
from collections import defaultdict

import os
import sys
import numpy as np
import matplotlib.pyplot as plt


def parse_yutiming(filename, n_nodes=20):

    hours = []
    times = []

    details = defaultdict(list)

    hour = None


    with open(filename) as yu:
        for i, line in enumerate(yu):

            if not line.strip():
                continue

            if line.startswith('Time for the setup of the model:'):
                pass
            elif line.startswith('Time for the cleanup of the model:'):
                pass

            elif line.startswith('Time for hour'):
                hour = int(line[13:21])
                time = float(line[31:])

                hours.append(hour)
                times.append(time)

            else:
                if hour is None:
                    continue

                label = line[:23].strip()
                for time in line[23:].split():
                    try:
                        details[label].append( float(time) )
                    except:
                        print(i)
                        print(line)
                        raise


    for k,v in details.items():
        details[k] = np.array(v).reshape(-1,n_nodes)


    return np.array(hours), np.array(times), details



def separate(d):
    io_groups = [
        'Input',
        'Output',
        'asynIO wait',
    ]
    io_times = np.sum([d[group].max(1) for group in io_groups], axis=0)

    phy_groups = [
        'Phy. Computations',
        'Communications  Phy',
        'Barrier Waiting Phy',
    ]
    phy_times = np.sum([d[group].max(1) for group in phy_groups], axis=0)

    dyn_groups = [
        'Dyn. Computations',
        'Fast Waves Comm.',
        'Fast Waves Barrier',
        'Communications  Dyn',
        'Barrier Waiting Dyn',
        'Global Comm     Dyn',
        'Barrier GlobCom Dyn',
        'wrapper copy in',
        'wrapper dycore step',
        'wrapper copy out',
        'wrapper copy swap',
    ]
    dyn_times = np.sum([d[group].max(1) for group in dyn_groups], axis=0)

    other_groups = [
        'Add. Computations',
    ]
    other_times = np.sum([d[group].max(1) for group in other_groups], axis=0)

    return dyn_times, phy_times, io_times, other_times







def load_all():
    path = '/project/d73/output/production_run_02/'

    hours = []
    times = []

    dyns = []
    phys = []
    ios = []
    oths = []

    for hour in (24 * np.arange(0,120,3)):
        print(hour)
        folder = '2015010100_%d_%d' % (hour, hour+72)
        h,t,d = parse_yutiming(os.path.join(path, folder, 'cosmo_run', 'YUTIMING'))

        dyn, phy, io, oth = separate(d)


        if hour == 0:
            h = np.insert(h, 0, 0)
            t = np.insert(t, 0, 0)
            dyn = np.insert(dyn, 0, 0)
            phy = np.insert(phy, 0, 0)
            io = np.insert(io, 0, 0)
            oth = np.insert(oth, 0, 0)

        hours.append(h)
        times.append(t)
        dyns.append((dyn))
        phys.append((phy))
        ios.append((io))
        oths.append((oth))


    return [np.array(v) for v in [hours, times, dyns, phys, ios, oths]]





def main():

    hours, times, dyns, phys, ios, oths = load_all()
    totals = dyns + phys + ios + oths
    oths = oths + phys


    sec2nh = 24 / 60**2 * 21

    fig = plt.figure(figsize=[6.31, 4.0])
    ax = fig.add_subplot(111)


    ax.plot(hours[:,0] / 24.0, sec2nh * np.median(totals, 1), 'k.-', label='total')
    ax.plot(hours[:,0] / 24.0, sec2nh * np.median(ios, 1), '.-', label='I/O')
    ax.plot(hours[:,0] / 24.0, sec2nh * np.median(dyns, 1), '.-', label='dynamics')
    ax.plot(hours[:,0] / 24.0, sec2nh * np.median(oths, 1), '.-', label='others')

    #ax.plot(hours[:,0] / 24.0, sec2nh * totals.sum(1) / 72.0, 'k.-', label='total')
    #ax.plot(hours[:,0] / 24.0, sec2nh * ios.sum(1) / 72.0, '.-', label='I/O')
    #ax.plot(hours[:,0] / 24.0, sec2nh * dyns.sum(1) / 72.0, '.-', label='dynamics')
    #ax.plot(hours[:,0] / 24.0, sec2nh * oths.sum(1) / 72.0, '.-', label='others')

    ax.set_xlabel('days')
    ax.set_ylabel('node hours per 24 h simulation')
    ax.set_ylim(0,50)
    ax.set_xlim(-2,117)

    ax.legend(loc=1, ncol=2)

    ax.grid(True)

    plt.tight_layout()



if __name__ == '__main__':
    main()
    plt.savefig(sys.argv[1])

