import os
import sys
import pandas
import json
from pathlib import Path
from typing import List
from nltk import word_tokenize
from snlp.mwes.am import calculate_am
from snlp.mwes.mwe_utils import replace_mwes, get_counts
from snlp import logger


class MWE(object):
    def __init__(
        self,
        df: pandas.DataFrame,
        text_column: str,
        mwe_types: List[str] = ["NC"],
        output_dir: str = "tmp",
        tokenize=False,
    ) -> None:
        """Provide functionalities around MWEs, for unsupervised extraction of MWEs from text and replacing
        them in the corpus.

        Args:
            df: DataFrame with a text_column that contains the corpus.
            text_col: Specifies the column of DataFrame that contains the corpus. 'text_column' must contain tokenized text.
            mwe_types: Types of MWEs. Can be a list containing any of ['NC', 'JNC'].
            output_dir: Output directory where counts, MWEs and corpus with replaced MWEs are stored.
            count_dir: Directory where count_file is sotred.
            count_file: File in which counts are sotred.
            mwe_dir: Directory where mwe_file is sotred.
            mwe_file: File in which MWEs are sotred.
            tokenize: Tokenize the content of 'text_column'.

        Returns:
            None
        """
        self.df = df
        self.text_col = text_column
        for mt in mwe_types:
            if mt not in ["NC", "JNC"]:
                raise ValueError(f"{mt} type is not recognized.")
        self.mwe_types = mwe_types

        self.output_dir = output_dir
        self.count_dir = os.path.join(self.output_dir, "counts")
        self.count_file = os.path.join(self.count_dir, "count_data.json")
        self.mwe_dir = os.path.join(self.output_dir, "mwes")
        self.mwe_file = os.path.join(self.mwe_dir, "mwe_data.json")
        
        Path(self.output_dir).mkdir(exist_ok=True)

        if tokenize:
            logger.info('"tokenize" flag set to True. This might lead to a slow instantiation.')
            self.df[text_column] = self.df[text_column].apply(self._tokenize)
        else:
            self._check_tokenized()

    def _tokenize(self, x):
        """Helper function to tokenize and join the results with a space.

        Args:
            None

        Returns:
            None
        """
        return " ".join(word_tokenize(x))

    def _check_tokenized(self) -> None:
        """Helper function to check if the content of text_column is tokenized.

        Args:
            None

        Returns:
            None
        """
        if self.df[self.text_col].shape[0] > 200:
            tests = self.df[self.text_col].sample(n=200).tolist()
        else:
            tests = self.df[self.text_col].sample(frac=0.8).tolist()
        num_pass = 0
        for t in tests:
            if " ".join(word_tokenize(t)) == t:
                num_pass += 1
        if float(num_pass) / float(len(tests)) < 0.8:
            logger.warning(
                f"It seems that the content of {self.text_col} in the input data frame is not (fully) tokenized.\nThis can lead to poor results. Consider re-instantiating your MWE instance with 'tokenize' flag set to True.\nNote that this might lead to a slower instantiation."
            )

    def build_counts(self, file_name: str=None) -> None:
        """Create various count files to be used by downstream methods 
        by calling snlp.mwes.counter.get_counts.

        Args:
            None

        Returns:
            None
        """
        logger.info("Creating counts...")
        res = get_counts(df=self.df, text_column=self.text_col, mwe_types=self.mwe_types)
        # Directory
        try:
            Path(self.count_dir).mkdir(exist_ok=True)
        except Exception as e:
            logger.error(e)
            raise e
        # File
        if not file_name:
            try:
                with open(self.count_file, "w") as file:
                    json.dump(res, file)
            except Exception as e:
                logger.error(e)
                raise e
        else:
            try:            
                with open(file_name, "w") as file:
                    json.dump(res, file)
            except Exception as e:
                logger.error(e)
                raise e
            self.count_file = file_name

    def extract_mwes(self, am: str = "pmi", file_name: str=None) -> None:
        """
        Args:
            mwe_types: Types of MWEs. Can be any of [NC, JNC]
            am: The association measure to be used. Can be any of [pmi, npmi]

        Returns:
            None
        """
        with open(self.count_file, "r") as file:
            count_data = json.load(file)
        logger.info(f"Extracting {self.mwe_types} based on {am}")
        mwe_am_dict = calculate_am(count_data=count_data, am=am, mwe_types=self.mwe_types)
        # Dir
        try:
            Path(self.mwe_dir).mkdir(exist_ok=True)
        except Exception as e:
            logger.error(e)
            raise e
        # File
        if not file_name:
            try:
                with open(self.mwe_file, "w") as file:
                    json.dump(mwe_am_dict, file)
            except Exception as e:
                logger.error(e)
                raise e
        else:
            try:
                with open(file_name, "w") as file:
                    json.dump(mwe_am_dict, file)
            except Exception as e:
                logger.error(e)
                raise e


if __name__ == "__main__":
    import pandas as pd
    from snlp.mwes.mwe_utils import replace_mwes

    imdb_train = pd.read_csv("data/imdb_train_sample.tsv", sep="\t", names=["label", "text"])
    mwe = MWE(df=imdb_train, mwe_types=["NC", "JNC"], text_column="text", tokenize=True)
    mwe.build_counts()
    mwe.extract_mwes(am="npmi")
    new_df = replace_mwes(
        path_to_mwes="tmp/mwes/mwe_data.json", mwe_types=["NC", "JNC"], df=imdb_train, text_column="text"
    )
    new_df.to_csv("tmp/new_df.csv", sep="\t")
