from string import atoi,atof,split
from output import viewTRIALS
import scitbx.array_family.flex

def underline(s):
  line = ''
  for i in s:
    line += '='
  print s
  print line

def frame(s):
  line = ''
  for i in s:
    line += '='
  print line
  print s
  print line


def solution_selector(s):
  nList = []
  qList = []
  try:
    raw1 = split(s,',')
    for element in raw1:
      raw2 = split(element,'-')
      if len(raw2) == 1:
        if '!' in raw2[0]:
          raw3 = split(raw2[0],'!')
          for x in raw3:
            try:
              qList.append(atoi(x))
            except:
              pass
        else:
          i_e = atoi(raw2[0])

          nList.append(i_e)
      else:
        if len(raw2) == 2:
          for i_e in xrange(atoi(raw2[0]),atoi(raw2[1])+1):
            nList.append(i_e)
        else:
          pass
    for element in qList:
      while element in nList:
        nList.remove(element)
    return nList

  except:
    print "Bad format"
    raise "Bad format"

def addense(i,dictENSEMBLES):
  for ensembleID in dictENSEMBLES.keys():
    PDBList = dictENSEMBLES[ensembleID]
    numberPDB = len(PDBList)
    print "\nENSEMBLE:", ensembleID, ",",numberPDB,"PDB file(s)"
    for i_PDB in xrange(numberPDB):
      pdb_file_name = PDBList[i_PDB][0]
      rms_seq = PDBList[i_PDB][1]
      rms_seq_value = PDBList[i_PDB][2]
      if rms_seq is None: # 2011-01-29 Read from REMARK records
        i.addENSE_PDB_CARD(ensembleID,pdb_file_name, True)
      elif rms_seq == 'RMS':
        i.addENSE_PDB_RMS(ensembleID,pdb_file_name,rms_seq_value)
      else:
        i.addENSE_PDB_ID(ensembleID,pdb_file_name,rms_seq_value)

def setCell(i,localLIST):
  i.setSPAC_HALL(localLIST[0])
  cell = localLIST[1]
  i.setCELL(cell[0],cell[1],cell[2],cell[3],cell[4],cell[5])

def setMTZ(i,localLIST):
  setCell(i,localLIST)
  i.setREFL(localLIST[4],localLIST[2],localLIST[3])
  i.addPROT(localLIST[5], 1)

def mtzit(i,localLIST):
  setCell(i,localLIST)
  i.setREFL(localLIST[4],localLIST[2],localLIST[3])
  i.addPROT(localLIST[5],1)

def selectKNOWN(i,criterion,selection,setLIST):
  nList = []
  if criterion == 'Solution Number' and selection is not None:
    nList = solution_selector(selection)
    
  if len(setLIST)>0:
    solutions = setLIST[0]
    likelihood = setLIST[1]
    trials = []
    for sols in xrange(len(solutions)):
      if selection is not None:
        if criterion == 'Likelihood Cutoff' and likelihood[sols] > atof(selection):
          trials.append(solutions[sols])
        if criterion == 'Solution Number' and int(sols)+1 in nList:
          trials.append(solutions[sols])
      if selection is None:
        trials.append(solutions[sols])
    if len(trials)<1:
      print "\nWARNING! NO SOLUTIONS SELECTED\n"

    print viewTRIALS(trials)
    print
    i.setSOLU(trials)

def selectTRIAL(i,criterion,selection,setLIST):
  nList = []
  if criterion == 'Solution Number' and selection is not None:
    nList = solution_selector(selection)
  if len(setLIST)>0:
    solutions = setLIST[3]
    likelihood = setLIST[1]
    mask = []
    count = 1
    for sols in xrange(len(solutions)):
      rlis = solutions[sols].RLIST
      sel = []
      for j in xrange(len(rlis)):
        if selection is not None:
          if criterion == 'Likelihood Cutoff':
            score = rlis[j].SCORE
            if score > atof(selection):
              sel.append(True)
            else:
              sel.append(False)
          if criterion == 'Solution Number':
            if count in nList:
              sel.append(True)
            else:
              sel.append(False)
            count += 1
        else:
          sel.append(True)
      mask.append(scitbx.array_family.flex.bool(sel))
    print mask
    print solutions
    i.setAndSelectSOLU(solutions,mask)
# AJM this bit needs rewriting
#void SOLU::setAndSelectSOLU(af::shared<mr_solution> mrset,af::shared<bool_array > selected)
#{
#  if (mrset.size() != selected.size())
#  CCP4base::store(InputError(INPUT,keywords,"MR set and selection flag not same size"));
#  for (int i = 0; i < mrset.size(); i++)
#    if (mrset[i].RLIST.size() != selected[i].size())
#    CCP4base::store(InputError(INPUT,keywords,"MR set and selection flag for rlist not same size"));
#
#  //if any one of the rlists are selected, then set is also selected
#  bool1D set_selected(selected.size(),false);
#  for (int i = 0; i < mrset.size(); i++)
#  {
#    //if it is a solution only, then it is selected
#    if (!mrset[i].RLIST.size()) set_selected[i] = true;
#    for (int r = 0; r < mrset[i].RLIST.size(); r++)
#      if (selected[i][r]) set_selected[i] = true;
#  }
#
#  //add the MRSETS with selected RLISTS to
#  for (int i = 0; i < mrset.size(); i++)
#    if (set_selected[i]) MRSET.push_back(mrset[i]);
#
#  //blank the RLISTS and re-add the ones that are selected
#  for (int i = 0; i < MRSET.size(); i++)
#    MRSET[i].RLIST.resize(0);
#
#  int set_number(0);
#  for (int i = 0; i < mrset.size(); i++)
#  if (set_selected[i])
#  {
#    for (int r = 0; r < mrset[i].RLIST.size(); r++)
#      if (selected[i][r]) MRSET[set_number].RLIST.push_back(mrset[i].RLIST[r]);
#    set_number++;
#  }
#
#}




if __name__ == '__main__':
  s = "1,2,3-6,7,9 - 15, !12, 2-4, !2"
  print solution_selector(s)
  underline("Hello")