import model
import utils
import numpy as np

class ThomasModel(model.Model):

    def __init__(self, dog_day_prob, message_prob, n_messengers = 1):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> m.probs
        {(1, 1, 1, 0, 0): 0.0015625000000000003, (1, 0, 1, 1, 0): 0.0046874999999999981, (1, 1, 0, 1, 1): 0.0046874999999999981, (1, 0, 0, 0, 0): 0.014062500000000011, (0, 0, 0, 0, 1): 0.12656249999999999, (1, 0, 1, 0, 0): 0.0046874999999999981, (0, 0, 1, 1, 1): 0.042187499999999989, (1, 1, 1, 1, 0): 0.0015625000000000003, (0, 1, 0, 1, 1): 0.042187499999999989, (1, 0, 0, 1, 0): 0.014062500000000011, (0, 0, 1, 0, 1): 0.042187499999999989, (0, 1, 1, 1, 0): 0.014062500000000011, (0, 0, 1, 0, 0): 0.042187499999999989, (1, 1, 1, 0, 1): 0.0015625000000000003, (0, 1, 0, 0, 0): 0.042187499999999989, (1, 0, 1, 1, 1): 0.0046874999999999981, (1, 1, 0, 1, 0): 0.0046874999999999981, (1, 1, 0, 0, 0): 0.0046874999999999981, (1, 0, 1, 0, 1): 0.0046874999999999981, (0, 0, 1, 1, 0): 0.042187499999999989, (0, 1, 1, 0, 0): 0.014062500000000011, (0, 0, 0, 1, 0): 0.12656249999999999, (1, 0, 0, 0, 1): 0.014062500000000011, (0, 0, 0, 0, 0): 0.12656249999999999, (1, 1, 1, 1, 1): 0.0015625000000000003, (0, 0, 0, 1, 1): 0.12656249999999999, (1, 0, 0, 1, 1): 0.014062500000000011, (0, 1, 0, 0, 1): 0.042187499999999989, (0, 1, 1, 0, 1): 0.014062500000000011, (0, 1, 0, 1, 0): 0.042187499999999989, (1, 1, 0, 0, 1): 0.0046874999999999981, (0, 1, 1, 1, 1): 0.014062500000000011}
        >>> utils.print_partition(m.partitions)
        player 0
        (1, 1, None, None)   [(1, 1, 0, 0, 0), (1, 1, 0, 0, 1), (1, 1, 1, 0, 0), (1, 1, 1, 0, 1)]
        (1, 1, 0, None)   [(1, 1, 0, 1, 0), (1, 1, 0, 1, 1)]
        (0, 1, 0, None)   [(0, 1, 0, 1, 0), (0, 1, 0, 1, 1)]
        (None, None, None, None)   [(0, 0, 0, 0, 0), (0, 0, 0, 0, 1), (0, 0, 0, 1, 0), (0, 0, 0, 1, 1), (0, 0, 1, 0, 0), (0, 0, 1, 0, 1), (0, 0, 1, 1, 0), (0, 0, 1, 1, 1), (1, 0, 0, 0, 0), (1, 0, 0, 0, 1), (1, 0, 0, 1, 0), (1, 0, 0, 1, 1), (1, 0, 1, 0, 0), (1, 0, 1, 0, 1), (1, 0, 1, 1, 0), (1, 0, 1, 1, 1)]
        (0, 1, 1, 1)   [(0, 1, 1, 1, 1)]
        (1, 1, 1, 1)   [(1, 1, 1, 1, 1)]
        (0, 1, 1, 0)   [(0, 1, 1, 1, 0)]
        (0, 1, None, None)   [(0, 1, 0, 0, 0), (0, 1, 0, 0, 1), (0, 1, 1, 0, 0), (0, 1, 1, 0, 1)]
        (1, 1, 1, 0)   [(1, 1, 1, 1, 0)]
        player 1
        (1, 0, 1, None)   [(1, 0, 1, 0, 0), (1, 0, 1, 0, 1), (1, 0, 1, 1, 0), (1, 0, 1, 1, 1)]
        (0, 0, 1, None)   [(0, 0, 1, 0, 0), (0, 0, 1, 0, 1), (0, 0, 1, 1, 0), (0, 0, 1, 1, 1)]
        (None, None, None, None)   [(0, 0, 0, 0, 0), (0, 0, 0, 0, 1), (0, 0, 0, 1, 0), (0, 0, 0, 1, 1), (0, 1, 0, 0, 0), (0, 1, 0, 0, 1), (0, 1, 0, 1, 0), (0, 1, 0, 1, 1), (1, 0, 0, 0, 0), (1, 0, 0, 0, 1), (1, 0, 0, 1, 0), (1, 0, 0, 1, 1), (1, 1, 0, 0, 0), (1, 1, 0, 0, 1), (1, 1, 0, 1, 0), (1, 1, 0, 1, 1)]
        (0, 1, 1, 1)   [(0, 1, 1, 1, 1)]
        (1, 1, 1, 1)   [(1, 1, 1, 1, 1)]
        (1, 1, 1, None)   [(1, 1, 1, 0, 0), (1, 1, 1, 1, 0)]
        (1, 1, 1, 0)   [(1, 1, 1, 0, 1)]
        (0, 1, 1, 0)   [(0, 1, 1, 0, 1)]
        (0, 1, 1, None)   [(0, 1, 1, 0, 0), (0, 1, 1, 1, 0)]
        """
        self.dog_day_prob = dog_day_prob
        self.message_prob = message_prob
        self.n_messengers = n_messengers
        self.message_length = 4
        super(ThomasModel, self).__init__([(0,1)]*(1 + self.message_length*self.n_messengers), 2)

    def run(self, variables):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> m.run((1,0,1,0,0))
        [(None, None, None, None), (1, 0, 1, None)]
        >>> m.run((1,0,1,1,0))
        [(None, None, None, None), (1, 0, 1, None)]
        >>> m.run((1,1,0,0,0))
        [(1, 1, None, None), (None, None, None, None)]
        >>> m.run((1,1,0,1,0))
        [(1, 1, 0, None), (None, None, None, None)]
        >>> m.run((1,1,1,1,0))
        [(1, 1, 1, 0), (1, 1, 1, None)]
        >>> m = ThomasModel(0.1, 0.25, 2)
        >>> m.run((1,1,1,1,0,0,1,0,0))
        [(1, 1, 1, 0, None, None, None, None), (1, 1, 1, None, 1, 0, 1, None)]
        >>> m.run((1,0,0,0,0,1,1,1,1))
        [(None, None, None, None, 1, 1, 1, 1), (None, None, None, None, 1, 1, 1, 1)]
        """

        out = [[] for i in range(self.n_players)]

        dog_day = variables[0]
        for i in range(self.n_messengers):
            this_out = self.messenger(dog_day, *variables[(1 + i*self.message_length):(1 + (i+1)*self.message_length)])
            for j in range(self.n_players):
                out[j] += this_out[j]

        out = [tuple(out[i]) for i in range(self.n_players)]

        return out

    # encodes the generative model of the messenger
    # d indicates if it is a "dog day"
    def messenger(self, d, m0, m1, tp, wtp):

        out = []

        if m0 == 1: # visiting player 0
            if tp == 1: # tell plan
                if m1 == 1: # visiting player 1
                    out += [[d, m0, m1, wtp]]
                else:
                    out += [[d, m0, m1, None]]
            else:
                out += [[d, m0, None, None]]
        else:
            out += [[None, None, None, None]]

        if m1 == 1: # visiting player 1
            if m0 == 1 and wtp == 1: # will tell player 0 plan to player 1, if it exists
                out += [[d, m0, m1, tp]]
            else:
                out += [[d, m0, m1, None]]
        else:
            out += [[None, None, None, None]]

        return out

    def evaluate(self, variables, log = False):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> np.abs(m.evaluate((1,0,1,0,0)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        >>> np.abs(np.exp(m.evaluate((1,0,1,0,0), True)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        >>> m = ThomasModel(0.1, 0.25, 2)
        >>> np.abs(np.exp(m.evaluate((1,0,1,0,0,1,1,0,0), True)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5 * 0.25 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        """

        log_p = 0

        log_p += utils.log_bernoulli(variables[0], self.dog_day_prob)
        for i in range(self.n_messengers):
            log_p += self.messenger_log_prob(*variables[(1 + i*self.message_length):(1 + (i+1)*self.message_length)])

        if log:
            return log_p
        else:
            return np.exp(log_p)

    def messenger_log_prob(self, m0, m1, tp, wtp):

        log_p = 0
        log_p += utils.log_bernoulli(m0, self.message_prob)
        log_p += utils.log_bernoulli(m1, self.message_prob)
        log_p += utils.log_bernoulli(tp, 0.5)
        log_p += utils.log_bernoulli(wtp, 0.5)

        return log_p

if __name__ == "__main__":
    import doctest
    doctest.testmod()


from collections import OrderedDict
from typing import Callable, List, Optional, Tuple, Dict

import numpy as np
from scipy import sparse as sps

from ._myfm import (
    FM,
    FMLearningConfig,
    ConfigBuilder,
    FMHyperParameters,
    LearningHistory,
    Predictor,
    RelationBlock,
    TaskType,
    create_train_fm,
)
from .base import (
    REAL,
    ArrayLike,
    MyFMBase,
    RegressorMixin,
    ClassifierMixin,
    check_data_consistency,
    std_cdf,
)

try:
    import pandas as pd
except:
    pd = None

class MyFMGibbsBase(
    MyFMBase[
        FM,
        FMHyperParameters,
        Predictor,
        LearningHistory,
    ]
):
    def _predict_core(
        self,
        X: Optional[ArrayLike],
        X_rel: List[RelationBlock] = [],
        n_workers: Optional[int] = None,
    ) -> np.ndarray:

        predictor = self._fetch_predictor()
        shape = check_data_consistency(X, X_rel)
        if X is None:
            X = sps.csr_matrix((shape, 0), dtype=REAL)
        else:
            X = sps.csr_matrix(X)
        if n_workers is None:
            return predictor.predict(X, X_rel)
        else:
            return predictor.predict_parallel(X, X_rel, n_workers)

    @classmethod
    def _train_core(
        cls,
        rank: int,
        init_stdev: float,
        X: sps.csr_matrix,
        X_rel: List[RelationBlock],
        y: np.ndarray,
        random_seed: int,
        config: FMLearningConfig,
        callback: Callable[[int, FM, FMHyperParameters, LearningHistory], bool],
    ) -> Tuple[Predictor, LearningHistory]:
        return create_train_fm(
            rank, init_stdev, X, X_rel, y, random_seed, config, callback
        )

    def get_hyper_trace(self) -> "pd.DataFrame":
        if pd is None:
            raise RuntimeError("Require pandas for get_hyper_trace.")

        if (self.n_groups_ is None) or (self.history_ is None):
            raise RuntimeError("Sampler not run yet.")

        columns = (
            ["alpha"]
            + ["mu_w[{}]".format(g) for g in range(self.n_groups_)]
            + ["lambda_w[{}]".format(g) for g in range(self.n_groups_)]
            + [
                "mu_V[{},{}]".format(g, r)
                for g in range(self.n_groups_)
                for r in range(self.rank)
            ]
            + [
                "lambda_V[{},{}]".format(g, r)
                for g in range(self.n_groups_)
                for r in range(self.rank)
            ]
        )

        res = []
        for hyper in self.history_.hypers:
            res.append(
                np.concatenate(
                    [
                        [hyper.alpha],
                        hyper.mu_w,
                        hyper.lambda_w,
                        hyper.mu_V.ravel(),
                        hyper.lambda_V.ravel(),
                    ]
                )
            )
        res = np.vstack(res)

        df = pd.DataFrame(res)
        df.columns = columns
        return df

class MyFMGibbsRegressor(RegressorMixin[FM, FMHyperParameters], MyFMGibbsBase):
    def fit(
        self,
        X: ArrayLike,
        y: np.ndarray,
        X_rel: List[RelationBlock] = [],
        X_test: Optional[ArrayLike] = None,
        y_test: Optional[np.ndarray] = None,
        X_rel_test: List[RelationBlock] = [],
        n_iter: int = 100,
        n_kept_samples: Optional[int] = None,
        grouping: Optional[List[int]] = None,
        group_shapes: Optional[List[int]] = None,
        callback: Optional[
            Callable[
                [int, FM, FMHyperParameters, LearningHistory],
                Tuple[bool, Optional[str]],
            ]
        ] = None,
        config_builder: Optional[ConfigBuilder] = None,
    ) -> "MyFMGibbsRegressor":
        """Performs Gibbs sampling to fit the data.

        Parameters
        ----------
        X : 2D array-like.
            Input variable.

        y : 1D array-like.
            Target variable.

        X_rel: list of RelationBlock, optional (default=[])
               Relation blocks which supplements X.

        n_iter : int, optional (default = 100)
            Iterations to perform.

        n_kept_samples: int, optional (default = None)
            The number of samples to store.
            If `None`, the value is set to `n_iter` - 5.

        grouping: Integer List, optional (default = None)
            If not `None`, this specifies which column of X belongs to which group.
            That is, if grouping[i] is g, then, :math:`w_i` and :math:`V_{i, r}`
            will be distributed according to
            :math:`\mathcal{N}(\mu_w[g], \lambda_w[g])` and :math:`\mathcal{N}(\mu_V[g, r], \lambda_V[g,r])`,
            respectively.
            If `None`, all the columns of X are assumed to belong to a single group, 0.

        group_shapes: Integer array, optional (default = None)
            If not `None`, this specifies each variable group's size.
            Ignored if grouping is not None.
            For example, if ``group_shapes = [n_1, n_2]``,
            this is equivalent to ``grouping = [0] * n_1 + [1] * n_2``

        callback: function(int, fm, hyper, history) -> (bool, str), optional(default = None)
            Called at the every end of each Gibbs iteration.
        """
        self._fit(
            X,
            y,
            X_rel=X_rel,
            X_test=X_test,
            X_rel_test=X_rel_test,
            y_test=y_test,
            n_iter=n_iter,
            n_kept_samples=n_kept_samples,
            grouping=grouping,
            callback=callback,
            group_shapes=group_shapes,
            config_builder=config_builder,
        )
        return self

    def predict(
        self,
        X: Optional[ArrayLike],
        X_rel: List[RelationBlock] = [],
        n_workers: Optional[int] = None,
    ) -> np.ndarray:
        """Make a prediction by compute the posterior predictive mean.

        Parameters
        ----------
        X : Optional[ArrayLike]
            Main table. When None, treated as a matrix with no column.
        X_rel : List[RelationBlock]
            Relations.
        n_workers : Optional[int], optional
            The number of threads to compute the posterior predictive mean, by default None

        Returns
        -------
        np.ndarray
            One-dimensional array of predictions.
        """
        return self._predict(X, X_rel, n_workers=n_workers)

class MyFMGibbsClassifier(ClassifierMixin[FM, FMHyperParameters], MyFMGibbsBase):
    r"""Bayesian Factorization Machines for binary classification tasks."""

    def fit(
        self,
        X: ArrayLike,
        y: np.ndarray,
        X_rel: List[RelationBlock] = [],
        X_test: Optional[ArrayLike] = None,
        y_test: Optional[np.ndarray] = None,
        X_rel_test: List[RelationBlock] = [],
        n_iter: int = 100,
        n_kept_samples: Optional[int] = None,
        grouping: Optional[List[int]] = None,
        group_shapes: Optional[List[int]] = None,
        callback: Optional[
            Callable[
                [int, FM, FMHyperParameters, LearningHistory],
                Tuple[bool, Optional[str]],
            ]
        ] = None,
        config_builder: Optional[ConfigBuilder] = None,
    ) -> "MyFMGibbsClassifier":
        """Performs Gibbs sampling to fit the data.

        Parameters
        ----------
        X : 2D array-like.
            Input variable.

        y : 1D array-like.
            Target variable.

        X_rel: list of RelationBlock, optional (default=[])
            Relation blocks which supplements X.

        n_iter : int, optional (default = 100)
            Iterations to perform.

        n_kept_samples: int, optional (default = None)
            The number of samples to store.
            If `None`, the value is set to `n_iter` - 5.

        grouping: Integer List, optional (default = None)
            If not `None`, this specifies which column of X belongs to which group.
            That is, if grouping[i] is g, then, :math:`w_i` and :math:`V_{i, r}`
            will be distributed according to
            :math:`\mathcal{N}(\mu_w[g], \lambda_w[g])` and :math:`\mathcal{N}(\mu_V[g, r], \lambda_V[g,r])`,
            respectively.
            If `None`, all the columns of X are assumed to belong to a single group, 0.

        group_shapes: Integer array, optional (default = None)
            If not `None`, this specifies each variable group's size.
            Ignored if grouping is not None.
            For example, if ``group_shapes = [n_1, n_2]``,
            this is equivalent to ``grouping = [0] * n_1 + [1] * n_2``

        callback: function(int, fm, hyper, history) -> (bool, str), optional(default = None)
            Called at the every end of each Gibbs iteration.
        """
        self._fit(
            X,
            y,
            X_rel=X_rel,
            X_test=X_test,
            X_rel_test=X_rel_test,
            y_test=y_test,
            n_iter=n_iter,
            n_kept_samples=n_kept_samples,
            grouping=grouping,
            callback=callback,
            group_shapes=group_shapes,
            config_builder=config_builder,
        )
        return self

    def predict(
        self,
        X: Optional[ArrayLike],
        X_rel: List[RelationBlock] = [],
        n_workers: Optional[int] = None,
    ) -> np.ndarray:
        """Based on the class probability, return binary classified outcome based on threshold = 0.5.
        If you want class probability instead, use `predict_proba` method.

        Parameters
        ----------
        Parameters
        ----------
        X : Optional[ArrayLike]
            When None, treated as a matrix with no column.
        X_rel : List[RelationBlock]
            Relations.
        n_workers : Optional[int], optional
            The number of threads to compute the posterior predictive mean, by default None

        Returns
        -------
        np.ndarray
            One-dimensional array of predicted outcomes.
        """
        return self._predict(X, X_rel, n_workers=n_workers)

    def predict_proba(
        self,
        X: Optional[ArrayLike],
        X_rel: List[RelationBlock] = [],
        n_workers: Optional[int] = None,
    ) -> np.ndarray:
        """Compute the probability that the outcome will be 1 based on posterior predictive mean.

        Parameters
        ----------
        Parameters
        ----------
        X : Optional[ArrayLike]
            When None, treated as a matrix with no column.
        X_rel : List[RelationBlock]
            Relations.
        n_workers : Optional[int], optional
            The number of threads to compute the posterior predictive mean, by default None

        Returns
        -------
        np.ndarray
            One-dimensional array of probabilities.

        """
        return self._predict_proba(X, X_rel, n_workers=n_workers)

class MyFMOrderedProbit(MyFMGibbsBase):
    """Bayesian Factorization Machines for Ordinal Regression Tasks."""

    @property
    def _task_type(self) -> TaskType:
        return TaskType.ORDERED

    def fit(
        self,
        X: ArrayLike,
        y: np.ndarray,
        X_rel: List[RelationBlock] = [],
        X_test: Optional[ArrayLike] = None,
        y_test: Optional[np.ndarray] = None,
        X_rel_test: List[RelationBlock] = [],
        n_iter: int = 100,
        n_kept_samples: Optional[int] = None,
        grouping: Optional[List[int]] = None,
        group_shapes: Optional[List[int]] = None,
        callback: Optional[
            Callable[
                [int, FM, FMHyperParameters, LearningHistory],
                Tuple[bool, Optional[str]],
            ]
        ] = None,
        cutpoint_group_configs: Optional[List[Tuple[int, np.ndarray]]] = None,
        callback_default_freq: int = 5,
    ) -> "MyFMOrderedProbit":
        config_builder = ConfigBuilder()
        y = np.asarray(y)
        if cutpoint_group_configs is None:
            n_class = y.max() + 1
            cutpoint_group_configs = [
                (int(n_class), np.arange(y.shape[0], dtype=np.int64))
            ]
        self.n_cutpoint_groups = len(cutpoint_group_configs)
        config_builder.set_cutpoint_groups(cutpoint_group_configs)
        super()._fit(
            X,
            y,
            X_rel=X_rel,
            X_test=X_test,
            y_test=y_test,
            X_rel_test=X_rel_test,
            n_iter=n_iter,
            n_kept_samples=n_kept_samples,
            grouping=grouping,
            callback=callback,
            group_shapes=group_shapes,
            config_builder=config_builder,
            callback_default_freq=callback_default_freq,
        )
        return self

    @classmethod
    def _score_to_class_prob(
        cls, score: np.ndarray, cutpoints: np.ndarray
    ) -> np.ndarray:
        score = std_cdf((cutpoints[np.newaxis, :] - score[:, np.newaxis]))
        score = np.hstack(
            [
                np.zeros((score.shape[0], 1), dtype=score.dtype),
                score,
                np.ones((score.shape[0], 1), dtype=score.dtype),
            ]
        )
        return score[:, 1:] - score[:, :-1]

    def _prepare_prediction_for_test(
        self, fm: FM, X: ArrayLike, X_rel: List[RelationBlock]
    ) -> np.ndarray:
        score = fm.predict_score(X, X_rel)
        return self._score_to_class_prob(score, fm.cutpoints[0])

    def _process_y(self, y: np.ndarray) -> np.ndarray:
        y_as_float = y.astype(np.float64)
        assert y.min() >= 0
        return y_as_float

    def _measure_score(cls, prediction: np.ndarray, y: np.ndarray) -> Dict[str, float]:
        result: Dict[str, float] = OrderedDict()
        accuracy = (np.argmax(prediction, axis=1) == y).mean()
        result["accuracy"] = accuracy
        log_loss = -np.log(
            prediction[np.arange(prediction.shape[0]), y.astype(np.int64)] + 1e-15
        ).mean()
        result["log_loss"] = log_loss
        return result

    def _status_report(cls, fm: FM, hyper: FMHyperParameters) -> str:
        log_str = "w0 = {:.2f}, ".format(fm.w0)
        if len(fm.cutpoints) == 1:
            log_str += "cutpoint = {} ".format(
                ["{:.3f}".format(c) for c in list(fm.cutpoints[0])]
            )
        return log_str

    def predict_proba(
        self,
        X: ArrayLike,
        X_rel: List[RelationBlock] = [],
        cutpoint_index: Optional[int] = None,
    ) -> np.ndarray:
        """Compute the ordinal class probability.

        Parameters
        ----------
        X : array_like
            The input data.
        X_rel : List[RelationBlock], optional
            Relational Block part of the data., by default []
        cutpoint_index : int, optional
            if not ``None`` and multiple cutpoints are enabled
            when ``fit``, compute the class probability
            based on the ``cutpoint_index``-th cutpoint, by default None.
            Must not be ``None`` when there are multiple cutpoints.

        Returns
        -------
        np.float
            The class probability
        """

        predictor = self._fetch_predictor()

        if cutpoint_index is None:
            if self.n_cutpoint_groups == 1:
                cutpoint_index = 0
            else:
                raise ValueError("specify the cutpoint index")

        X = sps.csr_matrix(X)
        if X.dtype != np.float64:
            X.data = X.data.astype(np.float64)
        p = 0

        for sample in predictor.samples:
            score = sample.predict_score(X, X_rel)
            p += self._score_to_class_prob(score, sample.cutpoints[cutpoint_index])
        return p / len(predictor.samples)

    def predict(
        self,
        X: ArrayLike,
        X_rel: List[RelationBlock] = [],
        cutpoint_index: Optional[int] = None,
    ) -> np.ndarray:
        """Predict the class outcome according to the class probability.

        Parameters
        ----------
        X : array_like
            The input data.
        X_rel : List[RelationBlock], optional
            Relational Block part of the data., by default []
        cutpoint_index : int, optional
            if not ``None`` and multiple cutpoints are enabled
            when ``fit``, compute the class probability
            based on the ``cutpoint_index``-th cutpoint, by default None
            Must not be ``None`` when there are multiple cutpoints.

        Returns
        -------
        np.int64
            The class prediction
        """

        return self.predict_proba(X, X_rel=X_rel, cutpoint_index=cutpoint_index).argmax(
            axis=1
        )


from collections import OrderedDict
from typing import Dict, Iterable, List, Set, Tuple

import numpy as np

from webdnn.backend.code_generator.allocator import MemoryLayout
from webdnn.backend.interface.graph_descriptor import IGraphDescriptor
from webdnn.backend.webassembly.kernel import Kernel
from webdnn.graph import traverse
from webdnn.graph.placeholder import Placeholder
from webdnn.graph.variable import Variable
from webdnn.graph.variables.attributes.constant import Constant
from webdnn.util import json

source_header = """
#include <stdlib.h>
#include <math.h>

float static_buffer[%%STATIC_SIZE%%];
float* dynamic_buffer = nullptr;

"""

source_init = """
extern "C" void init() {
    //static_buffer = (float*)malloc(%%STATIC_SIZE%% * sizeof(float));
}

extern "C" float* get_static_buffer(void) {
    return static_buffer;
}

extern "C" float* allocate_dynamic_buffer(int count) {
    if (dynamic_buffer) {
        free(dynamic_buffer);
        dynamic_buffer = nullptr;
    }
    dynamic_buffer = (float*)malloc(count * sizeof(float));
    return dynamic_buffer;
}

extern "C" float* get_dynamic_buffer(void) {
    return dynamic_buffer;
}
extern "C" void set_placeholder_value(int kernel_order, int offset, int value) {
    meta_buffers[kernel_order][offset] = value;
}
"""

source_exec = """
extern "C" void run() {
%%EXEC_LINES%%
}

"""

class GraphDescriptor(json.SerializableMixin, IGraphDescriptor):
    kernels: List[Kernel]
    memory_layout: MemoryLayout
    inputs: Iterable[Variable]
    outputs: Iterable[Variable]
    constants_encoding: str
    header_sources: Dict[str, str]
    footer_sources: Dict[str, str]
    required_heap: int
    licenses: Dict[str, str]

    def __init__(self,
                 kernels: List[Kernel],
                 memory_layout: MemoryLayout,
                 inputs: Iterable[Variable],
                 outputs: Iterable[Variable],
                 constants_encoding: str,
                 required_heap: int,
                 licenses: Dict[str, str]):
        self.kernels = kernels
        self.memory_layout = memory_layout
        self.inputs = inputs
        self.outputs = outputs
        self.constants_encoding = constants_encoding
        self.header_sources = OrderedDict()
        self.footer_sources = OrderedDict()
        self.required_heap = required_heap
        self.licenses = licenses

    def generate_top_source(self):
        self.header_sources["top"] = source_header \
            .replace("%%STATIC_SIZE%%", str(self.memory_layout.static_size))

    def generate_init_source(self):
        self.header_sources["init"] = source_init \
            .replace("%%STATIC_SIZE%%", str(self.memory_layout.static_size))

    # noinspection PyMethodMayBeStatic
    def generate_exec_line(self, kernel: Kernel, serial: int):
        line = f"{kernel.exec_info.entry_func_name}(meta_buf_{serial});\n"
        return line

    # noinspection PyMethodMayBeStatic
    def generate_kernel_metabuffer_initializer_line(self, kernel: Kernel, serial: int):
        line = "int meta_buf_" + str(serial) + "[] = {"
        # see as int32 and convert to 12345,67890
        line += ",".join(map(str, np.fromstring(kernel.exec_info.meta_buffer, dtype=np.int32).tolist()))
        line += "};\n"
        return line

    def generate_exec_source(self):
        lines = ""
        meta_buffer_initializer = ""
        for serial, kernel in enumerate(self.kernels):
            meta_buffer_initializer += self.generate_kernel_metabuffer_initializer_line(kernel, serial)
            lines += self.generate_exec_line(kernel, serial)
        meta_buffer_initializer += "int* meta_buffers[] = {"
        meta_buffer_initializer += ",".join(["meta_buf_" + str(serial) for serial in range(len(self.kernels))])
        meta_buffer_initializer += "};\n"
        self.footer_sources["exec"] = source_exec.replace("%%EXEC_LINES%%", lines)
        self.header_sources["meta_buffer_initializer"] = meta_buffer_initializer

    def concat_kernel_sources(self):
        func_sources = OrderedDict()

        for kernel in self.kernels:
            for func_name, source in kernel.func_sources.items():
                if func_name in func_sources:
                    assert func_sources[func_name] == source
                else:
                    func_sources[func_name] = source

        self.generate_top_source()
        self.generate_exec_source()
        self.generate_init_source()
        combined_source = \
            "".join(self.header_sources.values()) + \
            "\n".join(func_sources.values()) + \
            "".join(self.footer_sources.values())

        return combined_source

    def get_all_placeholders(self):
        unresolved_variables = []  # type: List[Tuple[int, Placeholder]]
        placeholders_set = set()  # type: Set[Placeholder]

        for kernel in self.kernels:
            unresolved_variables += kernel.exec_info.unresolved_value_list

        for offset, v in unresolved_variables:
            placeholders_set.update(v.get_depend_placeholders())

        placeholders = {p.label: None for p in placeholders_set}

        return placeholders

    def _to_serializable_(self):
        placeholders = self.get_all_placeholders()
        unresolved_value_lists = []
        for kernel in self.kernels:
            unresolved_value_list = kernel.exec_info.unresolved_value_list or []
            unresolved_value_lists.append([{"offset": v[0], "placeholder": v[1]} for v in unresolved_value_list])

        return {
            "weight_encoding": self.constants_encoding,
            "memory_layout": self.memory_layout,
            "placeholders": placeholders,
            "unresolved_value_lists": unresolved_value_lists,
            "inputs": [self.memory_layout[v].name for v in self.inputs if not traverse.check_attribute_match(v, Constant)],
            "outputs": [self.memory_layout[v].name for v in self.outputs],
            "licenses": self.licenses
        }


#!/usr/bin/env python
from warnings import warn

import numpy as np

from dataflow.lib.exporters import exports_json

from .refldata import ReflData, Intent
from .nexusref import load_nexus_entries, nexus_common, get_pol
from .nexusref import data_as, str_data
from .nexusref import TRAJECTORY_INTENTS
from .resolution import FWHM2sigma

def load_metadata(filename, file_obj=None):
    """
    Load the summary info for all entries in a NeXus file.
    """
    return load_nexus_entries(filename, file_obj=file_obj,
                              meta_only=True, entry_loader=Candor)

def load_entries(filename, file_obj=None, entries=None):
    #print("loading", filename, file_obj)
    return load_nexus_entries(filename, file_obj=file_obj, entries=entries,
                              meta_only=False, entry_loader=Candor)

#: Number of detector channels per detector tube on CANDOR
NUM_CHANNELS = 54

# CRUFT: these will be in the nexus/config.js eventually
S1_DISTANCE = -4335.86
S2_DISTANCE = -356.0
S3_DISTANCE = 356.0
DETECTOR_DISTANCE = 3496.0
MONO_WAVELENGTH = 4.75
MONO_WAVELENGTH_DISPERSION = 0.01

# CRUFT: these will be in the
_EFFICIENCY = None
def detector_efficiency():
    from pathlib import Path
    global _EFFICIENCY
    if _EFFICIENCY is None:
        path = Path(__file__).parent / 'DetectorWavelengths.csv'
        eff = np.loadtxt(path, delimiter=',')[:, 3]
        eff = eff / np.mean(eff)
        _EFFICIENCY = np.vstack((eff, eff)).T
    return _EFFICIENCY

class Candor(ReflData):
    """
    Candor data entry.

    See :class:`refldata.ReflData` for details.
    """
    format = "NeXus"
    probe = "neutron"

    def __init__(self, entry, entryname, filename):
        super().__init__()
        nexus_common(self, entry, entryname, filename)
        self.geometry = 'vertical'

    def load(self, entry):
        #print(entry['instrument'].values())
        das = entry['DAS_logs']
        n = self.points
        raw_intent = str_data(das, 'trajectoryData/_scanType')
        if raw_intent in TRAJECTORY_INTENTS:
            self.intent = TRAJECTORY_INTENTS[raw_intent]

        # Polarizers
        self.polarization = (
            get_pol(das, 'frontPolarization')
            + get_pol(das, 'backPolarization')
        )

        # Counts
        # Load counts early so we can tell whether channels are axis 1 or 2
        counts = data_as(das, 'areaDetector/counts', '', dtype='d')
        if counts is None or counts.size == 0:
            raise ValueError("Candor file '{self.path}' has no area detector data.".format(self=self))

        channels_at_end = (counts.shape[2] == NUM_CHANNELS)
        if channels_at_end:
            counts = np.swapaxes(counts, 1, 2)
        self.detector.counts = counts
        self.detector.counts_variance = counts.copy()
        self.detector.dims = counts.shape[1:]

        # Monochromator
        ismono = (str_data(das, 'monoTrans/key', 'OUT') == 'IN')
        if ismono:
            self.warn("monochromatic beams not yet supported for Candor reduction")
            self.monochromator.wavelength = data_as(entry, 'instrument/monochromator/wavelength', 'Ang', rep=n)
            # TODO: make sure wavelength_error is 1-sigma, not FWHM %
            self.monochromator.wavelength_resolution = data_as(entry, 'instrument/monochromator/wavelength_error', '', rep=n)
            # CRUFT: nexus config doesn't link in monochromator
            if self.monochromator.wavelength is None:
                self.monochromator.wavelength = data_as(das, 'mono/wavelength', 'Ang', rep=n)
            if self.monochromator.wavelength is None:
                self.warn("Wavelength is missing; using {MONO_WAVELENGTH} A".format(MONO_WAVELENGTH=MONO_WAVELENGTH))
                self.monochromator.wavelength = MONO_WAVELENGTH
            # TODO: generate wavelength_error directly with the nexus writer
            if self.monochromator.wavelength_resolution is None:
                dispersion = data_as(das, 'mono/wavelengthSpread', '', rep=n)
                if dispersion is not None:
                    self.monochromator.wavelength_resolution = \
                        FWHM2sigma(dispersion*self.monochromator.wavelength)
            if self.monochromator.wavelength_resolution is None:
                self.warn("Wavelength resolution is missing; using %.1f%% dL/L FWHM"
                        % (100*MONO_WAVELENGTH_DISPERSION,))
                self.monochromator.wavelength_resolution = \
                    FWHM2sigma(MONO_WAVELENGTH_DISPERSION*self.monochromator.wavelength)
            self.monochromator.wavelength = self.monochromator.wavelength[:, None, None]
            self.monochromator.wavelength_resolution = self.monochromator.wavelength_resolution[:, None, None]
        else:
            self.monochromator.wavelength = None
            self.monochromator.wavelength_resolution = None

        # Slit distances
        self.slit1.distance = data_as(entry, 'instrument/presample_slit1/distance', 'mm')
        self.slit2.distance = data_as(entry, 'instrument/presample_slit2/distance', 'mm')
        self.slit3.distance = data_as(entry, 'instrument/predetector_slit1/distance', 'mm')
        self.slit4.distance = data_as(entry, 'instrument/predetector_slit2/distance', 'mm')
        # CRUFT: old candor files don't populate instrument slits
        for slit, default in (
                (self.slit1, S1_DISTANCE),
                (self.slit2, S2_DISTANCE),
                (self.slit3, S3_DISTANCE),
                (self.slit4, DETECTOR_DISTANCE),
                ):
            if slit.distance is None:
                slit.distance = default

        # Slit openings
        # Note: Candor does not have slit4, but we translate the detector
        # mask value into a slit4 opening positioned right at the detector.
        # TODO: get values from instrument group instead of DAS_logs.
        for k, slit in enumerate((self.slit1, self.slit2, self.slit3, self.slit4)):
            x = 'slitAperture%d/softPosition'%(k+1)
            x_target = 'slitAperture%d/desiredSoftPosition'%(k+1)
            slit.x = data_as(das, x, 'mm', rep=n)
            slit.x_target = data_as(das, x_target, 'mm', rep=n)
            #y = 'vertSlitAperture%d/softPosition'%(k+1)
            #y_target = 'vertSlitAperture%d/desiredSoftPosition'%(k+1)
            #slit.y = data_as(das, y, 'mm', rep=n)
            #slit.y_target = data_as(das, y_target, 'mm', rep=n)
        # CRUFT: old files don't have the detector mask mapped
        # TODO: don't throw away datasets where the key is not defined:
        #       instead figure out an effective aperture size from motors or
        #       allow it to be overridden
        if self.slit4.x is None:
            mask_str = str_data(das, 'detectorMaskMap/key', '10')
            if not mask_str:
                mask_str = '10'
            try:
                mask = float(mask_str)
            except ValueError:
                raise ValueError("mask value {mask_str} cannot be converted to float".format(mask_str=mask_str))
            self.slit4.x_target = self.slit4.x = np.full(n, mask)

        # Detector
        wavelength = data_as(das, 'detectorTable/wavelengths', '')
        wavelength_spread = data_as(das, 'detectorTable/wavelengthSpreads', '')
        divergence = data_as(das, 'detectorTable/angularSpreads', '')
        efficiency = data_as(das, 'detectorTable/detectorEfficiencies', '')
        if divergence is None:
            divergence = np.repeat(0.01, wavelength.shape)
        if efficiency is None:
            efficiency = np.repeat(1.0, wavelength.shape)

        # TODO: check the orientation on detectorTable nodes.
        # For now they appear to be in "channels_at_end" order even if the
        # detector counts are in "banks_at_end" order.
        wavelength = wavelength.reshape(NUM_CHANNELS, -1)
        wavelength_spread = wavelength_spread.reshape(NUM_CHANNELS, -1)
        divergence = divergence.reshape(NUM_CHANNELS, -1)
        efficiency = efficiency.reshape(NUM_CHANNELS, -1)

        wavelength_resolution = FWHM2sigma(wavelength * wavelength_spread)
        if (efficiency == 1.0).all():
            efficiency = detector_efficiency()

        self.detector.wavelength = wavelength[None, :, :]
        self.detector.wavelength_resolution = wavelength_resolution[None, :, :]
        self.detector.efficiency = efficiency[None, :, :]
        self.angular_resolution = divergence[None, :, :]
        # TODO: sample broadening?

        # Angles
        self.sample.angle_x = data_as(das, 'sampleAngleMotor/softPosition', 'degree', rep=n)
        self.sample.angle_x_target = data_as(das, 'sampleAngleMotor/desiredSoftPosition', 'degree', rep=n)
        self.detector.angle_x = data_as(das, 'detectorTableMotor/softPosition', 'degree', rep=n)
        self.detector.angle_x_target = data_as(das, 'detectorTableMotor/desiredSoftPosition', 'degree', rep=n)
        self.detector.angle_x_offset = data_as(das, 'detectorTable/rowAngularOffsets', '')[0]

        #print("shapes", self.detector.counts.shape, self.detector.wavelength.shape, self.detector.efficiency.shape)
        #print("shapes", self.sample.angle_x.shape, self.detector.angle_x.shape, self.detector.angle_x_offset.shape)

    @property
    def Ti(self):
        return self.sample.angle_x[:, None, None]

    @property
    def Ti_target(self):
        return self.sample.angle_x_target[:, None, None]

    @property
    def Td(self):
        angle, offset = self.detector.angle_x, self.detector.angle_x_offset
        return angle[:, None, None] + offset[None, None, :]

    @property
    def Td_target(self):
        angle, offset = self.detector.angle_x_target, self.detector.angle_x_offset
        return angle[:, None, None] + offset[None, None, :]

    @property
    def Li(self):
        # Assume incident = reflected wavelength if no monochromator
        return (self.Ld if self.monochromator.wavelength is None
                else self.monochromator.wavelength)

    @property
    def Ld(self):
        return self.detector.wavelength  # Candor always has detector wavelength
        #return (self.Li if self.detector.wavelength is None
        #        else self.detector.wavelength)

    def plot(self, label=None):
        if label is None:
            label = self.name+self.polarization

        from matplotlib import pyplot as plt
        #v = np.log10(self.v + (self.v == 0))
        # set noise floor to ignore 5% of the data above zero
        v = self.v.copy()
        vmin = np.sort(v[v > 0])[int((v > 0).size * 0.01)]
        v[v<vmin] = vmin
        v = np.log10(v)

        channel = np.arange(1, NUM_CHANNELS+1)

        if True: # Qz - channel OR S1 - channel
            if Intent.isslit(self.intent):
                x, xlabel = self.slit1.x, "Slit 1 opening (mm)"
                x = x[:, None, None]
                x = np.concatenate((x, x), axis=2)
            else:
                x, xlabel = self.Qz, "Qz (1/Ang)"
            y, ylabel = np.vstack((channel, channel)).T[None, :, :], "channel"
            #print("Qz-channel", x.shape, y.shape, v.shape)

            #plt.figure()
            plt.subplot(211)
            plt.pcolormesh(x[:, :, 0].T, y[:, :, 0].T, v[:, :, 0].T)
            plt.xlabel(xlabel)
            plt.ylabel("{ylabel} - bank 0".format(ylabel=ylabel))
            plt.colorbar()
            plt.subplot(212)
            plt.pcolormesh(x[:, :, 1].T, y[:, :, 1].T, v[:, :, 1].T)
            plt.xlabel(xlabel)
            plt.ylabel("{ylabel} - bank 1".format(ylabel=ylabel))
            plt.colorbar()
            plt.suptitle('detector counts for {self.name}'.format(self=self))

        if False: # lambda-theta
            x, xlabel = self.detector.wavelength, 'detector wavelength (Ang)'
            y, ylabel = self.sample.angle_x, 'sample angle (degree)'
            #print("lambda-theta", x.shape, y.shape, v.shape)

            plt.figure()
            plt.subplot(211)
            plt.pcolormesh(edges(x[0, :, 0]), edges(y), v[:, :, 0])
            plt.xlabel(xlabel)
            plt.ylabel("{ylabel} - bank 0".format(ylabel=ylabel))
            plt.colorbar()
            plt.subplot(212)
            plt.pcolormesh(edges(x[0, :, 1]), edges(y), v[:, :, 1])
            plt.xlabel(xlabel)
            plt.ylabel("{ylabel} - bank 1".format(ylabel=ylabel))
            plt.colorbar()
            plt.suptitle('detector counts for {self.name}'.format(self=self))

        if False: # detector efficiency
            eff = self.detector.efficiency[0]
            #print("detector efficiency", eff.shape)

            plt.figure()
            plt.plot(channel, eff[:, 0], label='bank 0')
            plt.plot(channel, eff[:, 1], label='bank 1')
            plt.xlabel('channel number')
            plt.ylabel('wavelength efficiency')
            plt.suptitle('detectorTable.detectorEfficiencies for {self.name}'.format(self=self))
            plt.legend()

        if False: # detector wavelength
            y, ylabel = self.detector.wavelength, 'detector wavelength (Ang)'
            dy = self.detector.wavelength_resolution
            #print("detector wavelength", y.shape)

            plt.figure()
            plt.errorbar(channel, y[0, :, 0], yerr=dy[0, :, 0], label='bank 0')
            plt.errorbar(channel, y[0, :, 1], yerr=dy[0, :, 1], label='bank 1')
            plt.xlabel('channel number')
            plt.ylabel(ylabel)
            plt.suptitle('detectorTable.wavelength for {self.name}'.format(self=self))
            plt.legend()

    def get_axes(self):
        x, xlabel = self.detector.wavelength, "Wavelength (Ang)"
        #print("intent", self.intent)
        #if Intent.isslit(self.intent):
        #    y, ylabel = self.slit1.x, "S1 (mm)"
        #elif Intent.isspec(self.intent):
        #    y, ylabel = self.sample.angle_x, "Detector Angle (degrees)"
        if self.sample.angle_x.max() - self.sample.angle_x.min() > 0.01:
            y, ylabel = self.sample.angle_x, "Detector Angle (degrees)"
        elif self.slit1.x.max() - self.slit1.x.min() > 0.01:
            y, ylabel = self.slit1.x, "S1 (mm)"
        else:
            y, ylabel = np.arange(1, len(self.v)+1), "point"
        return (x, xlabel), (y, ylabel)

    def get_plottable(self):
        name = getattr(self, "name", "default_name")
        entry = getattr(self, "entry", "default_entry")
        def limits(vec, n):
            low, high = vec.min(), vec.max()
            delta = (high - low) / max(n-1, 1)
            # TODO: move range cleanup to plotter
            if delta == 0.:
                delta = vec[0]/10.
            return low - delta/2, high+delta/2
        data = self.v
        ny, nx, _ = data.shape  # nangles, nwavelengths, nbanks
        (x, xlabel), (y, ylabel) = self.get_axes()
        x, nx = np.hstack((x, x + (x.max()-x.min()))), 2*nx
        xmin, xmax = limits(x, nx)
        ymin, ymax = limits(y, ny)
        # TODO: self.detector.mask
        zmin, zmax = data.min(), data.max()
        # TODO: move range cleanup to plotter
        if zmin <= 0.:
            if (data > 0).any():
                zmin = data[data > 0].min()
            else:
                data[:] = zmin = 1e-10
            if zmin >= zmax:
                zmax = 10*zmin
        dims = {
            "xmin": xmin, "xmax": xmax, "xdim": nx,
            "ymin": ymin, "ymax": ymax, "ydim": ny,
            "zmin": zmin, "zmax": zmax,
        }
        # One z per detector bank
        #z = [data[..., k].ravel('F').tolist() for k in range(data.shape[-1])]
        # One z with banks back-to-back
        z = [data.ravel('F').tolist()]
        #print("data", data.shape, dims, len(z))
        plottable = {
            'type': '2d_multi',
            'dims': {'zmin': zmin, 'zmax': zmax},
            'datasets': [{'dims': dims, 'data': z[0]}],
            'entry': entry,
            'title': "%s:%s" % (name, entry),
            'options': {},
            'xlabel': xlabel,
            'ylabel': ylabel,
            'zlabel': 'Intensity (I)',
            'ztransform': 'log',
        }
        #print(plottable)
        return plottable

    # TODO: Define export format for partly reduced PSD data.
    @exports_json("json")
    def to_json_text(self):
        name = getattr(self, "name", "default_name")
        entry = getattr(self, "entry", "default_entry")
        return {
            "name": name,
            "entry": entry,
            "file_suffix": ".json",
            "value": self._toDict(),
            }

    # Kill column writer
    def to_column_text(self):
        pass

class QData(ReflData):
    def __init__(self, data, q, dq, v, dv):
        super().__init__()
        self.sample = data.sample
        self.probe = data.probe
        self.path = data.path
        self.uri = data.uri
        self.name = data.name
        self.entry = data.entry
        self.description = data.description
        self.polarization = data.polarization
        self.normbase = data.normbase
        self._q = q
        self._dq = dq
        self.v = v
        self.dv = dv
        self.intent = Intent.spec
        self.points = len(q)
        # TODO: set something for dL and dT
        self.angular_resolution = 1.
        self.scan_value = []
        self.scan_units = []
        self.scan_label = []

    @property
    def Qz(self):
        return self._q

    @property
    def Qx(self):
        return np.zeros_like(self._q)

    @property
    def dQ(self):
        return self._dq

def rebin(data, q):
    if data.normbase not in ("monitor", "time", "none"):
        raise ValueError("expected norm to be time, monitor or none")
    q_edges = edges(q, extended=True)
    datasets = []
    for bank in range(data.v.shape[2]):
        qz, dq, v, dv = _rebin_bank(data, bank, q_edges)
        #output = ReflData()
        #output.v, output.dv = v, dv
        #output.x, output.dx = q, dq
        datasets.append(QData(data, qz, dq, v, dv))

    # Only look at bank 0 for now
    return datasets[0]

def _rebin_bank(data, bank, q_edges):
    """
    Merge q points across channels and angles, returning q, dq, v, dv.

    Intensities (v, dv) are combined using poisson averaging.

    Q values (Q, dQ) for the combined measurements are weighted by intensity.
    This means that identical measurement conditions may give different
    (Q, dQ) depending on the randomness in the measured value.

    There is a simple expression for the moments of a mixture of distributions:
        T = (sum wk . T) / (sum wk)
        dT^2 = (sum wk . (dTk^2 + Tk^2)) / (sum wk) - T^2
    See: https://en.wikipedia.org/wiki/Mixture_density#Moments

    When rebinning Candor values as a function of Q then apply this directly
    to (Q, dQ) since the resolution function is curved in theta-lambda space.
    Note that the fitting function will need a weighted wavelength distribution
    for rare earths and for strongly absorbing samples where the scattering is
    wavelength dependent. It may suffice to compute the joint dQ and dL then
    assign dT so that it is consistent: this will underestimate the true dT,
    but it should be good enough for fitting. Alternatively, measure fewer
    angles for longer and do not merge points so that you can ignore
    wavelength variation within each theory value.
    """
    norm = data.normbase
    y, dy, q, dq = (val[:, :, bank].flatten()
                    for val in (data.v, data.dv, data.Qz, data.dQ))
    nbins = len(q_edges) - 1
    bin_index = np.searchsorted(q_edges, q)

    # Some bins may not have any points contributing, such as those before
    # and after, or those in the middle if the q-step is too fine.
    empty_q = (np.bincount(bin_index, minlength=nbins) == 0)
    # The following is cribbed from util.poisson_average, replacing
    # np.sum with np.bincount.
    # TODO: update poisson average so it handles grouping
    if norm == "none":
        bar_y = np.bincount(bin_index, weights=y, minlength=nbins)
        bar_dy = np.bincount(bin_index, weights=dy**2, minlength=nbins)
    else:
        # Counts must be positive for poisson averaging...
        y = y.copy()
        y[y < 0] = 0.
        dy = dy + (dy == 0) # protect against zero uncertainty
        monitors = y*(y+1)/dy**2 if norm == "monitor" else y/dy**2 # if "time"
        monitors[y == 0] = 1./dy[y == 0] # protect against zero counts
        counts = y*monitors
        combined_monitors = np.bincount(bin_index, weights=monitors, minlength=nbins)
        combined_counts = np.bincount(bin_index, weights=counts, minlength=nbins)
        combined_monitors += empty_q  # protect against division by zero
        bar_y = combined_counts/combined_monitors
        if norm == "time":
            bar_dy = bar_y * np.sqrt(1./combined_monitors)
        else:
            bar_dy = 1./combined_monitors * np.sqrt(1. + 1./combined_monitors)
            idx = (bar_y != 0)
            bar_dy[idx] = bar_y[idx] * np.sqrt(1./combined_counts[idx]
                                            + 1./combined_monitors[idx])

    # Find Q center and resolution, weighting by intensity
    # TODO: use intensity weighting when finding q centers?
    #weights = y
    w = np.ones_like(y)  # Weights must be positivie; use equal weights for now
    bar_w = np.bincount(bin_index, weights=w, minlength=nbins)
    bar_w += (bar_w == 0)  # protect against divide by zero
    bar_q = np.bincount(bin_index, weights=q*w, minlength=nbins)/bar_w
    # TODO: proper dQ calculation
    # Combined dq according to mixture distribution.
    #qsq = np.bincount(bin_index, weights=q**2*w, minlength=nbins)
    #dqsq = np.bincount(bin_index, weights=dq**2*w, minlength=nbins)
    #bar_dq = np.sqrt((dqsq + qsq)/bar_w - bar_q**2)
    # Combined dq according average of variance.
    #bar_dq = np.sqrt(np.bincount(bin_index, weights=dq*w, minlength=nbins)/bar_w)
    # Set dq to 1% for now...
    bar_dq = bar_q*0.01

    # Need to drop catch-all bins before and after q edges.
    # Also need to drop q bins which don't contain any values.
    keep = ~empty_q
    keep[0] = keep[-1] = False
    return bar_q[keep], bar_dq[keep], bar_y[keep], bar_dy[keep]

def edges(c, extended=False):
    r"""
    Linear bin edges given centers.

    If *extended* then create before/after bins so coverage is $(-\infty, \infty)$.
    """
    midpoints = (c[:-1]+c[1:])/2
    left = 2*c[0] - midpoints[0]
    right = 2*c[-1] - midpoints[-1]
    if extended:
        return np.hstack((-np.inf, left, midpoints, right, np.inf))
    else:
        return np.hstack((left, midpoints, right))

if __name__ == "__main__":
    from .nexusref import demo
    demo(loader=load_entries)


from __future__ import absolute_import
from typing import NoReturn
from .component import Component
import numpy as np
import sys
import shutil
from pathlib import Path
import os
from inversionson import InversionsonError
from salvus.mesh.unstructured_mesh import UnstructuredMesh
import h5py

class SalvusMeshComponent(Component):
    """
    Communications with Salvus Mesh.
    This will have to be done in a temporary way to begin with
    as it is not possible to make smoothiesem meshes through
    a nice config as things stand.

    :param infodict: Information related to inversion project
    :type infodict: Dictionary
    """

    def __init__(self, communicator, component_name):
        super(SalvusMeshComponent, self).__init__(communicator, component_name)
        self.meshes = Path(self.comm.project.lasif_root) / "MODELS"
        self.event_meshes = self.meshes / "EVENT_MESHES"
        self.average_meshes = self.meshes / "AVERAGE_MESHES"

    def create_mesh(self, event: str):
        """
        Create a smoothiesem mesh for an event. I'll keep refinements fixed
        for now.

        :param event: Name of event
        :type event: str
        """

        from salvus.mesh.simple_mesh import SmoothieSEM

        source_info = self.comm.lasif.get_source(event_name=event)
        if isinstance(source_info, list):
            source_info = source_info[0]
        sm = SmoothieSEM()
        sm.basic.model = "prem_ani_one_crust"
        sm.basic.min_period_in_seconds = self.comm.project.min_period
        sm.basic.elements_per_wavelength = 1.7
        sm.basic.number_of_lateral_elements = (
            self.comm.project.elem_per_quarter
        )
        sm.advanced.tensor_order = 4
        if self.comm.project.ellipticity:
            sm.spherical.ellipticity = 0.0033528106647474805
        if self.comm.project.ocean_loading["use"]:
            sm.ocean.bathymetry_file = self.comm.project.ocean_loading["file"]
            sm.ocean.bathymetry_varname = self.comm.project.ocean_loading[
                "variable"
            ]
            sm.ocean.ocean_layer_style = "loading"
            sm.ocean.ocean_layer_density = 1025.0
        if self.comm.project.topography["use"]:
            sm.topography.topography_file = self.comm.project.topography[
                "file"
            ]
            sm.topography.topography_varname = self.comm.project.topography[
                "variable"
            ]
        sm.source.latitude = source_info["latitude"]
        sm.source.longitude = source_info["longitude"]
        sm.refinement.lateral_refinements.append(
            {"theta_min": 40.0, "theta_max": 140.0, "r_min": 6250.0}
        )
        m = sm.create_mesh()
        mesh_file = self.event_meshes / event / "mesh.h5"
        if not os.path.exists(os.path.dirname(mesh_file)):
            os.makedirs(os.path.dirname(mesh_file))
        m.write_h5(mesh_file)

    def _check_if_mesh_has_field(
        self,
        check_mesh: str,
        field_name: str,
        elemental: bool,
        global_string: bool,
        side_sets: bool,
    ) -> bool:
        """
        Use h5py to quickly check whether field exists on mesh

        :param check_mesh: path to mesh to check
        :type check_mesh: str
        :param field_name: Name of field
        :type field_name: str
        :param elemental: Is field an elemental field
        :type elemental: bool
        :param global_string: Is it a global string
        :type global_string: bool
        :param side_sets: Are we checking for side sets? Provide the name,
        :type side_sets: str
        """
        with h5py.File(check_mesh, mode="r") as mesh:
            if global_string:
                global_strings = list(mesh["MODEL"].attrs.keys())
                if field_name in global_strings:
                    return True
                else:
                    return False
            if elemental:
                if "element_data" in mesh["MODEL"].keys():
                    elemental_fields = mesh["MODEL/element_data"].attrs.get(
                        "DIMENSION_LABELS"
                    )[1]
                    elemental_fields = (
                        elemental_fields[2:-2].replace(" ", "").split("|")
                    )
                    if field_name in elemental_fields:
                        return True
                    else:
                        return False
                else:
                    return False
            if side_sets:
                if "SIDE_SETS" in mesh.keys():
                    return True
                else:
                    return False
            else:
                # Here we assume it's an element_nodal_field
                nodal_fields = mesh["MODEL/data"].attrs.get(
                    "DIMENSION_LABELS"
                )[1]
                nodal_fields = nodal_fields[2:-2].replace(" ", "").split("|")
                if field_name in nodal_fields:
                    return True
                else:
                    return False

    def add_field_from_one_mesh_to_another(
        self,
        from_mesh: str,
        to_mesh: str,
        field_name: str,
        elemental: bool = False,
        global_string: bool = False,
        side_sets: bool = False,
        overwrite: bool = True,
    ):
        """
        Add one field from a specific mesh to another mesh. The two meshes
        need to have identical discretisations

        :param from_mesh: Path of mesh to copy field from
        :type from_mesh: str
        :param to_mesh: Path of mesh to copy field to
        :type to_mesh: str
        :param field_name: Name of the field to copy between them.
        :type field_name: str
        :param elemental: If the field is elemental make true, defaults to
            False
        :type elemental: bool, optional
        :param global_string: If the field is a global variable, defaults
            to False
        :type global_string: bool, optional
        :param overwrite: We check whether field is existing, if overwrite
            is True, we write it anyway, defaults to True
        :type bool, optional
        """
        import os
        import shutil

        # has_field = self._check_if_mesh_has_field(
        #     check_mesh=from_mesh,
        #     field_name=field_name,
        #     elemental=elemental,
        #     global_string=global_string,
        # )
        has_field = self._check_if_mesh_has_field(
            check_mesh=to_mesh,
            field_name=field_name,
            elemental=elemental,
            global_string=global_string,
            side_sets=side_sets,
        )
        if has_field and not overwrite:
            print(f"Field: {field_name} already exists on mesh")
            return
        attach_field = True
        if not os.path.exists(to_mesh):
            print(f"Mesh {to_mesh} does not exist. Will create new one.")
            shutil.copy(from_mesh, to_mesh)
            tm = UnstructuredMesh.from_h5(to_mesh)
            # tm.element_nodal_fields = {}
        else:
            tm = UnstructuredMesh.from_h5(to_mesh)
        fm = UnstructuredMesh.from_h5(from_mesh)
        if global_string:
            # if field_name in tm.global_strings.keys():
            #     if not overwrite:
            #         print(f"Field {field_name} already exists on mesh")
            #         return
            field = fm.global_strings[field_name]
            tm.attach_global_variable(name=field_name, data=field)
            tm.write_h5(to_mesh)
            print(f"Attached field {field_name} to mesh {to_mesh}")
            return
        elif elemental:
            # if field_name in tm.elemental_fields.keys():
            #     if not overwrite:
            #         print(f"Field {field_name} already exists on mesh")
            #         return
            field = fm.elemental_fields[field_name]
        elif side_sets:
            for side_set in fm.side_sets.keys():
                tm.define_side_set(
                    name=side_set,
                    element_ids=fm.side_sets[side_set][0],
                    side_ids=fm.side_sets[side_set][1]
                    )
                print(f"Attached side set {side_set} to mesh {to_mesh}")
            attach_field = False

        else:
            # if field_name in tm.element_nodal_fields.keys():
            #     if not overwrite:
            #         print(f"Field {field_name} already exists on mesh")
            #         return
            field = fm.element_nodal_fields[field_name]
        if attach_field:
            tm.attach_field(field_name, field)
            print(f"Attached field {field_name} to mesh {to_mesh}")
        tm.write_h5(to_mesh)

    def write_xdmf(self, filename: str):
        """
        A hacky way to write an xdmf file for the hdf5 file
        :param filename: path to hdf5 file
        :return:
        """

        mesh = UnstructuredMesh.from_h5(filename)
        mesh.write_h5(filename)

    def add_fluid_and_roi_from_lasif_mesh(self):
        """
        For some reason the salvus opt meshes don't have all the necessary info.
        I need this to get them simulation ready. I will write them into the
        lasif folder afterwards.
        As this is a quickfix, I will make it for my specific case.
        """
        import os
        import numpy as np

        initial_model = self.comm.lasif.lasif_comm.project.lasif_config[
            "domain_settings"
        ]["domain_file"]
        iteration = self.comm.project.current_iteration
        opt_mesh = os.path.join(
            self.comm.project.paths["salvus_opt"],
            "PHYSICAL_MODELS",
            f"{iteration}.h5",
        )
        m_opt = UnstructuredMesh.from_h5(opt_mesh)
        m_init = UnstructuredMesh.from_h5(initial_model)

        fluid = m_init.elemental_fields["fluid"]
        roi = np.abs(1.0 - fluid)

        m_opt.attach_field(name="fluid", data=fluid)
        m_opt.attach_field(name="ROI", data=roi)

        iteration_mesh = os.path.join(
            self.comm.project.lasif_root,
            "MODELS",
            f"ITERATION_{iteration}",
            "mesh.h5",
        )
        if not os.path.exists(os.path.dirname(iteration_mesh)):
            os.makedirs(os.path.dirname(iteration_mesh))
        m_opt.write_h5(iteration_mesh)

    def get_average_model(self, iteration_range: tuple) -> Path:
        """
        Get an average model between a list of iteration numbers.
        Can be used to get a smoother misfit curve for validation
        data set.

        :param iteration_range: From iteration to iteration tuple
        :type iterations: tuple
        """
        # I have to make sure the I am consistent with naming of things, might be a bit off there

        folder_name = f"it_{iteration_range[0]}_to_{iteration_range[1]}"
        full_path = self.average_meshes / folder_name / "mesh.h5"
        if not os.path.exists(full_path.parent):
            os.makedirs(full_path.parent)

        # We copy the newest mesh from SALVUS_OPT to LASIF and write the
        # average fields onto those.

        model = self.comm.salvus_opt.get_model_path()
        shutil.copy(model, full_path)

        m = UnstructuredMesh.from_h5(full_path)
        fields = m.element_nodal_fields
        new_fields = {}
        for field in fields.keys():
            new_fields[field] = np.zeros_like(fields[field])
        # m.element_nodal_fields = {}
        for iteration in range(iteration_range[0], iteration_range[1] + 1):
            it = self.comm.salvus_opt.get_name_for_accepted_iteration_number(
                number=iteration
            )
            model_path = self.comm.salvus_opt.get_model_path(iteration=it)
            m_tmp = UnstructuredMesh.from_h5(model_path)
            for field_name, field in new_fields.items():
                field += m_tmp.element_nodal_fields[field_name]

        for field_name, field in new_fields.items():
            field /= len(range(iteration_range[0], iteration_range[1] + 1))
            m.attach_field(field_name, field)
        m.write_h5(full_path)
        print(
            f"Wrote and average model of iteration {iteration_range[0]} to"
            f" iteration {iteration_range[1]} onto mesh: {full_path}"
        )

        return full_path

    def add_region_of_interest(self, event: str):
        """
        Region of interest is the region where the gradient is computed
        and outside the region, it is not computed.
        Currently we add the region of interest as an elemental field
        which is the oposite of the fluid field.

        :param event: Name of event
        :type event: str
        """

        mesh = self.comm.lasif.find_event_mesh(event)
        m = UnstructuredMesh.from_h5(mesh)
        mesh_layers = np.sort(np.unique(m.elemental_fields["layer"]))[
            ::-1
        ].astype(int)
        layers = m.elemental_fields["layer"]
        o_core_idx = layers[np.where(m.elemental_fields["fluid"] == 1)[0][0]]
        o_core_idx = np.where(mesh_layers == o_core_idx)[0][0]
        correct_layers = mesh_layers[o_core_idx:]
        roi = np.zeros_like(layers)
        for layer in correct_layers:
            roi = np.logical_or(roi, layers == layer)

        m.attach_field("ROI", roi)
        m.write_h5(mesh)

    def write_new_opt_fields_to_simulation_mesh(self):
        """
        Salvus opt makes a mesh which has the correct velocities but
        it does not have everything which is needed to run a simulation.
        We will thus write it's fields on to our simulation mesh.
        """
        if self.comm.project.meshes == "multi-mesh":
            raise InversionsonError(
                "Multi-mesh inversion should not use this function. Only "
                "Mono-mesh."
            )
        print("Writing new fields to simulation mesh")
        iteration = self.comm.project.current_iteration
        if "validation" in iteration:
            iteration = iteration[11:]  # We don't need a special mesh
        opt_model = os.path.join(
            self.comm.salvus_opt.models, f"{iteration}.h5"
        )
        simulation_mesh = self.comm.lasif.get_simulation_mesh(
            event_name=None, iteration="current"
        )

        sim_mesh_dir = os.path.dirname(simulation_mesh)
        success_file = os.path.join(sim_mesh_dir, "success.txt")

        if os.path.exists(simulation_mesh) and os.path.exists(success_file):
            print("Mesh already exists, will not add fields")
            return
        else:
            shutil.copy(
                self.comm.lasif.lasif_comm.project.lasif_config[
                    "domain_settings"
                ]["domain_file"],
                simulation_mesh,
            )

        with h5py.File(simulation_mesh, mode="r+") as f_new:
            with h5py.File(opt_model, mode="r") as f:
                dim_labels = (
                    f["MODEL/data"]
                    .attrs.get("DIMENSION_LABELS")[1][1:-1]
                    .replace(" ", "")
                    .split("|")
                )
                # This assumes the indices are the same in both files,
                # which seems to be the case as far as DP could tell.
                for param in self.comm.project.inversion_params:
                    print("Writing field:", param)
                    i = dim_labels.index(param)
                    f_new["MODEL/data"][:, i, :] = f["MODEL/data"][:, i, :]

        # When all fields are successfully copied write a file to indicate
        # success to prevent the issue that we continue
        # with the initial model when something here crashes unexpectedly.
        with open(success_file, "w") as text_file:
            text_file.write("All fields written successfully.")

    def sum_two_fields_on_a_mesh(
        self,
        mesh: str,
        fieldname_1: str,
        fieldname_2: str,
        newname: str = None,
        delete_old_fields: bool = False,
    ):
        """
        Take two fields on a mesh and sum them together. If no newname is
        given the summed field will be written into both the old fields.
        If newname is given the summed field will be written in there
        and if delete_old_fields is true they will be deleted of course.

        :param mesh: Path to mesh to be used
        :type mesh: str or Path
        :param fieldname_1: Name of field to be summed
        :type fieldname_1: str
        :param fieldname_2: Name of other field to be summed
        :type fieldname_2: str
        :param newname: Name of field to store summed field, defaults to None
        :type newname: str, optional
        :param delete_old_fields: Whether old fields should be deleted,
            defaults to False. Currently not implemented
        :type delete_old_fields: bool, optional
        """

        m = UnstructuredMesh.from_h5(mesh)

        available_fields = list(m.element_nodal_fields.keys())
        if fieldname_1 not in available_fields:
            raise InversionsonError(
                f"Field {fieldname_1} not available on mesh {mesh}. "
                f"Only available fields are: {available_fields}"
            )
        if fieldname_2 not in available_fields:
            raise InversionsonError(
                f"Field {fieldname_2} not available on mesh {mesh}. "
                f"Only available fields are: {available_fields}"
            )

        if delete_old_fields:
            if newname is not None:
                raise InversionsonError(
                    "If you want to delete old fields you need to write the "
                    "summed one into a new field"
                )

        summed_field = np.copy(m.element_nodal_fields[fieldname_1])
        summed_field += m.element_nodal_fields[fieldname_2]

        if newname is None:
            m.attach_field(fieldname_1, summed_field)
            m.attach_field(fieldname_2, summed_field)
            m.write_h5(mesh)

        else:
            m.attach_field(newname, summed_field)

    def fill_inversion_params_with_zeroes(self, mesh: str):
        """
        This is done because we don't interpolate every layer and then
        we want to make sure there is nothing sneaking into the gradients

        :param mesh: Path to mesh
        :type mesh: str
        """
        print("Filling inversion parameters with zeros before interpolation")
        m = UnstructuredMesh.from_h5(mesh)
        parameters = self.comm.project.inversion_params
        zero_nodal = np.zeros_like(m.element_nodal_fields[parameters[0]])

        for param in parameters:
            m.attach_field(param, zero_nodal)
        m.write_h5(mesh)


from tts_modules.encoder.models.dVecModel import DVecModel
from tts_modules.encoder.data.WavPreprocessor import StandardAudioPreprocessor
from tts_modules.encoder.data.Wav2MelTransform import StandardWav2MelTransform

import torch
import numpy as np
import yaml
import os

class SpeakerEncoderManager:
    def __init__(self, configs, model, checkpoint_path,  preprocessor=None, wav2mel=None):
        self.configs = configs
        self.preprocessor = preprocessor
        if preprocessor is None:
            self.preprocessor = StandardAudioPreprocessor(configs["AudioConfig"])
        self.wav2mel = wav2mel
        if wav2mel is None:
            self.wav2mel = StandardWav2MelTransform(configs["AudioConfig"])

        self.checkpoint_path = checkpoint_path
        self.current_embed = None
        with open(configs["AudioConfig"], "r") as ymlfile:
            self.AudioConfig = yaml.load(ymlfile)
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        if model is None:
            self.__init_dvec_model()
            self.__load_model()
        self.model = self.model.to(self.device)

    def __init_dvec_model(self):
        with open(self.configs["SpeakerEncoderConfig"], "r") as ymlfile:
            self.SpeakerEncoderConfig = yaml.load(ymlfile)
        self.model = DVecModel(self.device, self.device, self.SpeakerEncoderConfig)

    def process_speaker(self, speaker_speech_path, save_embeddings_path=None,
                        save_embeddings_speaker_name="test_speaker"):
        processed_wav = self.preprocessor.preprocess_wav(speaker_speech_path)

        embed = self.embed_utterance(processed_wav)
        self.current_embed = embed
        if save_embeddings_path is not None:
            self.save_embeddings(self, save_embeddings_path, save_embeddings_speaker_name)

        return embed

    def save_embeddings(self, save_embeddings_path,save_embeddings_speaker_name):
        np.save(os.path.join(save_embeddings_path,save_embeddings_speaker_name), self.current_embed)

    def __load_model(self):
        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state"])
        self.model.eval()

    def embed_utterance(self, wav, using_partials=True, return_partials=False):

        # Process the entire utterance if not using partials
        if not using_partials:
            # processed_wav = self.preprocessor.preprocess_wav(wav)
            frames = self.wav2mel.Wav2Mel(wav)
            frames = torch.from_numpy(frames[None, ...]).to(self.device)
            embed = self.model.forward(frames).detach().cpu().numpy()
            self.current_embed = embed[0]
            if return_partials:
                return embed[0], None, None

            return embed[0]

        # Compute where to split the utterance into partials and pad if necessary
        wave_slices, mel_slices = self.compute_partial_slices(len(wav))
        max_wave_length = wave_slices[-1].stop
        if max_wave_length >= len(wav):
            wav = np.pad(wav, (0, max_wave_length - len(wav)), "constant")

        # Split the utterance into partials
        # processed_wav = self.preprocessor.preprocess_wav(wav)
        frames = self.wav2mel.Wav2Mel(wav)
        frames_batch = np.array([frames[s] for s in mel_slices])
        frames = torch.from_numpy(frames_batch).to(self.device)
        partial_embeds = self.model.forward(frames).detach().cpu().numpy()

        # Compute the utterance embedding from the partial embeddings
        raw_embed = np.mean(partial_embeds, axis=0)
        embed = raw_embed / np.linalg.norm(raw_embed, 2)
        self.current_embed = embed
        if return_partials:
            return embed, partial_embeds, wave_slices
        return embed

    def compute_partial_slices(self, n_samples, min_pad_coverage=0.75, overlap=0.5):
        """
        Computes where to split an utterance waveform and its corresponding mel spectrogram to obtain
        partial utterances of <partial_utterance_n_frames> each. Both the waveform and the mel
        spectrogram slices are returned, so as to make each partial utterance waveform correspond to
        its spectrogram. This function assumes that the mel spectrogram parameters used are those
        defined in params_data.py.

        The returned ranges may be indexing further than the length of the waveform. It is
        recommended that you pad the waveform with zeros up to wave_slices[-1].stop.

        :param n_samples: the number of samples in the waveform
        :param partial_utterance_n_frames: the number of mel spectrogram frames in each partial
        utterance
        :param min_pad_coverage: when reaching the last partial utterance, it may or may not have
        enough frames. If at least <min_pad_coverage> of <partial_utterance_n_frames> are present,
        then the last partial utterance will be considered, as if we padded the audio. Otherwise,
        it will be discarded, as if we trimmed the audio. If there aren't enough frames for 1 partial
        utterance, this parameter is ignored so that the function always returns at least 1 slice.
        :param overlap: by how much the partial utterance should overlap. If set to 0, the partial
        utterances are entirely disjoint.
        :return: the waveform slices and mel spectrogram slices as lists of array slices. Index
        respectively the waveform and the mel spectrogram with these slices to obtain the partial
        utterances.
        """
        config = self.AudioConfig
        sampling_rate = config["SAMPLING_RATE"]
        mel_window_step = config["MEL_WINDOW_STEP"]
        partial_utterance_n_frames = config["PARTIAL_UTTERANCE_N_FRAMES"]

        assert 0 <= overlap < 1
        assert 0 < min_pad_coverage <= 1

        samples_per_frame = int((sampling_rate * mel_window_step / 1000))
        n_frames = int(np.ceil((n_samples + 1) / samples_per_frame))
        frame_step = max(int(np.round(partial_utterance_n_frames * (1 - overlap))), 1)

        # Compute the slices
        wav_slices, mel_slices = [], []
        steps = max(1, n_frames - partial_utterance_n_frames + frame_step + 1)
        for i in range(0, steps, frame_step):
            mel_range = np.array([i, i + partial_utterance_n_frames])
            wav_range = mel_range * samples_per_frame
            mel_slices.append(slice(*mel_range))
            wav_slices.append(slice(*wav_range))

        # Evaluate whether extra padding is warranted or not
        last_wav_range = wav_slices[-1]
        coverage = (n_samples - last_wav_range.start) / (last_wav_range.stop - last_wav_range.start)
        if coverage < min_pad_coverage and len(mel_slices) > 1:
            mel_slices = mel_slices[:-1]
            wav_slices = wav_slices[:-1]

        return wav_slices, mel_slices


from __future__ import absolute_import, division, print_function

import argparse
import logging
import os
import random
import sys

import numpy as np
import torch
from torch.utils.data import (DataLoader, RandomSampler, TensorDataset)
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
print( torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.version.cuda)
from models.modeling_bert import QuestionAnswering, Config
from utils.optimization import AdamW, WarmupLinearSchedule
from utils.tokenization import BertTokenizer
from utils.korquad_utils import read_squad_examples, convert_examples_to_features

if sys.version_info[0] == 2:
    import cPickle as pickle
else:
    import pickle

logger = logging.getLogger(__name__)

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def train(args, train_dataset, model):
    """ Train the model """
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = WarmupLinearSchedule(optimizer, warmup_steps=t_total*0.1, t_total=t_total)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        if args.fp16_opt_level == "O2":
            keep_batchnorm_fp32 = False
        else:
            keep_batchnorm_fp32 = True
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level, keep_batchnorm_fp32=keep_batchnorm_fp32)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        try:
            from apex.parallel import DistributedDataParallel as DDP
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
        model = DDP(model, message_size=250000000, gradient_predivide_factor=torch.distributed.get_world_size())

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    epochs = 0
    model.zero_grad()
    model.train()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Train(XX Epoch) Step(X/X) (loss=X.X)", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator):
            batch = tuple(t.to(args.device) for t in batch)  # multi-gpu does scattering it-self
            input_ids, input_mask, segment_ids, start_positions, end_positions = batch
            outputs = model(input_ids, segment_ids, input_mask, start_positions, end_positions)
            loss = outputs  # model outputs are always tuple in transformers (see doc)

            if args.n_gpu > 1:
                loss = loss.mean() # mean() to average on multi-gpu parallel (not distributed) training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                scheduler.step()  # Update learning rate schedule\
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1
                epoch_iterator.set_description(
                    "Train(%d Epoch) Step(%d / %d) (loss=%5.5f)" % (_, global_step, t_total, loss.item()))

        if args.local_rank in [-1, 0]:
            model_checkpoint = 'korquad_{0}_{1}_{2}_{3}.bin'.format(args.learning_rate,
                                                                    args.train_batch_size,
                                                                    epochs,
                                                                    int(args.num_train_epochs))
            logger.info(model_checkpoint)
            output_model_file = os.path.join(args.output_dir, model_checkpoint)
            if args.n_gpu > 1 or args.local_rank != -1:
                logger.info("** ** * Saving file * ** **(module)")
                torch.save(model.module.state_dict(), output_model_file)
            else:
                logger.info("** ** * Saving file * ** **")
                torch.save(model.state_dict(), output_model_file)
        epochs += 1
    logger.info("Training End!!!")

def load_and_cache_examples(args, tokenizer):
    # Load data features from cache or dataset file
    input_file = args.train_file
    cached_features_file = os.path.join(os.path.dirname(input_file), '_cached_{}_{}_{}'.format('train',
                                                                                str(args.max_seq_length),
                                                                                args.doc_stride))
    if os.path.exists(cached_features_file):
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", input_file)
        examples = read_squad_examples(input_file=args.train_file, is_training=True, version_2_with_negative=False)
        features = convert_examples_to_features(examples=examples,
                                                tokenizer=tokenizer,
                                                max_seq_length=args.max_seq_length,
                                                doc_stride=args.doc_stride,
                                                max_query_length=args.max_query_length,
                                                is_training=True)

        if args.local_rank in [-1, 0]:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file, pickle_protocol=pickle.HIGHEST_PROTOCOL)

    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
    all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)
    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_start_positions, all_end_positions)
    return dataset

def main():
    parser = argparse.ArgumentParser()
    # Parameters
    parser.add_argument("--output_dir", default='output', type=str,
                        help="The output directory where the model checkpoints and predictions will be written.")
    parser.add_argument("--checkpoint", default='pretrain_ckpt/large_v1_model.bin', type=str,
                        help="pre-trained model checkpoint")
    parser.add_argument("--config_file", default='data/large_config.json', type=str,
                        help="model configuration file")
    parser.add_argument("--vocab_file", default='data/large_v1_32k_vocab.txt', type=str,
                        help="tokenizer vocab file")

    parser.add_argument("--train_file", default='data/korquad/KorQuAD_v1.0_train.json', type=str,
                        help="SQuAD json for training. E.g., train-v1.1.json")

    parser.add_argument("--max_seq_length", default=512, type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. Sequences "
                             "longer than this will be truncated, and sequences shorter than this will be padded.")
    parser.add_argument("--doc_stride", default=128, type=int,
                        help="When splitting up a long document into chunks, how much stride to take between chunks.")
    parser.add_argument("--max_query_length", default=64, type=int,
                        help="The maximum number of tokens for the question. Questions longer than this will "
                             "be truncated to this length.")
    parser.add_argument("--max_answer_length", default=30, type=int,
                        help="The maximum length of an answer that can be generated. This is needed because the start "
                             "and end predictions are not conditioned on one another.")
    parser.add_argument("--per_gpu_train_batch_size", default=16, type=int,
                        help="Total batch size for training.")
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--num_train_epochs", default=4.0, type=float,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument("--adam_epsilon", default=1e-6, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--warmup_proportion", default=0.1, type=float,
                        help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% "
                             "of training.")
    parser.add_argument("--n_best_size", default=20, type=int,
                        help="The total number of n-best predictions to generate in the nbest_predictions.json "
                             "output file.")
    parser.add_argument("--verbose_logging", action='store_true',
                        help="If true, all of the warnings related to data processing will be printed. "
                             "A number of warnings are expected for a normal SQuAD evaluation.")
    parser.add_argument("--no_cuda", action='store_true',
                        help="Whether not to use CUDA when available")

    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--fp16',
                        action='store_true',
                        help="Whether to use 16-bit float precision instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O2',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    parser.add_argument("--local_rank", type=int, default=-1,
                        help="For distributed training: local_rank")
    parser.add_argument('--null_score_diff_threshold', type=float, default=0.0,
                        help="If null_score - best_non_null is greater than the threshold predict null.")

    args = parser.parse_args()

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend='nccl')
        args.n_gpu = 1
    args.device = device

    # Setup logging
    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt = '%m/%d/%Y %H:%M:%S',
                        level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)

    # Set seed
    set_seed(args)

    tokenizer = BertTokenizer(args.vocab_file,
                              max_len=args.max_seq_length,
                              do_basic_tokenize=True)
    # Prepare model
    config = Config.from_json_file(args.config_file)
    model = QuestionAnswering(config)
    model.bert.load_state_dict(torch.load(args.checkpoint))
    num_params = count_parameters(model)
    logger.info("Total Parameter: %d" % num_params)
    model.to(device)

    logger.info("Training hyper-parameters %s", args)

    # Training
    train_dataset = load_and_cache_examples(args, tokenizer)
    train(args, train_dataset, model)

if __name__ == "__main__":
    main()


"""This model is based partly on a similar model used by Thomas Davidson et al. in
the paper: Automated Hate Speech Detection and the Problem of Offensive Language.
Source: https://github.com/t-davidson/hate-speech-and-offensive-language"""

import numpy as np
import pandas as pd
import pickle
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedKFold, GridSearchCV
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix
from src.preprocess_utils import PreProcessUtils
from sklearn.utils import resample

import warnings
warnings.filterwarnings("ignore")
warnings.filterwarnings("ignore", category=DeprecationWarning)

PP = PreProcessUtils()

def lr_preprocessor(tweet):
    """
    Convert tweet to lowercase, stem words and remove or replace:
    1) Mentions i.e. @someone, 2) URLs i.e. https://www.twitter.com, 3) HTML entitiesRepeated,
    4) whitespace, 5) punctuation, 6) number, 7) stop words
    See PreProcessUtils Class for more details.
    :param tweet: text string
    :return pp_tweet: A pre-processed tweet
    """
    pp_tweet = tweet
    pp_tweet = PP.remove_noise(pp_tweet, mentions=True, replacement=' ')
    pp_tweet = PP.normalise(pp_tweet, numbers=True, stopwords=True, stem_words=True, replacement=' ')
    return pp_tweet

def lr_tokenizer(tweet):
    """Tokenize the tweet"""
    return PreProcessUtils.tokenize(tweet)

def train_model(data):
    """Train the model"""
    tweet_df = pd.read_csv(data)

    more_stopwords = ['rt', 'ff', 'tbt', 'ftw', 'dm']  # may add more later
    lr_stop_words = PP.stop_words(more_stopwords)

    vectorizer = TfidfVectorizer(
        preprocessor=lr_preprocessor,
        tokenizer=lr_tokenizer,
        stop_words=lr_stop_words,
        ngram_range=(1, 3),
        decode_error='replace',
        max_features=33000,
        max_df=0.75,
        min_df=5
    )

    X = pd.DataFrame(np.array(vectorizer.fit_transform(tweet_df.tweet).toarray()), dtype=np.float32)
    y = tweet_df['class'].astype(int)
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=33, test_size=0.1)

    # over sampling
    X = pd.concat([X_train, y_train], axis=1)
    not_hate = X[X['class']==0]
    hate = X[X['class']==1]
    hate = resample(hate, replace=True, n_samples=len(hate)*3, random_state=33)
    X = pd.concat([not_hate, hate])
    X_train = X.drop('class', axis=1)
    y_train = X['class']

    # best parameters found: C: 0.01, l2, saga
    c = [0.01]
    parameters = [{"penalty": ["l2"], "C": c, "solver": ["saga"]}]
    classifier = GridSearchCV(LogisticRegression(class_weight="balanced", max_iter=1000),
                              parameters,
                              scoring="f1_micro",
                              cv=StratifiedKFold(n_splits=5, random_state=33).split(X_train, y_train),
                              verbose=2)
    classifier.fit(X_train, y_train)

    print("best paramaters: ")
    print(classifier.best_params_)
    print("best estimator: ")
    print(classifier.best_estimator_)
    print("best score: ")
    print(classifier.best_score_)

    lr = classifier.best_estimator_

    lr.fit(X_train, y_train)
    # y_pred = lr.predict(X_test)
    # results = classification_report(y_test, y_pred)
    # print(results)
    # print()

    # conf_matrix = confusion_matrix(y_test, y_pred)
    # print(conf_matrix)

    return lr, vectorizer

if __name__ == '__main__':
    # train_model('data/private/td_zw_labeled_data.csv')
    model, vectorizer = train_model('data/private/td_zw_labeled_data.csv')

    # pickle.dump(model, open("models/logreg.pickle", "wb"))
    # pickle.dump(vectorizer, open("models/logreg_vectorizer.pickle", "wb"))

    test_pred_tweet = ["Listen up women, there is a reason why god made men stronger than you. "
                       "It's because rape is a force of good, that should be celebrated. IF a man "
                       "wants to fill you with his big cock, you should learn to : relax, orgasm and"
                       " thank him!!! Heil Hitler #bussdown",
                       "To all the little girls watching...never doubt that you are valuable and powerful "
                       "& deserving of every chance & opportunity in the world."]
    for t in test_pred_tweet:
        t = lr_preprocessor(t)
        print(t)
    print(test_pred_tweet)
    test_pred = vectorizer.transform(test_pred_tweet)

    print("test = %s" % test_pred)
    pred_result = model.predict(test_pred)
    print("model.predict = %s" % pred_result)

    pred_result = model.predict_proba(test_pred)
    print("model.predict_proba = %s" % pred_result)