import os, re, glob, sys
import phaser.test_phaser_parsers
from phaser import matching
from phaser import template
from cctbx import sgtbx
from cctbx import uctbx
import scitbx.math
import iotbx.pdb
from phaser import CalcCCFromMRsolutions
from phaser.utils2 import *



# workaround until ListEdit.cc is fixed
def AreMatching(mtzfname, templsolfnames, mrsolfname, modelfname, mrresolution=2.5):
    mtz = iotbx.mtz.object(mtzfname)
    myparm = mtz.crystals()[0].unit_cell_parameters()
    ucell = uctbx.unit_cell(parameters = myparm)
    space_group_hall = mtz.space_group_info().group()
    crystal = matching.CrystalInfo( space_group = space_group_hall, cell=ucell)
    mroot = iotbx.pdb.input(modelfname).construct_hierarchy()
    ensemble = matching.EnsembleInfo.from_hierarchy( root = mroot )

    matches = []
    foundit = False
# allow reading an array of template solution files
    templsols = TryParseSolfile([templsolfnames[0]])
    if templsols == False:
        return (foundit, matches)

# having read the first template solution file now extend the mrsolution object with
# solutions from additional soution files
    for i, templsolfname in enumerate(templsolfnames):
        if i == 0:
            continue

        if not TryParseSolfile([templsolfname]):
            continue

        templparser = phaser.test_phaser_parsers.SolutionFileParser(open(templsolfname))
        othersols = templparser.get_dot_sol()
        for sol in othersols:
            templsols.append(sol)
# all template solutions from all template solution files are now stored in templsols
#
    reference = []
    for tmpls in templsols:
        for peak in tmpls.KNOWN:
            reference.append(
                matching.Labelled.phaser_mr_solution(
                    rotation = peak.getR(),
                    translation = peak.getFracT(),
                    crystal = crystal,
                    ensemble = ensemble,
                    dmin = mrresolution
                    )
                )

    ref_origins = [ matching.OriginSearchEntity( r ) for r in reference ]
#
# read the mr solution file
    mrsols = TryParseSolfile([mrsolfname])
    if mrsols == False:
        return (foundit, matches)

    for ref_origin in ref_origins:
        match = []
        for sol in mrsols:
            trial = [
                matching.Labelled.phaser_mr_solution(
                    rotation = peak.getR(),
                    translation = peak.getFracT(),
                    crystal = crystal,
                    ensemble = ensemble,
                    dmin = 2.5
                    )
                for peak in sol.KNOWN
                ]
# try matching each template solution with each mr solution
            m = matching.phaser_style_equivalence_match(
              left = [ref_origin],
              right = [ matching.OriginSearchEntity( t ) for t in trial ]
              )
            if m:
                foundit = True
            match.append( m)
        matches.append(match)
#
    return (foundit, matches)


def TryRead(rankedfiles):
    mstr = None
    for fname in rankedfiles:
        try:
            if os.path.exists(str(fname)):
                mstr = eval(open(fname, "r").read())
        except Exception:
            mstr = None
            continue

    return mstr



def TryParseSolfile(rankedsolfiles, usedfile=[]):
    ret = []
    try:
        for solfname in rankedsolfiles:
            if os.path.exists(str(solfname)):
                mrparser = phaser.test_phaser_parsers.SolutionFileParser(open(solfname))
                ret = mrparser.get_dot_sol()
                usedfile.append(solfname)
                break
    except Exception, m:
        ret = []

    finally:
        return ret



class BLASTMRutils():
    def __init__(self, topfolder, parent, completedrunfname="", mylock = None):
        self.topfolder = topfolder
        self.completedrunfname = completedrunfname
        self.mylock = mylock
        self.parent = parent
# array that'll hold MR solution, template solution, FTF_scrores data, VRMS and LLG float values
        self.mrsol = [None,None,None,None,None,None,None,None,None,[]]



    def IsthereaPartialSolution(self, modeldir, partmodels, solutionfname = ""):
# match the templatesolution folder which is named as an alphanumerically sorted pdbs
# Returns 1 if template = solution otherwise 0. Returns -1 if solution file doesn't exist
        sortedpartmodel = ",".join(sorted(partmodels.split(",")))
        sortedpartmodelname = "{" + sortedpartmodel + "}"
        modeldirname = "{" + modeldir + "}"

        templmodeldir = os.path.join(self.topfolder,"TemplateSolutions", sortedpartmodel)
        templunsortmodeldir = os.path.join(self.topfolder,"TemplateSolutions", modeldir)
        self.mrsol = [None,None,None,None,None,None,None,None,None,[]]

        workpath = os.path.join(self.topfolder, "MR", modeldir)
        workpath = os.path.relpath(workpath, self.topfolder)
        calclabelmodelsname = self.parent.calclabel + modeldirname + partmodels
        rootname = os.path.join(workpath, calclabelmodelsname)
