"""Module for distribution and variable specifications."""
from __future__ import annotations

# from metasyn.util import VarSpec
from dataclasses import dataclass, field
from typing import Any, Optional, Union

from metasyn.distribution.base import BaseDistribution, BaseFitter
from metasyn.privacy import BasePrivacy, get_privacy
from metasyn.util import ALL_VAR_TYPES


@dataclass
class DistributionSpec():
    """Specification that determines which distribution is selected.

    It has the following attributes:
    - name: Which distribution is chosen.
    - unique: Whether the distribution should be unique.
    - parameters: The parameters of the distribution as defined by name.
    - fit_kwargs: Fitting keyword arguments to be used while fitting the distribution.
    - version: Version of the distribution to fit.
    """

    name: Optional[str] = None
    unique: Optional[bool] = None
    parameters: Optional[dict] = None
    fit_kwargs: dict = field(default_factory=dict)
    version: Optional[str] = None
    distribution: Optional[BaseDistribution] = None

    def __post_init__(self):
        if self.name is None:
            if self.version is not None:
                raise ValueError("Cannot create DistributionSpec with attribute 'version' but "
                                 "without attribute 'name'.")
            if self.parameters is not None:
                raise ValueError("Cannot create DistributionSpec with attribute 'parameters' but "
                                 "without attribute 'name'.")
            if len(self.fit_kwargs) > 0:
                raise ValueError("Cannot create DistributionSpec with attribute 'fit_kwargs' that"
                                 " is not empty but without attribute 'name'.")


    @classmethod
    def parse(cls, dist_spec: Optional[Union[dict, type[BaseDistribution], BaseDistribution,
                                             DistributionSpec, str]],
              unique: Optional[bool] = None,
              ) -> DistributionSpec:
        """Create a DistributionSpec instance from a variety of inputs.

        Parameters
        ----------
        dist_spec:
            Specification for the distribution in several types.
        unique:
            Whether the distribution is unique. This is only taken into account
            if dist_spec is None or a string.

        Returns
        -------
            A instantiated version of the dist_spec that has the DistributionSpec type.

        Raises
        ------
        TypeError
            If the input has the wrong type and cannot be parsed.
        """
        if isinstance(dist_spec, BaseDistribution):
            dist_dict = {key: value for key, value in dist_spec.to_dict().items()
                         if key in ["name", "version", "unique", "parameters"]}
            return cls(**dist_dict, distribution=dist_spec)
        if isinstance(dist_spec, str):
            return cls(name=dist_spec, unique=unique)
        if dist_spec is None:
            return cls(unique=unique)
        if isinstance(dist_spec, dict):
            return cls(**dist_spec)
        if isinstance(dist_spec, DistributionSpec):
            return dist_spec
        if issubclass(dist_spec, BaseDistribution):
            return cls(name=dist_spec.name, unique=dist_spec.unique)
        raise TypeError("Error parsing distribution specification of unknown type "
                        f"'{type(dist_spec)}' with value '{dist_spec}'")

    @property
    def fully_specified(self) -> bool:
        """Indicate whether the distribution is suitable for datafree creation.

        Returns
        -------
            A flag that indicates whether a distribution can be generated from the values
            that are specified (not None).
        """
        return self.name is not None and self.parameters is not None

    def get_creation_method(self, fitter: Optional[BaseFitter]) -> dict:
        """Create a dictionary on how the distribution was created.

        Parameters
        ----------
        privacy
            Privacy object with which the dictionary is being created.

        Returns
        -------
            Dictionary containing all the non-default settings for the creation method.
        """
        ret_dict: dict[str, Any] = {"created_by": "metasyn"}
        dist_dict = {var: getattr(self, var) for var in ["name", "unique", "parameters", "version"]
                     if getattr(self, var) is not None}
        if len(dist_dict) != 0:
            ret_dict["distribution"] = dist_dict
        fit_dict = {}
        if fitter is not None:
            fit_dict = fitter.to_dict()
        if len(self.fit_kwargs) > 0:
            fit_dict["fit_kwargs"] = self.fit_kwargs

        if len(fit_dict) != 0:
            ret_dict["fitter"] = fit_dict

        return ret_dict

    def __str__(self):
        return f"DistSpec<{self.name}, {self.unique}, {self.parameters}, {self.distribution}>"


