#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function)


try:
    import cPickle as pickle
except ImportError:
    import pickle

from collections import OrderedDict
import gzip
from itertools import product
import os

import numpy as np

from chemreac.util.pyutil import progress
from analytic_diffusion import integrate_rd

default_varied = OrderedDict([
    ('nstencil', [3, 5, 7]),  # categorical
    ('N', list(range(32, 385, 16))),  # continuous
    ('method', ['bdf', 'adams'])  # categorical
])

constant = dict(
    geom='f', nspecies=10,
    D=2e-3, t0=.25, tend=1.5, x0=0.0, xend=1.0, center=.5,
    nt=42, logt=False, logy=False, logx=False,
    random=False, p=0, a=2.7,
    linterpol=False, rinterpol=False, n_jac_diags=0, num_jacobian=False,
    atol=1e-8, rtol=1e-10,
    efield=True, random_seed=42, mobility=0.01,
    plot=False, savefig='None', verbose=False, yscale='linear',
    vline_limit=100, solver='cvode', iter_type='default',
    linear_solver='gmres', ilu_limit=1.0
)


def integrate(**kwargs):
    times_cpu, times_wall = [], []
    nrepeat = 21
    ndrop = 7
    for i in range(nrepeat):
        tout, yout, info, rmsd_over_atol, sys, rmsd = integrate_rd(**kwargs)
        times_cpu.append(info['time_cpu'])
        times_wall.append(info['time_wall'])
    for _ in range(ndrop):
        times_cpu.pop(times_cpu.index(max(times_cpu)))
        times_wall.pop(times_cpu.index(max(times_wall)))
    times_cpu = np.array(times_cpu)
    times_wall = np.array(times_wall)
    info['time_cpu'] = avg_cpu = np.average(times_cpu)
    info['time_wall'] = avg_wall = np.average(times_wall)
    info['dtime_cpu'] = np.sqrt(np.sum((times_cpu - avg_cpu)**2 /
                                       (len(times_cpu) - 1)))
    info['dtime_wall'] = np.sqrt(np.sum((times_wall - avg_wall)**2 /
                                        (len(times_wall) - 1)))
    info['tout'] = tout
    # info['yout'] = yout
    info['rmsd'] = rmsd
    info['rmsd_over_atol'] = np.sqrt(np.sum(rmsd_over_atol**2))
    for k in default_varied:
        info[k] = kwargs[k]
    return info


def main(varied=None, verbose=False):
    if varied is None:
        varied = default_varied
    results = {
        'varied_keys': list(default_varied.keys()),
        'varied_values': list(default_varied.values())
    }
    all_params = list(product(*varied.values()))
    for params in progress(all_params) if verbose else all_params:
        kwargs = constant.copy()
        kwargs.update(dict(zip(varied.keys(), params)))
        results[params] = integrate(**kwargs)
    basename = os.path.splitext(os.path.basename(__file__))[0]
    pickle.dump(results, gzip.open(basename+'.pkl', 'wb'))


if __name__ == '__main__':
    try:
        import argh
        argh.dispatch_command(main)
    except ImportError:
        import sys
        if len(sys.argv) > 1:
            print("Unable to process parameters, argh missing. "
                  "Run 'pip install --user argh' to fix.", file=sys.stderr)
            sys.exit(os.EX_USAGE)  # non-ok exit
        main()

@@ -1099,17 +1099,16 @@
 imes_wal
-l
  = %5B%5D, %5B
@@ -1306,17 +1306,16 @@
 imes_wal
-l
 .append(
@@ -1432,17 +1432,16 @@
 imes_wal
-l
 .pop(tim
@@ -1435,35 +1435,35 @@
 s_wal.pop(times_
-cpu
+wal
 .index(max(times
@@ -1466,17 +1466,16 @@
 imes_wal
-l
 )))%0A    
@@ -1515,25 +1515,24 @@
    times_wal
-l
  = np.array(
@@ -1532,33 +1532,32 @@
 .array(times_wal
-l
 )%0A    info%5B'time
@@ -1628,17 +1628,16 @@
  avg_wal
-l
  = np.av
@@ -1651,17 +1651,16 @@
 imes_wal
-l
 )%0A    in
@@ -1677,32 +1677,41 @@
 cpu'%5D = np.sqrt(
+%0A        
 np.sum((times_cp
@@ -1731,47 +1731,8 @@
 *2 /
-%0A                                      
  (le
@@ -1780,24 +1780,33 @@
 %5D = np.sqrt(
+%0A        
 np.sum((time
@@ -1810,17 +1810,16 @@
 imes_wal
-l
  - avg_w
@@ -1824,55 +1824,14 @@
 _wal
-l
 )**2 /
-%0A                                       
  (le
@@ -1841,17 +1841,16 @@
 imes_wal
-l
 ) - 1)))

fix test failure (fixup)