import os
wd="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/"
sc_dir="gcc_cisdistal_triplesparse/"

if not os.path.exists(wd+sc_dir):
    os.makedirs(wd+sc_dir)

print("Directory created or already present.")
        

import glob
import pysam
from collections import defaultdict
from collections import Counter
import pandas as pd
import os
import time

def read_pair_generator(bam):
    """
    Generate read pairs in a BAM file or within a region string.
    Reads are added to read_dict until a pair is found.
    """
    read_dict = defaultdict(lambda: [None, None])
    for read in bam.fetch(until_eof=True):
        qname = read.query_name
        if qname not in read_dict:
            if read.is_read1:
                read_dict[qname][0] = read
            else:
                read_dict[qname][1] = read
        else:
            if read.is_read1:
                yield read, read_dict[qname][1]
            else:
                yield read_dict[qname][0], read
            del read_dict[qname]

bam_dir="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/filtered_bam/" #working directory
bam_bulk=["CRC-4442CRC-4671_GCC_H.bbrd.q10.filt.bam"] #list of bam files
#Takes read 1 and read 2 paired by query name with the above function
#then rounds read 1 and read 2 starts to the nearest 1mb and takes that as the bin (drops the 6 0s for the bin names)
out_dict = defaultdict(dict)
chr_count=[]

i=0 #Counter just to check progress
start_time = time.time() #And set up a timer

#cis interactions first
for bulk in bam_bulk:
    cellline=bulk.split(".")[0]
    bam = pysam.AlignmentFile(bam_dir+bulk, 'rb')
    for read1, read2 in read_pair_generator(bam): #set paired reads together
        if read1 is not None and read2 is not None: #remove reads which dont pair
            if read1.reference_id==read2.reference_id: #ensure chromosomes are same
                if abs(read1.template_length)>=50000: #ignore reads <50kbp in length
                    i+=1
                    if i % 100000 == 0:
                        print("Processed "+str(i)+" distal cis reads in "+str(time.time()-start_time)+" seconds.")
                    if read1.reference_start < 1000000: #bin for 0 (less than 1mbp into chr)
                        r1_bin="0"
                    else:
                        r1_bin=str(read1.reference_start)[:-6] #bin for anything greater than 1mb, just cut last 6 digits off
                    if read2.reference_start < 1000000:
                        r2_bin="0"
                    else:
                        r2_bin=str(read2.reference_start)[:-6]
                    if int(r2_bin) < int(r1_bin): #quick sorting of bins
                      tmp=r2_bin
                      r2_bin=r1_bin
                      r1_bin=tmp
                    chr_out=bam.get_reference_name(read1.reference_id) #set chr name
                    chr_count.append(chr_out) #make sure chr name was seen before or added
                    cellid_out=read1.query_name.split(":")[0] #assign read to cellID
                    cellid_out=cellline+"_"+cellid_out
                    if cellid_out not in out_dict: #add any new cellIDs
                      out_dict[cellid_out][chr_out]=[[r1_bin,r2_bin]] #build out dict
                    if chr_out not in out_dict[cellid_out]:
                      out_dict[cellid_out][chr_out]=[[r1_bin,r2_bin]]
                    else:
                      out_dict[cellid_out][chr_out].append([r1_bin,r2_bin])

chr_count=len(Counter(chr_count))

wd="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/" #working directory
sc_dir="gcc_cisdistal_triplesparse/"

with open(wd+"scHiCluster.network","w") as fout:
  for cellid_out in out_dict: #generate sciHiCluster.network file (list of path/cellname)
      fout.write(wd+sc_dir+cellid_out+"\n")

