import gzip
from Bio import SeqIO
import sys
from Bio.SeqIO.QualityIO import FastqGeneralIterator

#set up plate
plate2_i7=['AGGCAGAA', 'TCCTGAGC', 'GGACTCCT', 'TAGGCATG', 'CTCTCTAC', 'CAGAGAGG', 'GCTACGCT', 'CGAGGCTG', 'AAGAGGCA', 'GTAGAGGA', 'GCTCATGA', 'ATCTCAGG', 'ACTCGCTA', 'GGAGCTAC', 'GCGTAGTA', 'CGGAGCCT']
plate2_i5=['TCTACTCT', 'CTCCTTAC', 'TATGCAGT', 'TACTCCTT', 'AGGCTTAG', 'ATTAGACG', 'CGGAGAGA', 'CTAGTCGA', 'AGCTAGAA', 'AGAGTCAA', 'AGATCGCA', 'AGCAGGAA', 'AGTCACTA', 'ATCCTGTA', 'ATTGAGGA', 'CAACCACA', 'GACTAGTA', 'CAATGGAA', 'CACTTCGA', 'CAGCGTTA', 'CATACCAA', 'CCAGTTCA', 'CCGAAGTA', 'CCGTGAGA']

idx_array=[]
i=1
for i5 in plate2_i5:
	for i7 in plate2_i7:
		idx_well=[i5,i7,"C"+str(i)]
		idx_array.append(idx_well)
		i+=1

fq1=sys.argv[1] #read argument fq1="~/fastq/230306_VH00219_371_AACJJFWM5/Undetermined_S0_L001_R1_001.fastq.gz"
fq2=sys.argv[2] #read argument fq2="~/fastq/230306_VH00219_371_AACJJFWM5/Undetermined_S0_L001_R2_001.fastq.gz"
idx3=sys.argv[3] #idx3=~/fastq/230306_VH00219_371_AACJJFWM5/Undetermined_S0_L001_I1_001.fastq.gz"
idx4=sys.argv[4] #idx4=~/fastq/230306_VH00219_371_AACJJFWM5/Undetermined_S0_L001_I2_001.fastq.gz"

i=0
#open fastq files, correct barcode read names then out fastq 1 and 2  with new read name
with gzip.open(fq1, "rt") as handle1:
	with gzip.open(fq2, "rt") as handle2:
		with gzip.open(idx3, "rt") as handle3:
			with gzip.open(idx4, "rt") as handle4:
				with open(fq1[:-9]+".barc.fastq", "w") as outfile_fq1:
					with open(fq2[:-9]+".barc.fastq", "w") as outfile_fq2:
						for (title1, seq1, qual1), (title2, seq2, qual2), (title3,seq3,qual3), (title3,seq4,qual4) in zip(FastqGeneralIterator(handle1), FastqGeneralIterator(handle2),FastqGeneralIterator(handle3),FastqGeneralIterator(handle4)):
							for j in idx_array:
								if seq3==j[1]+"AT" and seq4==j[0]+"GT":
									i+=1
									readname=j[2]+"_"+seq3+seq4
									outfile_fq1.write("@%s:%s\n%s\n+\n%s\n" % (readname, i, seq1, qual1))
									outfile_fq2.write("@%s:%s\n%s\n+\n%s\n" % (readname, i, seq2, qual2))

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import cooler
import cooltools
import os
import seaborn as sns
import cooltools.lib.plotting
import matplotlib
import bioframe

