import os
import typing
import numpy
from scipy import stats

from d3m import container, utils as d3m_utils
from d3m.metadata import base as metadata_base, hyperparams, params
from d3m.primitive_interfaces import base
from d3m.primitive_interfaces.unsupervised_learning import UnsupervisedLearnerPrimitiveBase

import common_primitives

Inputs = container.DataFrame
Outputs = container.List


class Params(params.Params):
    fitted: bool
    random_seed: int


class Hyperparams(hyperparams.Hyperparams):
    distribution = hyperparams.Enumeration(
        values=['uniform', 'binomial', 'geometric', 'poisson', 'normal'],
        default='uniform',
        semantic_types=['https://metadata.datadrivendiscovery.org/types/TuningParameter'],
        description="Distribution to be use on sampling"
    )
    sample_size = hyperparams.Bounded[float](
        lower=0.0,
        upper=1.0,
        default=0.8,
        description='Percentage of samples to use per output dataframe',
        semantic_types=['https://metadata.datadrivendiscovery.org/types/ControlParameter'],
    )

    n_output_dataframes = hyperparams.Bounded[int](
        lower=1,
        upper=None,
        default=1,
        semantic_types=['https://metadata.datadrivendiscovery.org/types/ControlParameter'],
        description='Number of dataframes generated by sampling.'
    )


class DataFrameSamplingPrimitive(UnsupervisedLearnerPrimitiveBase[Inputs, Outputs, Params, Hyperparams]):
    """
    A primitive which samples from a DataFrame and returns a List.
    The number of samples (for each sampled dataset) or the fraction and
    the length of list (number of sampled dataset) are specified.
    An extra hyperparameter indicating sampling distribution is allowed.
    If True targets are present on the input dataframe, the primitive will do sampling,
    otherwise it will return a copy of the input dataframe to all outptus.
    """

    metadata = metadata_base.PrimitiveMetadata(
        {
            'id': 'adc6e2e6-a477-41cf-8ce2-5dc1b475246e',
            'version': '0.1.0',
            'name': "DataFrame Sampling",
            # TODO add sampling as primitive name in D3M Core
            'python_path': 'd3m.primitives.data_transformation.dataframe_sampling.Common',
            'source': {
                'name': common_primitives.__author__,
                'contact': 'mailto:dmartinez05@tamu.edu',
                'uris': [
                    'https://gitlab.com/datadrivendiscovery/common-primitives/blob/master/common_primitives/dataframe_sampling.py',
                    'https://gitlab.com/datadrivendiscovery/common-primitives.git',
                ],
            },
            'installation': [{
                'type': metadata_base.PrimitiveInstallationType.PIP,
                'package': 'd3m-common-primitives',
                'version': '2022.2.17',
            }],
            'algorithm_types': [
                metadata_base.PrimitiveAlgorithmType.DATA_CONVERSION,
            ],
            'primitive_family': metadata_base.PrimitiveFamily.DATA_TRANSFORMATION,
        },
    )

    def __init__(self, *, hyperparams: Hyperparams, random_seed: int = 0) -> None:
        super().__init__(hyperparams=hyperparams, random_seed=random_seed)
        self._fitted: bool = False
        self._random_state = numpy.random.RandomState(self.random_seed)

    def set_training_data(self, *, inputs: Inputs) -> None:  # type: ignore
        pass

    def fit(self, *, timeout: float = None, iterations: int = None) -> base.CallResult[None]:
        return base.CallResult(None)

    def produce(self, *, inputs: Inputs, timeout: float = None, iterations: int = None) -> base.CallResult[Outputs]:
        sample_dataframes = []
        if not self._fitted:
            n_rows = len(inputs)
            weights = get_distribution_weights(n_rows, self.hyperparams['distribution'])

            n_new_rows = int(n_rows * self.hyperparams['sample_size'])
            new_metadata = inputs.metadata.update((), {'dimension': {'length': n_new_rows, }, })

            for i in range(self.hyperparams['n_output_dataframes']):
                sample_dataframe = inputs.sample(
                    n=n_new_rows, weights=weights, replace=False,
                    random_state=numpy.random.RandomState(self.random_seed)).sort_index().reset_index(drop=True)
                sample_dataframe.metadata = new_metadata
                sample_dataframes.append(sample_dataframe)
                self._fitted = True
        else:
            sample_dataframes = self._repeated_dataframes(inputs)

        outputs = container.List(sample_dataframes, generate_metadata=True)
        return base.CallResult(outputs)

    def _repeated_dataframes(self, inputs: Inputs) -> typing.List[Inputs]:
        sample_dataframes = []
        for i in range(self.hyperparams['n_output_dataframes']):
            sample_dataframes.append(inputs.copy())
        return sample_dataframes

    def get_params(self) -> Params:
        return Params(
            fitted=self._fitted,
            random_seed=self.random_seed
        )

    def set_params(self, *, params: Params) -> None:
        self._fitted = params['fitted']
        self.random_seed = params['random_seed']


def get_distribution_weights(n_rows: int, dist: str) -> typing.List[float]:
    if dist == 'uniform':
        return list([1 / n_rows] * n_rows)
    elif dist == 'binomial':
        x_index = list(range(0, n_rows))
        return list(stats.binom.pmf(x_index, n_rows, 0.5))
    elif dist == 'geometric':
        x_index = list(range(1, n_rows + 1))
        return list(stats.geom.pmf(x_index, 1/n_rows))
    elif dist == 'poisson':
        x_index = list(range(0, n_rows))
        return list(stats.poisson.pmf(x_index, n_rows / 2))
    elif dist == 'normal':
        x_index = list(range(0, n_rows))
        return list(stats.norm.pdf(x_index, (n_rows-1)/2, n_rows/6))
    else:
        raise ValueError('Distribution {} not supported'.format(dist))
