#!/usr/bin/env python
r"""
Parallel assembling and solving of a Poisson's equation, using commands for
interactive use.

Find :math:`u` such that:

.. math::
    \int_{\Omega} \nabla v \cdot \nabla u
    = \int_{\Omega} v f
    \;, \quad \forall s \;.

Important Notes
---------------

- This example requires petsc4py, mpi4py and (optionally) pymetis with their
  dependencies installed!
- This example generates a number of files - do not use an existing non-empty
  directory for the ``output_dir`` argument.
- Use the ``--clear`` option with care!

Notes
-----

- Each task is responsible for a subdomain consisting of a set of cells (a cell
  region).
- Each subdomain owns PETSc DOFs within a consecutive range.
- When both global and task-local variables exist, the task-local
  variables have ``_i`` suffix.
- This example does not use a nonlinear solver.
- This example can serve as a template for solving a linear single-field scalar
  problem - just replace the equations in :func:`create_local_problem()`.
- The command line options are saved into <output_dir>/options.txt file.

Usage Examples
--------------

See all options::

  $ python examples/diffusion/poisson_parallel_interactive.py -h

See PETSc options::

  $ python examples/diffusion/poisson_parallel_interactive.py -help

Single process run useful for debugging with :func:`debug()
<sfepy.base.base.debug>`::

  $ python examples/diffusion/poisson_parallel_interactive.py output-parallel

Parallel runs::

  $ mpiexec -n 3 python examples/diffusion/poisson_parallel_interactive.py output-parallel -2 --shape=101,101

  $ mpiexec -n 3 python examples/diffusion/poisson_parallel_interactive.py output-parallel -2 --shape=101,101 --metis

  $ mpiexec -n 5 python examples/diffusion/poisson_parallel_interactive.py output-parallel -2 --shape=101,101 --verify --metis -ksp_monitor -ksp_converged_reason

View the results using::

  $ python postproc.py output-parallel/sol.h5 --wireframe -b -d'u,plot_warp_scalar'
"""
from __future__ import absolute_import
from argparse import RawDescriptionHelpFormatter, ArgumentParser
import os
import sys
sys.path.append('.')
import csv

import numpy as nm
import matplotlib.pyplot as plt

from sfepy.base.base import output, Struct
from sfepy.base.ioutils import ensure_path, remove_files_patterns, save_options
from sfepy.base.timing import Timer
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.discrete.common.region import Region
from sfepy.discrete import (FieldVariable, Material, Integral, Function,
                            Equation, Equations, Problem, State)
from sfepy.discrete.conditions import Conditions, EssentialBC
from sfepy.discrete.evaluate import apply_ebc_to_matrix
from sfepy.terms import Term
from sfepy.solvers.ls import PETScKrylovSolver

import sfepy.parallel.parallel as pl
import sfepy.parallel.plot_parallel_dofs as ppd

def create_local_problem(omega_gi, order):
    """
    Local problem definition using a domain corresponding to the global region
    `omega_gi`.
    """
    mesh = omega_gi.domain.mesh

    # All tasks have the whole mesh.
    bbox = mesh.get_bounding_box()
    min_x, max_x = bbox[:, 0]
    eps_x = 1e-8 * (max_x - min_x)

    mesh_i = Mesh.from_region(omega_gi, mesh, localize=True)
    domain_i = FEDomain('domain_i', mesh_i)
    omega_i = domain_i.create_region('Omega', 'all')

    gamma1_i = domain_i.create_region('Gamma1',
                                      'vertices in (x < %.10f)'
                                      % (min_x + eps_x),
                                      'facet', allow_empty=True)
    gamma2_i = domain_i.create_region('Gamma2',
                                      'vertices in (x > %.10f)'
                                      % (max_x - eps_x),
                                      'facet', allow_empty=True)

    field_i = Field.from_args('fu', nm.float64, 1, omega_i,
                              approx_order=order)

    output('number of local field DOFs:', field_i.n_nod)

    u_i = FieldVariable('u_i', 'unknown', field_i)
    v_i = FieldVariable('v_i', 'test', field_i, primary_var_name='u_i')

    integral = Integral('i', order=2*order)

    mat = Material('m', lam=10, mu=5)
    t1 = Term.new('dw_laplace(m.lam, v_i, u_i)',
                  integral, omega_i, m=mat, v_i=v_i, u_i=u_i)

    def _get_load(coors):
        val = nm.ones_like(coors[:, 0])
        for coor in coors.T:
            val *= nm.sin(4 * nm.pi * coor)
        return val

    def get_load(ts, coors, mode=None, **kwargs):
        if mode == 'qp':
            return {'val' : _get_load(coors).reshape(coors.shape[0], 1, 1)}

    load = Material('load', function=Function('get_load', get_load))

    t2 = Term.new('dw_volume_lvf(load.val, v_i)',
                  integral, omega_i, load=load, v_i=v_i)

    eq = Equation('balance', t1 - 100 * t2)
    eqs = Equations([eq])

    ebc1 = EssentialBC('ebc1', gamma1_i, {'u_i.all' : 0.0})
    ebc2 = EssentialBC('ebc2', gamma2_i, {'u_i.all' : 0.1})

    pb = Problem('problem_i', equations=eqs, active_only=False)
    pb.time_update(ebcs=Conditions([ebc1, ebc2]))
    pb.update_materials()

    return pb

