# ProdSubmitFile.py

################################################################################

class ProdSubmitFileException:
  def __init__(self, msg):
    self.msg = msg
  def __str__(self):
    return self.msg

################################################################################

class ProdSubmitRequest:
  def __init__(self, group, intask, tag, nevt, fmt, nevperjob):
    space = " "
    if tag[0] == "r":
      priority = 500
    elif tag[0] == "t":
      priority = 550
      evperjob = 10000
    else:
      raise ProdSubmitFileException("Invalid tag: " + tag)
    eol = "\n"
    self.group = group
    self.tag = tag
    self.priority = int(priority)
    self.events_per_job = nevperjob
    self.formats = fmt
    self.num_events = int(nevt)
    if isinstance(intask, str):
      self.input_task_name = intask
    else:
      self.input_task_name = intask.name
    self.esd_fraction = 0.0
    if fmt.find("ESD")>=0:
      self.esd_fraction = 1.0
    self.project_mode = "default"

  def is_valid(self):
    if not isinstance(self.group, str):
      self.error = 1
      return False
    if not isinstance(self.tag, str):
      self.error = 2
      return False
    if not isinstance(self.priority, int):
      self.error = 3
      return False
    if not isinstance(self.events_per_job, int):
      self.error = 4
      return False
    if not isinstance(self.formats, str):
      self.error = 5
      return False
    if not isinstance(self.num_events, int):
      self.error = 6
      return False
    if not isinstance(self.input_task_name, str):
      self.error = 7
      return False
    self.error = 0
    return True

  # Merge another request into this one.
  def merge(self, req):
    if not self.is_valid(): return -1
    if not req.is_valid(): return -2
    if req.group != self.group: return -3
    if req.tag != self.tag: return -4
    if req.priority != self.priority: return -5
    if req.events_per_job != self.events_per_job: return -6
    if req.input_task_name != self.input_task_name: return -8
    if req.formats == self.formats:
      self.num_events += req.num_events
      return 0
    if self.formats != "AOD" and self.formats!="ESD" and self.formats!="AOD.ESD": return -9
    if req.formats != "AOD" and req.formats!="ESD" and req.formats!="AOD.ESD": return -10
    num_events = self.num_events + req.num_events
    self.esd_fraction = (self.num_events*self.esd_fraction + req.num_events*req.esd_fraction)/num_events
    self.num_events = num_events
    self.formats = "AOD.ESD"
    return 0

  def __str__(self):
    sout = ""
    sout += "            group: " + self.group
    sout += "\n"
    sout += "              tag: " + self.tag
    sout += "\n"
    sout += "         priority: " + str(self.priority)
    sout += "\n"
    sout += "   events_per_job: " + str(self.events_per_job)
    sout += "\n"
    sout += "          formats: " + self.formats
    sout += "\n"
    sout += "       num_events: " + str(self.num_events)
    sout += "\n"
    sout += "  input_task_name: " + self.input_task_name
    return sout

  def line(self):
    space = " "
    line = ""
    smode = "default"
    if self.formats.find("ESD")>=0 and self.formats!="ESD":
      if self.esd_fraction<0.0 or self.esd_fraction>1.0:
        raise ProdSubmitFileException("Invalid ESD fraction: " + str(self.esd_fraction))
      smode = "fraction_ESD=" + str(int(100.0*self.esd_fraction+0.495))
    line += "project_mode:" + smode + space
    line += "group:" + self.group + space
    line += "tag:" + self.tag + space
    line += "priority:" + str(self.priority) + space
    line += "events_per_job:" + str(self.events_per_job) + space
    line += "formats:" + self.formats + space
    line += "num_events:"
    if self.num_events>0:
      line += str(self.num_events) + space
    else:
      line += "all" + space
    line += "ds:" + self.input_task_name
    return line
  

################################################################################

class ProdSubmitFile:

  # Ctor.
  #   fname - Output file name.
  #   project - Project name, e.g. "mc08".
  #   type - task types (reco or merge)
  def __init__(self, project, type):
    self.project = project
    self.type = type
    self.fname = None
    self.cat_reco_requests = True
    if not self.is_valid(): return
    self.lines = []
    self.lines += [ "# ProdSys submission file" ]
    self.lines += [ "project:" + project ]
    self.reqs = []

  # Check this is a valid instance.
  def is_valid(self):
    if self.project == "": return False
    if self.type == "": return False
    return True

  # Add a request.
  def submit(self, group, intask, tag, nevt, fmt, nevperjob):
    req = ProdSubmitRequest(group, intask, tag, nevt, fmt, nevperjob)
    if not req.is_valid():
      raise ProdSubmitFileException("Invalid request: error " + str(req.error))
    addit = True
    if len(self.reqs):
      addit = False
      lastreq = self.reqs[len(self.reqs)-1]
      notmerged = lastreq.merge(req)
      if notmerged:
        addit = True
        print "Merge stat: " + str(notmerged)
    if addit: self.reqs += [ req ]

  # Function to create a name for a list submission file.
  def make_file_name(self):
    submit_file_dir = "/afs/usatlas.bnl.gov/users/dladams/WWW/prodhist/submit_lists/"
    idfile_name = submit_file_dir + "id.txt"
    idfile = open(idfile_name, "r")
    sid = idfile.read().strip()
    idfile.close()
    id = int(sid)
    id += 1
    sid = str(id)
    idfile = open(idfile_name, "w")
    idfile.write(sid+"\n")
    idfile.close()
    fname = submit_file_dir + self.project + "." + sid + ".dla_"
    if self.type == "reco":
      fname += "reco.simul.LIST.r635"
    elif self.type == "merge":
      fname += "merge.recon.LIST.t53"
    else:
      raise ProdSubmitFileException("Invalid type: " + self.type)
    return fname

  # Write file.
  def write(self, name =None, verbose =0):
    self.fname = ""
    if name == None:
      fname = self.make_file_name()
    else:
      fname = name
    myname = "ProdSubmitFile:write: "
    if not self.is_valid():
      raise ProdSubmitFileException(myname + "Invalid instance")
    if len(self.reqs) == 0:
      if verbose>-1: print myname + "No requests found"
      return 1
    for req in self.reqs:
      if req.is_valid():
        self.lines += [ req.line() ]
      else:
        raise ProdSubmitFileException(myname + "Invalid request")
    outfile = open(fname, "w")
    eol = "\n"
    for line in self.lines:
      outfile.write(line + eol)
    if verbose>0: print myname + str(len(self.reqs)) + " requests written to " + self.fname
    self.fname = fname
    return 0