import logging
import pickle
from typing import Optional

import numpy as np
import scanpy as sc
from scvi.model import SCVI
from sklearn.neighbors import KNeighborsClassifier


class SCVI_POPV:
    def __init__(
        self,
        batch_key: Optional[str] = "_batch_annotation",
        labels_key: Optional[str] = "_labels_annotation",
        max_epochs: Optional[int] = None,
        save_folder: Optional[str] = None,
        result_key: Optional[str] = "popv_knn_on_scvi_prediction",
        embedding_key: Optional[str] = "X_scvi_umap_popv",
        model_kwargs: Optional[dict] = {},
        classifier_dict: Optional[dict] = {},
        embedding_dict: Optional[dict] = {},
    ) -> None:
        """
        Class to compute KNN classifier after scVI integration.

        Parameters
        ----------
        batch_key
            Key in obs field of adata for batch information.
        labels_key
            Key in obs field of adata for cell-type information.
        max_epochs
            Number of epochs scvi is trained.
        result_key
            Key in obs in which celltype annotation results are stored.
        embedding_key
            Key in obsm in which UMAP embedding of integrated data is stored.
        model_kwargs
            Dictionary to supply non-default values for SCVI model. Options at scvi.model.SCVI
        classifier_dict
            Dictionary to supply non-default values for KNN classifier. Options at sklearn.neighbors.KNeighborsClassifier
        embedding_dict
            Dictionary to supply non-default values for UMAP embedding. Options at sc.tl.umap
        """
        self.batch_key = batch_key
        self.labels_key = labels_key
        self.result_key = result_key
        self.embedding_key = embedding_key

        self.max_epochs = max_epochs
        self.save_folder = save_folder

        self.model_kwargs = {
            "dropout_rate": 0.05,
            "dispersion": "gene",
            "n_layers": 3,
            "n_latent": 20,
            "gene_likelihood": "nb",
            "use_batch_norm": "none",
            "use_layer_norm": "both",
            "encode_covariates": True,
        }

        self.model_kwargs.update(model_kwargs)

        self.classifier_dict = {"weights": "uniform", "n_neighbors": 15}
        self.classifier_dict.update(classifier_dict)

        self.embedding_dict = {"min_dist": 0.3}
        self.embedding_dict.update(embedding_dict)

    def compute_integration(self, adata):
        logging.info("Integrating data with scvi")

        # Go through obs field with subsampling information and subsample label information.
        if "subsampled_labels" not in adata.obs.columns:
            adata.obs["subsampled_labels"] = [
                label if subsampled else adata.uns["unknown_celltype_label"]
                for label, subsampled in zip(
                    adata.obs["_labels_annotation"], adata.obs["_ref_subsample"]
                )
            ]
        adata.obs["subsampled_labels"] = adata.obs["subsampled_labels"].astype(
            "category"
        )

        if adata.uns["_pretrained_scvi_path"] is None:
            SCVI.setup_anndata(
                adata,
                batch_key=self.batch_key,
                labels_key="subsampled_labels",
                layer="scvi_counts",
            )
            model = SCVI(adata, **self.model_kwargs)
            logging.info("Training scvi offline.")
        else:
            query = adata[adata.obs["_dataset"] == "query"].copy()
            model = SCVI.load_query_data(query, adata.uns["_pretrained_scvi_path"])
            logging.info("Training scvi online.")

        if adata.uns["_prediction_mode"] == "fast":
            if self.max_epochs is None:
                self.max_epochs = 1
            model.train(
                max_epochs=self.max_epochs,
                train_size=0.9,
                use_gpu=adata.uns["_use_gpu"],
                plan_kwargs={"n_steps_kl_warmup": 1},
            )
        else:
            if self.max_epochs is None:
                self.max_epochs = min(round((20000 / adata.n_obs) * 200), 200)
            model.train(
                max_epochs=round(self.max_epochs),
                train_size=0.9,
                use_gpu=adata.uns["_use_gpu"],
                plan_kwargs={"n_epochs_kl_warmup": min(20, self.max_epochs)},
            )

            if (
                adata.uns["_save_path_trained_models"] is not None
                and adata.uns["_prediction_mode"] == "retrain"
            ):
                # Update scvi for scanvi.
                adata.uns["_pretrained_scvi_path"] = (
                    adata.uns["_save_path_trained_models"] + "/scvi"
                )
                model.save(
                    adata.uns["_save_path_trained_models"] + "/scvi",
                    save_anndata=False,
                    overwrite=True,
                )

        adata.obsm["X_scvi"] = model.get_latent_representation(adata)

    def predict(self, adata):
        logging.info(f'Saving knn on scvi results to adata.obs["{self.result_key}"]')

        if adata.uns["_prediction_mode"] == "retrain":
            ref_idx = adata.obs["_dataset"] == "ref"
            train_X = adata[ref_idx].obsm["X_scvi"]
            train_Y = adata[ref_idx].obs[self.labels_key].to_numpy()
            knn = KNeighborsClassifier(**self.classifier_dict)
            knn.fit(train_X, train_Y)
            if adata.uns["_save_path_trained_models"]:
                pickle.dump(
                    knn,
                    open(
                        adata.uns["_save_path_trained_models"]
                        + "scvi_knn_classifier.pkl",
                        "wb",
                    ),
                )
        else:
            knn = pickle.load(
                open(
                    adata.uns["_save_path_trained_models"] + "scvi_knn_classifier.pkl",
                    "rb",
                )
            )

        knn_pred = knn.predict(adata.obsm["X_scvi"])

        # save_results
        adata.obs[self.result_key] = knn_pred
        if adata.uns["_return_probabilities"]:
            adata.obs[self.result_key + "_probabilities"] = np.max(
                knn.predict_proba(adata.obsm["X_scvi"]), axis=1
            )

    def compute_embedding(self, adata):
        if adata.uns["_compute_embedding"]:
            logging.info(
                f'Saving UMAP of scvi results to adata.obs["{self.embedding_key}"]'
            )
            sc.pp.neighbors(adata, use_rep="X_scvi")
            adata.obsm[self.embedding_key] = sc.tl.umap(
                adata, copy=True, **self.embedding_dict
            ).obsm["X_umap"]
