# Radau ODE system
from __future__ import division, absolute_import, print_function

import imp

from .allimports import *
from PyDSTool.Generator import ODEsystem as ODEsystem
from .baseclasses import theGenSpecHelper, genDB, _pollInputs
from .mixins import CompiledMixin, full_path
from PyDSTool.utils import *
from PyDSTool.common import *
# for future cleanup of * imports
from PyDSTool import utils
from PyDSTool import common
from PyDSTool.ModelSpec import QuantSpec
from PyDSTool.integrator import integrator
import numpy as npy

# Other imports
from numpy import Inf, NaN, isfinite, int, int32, float, float64, \
    sometrue, alltrue, any, all, concatenate, transpose, array, zeros
import operator
from copy import copy, deepcopy

class radau(integrator):
    """Radau 5 specialization of the basic integrator class."""

    def __init__(self, modname, rhs='default_name', phaseDim=0, paramDim=0,
                 nAux=0, nEvents=0, nExtInputs=0, hasJac=0, hasJacP=0,
                 hasMass=0, extraSpace=0, defaultBound=1e8):
        integrator.__init__(self, rhs=rhs, phaseDim=phaseDim, paramDim=paramDim,
                            nAux=nAux, nEvents=nEvents, nExtInputs=nExtInputs, hasJac=hasJac,
                            hasJacP=hasJacP, hasMass=hasMass, extraSpace=extraSpace,
                            defaultBound=defaultBound)
        self.modname = modname
        try:
            self._integMod = imp.load_module(
                modname, *imp.find_module(modname, ["radau5_temp"]))
        except:
            print("Error in importing compiled vector field and integrator.")
            print("Did you compile the RHS C code?")
            raise
        # check module's directory
        assert 'Integrate' in dir(self._integMod), \
               "radau library does not contain Integrate()"

        self.safety = []
        self.jacRecompute = []
        self.newtonStop = []
        self.stepChangeLB = []
        self.stepChangeUB = []
        self.stepSizeLB = []
        self.stepSizeUB = []
        self.hessenberg = []
        self.maxNewton = []
        self.newtonStart = []
        self.index1dim = []
        self.index2dim = []
        self.index3dim = []
        self.stepSizeStrategy = []
        self.DAEstructureM1 = []
        self.DAEstructureM2 = []

        retval = self._integMod.InitBasic(self.phaseDim, self.paramDim, self.nAux,
                                          self.nEvents, self.nExtInputs, self.hasJac,
                                          self.hasJacP, self.hasMass, self.extraSpace)

        if retval[0] != 1:
            raise PyDSTool_InitError('Call to InitBasic failed! (radau)')

        self.initBasic = True

    def Run(self, hinit=0, hmax=1.0, checkAux=0, calcSpecTimes=0, verbose=0,
            safety=0.9, jacRecompute=0.001, newtonStop=-1, stepChangeLB=1,
            stepChangeUB=1.2, stepSizeLB=0.2, stepSizeUB=8.0, hessenberg=0,
            maxNewton=7, newtonStart=0, index1dim=-1, index2dim=0, index3dim=0,
            stepSizeStrategy=1, DAEstructureM1=0, DAEstructureM2=0, useJac=0, useMass=0):
        if not self.initBasic:
            raise PyDSTool_InitError('initBasic is False (radau)')
        if not self.initEvents:
            raise PyDSTool_InitError('initEvents is False (radau)')
        if not self.initIntegrate:
            raise PyDSTool_InitError('initInteg is False (radau)')
        if not self.setParams:
            raise PyDSTool_InitError('setParams is False (radau)')
        if self.nExtInputs > 0 and not self.initExtInputs:
            raise PyDSTool_InitError('initExtInputs is False (radau)')

        self.setRadauParams(hinit=hinit, hmax=hmax, checkAux=checkAux,
                            calcSpecTimes=calcSpecTimes,
                            verbose=verbose, safety=safety,
                            jacRecompute=jacRecompute, newtonStop=newtonStop,
                            stepChangeLB=stepChangeLB, stepChangeUB=stepChangeUB,
                            stepSizeLB=stepSizeLB, stepSizeUB=stepSizeUB,
                            hessenberg=hessenberg,maxNewton=maxNewton,
                            newtonStart=newtonStart, index1dim=index1dim,
                            index2dim=index2dim, index3dim=index3dim,
                            stepSizeStrategy=stepSizeStrategy,
                            DAEstructureM1=DAEstructureM1,
                            DAEstructureM2=DAEstructureM2,
                            useJac=useJac,useMass=useMass)

        # For a run, we want to ensure indices are set to 0
        self.Reset()
        T, P, A, Stats, H, Err, EvtT, EvtP = self._integMod.Integrate(self.ic,
                                                      self.t0,
                                                      self.hinit,
                                                      self.hmax,
                                                      self.safety,
                                                      self.jacRecompute,
                                                      self.newtonStop,
                                                      self.stepChangeLB,
                                                      self.stepChangeUB,
                                                      self.stepSizeLB,
                                                      self.stepSizeUB,
                                                      self.hessenberg,
                                                      self.maxNewton,
                                                      self.newtonStart,
                                                      self.index1dim,
                                                      self.index2dim,
                                                      self.index3dim,
                                                      self.stepSizeStrategy,
                                                      self.DAEstructureM1,
                                                      self.DAEstructureM2,
                                                      self.useJac,
                                                      self.useMass,
                                                      self.verbose,
                                                      self.checkAux,
                                                      self.calcSpecTimes)
        self.points = P
        self.times = T
        self.auxPoints = A
        self.eventTimes = EvtT
        self.eventPoints = EvtP
        self.errors = Err
        self.stats = Stats
        self.step = H

        try:
            self.lastTime = self.times[-1]
            self.lastPoint = [self.points[i][-1] for i in range(self.phaseDim)]
            self.lastStep = self.step
        except IndexError:
            self.lastTime = self.t0
            self.lastPoint = self.ic
            self.lastStep = self.hinit
        self.numRuns += 1
        self.canContinue = True

        return T, P, A, Stats, H, Err, EvtT, EvtP

    def Continue(self, tend, params=[], calcSpecTimes=0, verbose=0,
                 extInputChanged=False, extInputVals=[], extInputTimes=[],
                 bounds=[]):
        if not self.initBasic:
            raise PyDSTool_InitError('initBasic is False (radau)')
        if not self.initEvents:
            raise PyDSTool_InitError('initEvents is False (radau)')
        if not self.initIntegrate:
            raise PyDSTool_InitError('initInteg is False (radau)')
        if not self.setParams:
            raise PyDSTool_InitError('setParams is False (radau)')
        if self.nExtInputs > 0 and not self.initExtInputs:
            raise PyDSTool_InitError('initExtInputs is False (radau)')

        if not self.canContinue:
            raise PyDSTool_ContError('Unable to continue trajectory -- '
                        'have you run the integrator and reset events, etc?')

        self.setContParams(tend=tend, params=copy(params),
                           calcSpecTimes=calcSpecTimes, verbose=verbose, extInputChanged=extInputChanged,
                           extInputVals=copy(extInputVals), extInputTimes=copy(extInputTimes),
                           bounds=copy(bounds))

        # For a continue, we do not set indices to 0
        T, P, A, Stats, H, Err, EvtT, EvtP = \
                self._integMod.Integrate(self.lastPoint,
                                      self.lastTime,
                                      self.lastStep, self.hmax,
                                      self.safety,
                                      self.jacRecompute,
                                      self.newtonStop,
                                      self.stepChangeLB,
                                      self.stepChangeUB,
                                      self.stepSizeLB,
                                      self.stepSizeUB,
                                      self.hessenberg,
                                      self.maxNewton,
                                      self.newtonStart,
                                      self.index1dim,
                                      self.index2dim,
                                      self.index3dim,
                                      self.stepSizeStrategy,
                                      self.DAEstructureM1,
                                      self.DAEstructureM2,
                                      self.useJac,
                                      self.useMass,
                                      self.verbose,
                                      self.checkAux,
                                      self.calcSpecTimes)

        self.points = P
        self.times = T
        self.auxPoints = A
        self.eventTimes = EvtT
        self.eventPoints = EvtP
        self.errors = Err
        self.stats = Stats
        self.step = H

        try:
            self.lastTime = self.times[-1]
            self.lastPoint = [self.points[i][-1] for i in range(self.phaseDim)]
            self.lastStep = self.step
        except IndexError:
            self.lastTime = self.t0
            self.lastPoint = self.ic
            self.lastStep = self.hinit
        self.numRuns += 1
        self.numContinues += 1
        self.canContinue = True

        return T, P, A, Stats, H, Err, EvtT, EvtP

    def setRadauParams(self, hinit, hmax, checkAux, calcSpecTimes,
                       verbose, safety, jacRecompute, newtonStop,
                       stepChangeLB, stepChangeUB, stepSizeLB, stepSizeUB,
                       hessenberg, maxNewton, newtonStart, index1dim,
                       index2dim, index3dim, stepSizeStrategy,
                       DAEstructureM1, DAEstructureM2, useJac, useMass):
        useJac = int(useJac)
        useMass = int(useMass)
        checkAux = int(checkAux)
        calcSpecTimes = int(calcSpecTimes)
        hessenberg = int(hessenberg)

        if not isinstance(hinit, _num_types):
            raise TypeError("hinit must be int, float")

        if not isinstance(hmax, _num_types):
            raise TypeError("hmax must be int, float")

        if abs(hinit) > abs(hmax):
            raise ValueError("Abs value of hinit (%g) must be less than hmax (%g)"%(hinit,hmax))

        if not isinstance(checkAux, _int_types):
            raise TypeError("checkAux must be int")
        if checkAux not in (0,1):
            raise TypeError("checkAux must be 0 or 1")
        if checkAux == 1 and self.nAux <= 0:
            raise ValueError("checkAux cannot be 1 if nAux is 0")

        if not isinstance(verbose, _int_types):
            raise TypeError("verbose must be int")
        if verbose not in (0,1):
            if verbose >= 2:
                # interpret all greater values as 1
                verbose = 1
            else:
                raise TypeError("verbose must be 0 or 1")

        if not isinstance(calcSpecTimes, _int_types):
            raise TypeError("calcSpecTimes must be int")
        if calcSpecTimes not in (0,1):
            raise TypeError("calcSpecTimes must be 0 or 1")
        if calcSpecTimes == 1 and len(self.specTimes) <= 0:
            raise ValueError("calcSpecTimes cannot be 1 if specTimes is empty")

        if safety < 0:
            raise ValueError("safety must be non-negative")
        if jacRecompute <= 0.0:
            raise ValueError("jacRecompute must be positive")
        if newtonStop < 0:
            newtonStop = 0
        if stepChangeLB <= 0:
            raise ValueError("stepChangeLB must be positive")
        if stepChangeUB <= 0:
            raise ValueError("stepChangeUB must be positive")
        if stepSizeLB <= 0:
            raise ValueError("stepSizeLB must be positive")
        if stepSizeUB <= 0:
            raise ValueError("stepSizeUB must be positive")

        if stepChangeLB > stepChangeUB:   # was >= but this allows fac1=fac2=1
            raise ValueError("stepChangeLB must be less than stepChangeUB")
        if stepSizeLB >= stepSizeUB:
            raise ValueError("stepSizeLB must be less than stepSizeUB")

        if hessenberg not in (0,1):
            raise ValueError("hessenberg must be 0 or 1")
        if hessenberg == 1 and useMass != 0:
            raise ValueError("hessenberg form cannot be used for implicit systems (mass matrix)")
        if not isinstance(maxNewton, _int_types):
            raise TypeError("maxNewton must be int")
        if maxNewton <= 0:
            raise ValueError("maxNewton must be positive")

        if newtonStart not in (0,1):
            raise ValueError("newtonStart must be 0 or 1")

        if index1dim <= 0:
            index1dim = self.phaseDim
        if index2dim != 0:
            raise ValueError("Currently index2dim must be 0")
        if index3dim != 0:
            raise ValueError("Currently index3dim must be 0")

        if stepSizeStrategy not in (1,2):
            raise ValueError("stepSizeStrategy must be 1 or 2")

        if DAEstructureM1 != 0:
            raise ValueError("Currently DAEstructureM1 must be 0")
        if DAEstructureM2 != 0:
            raise ValueError("Currently DAEstructureM2 must be 0")

        if useJac not in (0,1):
            raise ValueError("useJac must be 0 or 1")
        if useMass not in (0,1):
            raise ValueError("useMass must be 0 or 1")

        if useJac == 1 and self.hasJac != 1:
            raise ValueError("useJac must be 0 if hasJac is not 1")
        if useMass == 1 and self.hasMass != 1:
            raise ValueError("useMass must be 0 if hasMass is not 1")

        self.hinit = hinit
        self.hmax = hmax
        self.safety = safety
        self.jacRecompute = jacRecompute
        self.newtonStop = newtonStop
        self.stepChangeLB = stepChangeLB
        self.stepChangeUB = stepChangeUB
        self.stepSizeLB = stepSizeLB
        self.stepSizeUB = stepSizeUB
        self.hessenberg = hessenberg
        self.maxNewton = maxNewton
        self.newtonStart = newtonStart
        self.index1dim = index1dim
        self.index2dim = index2dim
        self.index3dim = index3dim
        self.stepSizeStrategy = stepSizeStrategy
        self.DAEstructureM1 = DAEstructureM1
        self.DAEstructureM2 = DAEstructureM2
        self.useJac = useJac
        self.useMass = useMass
        self.verbose = verbose
        self.checkAux = checkAux
        self.calcSpecTimes = calcSpecTimes

