from contextlib import contextmanager
from ctypes import (CDLL, c_bool, c_int, c_int32, c_int64, c_double, c_char_p,
                    c_char, POINTER, Structure, c_void_p, create_string_buffer)
from warnings import warn

import numpy as np
from numpy.ctypeslib import as_array

from openmc.exceptions import AllocationError
from . import _dll
from .error import _error_handler
import openmc.capi


class _Bank(Structure):
    _fields_ = [('r', c_double*3),
                ('u', c_double*3),
                ('E', c_double),
                ('wgt', c_double),
                ('delayed_group', c_int),
                ('particle', c_int)]


# Define input type for numpy arrays that will be passed into C++ functions
# Must be an int or double array, with single dimension that is contiguous
_array_1d_int = np.ctypeslib.ndpointer(dtype=np.int32, ndim=1,
                                       flags='CONTIGUOUS')
_array_1d_dble = np.ctypeslib.ndpointer(dtype=np.double, ndim=1,
                                        flags='CONTIGUOUS')

_dll.openmc_calculate_volumes.restype = c_int
_dll.openmc_calculate_volumes.errcheck = _error_handler
_dll.openmc_finalize.restype = c_int
_dll.openmc_finalize.errcheck = _error_handler
_dll.openmc_find_cell.argtypes = [POINTER(c_double*3), POINTER(c_int32),
                                  POINTER(c_int32)]
_dll.openmc_find_cell.restype = c_int
_dll.openmc_find_cell.errcheck = _error_handler
_dll.openmc_hard_reset.restype = c_int
_dll.openmc_hard_reset.errcheck = _error_handler
_dll.openmc_init.argtypes = [c_int, POINTER(POINTER(c_char)), c_void_p]
_dll.openmc_init.restype = c_int
_dll.openmc_init.errcheck = _error_handler
_dll.openmc_get_keff.argtypes = [POINTER(c_double*2)]
_dll.openmc_get_keff.restype = c_int
_dll.openmc_get_keff.errcheck = _error_handler
_init_linsolver_argtypes = [_array_1d_int, c_int, _array_1d_int, c_int, c_int,
                            c_double, _array_1d_int, _array_1d_int]
_dll.openmc_initialize_linsolver.argtypes = _init_linsolver_argtypes
_dll.openmc_initialize_linsolver.restype = None
_dll.openmc_master.restype = c_bool
_dll.openmc_next_batch.argtypes = [POINTER(c_int)]
_dll.openmc_next_batch.restype = c_int
_dll.openmc_next_batch.errcheck = _error_handler
_dll.openmc_plot_geometry.restype = c_int
_dll.openmc_plot_geometry.restype = _error_handler
_dll.openmc_run.restype = c_int
_dll.openmc_run.errcheck = _error_handler
_dll.openmc_reset.restype = c_int
_dll.openmc_reset.errcheck = _error_handler
_run_linsolver_argtypes = [_array_1d_dble, _array_1d_dble, _array_1d_dble,
                           c_double]
_dll.openmc_run_linsolver.argtypes = _run_linsolver_argtypes
_dll.openmc_run_linsolver.restype = c_int
_dll.openmc_source_bank.argtypes = [POINTER(POINTER(_Bank)), POINTER(c_int64)]
_dll.openmc_source_bank.restype = c_int
_dll.openmc_source_bank.errcheck = _error_handler
_dll.openmc_simulation_init.restype = c_int
_dll.openmc_simulation_init.errcheck = _error_handler
_dll.openmc_simulation_finalize.restype = c_int
_dll.openmc_simulation_finalize.errcheck = _error_handler
_dll.openmc_statepoint_write.argtypes = [c_char_p, POINTER(c_bool)]
_dll.openmc_statepoint_write.restype = c_int
_dll.openmc_statepoint_write.errcheck = _error_handler


def calculate_volumes():
    """Run stochastic volume calculation"""
    _dll.openmc_calculate_volumes()


