# Copyright (c) 2022, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import os
import torch
import random
from pathlib import Path
from dataclasses import dataclass
from typing import Optional, List, Tuple, Any
from nemo.utils import logging
import nemo_chem
import warnings

__all__ = ['MolEncTokenizer', 'MolEncTokenizerBaseConfig', 'MolEncTokenizerFromVocabFileConfig', 'MolEncTokenizerFromSmilesConfig', 'DEFAULT_SEQ_LEN', 'DEFAULT_VOCAB_PATH', 'DEFAULT_MODEL_PATH']

DEFAULT_VOCAB_DIR = os.path.join(nemo_chem.__path__[0], '../models/vocab')
DEFAULT_VOCAB_PATH = os.path.join(DEFAULT_VOCAB_DIR, 'megamolbart.vocab')
DEFAULT_MODEL_PATH = os.path.join(DEFAULT_VOCAB_DIR, 'megamolbart.model')

# Defaults
DEFAULT_SEQ_LEN = 512
DEFAULT_CHEM_TOKEN_START = 272
DEFAULT_BEGIN_TOKEN = "^"
DEFAULT_END_TOKEN = "&"
DEFAULT_PAD_TOKEN = "<PAD>"
DEFAULT_UNK_TOKEN = "?"
DEFAULT_MASK_TOKEN = "<MASK>"
DEFAULT_SEP_TOKEN = "<SEP>"
DEFAULT_MASK_PROB = 0.15
DEFAULT_SHOW_MASK_TOKEN_PROB = 1.0
DEFAULT_MASK_SCHEME = "span"
DEFAULT_SPAN_LAMBDA = 3.0

# DEFAULT_REGEX = r"""\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9]"""
with open(DEFAULT_MODEL_PATH, 'r') as fh:
    DEFAULT_REGEX = fh.read().strip()
DEFAULT_REGEX = r"""{}""".format(DEFAULT_REGEX) # force literal string

@dataclass
class MolEncTokenizerBaseConfig():
    vocab: Tuple[str] = ('',)
    chem_token_idxs: Tuple[int] = (0,)
    prog: Any = ''
    begin_token: str = DEFAULT_BEGIN_TOKEN
    end_token: str = DEFAULT_END_TOKEN
    pad_token: str = DEFAULT_PAD_TOKEN
    unk_token: str = DEFAULT_UNK_TOKEN
    mask_token: str = DEFAULT_MASK_TOKEN
    sep_token: str = DEFAULT_SEP_TOKEN
    mask_prob: float = DEFAULT_MASK_PROB
    show_mask_token_prob: float = DEFAULT_SHOW_MASK_TOKEN_PROB
    mask_scheme: str = DEFAULT_MASK_SCHEME # TODO limit to possible options
    span_lambda: float = DEFAULT_SPAN_LAMBDA

@dataclass
class MolEncTokenizerFromVocabFileConfig():
    vocab_path: str = DEFAULT_VOCAB_PATH
    regex: str = DEFAULT_REGEX
    chem_tokens_start_idx: int = DEFAULT_CHEM_TOKEN_START
    pad_token_idx: int = 0
    unk_token_idx: int = 1
    begin_token_idx: int = 2
    end_token_idx: int = 3
    mask_token_idx: int = 4
    sep_token_idx: int = 5
    mask_prob: float = DEFAULT_MASK_PROB
    show_mask_token_prob: float = DEFAULT_SHOW_MASK_TOKEN_PROB
    mask_scheme: str = DEFAULT_MASK_SCHEME # TODO limit to possible options
    span_lambda: float = DEFAULT_SPAN_LAMBDA

@dataclass
class MolEncTokenizerFromSmilesConfig():
    smiles: Tuple[str] = ('c',)
    regex: str = DEFAULT_REGEX
    extra_tokens: Optional[List[str]] = None
    begin_token: str = DEFAULT_BEGIN_TOKEN
    end_token: str = DEFAULT_END_TOKEN
    pad_token: str = DEFAULT_PAD_TOKEN
    unk_token: str = DEFAULT_UNK_TOKEN
    mask_token: str = DEFAULT_MASK_TOKEN
    sep_token: str = DEFAULT_SEP_TOKEN
    mask_prob: float = DEFAULT_MASK_PROB
    show_mask_token_prob: float = DEFAULT_SHOW_MASK_TOKEN_PROB
    mask_scheme: str = DEFAULT_MASK_SCHEME # TODO limit to possible options
    span_lambda: float = DEFAULT_SPAN_LAMBDA

