import logging
import shutil
import h5py
import numpy

from copy import deepcopy
from typing import Dict, Any

from Bio import SeqIO
from pandas import read_csv, DataFrame
from tqdm import tqdm

from bio_embeddings.embed import name_to_embedder, EmbedderInterface
from bio_embeddings.utilities import (
    FileManagerInterface,
    InvalidParameterError,
    check_required,
    get_file_manager,
    get_model_directories_from_zip,
    get_model_file,
)
from bio_embeddings.utilities.backports import nullcontext

logger = logging.getLogger(__name__)


def _print_expected_file_sizes(
    embedder: EmbedderInterface, mapping_file: DataFrame, result_kwargs: Dict[str, Any]
) -> None:
    """
    Logs the lower bound size of embeddings_file and reduced_embedding_file

    :param embedder: the embedder being used
    :param mapping_file: the mapping file of the sequences
    :param result_kwargs: the kwargs passed to the pipeline --> will decide what to print

    :return: Nothing.
    """
    per_amino_acid_size_in_bytes = 4 * embedder.embedding_dimension * embedder.number_of_layers
    per_protein_size_in_bytes = 4 * embedder.embedding_dimension

    total_number_of_proteins = len(mapping_file)
    total_aa = mapping_file['sequence_length'].sum()

    embeddings_file_size_in_MB = per_amino_acid_size_in_bytes * total_aa * pow(10, -6)
    reduced_embeddings_file_size_in_MB = per_protein_size_in_bytes * total_number_of_proteins * pow(10, -6)

    required_space_in_MB = 0

    if result_kwargs.get("reduce") is True:
        logger.info(f"The minimum expected size for the reduced_embedding_file is "
                    f"{reduced_embeddings_file_size_in_MB:.3f}MB.")

        required_space_in_MB += reduced_embeddings_file_size_in_MB

    # TODO: calculate size of transformed embeddings?
    if not (result_kwargs.get("reduce") is True and result_kwargs.get("discard_per_amino_acid_embeddings") is True):
        logger.info(f"The minimum expected size for the embedding_file is {embeddings_file_size_in_MB:.3f}MB.")

        required_space_in_MB += embeddings_file_size_in_MB

    _, _, available_space_in_bytes = shutil.disk_usage(result_kwargs.get('prefix'))

    available_space_in_MB = available_space_in_bytes * pow(10, -6)

    if available_space_in_MB < required_space_in_MB:
        logger.warning(f"You are attempting to generate {required_space_in_MB:.3f}MB worth of embeddings, "
                       f"but only {available_space_in_MB:.3f}MB are available at "
                       f"the prefix({result_kwargs.get('prefix')}). \n"
                       f"We suggest you stop execution NOW and double check you have enough free space available. "
                       f"Alternatively, try reducing the input FASTA file.")
    else:
        logger.info(f"You are going to generate a total of {required_space_in_MB:.3f}MB of embeddings, and have "
                    f"{available_space_in_MB:.3f}MB available at {result_kwargs.get('prefix')}.")


def _get_reduced_embeddings_file_context(
    file_manager: FileManagerInterface, result_kwargs: Dict[str, Any]
):
    """
    :param file_manager: The FileManager derived class which will be used to create the file
    :param result_kwargs: A dictionary which will be updated in-place to include the path to the newly created file

    :return: a file context
    """

    # Create reduced embeddings file if set in params
    result_kwargs.setdefault("reduce", False)

    if result_kwargs["reduce"] is True:
        reduced_embeddings_file_path = file_manager.create_file(
            result_kwargs.get("prefix"),
            result_kwargs.get("stage_name"),
            "reduced_embeddings_file",
            extension=".h5",
        )
        result_kwargs["reduced_embeddings_file"] = reduced_embeddings_file_path
        return h5py.File(reduced_embeddings_file_path, "w")

    return nullcontext()


