"""
From a marker mapped spreadsheet with rows pairs as samples and columns as
marker mapped allele intensities for sample pairs in successive rows, that is as
quantile normalized A and B signal intensities for each sample, prompt the user
for an output directory and export the allele signals for all chromosomes
suitable for BEAGLECALL. Note BEAGLECALL does not work on male X or on Y.

The input spreadsheet is generated from the Affymetrix CEL import if you choose
to import intermediate steps. You must marker map the quantile normalized SNP
spreadsheet before running this script and transpose it.


Author: Christophe Lambert
Last Modified: January 20, 2009
"""
ghi.requireVersion('7.4')
import os
def exportBeagleCallAlleleSignals(ss,dm,outputDir):
    
    mmap = dm.markerMapChromosomes()
    numRows = dm.numRows()
    numCols = dm.numCols()
    chrdict = {}
    
    for x in mmap:
        if not chrdict.has_key(x):
            chrdict[x] = open(outputDir + "/chr" + x + ".signals","w")
    
    #Write file headers
    for x in chrdict:
        chrdict[x].write("marker\talleleA\tAlleleB")
        for sample in dm.rowLabels():
            chrdict[x].write("\t" + sample)
        chrdict[x].write("\n")  
  
    progress = ghi.progressDialog('Writing BEAGLECALL Allele Signals Files...', numCols)
    progress.show()
     
    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:
            if x == None:
                x = "NaN"
            outfile.write("\t" + str(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.FilterReal|ghi.const.FilterMapped) 
        ok = True
    except:
        ghi.message("Spreadsheet must have active real-valued\ncolumns containing allele intensity data.")
        
    if ok:
        outputDirIn = ghi.chooseDirectory("Choose directory for output files")
        if len(outputDirIn) != 0:
            exportBeagleCallAlleleSignals(ssIn,dmIn,outputDirIn)

except:
    ghi.error()