import numpy as np
import abc
import os
from typing import Optional, Dict, Tuple, Iterable

from common import common
from vocabularies import Code2VecVocabs, VocabType
from config import Config


class Code2VecModelBase(abc.ABC):
    def __init__(self, config: Config):
        self.config = config
        self.config.verify()

        self._log_creating_model()

        if not config.RELEASE:
            self._init_num_of_examples()
        self._log_model_configuration()
        self.vocabs = Code2VecVocabs(config)
        self.vocabs.target_vocab.get_index_to_word_lookup_table()  # just to initialize it (if not already initialized)
        self._load_or_create_inner_model()
        self._initialize()

    def _log_creating_model(self):
        self.log('')
        self.log('')
        self.log('---------------------------------------------------------------------')
        self.log('---------------------------------------------------------------------')
        self.log('---------------------- Creating code2vec model ----------------------')
        self.log('---------------------------------------------------------------------')
        self.log('---------------------------------------------------------------------')

    def _log_model_configuration(self):
        self.log('---------------------------------------------------------------------')
        self.log('----------------- Configuration - Hyper Parameters ------------------')
        longest_param_name_len = max(len(param_name) for param_name, _ in self.config)
        for param_name, param_val in self.config:
            self.log('{name: <{name_len}}{val}'.format(
                name=param_name, val=param_val, name_len=longest_param_name_len + 2))
        self.log('---------------------------------------------------------------------')

    @property
    def logger(self):
        return self.config.get_logger()

    def log(self, msg):
        self.logger.info(msg)

    def _init_num_of_examples(self):
        self.log('Checking number of examples ...')
        if self.config.is_training:
            self.config.NUM_TRAIN_EXAMPLES = self._get_num_of_examples_for_dataset(self.config.train_data_path)
            self.log('    Number of train examples: {}'.format(self.config.NUM_TRAIN_EXAMPLES))
        if self.config.is_testing:
            self.config.NUM_TEST_EXAMPLES = self._get_num_of_examples_for_dataset(self.config.TEST_DATA_PATH)
            self.log('    Number of test examples: {}'.format(self.config.NUM_TEST_EXAMPLES))

    @staticmethod
    def _get_num_of_examples_for_dataset(dataset_path: str) -> int:
        dataset_num_examples_file_path = dataset_path + '.num_examples'
        if os.path.isfile(dataset_num_examples_file_path):
            with open(dataset_num_examples_file_path, 'r') as file:
                num_examples_in_dataset = int(file.readline())
        else:
            num_examples_in_dataset = common.count_lines_in_file(dataset_path)
            with open(dataset_num_examples_file_path, 'w') as file:
                file.write(str(num_examples_in_dataset))
        return num_examples_in_dataset

    def load_or_build(self):
        self.vocabs = Code2VecVocabs(self.config)
        self._load_or_create_inner_model()

    def save(self, model_save_path=None):
        if model_save_path is None:
            model_save_path = self.config.MODEL_SAVE_PATH
        model_save_dir = '/'.join(model_save_path.split('/')[:-1])
        if not os.path.isdir(model_save_dir):
            os.makedirs(model_save_dir, exist_ok=True)
        self.vocabs.save(self.config.get_vocabularies_path_from_model_path(model_save_path))
        self._save_inner_model(model_save_path)

    def _write_code_vectors(self, file, code_vectors):
        for vec in code_vectors:
            file.write(' '.join(map(str, vec)) + '\n')

    def _get_attention_weight_per_context(
            self, path_source_strings: Iterable[str], path_strings: Iterable[str], path_target_strings: Iterable[str],
            attention_weights: Iterable[float]) -> Dict[Tuple[str, str, str], float]:
        attention_weights = np.squeeze(attention_weights, axis=-1)  # (max_contexts, )
        attention_per_context: Dict[Tuple[str, str, str], float] = {}
        # shape of path_source_strings, path_strings, path_target_strings, attention_weights is (max_contexts, )

        # iterate over contexts
        for path_source, path, path_target, weight in \
                zip(path_source_strings, path_strings, path_target_strings, attention_weights):
            string_context_triplet = (common.binary_to_string(path_source),
                                      common.binary_to_string(path),
                                      common.binary_to_string(path_target))
            attention_per_context[string_context_triplet] = weight
        return attention_per_context

    def close_session(self):
        # can be overridden by the implementation model class.
        # default implementation just does nothing.
        pass

    @abc.abstractmethod
    def evaluate(self) -> Optional[str]:
        ...

    def _load_or_create_inner_model(self):
        if self.config.is_loading:
            self._load_inner_model()
        else:
            self._create_inner_model()

    @abc.abstractmethod
    def _load_inner_model(self):
        ...

    def _create_inner_model(self):
        # can be overridden by the implementation model class.
        # default implementation just does nothing.
        pass

    def _initialize(self):
        # can be overridden by the implementation model class.
        # default implementation just does nothing.
        pass
