##################################
#                                #
# Last modified 2023/08/03       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math
import gzip
import numpy as np

def run():

    if len(sys.argv) < 8:
        print 'usage: python %s plus.wig minus.wig chrom.sizes window step minCov(either) pseudoCount outputfilename [-wigoutput]' % sys.argv[0]
        sys.exit(1)

    plusWig = sys.argv[1]
    minusWig = sys.argv[2]

    chrominfo = sys.argv[3]
    chromInfoDict = {}
    linelist = open(chrominfo)
    for line in linelist:
        fields = line.strip().split('\t')
        chr = fields[0]
        start = 0
        end = int(fields[1])
        chromInfoDict[chr] = end

    W = int(sys.argv[4])
    step = int(sys.argv[5])
    minCov = float(sys.argv[6])
    PseudoCount = float(sys.argv[7])
    outfilename = sys.argv[8]

    doWigOutput = False
    if '-wigoutput' in sys.argv:
        print 'will output bedGraph/wig'
        doWigOutput = True

    print W, step, minCov, PseudoCount

    ValuesDictPlus = {}
    ValuesDictMinus = {}
    if plusWig.endswith('.gz'):
        linelist = gzip.open(plusWig)
    else:
        linelist = open(plusWig)
    for line in linelist:
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        chr = fields[0]
        left = int(fields[1])
        right = int(fields[2])
        score = math.fabs(float(fields[3]))
        if ValuesDictPlus.has_key(chr):
            pass
        else:
            ValuesDictPlus[chr] = {}
            ValuesDictMinus[chr] = {}
            print chr
        for i in range(left,right):
            ValuesDictPlus[chr][i] = score

    if minusWig.endswith('.gz'):
        linelist = gzip.open(minusWig)
    else:
        linelist = open(minusWig)
    for line in linelist:
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        chr = fields[0]
        left = int(fields[1])
        right = int(fields[2])
        score = math.fabs(float(fields[3]))
        if ValuesDictMinus.has_key(chr):
            pass
        else:
            ValuesDictMinus[chr] = {}
            ValuesDictPlus[chr] = {}
            print chr
        for i in range(left,right):
            ValuesDictMinus[chr][i] = score

    outfile = open(outfilename,'w')
    if not doWigOutput:
         outlines = 'chr\tpos\tplus/minus_ratio\tplus\tminus\tplus+minus'

    for chr in chromInfoDict.keys():
        for i in range(int(W/2.),chromInfoDict[chr] - int(W/2.),step):
            plusSum = 0.0
            minusSum = 0.0
            for j in range(i - int(W/2.), i + int(W/2.)):
                if ValuesDictPlus[chr].has_key(j):
                    plusSum += ValuesDictPlus[chr][j]
                if ValuesDictMinus[chr].has_key(j):
                    minusSum += ValuesDictMinus[chr][j]
            plusAv = plusSum/W
            minusAv = minusSum/W
            if max(plusAv,minusAv) < minCov:
                continue
            else:
                ratio = (plusAv + PseudoCount)/(minusAv + PseudoCount)
            if doWigOutput:
#                print plusAv, minusAv, ratio
                outline = chr + '\t' + str(max(i-step/2,0) + '\t' + str(min(i+step/2,chromInfoDict[chr])) + '\t' + str(math.log(ratio,2))
                outfile.write(outline + '\n')
            else:
                outline = chr + '\t' + str(i) + '\t' + str(math.log(ratio,2))
                outline = outline + '\t' + str(plusAv)
                outline = outline + '\t' + str(minusAv)
                outline = outline + '\t' + str(plusAv + minusAv)
                outfile.write(outline + '\n')

    outfile.close()
   
run()
