#!/usr/bin/python
# -*- coding: utf-8 -*-

# cnv_tool.py is a program for calculation of Copy Number Variant Regions
# and checking the presence of CNVs in Database of Genomic Variants.

import sys, re, string, os, glob, math, time
from optparse import OptionParser, OptionGroup


class record(object):
  
  def __init__(self, rawline):
    
    # split line by tabs into list
    record_list = rawline.split("\t")
    
    self.rawline = rawline
    self.id_code = record_list[0]
    self.sample_nr = record_list[1]
    self.chromosome = int(record_list[2])
    self.start = int(record_list[3])
    self.end = int(record_list[4])
    self.length = record_list[5]
    self.cnv_type = record_list[6].strip()
    
    # these two are for calculating CNVRs
    if self.cnv_type=="0":
	self.cnv_types = cnv_type = {'0':1, '1':0, '3':0, '4':0}
    if self.cnv_type=="1":
	self.cnv_types = cnv_type = {'0':0, '1':1, '3':0, '4':0}
    if self.cnv_type=="3":
	self.cnv_types = cnv_type = {'0':0, '1':0, '3':1, '4':0}
    if self.cnv_type=="4":
	self.cnv_types = cnv_type = {'0':0, '1':0, '3':0, '4':1}
	
	
    self.n_cnvs = 1
    
    
  def __str__(self):
    return self.rawline

  __repr__ = __str__

  

# The function for parsing command line
def parse_commandline():
    usage = "%prog <cnv data> [options]"
    version = "1.4.1"
    description = \
        "%prog calculates CNV regions and checks CNVs against DBGV"
    epilog = \
        "Copyright (c) 2014 Svetlana Stoma"
    parser = OptionParser(usage=usage, description=description,
                     version=version, epilog=epilog)

    # options:
    
    parser.add_option("-c", "--cnv",  dest="cnvdata", metavar="<file>",
                     help="input data file of following format: ID code (kodas), sample nr (meginio nr), chromosome, start, end, length, copy number")
    parser.set_defaults(cnvdata=None)
    
    parser.add_option("-d", "--db", dest="dbgv", metavar="<file>",
                     help="database of genomic variants (default GRCh37_hg19_variants_2013-07-23.txt)")
    parser.set_defaults(dbgv="GRCh37_hg19_variants_2013-07-23.txt")
    
    parser.add_option("-f", "--find", "--check", dest="check_db", action="store_true",
                     help="find/check CNVs against Database of Genomic Variants")
    parser.set_defaults(check_db=False)
    
    parser.add_option("-g", "--genes", dest="genes", metavar="<file>",
                     help="file with UCSC genes from hg19: geneName, name, chrom, txStart, txEnd, status")
    parser.set_defaults(genes="ucsc_genes.txt")
    
    parser.add_option("-s", "--scan",  dest="scangenes", action="store_true",
                     help="for newly found CNVs check corresponding genes")
    parser.set_defaults(scangenes=False)
    
    parser.add_option("-r", "--regions",  dest="regions", action="store_true",
                     help="calculate regions of CNVs")
    parser.set_defaults(regions=False)
    
    parser.add_option("-e", "--cnvrs",  dest="cnvrs_file", type="str",
                     help="output file for calculated CNVRs")
    parser.set_defaults(cnvrs_file="cnvrs_"+time.strftime("%d_%m_%Y_")+time.strftime("%H_%M_%S")+".txt")

    parser.add_option("-o", "--output",  dest="output_file", type="str",
                     help="output file for CNVs checked against DBGV: output is divided into two files - inDB and notinDB CNVs")
    parser.set_defaults(output_file=time.strftime("%d_%m_%Y_")+time.strftime("%H_%M_%S")+".txt")

    parser.add_option("-a", "--cnvgenes",  dest="cnv_genes_file", type="str",
                     help="output file for genes that overlapped with given CNVs")
    parser.set_defaults(cnv_genes_file="cnv_genes_"+time.strftime("%d_%m_%Y_")+time.strftime("%H_%M_%S")+".txt")

    
    # get the options:
    (options, args) = parser.parse_args()

    if not options.cnvdata:
        print "No input file given. Exit."
        print ""
        parser.print_help()
        print ""
        exit(-1)

    if not options.check_db and not options.scangenes and not options.regions:
        
	print "Please check used options. Exit."
	print "Use -r (for finding CNVRs), -f (for DBGV search), or -s (for scanning for genes)"
	print ""
        parser.print_help()
        print ""
	exit(-1)
            
            
    # check for any leftover command line arguments:
    if len(args):
        warning("ignoring additional arguments "+str(args))
    
    # clean up (recommended):
    del(parser)
    return options


    
