"""
From a marker mapped spreadsheet with rows as samples and columns as marker mapped SNPs,
export the BEAGLECALL Genotype Probabilities file. Assumes all columns are marker mapped
SNPs and does not drop inactive rows or columns. Assumes genotypes are in A_B form.

Split output out by chromosome.

Author: Christophe Lambert
Last Modified: 01/20/2011
"""
ghi.requireVersion('7.4')
import os
def exportBeagleCallGenotypeProbabilities(ss,dm,outputDir):

    mmap = dm.markerMapChromosomes()
    numRows = dm.numRows()
    numCols = dm.numCols()
    chrdict = {}
    print outputDir
    for x in mmap:
        if not chrdict.has_key(x):
            chrdict[x] = open(outputDir + "/chr" + x + ".gprobs","w")
    
    #Write file headers
    for x in chrdict:
        chrdict[x].write("marker\talleleA\tAlleleB")
        for sample in dm.col(0):
            chrdict[x].write("\t" + sample + "\t" + sample + "\t" + sample)
        chrdict[x].write("\n")  
  
    progress = ghi.progressDialog('Writing beagle Genotype Probability Files...', numCols)
    progress.show()
     
    d = {}
    d["?_?"] = "\t0.333\t0.333\t0.333"
    d["A_A"] = "\t1\t0\t0"
    d["A_B"] = "\t0\t1\t0"
    d["B_B"] = "\t0\t0\t1"        
    for i in range(1,numCols+1):
        if progress.wasCanceled():
            for x in chrdict:
                chrdict[x].close()
                os.remove(chrdict[x].name)
            progress.finish()
            return
        progress.setProgress(i)
        v = dm.col(i)
        outfile = chrdict[mmap[i-1]]
        outfile.write(dm.cell(0,i) + "\tA\tB")
        for x in v:
            outfile.write(d[x])
        outfile.write("\n")
    
    progress.finish()
    for x in chrdict:
        chrdict[x].close()

    message = str(len(chrdict)) + " files written to\n" + outputDir
    ss.appendLog(message)
    ghi.message(message)

try:
    ssIn = ghi.getCurrentObject()
    if not ssIn.hasMarkerMap() or ssIn.getMarkerMapOrientation() == ghi.const.MapOrientationRows:
        ghi.message("Spreadsheet must have a marker map applied to the columns")
    ok = False
    try:
        dmIn = ssIn.dataModel(ghi.const.FilterGenotypic|ghi.const.FilterMapped) 
        ok = True
    except:
        ghi.message("Spreadsheet must have active genotypic columns.")
        
    if ok:
        outputDirIn = ghi.chooseDirectory("Choose directory for output files")
        if len(outputDirIn) != 0:
            exportBeagleCallGenotypeProbabilities(ssIn,dmIn,outputDirIn)

except:
    ghi.error()