def verify_save_dof_maps(field, cell_tasks, dof_maps, id_map, options,
                         verbose=False):
    vec = pl.verify_task_dof_maps(dof_maps, id_map, field, verbose=verbose)

    order = options.order
    mesh = field.domain.mesh

    sfield = Field.from_args('aux', nm.float64, 'scalar', field.region,
                             approx_order=order)
    aux = FieldVariable('aux', 'parameter', sfield,
                        primary_var_name='(set-to-None)')
    out = aux.create_output(vec,
                            linearization=Struct(kind='adaptive',
                                                 min_level=order-1,
                                                 max_level=order-1,
                                                 eps=1e-8))

    filename = os.path.join(options.output_dir,
                            'para-domains-dofs.h5')
    if field.is_higher_order():
        out['aux'].mesh.write(filename, out=out)

    else:
        mesh.write(filename, out=out)

    out = Struct(name='cells', mode='cell',
                 data=cell_tasks[:, None, None, None])
    filename = os.path.join(options.output_dir,
                            'para-domains-cells.h5')
    mesh.write(filename, out={'cells' : out})

def solve_problem(mesh_filename, options, comm):
    order = options.order

    rank, size = comm.Get_rank(), comm.Get_size()

    output('rank', rank, 'of', size)

    stats = Struct()
    timer = Timer('solve_timer')

    timer.start()
    mesh = Mesh.from_file(mesh_filename)
    stats.t_read_mesh = timer.stop()

    timer.start()
    if rank == 0:
        cell_tasks = pl.partition_mesh(mesh, size, use_metis=options.metis,
                                       verbose=True)

    else:
        cell_tasks = None

    stats.t_partition_mesh = timer.stop()

    output('creating global domain and field...')
    timer.start()

    domain = FEDomain('domain', mesh)
    omega = domain.create_region('Omega', 'all')
    field = Field.from_args('fu', nm.float64, 1, omega, approx_order=order)

    stats.t_create_global_fields = timer.stop()
    output('...done in', timer.dt)

    output('distributing field %s...' % field.name)
    timer.start()

    distribute = pl.distribute_fields_dofs
    lfds, gfds = distribute([field], cell_tasks,
                            is_overlap=True,
                            save_inter_regions=options.save_inter_regions,
                            output_dir=options.output_dir,
                            comm=comm, verbose=True)
    lfd = lfds[0]

    stats.t_distribute_fields_dofs = timer.stop()
    output('...done in', timer.dt)

    if rank == 0:
        dof_maps = gfds[0].dof_maps
        id_map = gfds[0].id_map

        if options.verify:
            verify_save_dof_maps(field, cell_tasks,
                                 dof_maps, id_map, options, verbose=True)

        if options.plot:
            ppd.plot_partitioning([None, None], field, cell_tasks, gfds[0],
                                  options.output_dir, size)

    output('creating local problem...')
    timer.start()

    omega_gi = Region.from_cells(lfd.cells, field.domain)
    omega_gi.finalize()
    omega_gi.update_shape()

    pb = create_local_problem(omega_gi, order)

    variables = pb.get_variables()
    eqs = pb.equations

    u_i = variables['u_i']
    field_i = u_i.field

    stats.t_create_local_problem = timer.stop()
    output('...done in', timer.dt)

    if options.plot:
        ppd.plot_local_dofs([None, None], field, field_i, omega_gi,
                            options.output_dir, rank)

    output('allocating global system...')
    timer.start()

    sizes, drange = 
pl.get_sizes(lfd.petsc_dofs_range, field.n_nod, 1)