"""
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')
        for var_name in self.ordered_state:
            self.adi.append_variable(self[var_name], active=active_only)

        if self.has_virtual_dcs:
            self.avdi = DofInfo('active_virtual_dof_info')
            for var_name in self.ordered_virtual:
                self.avdi.append_variable(self[var_name], active=active_only)

        else:
            self.avdi = self.adi

        self.has_eq_map = True

        return active_bcs

    def get_matrix_shape(self):
        if not self.has_eq_map:
            raise ValueError('call equation_mapping() first!')

        return (self.avdi.ptr[-1], self.adi.ptr[-1])

    def setup_initial_conditions(self, ics, functions):
        self.ics = ics
        self.ic_of_vars = self.ics.group_by_variables()

        for var_name in self.di.var_names:
            var = self[var_name]

            ics = self.ic_of_vars.get(var.name, None)
            if ics is None: continue

            var.setup_initial_conditions(ics, self.di, functions)

        for var_name in self.parameter:
            var = self[var_name]
            if hasattr(var, 'special') and ('ic' in var.special):
                setter, sargs, skwargs = var._get_setter('ic', functions)

                var.set_data(setter(*sargs, **skwargs))
                output('IC data of %s set by %s()' % (var.name, setter.name))

    def set_adof_conns(self, adof_conns):
        """
        Set all active DOF connectivities to `self` as well as relevant
        sub-dicts to the individual variables.
        """
        self.adof_conns = adof_conns

        for var in self:
            var.adof_conns = {}

        for key, val in six.iteritems(adof_conns):
            if key[0] in self.names:
                var = self[key[0]]
                var.adof_conns[key] = val

                var = var.get_dual()
                if var is not None:
                    var.adof_conns[key] = val

    def create_state_vector(self):
        vec = nm.zeros((self.di.ptr[-1],), dtype=self.dtype)
        return vec

    def create_stripped_state_vector(self):
        vec = nm.zeros((self.adi.ptr[-1],), dtype=self.dtype)
        return vec

    def apply_ebc(self, vec, force_values=None):
        """
        Apply essential (Dirichlet) and periodic boundary conditions
        defined for the state variables to vector `vec`.
        """
        for var in self.iter_state():
            var.apply_ebc(vec, self.di.indx[var.name].start, force_values)

    def apply_ic(self, vec, force_values=None):
        """
        Apply initial conditions defined for the state variables to
        vector `vec`.
        """
        for var in self.iter_state():
            var.apply_ic(vec, self.di.indx[var.name].start, force_values)

    def strip_state_vector(self, vec, follow_epbc=False, svec=None):
        """
        Get the reduced DOF vector, with EBC and PBC DOFs removed.

        Notes
        -----
        If 'follow_epbc' is True, values of EPBC master dofs are not simply
        thrown away, but added to the corresponding slave dofs, just like when
        assembling. For vectors with state (unknown) variables it should be set
        to False, for assembled vectors it should be set to True.
        """
        if svec is None:
            svec = nm.empty((self.adi.ptr[-1],), dtype=self.dtype)
        for var in self.iter_state():
            aindx = self.adi.indx[var.name]
            svec[aindx] = var.get_reduced(vec, self.di.indx[var.name].start,
                                          follow_epbc)
        return svec

    def make_full_vec(self, svec, force_value=None, vec=None):
        """
        Make a full DOF vector satisfying E(P)BCs from a reduced DOF
        vector.

        Parameters
        ----------
        svec : array
            The reduced DOF vector.
        force_value : float, optional
            Passing a `force_value` overrides the EBC values.
        vec : array, optional
            If given, the buffer for storing the result (zeroed).

        Returns
        -------
        vec : array
            The full DOF vector.
        """
        self.check_vector_size(svec, stripped=True)

        if self.has_lcbc:
            if self.has_lcbc_rhs:
                svec = self.mtx_lcbc * svec + self.vec_lcbc

            else:
                svec = self.mtx_lcbc * svec

        if vec is None:
            vec = self.create_state_vector()
        for var in self.iter_state():
            indx = self.di.indx[var.name]
            aindx = self.adi.indx[var.name]
            var.get_full(svec, aindx.start, force_value, vec, indx.start)

        return vec

    def has_ebc(self, vec, force_values=None):
        for var_name in self.di.var_names:
            eq_map = self[var_name].eq_map
            i0 = self.di.indx[var_name].start
            ii = i0 + eq_map.eq_ebc
            if force_values is None:
                if not nm.allclose(vec[ii], eq_map.val_ebc):
                    return False
            else:
                if isinstance(force_values, dict):
                    if not nm.allclose(vec[ii], force_values[var_name]):
                        return False
                else:
                    if not nm.allclose(vec[ii], force_values):
                        return False
            # EPBC.
            if not nm.allclose(vec[i0+eq_map.master], vec[i0+eq_map.slave]):
                return False
        return True

    def get_indx(self, var_name, stripped=False, allow_dual=False):
        var = self[var_name]

        if not var.is_state():
            if allow_dual and var.is_virtual():
                var_name = var.primary_var_name
            else:
                msg = '%s is not a state part' % var_name
                raise IndexError(msg)

        if stripped:
            return self.adi.indx[var_name]
        else:
            return self.di.indx[var_name]

    def check_vector_size(self, vec, stripped=False):
        """
        Check whether the shape of the DOF vector corresponds to the
        total number of DOFs of the state variables.

        Parameters
        ----------
        vec : array
            The vector of DOF values.
        stripped : bool
            If True, the size of the DOF vector should be reduced,
            i.e. without DOFs fixed by boundary conditions.
        """
        if not stripped:
            n_dof = self.di.get_n_dof_total()

            if vec.size != n_dof:
                msg = 'incompatible data size!' \
                      ' (%d (variables) == %d (DOF vector))' \
                      % (n_dof, vec.size)
                raise ValueError(msg)

        else:
            if self.has_lcbc:
                n_dof = self.lcdi.get_n_dof_total()

            else:
                n_dof = self.adi.get_n_dof_total()

            if vec.size != n_dof:
                msg = 'incompatible data size!' \
                      ' (%d (active variables) == %d (reduced DOF vector))' \
                      % (n_dof, vec.size)
                raise ValueError(msg)

    def get_state_part_view(self, state, var_name, stripped=False):
        self.check_vector_size(state, stripped=stripped)
        return state[self.get_indx(var_name, stripped)]

    def set_state_part(self, state, part, var_name, stripped=False):
        self.check_vector_size(state, stripped=stripped)
        state[self.get_indx(var_name, stripped)] = part

    def get_state_parts(self, vec=None):
        """
        Return parts of a state vector corresponding to individual state
        variables.

        Parameters
        ----------
        vec : array, optional
            The state vector. If not given, then the data stored in the
            variables are returned instead.

        Returns
        -------
        out : dict
            The dictionary of the state parts.
        """
        if vec is not None:
            self.check_vector_size(vec)

        out = {}
        for var in self.iter_state():
            if vec is None:
                out[var.name] = var()

            else:
                out[var.name] = vec[self.di.indx[var.name]]

        return out

    def set_data(self, data, step=0, ignore_unknown=False,
                 preserve_caches=False):
        """
        Set data (vectors of DOF values) of variables.

        Parameters
        ----------
        data : array
            The state vector or dictionary of {variable_name : data vector}.
        step : int, optional
            The time history step, 0 (default) = current.
        ignore_unknown : bool, optional
            Ignore unknown variable names if `data` is a dict.
        preserve_caches : bool
            If True, do not invalidate evaluate caches of variables.
        """
        if data is None: return

        if isinstance(data, dict):

            for key, val in six.iteritems(data):
                try:
                    var = self[key]

                except (ValueError, IndexError):
                    if ignore_unknown:
                        pass

                    else:
                        raise KeyError('unknown variable! (%s)' % key)

                else:
                    var.set_data(val, step=step,
                                 preserve_caches=preserve_caches)

        elif isinstance(data, nm.ndarray):
            self.check_vector_size(data)

            for ii in self.state:
                var = self[ii]
                var.set_data(data, self.di.indx[var.name], step=step,
                             preserve_caches=preserve_caches)

        else:
            raise ValueError('unknown data class! (%s)' % data.__class__)

    def set_from_state(self, var_names, state, var_names_state):
        """
        Set variables with names in `var_names` from state variables with names
        in `var_names_state` using DOF values in the state vector `state`.
        """
        self.check_vector_size(state)

        if isinstance(var_names, basestr):
            var_names = [var_names]
            var_names_state = [var_names_state]

        for ii, var_name in enumerate(var_names):
            var_name_state = var_names_state[ii]

            if self[var_name_state].is_state():
                self[var_name].set_data(state, self.di.indx[var_name_state])

            else:
                msg = '%s is not a state part' % var_name_state
                raise IndexError(msg)

    def state_to_output(self, vec, fill_value=None, var_info=None,
                        extend=True, linearization=None):
        """
        Convert a state vector to a dictionary of output data usable by
        Mesh.write().
        """
        di = self.di

        if var_info is None:
            self.check_vector_size(vec)

            var_info = {}
            for name in di.var_names:
                var_info[name] = (False, name)

        out = {}
        for key, indx in six.iteritems(di.indx):
            var = self[key]

            if key not in list(var_info.keys()): continue
            is_part, name = var_info[key]

            if is_part:
                aux = vec

            else:
                aux = vec[indx]

            out.update(var.create_output(aux, key=name, extend=extend,
                                         fill_value=fill_value,
                                         linearization=linearization))

        return out

    def iter_state(self, ordered=True):

        if ordered:
            for ii in self.ordered_state:
                yield self[ii]

        else:
            for ii in self.state:
                yield self[ii]

    def init_history(self):
        for var in self.iter_state():
            var.init_history()

    def time_update(self, ts, functions, verbose=True):
        if verbose:
            output('updating variables...')

        for var in self:
            var.time_update(ts, functions)

        if verbose:
            output('...done')

    def advance(self, ts):
        for var in self.iter_state():
            var.advance(ts)

class Variable(Struct):
    _count = 0
    _orders = []
    _all_var_names = set()

    @staticmethod
    def reset():
        Variable._count = 0
        Variable._orders = []
        Variable._all_var_names = set()

    @staticmethod
    def from_conf(key, conf, fields):
        aux = conf.kind.split()
        if len(aux) == 2:
            kind, family = aux

        elif len(aux) == 3:
            kind, family = aux[0], '_'.join(aux[1:])

        else:
            raise ValueError('variable kind is 2 or 3 words! (%s)' % conf.kind)

        history = conf.get('history', None)
        if history is not None:
            try:
                history = int(history)
                
assert_(history >= 0)