"""
Classes of variables for equations/terms.
"""
from __future__ import print_function
from __future__ import absolute_import
from collections import deque

import numpy as nm

from sfepy.base.base import (real_types, complex_types, assert_, get_default,
                             output, OneTypeList, Container, Struct, basestr,
                             iter_dict_of_lists)
from sfepy.base.timing import Timer
import sfepy.linalg as la
from sfepy.discrete.functions import Function
from sfepy.discrete.conditions import get_condition_value
from sfepy.discrete.integrals import Integral
from sfepy.discrete.common.dof_info import (DofInfo, EquationMap,
                                            expand_nodes_to_equations,
                                            is_active_bc)
from sfepy.discrete.fem.lcbc_operators import LCBCOperators
from sfepy.discrete.common.mappings import get_physical_qps
from sfepy.discrete.evaluate_variable import eval_real, eval_complex
import six
from six.moves import range

is_state = 0
is_virtual = 1
is_parameter = 2
is_field = 10

def create_adof_conns(conn_info, var_indx=None, active_only=True, verbose=True):
    """
    Create active DOF connectivities for all variables referenced in
    `conn_info`.

    If a variable has not the equation mapping, a trivial mapping is assumed
    and connectivity with all DOFs active is created.

    DOF connectivity key is a tuple ``(primary variable name, region name,
    type, is_trace flag)``.

    Notes
    -----
    If `active_only` is False, the DOF connectivities contain all DOFs, with
    the E(P)BC-constrained ones stored as `-1 - <DOF number>`, so that the full
    connectivities can be reconstructed for the matrix graph creation.
    """
    var_indx = get_default(var_indx, {})

    def _create(var, econn):
        offset = var_indx.get(var.name, slice(0, 0)).start
        if var.eq_map is None:
            eq = nm.arange(var.n_dof, dtype=nm.int32)

        else:
            if isinstance(var, DGFieldVariable):
                eq = nm.arange(var.n_dof, dtype=nm.int32)
            else:
                if active_only:
                    eq = var.eq_map.eq
                else:
                    eq = nm.arange(var.n_dof, dtype=nm.int32)
                    eq[var.eq_map.eq_ebc] = -1 - (var.eq_map.eq_ebc + offset)
                    eq[var.eq_map.master] = eq[var.eq_map.slave]

        adc = create_adof_conn(eq, econn, var.n_components, offset)

        return adc

    def _assign(adof_conns, info, region, var, field, is_trace):
        key = (var.name, region.name, info.dc_type.type, is_trace)
        if not key in adof_conns:
            econn = field.get_econn(info.dc_type, region, is_trace=is_trace)
            if econn is None: return

            adof_conns[key] = _create(var, econn)

        if info.is_trace:
            key = (var.name, region.name, info.dc_type.type, False)
            if not key in adof_conns:
                econn = field.get_econn(info.dc_type, region, is_trace=False)

                adof_conns[key] = _create(var, econn)

    if verbose:
        output('setting up dof connectivities...')
        timer = Timer(start=True)

    adof_conns = {}

    for key, ii, info in iter_dict_of_lists(conn_info, return_keys=True):
        if info.primary is not None:
            var = info.primary
            field = var.get_field()
            field.setup_extra_data(info.ps_tg, info, info.is_trace)

            region = info.get_region()
            _assign(adof_conns, info, region, var, field, info.is_trace)

        if info.has_virtual and not info.is_trace:
            var = info.virtual
            field = var.get_field()
            field.setup_extra_data(info.v_tg, info, False)

            aux = var.get_primary()
            var = aux if aux is not None else var

            region = info.get_region(can_trace=False)
            _assign(adof_conns, info, region, var, field, False)

    if verbose:
        output('...done in %.2f s' % timer.stop())

    return adof_conns

def create_adof_conn(eq, conn, dpn, offset):
    """
    Given a node connectivity, number of DOFs per node and equation mapping,
    create the active dof connectivity.

    Locally (in a connectivity row), the DOFs are stored DOF-by-DOF (u_0 in all
    local nodes, u_1 in all local nodes, ...).

    Globally (in a state vector), the DOFs are stored node-by-node (u_0, u_1,
    ..., u_X in node 0, u_0, u_1, ..., u_X in node 1, ...).
    """
    if dpn == 1:
        aux = nm.take(eq, conn)
        adc = aux + nm.asarray(offset * (aux >= 0), dtype=nm.int32)

    else:
        n_el, n_ep = conn.shape
        adc = nm.empty((n_el, n_ep * dpn), dtype=conn.dtype)
        ii = 0
        for idof in range(dpn):
            aux = nm.take(eq, dpn * conn + idof)
            adc[:, ii : ii + n_ep] = aux + nm.asarray(offset * (aux >= 0),
                                                      dtype=nm.int32)
            ii += n_ep

    return adc

