import dataclasses
import inspect
import pickle
import math
import random
from collections import OrderedDict
from dataclasses import Field, dataclass, fields
from functools import singledispatch, total_ordering
from pathlib import Path
from typing import (
    Any,
    ClassVar,
    Dict,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Type,
    TypeVar,
)
from simple_parsing import utils
from simple_parsing.helpers import Serializable
from logging import getLogger
from simple_parsing.utils import (
    compute_identity,
    dict_union,
    field_dict,
    get_type_arguments,
)

from .hparam import ValueOutsidePriorException, hparam, uniform, log_uniform
from .priors import Prior


logger = getLogger(__name__)
T = TypeVar("T")
HP = TypeVar("HP", bound="HyperParameters")

numpy_installed = False
try:
    import numpy as np

    numpy_installed = True
except ImportError:
    pass


@dataclass
class BoundInfo(Serializable):
    """Object used to provide the bounds as required by `GPyOpt`."""

    name: str
    # One of 'continuous', 'discrete' or 'bandit' (unsuported).
    type: str = "continuous"
    domain: Tuple[float, float] = (-math.inf, math.inf)


@dataclass
class HyperParameters(Serializable, decode_into_subclasses=True):  # type: ignore
    """Base class for dataclasses of HyperParameters."""

    # Class variable holding the random number generator used to create the
    # samples.
    if numpy_installed:
        np_rng: ClassVar[np.random.RandomState] = np.random
    else:
        rng: ClassVar[random.Random] = random.Random()

    def __post_init__(self):
        for name, f in field_dict(self).items():
            f: Field
            assert name == f.name
            value = getattr(self, name)
            # Apply any post-processing function, if applicable.
            if "postprocessing" in f.metadata:
                logger.debug(f"Post-processing of field {name}")
                try:
                    new_value = f.metadata["postprocessing"](value)
                except ValueOutsidePriorException as e:
                    raise ValueError(
                        f"Field '{name}' got value {repr(e.value)}, which is outside of the "
                        f"defined prior region: {e.prior}."
                    )
                setattr(self, name, new_value)

    @classmethod
    def field_names(cls) -> List[str]:
        return [f.name for f in fields(cls)]

    def id(self):
        return compute_identity(**self.to_dict())

    def seed(self, seed: Optional[int]) -> None:
        """TODO: Seed all priors with the given seed. (recursively if nested dataclasses
        are present.)
        """
        raise NotImplementedError("TODO")

    @classmethod
    def get_priors(cls) -> Dict[str, Prior]:
        """Returns a dictionary of the Priors for the hparam fields in this class."""
        priors_dict: Dict[str, Prior] = {}
        for field in fields(cls):
            # If a HyperParameters class contains another HyperParameters class as a field
            # we perform returned a flattened dict.
            if inspect.isclass(field.type) and issubclass(field.type, HyperParameters):
                priors_dict[field.name] = field.type.get_priors()
            else:
                prior: Optional[Prior] = field.metadata.get("prior")
                if prior:
                    priors_dict[field.name] = prior
        return priors_dict

    @classmethod
    def get_orion_space_dict(cls) -> Dict[str, str]:
        result: Dict[str, str] = {}
        for field in fields(cls):
            # If a HyperParameters class contains another HyperParameters class as a field
            # we perform returned a flattened dict.
            if inspect.isclass(field.type) and issubclass(field.type, HyperParameters):
                result[field.name] = field.type.get_orion_space_dict()
            else:
                prior: Optional[Prior] = field.metadata.get("prior")
                if prior:
                    result[field.name] = prior.get_orion_space_string()
        return result

    def get_orion_space(self) -> Dict[str, str]:
        """NOTE: This might be more useful in some cases than the above classmethod
        version, for example when a field is a different kind of dataclass than its
        annotation.
        """
        result: Dict[str, str] = {}
        for field in fields(self):
            value = getattr(self, field.name)
            # If a HyperParameters class contains another HyperParameters class as a field
            # we perform returned a flattened dict.
            if isinstance(value, HyperParameters):
                result[field.name] = value.get_orion_space()
            else:
                prior: Optional[Prior] = field.metadata.get("prior")
                if prior:
                    result[field.name] = prior.get_orion_space_string()