#
# NuGridpy - Tools for accessing and visualising NuGrid data.
#
# Copyright 2007 - 2014
 by the NuGrid Team.
# All rights reserved. See LICENSE.
#

'''
nugridse is a collection of plots of data in se-type h5 files.

Usage
=====

start by loading the module,

>>> import nugridse as mp

you can get help,

>>> help mp

next, initiate a se
 class instance,

>>> pt=mp.se('path/to/dir/with/se-files')

which would read all h5 files from the current directory; again, do
help(mp.se) or mp.se?  To get info for more options, as for example how
to specify a match pattern to select only certain files in the
directory; this may be useful if there are very many files and you want
only every 10th.  For example,

>>> pt =mp.se('.','M1.65Z0.020.00')

   Note: Initializing an instance with 120 files with 1000 packets each
   may take 3-4 minutes, which is too long for many people, therefore
   the initialization module will generate and write an index file.  If
   such an index file is present initialization will take less then a
   second.  The initialization module will automatically detect various
   situations in which the index file needs to be rewritten, for
   example if there is a new file.

look at the data that is available in your instance,

>>> pt.sedir
>>> pt.sefiles

The actual data is in pt.se and the following commands give you access
to header and cycle attributes as well as the cycle profile data:

>>> pt.se.cattrs
>>> pt.se.hattrs
>>> pt.se.dcols

Available cycles can be viewed with,

>>> pt.se.cycles

You can get any of the quantities via the 'get' method, which is
relatively smart to give you things in various ways.

>>> pt.get('rho')

would give you all rho vectors for all cycles, which is maybe more
than you want, so try

>>> pt.get(300,'rho')

to get the rho vector for cycle 300.  Instead of one cycle you may
also supply a cycle list for the first argument.

Use help(pt) (or whatever you instance is) for a full description of
what else your instance has to offer, which includes methods to work
with data.

   Note: A particularly nice feature is that you can work with instances
   of different types of data in a very flexible way, for example in
   lists:

   >>> cases=[pt1,pt2,pt3]
   >>> for this_case in cases:
   ...    do something with this_case

There are various plotting methods, including plotting abundance charts
(abu_chart), isotopic abundance distributions and the generic 'plot'
method that lets you plot every quantity against any other quantity that
can possibly be plotted.  Some of the methods (including the three just
mentioned) are available via the super class data_plot, and are equally
available in other python modules, as for example mesa.py or ppn.py.
Several methods accept lists of cycles which implies to create a series
of frames to be written to disk in order to make movies.

>>> pt.plot('mass','rho',fname=3000)

whereas,

>>> pt.plot('mass','rho',fname=[3000,4000])

will produce two png files, one for each cycle, do help(m.plot) for a
full list of options.

mppnp output allows to do plots with abundances:

>>> pt.iso_abund(23500)
>>> pt.abu_chart(23500)

Also for iso_abund and abu_chart: if instead of a single cycle the user
inputs a list of cycles, the method will then, instead of plotting them,
will then save a .png for each cycle.  Also if you just want a singular
plot saved, the user can input their cycle, in a list like [0].  And
that will save their plot.

'''
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

from builtins import zip
from builtins import filter
from builtins import next
from builtins import input
from builtins import str
from builtins import range
from past.utils import old_div

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as pl
import matplotlib.pylab as pyl
from matplotlib.ticker import MultipleLocator
from matplotlib.collections import Collection
from matplotlib.artist import allow_rasterization
from matplotlib.patches import PathPatch
from scipy import interpolate
import os
import re
import sys
import time
import glob

from .utils import *
from .data_plot import *
from . import h5T


def set_nugrid_path(path):
    '''
    This function sets the path to the NuGrid VOSpace directory as a
    global variable, so that it need only be set once during an inter-
    active session.
    '''
    global nugrid_path
    nugrid_path=path

def set_nice_params():
    fsize=18

    params = {'axes.labelsize':  fsize,
    #    'font.family':       'serif',
    'font.family':        'Times New Roman',
    'figure.facecolor':  'white',
    'text.fontsize':     fsize,
    'legend.fontsize':   fsize,
    'xtick.labelsize':   fsize*0.8,
    'ytick.labelsize':   fsize*0.8,
    'text.usetex':       False}
    pl.rcParams.update(params)

class se(DataPlot, Utils):
    """
    This class provides easy access to h5 files from the NuGrid project,
    along with some standard plots.

    Parameters
    ----------
    sedir : string, optional
        The directory on which the h5 files can be found.  The default
        is '.'.
    pattern : string, optional
        A string pattern that the file names have to contain in order
        to be read.  The default is ".h5".
    rewrite : boolean, optional
        Would the user like to rewrite the preprocessor file.  The
        default is False.
    mass : integer or float, optional
        The user may select a mass and metallicity instead of providing
        the sedir explicitly, if they are using the VOSpace data. If mass
        is provided then Z should also be provided.
        The default is None (i.e. user gives sedir explicitly)
    Z : float, optional
        See 'mass' above.
        The default is None (i.e. user gives sedir explicitly)
    type : string, optional
        If the user gives mass and metallicity instead of giving sedir
        explicitly, then the type of data should also be given here.
        For example, type='ppd_wind', 'see_wind', see_exp' or 'ppd_exp'
        The default is 'ppd_wind'.
    output : string, optional
        If the user gives mass and metallicity instead of giving sedir
        explicitly, then the type of data should also be given here.
        For example, out='restart', 'out' or 'surf'
        The default is 'out'.
    data_set : string, optional
            Coose your data  of 'set1' or 'set1ext'.  The default is 'set1ext'.
    exp_type :
        If type is 'see_exp' or 'ppd_exp' runs of exp_type are selected. For example
        'delay' and 'rapid' would be choices for Set1.
    verbose : boolean, optional
        If True, print more output

    Examples
    --------

    >>> f=nu.plot_tools('.','260')

    reads all h5 files with the string 260 in the name in the present
    directory.

    """

    sedir = ''
    sefiles = []
    se = []         # main data dictionary
    pattern=''

    def __init__(self, sedir='.', pattern='.h5', rewrite=False, mass=None, Z=None, type='ppd_wind', output='out', code_source='MES',exp_type='delay',data_set='set1ext',verbose=False):

        # seeker to find the data requested on VOspace:
        if mass is not None and Z is not None:
            try:
                print('nugrid_path = '+nugrid_path)
            except:
                raise IOError("nugrid_path has not been set. This is the path to the NuGrid VOSpace, e.g. /tmp/NuGrid. Set this using nugridse.set_nugrid_path('path')")

            # which set? [find nearest]

            if (data_set=='set1ext'):
                setsZs=[0.02,0.01,6.e-3,1.e-3,1.e-4]
                setsnames=['set1.2','set1.1','set1.3a','set1.4a','set1.5a']
            elif (data_set=='set1'):
                setsZs=[0.02,0.01]
                setsnames=['set1.2','set1.1']
            else:
               raise IOError("Sorry. Requested data_set not available. Choose between set1ext and set1.")                  
            idx=np.abs(np.array(setsZs)-Z).argmin()
            setname=setsnames[idx]
            realZ=setsZs[idx]

            print('closest set is '+setname+' (Z = '+str(realZ)+')')

            # try first data, then data-team:
        