# ========================================================================
# Copyright 2020 Emory University
#
# 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.
# ========================================================================

# -*- coding:utf-8 -*-
# Author: hankcs
import logging
import os
from abc import ABC, abstractmethod
from typing import Tuple, Union, List

from elit.common.constant import EOS, PAD
from elit.common.structure import SerializableDict, Configurable
from elit.common.vocab import Vocab
from elit.utils.io_util import get_resource, load_json
from elit.utils.reflection import classpath_of, str_to_type
from elit.utils.string_util import ispunct


class ToIndex(ABC):

    def __init__(self, vocab: Vocab = None) -> None:
        super().__init__()
        if vocab is None:
            vocab = Vocab()
        self.vocab = vocab

    @abstractmethod
    def __call__(self, sample):
        pass

    def save_vocab(self, save_dir, filename='vocab.json'):
        vocab = SerializableDict()
        vocab.update(self.vocab.to_dict())
        vocab.save_json(os.path.join(save_dir, filename))

    def load_vocab(self, save_dir, filename='vocab.json'):
        save_dir = get_resource(save_dir)
        vocab = SerializableDict()
        vocab.load_json(os.path.join(save_dir, filename))
        self.vocab.copy_from(vocab)


class FieldToIndex(ToIndex):

    def __init__(self, src, vocab: Vocab, dst=None) -> None:
        super().__init__(vocab)
        self.src = src
        if not dst:
            dst = f'{src}_id'
        self.dst = dst

    def __call__(self, sample: dict):
        sample[self.dst] = self.vocab(sample[self.src])
        return sample

    def save_vocab(self, save_dir, filename=None):
        if not filename:
            filename = f'{self.dst}_vocab.json'
        super().save_vocab(save_dir, filename)

    def load_vocab(self, save_dir, filename=None):
        if not filename:
            filename = f'{self.dst}_vocab.json'
        super().load_vocab(save_dir, filename)


class VocabList(list):

    def __init__(self, *fields) -> None:
        super().__init__()
        for each in fields:
            self.append(FieldToIndex(each))

    def append(self, item: Union[str, Tuple[str, Vocab], Tuple[str, str, Vocab], FieldToIndex]) -> None:
        if isinstance(item, str):
            item = FieldToIndex(item)
        elif isinstance(item, (list, tuple)):
            if len(item) == 2:
                item = FieldToIndex(src=item[0], vocab=item[1])
            elif len(item) == 3:
                item = FieldToIndex(src=item[0], dst=item[1], vocab=item[2])
            else:
                raise ValueError(f'Unsupported argument length: {item}')
        elif isinstance(item, FieldToIndex):
            pass
        else:
            raise ValueError(f'Unsupported argument type: {item}')
        super(self).append(item)

    def save_vocab(self, save_dir):
        for each in self:
            each.save_vocab(save_dir, None)

    def load_vocab(self, save_dir):
        for each in self:
            each.load_vocab(save_dir, None)


