#  Copyright (c) 2019-2020, RTE (https://www.rte-france.com)
#  See AUTHORS.txt
#  This Source Code Form is subject to the terms of the Apache License, version 2.0.
#  If a copy of the Apache License, version 2.0 was not distributed with this file, you can obtain one at http://www.apache.org/licenses/LICENSE-2.0.
#  SPDX-License-Identifier: Apache-2.0
#  This file is part of hadar-simulator, a python adequacy library for everyone.
import multiprocessing
from typing import Dict, Union

import numpy as np
import pandas as pd
from numpy.random.mtrand import randint

from hadar.workflow.pipeline import Pipeline, TO_SHUFFLER, Stage

__all__ = ['Shuffler', 'Timeline']


class Timeline:
    """
    Manage data used to generate timeline. Perform sampling too.
    """
    def __init__(self, data: np.ndarray = None, sampler = randint):
        """
        Instantiate.

        :param data: data to use for sampling
        :param sampler: sampler algorithm. Default it's np.random.randint
        """
        self.data = data
        self.sampler = sampler

    def sample(self, nb) -> np.ndarray:
        """
        Perform sampling. Compute data is needed before.

        :param nb: number of sampling
        :return: scenario matrix shape like (nb, horizon)
        """
        if self.data is None:
            self.data = self.compute()

        sampling = self.sampler(0, self.data.shape[0], nb)
        return self.data[sampling]

    def compute(self) -> np.ndarray:
        """
        Compute method called before sampling. For Timeline method just return data.

        :return: return data given in constructor
        """
        return self.data


class TimelinePipeline(Timeline):
    """
    Manage data timeline from pipeline generation.
    """
    def __init__(self, data: pd.DataFrame, pipeline: Pipeline, sampler = randint):
        """
        Instantiate.

        :param data: data for input pipeline
        :param pipeline: pipeline to use to generate data
        :param sampler: sampler algorithm. default it's np.random.randint
        """
        Timeline.__init__(self, sampler=sampler)

        self.df = data
        self.pipeline = pipeline

    def compute(self) -> np.ndarray:
        """
        Method called before sampling. For TimelinePipeline it will generate data by pipeline.

        :return: data generated by pipeline
        """
        res = self.pipeline(self.df)
        drop_columns = res.columns.get_level_values(1).unique().drop(TO_SHUFFLER)
        if drop_columns.any():
            res = res.drop(drop_columns, axis=1, level=1)
        return res.values.T


def compute(params):
    """
    Wrapper method to call Timeline.sample used by multiprocessing.Pool.

    :param params: (timeline, number of scenarios, timeline name)
    :return: (name, sampling)
    """
    tl, nb, name = params
    return name, tl.sample(nb)


class Shuffler:
    """
    Receive all data sources like raw matrix or pipeline.
    Schedule pipeline generation and shuffle all timeline to create scenarios.
    """
    def __init__(self, sampler=randint):
        """
        Instantiate.

        :param sampler: sampler algorithm. default it's np.random.randint
        """
        self.timelines = dict()
        self.sampler = sampler

    def add_data(self, name: str, data: np.ndarray):
        """
        Add raw data by numpy array.
        If you generate data by pipeline use add_pipeline. It will parallelize computation and manage swap.
        :param name: timeline name
        :param data: numpy array with shape as (scenario, horizon)
        :return: self
        """
        self.timelines[name] = Timeline(data, sampler=self.sampler)
        return self

    def add_pipeline(self, name: str, data: pd.DataFrame, pipeline: Pipeline):
        """
        Add data by pipeline and input data for pipeline.

        :param name: timeline name
        :param data: data to use as pipeline input
        :param pipeline: pipeline to generate data
        :return: self
        """
        data = Stage.standardize_column(data)
        pipeline.assert_computable(data)
        pipeline.assert_to_shuffler()

        self.timelines[name] = TimelinePipeline(data, pipeline, sampler=self.sampler)

    def shuffle(self, nb_scn):
        """
        Start pipeline generation and shuffle result to create scenario sampling.

        :param nb_scn: number of scenarios to sample
        :return:
        """
        # Compute pipelines
        pool = multiprocessing.Pool()
        res = pool.map(compute, ((tl, nb_scn, name) for name, tl in self.timelines.items()))
        return dict(res)


