#!/usr/bin/env python
r"""
Laplace equation with a field-dependent material parameter.

Find :math:`T(t)` for :math:`t \in [0, t_{\rm final}]` such that:

.. math::
   \int_{\Omega} c(T) \nabla s \cdot \nabla T
    = f
    \;, \quad \forall s \;.

where :math:`c(T)` is the :math:`T` dependent diffusion coefficient.
Each iteration calculates :math:`T` and adjusts :math:`c(T)`.
"""
from argparse import ArgumentParser
from sfepy.base.base import output, IndexedStruct
from sfepy.discrete import (FieldVariable, Material, Integral, Function,
                            Equation, Equations, Problem)
from sfepy.discrete.conditions import Conditions, EssentialBC
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.mesh.mesh_generators import gen_block_mesh
from sfepy.postprocess.viewer import Viewer
from sfepy.solvers.ls import ScipyDirect
from sfepy.solvers.nls import Newton
from sfepy.terms import Term

import numpy as np
import sys
sys.path.append('.')


help = {
    'show' : 'show the results figure',
}

def main():
    parser = ArgumentParser()
    parser.add_argument('--version', action='version', version='%(prog)s')
    parser.add_argument('-s', '--show',
                        action="store_true", dest='show',
                        default=False, help=help['show'])
    options = parser.parse_args()
    options.output_dir = './solutions/'

    # Mesh dimensions
    dims = [1, 1]
    # Mesh resolution: increase to improve accuracy
    shape = [21, 21]
    # Center of mesh
    center = [0, 0.5]

    mesh = gen_block_mesh(dims, shape, center, name='myPoissonBlock',
                          verbose=False)
    domain = FEDomain('domain', mesh)

    min_x, max_x = domain.get_mesh_bounding_box()[:,0]
    min_y, max_y = domain.get_mesh_bounding_box()[:,1]
    eps_x = 1e-8 * (max_x - min_x)
    eps_y = 1e-8 * (max_y - min_y)

    Omega  = domain.create_region('Omega', 'all')
    Left   = domain.create_region('Left',
                                  'vertices in x < %.10f' % (min_x + eps_x),
                                  'facet')
    Right  = domain.create_region('Right',
                                  'vertices in x > %.10f' % (max_x - eps_x),
                                  'facet')
    Bottom = domain.create_region('Bottom',
                                  'vertices in y < %.10f' % (min_y + eps_y),
                                  'facet')
    Top    = domain.create_region('Top',
                                  'vertices in y > %.10f' % (max_y - eps_y),
                                  'facet')

    field = Field.from_args('Temperature', np.float64, 'scalar', Omega, 
                            space='H1', poly_space_base='lagrange', 
                            approx_order=2)

    T = FieldVariable('T', 'unknown', field)
    s = FieldVariable('s', 'test', field, primary_var_name='T')

#    conductivity_fun = Function('conductivity_fun', get_conductivity)

    cond = Material('cond', val=  2.0) 
    flux = Material('flux', val=-10.0)
    G    = Material('G',    val= 20.0)
    insulated = Material('insulated', val=0.0)

    t0     = 0.0
    t1     = 1.0
    n_step = 1

    integral   = Integral('i', order=1)
    # Terms applying to main region
    # Lapacian term  
    tLaplacian = Term.new('dw_laplace(cond.val, s, T)', 
                          integral, Omega, cond=cond, T=T, s=s)
    # Source term/unit vol
    tVolInt    = Term.new('dw_volume_lvf(G.val, s)',  
                          integral, Omega, G=G, s=s)

    # Neumann BCs
    # insulated top
    tSurfInsT  = Term.new('dw_surface_integrate(insulated.val, s)',
                          integral, Top, insulated=insulated, s=s)
    # constant flux on l/r edges
    tSurfFluxL = Term.new('dw_surface_integrate(flux.val, s)', 
                          integral, Left, flux=flux, s=s)
    tSurfFluxR = Term.new('dw_surface_integrate(flux.val, s)', 
                          integral, Right, flux=flux, s=s)

    # Dirichelet BCs
    bc_l = EssentialBC('T1', Left,   {'T.0' : 2.0})
    bc_r = EssentialBC('T2', Right,  {'T.0' : 2.0})
    bc_b = EssentialBC('T3', Bottom, {'T.0' : 0.0})

    # Build equation(s) from above terms
    eq  = Equation('fluxBalance', (tLaplacian - tVolInt 
                                   - tSurfInsT))
    eqs = Equations([eq])

    ls         = ScipyDirect({})
    nls_status = IndexedStruct()
    nls        = Newton({}, lin_solver=ls, status=nls_status)

    pb = Problem('Temperature', equations=eqs, nls=nls, ls=ls)
    pb.save_regions_as_groups('regions')

    pb.time_update(ebcs=Conditions([bc_l, bc_r])) #, bc_b]))

    vec = pb.solve()
    print nls_status

    pb.save_state('./solutions/poisson_interactive.vtk', vec)

    if options.show:
        view = Viewer('./solutions/poisson_interactive.vtk')
        view(vector_mode='warp_norm', rel_scaling=2,
             is_scalar_bar=True, is_wireframe=True)

if __name__ == '__main__':
    main()