"""
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_thermal2km_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'])}},),
    'load' : 'get_pars',
}

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,),
    }
    
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.val, t)', 
                         mode='eval', copy_materials=False, verbose=False) 
    print flux 
    #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