##################################
#                                #
# Last modified 7/22/2009         # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math
import numpy
from cistematic.core import Genome
from cistematic.core.geneinfo import geneinfoDB
from commoncode import *
from sets import Set


def run():

    if len(sys.argv) < 6:
        print 'usage: python %s genome wigfilename1 wigfilename2 ratio minimumRPM outfilename' % sys.argv[0]
        sys.exit(1)
    
    cachePages = 2000000

    genome = sys.argv[1]
    inputfilename1 = sys.argv[2]
    inputfilename2 = sys.argv[3]
    ratio = float(sys.argv[4])
    minRPM = float(sys.argv[5])
    outfilename = sys.argv[6]

    outfile = open(outfilename, 'w')

    hg = Genome(genome)

    wig = open(inputfilename1)
    linelist = wig.readlines()
    linelist.remove(linelist[0])
    print len(linelist)
    wig1Dict={}
    for line in linelist:
        if line[0:3]!='chr':
            continue
        fields=line.strip().split(' ')
        if float(fields[3])<minRPM:
            continue
        if wig1Dict.has_key(fields[0]):
            for k in range(int(fields[1]), int(fields[2])):
                wig1Dict[fields[0]][k]=float(fields[3])
        else:
            wig1Dict[fields[0]]={}
            for k in range(int(fields[1]), int(fields[2])):
                wig1Dict[fields[0]][k]=float(fields[3])
    print 'finished processing first wig file'
    wig = open(inputfilename2)
    linelist = wig.readlines()
    linelist.remove(linelist[0])
    print len(linelist)
    wig2Dict={}
    for line in linelist:
        if line[0:3]!='chr':
            continue
        fields=line.strip().split(' ')
        if len(fields)<4:
           continue
        if float(fields[3])<minRPM:
            continue
        if wig2Dict.has_key(fields[0]):
            for k in range(int(fields[1]), int(fields[2])):
                wig2Dict[fields[0]][k]=float(fields[3])
        else:
            wig2Dict[fields[0]]={}
            for k in range(int(fields[1]), int(fields[2])):
                wig2Dict[fields[0]][k]=float(fields[3])
    print 'finished processing second wig file'

    for chr in wig1Dict:
        if chr not in wig2Dict.keys():
            wig2Dict[chr]={}
    for chr in wig2Dict:
        if chr not in wig1Dict.keys():
            wig1Dict[chr]={}

    wig1up={}
    wig2up={}
    for chr in wig1Dict:
        wig1up[chr]={}
        wig2up[chr]={}
        commonkeys = Set.intersection(Set(wig1Dict[chr].keys()),Set(wig2Dict[chr].keys())) 
        print chr, len(wig1Dict[chr].keys()), len(wig2Dict[chr].keys()), 'common: ', len(commonkeys)
        for position in wig1Dict[chr].keys():
            if wig2Dict[chr].has_key(position):
                if wig2Dict[chr][position]/(wig1Dict[chr][position]+0.001)>ratio:
                    wig2up[chr][position]=wig2Dict[chr][position]/(wig1Dict[chr][position]+0.001)
                    continue
                if wig1Dict[chr][position]/(wig2Dict[chr][position]+0.001)>ratio:
                    wig1up[chr][position]=wig1Dict[chr][position]/(wig2Dict[chr][position]+0.001)
                    continue
    wig1Dict={}
    wig2Dict={}

    print 'finished comparison of position values'

    wig1upBPstats={}
    wig2upBPstats={}
    wig1upBPstats['A']=[]
    wig1upBPstats['G']=[]
    wig1upBPstats['T']=[]
    wig1upBPstats['C']=[]
    wig1upBPstats['N']=[]
    wig2upBPstats['A']=[]
    wig2upBPstats['G']=[]
    wig2upBPstats['T']=[]
    wig2upBPstats['C']=[]
    wig2upBPstats['N']=[]
    for chr in wig1up.keys():
        if 'random' in chr:
            continue
        chromosome=chr[3:len(chr)]
        length=len(hg.getChromosomeSequence(chromosome))
        print chromosome
        for position in wig1up[chr].keys():
            if position >=length:
                continue
            bp=hg.sequence(chromosome,position,1)
            bp=bp.upper()
            wig1upBPstats[bp].append(wig1up[chr][position])

    for chr in wig2up.keys():
        if 'random' in chr:
            continue
        chromosome=chr[3:len(chr)]
        length=len(hg.getChromosomeSequence(chromosome))
        print chromosome
        for position in wig2up[chr].keys():
            if position >=length:
                continue
            bp=hg.sequence(chromosome,position,1)
            bp=bp.upper()
            wig2upBPstats[bp].append(wig2up[chr][position])


    outfile.write(inputfilename1+' high value positions stats\n')
    outfile.write('Base\tNumber\tMeanRPM\tVariance\n')
    for bp in wig1upBPstats:
        outline=bp+'\t'+str(len(wig1upBPstats[bp]))+'\t'+str(numpy.mean(wig1upBPstats[bp]))+'\t'+str(numpy.var(wig1upBPstats[bp]))+'\n'
        outfile.write(outline)
    outfile.write('\n')
    outfile.write('\n')
    outfile.write(inputfilename2+' high value positions stats\n')
    outfile.write('Base\tNumber\tMeanRPM\tVariance\n')
    for bp in wig2upBPstats:
        outline=bp+'\t'+str(len(wig2upBPstats[bp]))+'\t'+str(numpy.mean(wig2upBPstats[bp]))+'\t'+str(numpy.var(wig2upBPstats[bp]))+'\n'
        outfile.write(outline)

    print 'finished outputing stats'

run()