#!/usr/bin/env python2.7

""" - Main program code for "feyncop" - A program to calculate the coproduct 
on the Hopf algebra of Feynman graphs. """

# See also: http://people.physik.hu-berlin.de/~borinsky/

__author__ = "Michael Borinsky"
__email__ = "borinsky@physik.hu-berlin.de"
__copyright__ = "Copyright (C) 2014 Michael Borinsky"
__license__ = "MIT License"
__version__ = "1.0"

# Copyright (c) 2014 Michael Borinsky

# This program is distributed under the MIT License:

# Permission is hereby granted, free of charge, to any person obtaining a copy 
# of this software and associated documentation files (the "Software"), to 
# deal in the Software without restriction, including without limitation the 
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 
# sell copies of the Software, and to permit persons to whom the Software is 
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in 
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 
# IN THE SOFTWARE.

# If you have problems using this program, 
# please feel free to contact me!

from fractions import *
import copy
import argparse
import sys
import collections
from math import *
from stuff import *

import parsefg, outputfg

from graph import Graph
from hopf_graph import HopfGraph

def get_sub_graph_str( sub_graphs ):
    """Creates an output string for a list of subgraphs."""
    sub_graph_strs = [ "{" + ",".join([ "%d" % e for e in sg ] ) + "}" for sg in sub_graphs ]
    return "{" + ",".join( sub_graph_strs ) + "}"

def get_coproduct_str( g, dimension, ym ):
    """Creates an output string for the coproduct."""
    sum_str = []
    for sub_graphs in g.reduced_coproduct(dimension, ym):
        sub_graph_str = get_sub_graph_str( sub_graphs )

        sum_str+= [ "%s" % (sub_graph_str) ]

    return "D[%s,\n[%s]]" % ( g.get_graph_str(ym), ", ".join(sum_str) )

def get_coproduct_with_residues_str( g, dimension, ym ):
    """Creates an output string for the coproduct with cographs."""

    sum_str = []
    for sub_graphs, residue_graph in g.reduced_coproduct_with_residue_graph(dimension, ym):
        sub_graph_str = get_sub_graph_str( sub_graphs )

        sum_str+= [ "[%s,%s]" % (sub_graph_str, residue_graph.get_graph_str(ym)) ]

    return "D[%s,[\n%s]\n]" % ( g.get_graph_str(ym), ", \n".join(sum_str) )

def get_tensor_product_str( tp, ym ):
    """Creates an output string for the coproduct in the unlabeled tensor 
        product form."""

    sum_str = []
    sub_graphs, residue_graph = tp
    def nosymfac( g ):
        gp = copy.copy(g)
        gp.prepare_graph()
        gp.symmetry_factor = 0
        return gp
    def get_sg_str( g, m ):
        if m == 1:
            return "%s" % nosymfac(g).get_graph_str(ym)
        else:
            return "(%s)^%d" % (nosymfac(g).get_graph_str(ym),m)

    sub_graphs_str = "*".join( [ get_sg_str(g, m) for g,m in sub_graphs ] )
    return "T[ %s, %s ]" % (sub_graphs_str, nosymfac(residue_graph).get_graph_str(ym))

def get_primitive_str( g, dimension, ym ):
    """Creates an output string for the graph it is primitive."""

    if g.is_primitive( dimension ):
        return g.get_graph_str(ym)
    else:
        return ""

