#! /usr/bin/python

import os,sys,glob,errno,shutil,time,fnmatch    #argparse 
from optparse import OptionParser

#  helper to replace all sourceText in fileName with replaceText
def replacetext(fileName, sourceText, replaceText):
    file = open(fileName, "r") 
    text = file.read() 
    file.close()
    file = open(fileName, "w")
    file.write(text.replace(sourceText, replaceText))
    file.close() 

#  helper to build recursivly path
def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc: # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else: raise   
      
#  helper to find all files of with name in path
def find(name, path):
    for root, dirs, files in os.walk(path):
        if name in files:
            return os.path.join(root, name)      
 
#  helper to find all file paths which contain file names matching filepattern
def finddirs(filepattern, path):
    founddirs = []
    for root, dirs, files in os.walk(path):
        if fnmatch.filter(files, filepattern):
            founddirs.append(root)
    return founddirs
 
# fill the proc.dat file from BornAmplitudes.dat and VirtAmplitudes.dat.
def fillprocs(model,oras,orew):
  bornlist=[]
  virtlist=[]
  fileproc=open("proc.dat","w")
  fileproc.write("set fortran_compiler gfortran --no_save\n")
  fileproc.write("import model "+model+"\n")
  borns="BornAmplitudes.dat"
  virts="VirtAmplitudes.dat"
  first=True
  procnr=0
  virtlines=""
  bornlines=""
  


    
  minlegs=100
  legs=0
  for i in [borns, virts]:
    file = open(i, "r")
    for line in file:
      if (len(line.split(" "))<minlegs):
        minlegs=len(line.split(" "))
        for it in line.split(" "):
	  if it.replace("-","").isdigit():
	    legs+=1
      
    file.close()
  
   
    
  #conversion for heft model to go from (2QCD+1QED)->1HIG for each FS HIGGS.
  HIG=0
  if (model=="heft"):
    HIG=(int(oras)+int(orew)-legs+2)/2
    
    if (int(oras)+int(orew)-legs+2)%2!=0:
      print "Warning: No possible coupling power:(int(oras)+int(orew)-legs+2)%2!=0 "
      exit()
      return

  
  
  
  
  file = open(borns, "r")
  for line in file:
      #this assumes extra QCD emmissions
      addalphas=len(line.split(" "))-minlegs
      linetmp=line.rstrip()
      procnr+=1
      bornlist+=[str(procnr)]
      if first:
	if HIG ==0 :
	  bornlines+="generate "+linetmp+" QCD="+str(int(oras)+addalphas)+" QED="+str(orew)+" @"+str(procnr)+"\n"
	else:
	  bornlines+="generate "+linetmp+" HIG="+str(HIG)+" QCD="+str(int(oras)+addalphas-2*HIG)+" QED="+str(int(orew)-HIG)+" @"+str(procnr)+"\n"
	first=False
      else:
	if HIG ==0 :
	  bornlines+="add process "+linetmp+" QCD="+str(int(oras)+addalphas)+" QED="+str(orew)+" @"+str(procnr)+"\n"
	else:
	  bornlines+="add process "+linetmp+" HIG="+str(HIG)+" QCD="+str(int(oras)+addalphas-2*HIG)+" QED="+str(int(orew)-HIG)+" @"+str(procnr)+"\n"
	
  file.close()
  first=True
  file = open(virts, "r")
  for line in file:
      addalphas=len(line.split(" "))-minlegs
      linetmp=line.rstrip()+" QCD="+str(int(oras)+addalphas)+"  QED="+str(int(orew))+" [ virt=QCD ]"
      procnr+=1
      virtlist+=[str(procnr)]
      if first:
	virtlines+="generate "+linetmp+"  @"+str(procnr)+"\n"
	first=False
      else:
	virtlines+="add process "+linetmp+"  @"+str(procnr)+"\n"
  file.close()

  fileproc.write(bornlines)
  if virtlines!="" and bornlines!="":
     fileproc.write("output matchbox MG5 --postpone_model\n") 
  fileproc.write(virtlines)
  fileproc.write("output matchbox MG5 -f\n")
  fileproc.close()
  return bornlist,virtlist

