#usr/bin/python
import os 
import sqlite3
import csv


def read_files(trait_file):
    eqtls = []
    t_file = open(trait_file, 'r')
    t_eqtls = csv.reader(t_file, delimiter = '\t')
    next(t_eqtls, None)
    for row in t_eqtls:
        eqtls.append(row)

    return eqtls


def process(eqtl_dir):
    traits = {}
    all_tissues = {}
    trait_tissues = {}
    print('Reading trait eQTL files for...')
    if os.path.isdir(eqtl_dir):
        _,_,trait_files = next(os.walk(eqtl_dir), (None, None, []))
        for trait_file in trait_files:
            trait = trait_file[:len(trait_file)-4]
            traits[trait] = read_files(os.path.join(eqtl_dir, trait_file))
    
    for trait in traits:
        trait_tissues[trait] = {}
        for row in traits[trait]:
            gene = row[3]
            tissue = row[7]

            if tissue not in all_tissues:
                all_tissues[tissue] = {'#Genes':0, '%Genes':0.0, 'gene_list':[]}
            if tissue not in trait_tissues[trait]:
                trait_tissues[trait][tissue] = {'#Genes':0, '%Genes':0.0, 'gene_list':[]}
            all_tissues[tissue]['gene_list'].append(gene)
            trait_tissues[trait][tissue]['gene_list'].append(gene)

    summary_file = open('./tissue_summary_gaps.txt', 'w')
    summary = csv.writer(summary_file, delimiter = '\t')
    sheader = ('Tissue', '#Genes', '%Genes(6071)', 'Gene_list')
    summary.writerow(sheader)
    list_dir = './tissue_gene_list'
    if not os.path.isdir(list_dir):
        os.mkdir(list_dir)
    for tissue in all_tissues:
        gene_list = all_tissues[tissue]['gene_list']
        genes = gene_list[0]
        if len(gene_list) > 1:
            for i in range (1, len(gene_list)):
                genes += ' ' + gene_list[i]
        all_tissues[tissue]['#Genes'] = len(gene_list) 
        all_tissues[tissue]['%Genes'] = float(len(gene_list))/6071
        summary.writerow([tissue, all_tissues[tissue]['#Genes'], \
                              all_tissues[tissue]['%Genes'], genes])
        tissue_file = open(os.path.join(list_dir, tissue)+'.txt', 'w')
        #twriter = csv.writer(tissue_file, delimiter = ',')
        #twriter.writerows(gene_list)
        for gene in gene_list:
            tissue_file.write(gene + '\n')
        tissue_file.close()
    summary_file.close()


    for trait in trait_tissues:
        trait_file = open(os.path.join(tissues_dir, trait) + '.txt', 'w')
        twriter = csv.writer(trait_file, delimiter = '\t')
        theader = ('Tissue', '#Genes', '%Genes(#CatalogTissue)', 'Gene_list')
        twriter.writerow(theader)
        for tissue in trait_tissues[trait]:
            gene_list = trait_tissues[trait][tissue]['gene_list']
            genes = gene_list[0]
            if len(gene_list) > 1:
                for i in range (1, len(gene_list)):
                    genes += ', ' + gene_list[i]
            trait_tissues[trait][tissue]['#Genes'] = len(gene_list) 
            trait_tissues[trait][tissue]['%Genes'] = \
                float(len(gene_list)) /all_tissues[tissue]['#Genes']
            twriter.writerow([tissue, trait_tissues[trait][tissue]['#Genes'], \
                              trait_tissues[trait][tissue]['%Genes'], genes])
        trait_file.close()



if __name__ == "__main__":
    eqtl_dir = '../data/sig_eqtls'
    tissues_dir = '../data/tissues'
    if not os.path.isdir(tissues_dir):
        os.mkdir(tissues_dir)
    process(eqtl_dir)