def current_batch():
    """Return the current batch of the simulation.

    Returns
    -------
    int
        Current batch of the simulation

    """
    return c_int.in_dll(_dll, 'current_batch').value


def finalize():
    """Finalize simulation and free memory"""
    _dll.openmc_finalize()


def find_cell(xyz):
    """Find the cell at a given point

    Parameters
    ----------
    xyz : iterable of float
        Cartesian coordinates of position

    Returns
    -------
    openmc.capi.Cell
        Cell containing the point
    int
        If the cell at the given point is repeated in the geometry, this
        indicates which instance it is, i.e., 0 would be the first instance.

    """
    index = c_int32()
    instance = c_int32()
    _dll.openmc_find_cell((c_double*3)(*xyz), index, instance)
    return openmc.capi.Cell(index=index.value), instance.value


def find_material(xyz):
    """Find the material at a given point

    Parameters
    ----------
    xyz : iterable of float
        Cartesian coordinates of position

    Returns
    -------
    openmc.capi.Material or None
        Material containing the point, or None is no material is found

    """
    index = c_int32()
    instance = c_int32()
    _dll.openmc_find_cell((c_double*3)(*xyz), index, instance)

    mats = openmc.capi.Cell(index=index.value).fill
    if isinstance(mats, (openmc.capi.Material, type(None))):
        return mats
    else:
        return mats[instance.value]


def hard_reset():
    """Reset tallies, timers, and pseudo-random number generator state."""
    _dll.openmc_hard_reset()


def init(args=None, intracomm=None):
    """Initialize OpenMC

    Parameters
    ----------
    args : list of str
        Command-line arguments
    intracomm : mpi4py.MPI.Intracomm or None
        MPI intracommunicator

    """
    if args is not None:
        args = ['openmc'] + list(args)
        argc = len(args)

        # Create the argv array. Note that it is actually expected to be of
        # length argc + 1 with the final item being a null pointer.
        argv = (POINTER(c_char) * (argc + 1))()
        for i, arg in enumerate(args):
            argv[i] = create_string_buffer(arg.encode())
    else:
        argc = 0
        argv = None

    if intracomm is not None:
        # If an mpi4py communicator was passed, convert it to void* to be passed
        # to openmc_init
        try:
            from mpi4py import MPI
        except ImportError:
            intracomm = None
        else:
            address = MPI._addressof(intracomm)
            intracomm = c_void_p(address)

    _dll.openmc_init(argc, argv, intracomm)


def iter_batches():
    """Iterator over batches.

    This function returns a generator-iterator that allows Python code to be run
    between batches in an OpenMC simulation. It should be used in conjunction
    with :func:`openmc.capi.simulation_init` and
    :func:`openmc.capi.simulation_finalize`. For example:

    .. code-block:: Python

        with openmc.capi.run_in_memory():
            openmc.capi.simulation_init()
            for _ in openmc.capi.iter_batches():
                # Look at convergence of tallies, for example
                ...
            openmc.capi.simulation_finalize()

    See Also
    --------
    openmc.capi.next_batch

    """
    while True:
        # Run next batch
        status = next_batch()

        # Provide opportunity for user to perform action between batches
        yield

        # End the iteration
        if status != 0:
            break


def keff():
    """Return the calculated k-eigenvalue and its standard deviation.

    Returns
    -------
    tuple
        Mean k-eigenvalue and standard deviation of the mean

    """
    n = openmc.capi.num_realizations()
    if n > 3:
        # Use the combined estimator if there are enough realizations
        k = (c_double*2)()
        _dll.openmc_get_keff(k)
        return tuple(k)
    else:
        # Otherwise, return the tracklength estimator
        mean = c_double.in_dll(_dll, 'keff').value
        std_dev = c_double.in_dll(_dll, 'keff_std').value \
                  if n > 1 else np.inf