def expand_basis(basis, dpn):
    """
    Expand basis for variables with several components (DOFs per node), in a
    way compatible with :func:`create_adof_conn()`, according to `dpn`
    (DOF-per-node count).
    """
    n_c, n_bf = basis.shape[-2:]
    ebasis = nm.zeros(basis.shape[:2] + (dpn, n_bf * dpn), dtype=nm.float64)
    for ic in range(n_c):
        for ir in range(dpn):
            ebasis[..., n_c*ir+ic, ir*n_bf:(ir+1)*n_bf] = basis[..., ic, :]

    return ebasis

class Variables(Container):
    """
    Container holding instances of Variable.
    """

    @staticmethod
    def from_conf(conf, fields):
        """
        This method resets the variable counters for automatic order!
        """
        Variable.reset()

        obj = Variables()
        for key, val in six.iteritems(conf):
            var = Variable.from_conf(key, val, fields)

            obj[var.name] = var

        obj.setup_dtype()
        obj.setup_ordering()

        return obj

    def __init__(self, variables=None):
        Container.__init__(self, OneTypeList(Variable),
                           state=set(),
                           virtual=set(),
                           parameter=set(),
                           has_virtual_dcs=False,
                           has_lcbc=False,
                           has_lcbc_rhs=False,
                           has_eq_map=False,
                           ordered_state=[],
                           ordered_virtual=[])

        if variables is not None:
            for var in variables:
                self[var.name] = var

            self.setup_ordering()

        self.setup_dtype()

        self.adof_conns = {}

    def __setitem__(self, ii, var):
        Container.__setitem__(self, ii, var)

        if var.is_state():
            self.state.add(var.name)

        elif var.is_virtual():
            self.virtual.add(var.name)

        elif var.is_parameter():
            self.parameter.add(var.name)

        var._variables = self

        self.setup_ordering()
        self.setup_dof_info()

    def setup_dtype(self):
        """
        Setup data types of state variables - all have to be of the same
        data type, one of nm.float64 or nm.complex128.
        """
        dtypes = {nm.complex128 : 0, nm.float64 : 0}
        for var in self.iter_state(ordered=False):
            dtypes[var.dtype] += 1

        if dtypes[nm.float64] and dtypes[nm.complex128]:
            raise ValueError("All variables must have the same dtype!")

        elif dtypes[nm.float64]:
            self.dtype = nm.float64

        elif dtypes[nm.complex128]:
            self.dtype = nm.complex128

        else:
            self.dtype = None

    def link_duals(self):
        """
        Link state variables with corresponding virtual variables,
        and assign link to self to each variable instance.

        Usually, when solving a PDE in the weak form, each state
        variable has a corresponding virtual variable.
        """
        for ii in self.state:
            self[ii].dual_var_name = None

        for ii in self.virtual:
            vvar = self[ii]
            try:
                self[vvar.primary_var_name].dual_var_name = vvar.name
            except IndexError:
                pass

    def get_dual_names(self):
        """
        Get names of pairs of dual variables.

        Returns
        -------
        duals : dict
           The dual names as virtual name : state name pairs.
        """
        duals = {}
        for name in self.virtual:
            duals[name] = self[name].primary_var_name

        return duals

    def setup_ordering(self):
        """
        Setup ordering of variables.
        """
        self.link_duals()

        orders = []
        for var in self:
            try:
                orders.append(var._order)
            except:
                pass
        orders.sort()

        self.ordered_state = [None] * len(self.state)
        for var in self.iter_state(ordered=False):
            ii = orders.index(var._order)
            self.ordered_state[ii] = var.name

        self.ordered_virtual = [None] * len(self.virtual)
        ii = 0
        for var in self.iter_state(ordered=False):
            if var.dual_var_name is not None:
                self.ordered_virtual[ii] = var.dual_var_name
                ii += 1

    def has_virtuals(self):
        return len(self.virtual) > 0

    def setup_dof_info(self, make_virtual=False):
        """
        Setup global DOF information.
        """
        self.di = DofInfo('state_dof_info')
        for var_name in self.ordered_state:
            self.di.append_variable(self[var_name])

        if make_virtual:
            self.vdi = DofInfo('virtual_dof_info')
            for var_name in self.ordered_virtual:
                self.vdi.append_variable(self[var_name])

        else:
            self.vdi = self.di

    def setup_lcbc_operators(self, lcbcs, ts=None, functions=None):
        """
        Prepare linear combination BC operator matrix and right-hand side
        vector.
        """
        from sfepy.discrete.common.region import are_disjoint
        if lcbcs is None:
            self.lcdi = self.adi
            return

        self.lcbcs = lcbcs

        if (ts is None) or ((ts is not None) and (ts.step == 0)):
            regs = []
            var_names = []
            for bcs in self.lcbcs:
                for bc in bcs.iter_single():
                    vns = bc.get_var_names()

                    regs.append(bc.regions[0])
                    var_names.append(vns[0])
                    if bc.regions[1] is not None:
                        regs.append(bc.regions[1])
                        var_names.append(vns[1])

            for i0 in range(len(regs) - 1):
                for i1 in range(i0 + 1, len(regs)):
                    if ((var_names[i0] == var_names[i1])
                        and not are_disjoint(regs[i0], regs[i1])):
                        raise ValueError('regions %s and %s are not disjoint!'
                                         % (regs[i0].name, regs[i1].name))

        ops = LCBCOperators('lcbcs', self, functions=functions)

        for bcs in self.lcbcs:
            for bc in bcs.iter_single():
                vns = bc.get_var_names()
                dofs = [self[vn].dofs for vn in vns if vn is not None]
                bc.canonize_dof_names(*dofs)

                if not is_active_bc(bc, ts=ts, functions=functions):
                    continue

                output('lcbc:', bc.name)

                ops.add_from_bc(bc, ts)

        aux = ops.make_global_operator(self.adi)
        self.mtx_lcbc, self.vec_lcbc, self.lcdi = aux

        self.has_lcbc = self.mtx_lcbc is not None
        self.has_lcbc_rhs = self.vec_lcbc is not None

    def get_lcbc_operator(self):
        if self.has_lcbc:
            return self.mtx_lcbc

        else:
            raise ValueError('no LCBC defined!')

    def equation_mapping(self, ebcs, epbcs, ts, functions, problem=None,
                         active_only=True):
        """
        Create the mapping of active DOFs from/to all DOFs for all state
        variables.

        Parameters
        ----------
        ebcs : Conditions instance
            The essential (Dirichlet) boundary conditions.
        epbcs : Conditions instance
            The periodic boundary conditions.
        ts : TimeStepper instance
            The time stepper.
        functions : Functions instance
            The user functions for boundary conditions.
        problem : Problem instance, optional
            The problem that can be passed to user functions as a context.
        active_only : bool
            If True, the active DOF info ``self.adi`` uses the reduced (active
            DOFs only) numbering. Otherwise it is the same as ``self.di``.

        Returns
        -------
        active_bcs : set
            The set of boundary conditions active in the current time.
        """
        self.ebcs = ebcs
        self.epbcs = epbcs

        ##
        # Assing EBC, PBC to variables and regions.
        if ebcs is not None:
            self.bc_of_vars = self.ebcs.group_by_variables()

        else:
            self.bc_of_vars = {}

        if epbcs is not None:
            self.bc_of_vars = self.epbcs.group_by_variables(self.bc_of_vars)

        ##
        # List EBC nodes/dofs for each variable.
        active_bcs = set()
        for var_name in self.di.var_names:
            var = self[var_name]
            bcs = self.bc_of_vars.get(var.name, None)

            var_di = self.di.get_info(var_name)
            active = var.equation_mapping(bcs, var_di, ts, functions,
                                          problem=problem)
            active_bcs.update(active)

            if self.has_virtual_dcs:
                vvar = self[var.dual_var_name]
                vvar_di = self.vdi.get_info(var_name)
                active = vvar.equation_mapping(bcs, vvar_di, ts, functions,
                                               problem=problem)
                active_bcs.update(active)

        self.adi = 
DofInfo('active_state_dof_info')