# get LLG values from text files that holds float values rather than the int values in the sol files
        CorrCoeffname0 = os.path.join(workpath,self.parent.calclabel + "CalcCC_" + modeldirname + partmodels + ".txt")
        CorrCoeffname = os.path.join(workpath,"CalcCC_" + modeldirname + partmodels + ".txt")
        CorrCoeffname1 = os.path.join(workpath, self.parent.calclabel + "CalcCC_" + modeldirname + partmodels + ".txt")
        CorrCoeffname2 = os.path.join(workpath,"CalcCC_RNP_VRMS.txt")
        CorrCoeffname3 = os.path.join(workpath,"CalcCC__" + modeldirname + partmodels + ".txt")
        realLLGsVRMSfname = os.path.join(workpath, self.parent.calclabel + "LLG" + modeldirname +  partmodels + ".txt")
        #realLLGsVRMSfname1 = os.path.join(workpath,"LLG_RNP_VRMS" + modeldirname +  partmodels + ".txt")
        #realLLGsVRMSfname2= os.path.join(workpath,"LLGsFromRNP_VRMS.txt")
        #realLLGsVRMSfname3 = os.path.join(workpath,"LLG_RNP_VRMS.txt")
        CorrCoeffullresfname = os.path.join(workpath,"CalcCC_FullResRNP_VRMS_FullRes.txt")
        #realLLGsVRMSfullresfname = os.path.join(workpath,"LLG_FullRes_RNP_VRMS.txt")
        realLLGsfname1 = os.path.join(workpath,"LLGsFrom" + calclabelmodelsname + ".txt")
        realLLGsfname = os.path.join(workpath, self.parent.calclabel +"LLG" + modeldirname + partmodels + ".txt")
# override default value if supplied
        fname = os.path.join(workpath, solutionfname) if solutionfname != "" else ""
        fname0 = os.path.join(workpath, self.parent.calclabel + "_VRMS" + modeldirname +  partmodels + ".sol")
        fname1 = os.path.join(workpath, calclabelmodelsname + ".sol")
        #fname2 = os.path.join(workpath, self.parent.calclabel + "RNP_VRMS.sol")
        #fname3 = os.path.join(workpath, "RNP_VRMS" + modeldirname +  partmodels + ".sol")
        #fname4 = os.path.join(workpath, "RNP_VRMS.sol")
# allocate storage for the LLG float values
# 1st and 2nd are arrays for each MR solution. 3rd is the single value from RNP
# 4th and 5th are arrays for correlation coefficients
        self.mrsol[3] = [[None],[None],[None], None, None, None, None, [None]]
        #print "1.mrsol[3][6] = " +str(self.mrsol[3][6])
# store phaser solution in the 0th entry of the mrsol array
# initially solution files were given without vrms refinement.
# using RNP_VRMS.sol provides the vrms values in addition to the phaser solution
        mrsolfname = [None]
        self.mrsol[0] = TryParseSolfile([fname, fname0, fname1], mrsolfname)
        #if not self.mrsol[0]:
        #    return -1 # otherwise bail out

# open file with numbers of reflections
        FTFfname = rootname + ".FTF_Scores.txt"
        if os.path.exists(FTFfname) == False:
            fnames = glob.glob(os.path.join(workpath, "*.FTF_Scores.txt"))
            if len(fnames) > 0:
                FTFfname = fnames[0]

        if os.path.exists(FTFfname):
            FTFstr = open(FTFfname,"r").read()
# so we can calculate LLG per reflections
            if FTFstr != "":
                self.mrsol[2] = re.search("Number of selected reflections:\s+([0-9]+)", FTFstr).group(1)

# get ELLGs if present
        fname = os.path.join(workpath,"ELLG_*_%s.txt" %modeldir)
        ellgfname = glob.glob(fname)[0] if len(glob.glob(fname)) > 0 else ""
        ELLGtbl = ""
        if os.path.exists(ellgfname):
            ELLGtbl = open(ellgfname, "r").read()

        templfname = os.path.join(templmodeldir, self.parent.calclabel + "RNP" \
         + sortedpartmodelname + ".superposed.sol")
        templfname1 = os.path.join(templmodeldir, "RNP" + sortedpartmodelname \
         + ".superposed.sol")
        realtemplLLGfname = os.path.join(templmodeldir, "LLGsFromRNP"
         + sortedpartmodelname + ".txt")
        realtemplLLGfname1 = os.path.join(templmodeldir, self.parent.calclabel
         + "RNPLLG" + sortedpartmodelname + ".superposed.txt")

        templCorrCoeffname = os.path.join(templmodeldir, self.parent.calclabel + "RNPCalcCC_"
         + sortedpartmodelname + ".superposed.txt")
        templCorrCoeffname1 = os.path.join(templmodeldir, self.parent.calclabel + "RNPCalcCC_"
         + modeldirname + ".superposed.txt")