###Prepare Input Files Functions###
#read in eagleC called translocations, plot on top of diagonal
def prepare_eaglec_file(eaglec_in="clone_c3.SV.cnv.neoloopfinder.tsv"):
	""" Function to take in string file name of neoloopfinder output from eaglec_SV_CNV_neoout function (run above in page). Reads in files. 
	Sets transchr SVs to be above diagonal. Currently ignores cis SVs."""
	eaglec_dat=pd.read_csv(eaglec_in,sep="\t",names=["CHROM","CHR2","DIR","POS","CHR2POS","SVTYPE"])
	eaglec_dat_cis=eaglec_dat[eaglec_dat["CHROM"]==eaglec_dat["CHR2"]]
	for index, i in eaglec_dat_cis.iterrows():
		if eaglec_dat_cis.loc[index,"POS"] > eaglec_dat_cis.loc[index,"CHR2POS"]:
				print("Swapping "+eaglec_dat_cis.loc[index,"CHROM"]+" start and end cis SV.")
				tmp=eaglec_dat_cis.loc[index,:]
				eaglec_dat_cis.loc[index,"CHR2POS"]=tmp.POS
				eaglec_dat_cis.loc[index,"POS"]=tmp.CHR2POS
	eaglec_dat_trans=eaglec_dat[eaglec_dat["CHROM"]!=eaglec_dat["CHR2"]]
	for index, i in eaglec_dat_trans.iterrows():
		if (eaglec_dat_trans.loc[index,"CHROM"] in chr_number.keys()) & (eaglec_dat_trans.loc[index,"CHR2"] in chr_number.keys()):
			if chr_number[eaglec_dat_trans.loc[index,"CHROM"]] < chr_number[eaglec_dat_trans.loc[index,"CHR2"]]:
				print("Swapping "+eaglec_dat_trans.loc[index,"CHROM"]+" and "+eaglec_dat_trans.loc[index,"CHR2"])
				tmp=eaglec_dat_trans.loc[index,:]
				eaglec_dat_trans.loc[index,"CHROM"]=tmp.CHR2
				eaglec_dat_trans.loc[index,"CHR2"]=tmp.CHROM
				eaglec_dat_trans.loc[index,"POS"]=tmp.CHR2POS
				eaglec_dat_trans.loc[index,"CHR2POS"]=tmp.POS
	return([eaglec_dat_cis,eaglec_dat_trans])

#read in ont data and split for cis and trans interactions
def setup_ont_data(ont_in):
	ont_dat=pd.read_csv(ont_in,sep="\t")
	ont_dat_trans=ont_dat[(ont_dat['SVTYPE'] == "BND")]
	ont_dat_cis=ont_dat[((ont_dat['SVTYPE'].isin(['INV', 'DEL', 'DUP'])) & (ont_dat['SVLEN']>1000000))]
	for index, i in ont_dat_cis.iterrows(): #set chr and position order, so its always plotted below diagonal
		if ont_dat_cis.loc[index,"POS"] > ont_dat_cis.loc[index,"END"]:
				print("Swapping "+ont_dat_cis.loc[index,"CHROM"]+" start and end cis SV.")
				tmp=ont_dat_cis.loc[index,:]
				ont_dat_cis.loc[index,"END"]=tmp.POS
				ont_dat_cis.loc[index,"POS"]=tmp.END
	ont_dat_trans.loc[:,'CHR2POS']=[int(x.replace("[",",").replace("]",",").split(",")[1].split(":")[1]) for x in ont_dat_trans['ALT'].tolist()]
	for index, i in ont_dat_trans.iterrows(): #set chr and position order, so its always plotted below diagonal
		if (ont_dat_trans.loc[index,"CHROM"] in chr_number.keys()) & (ont_dat_trans.loc[index,"CHR2"] in chr_number.keys()):
			if chr_number[ont_dat_trans.loc[index,"CHROM"]]<chr_number[ont_dat_trans.loc[index,"CHR2"]]:
				print("Swapping "+ont_dat_trans.loc[index,"CHROM"]+" and "+ont_dat_trans.loc[index,"CHR2"])
				tmp=ont_dat_trans.loc[index,:]
				ont_dat_trans.loc[index,"CHROM"]=tmp.CHR2
				ont_dat_trans.loc[index,"CHR2"]=tmp.CHROM
				ont_dat_trans.loc[index,"POS"]=tmp.CHR2POS
				ont_dat_trans.loc[index,"CHR2POS"]=tmp.POS
	return([ont_dat_cis,ont_dat_trans])