class VocabDict(SerializableDict):

    def __init__(self, *args, **kwargs) -> None:
        vocabs = dict(kwargs)
        for each in args:
            vocabs[each] = Vocab()
        super().__init__(vocabs)

    def save_vocabs(self, save_dir, filename='vocabs.json'):
        vocabs = SerializableDict()
        for key, value in self.items():
            if isinstance(value, Vocab):
                vocabs[key] = value.to_dict()
        vocabs.save_json(os.path.join(save_dir, filename))

    def load_vocabs(self, save_dir, filename='vocabs.json', vocab_cls=Vocab):
        save_dir = get_resource(save_dir)
        vocabs = SerializableDict()
        vocabs.load_json(os.path.join(save_dir, filename))
        self._load_vocabs(self, vocabs, vocab_cls)

    @staticmethod
    def _load_vocabs(vd, vocabs: dict, vocab_cls=Vocab):
        """

        Args:
            vd:
            vocabs:
            vocab_cls: Default class for the new vocab
        """
        for key, value in vocabs.items():
            if 'idx_to_token' in value:
                cls = value.get('type', None)
                if cls:
                    cls = str_to_type(cls)
                else:
                    cls = vocab_cls
                vocab = cls()
                vocab.copy_from(value)
                vd[key] = vocab
            else:  # nested Vocab
                # noinspection PyTypeChecker
                vd[key] = nested = VocabDict()
                VocabDict._load_vocabs(nested, value, vocab_cls)

    def lock(self):
        for key, value in self.items():
            if isinstance(value, Vocab):
                value.lock()

    @property
    def mutable(self):
        status = [v.mutable for v in self.values() if isinstance(v, Vocab)]
        return len(status) == 0 or any(status)

    def __call__(self, sample: dict):
        for key, value in self.items():
            if isinstance(value, Vocab):
                field = sample.get(key, None)
                if field is not None:
                    sample[f'{key}_id'] = value(field)
        return sample

    def __getattr__(self, key):
        if key.startswith('__'):
            return dict.__getattr__(key)
        return self.__getitem__(key)

    def __setattr__(self, key, value):
        return self.__setitem__(key, value)

    def __getitem__(self, k: str) -> Vocab:
        return super().__getitem__(k)

    def __setitem__(self, k: str, v: Vocab) -> None:
        super().__setitem__(k, v)

    def summary(self, logger: logging.Logger = None):
        for key, value in self.items():
            if isinstance(value, Vocab):
                report = value.summary(verbose=False)
                if logger:
                    logger.info(f'{key}{report}')
                else:
                    print(f'{key}{report}')

    def put(self, **kwargs):
        for k, v in kwargs.items():
            self[k] = v


class NamedTransform(ABC):
    def __init__(self, src: str, dst: str = None) -> None:
        if dst is None:
            dst = src
        self.dst = dst
        self.src = src

    @abstractmethod
    def __call__(self, sample: dict) -> dict:
        return sample


class ConfigurableTransform(Configurable, ABC):
    @property
    def config(self):
        return dict([('classpath', classpath_of(self))] +
                    [(k, v) for k, v in self.__dict__.items() if not k.startswith('_')])

    @classmethod
    def from_config(cls, config: dict):
        """

        Args:
          config: 
          kwargs: 
          config: dict: 

        Returns:

        
        """
        cls = config.get('classpath', None)
        assert cls, f'{config} doesn\'t contain classpath field'
        cls = str_to_type(cls)
        config = dict(config)
        config.pop('classpath')
        return cls(**config)


class ConfigurableNamedTransform(NamedTransform, ConfigurableTransform, ABC):
    pass


class EmbeddingNamedTransform(ConfigurableNamedTransform, ABC):

    def __init__(self, output_dim: int, src: str, dst: str) -> None:
        super().__init__(src, dst)
        self.output_dim = output_dim


class RenameField(NamedTransform):

    def __call__(self, sample: dict):
        sample[self.dst] = sample.pop(self.src)
        return sample


class CopyField(object):
    def __init__(self, src, dst) -> None:
        self.dst = dst
        self.src = src

    def __call__(self, sample: dict) -> dict:
        sample[self.dst] = sample[self.src]
        return sample


class FilterField(object):
    def __init__(self, *keys) -> None:
        self.keys = keys

    def __call__(self, sample: dict):
        sample = dict((k, sample[k]) for k in self.keys)
        return sample


class TransformList(list):
    """Composes several transforms together.

    Args:
      transforms(list of ``Transform`` objects): list of transforms to compose.
    Example:

    Returns:

    >>> transforms.TransformList(
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> )
    """

    def __init__(self, *transforms) -> None:
        super().__init__()
        self.extend(transforms)

    def __call__(self, sample):
        for t in self:
            sample = t(sample)
        return sample

    def index_by_type(self, t):
        for i, trans in enumerate(self):
            if isinstance(trans, t):
                return i


class LowerCase(object):
    def __init__(self, src, dst=None) -> None:
        if dst is None:
            dst = src
        self.src = src
        self.dst = dst

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if isinstance(src, str):
            sample[self.dst] = src.lower()
        elif isinstance(src, list):
            sample[self.dst] = [x.lower() for x in src]
        return sample


