from __future__ import print_function, absolute_import, division, unicode_literals

import logging
from collections import namedtuple, deque, OrderedDict

import casadi as ca
import numpy as np
import itertools
from typing import Union, Dict, Iterable

from pymoca import ast
from pymoca.tree import TreeWalker, TreeListener, flatten

from .alias_relation import AliasRelation
from .model import Model, Variable, DelayArgument
from .mtensor import _MTensor, _new_mx

from ._options import _merge_default_options

logger = logging.getLogger("pymoca")

# TODO
#  - Nested for loops
#  - Delay operator on arbitrary expressions
#  - Pre operator

OP_MAP = {'*': "__mul__",
          '+': "__add__",
          "-": "__sub__",
          "/": "__div__",
          '^': "__pow__",
          '>': '__gt__',
          '<': '__lt__',
          '<=': '__le__',
          '>=': '__ge__',
          '!=': '__ne__',
          '==': '__eq__',
          "min": "fmin",
          "max": "fmax",
          "abs": "fabs",
          "and": "__mul__",
          "or": "__add__"}

ForLoopIndexedSymbol = namedtuple('ForLoopIndexedSymbol', ['tree', 'transpose', 'indices'])

# noinspection PyPep8Naming,PyUnresolvedReferences
class ForLoop:
    def __init__(self, generator, tree):
        self.tree = tree
        self.generator = generator
        i = tree.indices[0]
        e = i.expression
        start = e.start.value
        step = e.step.value
        stop = self.generator.get_integer(e.stop)
        self.values = np.arange(start, stop + step, step, dtype=np.int)
        self.index_variable = _new_mx(i.name)
        self.name = i.name
        self.indexed_symbols = OrderedDict()

    def register_indexed_symbol(self, e, index_function, transpose, tree, index_expr=None):
        if isinstance(index_expr, ca.MX) and index_expr is not self.index_variable:
            F = ca.Function('index_expr', [self.index_variable], [index_expr])
            # expr = lambda ar: np.array([F(a)[0] for a in ar], dtype=np.int)
            Fmap = F.map("map", self.generator.map_mode, len(self.values), [], [])
            res = Fmap.call([self.values])
            indices = np.array(res[0].T, dtype=np.int)
        else:
            indices = self.values
        self.indexed_symbols[e] = ForLoopIndexedSymbol(tree, transpose, index_function(indices - 1))

Assignment = namedtuple('Assignment', ['left', 'right'])

class GeneratorWalker(TreeWalker):
    """TreeWalker that skips processing of annotations"""

    def skip_child(self, tree: ast.Node, child_name: str) -> bool:
        skip = super().skip_child(tree, child_name)
        if isinstance(tree, ast.Class) and child_name == "annotation":
            return True
        return skip

    def order_keys(self, keys: Iterable[str]):
        # Symbols must come before classes, as we need to access symbol values when creating
        # CasADi interpolant functions.
        return sorted(keys, key=lambda attr: 0 if attr == 'symbols' else 1)