###Contact Map Plotting Functions###
#function to plot chromosomes
def trans_chr_plot(mat,chr_row,chr_col,ax_row,ax_col,axes,zmin,zmax,cmap,xlim,ylim,mat_bins,ont_dat_trans,ont_dat_cis,eaglec_dat_cis,eaglec_dat_trans):
	""" Function to plot chr by chr trans interactions (or chr by chr cis interactions if same chr given)
	Takes in c for cooler file and str for chr_row and chr_col (format "chr1","chrX", etc.)
	Takes in integers for ax_row and ax_col to add output plot as subplot.
	"""
	mat=np.log10(mat+1e-9)#adding small value for log
	ax=sns.heatmap(ax=axes[ax_row,ax_col],data=mat,cbar=False,yticklabels=False,square=False,xticklabels=False,vmin=zmin,vmax=zmax,cmap=cmap)
	ax.set_ylim(ylim,0)
	ax.set_xlim(0,xlim)
	if ax_col==0: #set chr name if it is first column
		axes[ax_row,ax_col].set_ylabel(chr_row,fontsize=24)
	print("Completed "+chr_row+" by "+chr_col+" plot")
	if (chr_col==chr_row): #cis SVs
		if chr_row in ont_dat_cis.CHROM.tolist(): #ONT data
			print("Adding ONT data cis annotations for "+chr_row)
			add_ont_sv_cis(ax,mat_bins,
				chr_row=chr_row,chr_col=chr_col,
				ont_dat_cis=ont_dat_cis)
		if chr_row in eaglec_dat_cis.CHROM.tolist(): #Eagle C data
			print("Adding EagleC data cis annotations for "+chr_row)
			add_eaglec_sv_cis(ax,mat_bins,
				chr_row=chr_row,chr_col=chr_col,
				eaglec_dat_cis=eaglec_dat_cis)
	else: #trans SVs
		if any((ont_dat_trans["CHROM"]==chr_row) & (ont_dat_trans["CHR2"]==chr_col)): #ont data
			print("Adding ONT data trans annotations for "+chr_row+" and "+chr_col)
			add_ont_sv_trans(ax,mat_bins,
				chr_row=chr_row,chr_col=chr_col,
				ont_dat_trans=ont_dat_trans)
		if any((eaglec_dat_trans["CHROM"]==chr_row) & (eaglec_dat_trans["CHR2"]==chr_col)): #eagle C data
			print("Adding EagleC trans annotations for "+chr_row+" and "+chr_col)
			add_eaglec_sv_trans(ax,mat_bins,
				chr_row=chr_row,chr_col=chr_col,
				eaglec_dat_trans=eaglec_dat_trans)
	return(ax)

#wrapper function to make panels and loop through chromosome pairs
def all_by_all_plot(ont_dat_cis, ont_dat_trans,eaglec_dat_trans,eaglec_dat_cis,cnv_dat,infile_name="clone_c0.bsorted.1mb.cool", chr_count=4, zmin=-3, zmax=-1, cmap="Reds", dpi=300): 
	"""Function to read in cooler file from given string name. And run all by all chr comparison"""
	#Read in Cooler File
	in_file=infile_name
	in_name=in_file.split(sep=".")[0]
	coolfile=in_file
	print("Reading in "+infile_name)
	c = cooler.Cooler(coolfile)
	#cooler.coarsen_cooler(coolfile,in_name+"_5mb.cool",factor=5,chunksize=5000000) #coarsen to 5mb
	#c = cooler.Cooler(in_name+"_5mb.cool")
	print("Balancing matrix.")
	cooler.balance_cooler(c,store=True,rescale_marginals=True,ignore_diags=1)#balance matrix ignore_diags=10,
	obs_mat = c.matrix()[:]
	chr_list=list(c.bins()[:]["chrom"].unique())
	out=''.join([in_name,'_all_by_all_log2_1Mb_obs.test.png'])
	out_pdf=''.join([in_name,'_all_by_all_log2_1Mb_obs.test.pdf'])
	#init subplot
	chr_in=len(chr_list)
	chr_sizes=pd.DataFrame(c.bins()[:]).groupby(["chrom"])["chrom"].count()
	chr_ratios=list(chr_sizes/chr_sizes[0])
	height_ratio=[chr_ratios[0]/4]+[chr_ratios[0]/4]+chr_ratios[0:chr_count]#added for insulation score and cnv rows
	width_ratio=chr_ratios[0:chr_count]
	print("Calculating Insulation Values.")
	resolution=cnv_dat.loc[0,"end"]-cnv_dat.loc[0,"start"]
	windows = [1*resolution,3*resolution, 5*resolution, 10*resolution]
	insulation_table = cooltools.insulation(c, windows, verbose=True)
	fig, axes = plt.subplots(nrows=chr_count+2, ncols=chr_count, figsize=(sum(height_ratio)*5, sum(width_ratio)*5),sharex='col', sharey='row',
		gridspec_kw={'height_ratios':height_ratio,
		'width_ratios': width_ratio}) #,sharex=True,sharey=True,
	print("Adding column annotations.")
	for j in range(0,chr_count):
			col_chrom=chr_list[j]
			ax_row=0 #0 row is for cnv data
			ax_col=j
			mat=c.matrix().fetch(col_chrom)
			mat_bins=c.bins()[:]
			mat_bins=mat_bins[mat_bins["chrom"].isin([col_chrom])]
			cnv_plot(mat=mat,
				chr_col=col_chrom,
				ax_row=ax_row,
				ax_col=ax_col,
				axes=axes,
				mat_bins=mat_bins,
				cnv_dat=cnv_dat)
			add_insulation_scores(mat=mat,
				chr_col=col_chrom,
				ax_row=ax_row+1,
				ax_col=ax_col,
				axes=axes,
				mat_bins=mat_bins,
				insulation_table=insulation_table,
				resolution=resolution,
				windows=windows)
	print("Adding subpanels of chr-chr contacts.")
	for i in range(0,chr_count):
			row_chrom=chr_list[i]
			ax_row=i+2
			xlim=chr_sizes[i]
			for j in range(0,chr_count):
				col_chrom=chr_list[j]
				ax_col=j
				ylim=chr_sizes[j]
				mat=c.matrix().fetch(row_chrom,col_chrom)
				mat_bins=c.bins()[:]
				mat_bins=mat_bins[mat_bins["chrom"].isin([row_chrom,col_chrom])]
				trans_chr_plot(mat=mat,
					chr_row=row_chrom,chr_col=col_chrom,
					ax_row=ax_row,ax_col=ax_col,axes=axes,
					zmin=zmin,zmax=zmax,cmap=cmap,xlim=xlim,ylim=ylim,
					mat_bins=mat_bins,
					ont_dat_trans=ont_dat_trans,ont_dat_cis=ont_dat_cis,
					eaglec_dat_trans=eaglec_dat_trans,eaglec_dat_cis=eaglec_dat_cis)
	plt.subplots_adjust(wspace=0.02, hspace=0.02)
	plt.savefig(out, dpi=dpi, format='png',bbox_inches='tight')
	plt.savefig(out_pdf, format='pdf',bbox_inches='tight')
	plt.close("all")
	return(insulation_table)

