#!/usr/bin/env python2.6
# -*- coding: UTF-8 -*-
#
# Copyright (C) 2010-2011 Yung-Yu Chen <yyc@solvcon.net>.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

"""
Convergence test for three-dimensional, anisotropic, linear elastic solver.
Tests are run against a two by two by two (meter) cube with 0.2, 0.15, 0.1 cell
edge sizes.  The meshes are generated by calling Cubit on the fly.  The wave
vector is set to (pi, pi, pi).  A command line tool "./go converge" is provided
to analyze the order of convergence of error after the arrangements are run.

Run the arrangements by ./go run elcv3d_200 elcv3d_150 elcv3d_100
"""

from solvcon.anchor import Anchor
from solvcon.hook import BlockHook
from solvcon.kerpak.cese import CesePeriodic
from solvcon.kerpak.lincese import PlaneWaveSolution
from solvcon.kerpak import elaslin
from solvcon.cmdutil import Command

###############################################################################
# Command line.
###############################################################################
class converge(Command):
    """
    Calculate and verify convergence.

    Must supply <delta> <M1>.
    """
    min_args = 0

    def __init__(self, env):
        from optparse import OptionGroup
        super(converge, self).__init__(env)
        op = self.op

        opg = OptionGroup(op, 'Convergence')
        opg.add_option("--wdir", action="store",
            dest="wdir", default="result", help="Working directory.")
        opg.add_option("--key", action="store",
            dest="key", default="L2", help="Linf or L2 norm.")
        opg.add_option("--idx", action="store", type=int,
            dest="idx", default=0, help="Index of variable: 0--8.")
        opg.add_option("--order", action="store", type=float,
            dest="order", default=None,
            help="Error-norm should converge at the rate, if given.")
        opg.add_option("--order-tolerance", action="store", type=float,
            dest="order_tolerance", default=0.4,
            help="The variation of converge order which can be tolerated.")
        opg.add_option("--stop-on-over", action="store_true",
            dest="stop_on_over", default=False,
            help="Raise ValueError if tolerance not met.")
        op.add_option_group(opg)
        self.opg_obshock = opg

    def __call__(self):
        import os, sys
        from pickle import load
        from math import log
        ops, args = self.opargs
        dat = [(mm,
            load(open(os.path.join(ops.wdir, 'elcv3d_%d_norm.pickle'%mm))))
            for mm in (200, 150, 100)]
        sys.stdout.write(
            'Convergence of %s error-norm at the %dth (0--8) variable:\n' % (
            ops.key, ops.idx))
        for ih in range(1, len(dat)):
            er = [dat[it][1][ops.key][ops.idx] for it in range(ih-1, ih+1)]
            hr = [float(dat[it][0])/1000 for it in range(ih-1, ih+1)]
            odr = log(er[1]/er[0])/log(hr[1]/hr[0])
            sys.stdout.write('  %6.4f -> %6.4f (m): %g' % (hr[0], hr[1], odr))
            if ops.order is not None:
                if abs(odr - ops.order) < ops.order_tolerance:
                    sys.stdout.write(' GOOD. Within')
                else:
                    if ops.stop_on_over:
                        raise ValueError('out of tolerance')
                    else:
                        sys.stdout.write(' BAD. Out of')
                sys.stdout.write(' %g +/- %g'%(ops.order, ops.order_tolerance))
            sys.stdout.write('\n')

################################################################################
# Mesh generation and boundary condition processor.
################################################################################
def mesher(cse):
    """
    Generate a cube according to journaling file cube.tmpl.
    """
    from solvcon.helper import Cubit
    try:
        itv = float(cse.io.basefn.split('_')[-1])/1000
    except ValueError:
        itv = 0.2
    cmds = open('cube.tmpl').read() % itv
    cmds = [cmd.strip() for cmd in cmds.strip().split('\n')]
    gn = Cubit(cmds, 3)()
    return gn.toblock(bcname_mapper=cse.condition.bcmap)

def match_periodic(blk):
    """
    Match periodic boundary condition.
    """
    from numpy import array
    from solvcon.boundcond import bctregy
    bct = bctregy.CesePeriodic
    bcmap = {
        'left': (
            bct, {
                'link': 'right',
                'ref': array([0,-2,-2], dtype='float64'),
            }
        ),
        'lower': (
            bct, {
                'link': 'upper',
                'ref': array([-2,0,-2], dtype='float64'),
            }
        ),
        'rear': (
            bct, {
                'link': 'front',
                'ref': array([-2,-2,0], dtype='float64'),
            }
        ),
    }
    bct.couple_all(blk, bcmap)