# noinspection PyPep8Naming,PyUnresolvedReferences
class Generator(TreeListener):
    def __init__(self, root: ast.Tree, class_name: str, options: Dict[str, bool]):
        super(Generator, self).__init__()
        self.src = {}
        self.model = Model()
        self.root = root
        c = self.root.classes[class_name]
        self.nodes = {c: {'time': self.model.time, 'finalTime': self.model.finalTime}}
        self.derivative = {}
        self.for_loops = deque()
        self.functions = {}
        self.entered_classes = deque()
        self.map_mode = 'inline' if options['unroll_loops'] else 'serial'
        self.function_mode = (True, False) if options['inline_functions'] else (False, True)
        self.delay_counter = 0

        # NOTE: Part of MTensor workaround.
        self._expand_vectors_enabled = options['expand_vectors']

    @property
    def current_class(self):
        return self.entered_classes[-1]

    def _ast_symbols_to_variables(self, ast_symbols, differentiate=False):
        variables = []
        for ast_symbol in ast_symbols:
            mx_symbol = self.get_mx(ast_symbol)
            modelica_shape = mx_symbol._modelica_shape
            if mx_symbol.is_empty():
                continue
            if differentiate:
                mx_symbol = self.get_derivative(mx_symbol)
                mx_symbol._modelica_shape = modelica_shape
            python_type = self.get_python_type(ast_symbol)
            variable = Variable(mx_symbol, python_type)
            if not differentiate:
                for a in ast.Symbol.ATTRIBUTES:
                    v = self.get_mx(getattr(ast_symbol, a))
                    if v is not None:
                        if isinstance(v, ca.DM) and all(x == (None,) for x in modelica_shape):
                            # Scalar numeric type that behaves like an array.
                            # Coerce to Pyhton type to avoid interpretation
                            # issues.
                            v = python_type(v)
                        elif isinstance(v, (float, int)) and not isinstance(v, python_type):
                            # We skip booleans for now, as users likely depend
                            # on them being integer/float-like.
                            try:
                                v = python_type(v)
                            except (OverflowError, ValueError):
                                # Cannot convert NaN/infs to integer
                                pass

                        setattr(variable, a, v)
                variable.prefixes = ast_symbol.prefixes
            variables.append(variable)
        return variables

    def enterClass(self, tree):
        logger.debug('enterClass {}'.format(tree.name))

        self.entered_classes.append(tree)
        self.nodes.setdefault(tree, {})

    def exitClass(self, tree):
        logger.debug('exitClass {}'.format(tree.name))

        if tree.type == 'function':
            # Already handled previously
            self.entered_classes.pop()
            return

        ode_states = []
        alg_states = []
        inputs = []
        constants = []
        parameters = []
        symbols = sorted(tree.symbols.values(), key=lambda x: x.order)
        for s in symbols:
            if 'constant' in s.prefixes:
                constants.append(s)
            elif 'parameter' in s.prefixes:
                parameters.append(s)
            elif 'input' in s.prefixes:
                inputs.append(s)
            elif 'state' in s.prefixes:
                ode_states.append(s)
            else:
                alg_states.append(s)

        self.model.states = self._ast_symbols_to_variables(ode_states)
        self.model.der_states = self._ast_symbols_to_variables(ode_states, differentiate=True)
        self.model.alg_states = self._ast_symbols_to_variables(alg_states)
        self.model.constants = self._ast_symbols_to_variables(constants)
        self.model.parameters = self._ast_symbols_to_variables(parameters)

        # We extend the input list, as it is already populated with delayed states.
        self.model.inputs.extend(self._ast_symbols_to_variables(inputs))

        # The outputs are a list of strings of state names. Specifying
        # multiple aliases of the same state is allowed.
        self.model.outputs = [v.symbol.name() for v in itertools.chain(self.model.states, self.model.alg_states) if 'output' in v.prefixes]

        def discard_empty(l):
            return list(filter(lambda x: not ca.MX(x).is_empty(), l))

        self.model.equations = discard_empty([self.get_mx(e) for e in tree.equations])
        self.model.initial_equations = discard_empty([self.get_mx(e) for e in tree.initial_equations])

        # TODO: check if it is valid (like not empty equation)
        self.model.constraints = [self.get_mx(e) for e in tree.constraints]

        if tree.type == 'optimization':
            self.model.optimization_attributes = {argument.value.component.name: self.src[argument.value.modifications[0]] for argument in tree.optimization_attributes.arguments}

        if len(tree.statements) + len(tree.initial_statements) > 0:
            raise NotImplementedError('Statements are currently supported inside functions only')

        self.entered_classes.pop()

    def exitArray(self, tree):
        self.src[tree] = [self.src[e] for e in tree.values]

    def exitPrimary(self, tree):
        self.src[tree] = tree.value

    def exitExpression(self, tree):
        if isinstance(tree.operator, ast.ComponentRef):
            op = tree.operator.name
        else:
            op = tree.operator

        if op == '*':
            op = 'mtimes'  # .* differs from *
        if op.startswith('.'):
            op = op[1:]

        logger.debug('exitExpression')

        n_operands = len(tree.operands)
        if op == 'der':
            v = self.get_mx(tree.operands[0])
            src = self.get_derivative(v)
        elif op == '-' and n_operands == 1:
            src = -self.get_mx(tree.operands[0])
        elif op == 'not' and n_operands == 1:
            src = ca.if_else(self.get_mx(tree.operands[0]), 0, 1, True)
        elif op == 'mtimes':
            assert n_operands >= 2
            src = self.get_mx(tree.operands[0])
            for i in tree.operands[1:]:
                src = ca.mtimes(src, self.get_mx(i))
        elif op == 'transpose' and n_operands == 1:
            src = self.get_mx(tree.operands[0]).T
        elif op == 'sum' and n_operands == 1:
            v = self.get_mx(tree.operands[0])
            src = ca.sum1(v)
        elif op == 'linspace' and n_operands == 3:
            a = self.get_mx(tree.operands[0])
            b = self.get_mx(tree.operands[1])
            n_steps = self.get_integer(tree.operands[2])
            src = ca.linspace(a, b, n_steps)
        elif op == 'fill' and n_operands == 2:
            val = self.get_mx(tree.operands[0])
            n_row = self.get_integer(tree.operands[1])
            src = val * ca.DM.ones(n_row)
        elif op == 'fill' and n_operands == 3:
            val = self.get_mx(tree.operands[0])
            n_row = self.get_integer(tree.operands[1])
            n_col = self.get_integer(tree.operands[2])
            src = val * ca.DM.ones(n_row, n_col)
        elif op == 'zeros' and n_operands == 1:
            n_row = self.get_integer(tree.operands[0])
            src = ca.DM.zeros(n_row)
        elif op == 'zeros' and n_operands == 2:
            n_row = self.get_integer(tree.operands[0])
            n_col = self.get_integer(tree.operands[1])
            src = ca.DM.zeros(n_row, n_col)
        elif op == 'ones' and n_operands == 1:
            n_row = self.get_integer(tree.operands[0])
            src = ca.DM.ones(n_row)
        elif op == 'ones' and n_operands == 2:
            n_row = self.get_integer(tree.operands[0])
            n_col = self.get_integer(tree.operands[1])
            src = ca.DM.ones(n_row, n_col)
        elif op == 'identity' and n_operands == 1:
            n = self.get_integer(tree.operands[0])
            src = ca.DM.eye(n)
        elif op == 'diagonal' and n_operands == 1:
            diag = self.get_mx(tree.operands[0])
            n = len(diag)
            indices = list(range(n))
            src = ca.DM.triplet(indices, indices, diag, n, n)
        elif op == 'cat':
            axis = self.get_integer(tree.operands[0])
            assert axis == 1, "Currently only concatenation on first axis is supported"

            entries = []
            for sym in [self.get_mx(op) for op in tree.operands[1:]]:
                if isinstance(sym, list):
                    for e in sym:
                        entries.append(e)
                else:
                    entries.append(sym)
            src = ca.vertcat(*entries)
        elif op == 'delay' and n_operands == 2:
            expr = self.get_mx(tree.operands[0])
            duration = self.get_mx(tree.operands[1])

            src = _new_mx('_pymoca_delay_{}'.format(self.delay_counter), *expr.size())
            self.delay_counter += 1

            for f in self.for_loops:
                syms = set(ca.symvar(expr))
                if syms.intersection(f.indexed_symbols):
                    f.register_indexed_symbol(src, lambda i: i, True, tree.operands[0], f.index_variable)

            self.model.delay_states.append(src.name())
            self.model.inputs.append(Variable(src))

            delay_argument = DelayArgument(expr, duration)
            self.model.delay_arguments.append(delay_argument)
        elif op == '_pymoca_interp1d' and n_operands >= 3 and n_operands <= 4:
            entered_class = self.entered_classes[-1]
            if isinstance(tree.operands[0], ast.ComponentRef):
                xp = self.get_mx(entered_class.symbols[tree.operands[0].name].value)
            else:
                xp = self.get_mx(tree.operands[0])
            if isinstance(tree.operands[1], ast.ComponentRef):
                yp = self.get_mx(entered_class.symbols[tree.operands[1].name].value)
            else:
                yp = self.get_mx(tree.operands[1])
            arg = self.get_mx(tree.operands[2])
            if n_operands == 4:
                assert isinstance(tree.operands[3], ast.Primary)
                mode = tree.operands[3].value
            else:
                mode = 'linear'
            func = ca.interpolant('interpolant', mode, [xp], yp)
            src = func(arg)
        elif op == '_pymoca_interp2d' and n_operands >= 5 and n_operands <= 6:
            entered_class = self.entered_classes[-1]
            if isinstance(tree.operands[0], ast.ComponentRef):
                xp = self.get_mx(entered_class.symbols[tree.operands[0].name].value)
            else:
                xp = self.get_mx(tree.operands[0])
            if isinstance(tree.operands[1], ast.ComponentRef):
                yp = self.get_mx(entered_class.symbols[tree.operands[1].name].value)
            else:
                yp = self.get_mx(tree.operands[1])
            if isinstance(tree.operands[2], ast.ComponentRef):
                zp = self.get_mx(entered_class.symbols[tree.operands[2].name].value)
            else:
                zp = self.get_mx(tree.operands[2])
            arg_1 = self.get_mx(tree.operands[3])
            arg_2 = self.get_mx(tree.operands[4])
            if n_operands == 6:
                assert isinstance(tree.operands[5], ast.Primary)
                mode = tree.operands[5].value
            else:
                mode = 'linear'
            func = ca.interpolant('interpolant', mode, [xp, yp], np.array(zp).ravel(order='F'))
            src = func(ca.vertcat(arg_1, arg_2))
        elif op in OP_MAP and n_operands == 2:
            lhs = ca.MX(self.get_mx(tree.operands[0]))
            rhs = ca.MX(self.get_mx(tree.operands[1]))
            lhs_op = getattr(lhs, OP_MAP[op])
            src = lhs_op(rhs)
        elif op in OP_MAP and n_operands == 1:
            lhs = ca.MX(self.get_mx(tree.operands[0]))
            lhs_op = getattr(lhs, OP_MAP[op])
            src = lhs_op()
        else:
            src = ca.MX(self.get_mx(tree.operands[0]))
            # Check for built-in operations, such as the
            # elementary functions, first.
            if hasattr(src, op) and n_operands <= 2:
                if n_operands == 1:
                    src = ca.MX(self.get_mx(tree.operands[0]))
                    src = getattr(src, op)()
                else:
                    lhs = ca.MX(self.get_mx(tree.operands[0]))
                    rhs = ca.MX(self.get_mx(tree.operands[1]))
                    lhs_op = getattr(lhs, op)
                    src = lhs_op(rhs)
            else:
                try: # Check if there is a component named as the operation. In that case we are dealing with a time access
                    # Should we check for symbol as well?
                    v = self.get_mx(ast.ComponentRef(name=op))
                    t = self.get_mx(tree.operands[0])
                    src = self.get_symbol_time_access(v, t)

                except KeyError:
                    func = self.get_function(op)
                    src = ca.vertcat(*func.call([self.get_mx(operand) for operand in tree.operands], *self.function_mode))

        self.src[tree] = src

    def exitIfExpression(self, tree):
        logger.debug('exitIfExpression')

        assert (len(tree.conditions) + 1 == len(tree.expressions))

        src = self.get_mx(tree.expressions[-1])
        for cond_index in range(len(tree.conditions)):
            cond = self.get_mx(tree.conditions[-(cond_index + 1)])
            expr1 = self.get_mx(tree.expressions[-(cond_index + 2)])

            src = ca.if_else(cond, expr1, src, True)

        self.src[tree] = src

    def exitConstraint(self, tree):
        logger.debug('exitConstraint')

        src_left = self.get_mx(tree.left)
        src_right = self.get_mx(tree.right)

        # Always return bigger equal constraints
        if tree.operand == '>=':
            self.src[tree] = src_left - src_right
        elif tree.operand == '<=':
            self.src[tree] = src_right - src_left
        else:
            raise Exception('Operand {} is not supported in constraints'.format(tree.operand))

    def exitEquation(self, tree):
        logger.debug('exitEquation')

        if isinstance(tree.left, list):
            src_left = ca.vertcat(*[self.get_mx(c) for c in tree.left])
        else:
            src_left = self.get_mx(tree.left)

        if isinstance(tree.right, list):
            src_right = ca.vertcat(*[self.get_mx(c) for c in tree.right])
        else:
            src_right = self.get_mx(tree.right)

        src_left = ca.MX(src_left)
        src_right = ca.MX(src_right)

        # According to the Modelica spec,
        # "It is possible to omit left hand side component references and/or truncate the left hand side list in order to discard outputs from a function call."
        if isinstance(tree.right, ast.Expression) and tree.right.operator in self.root.classes:
            if src_left.size1() < src_right.size1():
                src_right = src_right[0:src_left.size1()]
        if isinstance(tree.left, ast.Expression) and tree.left.operator in self.root.classes:
            if src_left.size1() > src_right.size1():
                src_left = src_left[0:src_right.size1()]

        # If dimensions between the lhs and rhs do not match, but the dimensions of lhs
        # and transposed rhs do match, transpose the rhs.
        if src_left.shape != src_right.shape and src_left.shape == src_right.shape[::-1]:
            src_right = ca.transpose(src_right)

        self.src[tree] = src_left - src_right

    def enterForEquation(self, tree):
        logger.debug('enterForEquation')

        self.for_loops.append(ForLoop(self, tree))

    def exitForEquation(self, tree):
        logger.debug('exitForEquation')

        f = self.for_loops.pop()
        if len(f.values) > 0:
            indexed_symbols = list(f.indexed_symbols.keys())
            args = [f.index_variable] + indexed_symbols
            expr = ca.vcat([ca.vec(self.get_mx(e)) for e in tree.equations])
            free_vars = ca.symvar(expr)

            arg_names = [arg.name() for arg in args]
            free_vars = [e for e in free_vars if e.name() not in arg_names]
            all_args = args + free_vars
            F = ca.Function('loop_body', all_args, [expr])

            indexed_symbols_full = []
            for k in indexed_symbols:
                s = f.indexed_symbols[k]
                indices = s.indices
                try:
                    i = self.model.delay_states.index(k.name())
                except ValueError:
                    orig_symbol = self.nodes[self.current_class][s.tree.name]
                else:
                    # We are missing a similarly shaped delayed symbol. Make a new one with the appropriate shape.
                    delay_symbol = self.model.delay_arguments[i]

                    # We need to figure out the shape of the expression that
                    # we are delaying. The symbols that can occur in the delay
                    # expression should have been encountered before this
                    # iteration of the loop. The assert statement below covers
                    # this.
                    delay_expr_args = free_vars + all_args[:len(indexed_symbols_full)+1]
                    assert set(ca.symvar(delay_symbol.expr)).issubset(delay_expr_args)

                    f_delay_expr = ca.Function('delay_expr', delay_expr_args, [delay_symbol.expr])
                    f_delay_map = f_delay_expr.map("map", self.map_mode, len(f.values), list(
                        range(len(free_vars))), [])
                    [res] = f_delay_map.call(free_vars + [f.values] + indexed_symbols_full)
                    res = res.T

                    # Make the symbol with the appropriate size, and replace the old symbol with the new one.
                    orig_symbol = _new_mx(k.name(), *res.size())
                    assert res.size1() == 1 or res.size2() == 1, "Slicing does not yet work with 2-D indices"
                    indices = slice(None, None)

                    model_input = next(x for x in self.model.inputs if x.symbol.name() == k.name())
                    model_input.symbol = orig_symbol
                    self.model.delay_arguments[i] = DelayArgument(res, delay_symbol.duration)

                indexed_symbol = orig_symbol[indices]
                if s.transpose:
                    indexed_symbol = ca.transpose(indexed_symbol)
                indexed_symbols_full.append(indexed_symbol)

            Fmap = F.map("map", self.map_mode, len(f.values), list(
                range(len(args), len(all_args))), [])
            res = Fmap.call([f.values] + indexed_symbols_full + free_vars)

            self.src[tree] = res[0].T
        else:
            self.src[tree] = ca.MX()

    def exitIfEquation(self, tree):
        logger.debug('exitIfEquation')

        # Check if every equation block contains the same number of equations
        if len(set((len(x) for x in tree.blocks))) != 1:
            raise Exception("Every branch in an if-equation needs the same number of equations.")

        # NOTE: We currently assume that we always have an else-clause. This
        # is not strictly necessary, see the Modelica Spec on if equations.
        assert tree.conditions[-1] == True

        src = ca.vertcat(*[self.get_mx(e) for e in tree.blocks[-1]])

        for cond_index in range(1, len(tree.conditions)):
            cond = self.get_mx(tree.conditions[-(cond_index + 1)])
            expr1 = ca.vertcat(*[self.get_mx(e) for e in tree.blocks[-(cond_index + 1)]])
            src = ca.if_else(cond, expr1, src, True)

        self.src[tree] = src

    def exitAssignmentStatement(self, tree):
        logger.debug('exitAssignmentStatement')

        all_assignments = []

        expr = self.get_mx(tree.right)
        for component_ref in tree.left:
            all_assignments.append(Assignment(self.get_mx(component_ref), expr))

        self.src[tree] = all_assignments

    def exitIfStatement(self, tree):
        logger.debug('exitIfStatement')

        # We assume an equal number of statements per branch.
        # Furthermore, we assume that every branch assigns to the same variables.
        assert len(set((len(x) for x in tree.blocks))) == 1

        # NOTE: We currently assume that we always have an else-clause. This
        # is not strictly necessary, see the Modelica Spec on if statements.
        assert tree.conditions[-1] == True

        expanded_blocks = OrderedDict()

        for b in tree.blocks:
            block_assignments = []
            for s in b:
                assignments = self.get_mx(s)
                for assignment in assignments:
                    expanded_blocks.setdefault(assignment.left, []).append(assignment.right)

        assert len(set((len(x) for x in expanded_blocks.values()))) == 1

        all_assignments = []

        for lhs, values in expanded_blocks.items():
            # Set default value to else block, and then loop in reverse over all branches
            src = values[-1]
            for cond, rhs in zip(tree.conditions[-2::-1], values[-2::-1]):
                cond = self.get_mx(cond)
                src = ca.if_else(cond, rhs, src, True)

            all_assignments.append(Assignment(lhs, src))

        self.src[tree] = all_assignments

    def enterForStatement(self, tree):
        logger.debug('enterForStatement')

        self.for_loops.append(ForLoop(self, tree))

    def exitForStatement(self, tree):
        logger.debug('exitForStatement')

        f = self.for_loops.pop()
        if len(f.values) > 0:
            indexed_symbols = list(f.indexed_symbols.keys())
            args = [f.index_variable] + indexed_symbols
            expr = ca.vcat([ca.vec(self.get_mx(e.right)) for e in tree.statements])
            free_vars = ca.symvar(expr)

            arg_names = [arg.name() for arg in args]
            free_vars = [e for e in free_vars if e.name() not in arg_names]
            all_args = args + free_vars
            F = ca.Function('loop_body', all_args, [expr])

            indexed_symbols_full = []
            for k in indexed_symbols:
                s = f.indexed_symbols[k]
                orig_symbol = self.nodes[self.current_class][s.tree.name]
                indexed_symbol = orig_symbol[s.indices]
                if s.transpose:
                    indexed_symbol = ca.transpose(indexed_symbol)
                indexed_symbols_full.append(indexed_symbol)

            Fmap = F.map("map", self.map_mode, len(f.values), list(
                range(len(args), len(all_args))), [])
            res = Fmap.call([f.values] + indexed_symbols_full + free_vars)

            # Split into a list of statements
            variables = [assignment.left for statement in tree.statements for assignment in self.get_mx(statement)]
            all_assignments = []
            for i in range(len(f.values)):
                for j, variable in enumerate(variables):
                    all_assignments.append(Assignment(variable, res[0][j, i].T))

            self.src[tree] = all_assignments
        else:
            self.src[tree] = []

    def get_integer(self, tree: Union[ast.Primary, ast.ComponentRef, ast.Expression, ast.Slice]) -> Union[int, ca.MX, np.ndarray]:
        # CasADi needs to know the dimensions of symbols at instantiation.
        # We therefore need a mechanism to evaluate expressions that define dimensions of symbols.
        if isinstance(tree, ast.Primary):
            return None if tree.value is None else int(tree.value)
        if isinstance(tree, ast.ComponentRef):
            s = self.current_class.symbols[tree.name]
            assert (s.type.name == 'Integer')
            return self.get_integer(s.value)
        if isinstance(tree, ast.Expression):
            # Make sure that the expression has been converted to MX by (re)visiting the
            # relevant part of the AST.
            ast_walker = TreeWalker()
            ast_walker.walk(self, tree)

            # Obtain expression
            expr = self.get_mx(tree)

            # Obtain the symbols it depends on
            free_vars = ca.symvar(expr)

            # Find the values of the symbols
            vals = []
            for free_var in free_vars:
                if free_var.is_symbolic():
                    if (len(self.for_loops) > 0) and (free_var.name() == self.for_loops[-1].name):
                        vals.append(self.for_loops[-1].index_variable)
                    else:
                        vals.append(self.get_integer(self.current_class.symbols[free_var.name()].value))

            # Evaluate the expression
            F = ca.Function('get_integer', free_vars, [expr])
            ret = F.call(vals, *self.function_mode)
            if ret[0].is_constant():
                # We managed to evaluate the expression.  Assume the result to be integer.
                return int(ret[0])
            else:
                # Expression depends on other symbols.  Could not extract integer value.
                return ret[0]
        if isinstance(tree, ast.Slice):
            start = self.get_integer(tree.start)
            step = self.get_integer(tree.step)
            stop = self.get_integer(tree.stop)
            return slice(start, stop, step)
        else:
            raise Exception('Unexpected node type {}'.format(tree.__class__.__name__))

    @staticmethod
    def get_python_type(tree):
        if tree.type.name == 'Boolean':
            return bool
        elif tree.type.name == 'Integer':
            return int
        else:
            return float

    def get_shape(self, tree):
        return [[self.get_integer(d) for d in d_list] for d_list in tree.dimensions]

    def get_symbol(self, tree):
        # Create symbol
        shape = self.get_shape(tree)

        if any(isinstance(x, slice) for var_shape in shape for x in var_shape):
            # Symbol has unspecified dimensions. Value is specified, and
            # carries the correct dimensions.

            # We should only get slices as dimensions for a symbol if one of
            # the dimensions is unspecified, i.e. None.
            assert None in (itertools.chain.from_iterable((x.start, x.stop)
                            for var_shape in shape for x in var_shape if isinstance(x, slice)))

            val_shape = np.array(self.src[tree.value]).shape

            # Check if specified dimensions agree between definition and value
            val_dim_i = -1
            for var_i, var_shape in enumerate(shape):
                for dim_i, dim_size in enumerate(var_shape):
                    if dim_size is None:
                        continue

                    val_dim_i += 1
                    if isinstance(dim_size, slice):
                        shape[var_i][dim_i] = val_shape[val_dim_i]
                        continue

                    if val_shape[val_dim_i] != dim_size:
                        raise Exception("Dimension {} of definition and value for symbol {} "
                                        "differs: {} != {}"
                                        .format(val_dim_i + 1, tree.name, dim_size,
                                                val_shape[val_dim_i]))

        tensor_shape = [d for var_shape in shape for d in var_shape if d is not None]
        if len(tensor_shape) > 2:
            # MX does not support this, so we have to use our own wrapper.
            if not self._expand_vectors_enabled:
                raise NotImplementedError("Cannot handle 3D+ arrays without setting 'expand_vectors'")
            s = _MTensor(tree.name, *tensor_shape)
        else:
            s = _new_mx(tree.name, *tensor_shape)

        # Make a notion of the original shape, as MX is always 2D (even for 1D symbols),
        # and for nested classes we want to remember at which symbols to place indices.
        s._modelica_shape = tuple([tuple(var_shape) for var_shape in shape])

        self.nodes[self.current_class][tree.name] = s
        return s

    def get_derivative(self, s):

        # Case 1: s is a constant, e.g. MX(5)
        if ca.MX(s).is_constant():
            return 0

        # Case 2: s is a symbol, e.g. MX(x)
        elif s.is_symbolic():
            if s.name() not in self.derivative:
                if len(self.for_loops) > 0 and s in self.for_loops[-1].indexed_symbols:
                    # Create a new indexed symbol, referencing to the for loop index inside the vector derivative symbol.
                    for_loop_symbol = self.for_loops[-1].indexed_symbols[s]
                    s_without_index = self.get_mx(ast.ComponentRef(name=for_loop_symbol.tree.name))
                    der_s_without_index = self.get_derivative(s_without_index)
                    if ca.MX(der_s_without_index).is_symbolic():
                        return self.get_indexed_symbol(ast.ComponentRef(name=der_s_without_index.name(), indices=for_loop_symbol.tree.indices), der_s_without_index)
                    else:
                        return 0
                else:
                    der_s = _new_mx("der({})".format(s.name()), s.size())
                    # If the derivative contains an expression (e.g. der(x + y)) this method is
                    # called with MX variables that are the result of a ca.symvar call. This
                    # ca.symvar call strips the _modelica_shape field from the MX variable,
                    # therefore we need to find the original MX to get the modelica shape.
                    der_s._modelica_shape = \
                        self.nodes[self.current_class][s.name()]._modelica_shape
                    self.derivative[s.name()] = der_s
                    self.nodes[self.current_class][der_s.name()] = der_s
                    return der_s
            else:
                return self.derivative[s.name()]

        # Case 3: s is an already indexed symbol, e.g. MX(x[1])
        elif s.is_op(ca.OP_GETNONZEROS) and s.dep().is_symbolic():
            slice_info = s.info()['slice']
            dep = s.dep()
            if dep.name() not in self.derivative:
                der_dep = _new_mx("der({})".format(dep.name()), dep.size())
                der_dep._modelica_shape = \
                    self.nodes[self.current_class][dep.name()]._modelica_shape
                self.derivative[dep.name()] = der_dep
                self.nodes[self.current_class][der_dep.name()] = der_dep
                return der_dep[slice_info['start']:slice_info['stop']:slice_info['step']]
            else:
                return self.derivative[dep.name()][slice_info['start']:slice_info['stop']:slice_info['step']]

        # Case 4: s is an expression that requires differentiation, e.g. MX(x2 * x2)
        # Need to do this sort of expansion: der(x1 * x2) = der(x1) * x2 + x1 * der(x2)
        else:
            # Differentiate expression using CasADi
            orig_deps = ca.symvar(s)
            deps = ca.vertcat(*orig_deps)
            J = ca.Function('J', [deps], [ca.jacobian(s, deps)])
            J_sparsity = J.sparsity_out(0)
            der_deps = [self.get_derivative(dep) if J_sparsity.has_nz(0, j) else ca.DM.zeros(dep.size()) for j, dep in enumerate(orig_deps)]
            return ca.mtimes(J(deps), ca.vertcat(*der_deps))

    def get_indexed_symbol(self, tree, s):
        assert len([dim for shape in s._modelica_shape for dim in shape if dim is not None]) <= 2,\
            "Dimensions higher than two are not yet supported"

        assert len(s._modelica_shape) >= len(tree.indices)

        # For nested variables where an equation is defined at one of the nested models,
        # the modelica shape will contain the shape for the whole nested variable, but the indices
        # will only contain the indices for the symbol in the nested model. We only use the last
        # part of _modelica_shape in this case.
        assert tree.indices
        shapes = s._modelica_shape[-len(tree.indices):]

        # Check whether we loop over an index of this symbol
        indices = []
        for_loop = None
        for i, (index_array, shape) in enumerate(zip(tree.indices, shapes)):
            if len(index_array) > len(shape):
                symbol_name = s.name() if len(tree.indices) == 1 \
                    else s.name().split('.')[i] + ' in nested symbol ' + s.name()
                raise ValueError('Too many indices found for symbol {}, check if the symbol has '
                                 'the correct dimensions.'.format(symbol_name))

            for index, dim in zip(index_array, shape):
                if index is None and dim is None:
                    continue

                sl = None

                if isinstance(index, ast.ComponentRef):
                    for f in self.for_loops:
                        if index.name == f.name:
                            # TODO support nested loops
                            for_loop = f
                            sl = for_loop.index_variable

                if sl is None:
                    sl = self.get_integer(index) if index is not None else None

                    if sl is None and dim is not None:
                        sl = slice(None, None, 1)
                    if sl is not None and dim is None:
                        symbol_name = s.name() if len(tree.indices) == 1 \
                            else s.name().split('.')[i] + ' in nested symbol ' + s.name()
                        raise ValueError('Symbol {} was given an index of {} but this symbol '
                                         'is not an array.'.format(symbol_name, sl))
                    elif isinstance(sl, int):
                        # Modelica indexing starts from one;  Python from zero.
                        if sl <= 0 or sl > dim:
                            symbol_name = s.name() if len(tree.indices) == 1 \
                                else s.name().split('.')[i] + ' in nested symbol ' + s.name()
                            raise ValueError("Index {} of symbol {} is out of bounds. "
                                             "Index should be in range [1,{}] "
                                             "(Modelica uses 1-based indexing)."
                                             .format(sl, symbol_name, dim))
                        sl = sl - 1
                    elif isinstance(sl, slice):
                        # Modelica indexing starts from one;  Python from zero.
                        sl = slice(None if sl.start is None else sl.start - 1, sl.stop, sl.step)
                    else:
                        for_loop = self.for_loops[-1]

                indices.append(sl)

        if for_loop is not None:
            if isinstance(indices[0], ca.MX):
                if len(indices) > 1:
                    s = s[:, indices[1]]
                    indexed_symbol = _new_mx('{}[{},{}]'.format(tree.name, for_loop.name, indices[1]), s.size2())
                    index_function = lambda i : (i, indices[1])
                else:
                    indexed_symbol = _new_mx('{}[{}]'.format(tree.name, for_loop.name))
                    index_function = lambda i : i

                # If the indexed symbol is empty, we know we do not have to
                # map the for loop over it
                if np.prod(s.shape) != 0:
                    for_loop.register_indexed_symbol(indexed_symbol, index_function, True, tree, indices[0])
            else:
                s = ca.transpose(s[indices[0], :])
                indexed_symbol = _new_mx('{}[{},{}]'.format(tree.name, indices[0], for_loop.name), s.size2())
                index_function = lambda i: (indices[0], i)
                if np.prod(s.shape) != 0:
                    for_loop.register_indexed_symbol(indexed_symbol, index_function, False, tree, indices[1])
            return indexed_symbol
        else:
            if len(indices) == 1:
                return s[indices[0]]
            else:
                return s[indices[0], indices[1]]

    def get_symbol_time_access(self, v, t):
        return _new_mx('{}({})'.format(v.name(), str(t)), v.size())

    def get_component(self, tree):
        # Check special symbols
        if tree.name == 'time':
            return self.model.time
        elif tree.name == 'finalTime':
            return self.model.finalTime
        else:
            for f in reversed(self.for_loops):
                if f.name == tree.name:
                    return f.index_variable

        # Check ordinary symbols
        symbol = self.current_class.symbols[tree.name]
        s = self.get_mx(symbol)
        if len([index for index_array in tree.indices
                for index in index_array if index is not None]) > 0:
            s = self.get_indexed_symbol(tree, s)
        return s

    def get_mx(self, tree: Union[ast.Symbol, ast.ComponentRef, ast.Expression]) -> ca.MX:
        """
        We pull components and symbols from the AST on demand.
        This is to ensure that parametrized vector dimensions can be resolved.  Vector
        dimensions need to be known at CasADi MX creation time.
        :param tree:
        :return:
        """
        if tree not in self.src:
            if isinstance(tree, ast.Symbol):
                s = self.get_symbol(tree)
            elif isinstance(tree, ast.ComponentRef):
                s = self.get_component(tree)
            else:
                raise Exception('Tried to look up expression before it was reached by the tree walker')
            self.src[tree] = s
        return self.src[tree]

    def get_function(self, function_name):
        if function_name in self.functions:
            return self.functions[function_name]

        try:
            tree = self.root.classes[function_name]
        except KeyError:
            raise Exception('Unknown function {}'.format(function_name))

        inputs = []
        outputs = []
        tmp = []
        for s in tree.symbols.values():
            src = self.get_mx(s)
            if 'input' in s.prefixes:
                inputs.append(src)
            elif 'output' in s.prefixes:
                outputs.append(src)
            else:
                tmp.append(src)

        # Store current variable values
        values = {}
        for variable in inputs:
            values[variable] = variable

        # Process statements in order
        for statement in tree.statements:
            src = self.get_mx(statement)
            for assignment in src:
                [values[assignment.left]] = ca.substitute([assignment.right], list(values.keys()), list(values.values()))

        output_expr = ca.substitute([values[output] for output in outputs], tmp, [values[t] for t in tmp])
        func = ca.Function(tree.name, inputs, output_expr)
        self.functions[function_name] = func

        return func