###Add circle patches over detected SVs Functions###
#adds ont circles around cis SVs
def add_ont_sv_cis(ax,mat_bins,chr_row,chr_col,ont_dat_cis,col="green"):
		"""Function to add SV from input ont_dat_cis on same chr. Adds as a circle around the contact map. """			
		for index, i in ont_dat_cis[ont_dat_cis["CHROM"]==chr_row].iterrows():
			if ont_dat_cis.loc[index,"SVTYPE"] == "INV":
				circle_col=col #inversions
			elif ont_dat_cis.loc[index,"SVTYPE"] == "DUP":
				circle_col=col #duplications
			else:
				circle_col=col #deletions
			chr_bin_start=mat_bins[(mat_bins.chrom==chr_row)].index.min()
			bin_annot=mat_bins.loc[(mat_bins.chrom==chr_row) & (mat_bins.start>=ont_dat_cis.loc[index,"POS"]) & (mat_bins.end<ont_dat_cis.loc[index,"END"])]
			bin_annot_x=bin_annot.index.min()-chr_bin_start
			bin_annot_y=bin_annot.index.max()-chr_bin_start
			patch=plt.Circle((bin_annot_x,bin_annot_y), alpha=0.5, radius=10, clip_box=False,linestyle=":", color=circle_col,lw=3,fill=False)
			ax.add_patch(patch)

#adds ont circles around trans SVs
def add_ont_sv_trans(ax,mat_bins,chr_row,chr_col,ont_dat_trans,col="green"):
		"""Function to add SV from input ont_dat_trans on chr-chr pairs. Adds as a circle around the contact map. """			
		for index, i in ont_dat_trans[(ont_dat_trans["CHROM"]==chr_row) & (ont_dat_trans["CHR2"]==chr_col)].iterrows():
			circle_col=col #translocations
			chr_bin_start_x=mat_bins[(mat_bins.chrom==chr_row)].index.min()
			chr_bin_start_y=mat_bins[(mat_bins.chrom==chr_col)].index.min()
			bin_annot_x=mat_bins.loc[(mat_bins.chrom==chr_row) & (mat_bins.start>=ont_dat_trans.loc[index,"POS"])]
			bin_annot_y=mat_bins.loc[(mat_bins.chrom==chr_col) & (mat_bins.start>=ont_dat_trans.loc[index,"CHR2POS"])]
			bin_annot_x=bin_annot_x.index.min()-chr_bin_start_x
			bin_annot_y=bin_annot_y.index.min()-chr_bin_start_y
			patch=plt.Circle((bin_annot_y,bin_annot_x), alpha=0.5,radius=10, clip_box=False,linestyle=":", color=circle_col,lw=3,fill=False)
			ax.add_patch(patch)

