import numpy as np

from sfepy.base.base import Struct
from sfepy.base.conf import ProblemConf
from sfepy.applications import PDESolverApp
from sfepy.fem import H1NodalVolumeField, FieldVariable

def getProblem():    
    def get_ebc(ts, coor, bc, problem):
        x = coor[:, 0]
        
        return x * x * x
        
    def get_f(ts, coors, mode = None, **kwargs):
        if mode == 'qp':
            x = coors[:, 0]
            y = coors[:, 1]
            z = coors[:, 2]
            
            val = 6 * x
            val.shape = (coors.shape[0], 1, 1)
            return {'val' : val}
            
    def post_process(out, pb, state, extend = False):       
        # evaluate gradient in nodes
        grad_data = pb.evaluate('ev_grad.i1.Omega(u)', mode = 'qp')
                
        grad_field = H1NodalVolumeField('grad', np.float64, (3, ),
                                        pb.domain.regions['Omega'])
        grad_var = FieldVariable('grad', 'parameter', grad_field, 3,
                                 primary_var_name='(set-to-None)')
        grad_var.data_from_qp(grad_data, pb.integrals['i1'])
        #out.update(grad_var.create_output())
                        
        return out
        
    problem = {    
        # OK with this mesh
        #'filename_mesh': 'box.mesh',
        
        # Error with this mesh
        'filename_mesh': 'sphere.mesh',
    
        # regions
        'regions' : {
            'Omega' : ('all', {}),
            'Gamma' : ('nodes of surface', {})
        },
    
        # materials    
        'materials' : {
            'm' : ({'val' : 1.0},),
            'f' : 'get_f'
        },

        # field and variables
        'fields' : {
            'potential' : ('real', (1,), 'Omega', 1)
        },

        'variables' : {
            'u' : ('unknown field', 'potential', 0),
            'v' : ('test field', 'potential', 'u'),
        },

        # boundary conditions
        'ebcs' : {
            'dirihlet' : ('Gamma', {'u.0' : 'get_ebc'})
        },
    
        # integrals    
        'integrals' : {
            'i1' : ('v', 2)
        },
    
        # equation
        'equations' : {
            'Potential' : """dw_laplace.i1.Omega(m.val, v, u) =
                             dw_volume_integrate.i1.Omega(f.val, v)"""
        },

        # solvers
        'solvers' : {
            'ls' : ('ls.scipy_direct', {}),
            'newton' : ('nls.newton', {
                'i_max'      : 10,
                'eps_a'      : 1e-10}),
        },
        
        'functions' : {
            'get_ebc': (get_ebc, ),
            'get_f' : (get_f, )
        },        
        
        # options    
        'options' : {
            'absolute_mesh_path' : True,
            'ls' : 'ls',
            'nls' : 'newton',
            'post_process_hook' : post_process,
            'save_results' : True
        }
    }
    
    return problem

options = Struct(
    output_filename_trunk = None,
    save_ebc = False,
    save_ebc_nodes = False,
    save_regions = False,
    save_field_meshes = False,
    save_regions_as_groups = False,
    solve_not = False,
)

conf = ProblemConf.from_dict(getProblem(), None)
app = PDESolverApp(conf, options, "FE:")
app()