#! /usr/bin/env python

"""\
Usage: %prog IN.slha [OUT.isa]

Convert an SLHA model/decay spectrum file to a HERWIG/ISAWIG SUSY spectrum input
file, in the format written by ISAWIG. The intention of this script is to allow
HERWIG to be used with spectra generated by codes other than ISAWIG,
e.g. SoftSUSY, SPheno, etc. And fundamentally because the author was fed up with
having to make ISASUSY compile with gfortran on platforms where CERNLIB wasn't
available!

Conversion based on the HERWIG SUSY specification format, from
http://www.hep.phy.cam.ac.uk/~richardn/HERWIG/ISAWIG/file.html
"""

__author__ = "Andy Buckley <andy.buckley@cern.ch"


## PDG MC ID codes mapped to HERWIG SUSY ID codes, based on
## http://www.hep.phy.cam.ac.uk/~richardn/HERWIG/ISAWIG/susycodes.html
PDGID2HERWIGID = {}
PDGID2HERWIGID[      25] = 203 ## HIGGSL0 (ADDED)
PDGID2HERWIGID[      26] = 203 ## HIGGSL0
PDGID2HERWIGID[      35] = 204 ## HIGGSH0
PDGID2HERWIGID[      36] = 205 ## HIGGSA0
PDGID2HERWIGID[      37] = 206 ## HIGGS+
PDGID2HERWIGID[     -37] = 207 ## HIGGS-
PDGID2HERWIGID[ 1000001] = 401 ## SSDLBR
PDGID2HERWIGID[-1000001] = 407 ## SSDLBR
PDGID2HERWIGID[ 1000002] = 402 ## SSULBR
PDGID2HERWIGID[-1000002] = 408 ## SSUL
PDGID2HERWIGID[ 1000003] = 403 ## SSSLBR
PDGID2HERWIGID[-1000003] = 409 ## SSSL
PDGID2HERWIGID[ 1000004] = 404 ## SSCLBR
PDGID2HERWIGID[-1000004] = 410 ## SSCL
PDGID2HERWIGID[ 1000005] = 405 ## SSB1BR
PDGID2HERWIGID[-1000005] = 411 ## SSB1
PDGID2HERWIGID[ 1000006] = 406 ## SST1BR
PDGID2HERWIGID[-1000006] = 412 ## SST1
PDGID2HERWIGID[ 2000001] = 413 ## SSDR
PDGID2HERWIGID[-2000001] = 419 ## SSDRBR
PDGID2HERWIGID[ 2000002] = 414 ## SSUR
PDGID2HERWIGID[-2000002] = 420 ## SSURBR
PDGID2HERWIGID[ 2000003] = 415 ## SSSR
PDGID2HERWIGID[-2000003] = 421 ## SSSRBR
PDGID2HERWIGID[ 2000004] = 416 ## SSCR
PDGID2HERWIGID[-2000004] = 422 ## SSCRBR
PDGID2HERWIGID[ 2000005] = 417 ## SSB2
PDGID2HERWIGID[-2000005] = 423 ## SSB2BR
PDGID2HERWIGID[ 2000006] = 418 ## SST2
PDGID2HERWIGID[-2000006] = 424 ## SST2BR
PDGID2HERWIGID[ 1000011] = 425 ## SSEL-
PDGID2HERWIGID[-1000011] = 431 ## SSEL+
PDGID2HERWIGID[ 1000012] = 426 ## SSNUEL
PDGID2HERWIGID[-1000012] = 432 ## SSNUELBR
PDGID2HERWIGID[ 1000013] = 427 ## SSMUL-
PDGID2HERWIGID[-1000013] = 433 ## SSMUL+
PDGID2HERWIGID[ 1000014] = 428 ## SSNUMUL
PDGID2HERWIGID[-1000014] = 434 ## SSNUMLBR
PDGID2HERWIGID[ 1000015] = 429 ## SSTAU1-
PDGID2HERWIGID[-1000015] = 435 ## SSTAU1+
PDGID2HERWIGID[ 1000016] = 430 ## SSNUTL
PDGID2HERWIGID[-1000016] = 436 ## SSNUTLBR
PDGID2HERWIGID[ 2000011] = 437 ## SSEL-
PDGID2HERWIGID[-2000011] = 443 ## SSEL+
PDGID2HERWIGID[ 2000012] = 438 ## SSNUEL
PDGID2HERWIGID[-2000012] = 444 ## SSNUELBR
PDGID2HERWIGID[ 2000013] = 439 ## SSMUL-
PDGID2HERWIGID[-2000013] = 445 ## SSMUL+
PDGID2HERWIGID[ 2000014] = 440 ## SSNUMUL
PDGID2HERWIGID[-2000014] = 446 ## SSNUMLBR
PDGID2HERWIGID[ 2000015] = 441 ## SSTAU1-
PDGID2HERWIGID[-2000015] = 447 ## SSTAU1+
PDGID2HERWIGID[ 2000016] = 442 ## SSNUTL
PDGID2HERWIGID[-2000016] = 448 ## SSNUTLBR
PDGID2HERWIGID[ 1000021] = 449 ## GLUINO
PDGID2HERWIGID[ 1000022] = 450 ## NTLINO1
PDGID2HERWIGID[ 1000023] = 451 ## NTLINO2
PDGID2HERWIGID[ 1000025] = 452 ## NTLINO3
PDGID2HERWIGID[ 1000035] = 453 ## NTLINO4
PDGID2HERWIGID[ 1000024] = 454 ## CHGINO1+
PDGID2HERWIGID[-1000024] = 456 ## CHGINO1-
PDGID2HERWIGID[ 1000037] = 455 ## CHGINO2+
PDGID2HERWIGID[-1000037] = 457 ## CHGINO2-
PDGID2HERWIGID[ 1000039] = 458 ## GRAVTINO