def build_matchbox_tmp(pwd,buildpath,absolute_links):
   cwd=os.getcwd() 
   os.chdir(pwd)
   mkdir_p(pwd+"/Herwig-scratch/MG_tmp/")
   if not buildpath.startswith("/"):
     buildpath=pwd+"/"+buildpath.lstrip("./")
   if not buildpath.endswith("/"):
     buildpath=buildpath + "/"
   resources=glob.glob(buildpath +"MG5/SubProcesses/MadLoop5_resources/*")
   resources+=glob.glob(buildpath +"MG5/Cards/*")
   resources+=glob.glob(buildpath +"MG5/Cards/SubProcesses/*")
   for i in resources:
     if not os.path.isfile( pwd+"/Herwig-scratch/MG_tmp/"+os.path.basename(i)) \
       and not os.path.islink( pwd+"/Herwig-scratch/MG_tmp/"+os.path.basename(i)):
         if not absolute_links:
             source=os.path.dirname(i)
             dest=pwd+"/Herwig-scratch/MG_tmp/"
             os.chdir(dest)
             os.symlink(os.path.relpath(source,dest)+"/"+os.path.basename(i),"./" + os.path.basename(i))
         else:
             os.symlink(i, pwd+"/Herwig-scratch/MG_tmp/"+os.path.basename(i))
   os.chdir(cwd)    


