################################################################################
#
# Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors
#
# This file is a part of the MadGraph5_aMC@NLO project, an application which 
# automatically generates Feynman diagrams and matrix elements for arbitrary
# high-energy processes in the Standard Model and beyond.
#
# It is subject to the MadGraph5_aMC@NLO license which should accompany this 
# distribution.
#
# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch
#
################################################################################

"""Methods and classes to export matrix elements to Python format."""

from __future__ import absolute_import
import fractions
import glob
import itertools
import logging
import os
import re
import shutil
import subprocess
import aloha

import madgraph.core.color_algebra as color
import madgraph.core.helas_objects as helas_objects
import madgraph.iolibs.drawing_eps as draw
import madgraph.iolibs.files as files
import madgraph.iolibs.helas_call_writers as helas_call_writers
import madgraph.iolibs.file_writers as writers
import madgraph.iolibs.template_files as Template
import madgraph.iolibs.ufo_expression_parsers as parsers
import madgraph.iolibs.group_subprocs as group_subprocs
from madgraph import MadGraph5Error, MG5DIR

import madgraph.various.misc as misc

import aloha.create_aloha as create_aloha
import aloha.aloha_writers as aloha_writers

_file_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] + '/'
logger = logging.getLogger('madgraph.export_python')


#===============================================================================
# ProcessExporterPython
#===============================================================================
class ProcessExporterPython(object):
    """Class to take care of exporting a set of matrix elements to
    Python format."""

    class ProcessExporterPythonError(Exception):
        pass
    
    def __init__(self, matrix_elements, python_helas_call_writer):
        """Initiate with matrix elements, helas call writer.
        Generate the process matrix element functions as strings."""

        self.config_maps = {}
        if isinstance(matrix_elements, helas_objects.HelasMultiProcess):
            self.matrix_elements = matrix_elements.get('matrix_elements')
        elif isinstance(matrix_elements,
                        group_subprocs.SubProcessGroup):
            self.config_maps = matrix_elements.get('diagram_maps')
            self.matrix_elements = matrix_elements.get('matrix_elements')
        elif isinstance(matrix_elements,
                        helas_objects.HelasMatrixElementList):
            self.matrix_elements = matrix_elements
        elif isinstance(matrix_elements,
                        helas_objects.HelasMatrixElement):
            self.matrix_elements = helas_objects.HelasMatrixElementList(\
                                                     [matrix_elements])
        if not self.matrix_elements:
            raise MadGraph5Error("No matrix elements to export")

        self.model = self.matrix_elements[0].get('processes')[0].get('model')

        self.helas_call_writer = python_helas_call_writer

        if not isinstance(self.helas_call_writer, helas_call_writers.PythonUFOHelasCallWriter):
            raise Exception("helas_call_writer not PythonUFOHelasCallWriter")

        self.matrix_methods = {}

    # Methods for generation of process file strings in Python

    #===========================================================================
    # write_python_process_cc_file
    #===========================================================================
    def get_python_matrix_methods(self, gauge_check=False):
        """Write the matrix element calculation method for the processes"""

        replace_dict = {}

        # Extract version number and date from VERSION file
        info_lines = self.get_mg5_info_lines()
        replace_dict['info_lines'] = info_lines

        for ime, matrix_element in enumerate(self.matrix_elements):
            process_string = matrix_element.get('processes')[0].shell_string()
            if process_string in self.matrix_methods:
                continue

            replace_dict['process_string'] = process_string

            # Extract number of external particles
            (nexternal, ninitial) = matrix_element.get_nexternal_ninitial()
            replace_dict['nexternal'] = nexternal

            # Extract ncomb
            ncomb = matrix_element.get_helicity_combinations()
            replace_dict['ncomb'] = ncomb

            # Extract helicity lines
            helicity_lines = self.get_helicity_matrix(matrix_element)
            replace_dict['helicity_lines'] = helicity_lines

            # Extract overall denominator
            # Averaging initial state color, spin, and identical FS particles
            den_factor_line = self.get_den_factor_line(matrix_element)
            replace_dict['den_factor_line'] = den_factor_line

            # Extract process info lines for all processes
            process_lines = self.get_process_info_lines(matrix_element)
            replace_dict['process_lines'] = process_lines
        
            # Extract ngraphs
            ngraphs = matrix_element.get_number_of_amplitudes()
            replace_dict['ngraphs'] = ngraphs

            # Extract ndiags
            ndiags = len(matrix_element.get('diagrams'))
            replace_dict['ndiags'] = ndiags

            # Extract helas calls
            helas_calls = self.helas_call_writer.get_matrix_element_calls(\
                                                    matrix_element, gauge_check)
            replace_dict['helas_calls'] = "\n        ".join(helas_calls)

            # Extract nwavefuncs
            nwavefuncs = matrix_element.get_number_of_wavefunctions()
            replace_dict['nwavefuncs'] = nwavefuncs

            # Extract ncolor
            ncolor = max(1, len(matrix_element.get('color_basis')))
            replace_dict['ncolor'] = ncolor

            # Extract model parameter lines
            model_parameter_lines = \
                                 self.get_model_parameter_lines(matrix_element)
            replace_dict['model_parameters'] = model_parameter_lines

            # Extract color data lines
            color_matrix_lines = self.get_color_matrix_lines(matrix_element)
            replace_dict['color_matrix_lines'] = \
                                               "\n        ".join(color_matrix_lines)



            # Extract JAMP lines
            jamp_lines = self.get_jamp_lines(matrix_element)
            replace_dict['jamp_lines'] = "\n        ".join(jamp_lines)

            # Extract amp2 lines
            amp2_lines = self.get_amp2_lines(matrix_element,
                                        self.config_maps.setdefault(ime, []))
            replace_dict['amp2_lines'] = '\n        '.join(amp2_lines)

            method_file = open(os.path.join(_file_path, \
                       'iolibs/template_files/matrix_method_python.inc')).read()
            method_file = method_file % replace_dict

            self.matrix_methods[process_string] = method_file

        return self.matrix_methods

    def get_helicity_matrix(self, matrix_element):
        """Return the Helicity matrix definition lines for this matrix element"""

        helicity_line = "helicities = [ \\\n        "
        helicity_line_list = []

        for helicities in matrix_element.get_helicity_matrix():
            helicity_line_list.append("[" + ",".join(['%d'] * len(helicities)) % \
                                      tuple(helicities) + "]")
            
        return helicity_line + ",\n        ".join(helicity_line_list) + "]"


    def get_den_factor_line(self, matrix_element):
        """Return the denominator factor line for this matrix element"""

        return "denominator = %d" % \
               matrix_element.get_denominator_factor()

    def get_color_matrix_lines(self, matrix_element):
        """Return the color matrix definition lines for this matrix element. Split
        rows in chunks of size n."""

        if not matrix_element.get('color_matrix'):
            return ["denom = [1.]", "cf = [[1.]];"]
        else:
            color_denominators = matrix_element.get('color_matrix').\
                                                 get_line_denominators()
            denom_string = "denom = [%s];" % \
                           ",".join(["%i" % denom for denom in color_denominators])

            matrix_strings = []
            my_cs = color.ColorString()
            for index, denominator in enumerate(color_denominators):
                # Then write the numerators for the matrix elements
                num_list = matrix_element.get('color_matrix').\
                                            get_line_numerators(index, denominator)

                matrix_strings.append("%s" % \
                                     ",".join(["%d" % i for i in num_list]))
            matrix_string = "cf = [[" + \
                            "],\n        [".join(matrix_strings) + "]];"
            return [denom_string, matrix_string]

    def get_jamp_lines(self, matrix_element):
        """Return the jamp = sum(fermionfactor * amp[i]) lines"""

        res_list = []

        for i, coeff_list in enumerate(matrix_element.get_color_amplitudes()):

            res = "jamp[%d] = " % i

            # Optimization: if all contributions to that color basis element have
            # the same coefficient (up to a sign), put it in front
            list_fracs = [abs(coefficient[0][1]) for coefficient in coeff_list]
            common_factor = False
            diff_fracs = misc.make_unique(list_fracs)
            if len(diff_fracs) == 1 and abs(diff_fracs[0]) != 1:
                common_factor = True
                global_factor = diff_fracs[0]
                res = res + '%s(' % coeff(1, global_factor, False, 0)

            for (coefficient, amp_number) in coeff_list:
                if common_factor:
                    res = res + "%samp[%d]" % (coeff(coefficient[0],
                                               coefficient[1] / abs(coefficient[1]),
                                               coefficient[2],
                                               coefficient[3]),
                                               amp_number - 1)
                else:
                    res = res + "%samp[%d]" % (coeff(coefficient[0],
                                               coefficient[1],
                                               coefficient[2],
                                               coefficient[3]),
                                               amp_number - 1)

            if common_factor:
                res = res + ')'

            res_list.append(res)

        return res_list

    def get_amp2_lines(self, matrix_element, config_map = []):
        """Return the amp2(i) = sum(amp for diag(i))^2 lines"""

        ret_lines = []
        # Get minimum legs in a vertex
        
        vert_list = [max(diag.get_vertex_leg_numbers()) for diag in \
           matrix_element.get('diagrams') if diag.get_vertex_leg_numbers()!=[]]
        minvert = min(vert_list) if vert_list!=[] else 0

        if config_map:
            # In this case, we need to sum up all amplitudes that have
            # identical topologies, as given by the config_map (which
            # gives the topology/config for each of the diagrams
            diagrams = matrix_element.get('diagrams')
            # Combine the diagrams with identical topologies
            config_to_diag_dict = {}
            for idiag, diag in enumerate(matrix_element.get('diagrams')):
                if config_map[idiag] == 0:
                    continue
                try:
                    config_to_diag_dict[config_map[idiag]].append(idiag)
                except KeyError:
                    config_to_diag_dict[config_map[idiag]] = [idiag]
            # Write out the AMP2s summing squares of amplitudes belonging
            # to eiher the same diagram or different diagrams with
            # identical propagator properties.  Note that we need to use
            # AMP2 number corresponding to the first diagram number used
            # for that AMP2.
            for config in config_to_diag_dict.keys():

                line = "self.amp2[%d]+=" % (config_to_diag_dict[config][0])

                line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
                                  {"num": a.get('number')-1} for a in \
                                  sum([diagrams[idiag].get('amplitudes') for \
                                       idiag in config_to_diag_dict[config]],
                                      [])])
                ret_lines.append(line)
            ret_lines.sort()
        else:
            wf_dict = {}
            vx_list = []
            optimization = 0
            for idiag, diag in enumerate(matrix_element.get('diagrams')):
                # Ignore any diagrams with 4-particle vertices.
                if diag.get_vertex_leg_numbers()!=[] and \
                                   max(diag.get_vertex_leg_numbers()) > minvert:
                    continue
                # Now write out the expression for AMP2, meaning the sum of
                # squared amplitudes belonging to the same diagram
                line = "self.amp2[%d]+=" % (idiag)
                line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
                                  {"num": a.get('number')-1} for a in \
                                  diag.get('amplitudes')])
                ret_lines.append(line)

        return ret_lines

    def get_mg5_info_lines(self):
        """Return info lines for MG5, suitable to place at beginning of
        Python files"""

        info = misc.get_pkg_info()
        info_lines = ""
        if info and 'version' in info and  'date' in info:
            info_lines = "#  MadGraph5_aMC@NLO v. %s, %s\n" % \
                         (info['version'], info['date'])
            info_lines = info_lines + \
                         "        #  By the MadGraph5_aMC@NLO Development Team\n" + \
                         "        #  Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"
        else:
            info_lines = "        #  by MadGraph5_aMC@NLO\n" + \
                         "        #  By the MadGraph5_aMC@NLO Development Team\n" + \
                         "        #  Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"        

        return info_lines

    def get_process_info_lines(self, matrix_element):
        """Return info lines describing the processes for this matrix element"""

        return"\n        ".join([ "# " + process.nice_string().replace('\n', '\n# * ') \
                         for process in matrix_element.get('processes')])


    def get_model_parameter_lines(self, matrix_element):
        """Return definitions for all model parameters used in this
        matrix element"""

        # Get all masses and widths used
        if aloha.complex_mass:
            parameters = [(wf.get('mass') == 'ZERO' or wf.get('width')=='ZERO') 
                          and wf.get('mass') or 'CMASS_%s' % wf.get('mass') 
                          for wf in \
                          matrix_element.get_all_wavefunctions()]
            parameters += [wf.get('mass') for wf in \
                      matrix_element.get_all_wavefunctions()]
        else:
            parameters = [wf.get('mass') for wf in \
                      matrix_element.get_all_wavefunctions()]
        parameters += [wf.get('width') for wf in \
                       matrix_element.get_all_wavefunctions()]
        parameters = misc.make_unique(parameters)
        if 'ZERO' in parameters:
            parameters.remove('ZERO')

        # Get all couplings used

        
        couplings = misc.make_unique([c.replace('-', '') for func \
                              in matrix_element.get_all_wavefunctions() + \
                              matrix_element.get_all_amplitudes() for c in func.get('coupling')
                              if func.get('mothers') ])
        
        return "\n        ".join([\
                         "%(param)s = model.get(\'parameter_dict\')[\"%(param)s\"]"\
                         % {"param": param} for param in sorted(parameters)]) + \
               "\n        " + "\n        ".join([\
                         "%(coup)s = model.get(\'coupling_dict\')[\"%(coup)s\"]"\
                              % {"coup": coup} for coup in sorted(couplings)])

#===============================================================================
# Global helper methods
#===============================================================================

def coeff(ff_number, frac, is_imaginary, Nc_power, Nc_value=3):
    """Returns a nicely formatted string for the coefficients in JAMP lines"""

    total_coeff = ff_number * frac * fractions.Fraction(Nc_value) ** Nc_power

    if total_coeff == 1:
        if is_imaginary:
            return '+complex(0,1)*'
        else:
            return '+'
    elif total_coeff == -1:
        if is_imaginary:
            return '-complex(0,1)*'
        else:
            return '-'

    res_str = '%+i.' % total_coeff.numerator

    if total_coeff.denominator != 1:
        # Check if total_coeff is an integer
        res_str = res_str + '/%i.' % total_coeff.denominator

    if is_imaginary:
        res_str = res_str + '*complex(0,1)'

    return res_str + '*'