class MolEncTokenizer:
    def __init__(
        self,
        vocab,
        chem_token_idxs,
        prog,
        begin_token=DEFAULT_BEGIN_TOKEN,
        end_token=DEFAULT_END_TOKEN,
        pad_token=DEFAULT_PAD_TOKEN,
        unk_token=DEFAULT_UNK_TOKEN,
        mask_token=DEFAULT_MASK_TOKEN,
        sep_token=DEFAULT_SEP_TOKEN,
        mask_prob=DEFAULT_MASK_PROB,
        show_mask_token_prob=DEFAULT_SHOW_MASK_TOKEN_PROB,
        mask_scheme=DEFAULT_MASK_SCHEME,
        span_lambda=DEFAULT_SPAN_LAMBDA
    ):
        """ Initialise the tokenizer

        Args:
            vocab (List[str]): Vocabulary for tokenizer
            chem_token_idxs (List[int]): List of idxs of chemical tokens
            prog (re.Pattern): Regex object for tokenizing
            begin_token (str): Token to use at start of each sequence
            end_token (str): Token to use at end of each sequence
            pad_token (str): Token to use when padding batches of sequences
            unk_token (str): Token to use for tokens which are not in the vocabulary
            mask_token (str): Token to use when masking pieces of the sequence
            sep_token (str): Token to use when sepatating two sentences
            mask_prob (float): Probability of token being masked when masking is enabled
            show_mask_token_prob (float): Probability of a masked token being replaced with mask token
            mask_scheme (str): Masking scheme used by the tokenizer when masking
            span_lambda (float): Mean for poisson distribution when sampling a span of tokens
        """

        self.vocab = {t: i for i, t in enumerate(vocab)}
        self.decode_vocab = {i: t for t, i in self.vocab.items()}
        self.chem_token_idxs = chem_token_idxs
        self.prog = prog

        self.begin_token = begin_token
        self.end_token = end_token
        self.pad_token = pad_token
        self.unk_token = unk_token
        self.mask_token = mask_token
        self.sep_token = sep_token

        self.mask_prob = mask_prob
        self.show_mask_token_prob = show_mask_token_prob
        self.mask_scheme = mask_scheme
        self.span_lambda = span_lambda

        self.unk_id = self.vocab[unk_token]
        self.unk_token_cnt = {}

    @staticmethod
    def from_vocab_file(
        vocab_path,
        regex=DEFAULT_REGEX,
        chem_tokens_start_idx=DEFAULT_CHEM_TOKEN_START,
        pad_token_idx=0,
        unk_token_idx=1,
        begin_token_idx=2,
        end_token_idx=3,
        mask_token_idx=4,
        sep_token_idx=5,
        mask_prob=DEFAULT_MASK_PROB,
        show_mask_token_prob=DEFAULT_SHOW_MASK_TOKEN_PROB,
        mask_scheme=DEFAULT_MASK_SCHEME,
        span_lambda=DEFAULT_SPAN_LAMBDA,
        **kwargs
    ):
        """ Load the tokenizer object from a vocab file and regex

        Reads a newline separated list of tokens from a file to use as the vocabulary
        Note: Assumes that the chemical tokens run from chem_tokens_start_idx to the end of the tokens list
              Anything after the defined tokens and before chem_tokens_start_idx is assumed to be an extra token
              and is added to the regex for tokenizing

        Args:
            vocab_path (str): Path to vocab file
            regex (str): Regex to use for tokenizing
            chem_tokens_start_idx (int): Index of the start of the chemical tokens in the tokens list

        Returns:
            MolEncTokenizer object
        """

        text = Path(vocab_path).read_text()
        tokens = text.split("\n")
        tokens = [t for t in tokens if t is not None and t != ""]

        token_idxs = [pad_token_idx, unk_token_idx, begin_token_idx, end_token_idx, mask_token_idx, sep_token_idx]
        extra_tokens_idxs = range(max(token_idxs) + 1, chem_tokens_start_idx)
        extra_tokens = [tokens[idx] for idx in extra_tokens_idxs]
        prog = MolEncTokenizer._get_compiled_regex(regex, extra_tokens)

        pad_token = tokens[pad_token_idx]
        unk_token = tokens[unk_token_idx]
        begin_token = tokens[begin_token_idx]
        end_token = tokens[end_token_idx]
        mask_token = tokens[mask_token_idx]
        sep_token = tokens[sep_token_idx]

        chem_tokens_idxs = list(range(chem_tokens_start_idx, len(tokens)))
        tokenizer = MolEncTokenizer(
            tokens,
            chem_tokens_idxs,
            prog,
            begin_token=begin_token,
            end_token=end_token,
            pad_token=pad_token,
            unk_token=unk_token,
            mask_token=mask_token,
            sep_token=sep_token,
            mask_prob=mask_prob,
            show_mask_token_prob=show_mask_token_prob,
            mask_scheme=mask_scheme,
            span_lambda=span_lambda
        )
        return tokenizer

    def vocab_size(self):
        """ Return the size of the vocab being used."""
        return len(self.vocab)

    @staticmethod
    def from_smiles(
        smiles,
        regex=DEFAULT_REGEX,
        extra_tokens=None,
        begin_token=DEFAULT_BEGIN_TOKEN,
        end_token=DEFAULT_END_TOKEN,
        pad_token=DEFAULT_PAD_TOKEN,
        unk_token=DEFAULT_UNK_TOKEN,
        mask_token=DEFAULT_MASK_TOKEN,
        sep_token=DEFAULT_SEP_TOKEN,
        mask_prob=DEFAULT_MASK_PROB,
        show_mask_token_prob=DEFAULT_SHOW_MASK_TOKEN_PROB,
        mask_scheme=DEFAULT_MASK_SCHEME,
        span_lambda=DEFAULT_SPAN_LAMBDA
    ):
        """ Build the tokenizer from smiles strings and a regex

        Args:
            smiles (List[str]): SMILES strings to use to build vocabulary
            regex (str): Regex to use for tokenizing
            extra_tokens (Optional[List[str]]): Additional tokens to add to the vocabulary that 
                                                may not appear in the SMILES strings
        """

        vocab = {
            pad_token: 0,
            unk_token: 1,
            begin_token: 2,
            end_token: 3,
            mask_token: 4,
            sep_token: 5
        }

        extra_tokens = [] if extra_tokens is None else extra_tokens
        [vocab.setdefault(token, len(vocab)) for token in extra_tokens]

        chem_start_idx = len(vocab)
        prog = MolEncTokenizer._get_compiled_regex(regex, extra_tokens)
        print(f"Chemistry tokens start at index {chem_start_idx}")

        for smi in smiles:
            for token in prog.findall(smi):
                vocab.setdefault(token, len(vocab))

        chem_token_idxs = list(range(chem_start_idx, len(vocab)))

        vocab = sorted(vocab.items(), key=lambda k_v: k_v[1])
        vocab = [key for key, val in vocab]

        tokenizer = MolEncTokenizer(
            vocab,
            chem_token_idxs,
            prog,
            begin_token=begin_token,
            end_token=end_token,
            pad_token=pad_token,
            unk_token=unk_token,
            mask_token=mask_token,
            sep_token=sep_token,
            mask_prob=mask_prob,
            show_mask_token_prob=show_mask_token_prob,
            mask_scheme=mask_scheme,
            span_lambda=span_lambda
        )
        return tokenizer

    def save_vocab(self, vocab_path):
        tokens = sorted(self.vocab.items(), key=lambda k_v: k_v[1])
        tokens = [key for key, val in tokens]

        tokens_str = ""
        for token in tokens:
            tokens_str += f"{token}\n"

        p = Path(vocab_path)
        p.write_text(tokens_str)

    def __len__(self):
        return len(self.vocab)

    def tokenize(self, sents1, sents2=None, mask=False, pad=False):
        # TODO this function needs cleanup
        if sents2:
            warnings.simplefilter('error', DeprecationWarning) # This functionality does not work, fail loudly
            warnings.warn('Sentence tokenization is not currently supported in MegaMolBART and will not produce correct results', category=DeprecationWarning)
        if pad is True:
            warnings.simplefilter('always', DeprecationWarning)
            warnings.warn('Padding by the tokenizer is not supported in this version of MegaMolBART and may not produce correct results', category=DeprecationWarning)

        if sents2 is not None and len(sents1) != len(sents2):
            raise ValueError("Sentence 1 batch and sentence 2 batch must have the same number of elements")

        tokens = self._regex_match(sents1)
        m_tokens, token_masks = self.mask_tokens(tokens, empty_mask=not mask)

        sent_masks = None
        if sents2 is not None:
            sents2_tokens = self._regex_match(sents2)
            sents2_m_tokens, sents2_masks = self.mask_tokens(sents2_tokens, empty_mask=not mask)
            tokens, sent_masks = self._concat_sentences(tokens, sents2_tokens, self.sep_token)
            m_tokens, _ = self._concat_sentences(m_tokens, sents2_m_tokens, self.sep_token)
            token_masks, _ = self._concat_sentences(token_masks, sents2_masks, False)

        # TODO this removes bos/eos addition since NeMo handles differently. 
        # Now handled in collate function
        # tokens = [[self.begin_token] + ts + [self.end_token] for ts in tokens] 
        # m_tokens = [[self.begin_token] + ts + [self.end_token] for ts in m_tokens]
        # token_masks = [[True] + ts + [True] for ts in token_masks]
        # sent_masks = [[1] + mask + [0] for mask in sent_masks] if sent_masks is not None else None

        output = {}

        if pad:
            tokens, orig_pad_masks = self._pad_seqs(tokens, self.pad_token)
            m_tokens, masked_pad_masks = self._pad_seqs(m_tokens, self.pad_token)
            token_masks, _ = self._pad_seqs(token_masks, False)
            sent_masks, _ = self._pad_seqs(sent_masks, False) if sent_masks is not None else (None, None)
            output["original_pad_masks"] = orig_pad_masks
            output["masked_pad_masks"] = masked_pad_masks

        output["original_tokens"] = tokens

        if mask:
            output["masked_tokens"] = m_tokens
            output["token_masks"] = token_masks

        if sent_masks is not None:
            output["sentence_masks"] = sent_masks

        return output

    def _regex_match(self, smiles):
        tokenized = []
        for smi in smiles:
            tokens = self.prog.findall(smi)
            tokenized.append(tokens)

        return tokenized

    @staticmethod
    def _get_compiled_regex(regex, extra_tokens):
        regex_string = r"("
        for token in extra_tokens:
            processed_token = token
            for special_character in "()[].|":
                processed_token = processed_token.replace(special_character, f"\\{special_character}")
            regex_string += processed_token + r"|"

        regex_string += regex + r"|"
        regex_string += r".)"
        return re.compile(regex_string)

    @staticmethod
    def _concat_sentences(tokens1, tokens2, sep):
        warnings.simplefilter('error', DeprecationWarning) # This function does not work, fail loudly
        warnings.warn('Sentence tokenization is not currently supported in MegaMolBART and will not produce correct results', category=DeprecationWarning)

        tokens = [ts1 + [sep] + ts2 for ts1, ts2 in zip(tokens1, tokens2)]
        sent_masks = [([1] * len(ts1)) + [1] + ([0] * len(ts2)) for ts1, ts2 in zip(tokens1, tokens2)]  # 1/True = Active, 0/False = Inactive
        return tokens, sent_masks

    def detokenize(self, tokens_list):
        new_tokens_list = []
        for tokens in tokens_list:
            if tokens[0] == self.begin_token:
                tokens = tokens[1:]

            # Remove any tokens after the end token (and end token) if it's there 
            if self.end_token in tokens:
                end_token_idx = tokens.index(self.end_token)
                tokens = tokens[:end_token_idx]

            new_tokens_list.append(tokens)

        strs = ["".join(tokens) for tokens in new_tokens_list]
        return strs

    def convert_tokens_to_ids(self, token_data):
        ids_list = []
        for tokens in token_data:
            for token in tokens:
                token_id = self.vocab.get(token)
                if token_id is None:
                    self._inc_in_dict(self.unk_token_cnt, token)

            ids = [self.vocab.get(token, self.unk_id) for token in tokens]
            ids_list.append(ids)

        return ids_list

    def convert_ids_to_tokens(self, token_ids):
        tokens_list = []
        for ids in token_ids:
            for token_id in ids:
                token = self.decode_vocab.get(token_id)
                if token is None:
                    raise ValueError(f"Token id {token_id} is not recognised")
 
            tokens = [self.decode_vocab.get(token_id) for token_id in ids]
            tokens_list.append(tokens)

        return tokens_list

    def print_unknown_tokens(self):
        print(f"{'Token':<10}Count")
        for token, cnt in self.unk_token_cnt.items():
            print(f"{token:<10}{cnt}")
    
        print()

    @staticmethod
    def _inc_in_dict(coll, item):
        cnt = coll.get(item, 0)
        cnt += 1
        coll[item] = cnt

    def mask_tokens(self, tokens, empty_mask=False):
        if empty_mask:
            mask = [[True] * len(ts) for ts in tokens]
            return tokens, mask

        masked_tokens = []
        token_masks = []

        for ts in tokens:
            if self.mask_scheme == "replace":
                masked, token_mask = self._mask_replace(ts)
            elif self.mask_scheme == "span":
                masked, token_mask = self._mask_span(ts)
            else:
                raise ValueError(f"Unrecognised mask scheme: {self.mask_scheme}")

            masked_tokens.append(masked)
            token_masks.append(token_mask)

        return masked_tokens, token_masks

    def _mask_replace(self, ts):
        mask_bools = [True, False]
        weights = [self.mask_prob, 1 - self.mask_prob]
        token_mask = random.choices(mask_bools, weights=weights, k=len(ts))
        masked = [self._mask_token(ts[i]) if m else ts[i] for i, m in enumerate(token_mask)]
        return masked, token_mask

    def _mask_span(self, ts):
        curr_token = 0
        masked = []
        token_mask = []

        mask_bools = [True, False]
        weights = [self.mask_prob, 1 - self.mask_prob]
        sampled_mask = random.choices(mask_bools, weights=weights, k=len(ts))

        while curr_token < len(ts):
            # If mask, sample from a poisson dist to get length of mask
            if sampled_mask[curr_token]:
                mask_len = torch.poisson(torch.tensor(self.span_lambda)).long().item()
                masked.append(self.mask_token)
                token_mask.append(True)
                curr_token += mask_len

            # Otherwise don't mask
            else:
                masked.append(ts[curr_token])
                token_mask.append(False)
                curr_token += 1

        return masked, token_mask

    def _mask_token(self, token):
        rand = random.random()
        if rand < self.show_mask_token_prob:
            return self.mask_token

        elif rand < self.show_mask_token_prob + ((1 - self.show_mask_token_prob) / 2):
            token_idx = random.choice(self.chem_token_idxs)
            return self.decode_vocab[token_idx]

        else:
            return token

    @staticmethod
    def _pad_seqs(seqs, pad_token):
        warnings.simplefilter('always', DeprecationWarning)
        warnings.warn('Padding by the tokenizer is not supported in this version of MegaMolBART and may not produce correct results', category=DeprecationWarning)
        
        pad_length = max([len(seq) for seq in seqs])
        padded = [seq + ([pad_token] * (pad_length - len(seq))) for seq in seqs]
        masks = [([1] * len(seq)) + ([0] * (pad_length - len(seq))) for seq in seqs]
        return padded, masks

    # NeMo compatbility
    @property
    def vocab_size(self):
        return len(self.vocab)

    def tokens_to_ids(self, token_data):
        return self.convert_tokens_to_ids(token_data)

    def ids_to_tokens(self, ids):
        return self.convert_ids_to_tokens(ids)

    def tokens_to_text(self, tokens_list):
        return self.detokenize(tokens_list)
        
    @property
    def pad_id(self):
        return self.tokens_to_ids([[self.pad_token]])[0][0]

    @property
    def bos_token(self):
        return self.begin_token

    @property
    def bos_id(self):
        return self.tokens_to_ids([[self.bos_token]])[0][0]

    @property
    def eos_token(self):
        return self.end_token

    @property
    def eos_id(self):
        return self.tokens_to_ids([[self.eos_token]])[0][0]

    @property
    def sep_id(self):
        return self.tokens_to_ids([[self.sep_token]])[0][0]

    @property
    def mask_id(self):
        return self.tokens_to_ids([[self.mask_token]])[0][0]

    # @property
    # def cls_id(self):
    #     return self.tokens_to_ids([[self.cls_token]])[0][0]