for cellid_out in out_dict: #generate [cellID].[chr].txt files in triple sparse format
    for chr_out in out_dict[cellid_out]:
      with open(wd+sc_dir+cellid_out+"_"+chr_out+".txt","w") as fout:
        dict = Counter([tuple(i) for i in out_dict[cellid_out][chr_out]])
        # Creating pandas dataframe
        Output = pd.DataFrame(data ={'bin1': list([item[0] for item in list(dict.keys())]),'bin2': list([item[1] for item in list(dict.keys())]),'count': list(dict.values())})
        Output.to_csv(fout,sep="\t",header=False,index=False)
        

import glob
import pysam
from collections import defaultdict
from collections import Counter
import pandas as pd
import os
import time

def read_pair_generator(bam):
    """
    Generate read pairs in a BAM file or within a region string.
    Reads are added to read_dict until a pair is found.
    """
    read_dict = defaultdict(lambda: [None, None])
    for read in bam.fetch(until_eof=True):
        qname = read.query_name
        if qname not in read_dict:
            if read.is_read1:
                read_dict[qname][0] = read
            else:
                read_dict[qname][1] = read
        else:
            if read.is_read1:
                yield read, read_dict[qname][1]
            else:
                yield read_dict[qname][0], read
            del read_dict[qname]

bam_dir="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/" #working directory
bam_bulk=["s3gcc_crc.bbrd.q10.filt.bam","s3gcc_cellline.bbrd.q10.filt.bam"] #list of bam files
#Takes read 1 and read 2 paired by query name with the above function
#then rounds read 1 and read 2 starts to the nearest 1mb and takes that as the bin (drops the 6 0s for the bin names)

distal_count={}
trans_count={}
local_count={}

#cis distal and trans interactions counting
for bulk in bam_bulk:
    cellline=bulk.split(".")[0]
    bam = pysam.AlignmentFile(bam_dir+bulk, 'rb')
    for read1, read2 in read_pair_generator(bam): #set paired reads together
        if read1 is not None and read2 is not None: #remove reads which dont pair
            cellid_out=read1.query_name.split(":")[0]
            if read1.reference_id==read2.reference_id: #ensure chromosomes are same
                if abs(read1.template_length)>=1000: #ignore reads <1kbp in length
                    if cellid_out in distal_count:
                        distal_count[cellid_out]+=1
                    else:
                        distal_count[cellid_out]=1
                else:
                    if cellid_out in local_count:
                        local_count[cellid_out]+=1
                    else:
                        local_count[cellid_out]=1
                    
            if read1.reference_id!=read2.reference_id: #ensure chromosomes are same
                    if cellid_out in trans_count:
                        trans_count[cellid_out]+=1
                    else:
                        trans_count[cellid_out]=1

                        
distal_df=pd.DataFrame.from_dict(distal_count,orient="index",columns=["distal_cis_interactions"])
distal_df.index.name = 'cellID'
distal_df.reset_index(inplace=True)

trans_df=pd.DataFrame.from_dict(trans_count,orient="index",columns=["trans_interactions"])
trans_df.index.name = 'cellID'
trans_df.reset_index(inplace=True)

local_df=pd.DataFrame.from_dict(local_count,orient="index",columns=["local_cis"])
local_df.index.name = 'cellID'
local_df.reset_index(inplace=True)

df = pd.merge(left=distal_df, right=trans_df, left_on='cellID', right_on='cellID')
df = pd.merge(left=df, right=local_df, left_on='cellID', right_on='cellID')
df.to_csv(bam_dir+'s3gcc_hiccontacts.csv', index=False)

import time
import numpy as np
from schicluster import *
from sklearn.cluster import KMeans, SpectralClustering
from sklearn.metrics.cluster import adjusted_rand_score as ARI
import csv
import glob
import pandas as pd

wd="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/" #working directory
sc_dir="gcc_cisdistal_triplesparse/"

network_file=wd+"scHiCluster.network"
network = open(network_file).read().splitlines()
network = np.loadtxt(network_file,dtype=np.str)