class Radau_ODEsystem(ODEsystem, CompiledMixin):
    """Wrapper for Radau integrator (with support for differential-algebraic equations).

    Uses C target language only for functional specifications"""
    _paraminfo = {'rtol': 'Relative error tolerance.',
                  'atol': 'Absolute error tolerance.',
                  'safety': 'Safety factor in the step size prediction, default 0.9.',
                  'max_step': 'Maximal step size, default tend-tstart.',
                  'init_step': 'Initial step size, default is a guess computed by the function init_step.',
                  'fac1': 'Parameter for step size selection; the new step size is chosen subject to the restriction  fac1 <= new_step/old_step <= fac2. Default value is 1.0.',
                  'fac2': 'Parameter for step size selection; the new step size is chosen subject to the restriction  fac1 <= new_step/old_step <= fac2. Default value is 1.2.',
                  'stepLB': '',
                  'stepUB': '',
                  'refine': 'Refine output by adding points interpolated using the RK4 polynomial (0, 1 or 2).',
                  'step_strategy': """Switch for step size strategy;
If step_strategy=1  mod. predictive controller (Gustafsson).
If step_strategy=2  classical step size control.
The default value (for step_strategy=0) is step_strategy=1.
the choice step_strategy=1 seems to produce safer results;
for simple problems, the choice step_strategy=2 produces
often slightly faster runs.""",
                  'jac_recompute': """Decides whether the Jacobian should be recomputed;
increase jac_recompute to 0.1 say, when Jacobian evaluations
are costly. for small systems jac_recompute should be smaller
(0.001, say). negative jac_recompute forces the code to
compute the Jacobian after every accepted step.
Default 0.001.""",
                  'newton_start': "",
                  'newton_stop': "",
                  'max_newton': "Maximum number of Newton iterations to take in solving the implicit system at each step (default 7)",
                  'DAEstructureM1': "",
                  'DAEstructureM2': "",
                  'hessenberg': "",
                  'index1dim': "",
                  'index2dim': "",
                  'index3dim': "",
                  'use_special': "Switch for using special times",
                  'specialtimes': "List of special times to use during integration",
                  'check_aux': "Switch",
                  'extraspace': ""
                  }

    def __init__(self, kw):
        """Use the nobuild key to postpone building of the library, e.g. in
        order to provide additional build options to makeLibSource and
        compileLib methods or to make changes to the C code by hand.
        No build options can be specified otherwise."""

        # delete because not covered in ODEsystem
        nobuild = kw.pop('nobuild', False)
        ODEsystem.__init__(self, kw)
        self._solver = None
        self.diagnostics._errorcodes = {
              0: 'Unrecognized error code returned (see stderr output)',
            -1 : 'input is not consistent',
            -2 : 'larger nmax is needed',
             2 : 'larger nmax or maxevtpts is probably needed (error raised by solout)',
            -3 : 'step size becomes too small',
            -4 : 'the matrix is repeatedly singular (interrupted)',
            -8 : 'The solution exceeded a magbound (poor choice of initial step)'}
        self.diagnostics.outputStatsInfo = {
            'last_step': 'Predicted step size of the last accepted step (useful for a subsequent call to radau).',
            'num_steps': 'Number of used steps.',
            'num_accept': 'Number of accepted steps.',
            'num_reject': 'Number of rejected steps.',
            'num_fcns': 'Number of function evaluations.',
            'num_jacs': 'Number of Jacobian evaluations.',
            'num_dec': 'Number of LU-decompositions',
            'num_subs': 'Number of forward-backward substitutions',
            'errorStatus': 'Error status on completion.'
             }

        # currently the final four of these params are for event handling
        algparams_def = {'poly_interp': False,
                        'init_step': 0,
                        'max_step': 0,
                        'rtol': [1e-9 for i in range(self.dimension)],
                        'atol': [1e-12 for i in range(self.dimension)],
                        'fac1': 1.0,
                        'fac2': 1.2,
                        'stepLB': 0.2,
                        'stepUB': 8.0,
                        'safety': 0.9,
                        'max_pts': 10000,
                        'refine': 0,
                        'maxbisect': [], # for events
                        'maxevtpts': 1000, # for events
                        'eventInt': [], # set using setEventInterval only
                        'eventDelay': [], # set using setEventDelay only
                        'eventTol': [], # set using setEventTol only
                        'use_special': 0,
                        'specialtimes': [],
                        'check_aux': 1,
                        'extraspace': 100,
                        'verbose': 0,
                        'jac_recompute': 0.001,
                        'step_strategy': 1,
                        'index1dim': -1,
                        'index2dim': 0,
                        'index3dim': 0,
                        'DAEstructureM1': 0,
                        'DAEstructureM2': 0,
                        'hessenberg': 0,
                        'newton_start': 0,
                        'newton_stop': -1,
                        'max_newton': 7,
                        'hasJac': 0,
                        'hasJacP': 0,
                        'checkBounds': self.checklevel
                        }
        for k, v in algparams_def.items():
            if k not in self.algparams:
                self.algparams[k] = v
        # verify that no additional keys are present in algparams, after
        # defaults are added above
        if len(self.algparams) != len(algparams_def):
            raise ValueError("Invalid keys present in algparams argument: " \
                     + str(remain(self.algparams.keys(),algparams_def.keys())))
        # Check for non-constant mass matrix
        if self.haveMass():
            mspec = self.funcspec.auxfns['massMatrix']
            lensig = len(mspec[1])
            body_str = mspec[0][lensig:].replace('\n','')
            qbody = QuantSpec('__body__', body_str, treatMultiRefs=False,
                              ignoreSpecial=['[',']','{','}'])
            self._const_massmat = intersect(['Y_','t'], qbody.usedSymbols) == []
        else:
            self._const_massmat = True

        self._prepareEventSpecs()
        self._inputVarList = []
        self._inputTimeList = []

        if nobuild:
            print("Build the library using the makeLib method, or in ")
            print("stages using the makeLibSource and compileLib methods.")
        else:
            self.makeLib()

    @property
    def integrator(self):
        return {
            'name': ('radau5' if self._const_massmat else 'radau5v', 'Radau'),
            'description': "Radau5 integrator" + \
            "" if self._const_massmat else " (version for non-constant mass matrices)",
            'src': ["radau5mod.c"],
            'cflags': ["-D__RADAU__"],
            'libs': [
                ('radau5', {
                    'sources': full_path(['radau5.f' if self._const_massmat else 'radau5v.f']),
                    'extra_f77_compile_args': utils.extra_arch_arg(['-w']),
                }),
                ('lapack_lite', {
                    'sources': full_path(['lapackc.f', 'lapack.f', 'dc_lapack.f']),
                    'extra_f77_compile_args': utils.extra_arch_arg(['-w']),
                })
            ],
        }

    def _prepareEventSpecs(self):
        eventActive = []
        eventTerm = []
        eventDir = []
        eventDelay = []
        eventTol = []
        maxbisect = []
        eventInt = []
        # convert event specs (term, active, etc.) into integparam specs
        self._eventNames = self.eventstruct.sortedEventNames()
        for evname in self._eventNames:
            ev = self.eventstruct.events[evname]
            assert isinstance(ev, LowLevelEvent), ("Radau can only "
                                                "accept low level events")
        # if event 'precise' flags set to False then set their tolerances
        # to be > max_step
        maxstep = self.algparams['max_step']
        for evname in self._eventNames:
            ev = self.eventstruct.events[evname]
            eventActive.append(int(ev.activeFlag))
            eventTerm.append(int(ev.termFlag))
            eventDir.append(ev.dircode)
            eventInt.append(ev.eventinterval)
            eventDelay.append(ev.eventdelay)
            if ev.preciseFlag:
                eventTol.append(ev.eventtol)
                maxbisect.append(ev.bisectlimit)
            else:
                eventTol.append(maxstep*1.5)
                maxbisect.append(1)
        self.algparams['eventTol'] = eventTol
        self.algparams['eventDelay'] = eventDelay
        self.algparams['eventInt'] = eventInt
        self.algparams['maxbisect'] = maxbisect
        self.algparams['eventActive'] = eventActive
        self.algparams['eventTerm'] = eventTerm
        self.algparams['eventDir'] = eventDir

    def compute(self, trajname, dirn='f', ics=None):
        continue_integ = ODEsystem.prepDirection(self, dirn)
        if ics is not None:
            self.set(ics=ics)
        self.validateICs()
        self.diagnostics.clearWarnings()
        self.diagnostics.clearErrors()
        if isinstance(self.algparams['rtol'], list):
            if len(self.algparams['rtol']) != self.dimension:
                raise ValueError('rtol list must have same length as phase dimension')
        else:
            rtol = self.algparams['rtol']
            self.algparams['rtol'] = [rtol for i in range(self.dimension)]
        if isinstance(self.algparams['atol'], list):
            if len(self.algparams['atol']) != self.dimension:
                raise ValueError('atol list must have same length as phase dimension')
        else:
            atol = self.algparams['atol']
            self.algparams['atol'] = [atol for i in range(self.dimension)]
        anames = self.funcspec.auxvars
        # Check i.c.'s are well defined (finite)
        self.checkInitialConditions()
        self.setEventICs(self.initialconditions, self.globalt0)
        # update event params in case changed since last run
        self._prepareEventSpecs()
        # Main integration
        t0 = self.indepvariable.depdomain[0]
        t1 = self.indepvariable.depdomain[1]
        plist = sortedDictValues(self.pars)
        self.algparams['hasJac'] = self.haveJacobian()
        self.algparams['hasJacP'] = self.haveJacobian_pars()
        self._ensure_solver()
        if self._dircode == 1:
            tbegin = t0
            tend = t1
        elif self._dircode == -1:
            # radau does reverse time integration simply by switching t0 and t1
            tbegin = t1
            tend = t0
        if len(self.algparams['specialtimes'])>0:
            use_special = self.algparams['use_special']
        else:
            use_special = 0
        bounds = [[],[]]  # lower, then upper
        for v in self.funcspec.vars:
            bds = self.xdomain[v]
            try:
                bounds[0].append(bds[0])
                bounds[1].append(bds[1])
            except TypeError:
                print("%r %s %r" % (v, type(bds), bds))
                print(self.xdomain)
                raise
        for p in self.funcspec.pars:
            bds = self.pdomain[p]
            try:
                bounds[0].append(bds[0])
                bounds[1].append(bds[1])
            except TypeError:
                print("%s %r" % (type(bds), bds))
                raise
        if continue_integ:
            x0 = self._solver.lastPoint
            # overwrite t0 from self.indepvariable.domain, but use its t1
            tbegin = self._solver.lastTime
            if abs(self._solver.lastStep) < abs(self.algparams['init_step']):
                self.algparams['init_step'] = self._solver.lastStep
            if abs(t1-tbegin) < abs(self.algparams['init_step']):
                raise ValueError("Integration end point too close to initial "
                                 "point")