def generate(ast_tree: ast.Tree, model_name: str, options: Dict[str, bool]=None) -> Model:
    """
    :param ast_tree: AST to generate from
    :param model_name: class to generate
    :param options: dictionary of generator options
    :return: casadi model
    """
    options = _merge_default_options(options)

    component_ref = ast.ComponentRef.from_string(model_name)
    ast_walker = GeneratorWalker()
    flat_tree = flatten(ast_tree, component_ref)
    component_ref_tuple = component_ref.to_tuple()
    casadi_gen = Generator(flat_tree, component_ref_tuple[-1], options)
    ast_walker.walk(casadi_gen, flat_tree)
    return casadi_gen.model


"""
A Python interface to CPTEC's Plume Rise Model.

TO DO: Create a base class name PLUMES and share code with the LockPlume.

"""

from numpy      import zeros, ones, meshgrid, linspace, any, \
                       pi, sin, cos, arccos, arange, array, \
                       savez, NaN, isnan

from scipy import optimize as opt

from datetime   import datetime, date, timedelta
from glob       import glob

from dozier     import DOZIER, granules
from PlumeRise_ import *       # f2py extension
from gfio       import GFIO

from pyobs           import NPZ, kde
from pyobs.binObs_   import binareas
from pyobs.minx      import MINXs
from MAPL.constants  import *

