# _*_ Coding: utf-8 -*-

# Iterates over the output generated by the Pytorch,
# TensorFlow and TensorFlow.js crawlers and provided a sorted
# output of every function.

import json
import re
import warnings
from translations import WORD_TRANSLATIONS, COMMAND_TRANSLATIONS

# Compiler() is the main parent class responsible for
# all the translation from the respective output files
# from all the crawlers.


class Compiler(object):
    # Create empyt array for TensorFlow, Pytorch and tensorflow.js
    # output files for separating the fucntions for each library.
    tf = []
    tfjs = []
    torch = []
    # Create dictionary for TensorFlow, TensorFlow.js and Pytorch.
    # These dictionary contain function name as header and then has
    # it's attributes, function calls and arguments.
    main_map = {'tf': {}, 'tfjs': {}, 'torch': {}}
    base_defs = set()

    def __init__(self):
        """This constructor calls all our data from tesnorFlow, Tensorflow.js
        and Pytorch crawler output files and creates instances
        for further usage.
        """
        with open('./output/tf/2.1.json') as tf_file, \
                open('./output/torch/1.4.0.json') as torch_file, \
                open('./output/tfjs/1.5.1.json') as tfjs_file:
            self.tf = json.load(tf_file)
            self.tfjs = json.load(tfjs_file)
            self.torch = json.load(torch_file)

        def normalize_func_name(self, name):
            """The normalize_func_name() function is a method which takes
             all the string values under function name from the target
            file and normalizizes them into header format.

            Args :
            name: Name of the fucntion from the output files from each
                    library.
             """
            # Regex rules for compiling a string to a Regex object.
            # Here the rules match on to two types of characters in a string.
            alpha = re.compile(r'[\W][a-zA-Z0-9]*')
            return alpha.sub('', name).lower()

    def generate_attrs(self, code):
        """The generate_attrs() function is a method which takes all
        the fucntion calls from the targeted code.

        Args:
            code: Initialization code or the arguments repective to the
            fucntion name form the output file of that library.
        """
        # Regex rules for compiling a string to a Regex object.
        # Here the rules match on to two types of groups in a given code.
        split_def = re.compile(r'^([\w\.]+)\((.*)\)')
        return split_def.match(code)[1].split('.')

    def populate_command(self, lib):
        """The populate_command() fucntion arranges all the output
        that is generated into a proper order and with appropriate header name.

        Args:
            lib: A dictionary consisting of both fucntion name and function
            calls separated according to the libararies that are considered.
        """
        for f in getattr(self, lib):
            # generates a header for the fucntion under which all the
            # related arguments are placed in the output file.
            nfunc = self.normalize_func_name(f['function_name'])
            # Generates the functiona call for the respective fucntion.
            f['attrs'] = self.generate_attrs(f['code'])
            # Generates all the related arguments to the function.
            f['args'] = self.hydrate_args(f['args'], f['kwargs'])
            # Generates the fucntion name for reference.
            f['name'] = f['function_name']

            if nfunc in self.main_map[lib]:
                existing_f = self.main_map[lib][nfunc][0]
                if len(f['attrs']) > len(existing_f['attrs']):
                    warnings.warn(
                        f'Ignoring op: {f}: {nfunc} is already defined in '
                        f'{lib} with shallower or '
                        f'same attrs path: {existing_f}'
                    )
                    continue
                else:
                    warnings.warn(
                        f'Overwriting op: {existing_f} with {f} '
                        'that has shallower attrs path'
                    )

            del f['kwargs']
            del f['code']
            del f['function_name']
            self.main_map[lib][nfunc] = [f]
            self.base_defs.add(nfunc)

    def load_base_defs(self):
        """The load_base_defs() function itirates through the libraries
            and generates a normalized output.

            Here base_defs() links to the dictionaries created by importing
            functions from the output files of the considered libraries.
        """
        self.populate_command('tf')
        self.populate_command('tfjs')
        self.populate_command('torch')

    def hydrate_args(self, base_args, base_kwargs):
        """The hydrate_args() fucntion maps the generated list and filters
        them out with selected arguments according to the operators and
        stores in variable as a list.

        Args:
            base_args: List of argumnets that are stored in the list
                and are passed to the fucntion.

            base_kwargs: List of keywords and variable-length arguments
                that are passed through the function.
        """
        base_args = list(filter(lambda a: a not in ['', '?'], base_args))
        base_kwargs = list(
            filter(lambda a: a[0] not in ['', '?'], base_kwargs)
        )
        # normal arguments
        ba = [
            {
                'name': self.normalize_func_name(a),
                'kwarg': False,
                'opt': a.endswith('?'),
            } for i, a in enumerate(base_args)
        ]
        # keyword arguments
        bk = [
            {
                'name': self.normalize_func_name(a[0]),
                'kwarg': True,
                'opt': True
            } for a in base_kwargs
        ]
        return ba + bk

    def match_arg_names(self, from_args, to_args, to_lang):
        """The match_arg_names() fucntion creates a loop that goes through
        all the fucntion and checks for the order issues between
        the fucntion considered.

        Example: There is a fucntion mapping to Pytorch matmul
        and TensorFlow matmul that may take same arguments but in
        different order.
        This fucntion compares the arguments between the two and fixes
        any issue related to order between the two.

        Args:
        from_args: list of arguments for calling a fuunction.

        to_args: list of arguments for the same function but for
            different library.

        to_lang: arguments for the fucntion calling.
        """
        for from_arg in from_args:
            try:
                to_arg = next(
                    m for m in to_args if m.get('name') == from_arg.get('name')
                )
                to_name = to_arg.get('name', None)
                if to_name is not None:
                    from_arg[to_lang] = to_name
            except Exception:
                to_name = WORD_TRANSLATIONS[to_lang].get(
                    from_arg['name'], None)
                if to_name is not None:
                    from_arg[to_lang] = to_name
        return from_args

    def load_translations(self, from_lang):
        """The load_translations() method that translates the list
        generated by all the crawlers with help of word translation.
        """
        langs = ['torch', 'tfjs', 'tf']
        langs.pop(langs.index(from_lang))

        for d in self.base_defs:
            # First perform any word level translations
            from_name = WORD_TRANSLATIONS[from_lang].get(d, False) or d

            # Check if translation exists in our from language
            if from_name not in self.main_map[from_lang]:
                continue

            # TODO: Treat as list
            from_def = self.main_map[from_lang][from_name][0]
            from_args = from_def['args']
            # Check if translatable to other langs
            for to_lang in langs:
                # Perform word level translation for target language
                to_name = WORD_TRANSLATIONS[to_lang].get(d, False) or d
                if to_name not in self.main_map[to_lang]:
                    continue
                # TODO: Treat as list
                to_def = self.main_map[to_lang][to_name][0]
                from_def[to_lang] = to_name

                # Format & match args
                to_args = to_def['args']
                from_def['args'] = self.match_arg_names(
                    from_args,
                    to_args,
                    to_lang
                )

    def slim_output(self, from_lang, to_lang):
        """The slim_output() fucntion creates dictionary from the
        language lists with a labeled header and returns it to
        the output variable as a list.

        Args:
            from_lang: list of arguments for calling a fuunction.

            to_lang: List of language
        """
        input_dict = {
            from_lang: self.main_map[from_lang],
            to_lang: self.main_map[to_lang]
        }
        output = {from_lang: {}, to_lang: {}}

        def hydrate_keys(input_dict, output, from_lang, to_lang):
            """The hydrate_keys() function generates dictionary from the
            given input variables by sorting them according to their labels.

            Args:
                input_dict: Dictionary conatining the cluster for the
                arguments for the function.

                output: Dictionary mapping to the language list
                    and label header.

                from_lang: list of arguments for calling a fuunction.

                to_lang: arguments for the fucntion calling.
            """
            for k, v in input_dict[from_lang].items():
                # TODO: Treat as list
                v = v[0]
                if v.get(to_lang, '') in input_dict[to_lang]:
                    output[from_lang][k] = [v]
            return output

        output = hydrate_keys(input_dict, output, from_lang, to_lang)
        output = hydrate_keys(input_dict, output, to_lang, from_lang)

        return output

    # The load_manual_translations() function maps the commands for
    # a certain fucntion into a dictioanry with labels for the commands.
    def load_manual_translations(self):
        for lang, commands in COMMAND_TRANSLATIONS.items():
            self.main_map[lang].update(commands)

    # The output_data() function maps all the compressed binareis into
    # a .json file.
    def output_data(self):
        with open('../../pythreepio/static/mapped_commands_full.json', 'w',
                  encoding='utf8') as f:
            # dumps entire compressed data into the file.
            json.dump(self.main_map, f, indent=4, ensure_ascii=False)
        # Compressed map with only torch & tfjs
        torch_tfjs_map = self.slim_output('torch', 'tfjs')
        with open('../../static/mapped_commands_torch_tfjs.json', 'w',
                  encoding='ascii') as f:
            # dumps tfjs and torch compressed data in to json file.
            json.dump(
                torch_tfjs_map,
                f,
                separators=(',', ':'),
                ensure_ascii=False
            )


def main():
    c = Compiler()  # The parent class called as variable
    c.load_base_defs()  # Caches all the functions form the libraries.
    # Creates dictionary from the mapped commands form the libraries.
    c.load_manual_translations()
    c.load_translations('tfjs')  # Generate list for TensorFlow functions.
    c.load_translations('torch')  # Generate list for Torch functions.
    c.load_translations('tf')  # Generate list for TensorFlow.js fucntions.
    c.output_data()  # Dumps all the created lists into the output .json file.A


if __name__ == '__main__':
    main()