# This function checks whether CNVs on one chromosome overlap. And when
# they do overlap the function joins them into one greater CNV that
# represents CNVR.
def check_overlap(records):
	rec_list = records # coppying records just in case
	cnvrs = []
	to_remove = []
	
	#cnv_type = {'0':0, '1':0, '2':0, '3':0, '4':0}
	
	while len(rec_list):
		rec = rec_list.pop(0) # pop first record
		second_records=[]
		
		for r in rec_list:
			count_r=0
			# second cnv is in first one
			if rec.start <= r.start and rec.end >= r.end: 
				cnvr_start = rec.start
				cnvr_end = rec.end
				count_r=+1
			  
			# second cnv starts in the middle of the first  
			if rec.start <= r.start and rec.end <= r.end and rec.end >= r.start:
				cnvr_start = rec.start
				cnvr_end = r.end
				count_r=+1
			
			# first starts in the middle of the second
			if rec.start >= r.start and rec.end >= r.end and rec.start <= r.end: 
				cnvr_start = r.start
				cnvr_end = rec.end
				count_r=+1
			
			# first cnv is in the second
			if rec.start >= r.start and rec.end <= r.end: 
				cnvr_start = r.start
				cnvr_end = r.end
				count_r=+1
			  
			# if count_r is positive, then there was an overlap
			if count_r>0:
				# expand first cnv (that was popped out)
				rec.start = cnvr_start
				rec.end = cnvr_end
				# update CNVR
				rec.cnv_types[r.cnv_type]+=1
				rec.n_cnvs+=1
				
				# and remove one that overlapped with the first one
				to_remove.append(r)
			
		  
		# remove cnv records the overlapped with the first cnv record  
		for r_cnv in to_remove:
			if r_cnv in rec_list:
				rec_list.remove(r_cnv)
		  
		cnvrs.append(rec)

	return cnvrs
 

    
# The function finds CNVRs. CNV are grouped by chromosome, and 
# CNVRs are formed from overlapping CNVs or stan-alone CNV.
def find_cnvr(records, cnvrs_output):

	# make dictionary of chromosomes for grouping CNVs by chromosome
	chromosomes = {}
	    
	# fill in the dictionary
	for every_record in records:
		if every_record.chromosome not in chromosomes:
			chromosomes[every_record.chromosome] = [every_record]
		else:
			# to add another record to the list for that chromosome
			chromosomes[every_record.chromosome].append(every_record)
		
	# search for cnv regions on every chromosome
	chromosomes_list = chromosomes.keys()
	chromosomes_cnvrs = {}
	    
	if len(chromosomes)!=0:
	      
		for chromosome in chromosomes_list:
			ch_records = chromosomes[chromosome]	# list of records for this chromosome
		  
			# check if there are at least two overlapping regions
			# use overlapped regions to build CNVRs
			chromosomes_cnvrs[chromosome] = check_overlap(ch_records)
	      
		open(cnvrs_output, "wb").write("CNVRs\tchr\tstart\tend\tlength\tn_cnvs\tlosses\tdeletions\tgains\tduplications\n")
		
		for c in chromosomes_cnvrs:
			for r in chromosomes_cnvrs[c]:
				r1 = str(r).split()
				n = str(r.n_cnvs)
				
				losses = str(r.cnv_types['0'])
				deletions = str(r.cnv_types['1'])
				gains = str(r.cnv_types['3'])
				duplications = str(r.cnv_types['4'])
				
				# for each CNVR we write following info: chr, start, end, length
				out = "CNVR\t"+r1[2]+"\t"+r1[3]+"\t"+r1[4]+"\t"+str(abs(int(r1[4])-int(r1[3])))+"\t"+n+"\t"+losses+"\t"+deletions+"\t"+gains+"\t"+duplications+"\n"
				open(cnvrs_output, "a+b").write(out)



