"""
Deploy semi-supervised PU machine learning models.

This module provides classes for training, testing, and deploying a PU
learning model for predicting material synthesizability. Utility functions
for plotting aid in visualizing and analyzing results.

References:
    [1] DOI: 10.1021/acsnano.8b08014
    [2] DOI: 10.1145/1401890.1401920
    [3] DOI: 10.1016/j.patrec.2013.06.010
"""

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import precision_recall_fscore_support
from sklearn.cluster import KMeans
from sklearn.mixture import GaussianMixture, BayesianGaussianMixture
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import precision_recall_curve
from sklearn.model_selection import RepeatedKFold
from sklearn.utils import resample

from mpl_toolkits.mplot3d import Axes3D

from monty.serialization import dumpfn

import pandas as pd
import seaborn as sns
import os
import pickle

import numpy as np
import matplotlib.pyplot as plt
from pylab import rcParams

__author__ = "Nathan C. Frey, Jin Wang"
__copyright__ = "MIT License"
__version__ = "0.0.1"
__maintainer__ = "Nathan C. Frey"
__email__ = "n.frey@seas.upenn.edu"
__status__ = "Development"
__date__ = "Aug 2017"


class PULearner:
    def __init__(self):
        """A machine learning model that predicts material synthesizability.

        Positive samples are experimentally synthesized materials. Unlabeled
        samples are not-yet synthesized materials.

        Features for training data might be generated by first-principles 
        (density functional theory) calculations, or structural or chemical
        data looked up from a table.

        Hyperparameters are initialized with sensible defaults, but any newly
        trained model should have hyperparams carefully converged.

        Attributes:
            pu_stats (dict): Outputs of cv_baggingDT
            df_U (DataFrame): Unlabeled data.
            df_P (DataFrame): Positive data.

            synth_scores (list): Synthesizability scores (between 0 and 1) of
                unlabeled samples.
            labels (list): Likely synthesizable (1) or not (0)
            feat_importances (DataFrame): Feature importances from trained
                decision tree classifiers. Index corresponds to feature index
                in original data. 

        """

        self.pu_stats = None
        self.df_U = None
        self.df_P = None
        self.synth_scores = None
        self.labels = None
        self.feat_importances = None

    def cv_baggingDT(self, pu_data, splits=10, repeats=10, bags=100, filename=""):
        """
        Train bagged decision tree base classifiers and do repeated 
        k-fold CV.

        Synthesizability scores (0 = not synthesizable, 1 = already
        synthesized) are generated for an unlabeled sample by averaging
        the scores from the ensemble of decision tree classifiers that
        have not been trained on that sample. 

        Args:
            pu_data (json): A file where each row describes a material.
                There MUST be a column called "PU_label" where a 1 value
                indicates a synthesized (positive) compound and a 0 value
                indicates an unlabeled compound.

            splits (int): Number of splits in k-fold CV.
            repeats (int): Number of repeated k-fold CV.
            bags (int): Number of bags in bootstrap aggregation.
            filename (string): Save model training results to file with
                filename ending in .json or .pkl.

        Returns:
            pu_stats (dict): Metrics and outputs of PU learning model
                training.

        """
        
        print("Start PU Learning.")

        # Preprocess data and set attributes
        df = pd.read_json(pu_data)
        df_P, df_U, X_P, X_U = self._process_pu_data(df)
        self.df_P = df_P
        self.df_U = df_U

        # Split data into training and test splits for k-fold CV
        kfold = RepeatedKFold(n_splits=splits, n_repeats=repeats, random_state=42)

        # Scores for PU learning (tpr = True Positive Rate)
        scores = []
        tprs = []

        # Predicted synthesis probability of CVed P and U sets
        prob_P = np.ones(shape=(X_P.shape[0], splits * repeats))
        prob_U = -np.ones(shape=(X_U.shape[0], splits * repeats))

        # Feature importance
        feat_rank = np.zeros(shape=(X_P.shape[1], splits * repeats))

        idsp = 0  # index of repeated k splits

        # Loop over P and U training/test samples
        for (ptrain, ptest), (utrain, utest) in zip(kfold.split(X_P), kfold.split(X_U)):

            # Number of P and U training samples
            N_ptrain = X_P[ptrain].shape[0]
            N_utrain = X_U[utrain].shape[0]

            d = X_P.shape[1]
            K = N_ptrain
            train_label = np.zeros(shape=(N_ptrain + K,))
            train_label[:N_ptrain] = 1.0  # Synthesized (positive)

            # Out of bag samples
            n_oob = np.zeros(shape=(N_utrain,))
            f_oob = np.zeros(shape=(N_utrain, 2))

            # Sums of probabilities of test sets
            f_ptest = np.zeros(shape=(X_P[ptest].shape[0], 2))
            f_utest = np.zeros(shape=(X_U[utest].shape[0], 2))

            # Bootstrap resampling for each bag
            for i in range(bags):
                bootstrap_sample = np.random.choice(
                    np.arange(N_utrain), replace=True, size=K
                )

                # Positive samples and bootstrapped unlabeled samples
                data_bootstrap = np.concatenate(
                    (X_P[ptrain], X_U[bootstrap_sample, :]), axis=0
                )

                # Train decision tree classifier
                model = DecisionTreeClassifier(
                    max_depth=None,
                    max_features=None,
                    criterion="gini",
                    class_weight="balanced",
                )

                model.fit(data_bootstrap, train_label)

                # Index for the oob samples
                idx_oob = sorted(
                    set(range(N_utrain)) - set(np.unique(bootstrap_sample))
                )

                # Transductive learning on oob samples
                f_oob[idx_oob] += model.predict_proba(X_U[utrain][idx_oob])
                n_oob[idx_oob] += 1
                f_ptest += model.predict_proba(X_P[ptest])
                f_utest += model.predict_proba(X_U[utest])
                feat_rank[:, idsp] = model.feature_importances_

            # Predicted synthesis probabilities of unlabeled samples
            predict_utrain = f_oob[:, 1] / n_oob

            # Predicted probabilities for P and U test sets
            predict_ptest = f_ptest[:, 1] / bags
            predict_utest = f_utest[:, 1] / bags

            # Find predicted positives
            true_pos = predict_ptest[np.where(predict_ptest > 0.5)].shape[0]
            u_pos = predict_utest[np.where(predict_utest > 0.5)].shape[0]

            N_ptest = X_P[ptest].shape[0]
            N_utest = X_U[utest].shape[0]

            # Predicted positive ratio in test set
            p_pred_pos = (true_pos + u_pos) / (N_ptest + N_utest) + 0.0001

            # Compute PU recall (TPR) and score metrics
            recall = true_pos / N_ptest
            score = recall ** 2 / p_pred_pos
            scores.append(score)
            tprs.append(recall)

            # Predicted probabilities
            prob_P[ptest, idsp] = predict_ptest
            prob_U[utrain, idsp] = predict_utrain
            prob_U[utest, idsp] = predict_utest
            idsp += 1

            # Progress update
            if (idsp + 1) % splits == 0:
                tpr_tmp = np.asarray(tprs[-splits - 1 : -1])
                print(
                    "Performed Repeated "
                    + str(splits)
                    + "-fold: "
                    + str(idsp // splits + 1)
                    + " out of "
                    + str(repeats)
                )
                print(
                    "True Positive Rate: %0.2f (+/- %0.2f)"
                    % (tpr_tmp.mean(), tpr_tmp.std() * 2)
                )

        # Predicted labels from k-fold CV
        label_U = np.zeros(shape=(X_U.shape[0], splits * repeats + 1), dtype=int)
        label_U[:, : splits * repeats][np.where(prob_U > 0.5)] = 1
        label_U[:, splits * repeats] = np.sum(
            label_U[:, : splits * repeats + 1], axis=1
        )

        tprs = np.asarray(tprs)
        scores = np.asarray(scores)

        # Metrics for each model in the k-folds
        label_U_rp = np.zeros(shape=(X_U.shape[0], repeats), dtype=int)
        prob_U_rp = np.zeros(shape=(X_U.shape[0], repeats))
        feat_rank_rp = np.zeros(shape=(X_U.shape[1], repeats))
        tpr_rp = np.zeros(shape=(repeats,))
        scores_rp = np.zeros(shape=(repeats,))
        labels = np.zeros(shape=(X_U.shape[0],))

        for i in range(repeats):
            prob_U_rp[:, i] = prob_U[:, i * splits : (i + 1) * splits].mean(axis=1)
            feat_rank_rp[:, i] = feat_rank[:, i * splits : (i + 1) * splits].mean(
                axis=1
            )
            tpr_rp[i] = tprs[i * splits : (i + 1) * splits].mean()
            scores_rp[i] = scores[i * splits : (i + 1) * splits].mean()

        label_U_rp[np.where(prob_U_rp > 0.5)] = 1
        prob = prob_U_rp.mean(axis=1)
        labels[np.where(prob > 0.5)] = 1

        # Get confidence interval of TPR for each kfold
        tpr_low, tpr_up = self.bootstrapCI(tpr_rp)
        scores_low, scores_up = self.bootstrapCI(scores_rp)

        # PU learning metrics
        metrics = np.asarray(
            [tpr_rp.mean(), tpr_low, tpr_up, scores_rp.mean(), scores_low, scores_up]
        )

        print("Accuracy: %0.2f" % (tpr_rp.mean()))
        print("95%% confidence interval: [%0.2f, %0.2f]" % (tpr_low, tpr_up))

        # Metrics and results from training / testing
        pu_stats = {
            "prob": prob,
            "labels": labels,
            "metrics": metrics,
            "prob_rp": prob_U_rp,
            "label_rp": label_U_rp,
            "tpr_rp": tpr_rp,
            "scores_rp": scores_rp,
            "feat_rank_rp": feat_rank_rp,
        }

        # Save results
        if filename:
            if filename.endswith(".json"):
                dumpfn(pu_stats, filename)
            if filename.endswith(".pkl"):
                with open(filename, "wb") as file:
                    pickle.dump(pu_stats, file, protocol=pickle.HIGHEST_PROTOCOL)

        self.pu_stats = pu_stats
        return pu_stats

    def bootstrapCI(self, data, ci=95, ns=10000):
        """Compute confidence interval of the TPR.

        Args:
            data (array): Array of TPRs for each kfold.
            ci (int): Confidence interval.
            ns (int): Number of bootstrap resamplings.

        Returns:
            lower (float): Lower endpoint of CI.
            upper (float): Upper endpoint of CI.
            
        """

        bs_rsample = []
        for _ in range(ns):
            rsample = resample(data, n_samples=len(data))
            bs_rsample.append(np.mean(rsample))

        bs_rsample = np.asarray(bs_rsample)
        lower = np.percentile(bs_rsample, (100 - ci) / 2)
        upper = np.percentile(bs_rsample, ci + (100 - ci) / 2)

        return lower, upper

    def corr_heatmap(self, num_feats=10, fname=""):
        """Plot correlation matrix between synthesizability and features.

        cv_baggingDT must be run first.

        Args:
            num_feats (int): How many features to consider.
            fname (str): Filename if correlation plot should be saved.

        Returns:
            None (generates plots)

        """

        pu_stats = self.pu_stats
        df_U = self.df_U
        df_U_copy = df_U.drop(columns=["PU_label"])

        # Get normalized, sorted & ranked list of most important features
        synth_scores = pu_stats["prob"]
        df_U_copy["synth_score"] = synth_scores

        # Make correlation matrix of top "num_feats" features
        corrmat = df_U_copy.corr()
        cols = corrmat.nlargest(num_feats, "synth_score")["synth_score"].index
        cm = np.corrcoef(df_U_copy[cols].values.T)

        sns.set(style='ticks')
        rcParams['figure.dpi'] = 300

        fig, ax = plt.subplots(1, 1)
        hm = sns.heatmap(
            cm,
            ax=ax,
            cbar=True,
            annot=True,
            square=True,
            fmt=".2f",
            annot_kws={"size": 7},
            yticklabels=cols.values,
            xticklabels=cols.values,
        )

        if fname:
            self.save_plot(fname + ".png", fig, ax)

    def get_feat_importances(self, plot_format=""):
        """Process output from PU learning k-fold cross validation.

        cv_baggingDT must be run first.

        If plot_format is specified, a feature importance plot will
        be saved.

        Args:
            plot_format (str): svg, png, or pdf file format for saving simple
                visualizations of feature importance and correlation. 

        """

        pu_stats = self.pu_stats

        # Feature importances for individual repetitions of kfold CV
        feat_rank_rp = pu_stats["feat_rank_rp"]
        feat_importances = np.sum(feat_rank_rp, axis=1)

        df_U = self.df_U
        df_U = df_U._get_numeric_data()
        df_U_copy = df_U.drop(columns=["PU_label"])
        feat_names = df_U_copy.columns

        # Index corresponds to feature in original data
        df_feat = pd.DataFrame(columns=["feature", "importance"])
        df_feat["feature"] = feat_names
        df_feat["importance"] = feat_importances

        # Sort by importance
        df_feat_sort = df_feat.sort_values(by="importance", ascending=False)
        max_value = df_feat["importance"].max()

        # Normalize to 1
        df_feat_sort["importance"] = df_feat_sort["importance"] / max_value

        # Set feature importance attribute
        self.feat_importances = df_feat

        if plot_format in ["svg", "pdf", "png"]:

            # Feature importance plot
            fig, ax = plt.subplots(figsize=(10, 4))
            with sns.axes_style(style="ticks"):
                sns.barplot(x="feature", y="importance", data=df_feat_sort)
            ax.set_xticklabels(
                ax.get_xticklabels(), rotation=45, ha="right", fontsize=7
            )
            filename = "feat_importance." + plot_format
            self.save_plot(filename, fig, ax)

    @staticmethod
    def _process_pu_data(data):
        """Utility method for processing input data.

        Args:
            data (DataFrame): Data with positive and unlabeled samples.

        Returns:
            X_P (array): Positive sample set.
            X_U (array): Unlabeled sample set.

        """

        df_P = data.query("PU_label == 1")  # Positive value is 1
        df_U = data.query("PU_label == 0")  # Unlabeled value is 0

        # Chop off PU label and drop non-numeric columns for sklearn
        X_P = np.asarray(df_P.drop(columns=["PU_label"])._get_numeric_data())
        X_U = np.asarray(df_U.drop(columns=["PU_label"])._get_numeric_data())

        return df_P, df_U, X_P, X_U

    @staticmethod
    def save_plot(filename, fig, ax):
        """Utility method for saving simple visualizations.

        Args:
            filename (str): Name ending in .svg, .png, or .pdf
            fig, ax (objects): Matplotlib objects.

        Returns:
            None

        """

        sns.set_style("ticks")
        fig.tight_layout()
        fig.savefig(filename)


class PUInteract:
    def __init__(self, df_parent, pu_parent, df_child, pu_child, merge_on=(), feats=()):
        """Consider parent and child phase PU learning scores.

        This class looks at PU learning scores for parent bulk
        compounds (e.g. layered h-BN) and scores of the child phases
        along with descriptors like exfoliation energy and changes
        in structural/electronic properties to predict (parent, child)
        pairs that can be synthesized.

        Parent and child must be linked by a column that allows the
        dataframes to be merged. There should also be additional features
        that characterize the structural and chemical differences between
        parents and children, e.g. changes in bond lengths, etc.

        Unsupervised clustering models are used to identify synthesizable 
        (parent/child) pairs.

        Args:
            df_parent (str): Parent data filename.
            pu_parent (dict): Output from PULearner.cv_baggingDT.
            df_child (str): Child data filename.
            pu_child (dict): Output from PULearner.cv_baggingDT.
            merge_on (tuple): Column name(s) on which to merge.
            feats (tuple): Column names to use as features. If empty, use all
                possible columns. 

        Attributes:
            merged_df (DataFrame): (Parent, child) pair data.
            X (array): Array representation of merged_df.

        Returns:
            None

        """

        df_parent = pd.read_json(df_parent)
        df_child = pd.read_json(df_child)

        # Set scores from PULearner
        df_parent["synth_score"] = 1
        df_child["synth_score"] = 1

        df_parent.loc[df_parent.eval("PU_label == 0"), "synth_score"] = pu_parent[
            "prob"
        ]
        df_child.loc[df_child.eval("PU_label == 0"), "synth_score"] = pu_child["prob"]

        # Merge parent and child dfs
        merge_on = list(merge_on)
        df = pd.merge(
            df_parent, df_child, on=merge_on, how="outer", suffixes=["_p", "_c"]
        )
        df.drop(columns=["PU_label_p", "PU_label_c"], inplace=True, axis=1)

        if feats:
            feat_names = [f + "_p" for f in feats] + [f + "_c" for f in feats]
            df = df[feat_names]

        self.merged_df = df
        self.X = np.array(df)

    def do_kmeans(self, n_clusters=2, seed=42):
        """Do k-means clustering on (parent, child) pairs.

        Args:
            n_clusters (int): Number of clusters.
            seed (int): Fix random seed for kmeans reproducibility.

        Returns:
            kmeans_output (dict): kmeans cluster centers, cluster labels for
                each (parent, child)

        """

        np.random.seed(seed)
        km = KMeans(n_clusters=n_clusters, random_state=seed)

        km.fit(self.X)
        kmeans_output = {
            "cluster_centers": km.cluster_centers_,
            "cluster_labels": km.labels_,
        }

        return kmeans_output

    def do_gmixture(self, n_components=2, seed=42):
        """
        Estimate parameters of a Gaussian mixture distribution of (parent,
            child) data.

        Args:
            n_components (int): Number of components in GMM.
            seed (int): Random seed.

        Returns:
            gmm_output (dict): Predicted labels of (parent, child) pairs and
                predicted posterior probabilities of each component.

        """

        np.random.seed(seed)
        gmm = GaussianMixture(
            n_components=n_components, random_state=seed, covariance_type="full"
        )

        gmm.fit(self.X)
        gmm_labels = gmm.predict(self.X)
        gmm_prob = gmm.predict_proba(self.X)[:, 0]
        gmm_output = {"gmm_labels": gmm_labels, "gmm_prob": gmm_prob}

        return gmm_output

    def do_bgm(self, n_components=6, seed=42):
        """Bayesian Gaussian Mixture.

        Infer the effective number of components in a Gaussian Mixture Model
            via variational Bayesian estimation.

        n_effective_componenents < n_components if the model sets some
            weights close to 0.

        Args:
            n_components (int): Number of components in GMM.
            seed (int): Random seed.

        Returns:
            bgm_output (dict): Labels and probabilities.

        """

        np.random.seed(seed)
        bgm = BayesianGaussianMixture(
            n_components=n_components,
            covariance_type="full",
            weight_concentration_prior=1e-2,
            weight_concentration_prior_type="dirichlet_process",
            mean_precision_prior=1e-2,
            init_params="random",
            max_iter=100,
            random_state=seed,
        )

        bgm.fit(self.X)
        bgm_labels = bgm.predict(self.X)
        bgm_prob = bgm.predict_proba(self.X)[:, 0]

        bgm_output = {"bgm_labels": bgm_labels, "bgm_prob": bgm_prob}

        return bgm_output
