import numpy as np
import rlberry.check_packages as check_packages
from rlberry.seeding.seeder import Seeder

if check_packages.TORCH_INSTALLED:
    import torch


def set_external_seed(seeder):
    """
    Set seeds of external libraries.

    To do:
        Check (torch seeding):
        https://github.com/pytorch/pytorch/issues/7068#issuecomment-487907668

    Parameters
    ---------
    seeder: seeding.Seeder or int
        Integer or Seeder object from which to generate random seeds.

    Examples
    --------
    >>> from rlberry.seeding import set_external_seed
    >>> set_external_seed(seeder)
    """
    if np.issubdtype(type(seeder), np.integer):
        seeder = Seeder(seeder)

    # seed torch
    if check_packages.TORCH_INSTALLED:
        torch.manual_seed(seeder.seed_seq.generate_state(1, dtype=np.uint32)[0])


def safe_reseed(obj, seeder, reseed_spaces=True):
    """
    Calls obj.reseed(seed_seq) method if available;
    If a obj.seed() method is available, call obj.seed(seed_val),
    where seed_val is generated by the seeder.
    Otherwise, does nothing.

    Parameters
    ----------
    obj : object
        Object to be reseeded.
    seeder: :class:`~rlberry.seeding.seeder.Seeder`
        Seeder object from which to generate random seeds.
    reseed_spaces: bool, default = True.
        If False, do not try to reseed observation_space and action_space (if
        they exist as attributes of `obj`).

    Returns
    -------
    True if reseeding was done, False otherwise.

    """
    reseeded = False
    try:
        obj.reseed(seeder)
        reseeded = True
    except AttributeError:
        seed_val = seeder.rng.integers(2**32).item()
        try:
            obj.seed(seed_val)
            reseeded = True
        except AttributeError:
            try:
                obj.reset(seed=seed_val)
                reseeded = True
            except AttributeError:
                reseeded = False

    # check if the object has observation and action spaces to be reseeded.
    if reseed_spaces:
        try:
            safe_reseed(obj.observation_space, seeder)
            safe_reseed(obj.action_space, seeder)
        except AttributeError:
            pass

    return reseeded