#adds ont circles around cis SVs
def add_eaglec_sv_cis(ax,mat_bins,chr_row,chr_col,eaglec_dat_cis,col="blue"):
		"""Function to add SV from input eagleC_dat_cis on same chr. Adds as a circle around the contact map. """			
		for index, i in eaglec_dat_cis[eaglec_dat_cis["CHROM"]==chr_row].iterrows():
			if eaglec_dat_cis.loc[index,"SVTYPE"] == "inversion":
				circle_col=col #inversions
			elif eaglec_dat_cis.loc[index,"SVTYPE"] == "duplication":
				circle_col=col #duplications
			else:
				circle_col=col #deletions and translocations
			chr_bin_start=mat_bins[(mat_bins.chrom==chr_row)].index.min()
			bin_annot=mat_bins.loc[(mat_bins.chrom==chr_row) & (mat_bins.start>=eaglec_dat_cis.loc[index,"POS"]) & (mat_bins.end<eaglec_dat_cis.loc[index,"CHR2POS"])]
			bin_annot_x=bin_annot.index.min()-chr_bin_start
			bin_annot_y=bin_annot.index.max()-chr_bin_start
			patch=plt.Circle((bin_annot_x,bin_annot_y), alpha=0.5,radius=10, clip_box=False,linestyle='--', color=circle_col,lw=3,fill=False)
			ax.add_patch(patch)

#adds eagleC circles around trans SVs
def add_eaglec_sv_trans(ax,mat_bins,chr_row,chr_col,eaglec_dat_trans,col="blue"):
		"""Function to add SV from input eagleC on chr-chr pairs. Adds as a circle around the contact map. """					
		for index, i in eaglec_dat_trans[(eaglec_dat_trans["CHROM"]==chr_row) & (eaglec_dat_trans["CHR2"]==chr_col)].iterrows():
			circle_col=col #translocations
			chr_bin_start_x=mat_bins[(mat_bins.chrom==chr_row)].index.min()
			chr_bin_start_y=mat_bins[(mat_bins.chrom==chr_col)].index.min()
			bin_annot_x=mat_bins.loc[(mat_bins.chrom==chr_row) & (mat_bins.start>=eaglec_dat_trans.loc[index,"POS"])]
			bin_annot_y=mat_bins.loc[(mat_bins.chrom==chr_col) & (mat_bins.start>=eaglec_dat_trans.loc[index,"CHR2POS"])]
			bin_annot_x=bin_annot_x.index.min()-chr_bin_start_x
			bin_annot_y=bin_annot_y.index.min()-chr_bin_start_y
			patch=plt.Circle((bin_annot_y,bin_annot_x), alpha=0.5,radius=10, clip_box=False,linestyle='--', color=circle_col,lw=3,fill=False)
			ax.add_patch(patch)

#adds bedgraph cnv over the heatmap plot
def setup_cnv_bedgraph(cnv_in,col):
	"""Function to read in CNV bedgraph as pandas DF"""
	cnv_dat=pd.read_csv(cnv_in,sep="\t",names=["chrom","start","end","cnv"])
	cnv_col_dict=dict(zip(["0","1","2","3","4","5","6"],col))
	cnv_dat["cnv_col"]=[cnv_col_dict[str(x)] for x in cnv_dat.cnv]
	return(cnv_dat)

#adds cnv plots as annotation in top row
def cnv_plot(mat,chr_col,ax_row,ax_col,axes,mat_bins,cnv_dat):
	"""Function to add CNV profiles on top row of contact heatmaps"""
	print("Adding CNV Profile of "+chr_col)
	cnv_dat_tmp=cnv_dat[cnv_dat["chrom"].isin([chr_col])].copy()
	cnv_dat_tmp["cnv_one"]=[1 for x in cnv_dat_tmp.cnv]
	sns.barplot(ax=axes[ax_row,ax_col],x=cnv_dat_tmp.start,y=cnv_dat_tmp.cnv_one,palette=cnv_dat_tmp.cnv_col,saturation=1,lw=0)
	axes[ax_row,ax_col].set_xticks([])
	axes[ax_row,ax_col].set_yticks([])
	if ax_col==0:
		axes[ax_row,ax_col].set_ylabel("Copy \n Number")
	else:
		axes[ax_row,ax_col].set_ylabel(None)
	axes[ax_row,ax_col].set_ylim([0,1])
	axes[ax_row,ax_col].set_title(chr_col,fontsize=24)
	return(axes[ax_row,ax_col])

