#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Analytic error scaling vs. number of bins
-----------------------------------------

:download:`examples/analytic_N_scaling.py` runs a series of integrations
with vaired parameters, :download:`examples/plot_analytic_N_scaling.py`
plots the error in the solutions vs. number of bins for each set of other
parameters. For example, one may expect different error scaling for the
different number of stencil points used.
(:math:`N^{-2}`, :math:`N^{-4}` and :math:`N^{-6}` for 3, 5 and 7 stencil
points respectively)

::

 $ python analytic_N_scaling.py --help

.. exec::
   echo "::\\n\\n"
   python examples/examples/analytic_N_scaling.py --help | sed "s/^/   /"

::

 $ python plot_analytic_N_scaling.py --help

.. exec::
   echo "::\\n\\n"
   python examples/examples/plot_analytic_N_scaling.py --help | sed "s/^/   /"


Here is an example generated by:

::

 $ python analytic_N_scaling.py --nNs 6
 $ python plot_analytic_N_scaling.py --savefig analytic_N_scaling.png


.. image:: ../_generated/analytic_N_scaling.png

"""

from __future__ import print_function, division, absolute_import

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

from chemreac.util.pyutil import progress

from analytic_diffusion import integrate_rd


def integrate(**kwargs):
    tout, yout, info, rmsd_over_atol, rd, rmsd = integrate_rd(**kwargs)
    info['rmsd'] = rmsd
    info['rmsd_over_atol'] = rmsd_over_atol
    return info

default_constant = dict(
    D=2e-3, t0=3.0, tend=7., x0=0.0, xend=1.0, center=None, nt=42, logt=False,
    logy=False, logx=False, random=False, p=0, a=0.2, linterpol=False,
    rinterpol=False, num_jacobian=False, method='bdf', integrator='cvode',
    linear_solver='gmres', atol=1e-8, rtol=1e-10, efield=False, random_seed=42,
    mobility=0.01, n_jac_diags=0, use_log2=False,
)


def main(plot=False, savefig='None', geoms='fcs', nNs=7, Ns=None,
         nspecies='1,2,3', nstencils='3,5,7', verbose=False, **kwargs):
    nstencils = [int(_) for _ in nstencils.split(',')]
    if Ns is None:
        Ns = [8*(2**i) for i in range(nNs)]
    else:
        Ns = list(map(int, Ns.split(',')))
        nNs = len(Ns)
    nspecies = list(map(int, nspecies.split(',')))
    varied = OrderedDict([('nspecies', nspecies), ('N', Ns), ('geom', geoms),
                          ('nstencil', nstencils)])

    kw1 = default_constant.copy()
    kw1.update(kwargs)

    results = {
        'varied_keys': list(varied.keys()),
        'varied_values': list(varied.values()),
    }

    all_params = list(product(*varied.values()))
    for params in progress(all_params) if verbose else all_params:
        kw2 = kw1.copy()
        kw2.update(dict(zip(varied.keys(), params)))
        results[params] = integrate(**kw2)
    basename = os.path.splitext(os.path.basename(__file__))[0]
    pickle.dump(results, gzip.open(basename + '.pkl', 'wb'))


if __name__ == '__main__':
    try:
        import argh
    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()
    else:
        argh.dispatch_command(main)