def _get_embeddings_file_context(
    file_manager: FileManagerInterface, result_kwargs: Dict[str, Any]
):
    """
    :param file_manager: The FileManager derived class which will be used to create the file
    :param result_kwargs: A dictionary which will be updated in-place to include the path to the newly created file

    :return: a file context
    """

    result_kwargs.setdefault("discard_per_amino_acid_embeddings", False)

    if result_kwargs["discard_per_amino_acid_embeddings"] is True:
        if result_kwargs.get("reduce", False) is False and result_kwargs.get("embeddings_transformer_function") is None:
            raise InvalidParameterError(
                "Cannot only have discard_per_amino_acid_embeddings: True. "
                "Either also set `reduce: True` or define an `embeddings_transformer_function`, or both."
            )
        return nullcontext()
    else:
        embeddings_file_path = file_manager.create_file(
            result_kwargs.get("prefix"),
            result_kwargs.get("stage_name"),
            "embeddings_file",
            extension=".h5",
        )
        result_kwargs["embeddings_file"] = embeddings_file_path
        return h5py.File(embeddings_file_path, "w")


def _get_transformed_embeddings_file_context(
    file_manager: FileManagerInterface, result_kwargs: Dict[str, Any]
):
    """
    :param file_manager: The FileManager derived class which will be used to create the file
    :param result_kwargs: A dictionary which will be updated in-place to include the path to the newly created file

    :return: a file context
    """

    result_kwargs.setdefault("embeddings_transformer_function", None)

    if result_kwargs["embeddings_transformer_function"] is not None:
        transformed_embeddings_file_path = file_manager.create_file(
            result_kwargs.get("prefix"),
            result_kwargs.get("stage_name"),
            "transformed_embeddings_file",
            extension=".h5",
        )
        result_kwargs["transformed_embeddings_file"] = transformed_embeddings_file_path
        return h5py.File(transformed_embeddings_file_path, "w")

    return nullcontext()


def _check_transform_embeddings_function(embedder: EmbedderInterface, result_kwargs: Dict[str, Any]):

    result_kwargs.setdefault("embeddings_transformer_function", None)

    if result_kwargs["embeddings_transformer_function"] is not None:
        try:
            transform_function = eval(result_kwargs["embeddings_transformer_function"], {}, {"np": numpy})
        except TypeError:
            raise InvalidParameterError(f"`embeddings_transformer_function` must be callable! \n"
                                        f"Instead is {result_kwargs['embeddings_transformer_function']}\n"
                                        f"Most likely you want a lambda function.")

        if not callable(transform_function):
            raise InvalidParameterError(f"`embeddings_transformer_function` must be callable! \n"
                                        f"Instead is {result_kwargs['embeddings_transformer_function']}\n"
                                        f"Most likely you want a lambda function.")

        template_embedding = embedder.embed("SEQVENCE")

        # Check that it works in principle
        try:
            transformed_template_embedding = transform_function(template_embedding)
        except:
            raise InvalidParameterError(f"`embeddings_transformer_function` must be valid callable! \n"
                                        f"Instead is {result_kwargs['embeddings_transformer_function']}\n"
                                        f"This function excepts when processing an embedding.")

        # Check that return can be cast to np.array
        try:
            numpy.array(transformed_template_embedding)
        except:
            raise InvalidParameterError(f"`embeddings_transformer_function` must be valid callable "
                                        f"returning numpy array compatible object! \n"
                                        f"Instead is {result_kwargs['embeddings_transformer_function']}\n"
                                        f"This function excepts when processing an embedding.")


def embed_and_write_batched(
    embedder: EmbedderInterface,
    file_manager: FileManagerInterface,
    result_kwargs: Dict[str, Any],
    half_precision: bool = False
) -> Dict[str, Any]:
    """ The shared code between the SeqVec, Albert, Bert and XLNet pipelines """
    # Lazy fasta file reader. The mapping file contains the corresponding ids in the same order
    sequences = (
        str(entry.seq)
        for entry in SeqIO.parse(result_kwargs["remapped_sequences_file"], "fasta")
    )
    # We want to read the unnamed column 0 as str (esp. with simple_remapping), which requires some workarounds
    # https://stackoverflow.com/a/29793294/3549270
    mapping_file = read_csv(result_kwargs["mapping_file"], index_col=0)
    mapping_file.index = mapping_file.index.astype('str')

    # Print the minimum required file sizes
    _print_expected_file_sizes(embedder, mapping_file, result_kwargs)

    # Get transformer function, if available
    transform_function = result_kwargs.get("embeddings_transformer_function", None)

    if transform_function:
        transform_function = eval(transform_function, {}, {"np": numpy})

    # Open embedding files or null contexts and iteratively save embeddings to file
    with _get_embeddings_file_context(
        file_manager, result_kwargs
    ) as embeddings_file, _get_reduced_embeddings_file_context(
        file_manager, result_kwargs
    ) as reduced_embeddings_file, _get_transformed_embeddings_file_context(
        file_manager, result_kwargs
    ) as transformed_embeddings_file:
        embedding_generator = embedder.embed_many(
            sequences, result_kwargs.get("max_amino_acids")
        )
        for sequence_id, original_id, embedding in zip(
            mapping_file.index,
            mapping_file["original_id"],
            tqdm(embedding_generator, total=len(mapping_file))
        ):
            # embedding: numpy.ndarray
            if half_precision:
                embedding = embedding.astype(numpy.float16)
            if result_kwargs.get("discard_per_amino_acid_embeddings") is False:
                dataset = embeddings_file.create_dataset(sequence_id, data=embedding)
                dataset.attrs["original_id"] = original_id
            if result_kwargs.get("reduce") is True:
                dataset = reduced_embeddings_file.create_dataset(
                    sequence_id, data=embedder.reduce_per_protein(embedding)
                )
                dataset.attrs["original_id"] = original_id
            if transform_function:
                dataset = transformed_embeddings_file.create_dataset(
                    sequence_id, data=numpy.array(transform_function(embedding))
                )
                dataset.attrs["original_id"] = original_id

    return result_kwargs