#            if self.inputs and self._extInputsChanged:
#                self._extInputsChanged = False
#                self._solver.setContParams(tend, plist,
#                                           use_special,
#                                           self.algparams['verbose'],
#                                           True, deepcopy(self._inputVarList),
#                                           deeppcopy(self._inputTimeList))
        else:
            if self._solver.numRuns > 0:
                self._solver.clearAll()
            x0 = sortedDictValues(self.initialconditions, self.funcspec.vars)
            self._solver.setInteg(maxpts=self.algparams['max_pts'],
                rtol=self.algparams['rtol'], atol=self.algparams['atol'])
            self._solver.setRunParams(ic=x0, params=plist,
                                  t0=tbegin, tend=tend, gt0=self.globalt0,
                                  refine=self.algparams['refine'],
                                  specTimes=self.algparams['specialtimes'],
                                  bounds=bounds)
        if self.inputs:
            # self._extInputsChanged if global t0 changed so that can
            # adjust times given to the integrator (it is blind to global t0
            # when accesses input variable times)
            self._ensure_inputs(self._extInputsChanged)
        # hinit only set if not continue_integ
        if len(anames)>0:
            check_aux = self.algparams['check_aux']
        else:
            check_aux = 0
        if self.algparams['max_step'] == 0:
            max_step = abs(tend-tbegin)
        else:
            max_step = self.algparams['max_step']
        init_step = self.algparams['init_step']
        if self._dircode == 1:
            if init_step < 0:
                init_step = -init_step
            if max_step < 0:
                max_step = -max_step
        else:
            if init_step > 0:
                init_step = -init_step
            if max_step > 0:
                max_step = -max_step
        if continue_integ:
            # record needed for bounds checking and truncation
            old_highest_ix = self._solver.points.shape[1]
            alltData, X, A, Stats, H, Err, Evtimes, \
                Evpoints = self._solver.Continue(tend, plist,
                                  use_special, self.algparams['verbose'],
                                  self._extInputsChanged,
                                  deepcopy(self._inputVarList),
                                  deepcopy(self._inputTimeList), bounds)
        else:
            old_highest_ix = 0
            self._solver.setEvents(eventActive=self.algparams['eventActive'],
                            eventTerm=self.algparams['eventTerm'],
                            eventDir=self.algparams['eventDir'],
                            eventDelay=self.algparams['eventDelay'],
                            eventInt=self.algparams['eventInt'],
                            eventTol=self.algparams['eventTol'],
                            maxevtpts=self.algparams['maxevtpts'],
                            maxbisect=self.algparams['maxbisect'])
            alltData, X, A, Stats, H, Err, Evtimes, \
                Evpoints = self._solver.Run(init_step,
                                    max_step,
                                    check_aux,
                                    use_special,
                                    self.algparams['verbose'],
                                    self.algparams['safety'],
                                    self.algparams['jac_recompute'],
                                    self.algparams['newton_stop'],
                                    self.algparams['fac1'],
                                    self.algparams['fac2'],
                                    self.algparams['stepLB'],
                                    self.algparams['stepUB'],
                                    self.algparams['hessenberg'],
                                    self.algparams['max_newton'],
                                    self.algparams['newton_start'],
                                    self.algparams['index1dim'],
                                    self.algparams['index2dim'],
                                    self.algparams['index3dim'],
                                    self.algparams['step_strategy'],
                                    self.algparams['DAEstructureM1'],
                                    self.algparams['DAEstructureM2'],
                                    self.haveJacobian(),
                                    self.haveMass())
        self._extInputsChanged = False    # reset this now
        self.diagnostics.outputStats = {'last_step': H,
                            'last_time': self._solver.lastTime,
                            'last_point': self._solver.lastPoint,
                            'num_fcns': Stats[0],
                            'num_jacs': Stats[1],
                            'num_steps': Stats[2],
                            'num_accept': Stats[3],
                            'num_reject': Stats[4],
                            'num_dec': Stats[5],
                            'num_subs': Stats[6],
                            'errorStatus': Err
                            }
        if self._dircode == -1:
            # reverse the array object (no reverse method!)
            alltData = alltData[::-1]
            X = X[:,::-1]
            if anames != []:
                A = A[:,::-1]
        xnames = self._var_ixmap
        # Package up computed trajectory in Variable variables
        # Add external inputs warnings to self.diagnostics.warnings, if any
##        for f in inputVarList:
##            for winfo in f.diagnostics.warnings:
##                self.diagnostics.warnings.append((W_NONTERMSTATEBD,
##                                     (winfo[0], f.name, winfo[1],
##                                      f.depdomain)))
        eventslist = self.eventstruct.query(['lowlevel', 'active'])
        termevents = self.eventstruct.query(['term'], eventslist)
        if self._eventNames != []:
            # build self.diagnostics.warnings because events happened --
            # and keep a record of which times terminal events happened because
            # Model.py's event handling procedure assumes multiple events
            # happening at one time are listed in one warning
            termevtimes = {}
            nontermevtimes = {}
            try:
                for evix in range(len(self._eventNames)):
                    if Evpoints[evix] is None:
                        continue
                    evname = self._eventNames[evix]
                    numevs = len(Evtimes[evix])
                    if self.algparams['eventTerm'][evix]:
                        if numevs > 1:
                            print("Event info: %r %r" % (Evpoints, Evtimes))
                        assert numevs <= 1, ("Internal error: more than one "
                                         "terminal event of same type found")
                        # For safety, we should assert that this event
                        # also appears in termevents, but we don't
                        if Evtimes[evix][0] in termevtimes.keys():
                            # append event name to this warning
                            warning_ix = termevtimes[Evtimes[evix][0]]
                            self.diagnostics.warnings[warning_ix][1][1].append(evname)
                        else:
                            # make new termevtime entry for the new warning
                            termevtimes[Evtimes[evix][0]] = \
                                       len(self.diagnostics.warnings)
                            self.diagnostics.warnings.append((W_TERMEVENT,
                                             (Evtimes[evix][0],
                                             [evname])))
                    else:
                        for ev in range(numevs):
                            if Evtimes[evix][ev] in nontermevtimes.keys():
                                # append event name to this warning
                                warning_ix = nontermevtimes[Evtimes[evix][ev]]
                                self.diagnostics.warnings[warning_ix][1][1].append(evname)
                            else:
                                # make new nontermevtime entry for the new warning
                                nontermevtimes[Evtimes[evix][ev]] = \
                                                    len(self.diagnostics.warnings)
                                self.diagnostics.warnings.append((W_NONTERMEVENT,
                                                 (Evtimes[evix][ev],
                                                  [evname])))
            except IndexError:
                print("Events returned from integrator are the wrong size.")
                print("  Did you change the system and not refresh the C " \
                      + "library using the forcelibrefresh() method?")
                raise
        termcount = 0
        for (w,i) in self.diagnostics.warnings:
            if w == W_TERMEVENT or w == W_TERMSTATEBD:
                if termcount > 0:
                    raise ValueError("Internal error: more than one terminal "
                                     "event found")
                termcount += 1
        # post-process check of variable bounds (if defined and algparams['checkBounds'] True)
        if self._dircode > 0:
            compare = operator.lt
            last_ix = Inf
        else:
            compare = operator.gt
            last_ix = -Inf
        highest_ix = X.shape[1]-1
        last_t = Inf
        if self.algparams['checkBounds'] > 0:
            # temp storage for repeatedly used object attributes (for lookup efficiency)
            depdomains = dict(zip(range(self.dimension),
                                  [self.variables[xn].depdomain for xn in xnames]))
            offender_ix = None
            for xi in range(self.dimension):
                if not any(depdomains[xi].isfinite()):
                    # no point in checking when the bounds are +/- infinity
                    continue
                next_last_ix = array_bounds_check(X[xi][old_highest_ix:],
                                    depdomains[xi], self._dircode) + old_highest_ix
                if compare(next_last_ix, last_ix):
                    # won't count as truncating unless the following checks
                    # hold
                    last_ix = next_last_ix
                    offender_ix = xi
            if not isfinite(last_ix) and last_ix < 0:
                # only use +Inf hereon to flag no truncation needed
                last_ix = Inf
            elif last_ix >= 0 and last_ix < highest_ix:
                # truncate data
                last_t = alltData[last_ix]
                print("Warning; domain bound reached (because algparams['checkBounds'] > 0)")
                self.diagnostics.warnings.append((W_TERMSTATEBD,
                                    (last_t, xnames[offender_ix],
                                     X[offender_ix, last_ix],
                                     depdomains[offender_ix].get())))
        # Create variables (self.variables contains no actual data)
        variables = copyVarDict(self.variables)
        # build event pointset information (reset previous trajectory's)
        # don't include events after any truncation due to state bound violation
        self.trajevents = {}
        for evix in range(len(self._eventNames)):
            evname = self._eventNames[evix]
            if Evpoints[evix] is None:
                self.trajevents[evname] = None
            else:
                try:
                    ev_a_list = []
                    for t in Evtimes[evix]:
                        tix = find(alltData, t)
                        ev_a_list.append(A[:,tix])
                    ev_array = concatenate((Evpoints[evix],
                                         transpose(array(ev_a_list, 'd'))))
                    del ev_a_list, tix
                except TypeError:
                    # A is empty
                    ev_array = Evpoints[evix]
                if last_ix >= 0 and last_ix < highest_ix:
                    # don't count last_ix = -1 which is the same as highest_ix
                    last_ev_tix = npy.argmax(Evtimes[evix] >= alltData[last_ix])
                    if last_ev_tix == 0 and Evtimes[evix][0] >= last_t:
                        # checks that there was actually a violation
                        # - so no events to record
                        self.trajevents[evname] = None
                    else:
                        # truncation needed
                        ev_array = ev_array[:, :last_ev_tix+1]
                        ev_times = Evtimes[evix][:last_ev_tix+1]
                        self.trajevents[evname] = Pointset({'coordnames': xnames+anames,
                                               'indepvarname': 't',
                                               'coordarray': ev_array,
                                               'indepvararray': ev_times})
                else:
                    # no truncation needed
                    self.trajevents[evname] = Pointset({'coordnames': xnames+anames,
                                               'indepvarname': 't',
                                               'coordarray': ev_array,
                                               'indepvararray': Evtimes[evix]})
        if last_ix >= 0 and last_ix < highest_ix:
            # truncate
            X = X[:, :last_ix]
            alltData = alltData[:last_ix]
        try:
            allxDataDict = dict(zip(xnames,X))
        except IndexError:
            print("Integration returned variable values of unexpected dimensions.")
            print("  Did you change the system and not refresh the C library" \
                  + " using the forcelibrefresh() method?")
            raise
        # storage of all auxiliary variable data
        anames = self.funcspec.auxvars
        try:
            if anames != []:
                if last_ix < highest_ix:
                    A = A[:, :last_ix]
                try:
                    allaDataDict = dict(zip(anames,A))
                except TypeError:
                    print("Internal error!  Type of A: %s" % type(A))
                    raise
        except IndexError:
            print("Integration returned auxiliary values of unexpected dimensions.")
            print("  Did you change the system and not refresh the C library" \
                  + " using the forcelibrefresh() method?")
            raise
        if int(Err) == 1 or (int(Err) == 2 and termcount == 1):
            # output OK
            if self.algparams['poly_interp']:
                rhsfn = self._solver.Rhs
                # when Dopri can output the Rhs values alongside variable
                # values then this won't be necessary
                dxvals = zeros((len(alltData),self.dimension),float)
                for tix, tval in enumerate(alltData):
                    # solver's Rhs function already contains the inputs so no
                    # need to recompute and provide here.
                    #i = _pollInputs(sortedDictValues(self.inputs), tval,
                    #                        self.checklevel)
                    # X is the output variable array, but rhsfn demands a list
                    dxvals[tix] = rhsfn(tval, list(X[:,tix]), plist)[0]
            for xi, x in enumerate(xnames):
                if len(alltData) > 1:
                    if self.algparams['poly_interp']:
                        interp = PiecewisePolynomial(alltData,
                                    array([allxDataDict[x], dxvals[:,xi]]).T, 2)
                    else:
                        interp = interp1d(alltData, allxDataDict[x])
                    variables[x] = Variable(interp, 't', x, x)
                else:
                    raise PyDSTool_ValueError("Fewer than 2 data points computed")
            for a in anames:
                if len(alltData) > 1:
                    variables[a] = Variable(interp1d(alltData,allaDataDict[a]),
                                             't', a, a)
                else:
                    raise PyDSTool_ValueError("Fewer than 2 data points computed")
            # final checks
            #self.validateSpec()
            self.defined = True
            return Trajectory(trajname, list(variables.values()),
                              abseps=self._abseps, globalt0=self.globalt0,
                              checklevel=self.checklevel,
                              FScompatibleNames=self._FScompatibleNames,
                              FScompatibleNamesInv=self._FScompatibleNamesInv,
                              events=self.trajevents,
                              modelNames=self.name,
                              modelEventStructs=self.eventstruct)
        else:
            try:
                diagnost_info = self.diagnostics._errorcodes[int(Err)]
            except TypeError:
                # errcode messed up from Radau
                print("Error code: %d" % Err)
                diagnost_info = self.diagnostics._errorcodes[0]
            if self._solver.verbose:
                info(self.diagnostics.outputStats, "Output statistics")
            self.defined = False
            # Did the solver run out of memory?
            if (len(alltData) == self.algparams['max_pts'] or \
                self.diagnostics.outputStats['num_steps'] >= self.algparams['max_pts']) \
                   and alltData[-1] < tend:
                print("max_pts algorithmic parameter too small: current " + \
                      "value is %i"%self.algparams['max_pts'])
