#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Subpackage for RNGs.

This file contains:
    - A Python wrapper for
      `P. L'Ecuyers' and Richard Simard's RngStreams <http://statmath.wu.ac.at/software/RngStreams/>`_.
      This code was based on https://github.com/adamzjw/RngStream-in-Python/blob/master/RngStream.py .
    - Some helper classes for working with RNGs.

By default, the MRG32k3a-based RngStreams are used for generating random
numbers in the building blocks, but this can be changed to Python's builtin
:mod:`random` module, or the :mod:`numpy.random` module, as is required by
the user. Call :meth:`setGeneratorType` to do so.

Warning:
    In order to use RngStreams, please make sure you have the required
    library installed. This consists of the compiled C code for RngStreams
    and should be located in the :file:`src/dependencies/RngStreams/c` folder.

    Whenever this library cannot be found, a warning will be printed and the
    builtin :mod:`random` module will be used instead.
"""

import sys, random
from ctypes import *
from pypdevsbbl.extra.fileutils import *
import numpy as np
from enum import Enum

class Generators(Enum):
    """Enum helper class to determine the RNG to use."""

    RANDOM = 0
    """Use Python's builtin RNG. (i.e. the :mod:`random` module)"""

    NUMPY = 1
    """Use Numpy's RNG."""

    RNGSTREAMS = 2
    """Use Pierre L'Ecuyers' and Richard Simard's RngStreams.
    
    See Also:
        http://statmath.wu.ac.at/software/RngStreams/
    """

GENERATE_WITH = Generators.RNGSTREAMS
"""Identifies the stream to use.

Set this value to the desired RNG stream, before the model is constructed.
Must be a value from the :class:`Generators` enum.
"""

def setGeneratorType(gen_type):
    """Sets the :attr:`GENERATE_WITH` variable.

    Call this function to validly set the generator type.

    See Also:
        - :attr:`GENERATE_WITH`
        - :class:`Generators`

    Note:
        May be changed during construction of RNG-dependent building blocks.
        This allows the combination of multiple generator types.
    """
    global GENERATE_WITH
    GENERATE_WITH = gen_type


# LOAD RNGSTREAMS
try:
    if 'win' in sys.platform: # Windows
        RngStream = CDLL(dependencies("RngStreams", "c", "librngstream.dll"))
    else:                     # Other OS
        RngStream = CDLL(dependencies("RngStreams", "c", "librngstream.so"))


    # define struct
    class RngStream_InfoState(Structure):
        _fields_ = [
            ('Cg', c_double * 6),
            ('Bg', c_double * 6),
            ('Ig', c_double * 6),
            ('Anti', c_int),
            ('IncPrec', c_int),
            ('name', c_char_p)
        ]

    # Set argument and return values
    RngStream.RngStream_SetPackageSeed.argtypes = [c_ulong * 6]
    RngStream.RngStream_SetPackageSeed.restype = c_int
    RngStream.RngStream_CreateStream.restype = POINTER(RngStream_InfoState)
    RngStream.RngStream_SetSeed.argtypes = [POINTER(RngStream_InfoState), c_ulong * 6]
    RngStream.RngStream_SetSeed.restype = c_int
    RngStream.RngStream_GetState.argtypes = [POINTER(RngStream_InfoState), c_ulong * 6]
    RngStream.RngStream_RandU01.restype = c_double

    def setSeed(gen, seed):
        """Sets a new seed for the RngStreams generator object.

        RngStream requires a seed that is of size 6. This function makes it
        easier in a sense that when it does not find enough values for its seed,
        it will automatically fill the seed until it is large enough.

        Args:
            gen (object):   The generator object (RngStreams).
            seed:           Either a numeric, or an iterable.
                            The new seed to use.

        Returns:
            -1 on faillure, 0 on success.

        Raises:
            ValueError: When an invalid seed was given.
        """
        assert isinstance(gen, POINTER(RngStream_InfoState))
        n = 6
        if issubclass(type(seed), (list, tuple)):
            if len(seed) == 0:
                raise ValueError("Seed list size cannot be of length 0!")
            if len(seed) < n:
                s = (seed*np.ceil(n/len(seed)))[:n]
            else:
                s = seed[:6]
        elif issubclass(type(seed), (int, float, complex)):
            s = [seed] * n
        else:
            raise ValueError("Invalid seed aquired: %s" % str(seed))
        SD_TYPE = c_ulong * 6
        return RngStream.RngStream_SetSeed(gen, SD_TYPE(*s))

except OSError as e:
    import warnings
    warnings.warn("%s. Cannot load RngStreams. Using 'random' module instead." % str(e))
    setGeneratorType(Generators.RANDOM)



def createStream(name, seed=None):
    """Creates a custom RNG stream, w.r.t. the requested generator.

    Args:
        name (str): The unique name in case of an RngStreams generator.
        seed:       Either a numeric, or an iterable. The seed to use.
                    When ``None``, no seed will be set. Defaults to
                    ``None``.
    """
    if GENERATE_WITH == Generators.RANDOM:
        stream = random.Random(seed)
    elif GENERATE_WITH == Generators.NUMPY:
        stream = np.random.RandomState(seed)
    elif GENERATE_WITH == Generators.RNGSTREAMS:
        stream = RngStream.RngStream_CreateStream(name)
        if seed is not None:
            setSeed(stream, seed)
    else:
        raise ValueError("Unable to determine stream to generate.")
    return stream


def U01(gen):
    """Return a uniformly distributed random variable within (0, 1).

    This function works no matter the selected generator.

    Args:
        gen (object):   The generator to use.

    Returns:
        A numpy.float64 in the range (0, 1).
    """
    if isinstance(gen, random.Random):
        value = gen.random()
    elif isinstance(gen, np.random.RandomState):
        value = gen.uniform(0, 1)
    elif isinstance(gen, POINTER(RngStream_InfoState)):
        value = RngStream.RngStream_RandU01(gen)
    else:
        raise ValueError("Invalid generator instance.")
    return np.float64(value)