#sizes (in bp) of autosomes in numerical order taken from UCSC hg38 
hg38dim=[248956422,242193529,198295559,190214555,181538259,170805979,159345973,145138636,138394717,133797422,135086622,133275309,114364328,107043718,101991189,90338345,83257441,80373285,58617616,64444167,46709983,50818468] #this is all autosomes, but I'm just using chr1-16 for now (others too sparse so it throws errors)
chrom = [str(i+1) for i in range(22)] #only using autosomes 1-22 due to sparsity

chromsize = {chrom[i]:hg38dim[i] for i in range(len(chrom))} #generate dicitonary of chromosome names and sizes

nc=5 #5 clusters?

chr_count = []
chr_names=["chr"+s for s in chrom]
#need to filter cells with less than a certain number of reads per chr
#building a pandas data frame for count per chr
for cellID in network:
    cellID_out=cellID.split("/")[-1]
    for file in glob.glob(cellID+"*"+".txt"):
        chr_out=file.split("/")[-1].split("_")[-1].split(".")[0]
        num_lines = sum(1 for line in open(file))
        chr_count.append([cellID_out,chr_out,num_lines])

chr_count=pd.DataFrame(chr_count,columns=["cellID","chr_out","contact_count"])
filter_list=[]

for i in chr_count.cellID.unique():
    tmp=chr_count.loc[chr_count['cellID'] == i]
    tmp=tmp[tmp["chr_out"].isin(chr_names)]
    if len(tmp)==len(chr_names) and tmp.contact_count.min()>1:
        filter_list.append(i)
        
#filter network to cells with atleast 2 reads per chr
#goal here is to group sparse cells for more thorough analysis

network_fil=["/".join(network[0].split("/")[0:-1])+"/"+i for i in filter_list] #filtering the network list to those that pass the count filter
len(network_fil)
label=["_".join(i.split("/")[-1].split("_")[0:2]) for i in network_fil]

start_time = time.time()
cluster,embedding=hicluster_cpu(network_fil,chromsize,nc=nc, pad=1, rp=0.5, prct=20, ncpus=20)
print(time.time() - start_time)
[ARI(label, KMeans(n_clusters = nc, n_init = 200).fit(embedding[:, :ndim]).labels_) for ndim in [2,5,10,20,30,40,50]]

out_dir=wd

#output annoted clusters
with open(wd+"scHiCluster.annot","w") as fout:
    for i in range(0,len(cluster)-1):
        fout.write(filter_list[i]+"\t"+str(cluster[i])+"\n")
        print(filter_list[i]+"\t"+str(cluster[i]))

        
#output embeddings
embed=pd.DataFrame(embedding,index=filter_list)
with open(wd+"scHiCluster.dims","w") as fout:
    embed.to_csv(fout,sep="\t",header=True,index=True)

import time
import numpy as np
from schicluster import *
from sklearn.cluster import KMeans, SpectralClustering
from sklearn.metrics.cluster import adjusted_rand_score as ARI
import csv
import glob
import pandas as pd

network_file_list=glob.glob("/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/for_gurkan/network*txt")

chromsize_file="/home/groups/oroaklab/adey_lab/projects/sciWGS/Public_Data/chromLengths.txt"
hg38dim=[248956422,242193529,198295559,190214555,181538259,170805979,159345973,145138636,138394717,133797422,135086622,133275309,114364328,107043718,101991189,90338345,83257441,80373285,58617616,64444167,46709983,50818468] #this is all autosomes, but I'm just using chr1-16 for now (others too sparse so it throws errors)
chrom = [str(i+1) for i in range(16)]
chromsize = {chrom[i]:hg38dim[i] for i in range(len(chrom))}

chr_count = []
chr_names=["chr"+s for s in chrom]

res=1000000

#this had to be corrected to work for CPU
#also added a loop for single cell topdom
def merge_cpu(network, c, res, pad=1, rp=0.5, prct=-1,chromsize=chromsize):
    ngene = int(chromsize[c] / res) + 1
    start_time = time.time()
    Q_sum = np.zeros(ngene * ngene)
    for cell in network:
        Q_sum = Q_sum + impute_cpu([cell, c, ngene, pad, rp])[1]
    end_time = time.time()
    print('Load and impute chromosome', c, 'take', end_time - start_time, 'seconds')
    Q_sum = Q_sum.reshape(ngene, ngene)
    return Q_sum