#                avstep = (self.algparams['init_step']+self.diagnostics.outputStats['last_step'])/2.
                if self.diagnostics.outputStats['last_time']-tbegin > 0:
                    ms = str(int(round(self.algparams['max_pts'] / \
                              (self.diagnostics.outputStats['last_time'] - \
                               tbegin)*(tend-tbegin))))
                else:
                    ms = 'Inf'
                print("(recommended value for this trajectory segment is " + \
                      "estimated to be %s (saved in diagnostics.errors attribute))"%str(ms))
                diagnost_info += " -- recommended value is " + ms
            self.diagnostics.errors.append((E_COMPUTFAIL,
                                    (self._solver.lastTime, diagnost_info)))
            raise PyDSTool_ExistError("No trajectory created")

    def Rhs(self, t, xdict, pdict=None, asarray=True):
        """asarray is an unused, dummy argument for compatibility with Model.Rhs"""
        # must convert names to FS-compatible as '.' sorts before letters
        # while '_' sorts after!
        x = sortedDictValues(filteredDict(self._FScompatibleNames(xdict),
                                          self.funcspec.vars))
        if pdict is None:
            pdict = self.pars
            # internal self.pars already is FS-compatible
            p = sortedDictValues(pdict)
        else:
            p = sortedDictValues(self._FScompatibleNames(pdict))
        i = _pollInputs(sortedDictValues(self.inputs),
                        t, self.checklevel)
        self._ensure_solver({'params': p, 't0': 0, 'tend': 1})
        self._ensure_inputs()
        return self._solver.Rhs(t, x, p+i)[0]

    def Jacobian(self, t, xdict, pdict=None, asarray=True):
        """asarray is an unused, dummy argument for compatibility with
        Model.Jacobian"""
        if self.haveJacobian():
            x = sortedDictValues(filteredDict(self._FScompatibleNames(xdict),
                                              self.funcspec.vars))
            if pdict is None:
                pdict = self.pars
                # internal self.pars already is FS-compatible
                p = sortedDictValues(pdict)
            else:
                p = sortedDictValues(self._FScompatibleNames(pdict))
            i = _pollInputs(sortedDictValues(self.inputs),
                            t, self.checklevel)
            self._ensure_solver({'params': p, 't0': 0, 'tend': 1})
            self._ensure_inputs()
            return self._solver.Jacobian(t, x, p+i)[0]
        else:
            raise PyDSTool_ExistError("Jacobian not defined")

    def JacobianP(self, t, xdict, pdict=None, asarray=True):
        """asarray is an unused, dummy argument for compatibility with
        Model.JacobianP"""
        if self.haveJacobian_pars():
            x = sortedDictValues(filteredDict(self._FScompatibleNames(xdict),
                                              self.funcspec.vars))
            if pdict is None:
                pdict = self.pars
                # internal self.pars already is FS-compatible
                p = sortedDictValues(pdict)
            else:
                p = sortedDictValues(self._FScompatibleNames(pdict))
            i = _pollInputs(sortedDictValues(self.inputs),
                            t, self.checklevel)
            self._ensure_solver({'params': p, 't0': 0, 'tend': 1})
            self._ensure_inputs()
            return self._solver.JacobianP(t, x, p+i)[0]
        else:
            raise PyDSTool_ExistError("Jacobian w.r.t. parameters not defined")

    def AuxVars(self, t, xdict, pdict=None, asarray=True):
        """asarray is an unused, dummy argument for compatibility with
        Model.AuxVars"""
        x = sortedDictValues(filteredDict(self._FScompatibleNames(xdict),
                                          self.funcspec.vars))
        if pdict is None:
            pdict = self.pars
            # internal self.pars already is FS-compatible
            p = sortedDictValues(pdict)
        else:
            p = sortedDictValues(self._FScompatibleNames(pdict))
        i = _pollInputs(sortedDictValues(self.inputs),
                        t, self.checklevel)
        self._ensure_solver({'params': p, 't0': 0, 'tend': 1})
        self._ensure_inputs()
        return self._solver.AuxFunc(t, x, p+i)[0]

    def MassMatrix(self, t, xdict, pdict=None, asarray=True):
        """asarray is an unused, dummy argument for compatibility with
        Model.MassMatrix"""
        if self.haveMass():
            x = sortedDictValues(filteredDict(self._FScompatibleNames(xdict),
                                              self.funcspec.vars))
            if pdict is None:
                pdict = self.pars
                # internal self.pars already is FS-compatible
                p = sortedDictValues(pdict)
            else:
                p = sortedDictValues(self._FScompatibleNames(pdict))
            i = _pollInputs(sortedDictValues(self.inputs),
                            t, self.checklevel)
            self._ensure_solver({'params': p, 't0': 0, 'tend': 1})
            self._ensure_inputs()
            return self._solver.MassMatrix(t, x, p+i)[0]
        else:
            raise PyDSTool_ExistError("Mass matrix not defined")

    def _ensure_solver(self, pars=None):
        if self._solver is None:
            x0 = sortedDictValues(filteredDict(self.initialconditions, self.funcspec.vars))
#            _integMod = self._ensureLoaded(self.modname)
            self._solver = radau(self.modname,
                                 rhs=self.name, phaseDim=self.dimension,
                                 paramDim=self.numpars,
                                 nAux=len(self.funcspec.auxvars),
                                 nEvents=len(self._eventNames),
                                 nExtInputs=len(self.inputs),
                                 hasJac=self.haveJacobian(),
                                 hasJacP=self.haveJacobian_pars(),
                                 hasMass=self.haveMass(),
                                 extraSpace=self.algparams['extraspace'])
            try:
                genDB.register(self)
            except PyDSTool_KeyError:
                errstr = "Generator " + self.name + ": this vector field's " +\
                         "DLL is already in use"
                raise RuntimeError(errstr)
            if pars is not None:
                # tend value doesn't matter
                self._solver.setRunParams(
                              ic=sortedDictValues(filteredDict(self.initialconditions,
                                                               self.funcspec.vars)),
                              params=pars['params'],
                              t0=pars['t0'], tend=pars['tend'],
                              gt0=self.globalt0,
                              refine=0, specTimes=[])

    def _ensure_inputs(self, force=False):
        if not self.inputs:
            return
        if force:
            listOK = False
        else:
            try:
                listOK = self._inputTimest0 == self.globalt0
            except AttributeError:
                # not yet defined, so proceed
                listOK = False
        if not listOK:
            self._inputVarList = []
            self._inputTimeList = []
            self._inputTimest0 = self.globalt0
            # inputVarList is a list of Variables or Pointsets
            for inp in sortedDictValues(self.inputs):
                if isinstance(inp, Variable):
                    pts = inp.getDataPoints()
                    if pts is None:
                        raise TypeError("Can only pass external input Variable objects if based on"
                                        " an underlying mesh")
                    else:
                        tvals = copy(pts[inp.indepvarname])
                        tvals -= self.globalt0
                    self._inputVarList.append(pts[inp.coordname].tolist())
                    self._inputTimeList.append(tvals.tolist())
                elif isinstance(inp, Pointset):
                    tvals = copy(inp.indepvararray)
                    tvals -= self.globalt0
                    self._inputVarList.append(inp[inp.coordname].tolist())
                    self._inputTimeList.append(tvals.tolist())
                else:
                    raise TypeError("Invalid type of input")
        if not self._solver.initExtInputs:
            self._solver.setExtInputs(True, deepcopy(self._inputVarList),
                                        deepcopy(self._inputTimeList))
        elif not listOK:
            self._solver.clearExtInputs()
            self._solver.setExtInputs(True, deepcopy(self._inputVarList),
                                    deepcopy(self._inputTimeList))
            self._solver.canContinue=True

    def __del__(self):
        genDB.unregister(self)
        ODEsystem.__del__(self)