################################################################################
# Basic configuration.
################################################################################
def elcv3d_base(casename=None, mtrlname='GaAs',
    al=20.0, be=40.0, ga=50.0, wtests=None, psteps=None, ssteps=None, **kw):
    """
    Fundamental configuration of the simulation and return the case object.

    @return: the created Case object.
    @rtype: solvcon.case.BlockCase
    """
    import os
    from numpy import pi, array, sin, cos, sqrt
    from solvcon.conf import env
    from solvcon.boundcond import bctregy
    from solvcon import hook, anchor
    from solvcon.solver import ALMOST_ZERO
    from solvcon.kerpak import cese
    from solvcon.kerpak.lincese import PlaneWaveHook
    from solvcon.kerpak.elaslin import mltregy, ElaslinPWSolution
    # set up BCs.
    bct = bctregy.BC
    bcmap = {
        'front': (bct, {}),
        'rear': (bct, {}),
        'left': (bct, {}),
        'right': (bct, {}),
        'upper': (bct, {}),
        'lower': (bct, {}),
    }
    # set up case.
    mtrl = mltregy[mtrlname](al=al*pi/180.0, be=be*pi/180.0, ga=ga*pi/180.0)
    basedir = os.path.join(os.path.abspath(os.getcwd()), 'result')
    cse = elaslin.ElaslinCase(basedir=basedir, rootdir=env.projdir,
        basefn=casename, mesher=mesher, bcmap=bcmap, bcmod=match_periodic,
        mtrldict={None: mtrl}, taylor=0.0, **kw)
    # statistical anchors for solvers.
    for name in 'Runtime', 'March', 'Tpool':
        cse.runhooks.append(getattr(anchor, name+'StatAnchor'))
    # informative hooks.
    cse.runhooks.append(hook.BlockInfoHook)
    cse.runhooks.append(hook.ProgressHook, psteps=psteps,
        linewidth=ssteps/psteps)
    cse.runhooks.append(cese.CflHook, fullstop=False, psteps=ssteps,
        cflmax=10.0, linewidth=ssteps/psteps)
    # initializer anchors..
    cse.runhooks.append(anchor.FillAnchor, keys=('soln',), value=ALMOST_ZERO)
    cse.runhooks.append(anchor.FillAnchor, keys=('dsoln',), value=0)
    ## plane wave solution.
    pws = list()
    for wvec, idx in wtests:
        pws.append(ElaslinPWSolution(amp=1.0,
            ctr=array([0,0,0], dtype='float64'), wvec=wvec, mtrl=mtrl, idx=idx
            ))
    cse.runhooks.append(PlaneWaveHook, psteps=ssteps, planewaves=pws)
    # analyzing/output anchors and hooks.
    #cse.runhooks.append(euler.ElaslinOAnchor)
    cse.runhooks.append(hook.PMarchSave, anames=[
            ('soln', False, -9),
            ('analytical', True, -9),
            ('difference', True, -9),
        ], fpdtype='float64', psteps=ssteps, compressor='gz')
    return cse

def elcv3d_skel(casename, div, std, **kw):
    from numpy import array, pi
    period = 2.353142528777195e-04
    return elcv3d_base(casename=casename,
        time_increment=period/div, steps_run=2*div, ssteps=div, psteps=std,
        wtests=(
            (array([1,1,1], dtype='float64')*pi, 8),
        ), **kw)

################################################################################
# The arrangement for demonstration.
################################################################################
@elaslin.ElaslinCase.register_arrangement
def elcv3d(casename, **kw):
    from numpy import array, pi
    period = 2.353142528777195e-04
    div = 32
    std = 1
    return elcv3d_base(casename=casename,
        time_increment=period/div, steps_run=2*div, ssteps=std, psteps=std,
        wtests=(
            (array([1,1,1], dtype='float64')*pi, 8),
        ), **kw)

################################################################################
# The arrangement for convergence test.
################################################################################
@elaslin.ElaslinCase.register_arrangement
def elcv3d_200(casename, **kw):
    return elcv3d_skel(casename=casename, div=32, std=1)
@elaslin.ElaslinCase.register_arrangement
def elcv3d_150(casename, **kw):
    return elcv3d_skel(casename=casename, div=40, std=1)
@elaslin.ElaslinCase.register_arrangement
def elcv3d_100(casename, **kw):
    return elcv3d_skel(casename=casename, div=64, std=2)

################################################################################
# Invoke SOLVCON workflow.
################################################################################
if __name__ == '__main__':
    import solvcon
    solvcon.go()