def main():
    """Main program section. Reads the options and parameters and starts the 
        relevant subroutines."""

    parser = argparse.ArgumentParser(description='Calculate the reduced coproduct of elements of the Hopf algebra of Feynman graphs')
    parser.add_argument('-D', '--dimension', dest='dimension', type=int, default=4, help='Dimension of the theory (default: 4)')
    parser.add_argument('-c', '--cographs', dest='cographs', action='store_true', help='Also print the cographs corresponding to the subgraphs. (default: false)')
    parser.add_argument('-u', '--unlabeled', dest='unlabeled', action='store_true', help='Calculate the coproduct of a sum of graphs with unlabeled subgraphs and cographs. (default: false)')
    parser.add_argument('-v', '--unlab_per_graph', dest='unlab_per_graph', action='store_true', help='Calculate the coproduct of a sum of graphs with unlabeled subgraphs and cographs. And give the output as a formal power series in graphs. (default: false)')
    parser.add_argument('-p', '--primitives', dest='primitives', action='store_true', help='Filter graphs for primitive/skeleton graphs. (default: false)')

    parser.add_argument('-k', '--check', dest='check', action='store_true', help='Check the sum of 1PI graphs identity for a certain class of graphs (all graphs of a given grade). (default: false)')
    parser.add_argument('-m', '--check_tadpoles', dest='check_tadpoles', action='store_true', help='Check the sum of non-tadpole 1PI graphs identity for a certain class of graphs (all graphs of a given grade). (default: false)')
    args = parser.parse_args()
        
    string = ""
    try:
        string = next(sys.stdin)
    except StopIteration:
        return

    var_name,ev,el = parsefg.parse_var_name(string)
    if var_name:
        if args.unlabeled:
            sys.stdout.write("%s_red_cop_unlab%s" % (var_name, string[ev:]))
        elif args.primitives:
            sys.stdout.write("%s_proj_to_prim%s" % (var_name, string[ev:]))
        elif args.unlab_per_graph:
            sys.stdout.write("%s_red_cop_unlab_pg%s" % (var_name, string[ev:]))
        else:
            sys.stdout.write("%s_red_cop%s" % (var_name, string[ev:]))
        string = string[el:]

    if args.unlabeled or args.check or args.check_tadpoles:
        graphs_fac = frozenset((g.unlabeled_graph,fac,ym) for g,fac,ym in parsefg.parse_input_lines( sys.stdin, None, string ))

        ym_glob = False

        coproducts = dict()
        for g,fac,ym in graphs_fac:
            for sub_graphs, residue_graph in g.reduced_coproduct_unlabeled( args.dimension, ym ):
                tensor_product = (tuple(sorted(sub_graphs.items())), residue_graph)

                coproducts.setdefault( tensor_product, Fraction(0,1) ) 
                coproducts[tensor_product] += fac
                ym_glob = ym_glob or ym

        if args.unlabeled:
            for tp, fac in coproducts.items():
                tp_str = get_tensor_product_str( tp, ym_glob )
                print outputfg.get_element_str( tp_str, fac )
        elif args.check_tadpoles:
            tensor_check(coproducts, tadpoles=False)
            print "Check successful!"
            return
        else:
            tensor_check(coproducts)
            print "Check successful!"
            return
    else:
        for g, fac, ym in parsefg.parse_input_lines( sys.stdin, sys.stdout, string ):
            if args.primitives:
                prim_str = get_primitive_str( g, args.dimension, ym )
                if prim_str:
                    print outputfg.get_element_str( prim_str, fac )
            elif args.unlab_per_graph:
                tps = collections.Counter( (tuple(sorted(sub_graphs.items())), residue_graph) for  sub_graphs, residue_graph in g.reduced_coproduct_unlabeled( args.dimension, ym ) )
                tps_str = " + ".join( "%d*%s" % ( fac, get_tensor_product_str(e , ym )) if fac != 1 else get_tensor_product_str(e , ym ) for e, fac in tps.items()  )
                t_str = "%s*(%s)" % ( g.get_graph_str(ym), tps_str )

                print outputfg.get_element_str( t_str, fac )
            else:
                cop_str = ""
                if args.cographs:
                    cop_str = get_coproduct_with_residues_str( g, args.dimension, ym )
                else:
                    cop_str = get_coproduct_str( g, args.dimension, ym )
               
                print outputfg.get_element_str( cop_str, fac ) 

    print ";"

def tensor_check(coproducts, tadpoles=True):
    """Performs non trivial check on given sum of 1PI diagrams. """
    for tensor_product, factor in coproducts.items():
        sub_graphs, residue_graph = tensor_product
        if not tadpoles:
            if any(sg.is_tadpole for sg,m in sub_graphs) or residue_graph.is_tadpole:
                continue

        print " --- "
        print "\t",
        for g, fac in sub_graphs:
            print '%d x' % fac, g,
        if not sub_graphs:
            print "1 ",
        print ""
        print "\t", residue_graph

        sg_mults = [ m for g,m in sub_graphs ]
        res_structure = collections.Counter()
        sym_factor_sg = 1
        edges_permutations = 1
        calc_factor = Fraction(1,1)
        for g,m in sub_graphs:
            sym_factor_sg *= factorial(m) * g.symmetry_factor**m

            r = g.residue_type 
            res_cntr = collections.Counter( r )
            for s in res_cntr.values():
                edges_permutations *= factorial(s)**m
            
            res_structure[ r ] += m

        calc_factor = Fraction( edges_permutations, sym_factor_sg*residue_graph.symmetry_factor )
        for r,m in res_structure.items():
            if len(r) == 2:
                w1,w2 = r
                if abs(w1) != abs(w2):
                    raise

                w = abs(w1)

                insert_edges = sum( 1 for e in residue_graph.internal_edges_set if residue_graph.edge_weights[e] == w )
                if not residue_graph.internal_vtcs_set:
                    insert_edges = sum( 1 for e in residue_graph.external_edges_set if residue_graph.edge_weights[e] == w )
                insertion_places_edges = binomial( insert_edges + m - 1, m ) * factorial(m)

                calc_factor *= insertion_places_edges
            else:
                insert_vtcs = sum( 1 for v in residue_graph.internal_vtcs_set if residue_graph.get_vtx_type(v) == r )
                insertion_places_vtcs = binomial( insert_vtcs, m ) * factorial(m)

                calc_factor *= insertion_places_vtcs

        print ' ', calc_factor, factor, "******************************** BAD *************************" if calc_factor != factor else ""
        if calc_factor != factor: 
            print "Check failed!"
            print "Maybe not all 1PI graphs of a given loop number where piped into feyncop or the dimension is wrong?",
            print "Please make sure that feyncop with the -u option is used! Handling of leg-fixed graphs is not possible in this check."
            sys.exit()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print ""
        print "Cancelled"