#adds detected TAD boundaries in second from top row
def add_insulation_scores(mat,chr_col,ax_row,ax_col,axes,mat_bins,insulation_table,resolution,windows):
	print("Adding insulation score for "+chr_col)
	region = (chr_col, min(mat_bins["start"]), max(mat_bins["end"]))
	insul_region = bioframe.select(insulation_table, region)
	insul_region=insul_region.assign(boundary_col=["black" if x is True else "white" for x in insul_region.loc[:,"is_boundary_"+str(resolution*3)]])
	insul_region=insul_region.assign(bound_one=[1 for x in range(0,insul_region.shape[0])])
	sns.barplot(ax=axes[ax_row,ax_col],x=insul_region.start,y=insul_region.bound_one,palette=insul_region.boundary_col,saturation=1,lw=0)
	axes[ax_row,ax_col].set_xticks([])
	axes[ax_row,ax_col].set_yticks([])
	if ax_col==0:
		axes[ax_row,ax_col].set_ylabel("TAD \n Boundaries")
	else:
		axes[ax_row,ax_col].set_ylabel(None)
	axes[ax_row,ax_col].set_ylim([0,1])
	print("Found "+str(sum(insul_region.loc[:,"is_boundary_"+str(resolution*3)]))+" boundaries in "+chr_col)
	return(axes[ax_row,ax_col])

os.chdir('/volumes/seq/projects/gccACT/230306_mdamb231_test/contacts')#wd 

#list chromosomes and set dictionary of order
L1=["chr"+str(x) for x in list(range(1,23))+["X","Y"]]
L2=list(range(1,25))
chr_number={k:v for k,v in zip(L1,L2)}

#read in ont data
ont_in="/volumes/seq/projects/gccACT/230808_mdamb231_ONT/20230726_1239_2D_PAO38369_output/20230726_1239_2D_PAO38369_output.1mb_SVs.tsv"
ont_dat=setup_ont_data(ont_in)

cnv_col_palette=["#053061","#4393c3","#f7f7f7","#fddbc7","#d6604d","#b2182b","#67001f"]

#make all by all plot for all clones, set up looping list
in_cool=[
"clone_c1.bsorted.1mb.cool",
"clone_c2.bsorted.1mb.cool",
"clone_c3.bsorted.1mb.cool"]

in_eaglec=[
"clone_c1.SV.cnv.neoloopfinder.tsv",
"clone_c2.SV.cnv.neoloopfinder.tsv",
"clone_c3.SV.cnv.neoloopfinder.tsv"]

in_bedgraph=[
"/volumes/seq/projects/gccACT/230306_mdamb231_test/clone_c1.cnv.1mb.segmented.bedgraph",
"/volumes/seq/projects/gccACT/230306_mdamb231_test/clone_c2.cnv.1mb.segmented.bedgraph",
"/volumes/seq/projects/gccACT/230306_mdamb231_test/clone_c3.cnv.1mb.segmented.bedgraph"]

for clone in range(0,len(in_cool)):
	eaglec_dat=prepare_eaglec_file(eaglec_in=in_eaglec[clone])
	cnv_dat=setup_cnv_bedgraph(in_bedgraph[clone],col=cnv_col_palette)
	all_by_all_plot(infile_name=in_cool[clone],
	chr_count=4,
	ont_dat_cis=ont_dat[0],
	ont_dat_trans=ont_dat[1],
	eaglec_dat_cis=eaglec_dat[0],
	eaglec_dat_trans=eaglec_dat[1],
	cnv_dat=cnv_dat,
	zmin=-3.5,
	cmap="fall")



eaglec_dat=prepare_eaglec_file(eaglec_in="clone_c1.SV.cnv.neoloopfinder.tsv")
cnv_dat=setup_cnv_bedgraph("/volumes/seq/projects/gccACT/230306_mdamb231_test/clone_c1.cnv.1mb.segmented.bedgraph",
	col=cnv_col_palette)
tab=all_by_all_plot(infile_name="clone_c1.bsorted.1mb.cool",
chr_count=4,
ont_dat_cis=ont_dat[0],
ont_dat_trans=ont_dat[1],
eaglec_dat_cis=eaglec_dat[0],
eaglec_dat_trans=eaglec_dat[1],
cnv_dat=cnv_dat,
zmin=-3.5,
cmap="Blues")