def search_dbgv(records, dbgv, genes, output_file) :
	#cnv_type = {'0': "Loss", '1': "Deletion", '2': "Norm", '3':"Gain", '4':"Duplication"}
	
	dbgv_list = open(dbgv).read().splitlines()
	dbgv_list.pop(0)  # removing the header
	
	# we will have two files: one for CNV records that are in DB, and other - not in DB
	cnvindb_output = "cnv_inDBGV_"+output_file
	open(cnvindb_output, "wb").write("code\tsample_nr\tchr\tstart\tend\tlength\tcnv_type"+"\t\t"+"variantaccession\tchr\tstart\tend\tvarianttype\tvariantsubtype\treference\tpubmedid\tmethod\tplatform\tmergedvariants\tsupportingvariants\tmergedorsample\tfrequency\tsamplesize\tobservedgains\tobservedlosses\tcohortdescription\tgenes\tsamples\n")
	
	cnvnotindb_output = "cnv_notinDBGV_"+output_file
	open(cnvnotindb_output, "wb").write("code\tsample_nr\tchr\tstart\tend\tlength\tcnv_type")
	
	
	not_in_db = []
	
	
	for record in records:
		
		if len(record.rawline):	# if not empty
		
			# don't write our CNV to the file of CNVs present in Db yet
			cnv_in_DB = False
			overlap_with_cnv_in_db = 0  
			  
			for variant in dbgv_list:
				variant_list = variant.split("\t")
					  
				# DB file heading looks like:
				# variantaccession	chr	start	end	varianttype	variantsubtype	reference	pubmedid	method	platform	mergedvariants	supportingvariants	mergedorsample	frequency	samplesize	observedgains	observedlosses	cohortdescription	genes	samples
				
				variant_chromosome = variant_list[1].split("_")[0]
				if variant_chromosome=="X" or variant_chromosome=="Y":
					variant_chromosome=23
				else:
					variant_chromosome = int(variant_chromosome)
				
				
				if record.chromosome==variant_chromosome:
					if (record.start >= int(variant_list[2]) and record.start <= int(variant_list[3])) or (record.end >= int(variant_list[2]) and record.end <= int(variant_list[3])) and variant_list[4]=="CNV":# and variant_list[5]==record_cnv_type:
						if cnv_in_DB==False:
							open(cnvindb_output, "a+b").write(record.rawline+"\n")
							cnv_in_DB=True
							
						# then that variant is in the DB
						overlap_with_cnv_in_db+=1
						
						# print to file of "overlapping CNVs"	
						open(cnvindb_output, "a+b").write("\t\t\t\t\t\t\t\t"+variant+"\n")
							
			# if out CNV did not overlap with any from DB, then it is a new one, and we save it in a file				
			if overlap_with_cnv_in_db==0:
				not_in_db.append(record)
				if options.scangenes==True:
					genes = ",".join(check_gene_coverage(record, genes))
					record_and_genes = record.rawline + "\t" + genes
					open(cnvnotindb_output, "a+b").write(record_and_genes+"\n")	# CHANGE HERE	
				else:
					open(cnvnotindb_output, "a+b").write(record.rawline+"\n")
      
      
      
# Function checks what genes overlap with given region (CNV)
def check_gene_coverage(cnv, genes):
	
	gene_lines = open(genes).read().splitlines()
	gene_lines.pop(0)
	
	cnv_genes=[]
	
	for gene_line in gene_lines:
		gene_record_list = gene_line.split("\t")
		gene_chr=0
		chrX=""
		chrY=""
		
		if gene_record_list[2][3:4]=="X":
			gene_chr = 23
			chrX="chrX,"
		  
		elif gene_record_list[2][3:4]=="Y":
			gene_chr = 23
			chrX="chrY,"
			
		else:
			g = gene_record_list[2].split("_")[0][3:]
			if g=="Un":
				pass
			else:
				gene_chr=int(g)
		
		if cnv.chromosome==gene_chr and (cnv.start >= int(gene_record_list[3]) and cnv.start <= int(gene_record_list[4])) or (cnv.end >= int(gene_record_list[3]) and cnv.end <= int(gene_record_list[4])):
			# then cnv overlaps with this gene
			
			# save gene for the cnv if it was not saved
			gene = gene_record_list[0]+"("+chrX+chrY+gene_record_list[5]+")"
			
			if gene not in cnv_genes:
				cnv_genes.append(gene)
		
	
	return cnv_genes
      
      

    
    

def main():
	# get command line options
	options = parse_commandline()
	print options
	
	# read file, where each row has followingcolumns:
	# ID code (kodas), sample nr (meginio nr), chromosome, start, end, length, copy number
	cnv_lines = open(options.cnvdata).read().splitlines()
	
	records = [record(i) for i in cnv_lines]
	

	# calculate CNVRs
	if options.regions==True:
		find_cnvr(records, options.cnvrs_file)
	

	# check with DB
	if options.check_db==True:
		search_dbgv(records, options.dbgv, options.genes, options.output_file)

	# if only scanning for genes is required
	if options.scangenes==True:
		
		open(options.cnv_genes_file, "wb").write("code\tsample_nr\tchr\tstart\tend\tlength\tcnv_type\tgenes\n")

		for cnv in records:
			genes = ",".join(check_gene_coverage(cnv, options.genes))
			record_and_genes = cnv.rawline + "\t" + genes
			open(options.cnv_genes_file, "a+b").write(record_and_genes+"\n")

    
           
if __name__ == "__main__":
    main()