# read those LLG float values from the text files
        self.mrsol[3][0] = TryRead([realLLGsfname, realLLGsfname1])
        self.mrsol[3][1] = TryRead([realLLGsVRMSfname]) #, realLLGsVRMSfname1, realLLGsVRMSfname2, realLLGsVRMSfname3])
        #self.mrsol[3][2] = TryRead([realLLGsVRMSfullresfname])
        self.mrsol[3][3] = TryRead([realtemplLLGfname1]) #, realtemplLLGfname])
        self.mrsol[3][4] = TryRead([CorrCoeffname0]) #, CorrCoeffname, CorrCoeffname1, CorrCoeffname2, CorrCoeffname3])
        #self.mrsol[3][5] = TryRead([CorrCoeffullresfname])
        self.mrsol[3][6] = TryRead([templCorrCoeffname1, templCorrCoeffname])
        #print "2.mrsol[3][6] = " +str(self.mrsol[3][6])
        #print "fname= " + templCorrCoeffname1
        #print "mrsol[3][6] = " +str(self.mrsol[3][6])

        if self.mrsol[3][6] != None and sortedpartmodelname != modeldirname:
# If order of models in the template solution equals the reverse order of an existing file
# CCs, LLGs and MLADs in nonexisting file are inferred from the existing one
# by reversing the order of the elements in the array of the individual components
            if isinstance( self.mrsol[3][6][0][2], tuple ):
# using old format
# [[0.3576, 0.0, (0.3105, 0.203, 0.8105), (0.2222, 0.136, 1.0067)]]
                lst2 = self.mrsol[3][6][0][2:]
                lst2.reverse()
                lst3 = []
                lst3.append(self.mrsol[3][6][0][0])
                lst3.append(self.mrsol[3][6][0][1])
                for e in lst2:
                    lst3.append(e)
                self.mrsol[3][6] = [lst3]
            else:
# using new format which includes LLG values
# [[0.3576, 0.0, 89.639, (0.3105, 0.203, 0.8105, 45.367), (0.2222, 0.136, 1.0067, 23.764)],
                lst2 = self.mrsol[3][6][0][3:]
                lst2.reverse()
                lst3 = []
                lst3.append(self.mrsol[3][6][0][0])
                lst3.append(self.mrsol[3][6][0][1])
                lst3.append(self.mrsol[3][6][0][2])
                for e in lst2:
                    lst3.append(e)
                self.mrsol[3][6] = [lst3]

        #print "\n" + templCorrCoeffname1
        #print self.mrsol[3][6]

        usedfile=[]
        self.mrsol[1] = TryParseSolfile([templfname, templfname1], usedfile)
        #import code; code.interact(local=locals())
        #print "[templfname, templfname1] = " + str([templfname, templfname1])
        #print "usedfile = " + str(usedfile)
        #if len(self.mrsol[1]) < 1:
        #    return -1

# We want the vrms values for the refinement with the full set of reflections
# Get these by storing the solution of the this refinement into the 7th entry of the mrsol array
        solfname = os.path.join(workpath, self.parent.calclabel + "RNP_VRMS_FullRes.sol")
        self.mrsol[7] = TryParseSolfile([solfname])
# Want Fsol and Bsol if these were calculated
# Open the BsolFsolRNPcalc_1RLH_2D16_A.txt file. TargetPDBid isn't defined here so we use a wildcard
        #Fsolfname = os.path.join(workpath, "BsolFsolRNPcalc.sol")
        #self.mrsol[8] = TryParseSolfile([Fsolfname])
# create EnsembleDefinitionDict of model pdbs using ensemble names as keywords
# get the pdbfile names from the mrinput file
        infile = open(templmodeldir + os.path.sep + self.parent.calclabel + "RNPMRinput.txt","r")
        templinput = infile.read()
        infile.close()
        pdbs = re.findall("ENSEMBLE\s+(.*?)\s+PDB\s+(.*?)\s+IDENTITY\s+(.+[^\r])(\r*)",
          templinput, re.DOTALL|re.MULTILINE|re.IGNORECASE) # to get rid of any \r character
        ed = template.EnsembleDefinition()
        for i in range(0,len(pdbs)):
            modelfname = os.path.join(self.topfolder, pdbs[i][1])
            ed.define(pdbs[i][0],{modelfname : 100})
