#!/usr/bin/env python

from types import MethodType
from subprocess import Popen
from subprocess import PIPE
import numpy as np
import shutil
import shlex
import os

top = '.'
out = '.bld'

# Set of compiler options.
DEBUG_OPTIONS = ['-O', '-Wall', '-Wline-truncation', '-Wcharacter-truncation']
DEBUG_OPTIONS += ['-Wsurprising', '-Waliasing', '-Wimplicit-interface']
DEBUG_OPTIONS += ['-Wunused-parameter', '-fwhole-file', '-fcheck=all']
DEBUG_OPTIONS += ['-fbacktrace', '-g', '-fmax-errors=1', '-ffpe-trap=invalid']
DEBUG_OPTIONS += ['-ffree-line-length-0']

PRODUCTION_OPTIONS = ['-O3', '-ffree-line-length-0']


def options(ctx):

    ctx.load('compiler_c')

    ctx.load('compiler_fc')

    ctx.add_option('--debug', action='store_true', dest='is_debug',
                   default=False, help='use debug compiler options')

    ctx.add_option('--without_mpi', action='store_true', dest='without_mpi',
        default=False, help='without parallelism')

    ctx.add_option('--without_fortran', action='store_true',
        dest='without_fortran', default=False, help='without FORTRAN')


def support_fortran(ctx):
    """ This function determines whether FORTRAN is supported or not.
    Some tests are duplicated.
    """

    ctx.start_msg('Support FORTRAN ')

    try:

        # Load compiler
        ctx.load('compiler_fc')

        try:
            assert ctx.env.FC_NAME == 'GFORTRAN'
        except AssertionError:
            raise ctx.errors.ConfigurationError

        # Compile a small FORTRAN program
        ctx.check_fortran()

        # Check for LAPACK library
        ctx.check_fc(lib='lapack')

    except ctx.errors.ConfigurationError:
        idx = 1
    else:
        idx = 0

    ctx.end_msg(['yes', 'no'][idx])

    return idx


def support_parallelism(ctx):
    """ This function determines whether parallelism is supported or not.
    Some tests are duplicated.
    """

    ctx.start_msg('Support PARALLELISM ')

    try:
        # Check for the existence of the compiler
        ctx.find_program('mpif90')

        # Check for MPICH implementation and version
        implementation, version = _get_mpif90_details(ctx)
        if implementation not in ['MPICH']:
            raise ctx.errors.ConfigurationError

        if int(version.split('.')[0]) < 3:
            raise ctx.errors.ConfigurationError

    except ctx.errors.ConfigurationError:
        idx = 1
    else:
        idx = 0

    ctx.end_msg(['yes', 'no'][idx])

    return idx


def configure(ctx):

    # Check for FORTRAN support
    ctx.support_fortran = MethodType(support_fortran, ctx)
    if ctx.options.without_fortran:
        ctx.env['FORTRAN'] = False
    else:
        ctx.env['FORTRAN'] = (ctx.support_fortran() == 0)

    # Checking for parallelism
    ctx.support_parallelism = MethodType(support_parallelism, ctx)
    if ctx.options.without_mpi or (not ctx.env['FORTRAN']):
        ctx.env['PARALLELISM'] = False
    else:
        ctx.env['PARALLELISM'] = (ctx.support_parallelism() == 0)
        if ctx.env['PARALLELISM']:
            ctx.env['FC'] = 'mpif90'
            ctx.env['COMPILER_FORTRAN'] = 'mpif90'

    if not ctx.env['FORTRAN']:
        return

    # Set up compilation (if required)
    is_debug = ctx.options.is_debug

    # The build is currently only tested for GFORTRAN.
    if is_debug:
        ctx.env.append_unique('FCFLAGS', DEBUG_OPTIONS)
    else:
        ctx.env.append_unique('FCFLAGS', PRODUCTION_OPTIONS)

    # Enable static libraries
    ctx.env.append_unique('STLIBPATH', ['../.bld/fortran'])
    ctx.env.append_unique('STLIB', ['resfort'])

    # Enable shared libraries
    ctx.env.append_unique('LIB', ['lapack'])

def build(ctx):

    ctx.add_pre_fun(pre)

    if ctx.env['FORTRAN']:

        ctx.recurse('fortran')

        ctx.add_group()

        ctx.recurse('tests/resources')

        ctx.add_post_fun(post)


def post(ctx):
    """ This function performs some housekeeping tasks. We keep all the
    executables in a single directory.
    """
    subdir = 'fortran/bin'

    # Copy the RESFORT executable back to package directory
    os.mkdir(subdir)
    shutil.copy('.bld/fortran/resfort_scalar', subdir)

    if ctx.env['PARALLELISM']:
        shutil.copy('.bld/fortran/resfort_parallel_master', subdir)
        shutil.copy('.bld/fortran/resfort_parallel_slave', subdir)

    # Copy the KW executable.
    src = '.bld/tests/resources/kw_dp3asim'
    shutil.copy(src, 'tests/resources')


def pre(ctx):
    """ Clean up some directories.
    """

    if os.path.exists('fortran/bin'):
        shutil.rmtree('fortran/bin')

    if os.path.exists('tests/resources/kw_dp3asim'):
        os.unlink('tests/resources/kw_dp3asim')


def _get_mpif90_details(ctx):

    output, _ = Popen(['mpif90', '-v'], stdout=PIPE, stderr=PIPE).communicate()

    # In PYTHON 2, output is already a string, but in PYTHON 3 we need to
    # decode it first.
    try:
        output = output.decode()
    except AttributeError:
        pass

    # Now we can process the information independent of PYTHON 2 or PYTHON 3.
    try:
        implementation, version = np.array(shlex.split(output))[[2, 4]]
    except:
        raise ctx.errors.ConfigurationError

    return implementation, version