#this also had to be fixed with chromsize as variable
def output_topdom(cell, c, Q, res, chromsize=chromsize):
    ngene, _ = Q.shape
    B = [['chr' + c, i * res, (i + 1) * res] for i in range(ngene)]
    B[-1][-1] = chromsize[c]
    C = np.concatenate((B, Q), axis=1)
    np.savetxt(cell + '_chr' + c + '.topdommatrix', C, fmt = '%s', delimiter = '\t')
    return

for network_file in network_file_list:
    network = open(network_file).read().splitlines()
    network = np.loadtxt(network_file,dtype=np.str)
    cell=network_file.split(".")[0]
    print("Processing:"+cell)
    for c in chromsize: #for whole network merged file
        Q = merge_cpu(network, c, res=1000000)
        output_topdom(cell, c, Q, res)
        output_sparse(cell, c, Q, res)

for network_file in network_file_list:
    print(network_file)
    for cell in open(network_file).read().splitlines():
        for c in chromsize:
            ngene = int(chromsize[c] / res) + 1
            start_time = time.time()
            pad=1
            rp=0.5
            Q = np.zeros(ngene * ngene)
            print("Processing"+" "+cell+" "+c)
            Q=impute_cpu([cell, c, ngene,pad,rp])[1]
            end_time = time.time()
            Q = Q.reshape(ngene, ngene)
            output_topdom(cell,c, Q, res)

import time
import numpy as np
from schicluster import *
from sklearn.cluster import KMeans, SpectralClustering
from sklearn.metrics.cluster import adjusted_rand_score as ARI
import csv
import glob
import pandas as pd
import re
from statsmodels.sandbox.stats.multicomp import multipletests as FDR

#using this filter function for subsetting lists
def Filter(string, substr):
    return [str for str in string
    if any(sub in str for sub in substr)]


#diff domain call from scihicluster, pretty extensively modified to work properly

def diff_dom(args):
    domlist, cluster, ctlist, res, c, chromsize = args
    start_time = time.time()
    ctdict = {x:i for i,x in enumerate(ctlist)}
    ngene = int(chromsize[c] / res) + 1
    bound = np.zeros((len(cluster), ngene))
    for i,cell in enumerate(domlist):
    	domain = np.loadtxt(cell, dtype = np.str)
    	C = domain[domain[:,-1]=='domain', :2].astype(int) / res
    	bound[i, np.array(list(set(C.flatten())),dtype=int)] += 1
    	
    cellfilter = (np.sum(bound, axis=1)>0)
    count = np.array([np.sum(np.logical_and([k==x for x in cluster], cellfilter)) for k in ctlist])
    prob = np.array([np.sum(bound[[x==k for x in cluster]], axis=0) for k in ctlist])
    ptmp, btmp = [],[]
    for i in range(ngene):
    	contig = [[prob[j,i], count[j]-prob[j,i]] for j in range(len(ctlist))]
    	if np.sum(np.sum(contig, axis=0)==0)==0:
    		ptmp.append(chi2_contingency(contig)[1])
    		btmp.append('\t'.join([c, str(i*res), str(i*res+res)]))
    		
    print(c, time.time()-start_time)
    return [bound, prob/count[:,None], btmp, ptmp]


network_file_list=glob.glob("/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/network*txt")

domlist=glob.glob("/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/gcc_cisdistal_triplesparse/*w5.domain")

chromsize_file="/home/groups/oroaklab/adey_lab/projects/sciWGS/Public_Data/chromLengths.txt"

