import time

import numpy as nm
import numpy.linalg as nla

from sfepy.base.base import output, get_default, Struct
from sfepy.base.log import Log, get_logging_conf
from sfepy.solvers.solvers import SolverMeta, NonlinearSolver
from nls import conv_test

class StabilizationFunction(Struct):
    """
    Definition of stabilization material function for the Oseen solver.

    Notes
    -----
    - tau_red <= 1.0; if tau is None: tau = tau_red * delta
    - diameter mode: 'edge': longest edge 'volume': volume-based, 'max': max. of
      previous
    """

    def __init__(self, name_map, gamma=None, delta=None, tau=None, tau_red=1.0,
                 tau_mul=1.0, delta_mul=1.0, gamma_mul=1.0,
                 diameter_mode='max'):
        Struct.__init__(self, name_map=name_map,
                        gamma=gamma, delta=delta, tau=tau,
                        tau_red=tau_red, tau_mul=tau_mul, delta_mul=delta_mul,
                        gamma_mul=gamma_mul, diameter_mode=diameter_mode)

    def setup(self, problem):
        """
        Setup common problem-dependent data.
        """
        variables = problem.get_variables()

        ns = self.name_map

        # Indices to the state vector.
        ii = {}
        ii['u'] = variables.get_indx(ns['u'])
        ii['us'] = variables.get_indx(ns['u'], stripped=True)
        ii['ps'] = variables.get_indx(ns['p'], stripped=True)
        self.indices = ii

        materials = problem.get_materials()

        # The viscosity.
        fluid_mat = materials[ns['fluid']]
        self.viscosity = fluid_mat.function()[ns['viscosity']]

        # The Friedrich's constant.
        self.c_friedrichs = problem.domain.get_diameter()
        self.sigma = 1e-12 # 1 / dt.

        self.b_norm = 1.0

    def get_maps(self):
        """
        Get the maps of names and indices of variables in state vector.
        """
        return self.name_map, self.indices

    def __call__(self, ts, coor, mode=None, term=None, problem=None,
                 b_norm=None, **kwargs):
        """
        The actual material function.
        """
        if mode != 'qp': return

        if not hasattr(self, 'viscosity'):
            self.setup(problem)

        ns = self.name_map

        # Update stored b_norm.
        self.b_norm = get_default(b_norm, self.b_norm)

        output('|b|_max (mat_fun):', self.b_norm)
        gamma = self.viscosity + self.b_norm * self.c_friedrichs

        data = {}
        if self.gamma is None:
            _gamma = self.gamma_mul * gamma

        else:
            _gamma = nm.asarray(self.gamma_mul * self.gamma, dtype=nm.float64)
        _gamma = nm.tile(_gamma, (coor.shape[0], 1, 1))