# Register this Generator with the database

symbolMapDict = {'abs': 'fabs', 'sign': 'signum', 'mod': 'fmod'}
# in future, provide appropriate mappings for libraries math,
# random, etc. (for now it's left to FuncSpec)
theGenSpecHelper.add(Radau_ODEsystem, symbolMapDict, 'c')


""" Starting point for DEM retrieval utilities.
"""
from math import pi, sin, cos
from os import unlink, close
from itertools import product
from tempfile import mkstemp
from sys import modules

import NED10m, NED100m, NED1km, SRTM1, SRTM3, VFP, Worldwide

from ModestMaps.Core import Coordinate
from TileStache.Geography import SphericalMercator
from TileStache.Core import Layer, Metatile
from TileStache.Config import Configuration
from TileStache.Caches import Disk

from osgeo import gdal, osr
from PIL import Image
import numpy

from .. import save_slope_aspect

# used to prevent clobbering in /vsimem/, see:
# http://osgeo-org.1803224.n2.nabble.com/gdal-dev-Outputting-to-vsimem-td6221295.html
vsimem_counter = 1

#
# Set up some useful projections.
#

osr.UseExceptions() # <-- otherwise errors will be silent and useless.

webmerc_proj = SphericalMercator()
webmerc_sref = osr.SpatialReference()
webmerc_sref.ImportFromProj4(webmerc_proj.srs)

class SeedingLayer (Layer):
    """ Tilestache-compatible seeding layer for preparing tiled data.

        Intended for use in hillup-seed.py script for preparing a tile directory.
    """
    def __init__(self, demdir, tiledir, tmpdir, source, size):
        """
        """
        cache = Disk(tiledir, dirs='safe')
        config = Configuration(cache, '.')
        Layer.__init__(self, config, SphericalMercator(), Metatile(), tile_height=size)

        self.provider = Provider(self, demdir, tmpdir, source)

    def name(self):
        return '.'

class Provider:
    """ TileStache provider for generating tiles of DEM slope and aspect data.

        Source parameter can be "srtm-ned" (default) or "ned-only".

        See http://tilestache.org/doc/#custom-providers for information
        on how the Provider object interacts with TileStache.
    """
    def __init__(self, layer, demdir, tmpdir=None, source='srtm-ned'):
        self.tmpdir = tmpdir
        self.demdir = demdir
        self.source = source

    def getTypeByExtension(self, ext):
        if ext.lower() != 'tiff':
            raise Exception()

        return 'image/tiff', 'TIFF'

    def renderArea(self, width, height, srs, xmin, ymin, xmax, ymax, zoom):
        """ Return an instance of SlopeAndAspect for requested area.
        """
        assert srs == webmerc_proj.srs # <-- good enough for now

        if self.source == 'srtm-ned':
            providers = choose_providers_srtm(zoom)

        elif self.source == 'ned-only':
            providers = choose_providers_ned(zoom)

        elif self.source == 'vfp':
            providers = [(VFP, 1)]

        elif self.source == 'worldwide':
            providers = [(Worldwide, 1)]

        else:
            providers = load_func_path(self.source)(zoom)

        assert sum([proportion for (mod, proportion) in providers]) == 1.0

        #
        # Prepare information for datasets of the desired extent and projection.
        #

        xres = (xmax - xmin) / width
        yres = (ymin - ymax) / height

        area_wkt = webmerc_sref.ExportToWkt()
        buffered_xform = xmin - xres, xres, 0, ymax - yres, 0, yres

        #
        # Reproject and merge DEM datasources into destination datasets.
        #

        driver = gdal.GetDriverByName('GTiff')

        composite_ds = make_empty_datasource(width+2, height+2, buffered_xform, area_wkt, self.tmpdir)
        proportion_complete = 0.

        for (module, proportion) in providers:

            cs2cs = osr.CoordinateTransformation(webmerc_sref, module.sref)

            # get a lat/lon bbox buffered by one pixel on all sides
            minlon, minlat, z = cs2cs.TransformPoint(xmin - xres, ymin + yres)
            maxlon, maxlat, z = cs2cs.TransformPoint(xmax + xres, ymax - yres)

            #
            # Keep a version of the composite without the
            # current layer applied for later alpha-blending.
            #
            do_blending = bool(proportion_complete > 0 and proportion < 1)

            if do_blending:
                composite_without = composite_ds.ReadAsArray()

            ds_args = minlon, minlat, maxlon, maxlat, self.demdir

            for ds_dem in module.datasources(*ds_args):

                # estimate the raster density across source DEM and output
                dem_samples = (maxlon - minlon) / ds_dem.GetGeoTransform()[1]
                area_pixels = (xmax - xmin) / composite_ds.GetGeoTransform()[1]

                if dem_samples > area_pixels:
                    # cubic looks better squeezing down
                    resample = gdal.GRA_Cubic
                else:
                    # cubic spline looks better stretching out
                    resample = gdal.GRA_CubicSpline

                gdal.ReprojectImage(ds_dem, composite_ds, ds_dem.GetProjection(), composite_ds.GetProjection(), resample)
                ds_dem = None

            #
            # Perform alpha-blending if needed.
            #
            if do_blending:
                proportion_with = proportion / (proportion_complete + proportion)
                proportion_without = 1 - proportion_with

                composite_with = composite_ds.ReadAsArray() * proportion_with
                composite_with += composite_without * proportion_without

                composite_ds.GetRasterBand(1).WriteArray(composite_with, 0, 0)

            proportion_complete += proportion

        elevation = composite_ds.ReadAsArray()

        unlink(composite_ds.GetFileList()[0])
        composite_ds = None

        #
        # Calculate and save slope and aspect.
        #

        slope, aspect = calculate_slope_aspect(elevation, xres, yres)

        tile_xform = xmin, xres, 0, ymax, 0, yres

        return SlopeAndAspect(self.tmpdir, slope, aspect, area_wkt, tile_xform)

class SlopeAndAspect:
    """ TileStache response object with PIL-like save() and crop() methods.

        This object knows only how to save two-band 8-bit GeoTIFFs.

        See http://tilestache.org/doc/#custom-providers for information
        on how the SlopeAndAspect object interacts with TileStache.
    """
    def __init__(self, tmpdir, slope, aspect, wkt, xform):
        """ Instantiate with array of slope and aspect, and minimal geographic information.
        """
        self.tmpdir = tmpdir

        self.slope = slope
        self.aspect = aspect

        self.w, self.h = self.slope.shape

        self.wkt = wkt
        self.xform = xform

    def save(self, output, format):
        """ Save a two-band GeoTIFF to output file-like object.
        """
        if format != 'TIFF':
            raise Exception('File format other than TIFF for slope and aspect: "%s"' % format)

        save_slope_aspect(self.slope, self.aspect, self.wkt, self.xform, output, self.tmpdir)

    def crop(self, box):
        """ Returns a rectangular region from the current image.

            Box is a 4-tuple with left, upper, right, and lower pixels.
            Not yet implemented!
        """
        raise NotImplementedError()

def choose_providers_srtm(zoom):
    """ Return a list of data sources and proportions for given zoom level.

        Each data source is a module such as SRTM1 or SRTM3, and the proportions
        must all add up to one. Return list has either one or two items.
    """
    if zoom <= SRTM3.ideal_zoom:
        return [(SRTM3, 1)]

    elif SRTM3.ideal_zoom < zoom and zoom < SRTM1.ideal_zoom:
        #bottom, top = SRTM3, SRTM1 # SRTM1 looks terrible
        bottom, top = SRTM3, NED10m

    elif zoom == SRTM1.ideal_zoom:
        #return [(SRTM1, 1)] # SRTM1 looks terrible
        bottom, top = SRTM3, NED10m

    elif SRTM1.ideal_zoom < zoom and zoom < NED10m.ideal_zoom:
        #bottom, top = SRTM1, NED10m # SRTM1 looks terrible
        bottom, top = SRTM3, NED10m

    elif zoom >= NED10m.ideal_zoom:
        return [(NED10m, 1)]

    difference = float(top.ideal_zoom) - float(bottom.ideal_zoom)
    proportion = 1. - (zoom - float(bottom.ideal_zoom)) / difference

    return [(bottom, proportion), (top, 1 - proportion)]

def choose_providers_ned(zoom):
    """ Return a list of data sources and proportions for given zoom level.

        Each data source is a module such as NED10m or NED1km, and the proportions
        must all add up to one. Return list has either one or two items.
    """
    if zoom <= NED1km.ideal_zoom:
        return [(NED1km, 1)]

    elif NED1km.ideal_zoom < zoom and zoom < NED100m.ideal_zoom:
        #bottom, top = NED1km, NED100m
        bottom, top = NED1km, NED100m

    elif zoom == NED100m.ideal_zoom:
        return [(NED100m, 1)]

    elif NED100m.ideal_zoom < zoom and zoom < NED10m.ideal_zoom:
        #bottom, top = NED100m, NED10m
        bottom, top = NED100m, NED10m

    elif zoom >= NED10m.ideal_zoom:
        return [(NED10m, 1)]

    difference = float(top.ideal_zoom) - float(bottom.ideal_zoom)
    proportion = 1. - (zoom - float(bottom.ideal_zoom)) / difference

    return [(bottom, proportion), (top, 1 - proportion)]

def make_empty_datasource(width, height, xform, wkt, tmpdir):
    '''
    '''
    driver = gdal.GetDriverByName('GTiff')
    handle, filename = mkstemp(dir=tmpdir, prefix='dem-tools-hillup-data-render-', suffix='.tif')
    close(handle)

    ds = driver.Create(filename, width, height, 1, gdal.GDT_Float32)
    ds.SetGeoTransform(xform)
    ds.SetProjection(wkt)

    ds.GetRasterBand(1).WriteArray(numpy.ones((width, height), numpy.float32) * -9999, 0, 0)
    ds.GetRasterBand(1).SetNoDataValue(-9999)

    return ds

def calculate_slope_aspect(elevation, xres, yres, z=1.0):
    """ Return a pair of arrays 2 pixels smaller than the input elevation array.

        Slope is returned in radians, from 0 for sheer face to pi/2 for
        flat ground. Aspect is returned in radians, counterclockwise from -pi
        at north around to pi.

        Logic here is borrowed from hillshade.cpp:
          http://www.perrygeo.net/wordpress/?p=7
    """
    width, height = elevation.shape[0] - 2, elevation.shape[1] - 2

    window = [z * elevation[row:(row + height), col:(col + width)]
              for (row, col)
              in product(range(3), range(3))]

    x = ((window[0] + window[3] + window[3] + window[6]) \
       - (window[2] + window[5] + window[5] + window[8])) \
      / (8.0 * xres);

    y = ((window[6] + window[7] + window[7] + window[8]) \
       - (window[0] + window[1] + window[1] + window[2])) \
      / (8.0 * yres);

    # in radians, from 0 to pi/2
    slope = pi/2 - numpy.arctan(numpy.sqrt(x*x + y*y))

    # in radians counterclockwise, from -pi at north back to pi
    aspect = numpy.arctan2(x, y)

    return slope, aspect

def load_func_path(funcpath):
    """ Load external function based on a path.

        Example funcpath: "Module.Submodule:Function".
    """
    modname, objname = funcpath.split(':', 1)

    __import__(modname)
    module = modules[modname]
    _func = eval(objname, module.__dict__)

    if _func is None:
        raise Exception('eval(%(objname)s) in %(modname)s came up None' % locals())

    return _func