import eta

__VERSION__ = 2.1
__CVSTAG__  = '@CVSTAG'
__AMISS__   = 1.E+20

Bioma = [ 'Tropical Forest',
          'Extra-Tropical Forest',
          'Savanna',
          'Grassland' ]

DAY = timedelta(seconds=60*60*24)

#----------------------------------------------------------------------------------------------

#----
class MINXs_PR(MINXs):

    """
    Extension of the MINXs class adding the Freitas Plume Rise
    functionality. This class handles non-gridded, observation
    location fires.

    Parabolic Vertical Mass Distribution (VMD)
    ------------------------------------------
    If using the parabolic VMD as in getVMD() below, the (z_c,delta) parameters can
    be computed from the (z_i,z_f,z_d) parameters returned by plumevmd().

     a) Like Saulo:
                     z_c   = (z_f+z_i)/2
                     delta = (z_f-z_i)/2

     b) Preserve bottom half:
                     z_c   = z_d
                     delta = z_d - z_i

     c) Preserve upper half:
                     z_c   = z_d
                     delta = z_f - z_d

                       ---

    """

    def getPlume1(self,i,
                  hflux_kW=None,
                  frp_MW=None,
                  area_m2=None,afac=None,
                  Area=None,
                  rad2conv=5.,
                  Nominal=False,
                  which='z_a',
                  Verbose=False):

        """
        Runs the Plume Rise extension to compute the extent of the plume.
        On input,

           i         ---  index to operate on
           frp_MW    ---  fire radiative power in MW
           area_m2   ---  firea are im m2
           afac      ---  scale area by this amount
           Area      ---  if specified, PR model will see this value, but heat flux
                          will be computed the regular way, based on *area*
           rad2conv  ---  factor for converting radiative to convective heat fluxes.
           which     ---  Return one of z_i, z_f, z_d, z_a or all (Tupple), where
                          z_i -- height of maximum W (bottom of plume)
                          z_d -- height of maximum detrainment
                          z_a -- average height in (z_i,z_f), weighted by -dw/dz
                          z_f -- height where w<1 (top of plume)
           Nominal   ---  by default, areas and heat fluxes are from a
                          Dozier type algorithm, If Nominal=True,
                             area = 20e4 (20 ha)
                             hflux = rad2conv * FRP / area

	"""

        area = area_m2
        ptop = 1. # Pascal

        # Default areas, heat flux
        # ------------------------
        if Nominal:
            d_area = 20e4  #  typical fire size in m2 (20 ha)
        else:
            d_area = 1e6 * self.mod14.farea[i]      # km2 --> m2

        # Area
        # ----
        if area == None:
            area = d_area # m2

        # FRP
        # ---
        if frp_MW == None:
            frp_MW = self.mod14.frp[i] # MW

        # Convective Heat Flux
        # --------------------
        if hflux_kW == None:
            hflux_kW = 1e3 * rad2conv * frp_MW / area # kW/m2

        # Scaled Area
        # -----------
        if afac!=None:
            if afac==0:
                afac = 1e-8
            area = afac * area
            hflux_kW = hflux_kW / afac

        u = self.sample.u[i]
        v = self.sample.v[i]
        T = self.sample.t[i]
        q = self.sample.qv[i]
        delp = self.sample.delp[i]

        if delp.min()<=0 or T.min()<=0:
            print "out of range: ", delp.min(), T.min()
            return NaN

        # Ensure arrays are top-down as in GEOS-5
        # ---------------------------------------
        if delp[0] > delp[-1]:
            u = u[::-1]
            v = v[::-1]
            T = T[::-1]
            q = q[::-1]
            delp = delp[::-1]

        # Override firea area
        # -------------------
        if Area is not None:
            area = Area

        # Run plume rise model
        # --------------------
        # p,z,k,rc = plume(u,v,T,q,delp,ptop,hflux_kW,area)
        z_i,z_d,z_a,z_f,z,w,rc = plumevmd(u,v,T,q,delp,ptop,hflux_kW,area)
        if rc:
            raise ValueError, "error on return from <plume>, rc = %d"%rc

        if z_i==-1.:
             z_i,z_d,z_a,z_f = (NaN,self.sample.pblh[i],NaN,NaN)

        if   which=='z_i': z = z_i
        elif which=='z_f': z = z_f
        elif which=='z_d': z = z_d
        elif which=='z_a': z = z_a
        else:              z = (z_i,z_d,z_a,z_f,z,w)

        if Verbose:
            print " ", self.tyme[i], "| %8.2f | %8.2f %8.2f | %8.2f %8.2f %8.2f %8.2f | %03d"%\
                (self.mod14.fdist[i], area/1e4, hflux_kW, z_i,z_d,z_a,z_f,i)

        return z

#---
    def getPlume(self,I=None,Verbose=True,**kwopts):

        """
        Runs the Plume Rise extension to compute the extent of each plume.
	    """

        z_plume = __AMISS__ * ones(self.N)

        R = arange(self.N)
        if I is not None:
            R = R[I]

        if Verbose:
            print ""
            print "                    Plume Height Estimates"
            print ""
            print "  --------------------|----------|-------------------|-------------------------------------------"
            print "                      | Distance |   Fire Properties |             Plume Height"
            print "    MINX Date/Time    |  to Fire |   Area   Heat Flx |    z_i      z_d      z_a      z_f"
            print "                      |    km    |    ha      kW     |    km       km       km       km"
            print "  --------------------|----------|-------------------|-------------------------------------------"

#       Loop over time
#       --------------
        for i in R:
            z_plume[i] = self.getPlume1(i,Verbose=Verbose,**kwopts)
        if Verbose:
            print "  --------------------|----------|-------------------|-------------------------------------------"

        return array(z_plume)

        if Verbose:
            print "  --------------------|----------|-------------------|----------"

#---
    def getFires(self,mod14_path='/home/adasilva/iesa/aerosol/data/MODIS/Level2/MOD14',
                      method='classic',
                      npzFile=None,Verbose=True):
        """
        Retrieves Level2 MOD14 fire data for each MINX fire.
        """
        from dozier import DOZIER

        self.mod14 = []
        d2r = pi / 180.
        a = MAPL_RADIUS/1000. # Earth radius in km
        dt = timedelta(seconds=5*60)

        self.mod14 = MOD14(self.N)

        if Verbose:
            print ""
            print "                   Fire Heat Flux Estimates"
            print ""
            print "  --------------------|----------|-------------------|----------"
            print "                      | Distance |    FRP Estimates  |   Fire"
            print "    MINX Date/Time    |  to Fire |   MINX      MODIS | Heat Flux"
            print "                      |    km    |    MW         MW  |  kW/m2"
            print "  --------------------|----------|-------------------|----------"

        for i in range(self.N):

            # Get fire granule for this particular day
            # ----------------------------------------
            t = self.tyme[i]
            p = mod14_path
            m = DOZIER(_getGran(t-dt,p) + _getGran(t,p) + _getGran(t+dt,p))

            # Select closest fires
            # --------------------
            x0 = cos(d2r*self.lat_f[i]) * cos(d2r*self.lon_f[i])
            y0 = cos(d2r*self.lat_f[i]) * sin(d2r*self.lon_f[i])
            z0 = sin(d2r*self.lat_f[i])
            dx = x0*cos(d2r*m.lat) * cos(d2r*m.lon)
            dy = y0*cos(d2r*m.lat) * sin(d2r*m.lon)
            dz = z0*sin(d2r*m.lat)
            s  = a * arccos(dx+dy+dz) # great circle distance
            j = s.argmin()

            # Estimate fire heat flux
            # -----------------------
            if method=='bimodal':
                m.bimodal_u()   # bimodal Dozier, need to provide trasnmittance
                self.mod14.hflux[i] = m.h_F[j] # flaming Radiative Heat Flux
                self.mod14.fdist[i] = s[j]
                self.mod14.frp[i]  = m.pow_F[j]
                self.mod14.pixar[i] = m.pixar[j]
                self.mod14.farea[i] = m.r_F[j]
                self.mod14.qa[i] = m.m[j] # bolean
            else:
                m.classic_var() # classic Dozier, need to provide trasnmittance
                self.mod14.hflux[i] = m.hflux[j] # Radiative Heat Flux
                self.mod14.fdist[i] = s[j]
                self.mod14.frp[i]  = m.pow[j]
                self.mod14.pixar[i] = m.pixar[j]
                self.mod14.farea[i] = m.farea[j]
                self.mod14.qa[i] = m.m[j] # bolean
            if Verbose:
                print " ", t, "| %8.2f | %8.2f %8.2f | %8.2f"%\
                    (self.mod14.fdist[i], self.mod14.frp[i], self.mod14.frp[i], \
                     self.mod14.hflux[i] )

        # Save fire properties in NPZ file for later
        # ------------------------------------------
        if npzFile!=None:
            savez(npzFile,**self.mod14.__dict__)

        if Verbose:
            print "  --------------------|----------|-------------------|----------"

