"""
Laplace equation using the short syntax of keywords.

See :ref:`diffusion-poisson` for the long syntax version.

Find :math:`t` such that:

.. math::
    \int_{\Omega} c \nabla s \cdot \nabla t
    = 0
    \;, \quad \forall s \;.
"""
import numpy as np

filename_mesh = 'flnders_thermal_elm.mesh'

prms = {      'basin':{'rho':2300.0, 'cp':900.0,  'k':2.0},
        'upper_crust':{'rho':2800.0, 'cp':850.0,  'k':3.0},
        'lower_crust':{'rho':2900.0, 'cp':900.0,  'k':2.3},
             'mantle':{'rho':3400.0, 'cp':1000.0, 'k':4.0}}


materials = {
    'flux' : ({'val' : 0.03},),
    'coef_huh' : ({'val' : {'basin'  : 1.,
                            'upper_crust': 2.,
                            'lower_crust': 3.,
                            'mantle'     : 4.}},),
    'coef_old' : ({'val' : 1.0},),
    ## 'coef'     : ({'val' : {'basin'  : prms['basin']['k']/\
    ##                     (prms['basin']['rho']*prms['basin']['cp']),

    ##                     'upper_crust': prms['upper_crust']['k']/\
    ##                     (prms['upper_crust']['rho']*prms['upper_crust']['cp']),

    ##                         'lower_crust': prms['lower_crust']['k']/\
    ##                     (prms['lower_crust']['rho']*prms['lower_crust']['cp']),

    ##                         'mantle'     : prms['mantle']['k']/\
    ##                     (prms['mantle']['rho']*prms['mantle']['cp'])}},),
    'coef' : 'get_coef',
    'load' : 'get_pars',
}

def get_coef(ts, coors, mode=None, **kwargs):
    if mode != 'qp': return

    indx = kwargs['group_indx']
    term = kwargs['term']
    regions = kwargs['problem'].domain.regions

    val = np.empty((coors.shape[0], 1, 1), dtype=np.float64)

    print term.region.name
    for key, par in prms.iteritems():
        region = regions[key]
        for ig in term.region.igs:
            if ig not in region.igs: continue
            print region.name, indx[ig]
            val[indx[ig]] = par['k'] / (par['rho'] * par['cp'])

    return {'val' : val, 'valI' : val * np.eye(2)}

def get_pars(ts, coors, mode=None, **kwargs):
    """
    Evaluate the coefficient `load.f` in quadrature points `coors` using a
    function of space.

    For scalar parameters, the shape has to be set to `(coors.shape[0], 1, 1)`.
    """
    if mode == 'qp':
        q_anomaly  = 2.5e3
        q_background = 1.0
        q1 = q_anomaly/(prms['upper_crust']['rho']*prms['upper_crust']['rho'])
        q2 = q_background/(prms['upper_crust']['rho']*prms['upper_crust']['rho'])
        hx = 5000.
        hy = 5000.
        x = coors[:, 0]
        y = coors[:, 1]
        val = -q2-q1*np.exp(-((y+5.)/hy)**2-(x/hx)**2)
        val.shape = (coors.shape[0], 1, 1)
        return {'f' : val}


regions = {
    'Omega' : 'all', # or 'cells of group 6'
    'Gamma_Top' : ('vertices in (y > -0.1)', 'facet'),
    'Gamma_Bottom' : ('vertices of group 9', 'facet'),
    'basin'       : 'cells of group 4',
    'upper_crust' : 'cells of group 3',
    'lower_crust' : 'cells of group 2',
    'mantle'      : 'cells of group 1'
}

fields = {
    'temperature' : ('real', 1, 'Omega', 1),
}

variables = {
    't' : ('unknown field', 'temperature', 0),
    's' : ('test field',    'temperature', 't'),
}

ebcs = {
    't1' : ('Gamma_Top', {'t.0' : 0.0}),
}

equations = {
    'Temperature' : """dw_laplace.2.Omega( coef.val, s, t )
    = dw_surface_integrate.2.Gamma_Bottom( flux.val, s)
      - dw_volume_integrate.2.Omega( load.f, s )"""
#    = 0 """
}
#    't2' : ('Gamma_Bottom', {'t.0' : 1000.0}),


solvers = {
    'ls' : ('ls.scipy_direct', {}),
    'newton' : ('nls.newton',
                {'i_max'      : 1,
                 'eps_a'      : 1e-10,
    }),
}

functions = {
    'get_pars' : (get_pars,),
    'get_coef' : (get_coef,),
    }

options = {
    'nls' : 'newton',
    'ls' : 'ls',
    'post_process_hook' : 'post_process',
}

def post_process(out, problem, state, extend=False):
    """
    Calculate gradient of the solution.
    """
    flux = problem.evaluate('d_surface_flux.2.Gamma_Top(coef.valI, t)',
                         mode='eval', copy_materials=False, verbose=False)
    print flux
    flux_faces = problem.evaluate('d_surface_flux.2.Gamma_Top(coef.valI, t)',
                                  mode='el',
                                  copy_materials=False)
    print flux_faces.shape
    #from sfepy.discrete.fem.fields_base import create_expression_output

    #aux = create_expression_output('ev_grad.ie.Elements( t )',
    #                               'grad', 'temperature',
    #                               pb.fields, pb.get_materials(),
    #                               pb.get_variables(), functions=pb.functions,
    #                               mode='qp', verbose=False,
    #                               min_level=0, max_level=5, eps=1e-3)
    #out.update(aux)

    return out