#!/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 __future__ import print_function, absolute_import
from argparse import ArgumentParser
from sfepy import data_dir
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.discrete.fem.meshio import UserMeshIO
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('.')


def get_conductivity(ts, coors, problem, equations=None, mode=None, **kwargs):
    """
    Calculates the conductivity as 2+10*T and returns it.
    This relation results in larger T gradients where T is small.
    """
    if mode == 'qp':
        # T-field values in quadrature points coordinates given by integral i
        # - they are the same as in `coors` argument.
        T_values = problem.evaluate('ev_volume_integrate.i.Omega(T)',
                                    mode='qp', verbose=False)
        eps = 1e-8
        val = 2.0 + eps * T_values #10.0 * (T_values + 2.0)

        output('conductivity: min:', val.min(), 'max:', val.max())

        val.shape = (val.shape[0] * val.shape[1], 1, 1)
        return {'val' : val}


def initial_guess(coors, flux):
    vec_x0 = [] #np.array()
    return vec_x0


def mesh_hook(mesh, mode):
    """
    Generate the block mesh.
    """
    if mode == 'read':
        mesh = gen_block_mesh(dims, shape, center, name='myPoissonBlock',
                              verbose=False)
        return mesh

    elif mode == 'write':
        pass


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

def main():
    from sfepy import data_dir

    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)

    # mesh = Mesh.from_file(data_dir + '/meshes/2d/rectangle_tri.mesh')
    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, 'vector', Omega, approx_order=2)

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

    coef = Material('coef', val=[2.0]) #'get_conductivity')
    flux = Material('flux', val=[-10.0])
    insulated = Material('insulated', val=[[-0.0]])

    t0 = 0.0
    t1 = 0.1
    n_step = 21

    integral   = Integral('i', order=1)
    tLaplacian = Term.new('dw_laplace(coef.val, s, T)', 
                          integral, Omega, coef=coef, T=T, s=s)
    tSurfIntT  = Term.new('dw_surface_integrate(insulated.val, s)',
                          integral, Top, insulated=insulated, T=T, s=s)
    tSurfIntB  = Term.new('dw_surface_integrate(insulated.val, s)',
                          integral, Bottom, insulated=insulated, T=T, s=s)
    # tSurfIntR  = Term.new('dw_surface_integrate(flux.val, s)', 
    #                       integral, Right, flux=flux, T=T, s=s)

    eq  = Equation('Temperature', tLaplacian - tSurfIntT - tSurfIntB)
    eqs = Equations([eq])

    T1 = EssentialBC('T1', Left,  {'T.0' : 0.0})
    T2 = EssentialBC('T2', Right, {'T.0' : 2.0})

#    conductivity_fun = Function('get_conductivity', get_conductivity, {'T.0'),

# ics = {
#     'ic' : ('Omega', {'T.0' : 0.0}),
# }

    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([T1, T2]))

    vec = pb.solve()
    print(nls_status)

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

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

if __name__ == '__main__':
    main()

#         'i_max' : 1,
#         'eps_a' : 1e-10,
#         'eps_r' : 1.0,
#         'vec_x0' : 'initial_guess',
#     }),
#     'ts' : ('ts.simple', {
#         't0' : t0,
#         't1' : t1,
#         'dt' : None,
#         'n_step' : n_step, # has precedence over dt!
#         'quasistatic' : True,
#     }),
# }


# options = {
#     'nls' : 'newton',
#     'ls' : 'ls',
#     'ts' : 'ts',
#     'save_steps' : -1,
#     'output_dir' : output_dir,
# }