#---
    def getOptBrute(self,I=None,Verbose=True,**kwopts):

        """
        Runs the Plume Rise extension to compute brute force optimal value of (hflux,area)
        to match MISR plume height.
	    """

        z_opt = __AMISS__ * ones(self.N)
        h_opt = __AMISS__ * ones(self.N)
        a_opt = __AMISS__ * ones(self.N)

        R = arange(self.N)
        if I is not None:
            R = R[I]

        Hflux_kW = linspace(1,100,10)
        Area = linspace(0.1e4,2e4,10)

        print Hflux_kW
        print Area

        if Verbose:
            print ""
            print "                    Plume Height Estimates"
            print ""
            print "  --------------------|----------|-------------------|----------"
            print "                      | Observed |    Opt Properties |  Optimal "
            print "    MINX Date/Time    |  Height  |   Area   Heat Flx |  Height"
            print "                      |    km    |    ha      kW     |    km"
            print "  --------------------|----------|-------------------|----------"

#       Loop over time
#       --------------
        for i in R:
            z = m.z[i]
            e = 1e20 # overestimate
            for hflux_kW in Hflux_kW:
                for area in Area:
                    z_ = self.getPlume1(i,Verbose=False,hflux_kW=hflux_kW,area=area)
                    if isnan(z)==False:
                        e_ = (z-z_)**2
                        if e_<e:
                            e = e_
                            z_opt[i] = z_
                            h_opt[i] = hflux_kW
                            a_opt[i] = area
            if e<1e20:
                if Verbose:
                    print " ", self.tyme[i], "| %8.2f | %8.2f %8.2f | %8.2f"%\
                          (z, a_opt[i]/1e4, h_opt[i], z_opt[i])

        if Verbose:
            print "  --------------------|----------|-------------------|----------"

        return (z_opt,h_opt,a_opt)

#---
    def getOptF(self,Verbose=True,area_m2=None):
        """
        Find optimal fire modified bstar to match observed plume height.
        """

        if Verbose:
            print ""
            print "    Fire Modified bstar Optimization"
            print ""
            print "Plume | f_opt  |    J     |  Ni | Nf"
            print "------|--------|----------|-----|-----"

        self.f_opt = ones(self.N)
        self.z_opt = ones(self.N)
        for i in range(self.N):
#            xmin, fval, iter, fcalls = opt.brent(CostFuncF,args=(self,i),brack=(1.,4),full_output=True)
            xmin, fval, iter, fcalls = opt.brent(CostFuncF,args=(self,i,area_m2),full_output=True)
            if isnan(fval):
                self.f_opt[i] = NaN
                self.z_opt[i] = NaN
            else:
                rad2conv = xmin**2
                self.f_opt[i] = rad2conv
                self.z_opt[i] = self.getPlume1(i,rad2conv=self.f_opt[i],area_m2=area_m2)

            if Verbose:
                print "%5d | %6.2f | %8.2f | %3d | %3d "%(i, rad2conv, fval, iter, fcalls)

        if Verbose:
            print "------|--------|----------|-----|-----"

#---
    def getOptA(self,Verbose=True,rad2conv=5):
        """
        Find optimal fire modified afac to match observed plume height.
        """

        if Verbose:
            print ""
            print "    Fire Modified afac Optimization"
            print ""
            print "Plume | f_opt  |    J     |  Ni | Nf"
            print "------|--------|----------|-----|-----"

        self.a_opt = ones(self.N)
        self.z_opt = ones(self.N)
        for i in range(self.N):
            xmin, fval, iter, fcalls = opt.brent(CostFuncA,args=(self,i,rad2conv),full_output=True)
            afac = xmin**2
            self.a_opt[i] = afac
            self.z_opt[i] = self.getPlume1(i,afac=self.a_opt[i],rad2conv=rad2conv)
            if Verbose:
                print "%5d | %6.2f | %8.2f | %3d | %3d "%(i, afac, fval, iter, fcalls)

        if Verbose:
            print "------|--------|----------|-----|-----"

#---
    def getOptFanneal(self,Verbose=True):
        """
        Find optimal fire hflux scaling to match observed plume height.
        """

        if Verbose:
            print ""
            print "    Fire Modified bstar Optimization"
            print ""
            print "Plume | f_opt  |    J     |  Ni | Nf"
            print "------|--------|----------|-----|-----"

        self.f_opt = ones(self.N)
        self.z_opt = ones(self.N)
        for i in range(self.N):
            xmin, fval, T, fcalls, iter, accept, retval = opt.anneal(CostFuncF,1.0,args=(self,i),full_output=True)
            ffac = xmin**2
            self.f_opt[i] = ffac
            self.z_opt[i] = self.getPlume1(i,ffac=self.f_opt[i])
            if Verbose:
                print "%5d | %6.2f | %8.2f | %3d | %3d "%(i, ffac, fval, iter, fcalls)

        if Verbose:
            print "------|--------|----------|-----|-----"

#---
    def getOptFbnd(self,Verbose=True):
        """
        Find optimal fire hflux scaling to match observed plume height.
        """

        if Verbose:
            print ""
            print "    Fire Modified bstar Optimization"
            print ""
            print "Plume | f_opt  |    J     |  Ni | Nf"
            print "------|--------|----------|-----|-----"

        self.f_opt = ones(self.N)
        self.z_opt = ones(self.N)
        for i in range(self.N):
            xmin,fval,ier,fcalls  = opt.fminbound(CostFuncF,0.,2.,args=(self,i),full_output=True)
            ffac = xmin**2
            self.f_opt[i] = ffac
            self.z_opt[i] = self.getPlume1(i,rad2conv=self.f_opt[i])
            if Verbose:
                print "%5d | %6.2f | %8.2f | %3d | %3d "%(i, ffac, fval, fcalls, fcalls)

        if Verbose:
            print "------|--------|----------|-----|-----"

#----
def CostFuncF(f,m,i,area_m2):
    z = m.getPlume1(i,rad2conv=f**2,area_m2=area_m2)
    return 1e-6 * (m.z[i]-z)**2

def CostFuncA(f,m,i,rad2conv):
    z = m.getPlume1(i,afac=f**2,rad2conv=rad2conv)
    return 1e-6 * (m.z[i]-z)**2

#----------------------------------------------------------------------------------------------

class MOD14(object):
    def __init__(self,N):
        """
        Simple container class for fire properties.
        """
        self.hflux = __AMISS__ * ones(N)  # fire heat flux estimate
        self.fdist = __AMISS__ * ones(N)  # distance to detected fire
        self.frp   = __AMISS__ * ones(N)  # nearest MODIS FRP
        self.pixar = __AMISS__ * ones(N)  # pixel area
        self.farea = __AMISS__ * ones(N)  # fire area
        self.qa    = __AMISS__ * ones(N)  # quality flag
        return

#----------------------------------------------------------------------------------------------
class PLUME_L2(DOZIER):

    """
    Extension of the MxD14,IGBP and DOZIER classes, adding the
    Plume Rise functionality. This class handles non-gridded,
    observation location fires.
    """

    def getPlume1(self,i,Verbose=False,rad2conv=5,area_m2=None):
        """
        Compute plume height for the ith fire.
        """

        # Fire properties
        # ---------------
        if area_m2 is None:
            area = 1e6 * self.sample.farea[i] # km2 --> m2
        else:
            area = area_m2
        hflux_kW = 1e3 * rad2conv * self.sample.pow[i] / area

        # Meteorology
        # -----------
        ptop = 1. # Pascal
        u = self.sample.u[i]
        v = self.sample.v[i]
        T = self.sample.t[i]
        q = self.sample.qv[i]
        delp = self.sample.delp[i]

        if delp.min()<=0 or T.min()<=0:
            print "out of range: ", delp.min(), T.min()
            return None

        # Ensure arrays are top-down as in GEOS-5
        # ---------------------------------------
        if delp[0] > delp[-1]:
            u = u[::-1]
            v = v[::-1]
            T = T[::-1]
            q = q[::-1]
            delp = delp[::-1]

        # Run plume rise model
        # --------------------
        z_i,z_d,z_a,z_f,z,w,rc = plumevmd(u,v,T,q,delp,ptop,hflux_kW,area)
        if rc:
            raise ValueError, "error on return from <plume>, rc = %d"%rc

        if Verbose:
            print "| %8.2f %8.2f %s | %8.2f %8.2f | %8.2f %8.2f %8.2f %8.2f | %03d"%\
                (self.lon[i],self.lat[i], str(self.tyme[i]), area/1e4, hflux_kW, z_i,z_d,z_a,z_f,i)

        return (z_i,z_d,z_a,z_f,z,w)

#---
    def getPlume(self,algo='dozier',Verbose=False,**kwargs):

        """
        Runs the Plume Rise extension to compute the extent of the plume for each fire.
        It is assumed that the necessary met fields have already been loaded in attribute
        *sample* by method sampleFP or alternative.

        """

        I = self.sample.I             # spatial subsetting
        if len(I) != self.lon.size:
                raise ValueError, 'sampling data appear inconsistent'

        N = len(self.lon)     # all obs
        n = len(self.lon[I])  # reduced set in case of regional subsetting
        self.sample.z_i = __AMISS__ * ones(n,dtype='float32')
        self.sample.z_d = __AMISS__ * ones(n,dtype='float32')
        self.sample.z_a = __AMISS__ * ones(n,dtype='float32')
        self.sample.z_f = __AMISS__ * ones(n,dtype='float32')

        # Gather fire properties
        # ----------------------
        self.sample.pow   = self.pow[I]
        self.sample.m     = self.m[I]

        # Assume that fire area has alread been estimated by classic_var()
        # ----------------------------------------------------------------
        if algo is not None:
            self.sample.farea = self.farea[I]
            if self.algo != algo:
                raise ValueError, 'only dozier algorithm supportted'

        if Verbose:
            print ""
            print "                         Plume Height Estimates"
            print ""
            print "  --------------------|-----------------------|-------------------------------------------"
            print "                      |    Fire Properties    |             Plume Height"
            print "    Lon  Lat  Time    |   Area     Conv Pwr   |    z_i      z_d      z_a      z_f"
            print "                      |    ha         kW      |    km       km       km       km"
            print "  --------------------|-----------------------|-------------------------------------------"

#       Loop over time
#       --------------
        s = self.sample # shorthand
        for i in range(n):

#           Interpolate met fields to fire location
#           ---------------------------------------
            if not self.sample.m[i]: continue   # skip over bad points

#           Compute plume rise for this fire
#           --------------------------------
            s.z_i[i],s.z_d[i],s.z_a[i],s.z_f[i],z,q = self.getPlume1(i,Verbose=Verbose,**kwargs)

#---
    def getPlumeDeprecated(self,met):

        """
        Runs the Plume Rise extension to compute the extent of the plume.
        On input,

        met --- MET object for computing Met fields

        Notice that the Dozier must have been run and produced the farea
        attribute. In addition, the veg attribute with the biome type must
        have been defined as well.
	"""

        if self.algo != 'dozier':
            raise ValueError, 'only Dozier algorithm supported for now'

        if self.veg is None:
            raise ValueError, 'veg attribute with biome type has not been defined'

