#!/usr/bin/env python3

"""
Demonstration and testing of checkpointing and restarting CARP.
"""
import os

EXAMPLE_DESCRIPTIVE_NAME = 'Checkpoint and Restart'
EXAMPLE_AUTHOR = 'Gernot Plank <gernot.plank@medunigraz.at>'
EXAMPLE_DIR = os.path.dirname(__file__)

from datetime import date

from carputils import settings
from carputils import tools
from carputils import mesh
from carputils import testing

def parser():
    parser = tools.standard_parser()

    parser.add_argument('--experiment',
                        default='checkpoint',
                        choices=['checkpoint',
                                 'restart' ],
                        help='define experiment type')
    parser.add_argument('--np_restart',
                        type=int,
                        help='number of processors when restarting (default: '
                            +'same as checkpoint simulation)')
    return parser

def jobID(args):
    """
    Generate name of top level output directory.
    """
    today = date.today()
    jobID = '{}_{}_{}_np{}'.format(today.isoformat(), args.experiment,
                                   args.flv, args.np)
    if args.experiment == 'restart' and args.np_restart is not None:
        jobID += '_nprst{}'.format(args.np_restart)
    return jobID

@tools.carpexample(parser, jobID)
def run(args, job):

    # Restart np defaults to checkpoint np
    if args.np_restart is None:
        args.np_restart = args.np

    # Generate mesh
    # Block which is thin in z direction
    geom = mesh.Block(size=(5, 5, 0.1))
    # Set fibre angle to 0, sheet angle to 0
    geom.set_fibres(0, 0, 90, 90)

    # Define cylinder region
    reg = mesh.CylinderRegion(start=(0,0,-0.05), axis=(0,0,0.1), 
                              radius=2.5, tag=2)
    geom.add_region(reg)

    # Generate and return base name
    meshname = mesh.generate(geom)

    # define default view settings
    view = os.path.join(EXAMPLE_DIR, 'checkpoint.mshz')
 
    # Simulation output directories
    chkID = job.ID
    # Restarting? Put in subdirectories
    if args.experiment == 'restart':
        chkID = os.path.join(job.ID, 'checkpoint')
        rstID = os.path.join(job.ID, 'restart')

    # Add all the non-general arguments
    cmd  = tools.carp_cmd(os.path.join(EXAMPLE_DIR, 'checkpoint.par'))
    cmd += ['-meshname', meshname]

    if args.visualize:
        cmd += ['-gridout_i', 3,
                '-gridout_e', 3]
   
    # Run checkpoint simulation
    chk_cmd = cmd + ['-simID', chkID]
    job.carp(chk_cmd)

    start_statef = []
    if args.experiment == 'restart':
        start_statef = ['-start_statef',
                        os.path.join(chkID, 'state.chkpt.5.00')]

    if args.experiment == 'restart':
        # Overwrite nprocs
        args.np = args.np_restart
        # Run same command, but with restarting from checkpoint
        rst_cmd = cmd + ['-simID',        rstID,
                         '-start_statef', os.path.join(chkID, 'state.chkpt.5.00'),
                         '-num_tsav',     0]
        job.carp(rst_cmd)

    # Do visualization
    if args.visualize and not settings.platform.BATCH:

        geom = os.path.join(chkID, os.path.basename(meshname)+'_i')
        data = os.path.join(chkID, 'vm.igb')
        view = os.path.join(EXAMPLE_DIR, 'checkpoint.mshz')
        
        job.gunzip(data)
        job.meshalyzer(geom, data, view)

        if args.experiment == 'restart':
            geom = os.path.join(rstID, os.path.basename(meshname)+'_i')
            data = os.path.join(rstID, 'vm.igb')
            view = os.path.join(EXAMPLE_DIR, 'restart.mshz')
        
            job.gunzip(data)
            job.meshalyzer(geom, data, view)


# ---SERIAL---------------------------------------------------------------------
test_restart = testing.Test('restart', 
                            run, ['--experiment', 'restart' ],
                            tags=[testing.tag.FAST, testing.tag.SERIAL])
test_restart.add_filecmp_check('restart/vm.igb.gz', testing.max_error, 0.001)
# ---
test_checkpoint = testing.Test('checkpoint', run, 
                              ['--experiment', 'checkpoint'],
                              tags=[testing.tag.FAST, testing.tag.SERIAL])
test_checkpoint.add_filecmp_check('vm.igb.gz', testing.max_error, 0.001)
# ---PARALLEL-------------------------------------------------------------------
test_restart_parallel = testing.Test('restart_parallel', run, 
                                    ['--experiment', 'restart', '--np', '2'],
                                     refdir='restart', 
                                     tags=[testing.tag.SHORT, testing.tag.PARALLEL])
test_restart_parallel.add_filecmp_check('restart/vm.igb.gz', testing.max_error, 0.001)
test_restart_parallel.disable_reference_generation()
# ---
test_checkpoint_parallel = testing.Test('checkpoint_parallel', run, 
                                       ['--experiment', 'checkpoint', '--np', '2'],
                                        refdir='checkpoint', 
                                        tags=[testing.tag.SHORT, testing.tag.PARALLEL])
test_checkpoint_parallel.add_filecmp_check('vm.igb.gz', testing.max_error, 0.001)
test_checkpoint_parallel.disable_reference_generation()

__tests__ = [test_restart, test_restart_parallel, 
             test_checkpoint, test_checkpoint_parallel]

if __name__ == '__main__':
    run()