hg38dim=[248956422,242193529,198295559,190214555,181538259,170805979,159345973,145138636,138394717,133797422,135086622,133275309,114364328,107043718,101991189,90338345,83257441,80373285,58617616,64444167,46709983,50818468] #this is all autosomes, but I'm just using chr1-16 for now (others too sparse so it throws errors)
chrom = [str(i+1) for i in range(16)]
chromsize = {chrom[i]:hg38dim[i] for i in range(len(chrom))}

chr_count = []
chr_names=["chr"+s for s in chrom]

res=1000000

for chrom  in chromsize:
    chrom_filt_name="_chr"+chrom+"." #set up chr name
    domlist_chr=[i for i in domlist if chrom_filt_name in i] #first filter topdom list to the chromosome
    domlist_cell_order=[] #initiate empty topdom list for each chr
    cluster_list=[] #initiate empty cluster list for each chr
    for cluster in network_file_list: #loop through each cluster network file
        clus_name=[cluster.split("/")[-1].split(".")[0]] #set cluster name
        cluster_list=cluster_list+(clus_name)*len(open(cluster).read().splitlines()) #make a list of cluster assignment per cell by multiplying cluster name by number of cells
        domlist_cell_order=domlist_cell_order+Filter(domlist_chr,open(cluster).read().splitlines()) #make list of topdoms with appropriate cluster assignment and chr
    celltype_list=list(set(cluster_list)) #set celltype_list by unique celltypes
    sc_dom, dom_prob, bins, pvalue=diff_dom([domlist_cell_order,cluster_list,celltype_list,res,chrom,chromsize])
    bin_names=["chr"+chrom+":"+str(i*res)+"-"+str((i+1)*res) for i in range(np.shape(sc_dom)[1])]
    cellid_names=[cell.split("/")[-1].split("_chr")[0] for cell in domlist_cell_order]
    sc_dom_df=pd.DataFrame(sc_dom,index=cellid_names,columns=bin_names)
    dom_prob_df=pd.DataFrame(dom_prob,index=celltype_list,columns=bin_names)
    pval_df=pd.DataFrame(pvalue,index=[col.split("\t")[0]+":"+col.split("\t")[1]+"-"+col.split("\t")[2] for col in bins],columns=["pvalue"])
    sc_dom_df.to_csv("chr"+chrom+"_scDomainBoundaries.csv")
    dom_prob_df.to_csv("chr"+chrom+"_scDomainProbabilities.csv")
    pval_df.to_csv("chr"+chrom+"_scDomainProbabilities_pval.csv")

#merge raw matrices for whole chr comparisons

import glob
import pysam
import pandas as pd
from collections import defaultdict
import time
import sys, os

# Disable
def blockPrint():
    sys.stdout = open(os.devnull, 'w')


# Restore
def enablePrint():
    sys.stdout = sys.__stdout__


def read_pair_generator(bam):
    """
    Generate read pairs in a BAM file or within a region string.
    Reads are added to read_dict until a pair is found.
    """
    read_dict = defaultdict(lambda: [None, None])
    for read in bam.fetch(until_eof=True):
        qname = read.query_name
        if qname not in read_dict:
            if read.is_read1:
                read_dict[qname][0] = read
            else:
                read_dict[qname][1] = read
        else:
            if read.is_read1:
                yield read, read_dict[qname][1]
            else:
                yield read_dict[qname][0], read
            del read_dict[qname]

    
with open("/home/groups/oroaklab/adey_lab/projects/sciWGS/200522_s3WGS_CRC/gcc_cisdistal_triplesparse/scHiCluster.annot","r") as annot:
    annot_df=pd.DataFrame([line.split() for line in annot],columns=["cellID","cluster"])
    annot_df["cluster"]=[x[1] for x in annot_df["cellID"].str.split(pat="_")] #using cell lines



annot_df["cellname"]=[line.split("_")[2] for line in annot_df["cellID"]]