#       Initialize Plume Rise ranges
#       ----------------------------
        ntd = met.ntd
        N = self.lon.size
        self.p_plume = zeros((ntd,N,2))
        self.k_plume = zeros((ntd,N,2))
        self.z_plume = zeros((ntd,N,2))
        yyyy = self.yyyy[N/2]
        jjj = self.jjj[N/2]

#       Loop over time
#       --------------
        for t in range(1,ntd+1):

#           Interpolate met fields to fire location
#           ---------------------------------------
            met.interp(yyyy,jjj,t,lon=self.lon,lat=self.lat)

#           Compute plume rise for this time
#           --------------------------------
            farea = self.r_F * self.farea # reduce area by flaming fraction
            p, z, k = getPlume(farea,self.veg,met,t,ntd,self.verb)

#           Save in the approapriate containers
#           -----------------------------------
            self.p_plume[t-1,:,:] = p[:,:]
            self.z_plume[t-1,:,:] = z[:,:]
            self.k_plume[t-1,:,:] = k[:,:]

#...............................................................................

class PLUME_L3(object):

    """
    Extension of the MxD14,IGBP and DOZIER classes, adding the
    Plume Rise functionality. This class handles non-gridded,
    observation location fires.
    """

    def __init__(self,plume,refine=4,res=None):
        """
        Create a gridded Plume Rise object from a Level 2
        PLUME_L2 object *plume*. The grid resolution is
        specified by

        refine  -- refinement level for a base 4x5 GEOS-5 grid
                       refine=1  produces a   4  x  5    grid
                       refine=2  produces a   2  x2.50   grid
                       refine=4  produces a   1  x1,25   grid
                       refine=8  produces a  0.50x0.625  grid
                       refine=16 produces a  0.25x0.3125 grid

        Alternatively, one can specify the grid resolution with a
        single letter:

        res     -- single letter denoting GEOS-5 resolution,
                       res='a'  produces a   4  x  5    grid
                       res='b'  produces a   2  x2.50   grid
                       res='c'  produces a   1  x1,25   grid
                       res='d'  produces a  0.50x0.625  grid
                       res='e'  produces a  0.25x0.3125 grid

                   NOTE: *res*, if specified, supersedes *refine*.

        After initialization only the FRP weighted average area is
        set, for each gridbox/biome, along with the coordinates of the
        global grid.

        """

        N = plume.lon.size
        self.verb = plume.verb

#       Output grid resolution
#       ----------------------
        if res is not None:
            if res=='a': refine = 1
            if res=='b': refine = 2
            if res=='c': refine = 4
            if res=='d': refine = 8
            if res=='e': refine = 16

#       Lat lon grid
#       ------------
        dx = 5. / refine
        dy = 4. / refine
        im = int(360. / dx)
        jm = int(180. / dy + 1)
        self.im = im
        self.jm = jm
        self.glon = linspace(-180.,180.,im,endpoint=False)
        self.glat = linspace(-90.,90.,jm)
        Lat, Lon  = meshgrid(self.glat,self.glon)  # shape should be (im,jm)

        self.yyyy = plume.yyyy[N/2]
        self.jjj  = plume.jjj[N/2]
        self.date = date((int(self.yyyy),1,1)) + (int(self.jjj) - 1)*DAY
        self.col  = plume.col

#       Supperobed fire attributes for each biome
#       These will have 1D arrays for each biome, using
#       a standard sparse matrix storage
#       -----------------------------------------------
        self.bioma = [1,2,3,4]
        NONE       = [None,None,None,None]
        self.idx   = NONE[:]  # non-zero indices for each biome
        self.area  = NONE[:]  # non-zero areas   for
        self.r_F   = NONE[:]  # corresponding flaming fraction
        self.lon   = NONE[:]  # corresponding lon
        self.lat   = NONE[:]  # corresponding lat

#       Plume extent to be filled later
#       -------------------------------
        self.p_plume = NONE[:]
        self.k_plume = NONE[:]
        self.z_plume = NONE[:]

#       Grid box average of fire flaming area, weighted by FRP
#       -----------------------------------------------------
        for b in self.bioma:

#           Compute average area in gridbox for this biome
#           ----------------------------------------------
            Frac = zeros((im,jm))
            Area = zeros((im,jm))
            FRP  = zeros((im,jm))
            i = (plume.veg==b)
            if any(i):
                blon = plume.lon[i]
                blat = plume.lat[i]
                bfrac = plume.r_F[i] * plume.pow[i]
                barea = plume.r_F[i] * plume.farea[i] * plume.pow[i] # notice flaming fraction
                bfrp  = plume.pow[i]
                Area +=  binareas(blon,blat,barea,im,jm) # to be normalized
                Frac +=  binareas(blon,blat,bfrac,im,jm) # to be normalized
                FRP  +=  binareas(blon,blat,bfrp, im,jm)
                I = (FRP>0.0)
                if any(I):
                    Area[I] = Area[I] / FRP[I]
                    Frac[I] = Frac[I] / FRP[I]

#           Use sparse matrix storage scheme
#           --------------------------------
            I = Area.nonzero()
            if any(I):
                self.area[b-1] = Area[I]
                self.r_F[b-1]  = Frac[I] # average flaming/total energy fraction
                self.lon[b-1]  = Lon[I]
                self.lat[b-1]  = Lat[I]
                self.idx[b-1]  = I       # save indices for going to global grid

    def getPlume(self):

        """
        Runs the Plume Rise extension to compute the extent of the plume.
        On input,

        met --- MET object for computing Met Fields at obs location.

	"""

        ntd = met.ntd # number of time steps per day
        self.ntd = ntd

#       Loop over bioma
#       ---------------
        for i in range(len(self.bioma)):

#           No data for this bioma, nothing to do
#           -------------------------------------
            if self.idx[i] is None:
                if self.verb>0:
                    print "[x] no data for %s"%Bioma[i]
                continue

            lon = self.lon[i]
            lat = self.lat[i]
            area = self.area[i]
            N = lon.size
            veg = self.bioma[i] * ones(N)

            p_plume = zeros((ntd,2,N))
            z_plume = zeros((ntd,2,N))
            k_plume = zeros((ntd,2,N))

            if self.verb>0:
                print "[ ] got %d burning gridboxes in %s"%(N,Bioma[i])

#           Loop over time
#           --------------
            for t in range(1,ntd+1):

#               Interpolate met fields to fire locations
#               ----------------------------------------
                met.interp(self.yyyy,self.jjj,t,lon=lon,lat=lat)

#               Compute plume rise for this time, biome
#               ---------------------------------------
                p, z, k = getPlume(area,veg,met,t,ntd,self.verb)

#               Save in the approapriate containers
#               -----------------------------------
                p_plume[t-1,:,:] = p.T[:,:]
                z_plume[t-1,:,:] = z.T[:,:]
                k_plume[t-1,:,:] = k.T[:,:]

#           Plume extent for this biome (sparse storage)
#           --------------------------------------------
            self.p_plume[i] = p_plume
            self.z_plume[i] = z_plume
            self.k_plume[i] = k_plume

#---
    def write(self,filename=None,dir='.',expid='qfed2',tag=None):
       """
       Writes gridded Area and FRP to file.
       """

       vtitle = {}
       vtitle['fa'] = 'Flaming Area'
       vtitle['ff'] = 'Fraction of Flaming Energy'
       vtitle['p2'] = 'Plume Bottom Pressure'
       vtitle['p1'] = 'Plume Top Pressure'
       vtitle['z2'] = 'Plume Bottom Height'
       vtitle['z1'] = 'Plume Top Height'
       vtitle['k2'] = 'Plume Bottom Vertical Index'
       vtitle['k1'] = 'Plume Top Vertical Index'

       vunits = {}
       vunits['fa'] = 'km2'
       vunits['ff'] = '1'
       vunits['p2'] = 'Pa'
       vunits['p1'] = 'Pa'
       vunits['z2'] = 'meter'
       vunits['z1'] = 'meter'
       vunits['k2'] = '1'
       vunits['k1'] = '1'

       btitle = {}
       btitle['tf'] = 'Tropical Forest'
       btitle['xf'] = 'Extra-Tropical Forest'
       btitle['sv'] = 'Savanna'
       btitle['gl'] = 'Grassland'

#      Create master variable list
#      ---------------------------
       Vname  = []
       Vtitle = []
       Vunits = []
       for v in vtitle.keys():
           vt = vtitle[v]
           vu = vunits[v]
           for b in btitle.keys():
               bt = btitle[b]
               Vname.append(v+'_'+b)
               Vtitle.append(vt+' ('+bt+')')
               Vunits.append(v)

#      Global metadata
#      ---------------
       title = 'QFED Level3c v%3.1f (%s) Gridded Plume Rise Estimates' % (__VERSION__, _getTagName(tag))
       source = 'NASA/GSFC/GMAO GEOS-5 Aerosol Group'
       contact = 'arlindo.dasilva@nasa.gov'

#      Time/date handling
#      ------------------
       if self.date is None:
           print "[x] did not find matching files, skipped writing an output file"
           return

       if 24%self.ntd != 0:
           raise ValueError,"invalid number of times per day (%d),"%self.ntd\
                 +"it must be a divisor of 24."
       else:
           dT = 240000/self.ntd # timestep in hhmmss format
           NHMS = range(0,240000,dT)

       nymd = 10000*self.date.year + 100*self.date.month + self.date.day
       nhms = NHMS[0]
       col = self.col

#      Create output file name
#      -----------------------
       if filename is None:
           filename = '%s/%s.plumerise.%s.%d.nc'%(dir,expid,col,nymd)
       self.filename = filename
       f = GFIO()
       f.create(filename,Vname, nymd, nhms,
                lon=self.glon, lat=self.glat,
                vtitle=Vtitle, vunits=Vunits,
                timinc=dT, amiss=__AMISS__,
                title=title, source=source, contact=contact)

#      Write out Plume Rise variables
#      ------------------------------
       d = (self.im,self.jm)
       for t in range(self.ntd):
           nhms = NHMS[t]
           b = 0
           for bn in btitle.keys():
               I = self.idx[b]
               f_area  = self.area[b]
               f_frac  = self.r_F[b]
               p_plume = self.p_plume[b]
               z_plume = self.z_plume[b]
               k_plume = self.k_plume[b]
               _writeOne(f,'fa_'+bn,nymd,nhms,I,f_area, t,0,d)
               _writeOne(f,'ff_'+bn,nymd,nhms,I,f_frac, t,0,d)
               _writeOne(f,'p1_'+bn,nymd,nhms,I,p_plume,t,0,d)
               _writeOne(f,'p1_'+bn,nymd,nhms,I,p_plume,t,0,d)
               _writeOne(f,'p2_'+bn,nymd,nhms,I,p_plume,t,1,d)
               _writeOne(f,'z1_'+bn,nymd,nhms,I,z_plume,t,0,d)
               _writeOne(f,'z2_'+bn,nymd,nhms,I,z_plume,t,1,d)
               _writeOne(f,'k1_'+bn,nymd,nhms,I,k_plume,t,0,d)
               _writeOne(f,'k2_'+bn,nymd,nhms,I,k_plume,t,1,d)
               b += 1

       try:
           f.close()
       except:
           pass

       if self.verb >=1:
           print "[w] Wrote file "+filename

#..............................................................................

#
#                                    Static Methods
#                                    --------------
#