#!/usr/bin/env python
"""Random sequences and random evolution of sequences in a tree"""

import bisect

import numpy

__author__ = "Peter Maxwell"
__copyright__ = "Copyright 2007-2020, The Cogent Project"
__credits__ = ["Peter Maxwell"]
__license__ = "BSD-3"
__version__ = "2020.6.30a"
__maintainer__ = "Peter Maxwell"
__email__ = "pm67nz@gmail.com"
__status__ = "Production"

def argpicks(freqs, random_series):
    partition = numpy.add.accumulate(freqs)
    assert abs(partition[-1] - 1.0) < 1e-6, (freqs, partition)
    while True:
        x = random_series.uniform(0.0, 1.0)
        i = bisect.bisect_left(partition, x)
        yield i

def argpick(freqs, random_series):
    return next(argpicks(freqs, random_series))

def _randomMotifGenerator(random_series, motif_probs):
    motifs = list(motif_probs.keys())
    freqs = [motif_probs[m] for m in motifs]
    for i in argpicks(freqs, random_series):
        yield motifs[i]

def evolve_sequence(
    random_series, motifs, parent_seq, site_cats, psubs, preserved_sites=()
):
    """Evolve a new sequence derived from parent_seq.  Uses psubs[site_cats[i]]
    to pick a new motif derived from parent_seq[i]"""
    seq = []
    randomMotifSources = {}
    for (i, parent_motif) in enumerate(parent_seq):
        if i in preserved_sites:
            edge_motif = preserved_sites[i]
        else:
            if parent_motif not in randomMotifSources:
                mprobs = {}
                parent_motif_index = motifs.index(parent_motif)
                site_cat = site_cats[i]
                psub = psubs[site_cat]
                for (dest_motif_index, dest_motif) in enumerate(motifs):
                    prob = psub[parent_motif_index, dest_motif_index]
                    mprobs[dest_motif] = prob
                randomMotifSources[site_cat, parent_motif] = _randomMotifGenerator(
                    random_series, mprobs
                )
            edge_motif = next(randomMotifSources[site_cat, parent_motif])
        seq.append(edge_motif)
    return seq

def random_sequence(random_series, motif_probs, sequence_length):
    getRootRandomMotif = _randomMotifGenerator(random_series, motif_probs).__next__
    return [getRootRandomMotif() for i in range(sequence_length)]

class AlignmentEvolver(object):
    # Encapsulates settings that are constant throughout the recursive generation
    # of a synthetic alignment.

    def __init__(
        self,
        random_series,
        orig_ambig,
        exclude_internal,
        bin_names,
        site_bins,
        psub_for,
        motifs,
    ):
        self.random_series = random_series
        self.orig_ambig = orig_ambig
        self.exclude_internal = exclude_internal
        self.bin_names = bin_names
        self.site_bins = site_bins
        self.psub_for = psub_for
        self.motifs = motifs

    def __call__(self, tree, root_sequence):
        # probsd = dict(enumerate(self.bin_probs))
        # bprobs = _randomMotifGenerator(self.random_series, probsd)
        # site_bins = [bprobs.next() for c in range(len(root_sequence))]
        return self.generate_simulated_seqs(tree, root_sequence)

    def generate_simulated_seqs(self, parent, parent_seq):
        """recursively generate the descendant sequences by descending the tree
        from root.
        Each child will be set by mutating the parent motif based on the probs
        in the psub matrix of this edge.

        random_series - get a random numer 0-1 by calling random_series.random()
        length - the desired alignment length
        parent - the edge structure.
        parent_seq - the corresponding sequence. This will be mutated for each
        of its children, based on their psub matricies.
        """

        # This depends on parameter names 'mprobs', 'alignment2', 'bprobs' and
        # 'psubs'.  Might be better to integrate it into likelihood_calculation.

        if self.exclude_internal and parent.children:
            simulated_sequences = {}
        else:
            simulated_sequences = {parent.name: "".join(parent_seq)}

        for edge in parent.children:
            # The result for this edge - a list of motifs

            # Keep original ambiguity codes
            if edge.name in self.orig_ambig:
                orig_seq_ambig = self.orig_ambig[edge.name]
            else:
                orig_seq_ambig = {}

            # Matrix of substitution probabilities
            psubs = [self.psub_for(edge.name, bin) for bin in self.bin_names]

            # Make the semi-random sequence for this edge.
            edge_seq = evolve_sequence(
                self.random_series,
                self.motifs,
                parent_seq,
                self.site_bins,
                psubs,
                orig_seq_ambig,
            )

            # Pass this new edge sequence on down the tree
            descendant_sequences = self.generate_simulated_seqs(edge, edge_seq)
            simulated_sequences.update(descendant_sequences)

        return simulated_sequences


import xarray as xr
import pandas as pd
import numpy as np
import os
import numba as nb
import time
import dask.array as da

def compute_all(work_dir, memory_to_allocate_gb, date_strings):
    array_size_bytes = 3060  # Based on 15 x 51 member array
    memory_to_allocate_bytes = memory_to_allocate_gb * 1e9

    files = [os.path.join(work_dir, i + ".nc") for i in date_strings]

    chunk_size = int(np.floor(memory_to_allocate_bytes / ((array_size_bytes * len(files)) + len(files))))

    print("Chunk Size:", chunk_size)

    list_of_dask_q_arrays = []
    list_of_dask_init_arrays = []

    # Creating a large dask array with all of the data in it
    start = time.time()
    for file in files:
        ds = xr.open_dataset(file, chunks={"rivid": chunk_size})

        tmp_dask_q_array = ds["Qout"].data
        list_of_dask_q_arrays.append(tmp_dask_q_array)

        tmp_dask_init_array = ds["initialization_values"].data
        list_of_dask_init_arrays.append(tmp_dask_init_array)

        ds.close()
    end = time.time()

    big_dask_q_array = da.stack(list_of_dask_q_arrays)
    big_dask_init_array = da.stack(list_of_dask_init_arrays)

    print(big_dask_q_array.shape)
    print(big_dask_init_array.shape)

    print("Time to create dask arrays: ", end - start)

    # Retrieving the number of streams and their corresponding Rivids
    tmp_dataset = xr.open_dataset(files[0])

    num_of_streams = tmp_dataset['rivid'].size
    rivids = tmp_dataset['rivid'].data

    tmp_dataset.close()

    num_chunk_iterations = int(np.ceil(num_of_streams / chunk_size))
    start_chunk = 0
    end_chunk = chunk_size
    list_of_tuples_with_metrics = []

    for chunk_number in range(num_chunk_iterations):

        start = time.time()
        big_forecast_data_array = np.asarray(big_dask_q_array[:, start_chunk:end_chunk, :, :])
        big_init_data_array = np.asarray(big_dask_init_array[:, start_chunk:end_chunk])
        end = time.time()
        print("Time to read from disk:", end - start)

        rivids_chunk = rivids[start_chunk:end_chunk]

        start = time.time()
        results_array = numba_calculate_metrics(
            big_forecast_data_array, big_init_data_array, len(files), big_forecast_data_array.shape[1], 15
        )
        end = time.time()
        print("Numba Calculation Time: ", end - start)

        for rivid in range(results_array.shape[1]):
            for forecast_day in range(results_array.shape[0]):
                tmp_array = results_array[forecast_day, rivid, :]
                tuple_to_append = (rivids_chunk[rivid], '{} Day Forecast'.format(str(forecast_day + 1).zfill(2)),
                                   tmp_array[0], tmp_array[1], tmp_array[2])
                list_of_tuples_with_metrics.append(tuple_to_append)

        start_chunk += chunk_size
        end_chunk += chunk_size

    final_df = pd.DataFrame(list_of_tuples_with_metrics,
                            columns=['Rivid', 'Forecast Day', 'CRPS', 'CRPS BENCH', 'CRPSS'])
    final_df.to_csv(r'/Users/wade/PycharmProjects/Forecast_Validation/South_America_Test_DF.csv', index=False)

@nb.njit(parallel=True)
def numba_calculate_metrics(forecast_array, initialization_array, number_of_start_dates, number_of_streams,
                            num_forecast_days):
    """
    Parameters
    ----------

    forecast_array: 4D ndarray
        A 4 dimensional numPy array with the following dimensions: 1) Start Date Number (365 if there are a year's
        forecasts), 2) Unique stream ID, 3) Forecast Days (e.g. 1-15 in a 15 day forecast), 4) Ensembles

    initialization_array: 2D ndarray
        A 2 dimenensional NumPy array with the following dimensions: 1) Start Dates, 2) Unique stream ID

    number_of_start_dates: The number of start dates for the analysis to perform

    number_of_streams:
        The number of streams in the analysis

    num_forecast_days:
        The number of forecast days in the analysis

    Returns
    -------
    ndarray
        An ndarray with the folowing dimenstions:
        1) Forecast Day: 1-15 in the case of a 15 day forecast
        2) Rivid: The stream unique ID
        3) Metrics: CRPS, CRPS_BENCH, CRPSS

    """
    return_array = np.zeros((num_forecast_days, number_of_streams, 3), dtype=np.float32)

    for stream in nb.prange(number_of_streams):
        for forecast_day in range(num_forecast_days):
            initialization_vals = initialization_array[(forecast_day + 1):, stream]
            # np.savetxt("init_test.txt", initialization_vals)
            forecasts = forecast_array[:(number_of_start_dates - (forecast_day + 1)), stream, forecast_day, :]
            # np.savetxt("forecasts_test.txt", forecasts)
            benchmark_forecasts = initialization_array[:(number_of_start_dates - (forecast_day + 1)), stream]
            # np.savetxt("benchmark_forecasts_test.txt", benchmark_forecasts)

            crps = ens_crps(initialization_vals, forecasts)
            crps_bench = mae(initialization_vals, benchmark_forecasts)
            if crps_bench == 0:
                crpss = np.inf
                print("Warning: Division by zero on: ", stream)
            else:
                crpss = 1 - crps / crps_bench

            return_array[forecast_day, stream, 0] = crps
            return_array[forecast_day, stream, 1] = crps_bench
            return_array[forecast_day, stream, 2] = crpss

            # print(crps, crps_bench, crpss)
        if (stream % 1000) == 0:
            print("Count: ", stream)

    return return_array

@nb.njit()
def mae(sim, obs):
    return np.mean(np.abs(sim - obs))

@nb.njit()
def ens_crps(obs, fcst_ens, adj=np.nan):

    rows = obs.size
    cols = fcst_ens.shape[1]

    col_len_array = np.ones(rows) * cols
    sad_ens_half = np.zeros(rows)
    sad_obs = np.zeros(rows)
    crps = np.zeros(rows)

    crps = numba_crps(
        fcst_ens, obs, rows, cols, col_len_array, sad_ens_half, sad_obs, crps, np.float64(adj)
    )

    # Calc mean crps as simple mean across crps[i]
    crps_mean = np.mean(crps)

    return crps_mean

