(**
                     AddMomenta.m

(utility for CreateFeynAmp)
**)

BeginPackage["AddMomenta`","FeynArts`"]

AddMomenta::usage = " AddMomenta[ Topology[s][...] ] returns the \n
topology  with momenta p[i],q[i] appended to extern + intern lines \n
respectively. \n
AddMomenta on a TopologyList applies AddMomenta to every element."

Begin["`Private`"]

                          (** decide wether external or internal **)

ExternQ[ Propagator[ ty_ ][ __ ] ] := MemberQ[ { ex, inc , out } , ty ]

IncomingQ[ Propagator[ ty_ ][ __ ] ] := TrueQ[ ty===inc ]

InternQ[ Propagator[ ty_ ][ __ ] ] := FreeQ[ { ex, inc , out } , ty ]

	                               (** external propagators **)

NrOfInc[ t:Topology[___][__] ] := Length[List@@Select[t, IncomingQ ]]

momput1[ Propagator[ty_][ e[n_] , to_[m_] , a___ ] ] := 
	Propagator[ty][e[n],to[m],a,p[n]]

momput1[ Propagator[ty_][ to_[m_] , e[n_] , a___ ] ] := 
	Propagator[ty][to[m],e[n],a,-p[n]]

momput1[ Propagator[inc][ e[n_] , to_[m_] , a___ ], Nri_ ] := 
	Propagator[inc][e[n],to[m],a,p[n]]

momput1[ Propagator[inc][ to_[m_] , e[n_] , a___ ], Nri_ ] := 
	Propagator[inc][to[m],e[n],a,-p[n]]

momput1[ Propagator[out][ e[n_] , to_[m_] , a___ ], Nri_ ] := 
	Propagator[out][e[n],to[m],a,-k[n-Nri]]

momput1[ Propagator[out][ to_[m_] , e[n_] , a___ ], Nri_ ] := 
	Propagator[out][to[m],e[n],a,k[n-Nri]]

ext[ t:Topology[___][__] ] := 
      If[ NrOfInc[t]===0,
          momput1/@ (List@@Select[ t , ExternQ ]),
          momput1[ # , NrOfInc[t] ]& /@ (List@@Select[ t , ExternQ ])  ]

                                        (** internal propagators **)

momput2[ x_List ] :=  Apply[Append,#]& /@ 
		      Transpose[{x,Array[q,Length[x]]}]

int[ t:Topology[___][__] ] := momput2[ List@@ Select[ t , InternQ ] ]

                                            (** ext + int **)

AppendMomenta[ t:Topology[s___][__] ] := Join[ ext[t] , int[t] ]

                             (** construct the "momentum vertex" **)

mom[ver_, Propagator[_][ver_ , _ , ___ , m_]] := -m

mom[ver_, Propagator[_][ _ , ver_, ___ , m_]] := m

mom[ver_, Propagator[_][ver_,ver_, ___ , m_]] := {-m , m}

Vertices[ tt_List ] := Vertices[ Topology[1]@@ tt ]

Vertices[ tt:Topology[ elem__ ]  ] := Vertices[ Topology[1][elem] ]

Vertices[ tt:Topology[_][___] ] :=
  Block[{flat,three,four},
	 flat = Flatten[ tt/.{Topology[_]->List,
			      SpinorChain->List,
			      Propagator[_]->List} ];
         three = Cases[ flat, v3[_] ];
	 four = Cases[ flat, v4[_] ];
	 Union[ three,four ]
       ]

MakeVertices[t_List] := 
     Function[z, 
       Expand[ Flatten[ mom[z,#1]& /@ Select[t, MemberQ[#1,z]& ] ] ]
             ] /@  Vertices[t]  

        	       (** search momentum to solve equation **)

firstq[vert_List] := If[ FreeQ[vert/.List->Plus,q],
                         {},
                         Cases[ vert/.List->Plus, 
                                _.q[j_] -> q[j] ][[1]]
		       ]

                     (** solve equation for first momentum "q" **)

MomentumReplacement[topo_List , i_ ] := 
    If[ firstq[ topo[[i]] ]==={},
       { { p[1]->p[1] } },
       Solve[   Apply[ Plus,topo[[i]] ] == 0 , firstq[ topo[[i]] ]   ]
      ] 

              (** sum of all external momenta (=0) **)

momsum[ t:Topology[s___][__] ] := 
        Plus@@ (Join[
                Cases[t, Propagator[ex][___,e[n_],___] -> p[n] ],
                Cases[t, Propagator[inc][___,e[n_],___] -> p[n] ],
                Cases[t, Propagator[out][___,e[n_],___] -> 
                       -k[n-NrOfInc[t]] ]
               ])

(** insert momentum conservation for every vertex **)

PutMomenta[ topo_List , n_Integer?Positive ] := Block[{ newt },
       newt = PutMomenta[ topo , n-1 ];
       newt /. Flatten[ MomentumReplacement[ MakeVertices[newt], n ]]]

PutMomenta[ topo_List , 0 ] := topo

PutMomenta[ t:Topology[s___][__] ] := 
      ExpandAll[
      Topology[s]@@PutMomenta[ AppendMomenta[t] , Length[Vertices[t]] 
			   ] /. momsum[t]->0
      ]

PutMomenta[ t:Topology[ pr__ , v:Insertions[ ___ ] ]  ] :=
     Append[ Topology@@ PutMomenta[ Topology[1]@@ t ] , v ]

                  (** exported function : momentum conservation **)
	          (** and renumbering the internal momenta      **)

AddMomenta[ thetop_ ] := 
   Block[ {tt,theq,newq} ,
	 tt = PutMomenta[ thetop ];
         theq = Cases[ Flatten[List@@ tt/.Propagator[_]->List], q[_] ];
         newq = Array[ q , Length[theq] ];
	 tt /. Thread[ Rule[ theq , newq  ] ]
	]

AddMomenta[ ttt_TopologyList ] := AddMomenta /@ ttt
AddMomenta[ ttt:TopologyList[__][___] ] := AddMomenta /@ ttt

End[]

EndPackage[]