def _writeOne(f,vname,nymd,nhms,I,S,t,k,d):
    """
    Write one sparse variable to a GFIO file.
    """
    A = zeros(d) + __AMISS__
    if I is not None:
        if len(S.shape)==3:
            A[I] = S[t,k,:]
        elif len(S.shape)==1:
            A[I] = S[:]
        else:
            raise ValueError, 'invalid S rank = %d'%len(S.shape)

    f.write(vname,nymd,nhms,A)

def getPlumeDeprecated(farea,veg,met,t,ntd,Verb=0):

    """
    Runs the Plume Rise extension to compute the extent of the plume.

          p, z, k = getPlume(farea,veg,met,t,ntd)

    where p, z and k are nd-arrays of shape(N,2), N being the
    number of observations (as in met.lon.size). On input,

    farea --- (flaming) fire area
    veg   --- biome type

    """

    N = met.lon.size
    nominal_area = 1.e6 # 1 km^2: pixar and farea are in units of km2
    km = met.lev.size
    ptop = met.ptop
    ktop = met.ktop # 1-offset

    if Verb:
        if N>100:
            Np = range(0,N,N/10)
        elif N>10:
            Np = range(0,N,N/10)
        else:
            Np = range(N)
        print ""
        print "                   Plume Rise Estimation for t=%d"%t
        print "                   ------------------------------"
        print ""
        print "  %  |    Lon    Lat  b |   p_bot    p_top  |  z_bot z_top  |  k   k"
        print "     |    deg    deg    |    mb       mb    |   km     km   | bot top"
        print "---- |  ------ ------ - | -------- -------- | ------ ------ | --- ---"

#   Allocate space
#   --------------
    p_plume = zeros((N,2))
    k_plume = zeros((N,2))
    z_plume = zeros((N,2))

#   Compute plume extent, one fire at a time
#   ----------------------------------------
    for i in range(N):

        u = met.fields['u'][i]
        v = met.fields['v'][i]
        T = met.fields['t'][i]
        q = met.fields['qv'][i]
        delp = met.fields['delp'][i]

#       Ensure arrays are top-down as in GEOS-5
#       ---------------------------------------
        if delp[0] > delp[-1]:
            u = u[::-1]
            v = v[::-1]
            T = T[::-1]
            q = q[::-1]
            delp = delp[::-1]

#       Units:
#           farea - km2 (must multiply by nominal area for m2)
#            area  - m2 as required by plume rise model
#       ------------------------------------------------------
        area = farea[i] * nominal_area
        veg_ = veg[i]

#       Run plume rise model
#       --------------------
        p1, p2, z1, z2, k1, k2, rc = \
            biome(u, v, T, q, delp, ptop, area, veg_)

        k1, k2 = (k1+ktop-1, k2+ktop-1)

        p_plume[i,:] = (p1, p2)
        k_plume[i,:] = (k1, k2)
        z_plume[i,:] = (z1, z2)

        if Verb:
            if i in Np:
                ip = int(0.5+100.*i/N)
                print "%3d%% | %7.2f %6.2f %d | %8.2f %8.2f | %6.2f %6.2f | %3d %3d "%\
                      (ip,met.lon[i],met.lat[i],veg[i], \
                       p2/100,p1/100,z2/1000,z1/1000,k2,k1)

    return (p_plume, z_plume, k_plume)

def _getTagName(tag):
    if tag != None:
        tag_name = tag
    else:
        if __CVSTAG__ not in (None, ''):
            tag_name = __CVSTAG__
        else:
            tag_name = 'unknown'

    return tag_name

#----
def _getGran(t,mod14_path):
    doy = t.date().toordinal() - date(t.year-1,12,31).toordinal()
    hhmm = "%02d%02d"%(t.hour,5*(t.minute/5))
    patt = mod14_path+'/%4d/%03d/MOD14.A%4d%03d.%s.005.*.hdf'%(t.year,doy,t.year,doy,hhmm)
#    print '--> ', patt
    return glob(patt)

#---

#..............................................................................

def plotMINX(m,imfile=None,Title=None):

    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    I_ = m.z>0
    I = m.z>(m.sample.pblh+250)
    J = m.zm>(m.sample.pblh+250)
    K = m.zt>(m.sample.pblh+250)

    figure(dpi=120)
    plot([0,5],[0,5],'k') # 1:1 line
    plot(m.sample.pblh[I_]/1000,m.z[I_]/1000, 'bo',label='Mode Height < PBL')
    plot(m.sample.pblh[I]/1000,m.z[I]/1000, 'co',label='Mode Height > PBL')
    plot(m.sample.pblh[K]/1000,m.zt[K]/1000, 'ro',label='95%-ile Height > PBL')

    print "Percent above PBL: ",100.*len(m.z[I])/len(m.z[I_])

    x, ya, yb = m.sample.pblh[K]/1000, m.zt[K]/1000, m.z[K]/1000
    for i in range(len(x)):
        plot([x[i],x[i]],[ya[i],yb[i]],'k')

    xlabel('GEOS-5 PBL Height AGL [km]')
    ylabel('MINX Plume Height AGL [km]')
    legend(loc='upper right',fontsize='medium')
    grid()

    if Title is not None:
        title(Title)

    if imfile is not None:
        savefig(imfile,bbox_inches='tight')

def plotAreaFRP(m,z_plume,imfile=None,Title=None):
    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    figure(dpi=120)
    J = z_plume>(m.sample.pblh+250)
    area = m.mod14.farea * m.mod14.pixar * 1e6 / 1e4
    plot(area,m.mod14.frp,'bo')
    plot(area[J],m.mod14.frp[J],'ro')
    grid()
    xlabel(r'Fire Area [Ha]')
    ylabel('FRP/Area [MW]')

def plotPR(m,z_plume,imfile=None,Title=None):
    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    I = m.z>(m.sample.pblh+250)
    J = z_plume>(m.sample.pblh+250)

    figure(dpi=120)
    plot([0.5,4],[0.5,4],'k')
    plot(m.sample.pblh/1000,z_plume/1000, 'bo',label='Modeled Below PBL')
    plot(m.sample.pblh[J]/1000,z_plume[J]/1000,'co',label='Modeled above PBL')
    #plot(m.sample.pblh[I]/1000,z_plume[I]/1000,'ro',label='Observed above PBL')
    plot(m.sample.pblh[I]/1000,m.z[I]/1000,'ro',label='Observed above PBL')

    x, ya, yb = m.sample.pblh[I]/1000, z_plume[I]/1000, m.z[I]/1000
    for i in range(len(x)):
        plot([x[i],x[i]],[ya[i],yb[i]],'k')

    xlabel('GEOS-5 PBL Height AGL [km]')
    ylabel('Plume Height AGL [km]')
    legend(loc='lower right')
    grid()

    if Title is not None:
        title(Title)

    if imfile is not None:
        savefig(imfile,bbox_inches='tight')

def scatOpt(m,imfile=None,Title=None,norm=True,xymax=1.9999):

    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    z     = m.z/m.sample.pblh
    z_opt = m.z_opt/m.sample.pblh

    I = z>0
    J = z_opt>1
    K = (z>1)&(z_opt>1)

    if not norm:
        z     = m.z/1000
        z_opt = m.z_opt/1000

    figure(dpi=120)
    if norm:
        plot([0,xymax],[0,xymax],'k')
        plot([0,xymax],[1,1],'k')
        plot([1,1],[0,xymax],'k')
    plot(z[I],z_opt[I],'yo',label=r'$z_{OBS}<z_{PBL}$, $z_{OPT}<z_{PBL}$')
    plot(z[J],z_opt[J],'ro',label=r'$z_{OBS}<z_{PBL}$, $z_{OPT}>z_{PBL}$')
    plot(z[K],z_opt[K],'go',label=r'$z_{OBS}>z_{PBL}$, $z_{OPT}>z_{PBL}$')

    if norm:
        xlabel(r'$z_{OBS}/z_{PBL}$')
        ylabel(r'$z_{OPT}/z_{PBL}$')
    else:
        xlabel(r'$z_{OBS}$')
        ylabel(r'$z_{OPT}$')

    legend(loc='lower right')
    grid()

    if Title is not None:
        title(Title)

    if imfile is not None:
        savefig(imfile,bbox_inches='tight')

def scatOptF(m,imfile=None,Title=None,norm=True,xymax=1.9999):

    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig, loglog

    z     = m.z/m.sample.pblh
    z_opt = m.z_opt/m.sample.pblh

    I = (z>0)&(isnan(m.f_opt)==False)#&(m.f_opt<100)&(m.mod14.frp<100)
    J = I&(z_opt>1)
    K = I&(z>1)&(z_opt>1)

    if not norm:
        z     = m.z/1000
        z_opt = m.z_opt/1000

    figure(dpi=120)
    loglog(m.mod14.frp[I],m.f_opt[I],'yo',label=r'$z_{OBS}<z_{PBL}$, $z_{OPT}<z_{PBL}$')
    loglog(m.mod14.frp[J],m.f_opt[J],'ro',label=r'$z_{OBS}<z_{PBL}$, $z_{OPT}>z_{PBL}$')
    loglog(m.mod14.frp[K],m.f_opt[K],'go',label=r'$z_{OBS}>z_{PBL}$, $z_{OPT}>z_{PBL}$')

    xlabel('FRP [MW]')
    ylabel(r'$\gamma$')
    #legend(loc='upper right')
    grid()

    if Title is not None:
        title(Title)

    if imfile is not None:
        savefig(imfile,bbox_inches='tight')

def scatPR(m,z_plume,imfile=None,Title=None):

    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    z     = m.z/m.sample.pblh
    z_pr = z_plume/m.sample.pblh

    I = z>0
    J = (z_pr>1)|((z_pr<1)&(z>1))
    K = (z>1)&(z_pr>1)

    figure(dpi=120)
    plot([0,1.5999],[0,1.5999],'k')
    plot([0,2.999],[1,1],'k')
    plot([1,1],[0,1.5999],'k')
    plot(z[I],z_pr[I],'yo',label=r'$z_{OBS}<z_{PBL}$, $z_{PR}<z_{PBL}$')
    plot(z[J],z_pr[J],'ro',label=r'$z_{OBS}<z_{PBL}$, $z_{PR}>z_{PBL}$')
    plot(z[K],z_pr[K],'go',label=r'$z_{OBS}>z_{PBL}$, $z_{PR}>z_{PBL}$')

    xlabel(r'$z_{OBS}/z_{PBL}$')
    ylabel(r'$z_{PR}/z_{PBL}$')

    legend(loc='lower right')
    grid()

    if Title is not None:
        title(Title)

    if imfile is not None:
        savefig(imfile,bbox_inches='tight')

#--

def plotKDE(s,s2=None,Title='',imfile=None):

    from matplotlib.pyplot import plot, legend, xlabel, ylabel, title, grid, figure, savefig

    J = (s.pblh>100)&(s.z_a>0)
    r = s.z_a/s.pblh
    if s2 is not None:
        J2 = (s2.pblh>100)&(s2.z_a>0)
        r2 = s2.z_a/s2.pblh

    # KDE 1d
    # ------
    figure(dpi=120)
    bins, P = kde.calc_kde1d(r[J],range=(0,5))
    plot(bins,P,'b',linewidth=2,label='Terra')
    if s2 is not None:
        bins, P2 = kde.calc_kde1d(r2[J2],range=(0,5))
        plot(bins,P2,'r',linewidth=2,label='Aqua')
        legend(loc='upper right')
    grid()
    xlabel(r'$z_{PR}/z_{PBL}$')
    ylabel('p.d.f.')
    if Title is not None:
        title(Title)
    if imfile is not None:
        savefig('kde1d.'+imfile,bbox_inches='tight')

    if s2 is not None: return

    # KDE 2d
    # ------
    x, y, P = kde.calc_kde2d(s.pblh[J]/1000,s.z_a[J]/1000,x_range=(0,3),y_range=(0,3))
    kde.plot_kde2d(x,y,P,dpi=120,Title=Title,
                   xLabel=r'$z_{PBL}$', yLabel=r'$z_{PR}$')
    grid()
    if imfile is not None:
        savefig('kde2d.'+imfile,bbox_inches='tight')