@nb.njit()
def numba_crps(ens, obs, rows, cols, col_len_array, sad_ens_half, sad_obs, crps, adj):
    for i in range(rows):
        the_obs = obs[i]
        the_ens = ens[i, :]
        the_ens = np.sort(the_ens)
        sum_xj = 0.
        sum_jxj = 0.

        j = 0
        while j < cols:
            sad_obs[i] += np.abs(the_ens[j] - the_obs)
            sum_xj += the_ens[j]
            sum_jxj += (j + 1) * the_ens[j]
            j += 1

        sad_ens_half[i] = 2.0 * sum_jxj - (col_len_array[i] + 1) * sum_xj

    if np.isnan(adj):
        for i in range(rows):
            crps[i] = sad_obs[i] / col_len_array[i] - sad_ens_half[i] / \
                      (col_len_array[i] * col_len_array[i])
    elif adj > 1:
        for i in range(rows):
            crps[i] = sad_obs[i] / col_len_array[i] - sad_ens_half[i] / \
                      (col_len_array[i] * (col_len_array[i] - 1)) * (1 - 1 / adj)
    elif adj == 1:
        for i in range(rows):
            crps[i] = sad_obs[i] / col_len_array[i]
    else:
        for i in range(rows):
            crps[i] = np.nan

    return crps

if __name__ == "__main__":

    starting_date = "2018-08-19"
    ending_date = "2018-12-16"

    dates_range = pd.date_range(starting_date, ending_date)
    dates_strings = dates_range.strftime("%Y%m%d").tolist()

    workspace = r'/Users/wade/Documents/South_America_Forecasts'
    MEMORY_TO_ALLOCATE = 1.0  # GB

    start = time.time()
    compute_all(workspace, MEMORY_TO_ALLOCATE, dates_strings)
    end = time.time()
    print(end - start)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy
from PIL import Image, ImageFile
__version__ = "1.0"

"""Visualization of 2-dimensional numpy arrays."""

def array2im(a, brightness=1, contrast=1):
    P = (a.real * (a.real > 0) * brightness) ** contrast
    N = (-a.real * (a.real < 0) * brightness) ** contrast
    R = numpy.array(
        (256 * (1 - numpy.float64(0.5) ** P)).clip(max=255), dtype=numpy.uint8)[::-1]
    G = numpy.zeros(R.shape, dtype=numpy.uint8)
    B = numpy.array(
        (256 * (1 - numpy.float64(0.5) ** N)).clip(max=255), dtype=numpy.uint8)[::-1]
    return Image.fromarray(numpy.array([R, G, B]).swapaxes(0, 1).swapaxes(1, 2))

if __name__ == "__main__":
    import argparse
    import struct
    import os
    import sys
    import time
    from numpy import dtype, fromfile
    from numpy import save as npsave
    from numpy import load as npload
    import warnings
    warnings.filterwarnings('ignore')

    parser = argparse.ArgumentParser(
        description="Converts a two-dimensional numpy array into an image.")
    parser.add_argument("--in", "-i", dest="input", action="store",
                        help="Input file. Default: <stdin>.", default="-")
    parser.add_argument("--out", "-o", dest="output", action="store",
                        help="Output file. Default: <stdout>.", default="-")
    parser.add_argument("--brightness", "-b", dest="brightness",
                        action="store", help="Adjust brightness. Default: 1.", default=1)
    parser.add_argument("--contrast", "-c", dest="contrast",
                        action="store", help="Adjust contrast. Default: 1.", default=1)
    parser.add_argument("--quality", "-q", dest="quality", action="store",
                        help="JPEG quality, between 1 and 100. Only valid when saving image as a JPEG. Default: 95.", default=95)
    parser.add_argument("--progressive", "-p", dest="progressive", action="store_const",
                        help="Save progressive JPEG. Only valid when saving image as a JPEG.", const=True)
    parser.add_argument("--overwrite", "-y", dest="overwrite",
                        action='store_const', help="Overwrite output file, if it exists.", const=True)
    parser.add_argument("--no-overwrite", "-n", dest="nooverwrite", action='store_const',
                        help="Do not overwrite output file, if it exists.", const=True)
    args = parser.parse_args()

    msgout = sys.stdout if sys.stdout.isatty() else sys.stderr

    if args.overwrite and args.nooverwrite:
        print >>sys.stderr, "Cannot specify both '-y' and '-n'!"
        sys.exit()

    if args.input == "-":
        if sys.stdin.isatty():
            print >>sys.stderr, "Surely you are not typing the raw data into the terminal. Please specify input file, or pipe input from another program.\n"
            parser.print_help(sys.stderr)
            sys.exit()
        infile = sys.stdin
    else:
        infile = open(args.input, "rb")

    if args.output == "-":
        if sys.stdout.isatty():
            print >>sys.stderr, "Cowardly refusing to write binary data to terminal. Please specify output file, or redirect output to a pipe.\n"
            parser.print_help(sys.stderr)
            sys.exit()
        outfile = sys.stdout
        print >>sys.stderr, "Writing data to <stdout>."
    else:
        if os.path.exists(args.output):
            if args.nooverwrite:
                print >>sys.stderr, "Error: Output file '%s' exists. Terminating because '-n' was specified." % args.output
                sys.exit()
            elif args.overwrite:
                print >>msgout, "Warning: Output file '%s' exists. Overwriting because '-y' was specified." % args.output
            elif sys.stdin.isatty() and sys.stdout.isatty():
                overwrite = raw_input(
                    "Warning: Output file '%s' exists. Do you wish to overwrite file? (Y/N) " % args.output)
                while overwrite.upper() not in ("Y", "N", "YES", "NO"):
                    print >>msgout, "Invalid answer: '%s'" % overwrite
                    overwrite = raw_input(
                        "Warning: Output file '%s' exists. Do you wish to overwrite file? (Y/N) " % args.output)

                if overwrite.upper() in ("Y", "YES"):
                    print >>msgout, "Overwriting '%s'." % args.output
                    print >>msgout, ""
                elif overwrite.upper() in ("N", "NO"):
                    print >>msgout, "Operation aborted."
                    sys.exit()
            else:
                print >>sys.stderr, "Operation aborted. Cowardly refusing to overwrite '%s'." % args.output
                sys.exit()

        outfile = args.output

    tag = npload(infile)
    metadata = npload(infile)
    data = npload(infile)
    if infile is not sys.stdin:
        infile.close()

    if len(data.shape) != 2:
        print >>sys.stderr, "Expected a two-dimensional array. Got an array with shape %s instead." % (
            data.shape,)
        sys.exit()

    kind_str = {
        "f": "floating point",
        "i": "integer",
        "u": "unsigned integer",
        "b": "boolean",
        "c": "complex"
    }

    if data.dtype.kind in "iub":
        print >>sys.stderr, "Error: Data type '%s' not supported." % kind_str[
            data.dtype.kind]
        sys.exit()
    elif data.dtype.kind == "c":
        print >>msgout, "Warning: Discarding imaginary part of 'complex' array."

    try:
        args.brightness = float(args.brightness)
    except ValueError:
        print >>sys.stderr, "Bad parameter for brightness. Expected floating point or integer, got '%s' instead." % args.brightness
        sys.exit()

    try:
        args.contrast = float(args.contrast)
    except ValueError:
        print >>sys.stderr, "Bad parameter for contrast. Expected floating point or integer, got '%s' instead." % args.contrast
        sys.exit()

    try:
        args.quality = int(args.quality)
    except ValueError:
        print >>sys.stderr, "Bad parameter for quality. Expected integer between 1 and 100, got '%s' instead." % args.quality
        sys.exit()

    if not (1 <= args.quality <= 100):
        print >>sys.stderr, "Bad parameter for quality. Expected integer between 1 and 100, got '%s' instead." % args.quality
        sys.exit()

    H, W = data.shape

    t0 = time.time()
    print >>msgout, u"Saving %d×%d image to '%s'..." % (W, H, args.output),
    msgout.flush()
    img = array2im(data, brightness=args.brightness, contrast=args.contrast)
    ImageFile.MAXBLOCK = img.size[0] * img.size[1]
    img.save(outfile, progressive=args.progressive, quality=args.quality)
    print >>msgout, "%.2f seconds" % (time.time() - t0)


''' FSL IO '''

from __future__ import with_statement

import os
from os.path import join as pjoin
from subprocess import Popen,PIPE

import numpy as np
import numpy.linalg as npl
from numpy import newaxis

from scipy.ndimage import map_coordinates as mc
from scipy.ndimage import affine_transform
from dipy.io.dpy import Dpy

import nibabel as nib
from nibabel.tmpdirs import InTemporaryDirectory

_VAL_FMT = '   %e'

class FSLError(Exception):
    """ Class signals error in FSL processing """

def have_flirt():
    """ Return True if we can call flirt without error

    Relies on the fact that flirt produces text on stdout when called with no
    arguments
    """
    p = Popen('flirt', stdout=PIPE, stderr=PIPE, shell=True)
    stdout, stderr = p.communicate()
    return stdout != ''

def write_bvals_bvecs(bvals, bvecs, outpath=None, prefix=''):
    ''' Write FSL FDT bvals and bvecs files

    Parameters
    -------------
    bvals : (N,) sequence
       Vector with diffusion gradient strength (one per diffusion
       acquisition, N=no of acquisitions)
    bvecs : (N, 3) array-like
       diffusion gradient directions
    outpath : None or str
       path to write FDT bvals, bvecs text files
       None results in current working directory.
    prefix : str
       prefix for bvals, bvecs files in directory.  Defaults to ''
    '''
    if outpath is None:
        outpath = os.getcwd()
    bvals = tuple(bvals)
    bvecs = np.asarray(bvecs)
    bvecs[np.isnan(bvecs)] = 0
    N = len(bvals)
    fname = pjoin(outpath, prefix + 'bvals')
    fmt = _VAL_FMT * N + '\n'
    open(fname, 'wt').write(fmt % bvals)
    fname = pjoin(outpath, prefix + 'bvecs')
    bvf = open(fname, 'wt')
    for dim_vals in bvecs.T:
        bvf.write(fmt % tuple(dim_vals))

def flirt2aff(mat, in_img, ref_img):
    """ Transform from `in_img` voxels to `ref_img` voxels given `mat`

    Parameters
    ----------
    mat : (4,4) array
        contents (as array) of output ``-omat`` transformation file from flirt
    in_img : img
        image passed (as filename) to flirt as ``-in`` image
    ref_img : img
        image passed (as filename) to flirt as ``-ref`` image

    Returns
    -------
    aff : (4,4) array
        Transform from voxel coordinates in ``in_img`` to voxel coordinates in
        ``ref_img``

    Notes
    -----
    Thanks to Mark Jenkinson and Jesper Andersson for the correct statements
    here, apologies for any errors we've added.

    ``flirt`` registers an ``in`` image to a ``ref`` image.  It can produce
    (with the ``-omat`` option) - a 4 x 4 affine matrix giving the mapping from
    *inspace* to *refspace*.

    The rest of this note is to specify what *inspace* and *refspace* are.

    In what follows, a *voxtrans* for an image is the 4 by 4 affine
    ``np.diag([vox_i, vox_j, vox_k, 1])`` where ``vox_i`` etc are the voxel
    sizes for the first second and third voxel dimension.  ``vox_i`` etc are
    always positive.

    If the input image has an affine with a negative determinant, then the
    mapping from voxel coordinates in the input image to *inspace* is simply
    *voxtrans* for the input image.  If the reference image has a negative
    determinant, the mapping from voxel space in the reference image to
    *refspace* is simply *voxtrans* for the reference image.

    A negative determinant for the image affine is the common case, of an image
    with a x voxel flip.  Analyze images don't store affines and flirt assumes a
    negative determinant in these cases.

    For positive determinant affines, flirt starts *inspace* and / or *refspace*
    with an x voxel flip.  The mapping implied for an x voxel flip for image
    with shape (N_i, N_j, N_k) is:

        [[-1, 0, 0, N_i - 1],
         [ 0, 1, 0,       0],
         [ 0, 0, 1,       0],
         [ 0, 0, 0,       1]]

    If the input image has an affine with a positive determinant, then mapping
    from input image voxel coordinates to *inspace* is ``np.dot(input_voxtrans,
    input_x_flip)`` - where ``input_x_flip`` is the matrix above with ``N_i``
    given by the input image first axis length.  Similarly the mapping from
    reference voxel coordinates to *refspace*, if the reference image has a
    positive determinant, is ``np.dot(ref_voxtrans, ref_x_flip)`` - where
    ``ref_x_flip`` is the matrix above with ``N_i`` given by the reference image
    first axis length.
    """
    in_hdr = in_img.get_header()
    ref_hdr = ref_img.get_header()
    # get_zooms gets the positive voxel sizes as returned in the header
    inspace = np.diag(in_hdr.get_zooms() + (1,))
    refspace = np.diag(ref_hdr.get_zooms() + (1,))
    if npl.det(in_img.get_affine())>=0:
        inspace = np.dot(inspace, _x_flipper(in_hdr.get_data_shape()[0]))
    if npl.det(ref_img.get_affine())>=0:
        refspace = np.dot(refspace, _x_flipper(ref_hdr.get_data_shape()[0]))
    # Return voxel to voxel mapping
    return np.dot(npl.inv(refspace), np.dot(mat, inspace))