import sys, optparse
parser = optparse.OptionParser(usage=__doc__)
opts, args = parser.parse_args()
if len(args) < 1 or len(args) > 2:
    parser.print_help()
    sys.exit(1)

## Choose output file
import os
o = os.path.basename(args[0])
if "." in o:
    o = o[:o.rindex(".")]
opts.OUTFILE = o + ".isa"
if len(args) == 2:
    opts.OUTFILE = args[1]


## Read spectrum file
import pyslha
BLOCKS, DECAYS = pyslha.readSLHAFile(args[0])
MASSES = BLOCKS["MASS"].entries
# PIDS = MASSES.keys()


## Init output string
out = ""

## First write out masses section:
##   Number of SUSY + top particles
##   IDHW, RMASS(IDHW), RLTIM(IDHW)
##   repeated for each particle
## IDHW is the HERWIG identity code.
## RMASS and RTLIM are the mass in GeV, and lifetime in seconds respectively.
massout = ""
for pid in MASSES.keys():
    lifetime = -1
    try:
        width = DECAYS[pid].totalwidth
        if width and width > 0:
            lifetime = 1.0/(width * 1.51926778e24) ## lifetime in seconds == hbar/width in GeV
    except:
        pass
    massout += "%d %e %e\n" % (PDGID2HERWIGID.get(pid, pid), MASSES[pid], lifetime)
out += "%d\n" % massout.count("\n")
out += massout


assert(len(MASSES) == len(DECAYS))


## Next each particles decay modes together with their branching ratios and matrix element codes
##   Number of decay modes for a given particle (IDK)
##     IDK(*), BRFRAC(*), NME(*) & IDKPRD(1-5,*)
##     repeated for each mode.
##   Repeated for each particle.
## IDK is the HERWIG code for the decaying particle, BRFRAC is the branching ratio of
## the decay mode. NME is a code for the matrix element to be used, either from the
## SUSY elements or the main HERWIG MEs. IDKPRD are the HERWIG identity codes of the decay products.
for i, pid in enumerate(DECAYS.keys()):
    # if not DECAYS.has_key(pid):
    #     continue
    decayout = "%d\n" % len(DECAYS[pid].decays)
    hwid = PDGID2HERWIGID.get(pid, pid)
    #decayout += "@@@@ %d %d %d\n" % (i, pid, hwid)
    for i_d, d in enumerate(DECAYS[pid].decays):
        ## TODO: Identify decay matrix element to use
        #decayout += str(i_d + 1) + " "
        decayout += "%d %e 0 " % (hwid, d.br)
        ## TODO: Order decay products as required
        ids = [0,0,0,0,0]
        for i, pid in enumerate(d.ids):
            ids[i] = pid
        ids = map(str, ids)
        decayout += " ".join(ids) + "\n"
    out += decayout

## Now the SUSY parameters
## TANB, ALPHAH:
out += "%e %e\n" % (BLOCKS["MINPAR"].entries[3], BLOCKS["ALPHA"].entries)
## Neutralino mixing matrix
nmix = BLOCKS["NMIX"].entries
for i in xrange(1, 5):
    out += "%e %e %e %e\n" % (nmix[i][1], nmix[i][2], nmix[i][3], nmix[i][4])
## Chargino mixing matrices V and U
vmix = BLOCKS["VMIX"].entries
out += "%e %e %e %e\n" % (vmix[1][1], vmix[1][2], vmix[2][1], vmix[2][2])
umix = BLOCKS["UMIX"].entries
out += "%e %e %e %e\n" % (umix[1][1], umix[1][2], umix[2][1], umix[2][2])
# THETAT,THETAB,THETAL
import math
out += "%e %e %e\n" % (math.acos(BLOCKS["STOPMIX"].entries[1][1]),
                       math.acos(BLOCKS["SBOTMIX"].entries[1][1]),
                       math.acos(BLOCKS["STAUMIX"].entries[1][1]))
# ATSS,ABSS,ALSS
out += "%e %e %e\n" % (BLOCKS["AU"].entries[3][3],
                       BLOCKS["AD"].entries[3][3],
                       BLOCKS["AE"].entries[3][3])
# MUSS == sign(mu)
out += "%f\n" % BLOCKS["MINPAR"].entries[4]


## TODO: Handle RPV SUSY


## Write it out
f = sys.stdout
if opts.OUTFILE != "-":
    f = open(opts.OUTFILE, "w")
f.write(out)
if f is not sys.stdout:
    f.close()