def aveVMD(s):
    """
    Compute average vertical mass distribution, weighted by FRP.
    """
    J = (s.pblh>100)&(s.z_a>0)

    z_c   = s.z_d[J]/1000
    delta = (s.z_f[J] - s.z_d[J])/1000
    pow = s.pow[J]
    pblh = s.pblh[J]/1000

    nz = 100
    nf = len(pblh)
    z = linspace(0,5,nz)

    # PBL mass distribution
    # ---------------------
    V = zeros(nz)
    for i in range(nf):
        v = zeros(nz)
        v[z<pblh[i]] = 1
        v = pow[i] * v / sum(v)
        V += v
    v_pbl = V / pow.sum()

    return v_pbl

    v = getvmd(z,z_c,delta)

#-------------------------------

if __name__ == "__main__":

    import crtmmodis_

    m = MINXs_PR('/Users/adasilva/workspace/misrPlumes/western_fires_2013/*.txt')
    m.sampleLoadz('/Users/adasilva/workspace/Data_Analysis/AGU-2014/seac4rs_01.npz')
    m.mod14 = NPZ(('/Users/adasilva/workspace/Data_Analysis/AGU-2014/mod14_seac4rs.npz',))

def _allfires():

    print "ok"
    topdir = '/Volumes/ArlindoSD' # SDXC card

    t1 = datetime(2013,8,2)
    t2 = datetime(2013,8,31)

    #for p in ( 'MOD14', 'MYD14'):
    for p in ( 'MYD14',):

        print "Loading fires"
        f = PLUME_L2(None)
        f.restart('%s.fires_nam.%4d-%02d.npz'%(p,t1.year,t1.month))

        f.m = f.qc>0

        print "Loading Meteorology"
        f.sampleLoadz('%s/AGU-2014/%s.sample_nam.%4d-%02d.npz'%(topdir,p,t1.year,t1.month))

        I_na = (f.lon>-170)&(f.lon<-50)&(f.lat>15)&(f.lat<80)
        f.sample.I = I_na

        f.getPlume(algo=None,Verbose=True,area_m2=10e4,rad2conv=5)

        f.sample.pblh = f.sample.pblh[I_na]
        savez('%s.sample_a1_r10_r5.2013-08.npz'%p,**f.sample.__dict__)

def xxxxx():

#--

        Files = granules(t1,t2,product=p,rootdir=topdir+'/MODIS')

        print 'Number of files: ', len(Files)
        f = PLUME_L2(Files,Verb=1)

        print "Computing fire properties"
        ###f.classic_var()

        print "Checkpointing"
        f.checkpoint('%s.fires_nam.%4d-%02d.npz'%(p,t1.year,t1.month))

def minx_test():

    m = MINXs_PR('/Users/adasilva/workspace/misrPlumes/western_fires_2013/*.txt')
    m.sampleLoadz('/Users/adasilva/workspace/Data_Analysis/AGU-2014/seac4rs_01.npz')
    m.mod14 = NPZ(('/Users/adasilva/workspace/Data_Analysis/AGU-2014/mod14_seac4rs.npz',))
    #plotMINX(m)

    #z_plume = m.getPlume()

def arctas():
    # m.getFires(npzFile='mod14_seac4rs.npz')
#    m = MINXs_PR('/Users/adasilva/workspace/misrPlumes/canada2008/Plumes*.txt')

    m = MINXs_PR('/Users/adasilva/workspace/misrPlumes/canada2008/Plumes*.txt')
    m.sampleLoadz('merraero.npz')
    m.mod14 = NPZ('mod14.npz')

#    z_plume = m.getPlume()

def Mac():
    m = MINXs_PR('/Users/adasilva/workspace.local/misrPlumes/canada2008/Plumes*.txt')
#    m.sampleLoadz('/Users/adasilva/workspace.local/misrPlumes/canada2008/merra.npz')
#    m.sampleLoadz('merra.npz')
#    m.mod14 = NPZ('mod14.npz')
#    z_plume = m.getPlume()


# Copyright (c) Anand Patil, 2007

__docformat__ = 'reStructuredText'

import pymc as pm
import linalg_utils
import copy
import types
import numpy as np
from gp_submodel import *
import warnings

from Realization import Realization
from Mean import Mean
from Covariance import Covariance
from GPutils import observe, regularize_array

__all__ = ['wrap_metropolis_for_gp_parents', 'GPEvaluationGibbs', 'GPParentAdaptiveMetropolis']

def wrap_metropolis_for_gp_parents(metro_class):
    """
    Wraps Metropolis step methods so they can handle extended parents of
    Gaussian processes.
    """
    class wrapper(metro_class):
        def __init__(self, stochastic, *args, **kwds):

            self.metro_class.__init__(self, stochastic, *args, **kwds)

            # Remove f from the set that will be used to compute logp_plus_loglike.
            self.markov_blanket_no_f = filter(lambda x: not isinstance(x, GaussianProcess), self.markov_blanket)
            self.fs = filter(lambda x: isinstance(x, GaussianProcess), self.markov_blanket)
            self.fr_checks = [f.submodel.fr_check for f in self.fs]

        def get_logp_plus_loglike(self):
            return pm.logp_of_set(self.markov_blanket_no_f)
        logp_plus_loglike = property(get_logp_plus_loglike)

        def propose(self):
            self.metro_class.propose(self)
            try:
                # First make sure none of the stochastics handled by metro_method forbid their current values.
                for s in self.stochastics:
                    s.logp
                # Then make sure the covariances are all still full-rank on the observation locations.
                for frc in self.fr_checks:
                    frc.logp
                for f in self.fs:
                    f.rand()
                self.f_proposed = True
            except pm.ZeroProbability:
                self.f_proposed = False

        def reject(self):
            self.metro_class.reject(self)
            if self.f_proposed:
                for f in self.fs:
                    f.revert()

        @staticmethod
        def competence(stochastic, metro_class=metro_class):
            if any([isinstance(child, GaussianProcess) for child in stochastic.extended_children]):
                return metro_class.competence(stochastic)+.01
            else:
                return 0

    wrapper.__name__ = 'GPParent%s'%metro_class.__name__
    wrapper.metro_class = metro_class
    wrapper.__doc__ = """A modified version of class %s that handles parents of Gaussian processes.
Docstring of class %s: \n\n%s"""%(metro_class.__name__,metro_class.__name__,metro_class.__doc__)

    return wrapper

# Wrap all registered Metropolis step methods to use GP parents.
new_sm_dict = {}
filtered_registry = filter(lambda x: issubclass(x, pm.Metropolis), pm.StepMethodRegistry)
for sm in filtered_registry:
    wrapped_method = wrap_metropolis_for_gp_parents(sm)
    new_sm_dict[wrapped_method.__name__] = wrapped_method
GPParentAdaptiveMetropolis = wrap_metropolis_for_gp_parents(pm.AdaptiveMetropolis)
__all__ += new_sm_dict.keys()
locals().update(new_sm_dict)

class GPEvaluationGibbs(pm.Metropolis):
    """
    Updates a GP evaluation f_eval. Assumes the only children of f_eval
    are as distributed follows:

    eps_p_f ~ Normal(f_eval[ti], 1./V)

    or

    eps_p_f ~ Normal(f_eval, 1./V)

    if ti is None.
    """
    def __init__(self, submod, V, eps_p_f, ti=None, tally=True, verbose=0):

        self.f_eval = submod.f_eval
        self.f = submod.f
        pm.StepMethod.__init__(self, [self.f, self.f_eval], tally=tally)

        self.children_no_data = copy.copy(self.children)
        if isinstance(eps_p_f, pm.Variable):
            self.children_no_data.discard(eps_p_f)
        else:
            for epf in eps_p_f:
                self.children_no_data.discard(epf)

        self.V = V
        self.C_eval = submod.C_eval
        self.M_eval = submod.M_eval
        self.S_eval = submod.S_eval
        self.eps_p_f = eps_p_f

        M_eval_shape = pm.utils.value(self.M_eval).shape
        C_eval_shape = pm.utils.value(self.C_eval).shape
        self.ti = ti or np.arange(M_eval_shape[0])

        # Work arrays
        self.scratch1 = np.asmatrix(np.empty(C_eval_shape, order='F'))
        self.scratch2 = np.asmatrix(np.empty(C_eval_shape, order='F'))
        self.scratch3 = np.empty(M_eval_shape)

        # Initialize hidden attributes
        self.accepted = 0.
        self.rejected = 0.
        self._state = ['rejected', 'accepted', 'proposal_distribution']
        self._tuning_info = []
        self.proposal_distribution=None

    def get_logp(self):
        return 0.
    logp = property(get_logp)

    def get_loglike(self):
        return pm.utils.logp_of_set(self.children_no_data)
    loglike = property(get_loglike)

    def get_logp_plus_loglike(self):
        return self.get_loglike()
    logp_plus_loglike = property(get_logp_plus_loglike)

    def reject(self):
        self.rejected += 1
        if self.verbose:
            print self._id + ' rejecting'
        # Revert the field evaluation and the rest of the field.
        self.f_eval.revert()
        self.f.revert()

    def tune(self, verbose=0):
        return False

    def propose(self):
        if self.verbose:
            print self._id + ' proposing'

        fc = pm.gp.fast_matrix_copy

        eps_p_f = pm.utils.value(self.eps_p_f)
        f = pm.utils.value(self.f_eval)
        for i in xrange(len(self.scratch3)):
            self.scratch3[i] = np.sum(eps_p_f[self.ti[i]] - f[i])

        # Compute Cholesky factor of covariance of eps_p_f, C(x,x) + V
        C_eval_value = pm.utils.value(self.C_eval)
        C_eval_shape = C_eval_value.shape

        # Get the Cholesky factor of C_eval, plus the nugget.
        # I don't think you can use S_eval for speed, unfortunately.
        in_chol = fc(C_eval_value, self.scratch1)
        for i in xrange(pm.utils.value(C_eval_shape)[0]):
            in_chol[i,i] += pm.utils.value(self.V) / np.alen(self.ti[i])
        info = pm.gp.linalg_utils.dpotrf_wrap(in_chol)
        if info > 0:
            raise np.linalg.LinAlgError

        # Compute covariance of f conditional on eps_p_f.
        offdiag = fc(C_eval_value, self.scratch2)
        offdiag = pm.gp.trisolve(in_chol, offdiag, uplo='U', transa='T', inplace=True)

        C_step = offdiag.T * offdiag
        C_step *= -1
        C_step += C_eval_value

        # Compute mean of f conditional on eps_p_f.
        for i in xrange(len(self.scratch3)):
            self.scratch3[i] = np.mean(eps_p_f[self.ti[i]])
        m_step = pm.utils.value(self.M_eval) + np.dot(offdiag.T, pm.gp.trisolve(in_chol,(self.scratch3 - self.M_eval.value),uplo='U',transa='T')).view(np.ndarray).ravel()

        sig_step = C_step
        info = pm.gp.linalg_utils.dpotrf_wrap(C_step.T)
        if info > 0:
            warnings.warn('Full conditional covariance was not positive definite.')
            return

        # Update value of f.
        self.f_eval.value = m_step+np.dot(sig_step,np.random.normal(size=sig_step.shape[1])).view(np.ndarray).ravel()
        # Propose the rest of the field from its conditional prior.
        self.f.rand()