# @author Enric Tejedor
# @date 2021-07

################################################################################
# Copyright (C) 1995-2021, Rene Brun and Fons Rademakers.                      #
# All rights reserved.                                                         #
#                                                                              #
# For the licensing terms see $ROOTSYS/LICENSE.                                #
# For the list of contributors see $ROOTSYS/README/CREDITS.                    #
################################################################################

import os
from collections import namedtuple

import ROOT
RunGraphs = ROOT.RDF.RunGraphs

class CppWorkflow(object):
    '''
    Class that encapsulates the generation of the code of an RDataFrame workflow
    in C++, together with its compilation into a shared library and execution.

    This class is used by worker processes to execute in C++ the RDataFrame
    graph they receive. This is done for the sake of performance, since running
    the graph from Python necessarily relies on jitted code, which is less
    optimized and thus slower than a shared library compiled with ACLiC.

    Attributes:
        _cached_wfs (dict): uses as key the code of workflow functions that have
            been already compiled and loaded by the current process, while the
            value is the id of a given workflow function. Used to prevent
            recompilation of already executed workflow functions.

        _FUNCTION_NAME (string): name of the function that encapsulates the
            RDataFrame graph creation

        _FUNCTION_NAMESPACE (string): namespace of the function that
            encapsulates the RDataFrame graph creation

        _graph_nodes (string): statements that add nodes to the graph.

        _includes (string): include statements needed by the workflow.

        _lambdas (string): lambda functions used by the workflow.

        _lambda_id (int): counter used to generate ids for each defined lambda
            function.

        _node_id (int): counter used to generate ids for each graph node.

        _PyActionData (named tuple type): pair `(res_id,operation)`, where
            `res_id` is the Python action result index in the vector of
            workflow results and `operation` is the Operation object
            corresponding to the action.

        _py_actions (list): list that contains _PyActionData objects.

        _res_ptr_id (int): counter used to generate ids for each result
            generated by graph actions.

        _SnapshotData (named tuple type): pair `(res_id,filename)`, where
            `res_id` is the Snapshot result index in the vector of workflow
            results and `filename` is its modified output file name.

        _snapshots (list): list that contains _SnapshotData objects

        _wf_id (int): used to assign new ids to workflow functions
    '''

    _FUNCTION_NAME = '__RDF_WORKFLOW_FUNCTION__'
    _FUNCTION_NAMESPACE = 'DistRDF_Internal'

    _cached_wfs = {}
    _wf_id = 0

    _SnapshotData = namedtuple('SnapshotData', ['res_id', 'filename'])
    _PyActionData = namedtuple('PyActionData', ['res_id', 'operation'])

    def __init__(self, head_node, range_id):
        '''
        Generates the C++ code of an RDF workflow that corresponds to the
        received graph and data range.

        Args:
            head_node (Node): head node of a graph that represents an RDF
                workflow.
            range_id (int): id of the data range to be processed by this
                workflow. Needed to assign a name to a partial Snapshot output
                file.
        '''

        self._includes = '''
#include "ROOT/RDataFrame.hxx"
#include "ROOT/RResultHandle.hxx"
#include "ROOT/RDFHelpers.hxx"

#include <utility>
#include <vector>
#include <tuple>
'''

        self._lambdas = ''
        self._lambda_id = 0

        self._graph_nodes = ''
        self._node_id = 1  # 0 is the head node we receive

        self._res_ptr_id = 0

        self._snapshots = []

        self._py_actions = []

        # Generate the C++ workflow.
        # Recurse over children nodes of received graph head node
        head_node_id = 0
        for child_node in head_node.children:
            self._explore_graph(child_node, range_id, head_node_id)

    def _explore_graph(self, node, range_id, parent_id):
        """
        Recursively traverses the graph nodes in DFS order and, for each of
        them, adds a new node to the C++ workflow.

        Args:
            node (Node): object that contains the information to add the
                corresponding node to the C++ workflow.
            range_id (int): id of the current range. Needed to assign a name to a
                partial Snapshot output file.
            parent_id (int): id of the parent node in the C++ workflow.
        """
        node_id = self.add_node(node.operation, range_id, parent_id)

        for child_node in node.children:
            self._explore_graph(child_node, range_id, node_id)

    def add_include(self, header):
        '''
        Adds a new include statement.

        Args:
            header (str): header to be included.
        '''

        self._includes += '\n#include "{}"'.format(header)

    def add_lambda(self, lambda_code):
        '''
        Adds a new lambda to be defined, which is needed by some operation
        in the workflow.

        Args:
            lambda_code (str): code of the lambda to be defined.
        '''

        self._lambdas += '\n  auto rdf_lambda{n} = {code};' \
                        .format(n=self._lambda_id, code=lambda_code)
        self._lambda_id += 1

    def add_node(self, operation, range_id, parent_id):
        '''
        Adds the corresponding statement to add a new node in the RDataFrame
        graph. What the statement returns depends on the type of the
        operation:
        - Transformation: the statement produces a new dataset node.
        - Action: the statement returns an RResultPtr, which is added to a
        vector of results to be returned at the end of the workflow
        generation function. The type of the elements of the vector is
        ROOT::RDF::RResultHandle to do type erasure since there could be
        multiple booked actions with results of different types (e.g. integers,
        TH1D, etc.).

        Args:
            operation (Operation): object representing the operation to be added
            to the graph.

            range_id (int): id of the current range. Needed to assign a name to
            a partial Snapshot output file.

            parent_id (int): id of the parent node in this workflow. Used to
                connect the new node to be added with its parent.

        Returns:
            int: identifier of the node that was added to the graph
            corresponding to its DFS order, if the operation is a
            transformation, or None, if the operation is an action.
        '''

        # Operations that need special treatment
        if operation.name == 'AsNumpy':
            self._handle_asnumpy(operation, parent_id)
            return None # nothing else to do

        if operation.name == 'Snapshot':
            self._handle_snapshot(operation, range_id) # this modifies operation.args

        # Generate the code of the call that creates the new node
        op_call = 'rdf{n}.{op}{templ_args}({args});' \
                  .format(n=parent_id, op=operation.name, \
                          templ_args=self._get_template_args(operation), \
                          args=self._get_call_args(operation))

        if operation.is_transformation():
            new_node_id = self._node_id
            self._graph_nodes += '\n  auto rdf{n} = {call}' \
                                .format(n=new_node_id, call=op_call)
            self._node_id += 1
            return new_node_id

        # Else it's an action or instant action
        self._graph_nodes += '\n  auto res_ptr{n} = {call}' \
                            .format(n=self._res_ptr_id, call=op_call)

        # The result is stored in the vector of results to be returned
        self._graph_nodes += '\n  result_handles.emplace_back(res_ptr{});' \
                            .format(self._res_ptr_id)

        # The result type is stored in the vector of result types to be
        # returned
        self._graph_nodes += \
            '\n  auto c{0} = TClass::GetClass(typeid(res_ptr{0}));' \
            '\n  if (c{0} == nullptr)' \
            '\n    throw std::runtime_error(' \
            '\n    "Cannot get type of result {0} of action {1} during "' \
            '\n    "generation of RDF C++ workflow");' \
            '\n  result_types.emplace_back(c{0}->GetName());' \
            .format(self._res_ptr_id, operation.name)

        self._res_ptr_id += 1

    def _handle_snapshot(self, operation, range_id):
        '''
        Does two extra settings needed for Snapshot nodes:
        - Modifies the output file name to be of the form `filename_rangeid`,
        since it is a partial snapshot for a given range.
        - Stores the index of the returned vector<RResultHandle> in which the
        result of this Snapshot is stored, together with the modified file
        path.

        Args:
            operation (Operation): object representing the operation to be added
                to the graph.

            range_id (int): id of the current range. Needed to assign a name to
                a partial Snapshot output file.
        '''

        # Modify file name
        filename = operation.args[1].partition('.root')[0]
        path_with_range = '{filename}_{rangeid}.root' \
                          .format(filename=filename, rangeid=range_id)
        operation.args[1] = path_with_range

        # Store Snapshot result index -> path
        self._snapshots.append(
            CppWorkflow._SnapshotData(self._res_ptr_id, path_with_range))

    def _handle_asnumpy(self, operation, parent_id):
        '''
        Since AsNumpy is a Python-only action, it can't be included in the
        C++ workflow built by this class. Therefore, this function takes care
        of saving the RDF node, generated in C++, on which an AsNumpy action
        should be applied from Python.

        Args:
            operation (Operation): object representing the AsNumpy operation
            
            parent_id (int): id of the parent node in this workflow. Used to
                save that node in the vector of output nodes.
        '''

        # Store DFS-order index of the AsNumpy operation, together with the
        # operation information, for later invocation from Python
        self._py_actions.append(CppWorkflow._PyActionData(self._res_ptr_id,
                                                         operation))
        self._res_ptr_id += 1

        # Save parent RDF node to run AsNumpy on it later from Python
        self._graph_nodes += \
            '\n  output_nodes.push_back(ROOT::RDF::AsRNode(rdf{}));' \
            .format(parent_id)

        # Add placeholders to the result lists
        self._graph_nodes += '\n  result_handles.emplace_back();'
        self._graph_nodes += '\n  result_types.emplace_back();'

    def _get_template_args(self, operation):
        '''
        Gets the template arguments with which to generate the call to a given
        operation.

        Args:
            operation (Operation): object representing the operation whose
                template arguments need to be returned.

        Returns:
            string: template arguments for this operation.
        '''

        # TODO: generate templated operations when possible, e.g. Max<double>

        return ''

    def _get_call_args(self, operation):
        '''
        Gets the arguments with which to generate the call to a given operation.

        Args:
            operation (Operation): object representing the operation whose
                call arguments need to be returned.

        Returns:
            string: call arguments for this operation.
        '''

        # TODO
        # - Do a more thorough type conversion
        # - Use RDF helper functions to convert jitted strings to lambdas

        args = ""

        # Argument type conversion
        for narg, arg in enumerate(operation.args):
            if (narg > 0):
                args += ', '

            if isinstance(arg, str):
                args += '"{}"'.format(arg)
            elif isinstance(arg, tuple):
                args += '{'
                for nelem, elem in enumerate(arg):
                    if nelem > 0:
                        args += ','
                    if isinstance(elem, str):
                        args += '"{}"'.format(elem)
                    else:
                        args += '{}'.format(elem)
                args += '}'
            elif isinstance(arg, ROOT.RDF.RSnapshotOptions):
                ROOT.Warning('DistRDF',
                             'Processing of RSnapshotOptions not implemented')

        # Make Snapshot lazy
        # TODO: Do a proper processing of the args (user might have specified
        # her own options object)
        # Remove warning above when this is implemented
        if operation.name == 'Snapshot':
            args += ', "", lazy_options'

        return args

    def execute(self, rdf):
        '''
        Compiles the workflow generation code and executes it.

        Args:
            rdf (ROOT::RDF::RNode): object that represents the dataset on
                which to execute the workflow.

        Returns:
            tuple: the first element is the list of results of the actions in
                the C++ workflow, the second element is the list of result types
                corresponding to those actions.
        '''

        wf_id = self._compile()
        return self._run_function(rdf, wf_id)

    def _compile(self):
        '''
        Generates the workflow code C++ file and compiles it with ACLiC
        into a shared library. The library is also loaded as part of the
        `TSystem::CompileMacro` call.

        The name of the generated C++ file contains both a hash of its
        code and the ID of the process that created it. This is done to
        prevent clashes between multiple (non-sandboxed) worker processes
        that try to write to the same file concurrently.

        A class-level cache keeps track of the workflows that have been already
        compiled to prevent unncessary recompilation (e.g. when a worker
        process runs multiple times the same workflow).

        Returns:
            int: the id of the workflow function to be executed. Such id is
                appended to CppWorkflow._FUNCTION_NAME to prevent name clashes
                (a worker process might compile and load multiple workflow
                functions).
        '''

        # TODO: Make this function thread-safe? To support Dask threaded
        # workers

        code = self._get_code()
        this_wf_id = CppWorkflow._cached_wfs.get(code)
        if this_wf_id is not None:
            # We already compiled and loaded a workflow function with this
            # code. Return the id of that function
            return this_wf_id

        # We are trying to run this workflow for the first time in this
        # process. First dump the code in a file with the right function name
        this_wf_id = CppWorkflow._wf_id
        cpp_file_name = 'rdfworkflow_{wf_id}_{pid}.cxx' \
                        .format(wf_id=this_wf_id, pid=os.getpid())
        final_code = code.replace(CppWorkflow._FUNCTION_NAME,
                                  CppWorkflow._FUNCTION_NAME + str(this_wf_id),
                                  1)
        with open(cpp_file_name, 'w') as f:
            f.write(final_code)

        # Now compile and load the code
        if not ROOT.gSystem.CompileMacro(cpp_file_name, 'O'):
            raise RuntimeError(
                'Error compiling the RDataFrame workflow file: {}' \
                .format(cpp_file_name))

        # Let the cache know there is a new workflow
        CppWorkflow._cached_wfs[code] = this_wf_id
        CppWorkflow._wf_id += 1

        return this_wf_id

    def _run_function(self, rdf, wf_id):
        '''
        Runs the workflow generation function.

        Args:
            rdf (ROOT::RDF::RNode): object that represents the dataset on
                which to execute the workflow.
            wf_id (int): identifier of the workflow function to be executed.

        Returns:
            tuple: the first element is the list of results of the actions in
                the C++ workflow, the second element is the list of result types
                corresponding to those actions.
        '''

        ns = getattr(ROOT, CppWorkflow._FUNCTION_NAMESPACE)
        func = getattr(ns, CppWorkflow._FUNCTION_NAME + str(wf_id))

        # Run the workflow generator function
        vectors = func(rdf) # need to keep the tuple alive
        v_results, v_res_types, v_nodes = vectors

        # Convert the vector of results into a list so that we can mix
        # different types in it.
        # We copy the results since the life of the original ones is tied to
        # that of the vector
        results = [ ROOT.RDF.RResultHandle(res) for res in v_results ]

        # Strip out the ROOT::RDF::RResultPtr<> part of the type
        def get_result_type(s):
            if s.empty():
                # Python-only actions have an empty return type in C++
                return ''

            s = str(s)
            pos = s.find('<')
            if pos == -1:
                raise RuntimeError(
                    'Error parsing the result types of RDataFrame workflow')
            return s[pos+1:-1].strip()

        res_types = [ get_result_type(elem) for elem in v_res_types ]

        # Add Python-only actions on their corresponding nodes
        for (i, operation), n in zip(self._py_actions, v_nodes):
            operation.kwargs['lazy'] = True  # make it lazy
            results[i] = getattr(n, operation.name)(*operation.args, **operation.kwargs)

        if v_results:
            # We trigger the event loop here, so make sure we release the GIL
            old_rg = RunGraphs.__release_gil__
            RunGraphs.__release_gil__ = True
            RunGraphs(v_results)
            RunGraphs.__release_gil__ = old_rg

        # Replace the RResultHandle of each Snapshot by its modified output
        # path, since the latter is what we actually need in the reducer
        for i, path in self._snapshots:
            results[i] = [path]
            res_types[i] = None # placeholder

        # Replace the future-like result of every Python-only action (e.g.
        # AsNumpyResult) by its actual value
        for i, operation in self._py_actions:
            results[i] = results[i].GetValue()

        return results, res_types

    def __repr__(self):
        '''
        Generates a string representation for this C++ workflow.

        Returns:
            string: code of this C++ workflow.
        '''

        return self._get_code()

    def _get_code(self):
        '''
        Composes the workflow generation code from the different attributes
        of this class. The resulting code contains a function that will be
        called to generate the RDataFrame graph. Such function returns a tuple
        of three elements:
        1. A vector of results of the graph actions.
        2. A vector with the result types of those actions.
        3. A vector of RDF nodes that will be used in Python to invoke
        Python-only actions on them (e.g. `AsNumpy`).
        '''

        code = '''
{includes}

namespace {namespace} {{

using CppWorkflowResult = std::tuple<std::vector<ROOT::RDF::RResultHandle>,
                          std::vector<std::string>,
                          std::vector<ROOT::RDF::RNode>>;

CppWorkflowResult {func_name}(ROOT::RDF::RNode &rdf0)
{{
  std::vector<ROOT::RDF::RResultHandle> result_handles;
  std::vector<std::string> result_types;
  std::vector<ROOT::RDF::RNode> output_nodes;

  // To make Snapshots lazy
  ROOT::RDF::RSnapshotOptions lazy_options;
  lazy_options.fLazy = true;

{lambdas}

{nodes}

  return {{ std::move(result_handles), std::move(result_types), std::move(output_nodes) }};
}}

}}
'''.format(func_name=CppWorkflow._FUNCTION_NAME,
           namespace=CppWorkflow._FUNCTION_NAMESPACE,
           includes=self._includes,
           lambdas=self._lambdas,
           nodes=self._graph_nodes)

        return code