eaglec_dat=prepare_eaglec_file(eaglec_in="clone_c3.SV.cnv.neoloopfinder.tsv")
cnv_dat=setup_cnv_bedgraph("/volumes/seq/projects/gccACT/230306_mdamb231_test/clone_c4.cnv.1mb.segmented.bedgraph",
	col=cnv_col_palette)
tab=all_by_all_plot(infile_name="clone_c4.bsorted.1mb.cool",
chr_count=4,
ont_dat_cis=ont_dat[0],
ont_dat_trans=ont_dat[1],
eaglec_dat_cis=eaglec_dat[0],
eaglec_dat_trans=eaglec_dat[1],
cnv_dat=cnv_dat,
zmin=-3.5,
cmap="Reds")
#additional cmap color pallets here
#https://www.practicalpythonfordatascience.com/ap_seaborn_palette
#some good ones! magma, mako, rocket, twilight, viridis, vlag

#from here: https://cooltools.readthedocs.io/en/latest/notebooks/insulation_and_boundaries.html
# import standard python libraries
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import cooler
import cooltools.lib.plotting
from cooltools import insulation
import itertools
from matplotlib.ticker import EngFormatter
from matplotlib.colors import LogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
import bioframe
import os 
os.chdir('/volumes/seq/projects/gccACT/230306_mdamb231_test/contacts')#wd 

resolution=500000
clone_dir="/volumes/seq/projects/gccACT/230306_mdamb231_test/contacts"
clone="clone_c3"
coolfile=clone_dir+"/"+clone+".bsorted.500kb.cool"
c = cooler.Cooler(coolfile)
windows = [1*resolution,3*resolution, 5*resolution, 10*resolution]
insulation_table = insulation(c, windows, verbose=True)
bp_formatter = EngFormatter('b')

# Functions to help with plotting
def pcolormesh_45deg(ax, matrix_c, start=0, resolution=1, *args, **kwargs):
    start_pos_vector = [start+resolution*i for i in range(len(matrix_c)+1)]
    n = matrix_c.shape[0]
    t = np.array([[1, 0.5], [-1, 0.5]])
    matrix_a = np.dot(np.array([(i[1], i[0])
                                for i in itertools.product(start_pos_vector[::-1],
                                                           start_pos_vector)]), t)
    x = matrix_a[:, 1].reshape(n + 1, n + 1)
    y = matrix_a[:, 0].reshape(n + 1, n + 1)
    im = ax.pcolormesh(x, y, np.flipud(matrix_c), *args, **kwargs)
    im.set_rasterized(True)
    return im

def format_ticks(ax, x=True, y=True, rotate=True):
    if y:
        ax.yaxis.set_major_formatter(bp_formatter)
    if x:
        ax.xaxis.set_major_formatter(bp_formatter)
        ax.xaxis.tick_bottom()
    if rotate:
        ax.tick_params(axis='x',rotation=45)


plt.rcParams['font.size'] = 12
start = 1000000
end = start+ 100*windows[0]
region = ('chr4', start, end)
norm = LogNorm(vmax=0.1, vmin=0.001)
data = c.matrix(balance=True).fetch(region)
f, ax = plt.subplots(figsize=(18, 6))
im = pcolormesh_45deg(ax, data, start=region[1], resolution=resolution, norm=norm, cmap='fall')
ax.set_aspect(0.5)
ax.set_ylim(0, 10*windows[0])
format_ticks(ax, rotate=False)
ax.xaxis.set_visible(False)

divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="1%", pad=0.1, aspect=6)
plt.colorbar(im, cax=cax)

insul_region = bioframe.select(insulation_table, region)

ins_ax = divider.append_axes("bottom", size="50%", pad=0., sharex=ax)
ins_ax.set_prop_cycle(plt.cycler("color", plt.cm.plasma(np.linspace(0,1,5))))
ins_ax.plot(insul_region[['start', 'end']].mean(axis=1),
            insul_region['log2_insulation_score_'+str(windows[0])],
            label=f'Window {windows[0]} bp')

ins_ax.legend(bbox_to_anchor=(0., -1), loc='lower left', ncol=4);

format_ticks(ins_ax, y=False, rotate=False)
ax.set_xlim(region[1], region[2])

for res in windows[1:]:
    ins_ax.plot(insul_region[['start', 'end']].mean(axis=1), insul_region[f'log2_insulation_score_{res}'], label=f'Window {res} bp')

ins_ax.legend(bbox_to_anchor=(0., -1), loc='lower left', ncol=4);
f
plt.savefig("insulation_test.png", format='png',bbox_inches="tight")
plt.close("all")