bam_dir="/home/groups/oroaklab/adey_lab/projects/sciWGS/200522_s3WGS_CRC/" #working directory
bam_bulk=["hg38.s3GCC_4671.bbrd.q10.filt.bam","hg38.s3GCC_4442.bbrd.q10.filt.bam"] #list of bam files
#Takes read 1 and read 2 paired by query name with the above function

dir_out="/home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/for_gurkan/"

i=0 # start counter
start_time = time.time() # start timer
chr_set=["chr"+str(i) for i in range(1,23)]

blockPrint()
for clus in annot_df["cluster"].unique():
    cell_set=list(annot_df.loc[annot_df["cluster"]==str(clus)]["cellID"])
    with open(dir_out+"s3GCC_clus"+clus+".bulk","w") as fout:
        for bulk in bam_bulk:
            cellline=bulk.split(".")[1]
            bam = pysam.AlignmentFile(bam_dir+bulk, 'rb')
            for read1, read2 in read_pair_generator(bam): #set paired reads together
                if read1 is not None and read2 is not None: #remove reads which dont pair
                    chr1=bam.get_reference_name(read1.reference_id) #set chr1 name
                    chr2=bam.get_reference_name(read2.reference_id) #set chr2 name
                    if (chr1 in chr_set) and (chr2 in chr_set):
                        if ((chr1 == chr2) and (abs(read1.template_length)>=1000)) or (chr1 is not chr2): #so same chr >1kb or diff chr
                            i+=1
                            if i % 10000 == 0:
                                enablePrint()
                                print("Processed "+str(i)+" distal reads in "+ str(time.time()-start_time) +" seconds.")
                                blockPrint()
                            str1="+"
                            str2="+"
                            if read1.is_reverse:
                                str1="-"
                            if read2.is_reverse:
                                str2="-"
                            pos1=read1.reference_start
                            pos2=read2.reference_start
                            cellid_out=cellline+"_"+read1.query_name.split(":")[0] #assign read to cellID
                            if cellid_out in cell_set: #check to make sure it was a cellID used in clustering
                                out_list=[read1.query_name,chr1,pos1,str1,chr2,pos2,str2]
                                fout.write("\t".join(str(item) for item in out_list)+"\n")


#1. Read Name (can be blank)
#2. chromosome for read 1
#3. positions for read 1 (5' end of read, one-indexed)
#4. strand of read 1 (+ or -)
#5. chromosome for read 2
#6. positions for read 2 (5' end of read, one-indexed)
#7. strand of read 2 (+ or -)

import subprocess as sp
### Using awk for final formatting of text into HOMER input.
cmd= """for i in `ls /home/groups/oroaklab/adey_lab/projects/sciWGS/200730_s3FinalAnalysis/s3gcc_data/for_gurkan/*bulk`; do awk 'OFS="\t" {if ($2 > $5){ print $1,$5,$6,$7,$2,$3,$4} else {print}}' $i | awk 'OFS="\t" {if (($3 > $6) && ($2 == $5)) {print $1,$5,$6,$7,$2,$3,$4} else {print}}' | sort -T . -S 2G -k2,2d -k5,5d -k3,3n -k6,6n > $i.sorted.hicSummary ; done"""

p = sp.Popen(cmd, stdin=sp.PIPE, stdout = sp.PIPE, stderr = sp.PIPE,shell=True)
p.wait()
### Using Homer for generation of the raw formatted HiC matrix
#Resolution of 2.5MB and uwing 20CPUs
cmd="""for i in clus4442 clus4671;
do /home/groups/oroaklab/src/homer/bin/makeTagDirectory s3GCC_${i} -format HiCsummary s3GCC_${i}.bulk.sorted.hicSummary;
/home/groups/oroaklab/src/homer/bin/analyzeHiC s3GCC_${i} -res 2500000 -raw -cpu 20 > $i.hic.matrix; done"""

p = sp.Popen(cmd, stdin=sp.PIPE, stdout = sp.PIPE, stderr = sp.PIPE,shell=True)
p.wait()