# Source code for pypdevsbbl.extra.rngstreams

    #  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)