class VarSpec():  # pylint: disable=too-few-public-methods
    """Data class for storing the specifications for variables.

    Parameters
    ----------
    name:
        Name of the variable/column.
    distribution:
        Distribution to use for fitting/finding the distribution.
        Leave at None to allow metasyn to find the most suitable distribution
        automatically.

        >>> # Use normal distribution
        >>> distribution="normal"
        >>> # Use normal distribution with mean 0, standard deviation 1
        >>> distribution=NormalDistribution(0, 1)

    unique:
        To set a column to be unique/key.
        This is only available for the integer and string datatypes. Setting a variable
        to unique ensures that the synthetic values generated for this variable are unique.
        This is useful for ID or primary key variables, for example. The parameter...
        is ignored when the distribution is set manually. For example:
        {"unique": True}, which sets the variable to be unique or {"unique": False} which
        forces the variable to be not unique. If the uniqueness is not specified, it is
        assumed to be not unique, but gives a warning if metasyn thinks it should be.
    privacy:
        Set the privacy level for a variable, e.g.: DifferentialPrivacy(epsilon=10).
    prop_missing:
        Proportion of missing values for a variable.
    description:
        Set the description of a variable.
    data_free:
        Whether this variable/column is to be generated from scratch or from an existing column
        in the dataframe.
    var_type:
        Manually set the variable type of the columns (used mainly for data_free columns).
    """

    def __init__(
            self,
            name: str,
            distribution: Optional[Union[dict, type[BaseDistribution], BaseDistribution,
                                                DistributionSpec, str]] = None,
            unique=None,
            privacy: Optional[BasePrivacy] = None,
            prop_missing: Optional[float] = None,
            description: Optional[str] = None,
            data_free: Optional[bool] = None,
            var_type: Optional[str] = None):

        self.name = name
        self.dist_spec = DistributionSpec.parse(distribution, unique)
        self.privacy = privacy
        self.prop_missing = prop_missing
        self.description = description
        self.data_free = data_free
        self.var_type = var_type
        self.__post_init__()

    def __post_init__(self):
        # Convert the the privacy attribute if it is a dictionary.
        if isinstance(self.privacy, dict):
            self.privacy = get_privacy(**self.privacy)
        if self.data_free and not self.dist_spec.fully_specified:
            raise ValueError("Error creating variable specification: data free variable should have"
                            f" 'name' and 'parameters'. {self}")
        if self.var_type is not None and self.var_type not in ALL_VAR_TYPES:
            raise ValueError(f"Cannot create variable '{self.name}': unknown variable type "
                             f"'{self.var_type}'. Choose from {ALL_VAR_TYPES}.")

    @classmethod
    def from_dict(cls, var_dict: dict) -> VarSpec:
        """Create a variable specification from a dictionary.

        Parameters
        ----------
        var_dict
            Dictionary to parse the specification from.

        Returns
        -------
            A VarSpec instance.
        """
        return cls(**var_dict)


@dataclass
class VarDefaults():
    """Dataclass for variable defaults.

    Parameters
    ----------
    data_free:
        Whether the variable is completely synthetic or is based on real data.
    prop_missing:
        Proportion of missing values.
    distribution:
        Dictionary containing default distributions for each variable type.
    privacy:
        Privacy to be used by default for estimating distributions.
    """

    data_free: bool = False
    prop_missing: Optional[float] = None
    distribution: Optional[dict] = None
    privacy: Optional[BasePrivacy] = None

    def __post_init__(self):
        if isinstance(self.privacy, dict):
            self.privacy = get_privacy(**self.privacy)
        elif not isinstance(self.privacy, BasePrivacy):
            self.privacy = get_privacy("none")