class ToChar(object):
    def __init__(self, src, dst='char', max_word_length=None, min_word_length=None, pad=PAD) -> None:
        if dst is None:
            dst = src
        self.src = src
        self.dst = dst
        self.max_word_length = max_word_length
        self.min_word_length = min_word_length
        self.pad = pad

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if isinstance(src, str):
            sample[self.dst] = self.to_chars(src)
        elif isinstance(src, list):
            sample[self.dst] = [self.to_chars(x) for x in src]
        return sample

    def to_chars(self, word: str):
        chars = list(word)
        if self.min_word_length and len(chars) < self.min_word_length:
            chars = chars + [self.pad] * (self.min_word_length - len(chars))
        if self.max_word_length:
            chars = chars[:self.max_word_length]
        return chars


class AppendEOS(NamedTransform):

    def __init__(self, src: str, dst: str = None, eos=EOS) -> None:
        super().__init__(src, dst)
        self.eos = eos

    def __call__(self, sample: dict) -> dict:
        sample[self.dst] = sample[self.src] + [self.eos]
        return sample


class WhitespaceTokenizer(NamedTransform):

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if isinstance(src, str):
            sample[self.dst] = self.tokenize(src)
        elif isinstance(src, list):
            sample[self.dst] = [self.tokenize(x) for x in src]
        return sample

    @staticmethod
    def tokenize(text: str):
        return text.split()


class NormalizeDigit(object):
    def __init__(self, src, dst=None) -> None:
        if dst is None:
            dst = src
        self.src = src
        self.dst = dst

    @staticmethod
    def transform(word: str):
        new_word = ""
        for char in word:
            if char.isdigit():
                new_word += '0'
            else:
                new_word += char
        return new_word

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if isinstance(src, str):
            sample[self.dst] = self.transform(src)
        elif isinstance(src, list):
            sample[self.dst] = [self.transform(x) for x in src]
        return sample


class Bigram(NamedTransform):

    def __init__(self, src: str, dst: str = None) -> None:
        if not dst:
            dst = f'{src}_bigram'
        super().__init__(src, dst)

    def __call__(self, sample: dict) -> dict:
        src: List = sample[self.src]
        dst = src + [EOS]
        dst = [dst[i] + dst[i + 1] for i in range(len(src))]
        sample[self.dst] = dst
        return sample


class FieldLength(NamedTransform):

    def __init__(self, src: str, dst: str = None, delta=0) -> None:
        self.delta = delta
        if not dst:
            dst = f'{src}_length'
        super().__init__(src, dst)

    def __call__(self, sample: dict) -> dict:
        sample[self.dst] = len(sample[self.src]) + self.delta
        return sample


class BMESOtoIOBES(object):
    def __init__(self, field='tag') -> None:
        self.field = field

    def __call__(self, sample: dict) -> dict:
        sample[self.field] = [self.convert(y) for y in sample[self.field]]
        return sample

    @staticmethod
    def convert(y: str):
        if y.startswith('M-'):
            return 'I-'
        return y


class NormalizeToken(ConfigurableNamedTransform):

    def __init__(self, mapper: Union[str, dict], src: str, dst: str = None) -> None:
        super().__init__(src, dst)
        self.mapper = mapper
        if isinstance(mapper, str):
            mapper = get_resource(mapper)
        if isinstance(mapper, str):
            self._table = load_json(mapper)
        elif isinstance(mapper, dict):
            self._table = mapper
        else:
            raise ValueError(f'Unrecognized mapper type {mapper}')

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if self.src == self.dst:
            sample[f'{self.src}_'] = src
        if isinstance(src, str):
            src = self.convert(src)
        else:
            src = [self.convert(x) for x in src]
        sample[self.dst] = src
        return sample

    def convert(self, token) -> str:
        return self._table.get(token, token)


class PunctuationMask(ConfigurableNamedTransform):
    def __init__(self, src: str, dst: str = None) -> None:
        """Mask out all punctuations (set mask of punctuations to False)

        Args:
          src:
          dst:

        Returns:

        """
        if not dst:
            dst = f'{src}_punct_mask'
        super().__init__(src, dst)

    def __call__(self, sample: dict) -> dict:
        src = sample[self.src]
        if isinstance(src, str):
            dst = not ispunct(src)
        else:
            dst = [not ispunct(x) for x in src]
        sample[self.dst] = dst
        return sample


class NormalizeCharacter(NormalizeToken):
    def convert(self, token) -> str:
        return ''.join([NormalizeToken.convert(self, c) for c in token])