parser = OptionParser()
parser.add_option("-a", "--buildpath", dest="buildpath",help="Do not use this script. Only for Herwig internal use. ")
parser.add_option("-b", "--build", action="store_true", dest="build", default=True,help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-c", "--madgraph", dest="madgraph",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-d", "--runpath", dest="runpath",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-e", "--model", dest="model",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-f", "--orderas", dest="orderas",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-g", "--orderew", dest="orderew",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-i", "--datadir",dest="datadir",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-I", "--includedir",dest="includedir",help="Do not use this script. Only for Herwig internal use.")
parser.add_option("-l", "--absolute-links",action="store_true", dest="absolute_links", default=False,\
                  help="Do not use this script. Only for Herwig internal use.")


(options, args) = parser.parse_args()






#parser = argparse.ArgumentParser()
#parser.add_argument('--buildpath', help='installpath')
#parser.add_argument('--build', help='build', action="store_true")
#parser.add_argument('--madgraph', help='madgraph_installpath')
#parser.add_argument('--runpath', help='runpath')
#parser.add_argument('--model', help='model')
#parser.add_argument('--orderas', help='orderas')
#parser.add_argument('--orderew', help='orderew')
#parser.add_argument('--datadir', help='datadir')
#args = parser.parse_args()






pwd=os.getcwd() 


param_card=""
mkdir_p(pwd+"/Herwig-scratch/MG_tmp/")






if options.model=="loop_sm" or options.model=="heft":

  if options.model=="loop_sm":
    param_card="param_card.dat"
  else:
    param_card="param_card_"+options.model+".dat"  
    
  file = open("%s/MadGraphInterface/%s.in" % (options.datadir,param_card) , "r") 
  paramcard = file.read() 
  file.close()
  file = open(options.runpath+"/"+param_card, "w")

  params=open(options.runpath+"/MG-Parameter.dat", "r")

  for line in params:
    a=line.split()
    paramcard=paramcard.replace(a[0],a[1])
  params.close()
  file.write(paramcard)
  file.close()   
elif  options.model.startswith("/"):
  
  os.system("python %s/write_param_card.py " % options.model)
  
  
  
else:
  print "---------------------------------------------------------------"
  print "---------------------------------------------------------------"
  print "Warning: The model set for the MadGraph Interface "
  print "         needs a parameter setting by hand."
  print "         Please fill the param_card_"+options.model+".dat"
  print "         with your favourite assumptions."
  print "         And make sure Herwig uses the same parameters."
  print "---------------------------------------------------------------"
  print "---------------------------------------------------------------"
  if os.path.isfile(options.buildpath +"/MG5/Cards/param_card.dat") and not os.path.isfile(options.runpath+"/"+"param_card_"+options.model+".dat"):
    shutil.copyfile(options.buildpath +"/MG5/Cards/param_card.dat", options.runpath+"/"+"param_card_"+options.model+".dat")
  time.sleep(1)

if not os.path.isdir(options.buildpath):
   print "The MadGraph Install path was not existend. It has been created for you."
   print "Just start Herwig read again.."
   mkdir_p(options.buildpath)
   exit()

os.chdir(options.buildpath)
if os.path.isfile("InterfaceMadGraph.so"):
  build_matchbox_tmp(pwd,options.buildpath,options.absolute_links)
  exit()
 
Bornlist,Virtlist=fillprocs(options.model,options.orderas,options.orderew)

if not options.madgraph and not os.path.isfile("InterfaceMadGraph.so"):
  print "*** warning *** MadGraph build failed, check logfile for details"
  print "Known issue: If this is your first NLO calculation with pure Madgraph Amplitudes"
  print "             the CutTools compilation can result in a non usable configuration."
  print "             Please open $HERWIG_ENV/opt/madgraph/vendor/CutTools/makefile"
  print "             and add a \\ after FC=gfortran... in the ARGS variable."
  print "             Then run make clean && make in the CutTools folder. " 
  exit()

os.system("python "+options.madgraph+"/mg5_aMC proc.dat")


routines=[["","BORN(momenta,hel)"],
	 ["","SLOOPMATRIX(momenta,virt)"],
	 ["","GET_JAMP(color,Jamp)"],
	 ["","GET_LNJAMP(color,Jamp)"],
	 ["","GET_NCOL(color)"],
	 ["","GET_NCOLOR(i,j,color)"]]

for routine in routines:
  for i in Bornlist + list(set(Virtlist) - set(Bornlist)):
    if routine[1]=="Virt(amp)" or routine[1]=="SLOOPMATRIX(momenta,virt)" and i not in  Virtlist:
      continue
    if routine[0]=="":
       routine[0]+="         SELECT CASE (proc) \n"
       routine[0]+="         CASE("+i+") \n            CALL "
       routine[0]+=   "MG5_"+i+"_"+routine[1]+"\n"
    else:
       routine[0]+="         CASE("+i+") \n"\
               "            CALL "
       routine[0]+=   "MG5_"+i+"_"+routine[1]+"\n"
  if routine[0]!="":	     
       routine[0]+="         CASE DEFAULT\n"
       routine[0]+="             WRITE(*,*) '##W02A WARNING No id found '\n"
       routine[0]+="         END SELECT    \n"
 
shutil.copyfile("%s/MadGraphInterface/InterfaceMadGraph.f.in" % options.datadir, "InterfaceMadGraph.f")

replacetext("InterfaceMadGraph.f","MG_CalculateBORNtxt",routines[0][0])
replacetext("InterfaceMadGraph.f","MG_CalculateVIRTtxt",routines[1][0])
replacetext("InterfaceMadGraph.f","MG_Jamptxt",     routines[2][0])
replacetext("InterfaceMadGraph.f","MG_LNJamptxt",     routines[3][0])
replacetext("InterfaceMadGraph.f","MG_NColtxt",     routines[4][0])
replacetext("InterfaceMadGraph.f","MG_ColourMattxt",routines[5][0])

MG_vxxxxxtxt=""
if routines[1][0]!="":
  MG_vxxxxxtxt="""       subroutine  MG_vxxxxx(p, n,inc,VC)
     $   bind(c, name='MG_vxxxxx')
         IMPLICIT NONE
         double precision p(0:3)
         double precision n(0:3)
         INTEGER inc
         double precision VC(0:7)
         double complex  VCtmp(8)
         call vxxxxx(p, 0d0,1,inc ,VCtmp)
         VC(0)= real(VCtmp(5))
         VC(1)=aimag(VCtmp(5))  
         VC(2)= real(VCtmp(6))
         VC(3)=aimag(VCtmp(6))  
         VC(4)= real(VCtmp(7))
         VC(5)=aimag(VCtmp(7))  
         VC(6)= real(VCtmp(8))
         VC(7)=aimag(VCtmp(8))  
       END"""
else:
  MG_vxxxxxtxt="""       subroutine  MG_vxxxxx(p, n,inc,VC)
     $   bind(c, name='MG_vxxxxx')
         IMPLICIT NONE
         double precision p(0:3)
         double precision n(0:3)
         INTEGER inc
         double precision VC(0:7)
         double complex  VCtmp(6)
         call vxxxxx(p, 0d0,1,inc ,VCtmp)
         VC(0)= real(VCtmp(3))
         VC(1)=aimag(VCtmp(3))  
         VC(2)= real(VCtmp(4))
         VC(3)=aimag(VCtmp(4))  
         VC(4)= real(VCtmp(5))
         VC(5)=aimag(VCtmp(5))  
         VC(6)= real(VCtmp(6))
         VC(7)=aimag(VCtmp(6))  
       END"""
replacetext("InterfaceMadGraph.f","MG_vxxxxxtxt",MG_vxxxxxtxt)


make=" "
fortanfiles=glob.glob('*/*/*.f')+glob.glob('*/*/*/*.f')

for i in fortanfiles:
  if "check_sa" not in i and "f2py_wrapper" not in i:
    if not os.path.islink(i):
       make += " "+i+"\\\n                  "

incfiles=glob.glob('*/*/*.inc')+glob.glob('*/*/*/*.inc')

coefdir=""
for i in incfiles:
  if "nexternal.inc" in i:
   coefdir+=" -I"+i.replace("nexternal.inc"," ")


file=open("makefile","w")
file.write("include MG5/Source/make_opts  ")
if Virtlist!=[]:

    collierlib=options.madgraph
    collierlib=collierlib.replace("bin","HEPTools/collier")
    onelooplib=options.madgraph
    onelooplib=onelooplib.replace("bin","HEPTools/oneloop")

    file.write("\nLIBDIR = MG5/lib")
    file.write("\nLINKLIBS =-L$(LIBDIR) -lcts -liregi ")
    if (os.path.exists("MG5/lib/ninja_lib")):
	file.write("\nLINKLIBS += -L$(LIBDIR)/ninja_lib -lninja   ")
    if (os.path.exists(collierlib)):
        file.write("\nLINKLIBS += -L%s -lcollier "%(collierlib))
    if (os.path.exists(onelooplib)):
        file.write("\nLINKLIBS += -L%s -lavh_olo "%(onelooplib))
    if (os.path.exists("MG5/lib/golem95_lib")):
        file.write("\nLINKLIBS += -L$(LIBDIR)/golem95_lib -lgolem") 
     
    
file.write("\nPROCESS= InterfaceMadGraph.f "+make+"\n\nall:  \n\t gfortran  -g -O2 -w -fbounds-check -ffixed-line-length-132 -fPIC -fno-f2c -shared -s -o  InterfaceMadGraph.so -IMG5/SubProcesses/" )
if Virtlist!=[]:
  if (os.path.exists("MG5/lib/golem95_include")):
      file.write(" -IMG5/lib/golem95_include ")
  if (os.path.exists("MG5/lib/ninja_include")):
      file.write(" -IMG5/lib/ninja_include ")


  incldir=options.madgraph.replace("bin","HEPTools/include")
  
  file.write(" -I%s -I./ "%incldir)
  # Find all .mod files also in /usr/include if golem was build there.
  # There can be an error message in the MadGraph output to add the golem include path to the makefiles.
  # Usually MadGraph finds the path if its Golem was build in an separate dictionary.
  # Our bootstrap script installs golem with gosam beside boost. Here MadGraph creates a  link (->errormessage).
  # If we can find the modfiles easily the user doesn't need to change the makefiles. 
  moddirs=finddirs('*.mod',options.includedir)
  for moddir in moddirs:
    file.write(" -I%s " % moddir)
  if os.path.isdir("/usr/include"):
    moddirs=finddirs('*.mod',"/usr/include")
    for moddir in moddirs:
      file.write(" -I%s " % moddir)

if coefdir != "":
   file.write(coefdir)
file.write("   $(PROCESS) $(LINKLIBS) ")
file.close()


os.chdir(pwd)
os.chdir(options.buildpath)
replacetext("MG5/Source/MODEL/lha_read.f", "ident_card.dat","Herwig-scratch/MG_tmp/ident_card.dat")
replacetext("MG5/Source/MODEL/lha_read.f", "param.log","Herwig-scratch/MG_tmp/param.log")
if Virtlist!=[]:
  replacetext("MG5/SubProcesses/MadLoopCommons.f", "PREFIX='./'","PREFIX='./Herwig-scratch/MG_tmp/'")

os.system("make")
if not os.path.isfile("InterfaceMadGraph.so"):
   print "Second trial to make MadGraph Interface. "
   print "Needed if new .mod files are produced by make." 
   os.system("make")

build_matchbox_tmp(pwd,options.buildpath,options.absolute_links)

