from collections import Callable, OrderedDict
from glob import glob
from pathlib import Path
from types import SimpleNamespace

import os
import streamlit as st
import torch
from transformers import pipeline, set_seed, TextGenerationPipeline, \
    AutoModelForCausalLM, AutoTokenizer

from app.globals import MODEL_SOURCES, MERCURY_MODELS, MERCURY_PATHS, HUGGINGFACE_MODELS


class TextGenerator:
    """
    TextGenerator encapsulates a pre-trained language model used for generating text.
    """

    def __init__(
            self,
            model_source: str,
            model_name: str,
            checkpoint: str,
            checkpoint_path: Path,
            seed: int = 42,
            device: str = None,
    ):

        # store the parameters
        self.model_source = model_source
        self.model_name = model_name
        self.checkpoint = checkpoint
        self.checkpoint_path = checkpoint_path

        # set the seed for generation
        self._seed = seed
        set_seed(seed)

        # set the device
        self._device = device
        if device is None:
            self._device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # load the language model
        self.generator = self.load_generator()

    def get_checkpoint_info(self):
        return {
            'model_source': self.model_source,
            'model_name': self.model_name,
            'checkpoint': self.checkpoint,
            'checkpoint_path': self.checkpoint_path,
        }

    def get_checkpoint_info_string(self):
        if self.model_source == "Mercury":
            return f'Source: {self.model_source}\n' \
                   f'Model: {self.model_name}\n' \
                   f'Checkpoint: {self.checkpoint}'
        elif self.model_source == 'Huggingface':
            return f'Source: {self.model_source}\n' \
                   f'Model: {self.model_name}'
        else:
            raise NotImplementedError(
                f"Model source {self.model_source} not recognized."
            )

    def _load_mercury(self, checkpoint_path: Path):
        """
        Load a language model and its tokenizer from a HF checkpoint.

        :param checkpoint_path: a single checkpoint directory (e.g.
        /path/to/checkpoint-1000)
        """
        assert (
            checkpoint_path.is_dir()
        ), f"`checkpoint_path` must be a directory but {checkpoint_path} " \
           f"is not a directory."
        assert (
                checkpoint_path / "config.json"
        ).exists(), "Checkpoint directory must contain a `config.json` file."
        assert (
                checkpoint_path / "tokenizer_config.json"
        ).exists(), "Checkpoint directory must contain a `tokenizer_config.json` file."

        # Load things
        model = AutoModelForCausalLM.from_pretrained(f"{checkpoint_path}"). \
            eval().to(self._device)
        tokenizer = AutoTokenizer.from_pretrained(f"{checkpoint_path}")

        return pipeline("text-generation", model=model, tokenizer=tokenizer,
                        device={'cpu': -1, 'cuda': 0}[self._device])

    def load_generator(self) -> TextGenerationPipeline:
        """
        Load the language model.
        """

        if self.model_source == 'Mercury':
            return self._load_mercury(self.checkpoint_path)
        elif self.model_source == 'Huggingface':
            # model: gpt2, gpt2-large
            return pipeline(
                'text-generation',
                model=HUGGINGFACE_MODELS[self.model_name],
                device={'cpu': -1, 'cuda': 0}[self._device],
            )
        else:
            raise NotImplementedError(
                f"Model source {self.model_source} not recognized."
            )

    def generate_text(
            self,
            starting_text: str,
            max_length: int = 100,
            num_return_sequences: int = 1,
            temperature: float = 1.0,
            top_p: float = 1.0
    ) -> str:
        """
        Generate text using the language model.
        """
        return self.generator(
            starting_text,
            max_length=max_length,
            num_return_sequences=num_return_sequences,
            temperature=temperature,
            top_p=top_p
        )[0]['generated_text']


@st.cache(allow_output_mutation=True)
def instantiate_generator(
        model_source: str,
        model_name: str,
        checkpoint: str,
        checkpoint_path: Path,
        seed: int = 42,
        device: str = None,
):
    """
    Create a generator.
    """
    return TextGenerator(
        model_source=model_source,
        model_name=model_name,
        checkpoint=checkpoint,
        checkpoint_path=checkpoint_path,
        seed=seed,
        device=device,
    )


def get_checkpoints(
        path: Path,
        is_checkpoint_fn: Callable = None,
        step_extraction_fn: Callable = None,
) -> OrderedDict:
    """Given the path to a directory that contains one or more checkpoints for
    a single training run, return all checkpoints in the directory.
    :param path: directory where checkpoints are stored
    :param is_checkpoint_fn: given a path, checks if it is a valid checkpoint directory
    :param step_extraction_fn: given a valid checkpoint directory, returns the step
    at which the checkpoint was stored
    """

    # Get the list of all files and folders in the directory
    print(path)
    candidate_paths = [Path(p) for p in glob(str(path / "*"))]
    print(candidate_paths)
    # Filter out only the checkpoints
    if not is_checkpoint_fn:
        is_checkpoint_fn = lambda p: p.resolve().name.startswith("checkpoint-")
    checkpoint_paths = [p for p in candidate_paths if is_checkpoint_fn(p)]

    # Function to extract the step that a checkpoint corresponds to
    if not step_extraction_fn:
        step_extraction_fn = lambda p: int(p.resolve().name.split("-")[-1])

    # Construct an ordered dictionary mapping {step: checkpoint}, sorted by step
    return OrderedDict(
        sorted(
            [(step_extraction_fn(p), p) for p in checkpoint_paths], key=lambda k: k[0]
        )
    )


def checkpoint_widget():
    """Widget for selecting a model and checkpoint."""
    # Check if mercury paths exists before offering as option
    for model, path in MERCURY_PATHS.items():
        if not os.path.exists(path):
            if 'Mercury' in MODEL_SOURCES:
                MODEL_SOURCES.remove('Mercury')

    # Select the model source
    model_source = st.sidebar.radio("Model Source", options=MODEL_SOURCES)

    if model_source == 'Mercury':
        # Hardcode the models that are supported by Mercury
        model_name = st.sidebar.selectbox(
            "Model", options=list(MERCURY_MODELS.keys())
        )

        # Select from the list of available model checkpoints
        checkpoints_available = get_checkpoints(
            path=Path(MERCURY_PATHS[MERCURY_MODELS[model_name]]),
        )
        checkpoint = st.sidebar.selectbox(
            "Checkpoint Step", options=list(checkpoints_available.keys())
        )
        checkpoint_path = checkpoints_available[checkpoint]

    elif model_source == 'Huggingface':
        # Hardcode the models that are supported by Huggingface
        model_name = st.sidebar.selectbox(
            "Model", options=list(HUGGINGFACE_MODELS.keys())
        )

        # No checkpoints available for Huggingface models
        checkpoint = None
        checkpoint_path = None
    else:
        raise NotImplementedError("Model source not recognized.")

    return SimpleNamespace(
        model_source=model_source,
        model_name=model_name,
        checkpoint=checkpoint,
        checkpoint_path=checkpoint_path,
    )