def _x_flipper(N_i):
    flipr = np.diag([-1, 1, 1, 1])
    flipr[0,3] = N_i - 1
    return flipr

def flirt2aff_files(matfile, in_fname, ref_fname):
    """ Map from `in_fname` image voxels to `ref_fname` voxels given `matfile`

    See :func:`flirt2aff` docstring for details.

    Parameters
    ------------
    matfile : str
        filename of output ``-omat`` transformation file from flirt
    in_fname : str
        filename for image passed to flirt as ``-in`` image
    ref_fname : str
        filename for image passed to flirt as ``-ref`` image

    Returns
    -------
    aff : (4,4) array
        Transform from voxel coordinates in image for ``in_fname`` to voxel
        coordinates in image for ``ref_fname``
    """
    mat = np.loadtxt(matfile)
    in_img = nib.load(in_fname)
    ref_img = nib.load(ref_fname)
    return flirt2aff(mat, in_img, ref_img)

def warp_displacements(ffa,flaff,fdis,fref,ffaw,order=1):
    ''' Warp an image using fsl displacements

    Parameters
    ------------
    ffa : filename of nifti to be warped
    flaff : filename of .mat  (flirt)
    fdis :  filename of displacements (fnirtfileutils)
    fref : filename of reference volume e.g. (FMRIB58_FA_1mm.nii.gz)
    ffaw : filename for the output warped image
    '''
    refaff=nib.load(fref).get_affine()
    disdata=nib.load(fdis).get_data()
    imgfa=nib.load(ffa)
    fadata=imgfa.get_data()
    fazooms=imgfa.get_header().get_zooms()
    #from fa index to ref index
    res=flirt2aff_files(flaff,ffa,fref)
    #from ref index to fa index
    ires=np.linalg.inv(res)
    #create the 4d volume which has the indices for the reference image
    reftmp=np.zeros(disdata.shape)
    '''
    #create the grid indices for the reference
    #refinds = np.ndindex(disdata.shape[:3])
    for ijk_t in refinds:
        i,j,k = ijk_t
        reftmp[i,j,k,0]=i
        reftmp[i,j,k,1]=j
        reftmp[i,j,k,2]=k
    '''
    #same as commented above but much faster
    reftmp[...,0] = np.arange(disdata.shape[0])[:,newaxis,newaxis]
    reftmp[...,1] = np.arange(disdata.shape[1])[newaxis,:,newaxis]
    reftmp[...,2] = np.arange(disdata.shape[2])[newaxis,newaxis,:]

    #affine transform from reference index to the fa index
    A = np.dot(reftmp,ires[:3,:3].T)+ires[:3,3]
    #add the displacements but first devide them by the voxel sizes
    A2=A+disdata/fazooms
    #hold the displacements' shape reshaping
    di,dj,dk,dl=disdata.shape
    #do the interpolation using map coordinates
    #the list of points where the interpolation is done given by the reshaped in 2D A2 (list of 3d points in fa index)
    W=mc(fadata,A2.reshape(di*dj*dk,dl).T,order=order).reshape(di,dj,dk)
    #save the warped image
    Wimg=nib.Nifti1Image(W,refaff)
    nib.save(Wimg,ffaw)

def warp_displacements_tracks(fdpy,ffa,fmat,finv,fdis,fdisa,fref,fdpyw):
    """ Warp tracks from native space to the FMRIB58/MNI space

    We use here the fsl displacements. Have a look at create_displacements to
    see an example of how to use these displacements.

    Parameters
    ------------
    fdpy : filename of the .dpy file with the tractography
    ffa : filename of nifti to be warped
    fmat : filename of .mat  (flirt)
    fdis :  filename of displacements (fnirtfileutils)
    fdisa :  filename of displacements (fnirtfileutils + affine)
    finv : filename of invwarp displacements (invwarp)
    fref : filename of reference volume e.g. (FMRIB58_FA_1mm.nii.gz)
    fdpyw : filename of the warped tractography

    See also
    -----------
    dipy.external.fsl.create_displacements

    """

    #read the tracks from the image space
    dpr=Dpy(fdpy,'r')
    T=dpr.read_tracks()
    dpr.close()

    #copy them in a new file
    dpw=Dpy(fdpyw,'w',compression=1)
    dpw.write_tracks(T)
    dpw.close()

    #from fa index to ref index
    res=flirt2aff_files(fmat,ffa,fref)

    #load the reference img
    imgref=nib.load(fref)
    refaff=imgref.get_affine()

    #load the invwarp displacements
    imginvw=nib.load(finv)
    invwdata=imginvw.get_data()
    invwaff = imginvw.get_affine()

    #load the forward displacements
    imgdis=nib.load(fdis)
    disdata=imgdis.get_data()

    #load the forward displacements + affine
    imgdis2=nib.load(fdisa)
    disdata2=imgdis2.get_data()

    #from their difference create the affine
    disaff=disdata2-disdata

    del disdata
    del disdata2

    shape=nib.load(ffa).get_data().shape

    #transform the displacements affine back to image space
    disaff0=affine_transform(disaff[...,0],res[:3,:3],res[:3,3],shape,order=1)
    disaff1=affine_transform(disaff[...,1],res[:3,:3],res[:3,3],shape,order=1)
    disaff2=affine_transform(disaff[...,2],res[:3,:3],res[:3,3],shape,order=1)

    #remove the transformed affine from the invwarp displacements
    di=invwdata[:,:,:,0] + disaff0
    dj=invwdata[:,:,:,1] + disaff1
    dk=invwdata[:,:,:,2] + disaff2

    dprw=Dpy(fdpyw,'r+')
    rows=len(dprw.f.root.streamlines.tracks)
    blocks=np.round(np.linspace(0,rows,10)).astype(int)#lets work in blocks
    #print rows
    for i in range(len(blocks)-1):
        #print blocks[i],blocks[i+1]
        #copy a lot of tracks together
        caboodle=dprw.f.root.streamlines.tracks[blocks[i]:blocks[i+1]]
        mci=mc(di,caboodle.T,order=1) #interpolations for i displacement
        mcj=mc(dj,caboodle.T,order=1) #interpolations for j displacement
        mck=mc(dk,caboodle.T,order=1) #interpolations for k displacement
        D=np.vstack((mci,mcj,mck)).T
        #go back to mni image space
        WI2=np.dot(caboodle,res[:3,:3].T)+res[:3,3]+D
        #and then to mni world space
        caboodlew=np.dot(WI2,refaff[:3,:3].T)+refaff[:3,3]
        #write back
        dprw.f.root.streamlines.tracks[blocks[i]:blocks[i+1]]=caboodlew.astype('f4')
    dprw.close()

def pipe(cmd):
    """ A tine pipeline system to run external tools.

    For more advanced pipelining use nipype http://www.nipy.org/nipype
    """
    p = Popen(cmd, shell=True,stdout=PIPE,stderr=PIPE)
    sto=p.stdout.readlines()
    ste=p.stderr.readlines()
    print(sto)
    print(ste)

def dcm2nii(dname,outdir,filt='*.dcm',options='-d n -g n -i n -o'):
    cmd='dcm2nii '+options +' ' + outdir +' ' + dname + '/' + filt
    print(cmd)
    pipe(cmd)

def eddy_correct(in_nii,out_nii,ref=0):
    cmd='eddy_correct '+in_nii+' '+ out_nii + ' '+str(ref)
    print(cmd)
    pipe(cmd)

def bet(in_nii,out_nii,options=' -F -f .2 -g 0'):
    cmd='bet '+in_nii+' '+ out_nii + options
    print(cmd)
    pipe(cmd)

def run_flirt_imgs(in_img, ref_img, dof=6, flags=''):
    """ Run flirt on nibabel images, returning affine

    Parameters
    ----------
    in_img : `SpatialImage'
        image to register
    ref_img : `SpatialImage`
        image to register to
    dof : int, optional
        degrees of freedom for registration (default 6)
    flags : str, optional
        other flags to pass to flirt command string

    Returns
    -------
    in_vox2out_vox : (4,4) ndarray
        affine such that, if [i, j, k] is a coordinate in voxels in the
        `in_img`, and [p, q, r] are the equivalent voxel coordinates in the
        reference image, then [p, q, r] = np.dot(in_vox2out_vox[:3,:3]), [i, j,
        k] + in_vox2out_vox[:3,3])
    """
    omat = 'reg.mat'
    with InTemporaryDirectory():
        nib.save(in_img, 'in.nii')
        nib.save(ref_img, 'ref.nii')
        cmd = 'flirt %s -dof %d -in in.nii -ref ref.nii -omat %s' % (
            flags, dof, omat)
        proc = Popen(cmd, shell=True,stdout=PIPE,stderr=PIPE)
        stdout, stderr = proc.communicate()
        if not os.path.isfile(omat):
            raise FSLError('Command "%s" failed somehow - stdout: %s\n'
                           'and stderr: %s\n' % (cmd, stdout, stderr))
        res = np.loadtxt(omat)
    return flirt2aff(res, in_img, ref_img)

def apply_warp(in_nii,affine_mat,nonlin_nii,out_nii):
    cmd='applywarp --ref=${FSLDIR}/data/standard/FMRIB58_FA_1mm --in='+in_nii+' --warp='+nonlin_nii+' --out='+out_nii
    print(cmd)
    pipe(cmd)

def create_displacements(in_nii,affine_mat,nonlin_nii,invw_nii,disp_nii,dispa_nii):
    commands=[]
    commands.append('flirt -ref ${FSLDIR}/data/standard/FMRIB58_FA_1mm -in '+in_nii+' -omat ' + affine_mat)
    commands.append('fnirt --in='+in_nii+' --aff='+affine_mat+' --cout='+nonlin_nii+' --config=FA_2_FMRIB58_1mm')
    commands.append('invwarp --ref='+in_nii+' --warp='+nonlin_nii+' --out='+invw_nii)
    commands.append('fnirtfileutils --in='+nonlin_nii+' --ref=${FSLDIR}/data/standard/FMRIB58_FA_1mm --out='+disp_nii)
    commands.append('fnirtfileutils --in='+nonlin_nii+' --ref=${FSLDIR}/data/standard/FMRIB58_FA_1mm --out='+dispa_nii + ' --withaff')
    for c in commands:
        print(c)
        pipe(c)