# get mtz file
        m = templinput.find("HKLIN")
        n = templinput.find("\n",m)
        mtzname = re.findall("HKLIN\s+(\S+)", templinput[m:n], re.MULTILINE)[0]
# define crystallographic matches from spacegroup info in mtz and appropriate resolutions
# for the template matching by reading the dimensions of the models
        fname = os.path.join(self.topfolder,mtzname)
        fname = PosixSlashFilePath(fname)
        fname = RecaseFilename2(fname)

        #em = template.EquivalentMatcher.from_mtz(fname)
# get all matched template solutions
        #solmatch = em.equivalents_of(templsols[0], self.mrsol[0], ed)
        #return len(solmatch)
        modelfname = PosixSlashFilePath(modelfname)
        matches = AreMatching(fname, [templfname], mrsolfname[0], modelfname)
# mrsol[9] is the ELLGtable values read as a string. Convert to a python list
        if ELLGtbl != "":
            for k, line in enumerate(ELLGtbl.split("\n")):
                if k == 0:
                    continue
                llst = line.split("\t")
                if len(llst) < 9:
                    continue
# if template solution was wrong evaluate correctness of ELLG solutions by also comparing with the mrsolution
                ellgsolfname = os.path.join(workpath,"ELLG_MRcalcs_%d.sol" %int(float(llst[0])))
                #mrresolution = float(llst[8])
                #ellgmatches = AreMatching(fname, [templfname, mrsolfname], ellgsolfname,
                #  modelfname, mrresolution)
                ellgmatches = AreMatching(fname, [templfname, mrsolfname[0]], ellgsolfname,
                  modelfname)
                #print "ellgmatches = " + str(ellgmatches)
# append result to llst
                llst.append(int(ellgmatches[0]))
                self.mrsol[9].append(llst)

        return matches[0]




    def NoticeCompletedPhaserRun(self,workdir, partmodels,
                     nPartSolWithThisStrategy):
        bPartSolWithThisStrategy = True
        if nPartSolWithThisStrategy < 1:
            bPartSolWithThisStrategy = False

        if self.mylock != None:
            self.mylock.acquire()

        completedruns = open(self.completedrunfname,"a")
        completedruns.write("%s {%s} %s\n" %(workdir, partmodels,
                     str(bPartSolWithThisStrategy)))
        completedruns.close()

        if self.mylock != None:
            self.mylock.release()



    def PhaserRanPartModelsBefore(self,founddir,partmodels,bsolved,doneruns=""):
        if self.completedrunfname != "" and os.path.exists(self.completedrunfname) == False:
            return False

        if self.mylock != None:
            self.mylock.acquire()

        if doneruns == "":
            completedruns = open(self.completedrunfname,"r")
            doneruns = completedruns.read()
            completedruns.close()
# Hack: to allow picking up template solutions of first component contrive these from the doneruns file
        firstruns = ""
        for e in doneruns.split("\n"):
            m = re.search("(.*)(\{.*\})(.*)", e, re.MULTILINE)
            if m is not None:
                for partmdl in m.group(1).strip().split(","):
                    firstruns += m.group(1) + "{" + partmdl + "}" + m.group(3)

        doneruns += firstruns

        if self.mylock != None:
            self.mylock.release()
# assuming the file looks like
# 1I1A_D,3GIZ_H,1NC2_A {1I1A_D} True
# 1W72_H,1NC2_A,1I1A_D {1W72_H} False
# 1I1A_D,1W72_H,1NC2_A {1I1A_D,1W72_H} True
#
        bsolved[0] = False
        m = doneruns.find("{" + partmodels + "}", 0)
        if m > 0:
# note what calculation was used for searching on the partial model
            n = doneruns.rfind("\n",0,m)
            if n == -1 and m > 0:
                n = 0  # found it in the very first line
            #print "partmodels= %s m= %d, n= %d, doneruns[n:m]= %s" % (partmodels,m,n,doneruns[n:m])
            founddir[0] = doneruns[n:m].split()[0]

            n = doneruns.find("\n",m)
        # find the bool value from "{1I1A_D} True"
            if doneruns[m:n].split(" ")[1] == "True":
                bsolved[0] = True
            else:
                bsolved[0] = False

            return True
        else:
            return False