# Some of this might not be available when installed without the `all` extra
ALL_PROTOCOLS = [
    "bepler",
    "cpcprot",
    "esm",
    "plus_rnn",
    "prottrans_albert_bfd",
    "prottrans_bert_bfd",
    "prottrans_t5_bfd",
    "prottrans_xlnet_uniref100",
    "seqvec",
    "unirep",
]

# TODO: 10000 is a random guess
# There remainder was measured for a GTX 1080 with 8GB memory
DEFAULT_MAX_AMINO_ACIDS = {
    "bepler": 10000,
    "cpcprot": 10000,
    "esm": 10000,
    "plus_rnn": 10000,
    "prottrans_albert_bfd": 3035,
    "prottrans_bert_bfd": 6024,
    "prottrans_t5_bfd": 1000,
    "prottrans_xlnet_uniref100": 4000,
    "seqvec": 15000,
    "unirep": 10000,
}


def run(**kwargs):
    """
    Run embedding protocol

    Parameters
    ----------
    kwargs arguments (* denotes optional):
        sequences_file: Where sequences live
        prefix: Output prefix for all generated files
        protocol: Which embedder to use
        mapping_file: the mapping file generated by the pipeline when remapping indexes
        stage_name: The stage name

    Returns
    -------
    Dictionary with results of stage
    """
    check_required(
        kwargs,
        ["protocol", "prefix", "stage_name", "remapped_sequences_file", "mapping_file"],
    )

    if kwargs["protocol"] not in name_to_embedder:
        if kwargs["protocol"] in ALL_PROTOCOLS:
            raise InvalidParameterError(
                f"The extra for the protocol {kwargs['protocol']} is missing. "
                "See https://docs.bioembeddings.com/#installation on how to install all extras"
            )
        raise InvalidParameterError(
            "Invalid protocol selection: {}. Valid protocols are: {}".format(
                kwargs["protocol"], ", ".join(name_to_embedder.keys())
            )
        )

    embedder_class = name_to_embedder[kwargs["protocol"]]

    if kwargs["protocol"] == "unirep" and kwargs.get("use_cpu") is not None:
        raise InvalidParameterError("UniRep does not support configuring `use_cpu`")

    result_kwargs = deepcopy(kwargs)

    # Download necessary files if needed
    # noinspection PyProtectedMember
    for file in embedder_class._necessary_files:
        if not result_kwargs.get(file):
            result_kwargs[file] = get_model_file(model=embedder_class.name, file=file)

    # noinspection PyProtectedMember
    for directory in embedder_class._necessary_directories:
        if not result_kwargs.get(directory):
            result_kwargs[directory] = get_model_directories_from_zip(
                model=embedder_class.name, directory=directory
            )

    result_kwargs.setdefault("max_amino_acids", DEFAULT_MAX_AMINO_ACIDS[kwargs["protocol"]])

    file_manager = get_file_manager(**kwargs)
    embedder: EmbedderInterface = embedder_class(**result_kwargs)
    _check_transform_embeddings_function(embedder, result_kwargs)

    return embed_and_write_batched(embedder, file_manager, result_kwargs, kwargs.get("half_precision", False))
