# Time-stamp: <2008-06-03 15:30:49 Tao Liu>

"""Module Description: normalize based on GC content, and write
normalized data to *.normalized.txt files

"""

# ------------------------------------
# python modules
# ------------------------------------
import os
import math
import csv
from MA2C import GcBinsReader
from MA2C import SimpleMeanVariance
from MA2C import RobustMeanVariance
from MA2C import MeansPool
from MA2C import CsvFilter

# ------------------------------------
# Misc functions
# ------------------------------------

def _readTpmapFile(tpmapFileName):
    result = {}
    
    tpmapReader = csv.DictReader(CsvFilter.Filter(tpmapFileName), delimiter = '\t')
    
    for row in tpmapReader:
        x = int(row["X"])
        y = int(row["Y"])
        gc = int(row["GC"])
        
        result[(x, y)] = gc
    
    del tpmapReader
    
    return result

def _computeStd(v, m):
    std = 0.0
    
    for (x, y, n) in v:
        t = n - m
        std += t * t
        
    return math.sqrt(std / len(v))


def generate(tagFileInfo):
    
    tpmapCache = None
    
    tpmapFileName = tagFileInfo["output"]["tpmap"]
    
    chips = tagFileInfo["sample"]["chips"]

    for chipId in chips:
        rawFileName = tagFileInfo["output"]["chips"][chipId]["raw"]
        normalizedFileName = tagFileInfo["output"]["chips"][chipId]["normalized"]
        
#         if os.path.exists(normalizedFileName) \
#                 and os.stat(tpmapFileName).st_mtime < os.stat(normalizedFileName).st_mtime \
#                 and os.stat(rawFileName).st_mtime < os.stat(normalizedFileName).st_mtime \
#                 and os.stat(tagFileInfo["tagFile"]).st_mtime < os.stat(normalizedFileName).st_mtime:
#             continue
        
        if tpmapCache == None:
            tpmapCache = _readTpmapFile(tpmapFileName)
        
        (ipGcBin, inputGcBin) = GcBinsReader.read(tpmapFileName, rawFileName)
        
        (ipMean, inputMean, std) = ({}, {}, {})
        
        if tagFileInfo["normalization"]["method"] == "simple":
            (ipMean, inputMean, std) = SimpleMeanVariance.compute(ipGcBin, inputGcBin)
        elif tagFileInfo["normalization"]["method"] == "robust":
            (ipMean, inputMean, std) = RobustMeanVariance.compute(ipGcBin, inputGcBin, tagFileInfo["normalization"]["c"])
        else:
            raise Exception("Unknown normalization method %s!" % tagFileInfo["normalization"]["method"])
        
        MeansPool.pool(ipMean, inputMean, std)
        
        rawReader = csv.DictReader(CsvFilter.Filter(rawFileName), delimiter = '\t')
        
        normalizedData = []
        
        mean = 0.0
        
        for row in rawReader:
            x = int(row["X"])
            y = int(row["Y"])
            ip = float(row["IP"])
            input = float(row["INPUT"])
            
            ipSignal = math.log(ip, 2)
            inputSignal = math.log(input, 2)
            
            gc = tpmapCache[(x, y)]
            
            normalized = (ipSignal - inputSignal - ipMean[gc] + inputMean[gc]) / std[gc]
            
            mean += normalized
            
            normalizedData.append((x, y, normalized))
            
        std = _computeStd(normalizedData, mean / len(normalizedData))
        
        # writing data
        if not os.path.exists(tagFileInfo["output"]["dir"]):
            os.mkdir(tagFileInfo["output"]["dir"])

        f = open(normalizedFileName, "w")
        
        try:
            f.write("#DESIGN_ID\t%s\n" % tagFileInfo["sample"]["designId"])
            f.write("X\tY\tNormalizedLog2Ratio\n")
            
            for (x, y, n) in normalizedData:
                finalNormalized = n / std
            
                f.write("%d\t%d\t%f\n" % (x, y, finalNormalized))
        finally:
            f.close()
            
        del rawReader

        # write R script file
        raw_name = normalizedFileName.rstrip("txt")
        r_filename = raw_name+"r"
        pdf_filename = os.path.basename(raw_name)+"pdf"
        f = open(r_filename,"w")
        f.write("# R script generated by MA2C. Usage: $ R --vanilla < %s\n" % (os.path.basename(r_filename)))
        f.write("d <- read.table(\"%s\",header=T)\n" % (os.path.basename(normalizedFileName)))
        f.write("h <- hist(d$N,breaks=1000,plot=F)\n")
        f.write("pdf(\"%s\",width=8,height=6)\n" % (pdf_filename))
        f.write("plot(h$mid,h$counts,xlab=\"Normalized log-ratio\",ylab=\"Counts\",main=\"Histogram of %s\",type='l')\n" % (chipId))
        f.write("dev.off()")
    return True