f, ax = plt.subplots(figsize=(20, 10))
im = pcolormesh_45deg(ax, data, start=region[1], resolution=resolution, norm=norm, cmap='fall')
ax.set_aspect(0.5)
ax.set_ylim(0, 10*windows[0])
format_ticks(ax, rotate=False)
ax.xaxis.set_visible(False)

divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="1%", pad=0.1, aspect=6)
plt.colorbar(im, cax=cax)

insul_region = bioframe.select(insulation_table, region)

ins_ax = divider.append_axes("bottom", size="50%", pad=0., sharex=ax)

ins_ax.plot(insul_region[['start', 'end']].mean(axis=1),
            insul_region[f'log2_insulation_score_{windows[0]}'], label=f'Window {windows[0]} bp')

boundaries = insul_region[~np.isnan(insul_region[f'boundary_strength_{windows[0]}'])]
weak_boundaries = boundaries[~boundaries[f'is_boundary_{windows[0]}']]
strong_boundaries = boundaries[boundaries[f'is_boundary_{windows[0]}']]
ins_ax.scatter(weak_boundaries[['start', 'end']].mean(axis=1),
            weak_boundaries[f'log2_insulation_score_{windows[0]}'], label='Weak boundaries')
ins_ax.scatter(strong_boundaries[['start', 'end']].mean(axis=1),
            strong_boundaries[f'log2_insulation_score_{windows[0]}'], label='Strong boundaries')

ins_ax.legend(bbox_to_anchor=(0., -1), loc='lower left', ncol=4);

format_ticks(ins_ax, y=False, rotate=False)
ax.set_xlim(region[1], region[2])

plt.savefig("insulation_test.png", format='png',bbox_inches="tight")
plt.close("all")

histkwargs = dict(
    bins=10**np.linspace(-4,1,200),
    histtype='step',
    lw=2,
)

f, axs = plt.subplots(len(windows),1, sharex=True, figsize=(6,6), constrained_layout=True)
for i, (w, ax) in enumerate(zip(windows, axs)):
    ax.hist(
        insulation_table[f'boundary_strength_{w}'],
        **histkwargs
    )
    ax.text(0.02, 0.9,
             f'Window {w//1000}kb',
             ha='left',
             va='top',
             transform=ax.transAxes)
    ax.set(
        xscale='log',
        ylabel='# boundaries'
    )

axs[-1].set(xlabel='Boundary strength');
plt.savefig("insulation_boundary_test.png", format='png',bbox_inches="tight")
plt.close("all")


from skimage.filters import threshold_li, threshold_otsu

f, axs = plt.subplots(len(windows), 1, sharex=True, figsize=(6,6), constrained_layout=True)
thresholds_li = {}
thresholds_otsu = {}
for i, (w, ax) in enumerate(zip(windows, axs)):
    ax.hist(insulation_table[f'boundary_strength_{w}'],**histkwargs)
    thresholds_li[w] = threshold_li(insulation_table[f'boundary_strength_{w}'].dropna().values)
    thresholds_otsu[w] = threshold_otsu(insulation_table[f'boundary_strength_{w}'].dropna().values)
    n_boundaries_li = (insulation_table[f'boundary_strength_{w}'].dropna()>=thresholds_li[w]).sum()
    n_boundaries_otsu = (insulation_table[f'boundary_strength_{w}'].dropna()>=thresholds_otsu[w]).sum()
    ax.axvline(thresholds_li[w], c='green')
    ax.axvline(thresholds_otsu[w], c='magenta')
    ax.text(0.01, 0.9, f'Window {w//1000}kb', ha='left', va='top', transform=ax.transAxes) 
    ax.text(0.01, 0.7, f'{n_boundaries_otsu} boundaries (Otsu)', c='magenta', ha='left', va='top', transform=ax.transAxes)
    ax.text(0.01, 0.5, f'{n_boundaries_li} boundaries (Li)', c='green', ha='left', va='top', transform=ax.transAxes)
    ax.set(xscale='log', ylabel='# boundaries')
axs[-1].set(xlabel='Boundary strength')
plt.savefig("insulation_boundary_thresholding_test.png", format='png',bbox_inches="tight")
plt.close("all")

# Download CTCF bigwig file. The file size is 592 Mb, so the download might take a while:
ctcf_fc_file = cooltools.download_data("HFF_CTCF_fc", cache=True, data_dir="/volumes/USR2/Ryan/ref/refdata-cellranger-arc-GRCh38-2020-A-2.0.0/fasta")