#!/usr/bin/env python

import os
import re
from genomicode import genefinder, timer, Matrix, matrixlib
from genomicode import parselib, filelib, arrayannot

# retrieve_all_dates
# retrieve_diseases

# Not quite a good idea.
datatype_match = {
    'RSEM_genes' : (
        'Merge_rnaseqv2__illuminahiseq_rnaseqv2__unc_edu__Level_3__' +
        'RSEM_genes_normalized__data.Level_3'),
    'RSEM_exons' : 'exon_expression__data.Level_3',
    'RSEM_isoforms' : (
        'Merge_rnaseqv2__illuminahiseq_rnaseqv2__unc_edu__Level_3__' +
        'RSEM_isoforms_normalized__data.Level_3'),
    
    'humanmethylation450' : (
        'Merge_methylation__humanmethylation450__jhu_usc_edu__Level_3__' +
        'within_bioassay_data_set_function__data.Level_3'),
    'mirnaseq' : (
        'Merge_mirnaseq__illuminahiseq_mirnaseq__bcgsc_ca__Level_3__' +
        'miR_gene_expression__data.Level_3'),
    
    'clinical' : 'Merge_Clinical.Level_1',
    'rppa' : '.RPPA_AnnotateWithGene.Level_3',
    'cnv_gistic2' : 'CopyNumber_Gistic2.Level_4',
    
    'agilent' : [
        'Merge_transcriptome__agilentg4502a_07_1__unc_edu__Level_3__unc_lowess_normalization_gene_level__data.Level_3',
        'Merge_transcriptome__agilentg4502a_07_2__unc_edu__Level_3__unc_lowess_normalization_gene_level__data.Level_3',
        ],
    'affymetrix' : (
        'Merge_transcriptome__ht_hg_u133a__broad_mit_edu__Level_3__'
        'gene_rma__data.Level_3'),
    "mutation" : "Mutation_Packager_Calls.Level_3",
    }

datatype2resource = {
    "RSEM_genes" : "stddata",
    "RSEM_exons" : "stddata",
    "humanmethylation450" : "stddata",
    "mirnaseq" : "stddata",
    "clinical" : "stddata",
    "rppa" : "stddata",
    "RSEM_isoforms" : "stddata",
    "cnv_gistic2" : "analyses",
    "agilent" : "stddata",
    "affymetrix" : "stddata",
    "mutation" : "stddata",
    }

resources = ["stddata", "analyses"]

URL2HTML = {}
def read_url(url):
    global URL2HTML
    import urllib2
    
    if url not in URL2HTML:
        #print "Reading %s" % url; sys.stdout.flush()
        timer.wait(2, 'tcga')
        response = urllib2.urlopen(url)
        html = response.read()
        URL2HTML[url] = html
    return URL2HTML[url]

def retieve_dates_from_resource(resource):
    #import pprint  ###
    # Return a dictionary of the date -> url with given resource.
    #url points to an HTML page for the runs for a specific date.  It has a table of the
    # diseases and HREFs to the data for the diseases at that date.
    #download_url = None
    url = None
    if resource == 'stddata':
        dashboard_url = (
        'https://confluence.broadinstitute.org/display/GDAC/Dashboard-Stddata')
        url = 'http://gdac.broadinstitute.org/runs/info/stddata__runs_list.html'
    elif resource == 'analyses':
        dashboard_url =(
            'https://confluence.broadinstitute.org/display/GDAC/Dashboard-Analyses')
        url = 'http://gdac.broadinstitute.org/runs/info/analyses__runs_list.html'
    else:
        raise ValueError('we do not recognize the resource %s' %resource)
    # The latest data is listed on the "Dashboard" page.  The old data
    # is shown on a different one.  Download them separately.
    all_dates = {}
    urls = read_and_extract_urls(dashboard_url)
    diseases, date = get_disease_lastest(urls,resource)
    date = date.replace("_", "")
    all_dates[date] = dashboard_url

    # Read the old data.
    html = read_url(url)
    link_tags = parselib.get_tags_and_contents(html,'a')
    for i in link_tags:
        url = parselib.get_href(i)
        date = re.search(r'[0-9]{4}_[0-9]{2}_[0-9]{2}', url)
        if not date:
            continue
        date = date.group(0)
        date = date.replace("_", "")
        all_dates[date] = url
    assert all_dates, "No dates found"
    return all_dates


def retrieve_all_dates():
    # Return a dictionary of the resource-> date and date is a dict which date -> url.
    # url points to an HTML page for the runs for a specific date.  It has a table of the
    # diseases and HREFs to the data for the diseases at that date.
    resources = ['stddata','analyses']
    all_dates = {}  # resource->date
    all_dates_list = []
    for resource in resources:
        all_date_from_resource = retieve_dates_from_resource(resource)
        all_dates[resource] = all_date_from_resource
        all_dates_list.extend(all_date_from_resource.keys())
    all_dates_list = list(set(all_dates_list))
    
    all_dates_list = sorted(all_dates_list)
    return all_dates, all_dates_list


def extract_all_hrefs(html):
    x = parselib.get_tags_and_contents(html,'a')
    x = [parselib.get_href(x) for x in x]
    x = [x for x in x if x]
    x = [x for x in x if not x.startswith("#")]
    return x

def retrieve_disease_resource(date,resource):
    all_dates_from_resource = retieve_dates_from_resource(resource)
    if date not in all_dates_from_resource:
        return None
    url = all_dates_from_resource[date]
    # URL:
    # http://gdac.broadinstitute.org/runs/stddata__2014_07_15/data/ACC/20140715
    pattern = re.compile(
        r'http://%s/runs/%s__[0-9_]{10}/data/([A-Z]+)/([0-9]{8})' % (
            "gdac.broadinstitute.org",resource), re.IGNORECASE)
    diseases = []
    html = read_url(url)
    for href in extract_all_hrefs(html):
        m = pattern.match(href)
        if not m:
            continue
        x = m.group(1)
        diseases.append(x)
    return diseases

def retrieve_diseases(date):
    all_diseases = []
    for resource in resources:
        disease = retrieve_disease_resource(date,resource)
        if disease:
            all_diseases.extend(disease)
    all_diseases = list(set(all_diseases))
    assert all_diseases, "could not find diseases"
    return all_diseases


def download_file(disease, date, datatype):
    assert len(date) == 8
    long_date = "%s_%s_%s" % (date[:4], date[4:6], date[6:8])
    resource = datatype2resource[datatype]
    # URL:
    # http://gdac.broadinstitute.org/runs/stddata__2014_07_15/data/ACC/20140715
    link = "http://%s/runs/%s__%s/data/%s/%s/" % (
        "gdac.broadinstitute.org", resource, long_date, disease, date)
    newlinks = get_all_datas_on_page(link)
    if isinstance(datatype_match[datatype],str):
        match_items = [datatype_match[datatype]]
    elif isinstance(datatype_match[datatype],list):
        match_items = datatype_match[datatype]
    else:
        raise ValueError('cannot recognize datatype in datatype_match dict')
    resultlinks = []
    for newlink in newlinks:
        for match_item in match_items:
            if match_item in newlink and disease+'-FFPE' not in newlink:
                data = read_url(link+newlink)
                with open(newlink, "wb") as code:
                     code.write(data)
                print 'finished download %s' % newlink
                resultlinks.append(newlink)
    if resultlinks:
        return resultlinks
    raise AssertionError, "download fails"


def get_data_type_resource(disease, date, resource):
    # URL:
    # http://gdac.broadinstitute.org/runs/stddata__2014_07_15/data/ACC/20140715
    long_date = "%s_%s_%s" % (date[:4], date[4:6], date[6:8])
    link = "http://%s/runs/%s__%s/data/%s/%s/" % (
        "gdac.broadinstitute.org", resource, long_date, disease, date)
    newlinks = get_all_datas_on_page(link)
    result = []
    for newlink in newlinks:
        for datatype in datatype_match:
            if datatype2resource[datatype] != resource:
                continue
            # Can be string or list of strings.
            names = datatype_match[datatype]
            if type(names) is type(""):
                names = [names]
            found = False
            for name in names:
                if newlink.find(name) >= 0:
                    found = True
                    break
            if not found:
                continue
            #if datatype_match[datatype] not in newlink:
            #    continue
            # Brittle
            if disease+'-FFPE' in newlink:
                continue
            assert datatype not in result, "dup datatype"
            result.append(datatype)
    return result
    
def get_data_type(disease, date):
    assert len(date) == 8
    #long_date = "%s_%s_%s" % (date[:4], date[4:6], date[6:8])
    result = []
    for resource in resources:
       result.extend(get_data_type_resource(disease,date,resource))
    return result
                
            
def get_disease_lastest(urls,resource):
    # Return list of diseases, latest_date
    diseases = []
    pattern = re.compile(
        r'http://gdac.broadinstitute.org/runs/%s__[0-9_]*/[A-Z]*.html'%resource)
    lastest = None
    for url in urls:
        disease_link = re.findall(pattern, url)
        for link in disease_link:
            x = re.search(r'[A-Z]+', link)
            if not x:
                continue
            disease = x.group(0)
            assert disease not in diseases
            diseases.append(disease)

            x = re.search(r'[0-9]{4}_[0-9]{2}_[0-9]{2}', link)
            assert x, "date not found"
            date = x.group(0)
            assert not lastest or lastest == date
            lastest = date
    return diseases, lastest 


def get_all_datas_on_page(page):
    html = read_url(page)
    links = []
    link_tags = parselib.get_tags_and_contents(html,'a')
    for i in link_tags:
        url = parselib.get_href(i)
        # This is brittle.
        if url.startswith("gdac.broadinstitute.org_") and url.endswith(".gz"):
            links.append(url)
    return links


def read_and_extract_urls(page):
    # Read a webpage a return a list of the URLs on that page.
    html = read_url(page)
    x = parselib.get_tags_and_contents(html,'a')
    x = [parselib.get_href(x) for x in x]
    x = [x for x in x if x]
    return x

def merge_files(input_list, outfile):
    """input two files and merge, write to the outfile"""
    import arrayio

    assert len(input_list) == 2
    A_file = input_list[0]
    B_file = input_list[1]
    M_A = arrayio.read(A_file)
    M_B = arrayio.read(B_file)
    assert arrayio.tab_delimited_format.is_matrix(M_A)
    assert arrayio.tab_delimited_format.is_matrix(M_B)
    [M_A, M_B] = matrixlib.align_rows(M_A, M_B)
    assert M_A.nrow() > 0, 'there is no common genes between two files'
    X = []
    for i in range(M_A.dim()[0]):
        x = M_A._X[i] + M_B._X[i]
        X.append(x)
    row_names = M_A._row_names
    row_order = M_A._row_order
    col_names = {}
    for name in M_A._col_names:
        if name not in M_B._col_names:
            continue
        newsample_list = []
        for sample in M_B._col_names[name]:
            if sample in M_A._col_names[name]:
                newsample = sample + '_2'
            else:
                newsample = sample
            newsample_list.append(newsample)
        #x = M_A._col_names[name] + M_B._col_names[name]
        x = M_A._col_names[name] + newsample_list
        col_names[name] = x
    M_c = Matrix.InMemoryMatrix(X, row_names, col_names, row_order)
    handle = file(outfile,'w')
    arrayio.tab_delimited_format.write(M_c, handle)
    handle.close()


def merge_rppa_files(in_files, out_file):
    import shutil
    from genomicode import filelib
    
    assert len(in_files) == 2
    x1 = [x for x in in_files if x.endswith(".antibody_annotation.txt")]
    x2 = [x for x in in_files if x.endswith(".rppa.txt")]
    assert len(x1) == 1
    assert len(x2) == 1
    annotation_file = x1[0]
    data_file = x2[0]

    # Actually, just return the data_file.  It contains all the
    # information we need.
    shutil.copy2(data_file, out_file)
    return

    # OV.antibody_annotation.txt
    # Gene Name  Composite Element REF
    # YWHAB      14-3-3_beta
    # YWHAE      14-3-3_epsilon
    # YWHAZ      14-3-3_zeta
    # EIF4EBP1   4E-BP1
    # EIF4EBP1   4E-BP1_pS65

    # OV.rppa.txt
    # Composite.Element.REF     TCGA-04-1335-01A-21-1561-20
    # YWHAB|14-3-3_beta         -0.00855276625000018
    # YWHAE|14-3-3_epsilon      0.05985423025
    # YWHAZ|14-3-3_zeta         -0.04074335825
    # EIF4EBP1|4E-BP1           -0.62276845725
    # EIF4EBP1|4E-BP1_pS65      0.00776960074999994
    # EIF4EBP1|4E-BP1_pT37_T46  -0.04959447325

    # Make sure these files are aligned properly.
    M1 = [x for x in filelib.read_cols(annotation_file)]
    M2 = [x for x in filelib.read_cols(data_file)]
    assert M1 and M2
    assert M1[0][0] == "Gene Name"
    assert M1[0][1] == "Composite Element REF"
    assert M2[0][0] == "Composite.Element.REF"
    assert len(M1) == len(M2)

    # Make sure the header names don't conflict.
    M1[0][1] = "Antibody"

    for i in range(1, len(M1)):
        name1 = M1[i][0]
        x = M2[i][0]
        x = x.split("|")
        assert len(x) == 2
        name2, antibody = x
        assert name1 == name2

    M = []
    for i in range(len(M1)):
        x = M1[i] + M2[i]
        M.append(x)

    handle = open(out_file, 'w')
    for x in M:
        print >>handle, "\t".join(x)
    
    
def merge_mutation_files(in_files, out_file):
    # Merge the maf files.
    x = [x for x in in_files if x.endswith(".maf.txt")]
    assert x, "No maf files"
    in_files = x

    outhandle = open(out_file, 'w')
    
    header = None
    for filename in in_files:
        handle = open(filename)
        line = handle.readline()
        if header is None:
            header = line
            print >>outhandle, line,
        assert line == header

        for line in handle:
            print >>outhandle, line,

            
def extract_and_merge_files(gzfile_list, datatype, resource):
    assert gzfile_list is not None
    result = []
    for gzfile in gzfile_list:
        # BUG: This fails if passing an unzipped directory.
        if not gzfile.endswith('tar.gz'):
            return gzfile
        import tarfile
        tfile = tarfile.open(gzfile, 'r:gz')
        gzname = os.path.split(gzfile)[-1]
        newdir = os.path.join(os.getcwd(),gzname[:-7])
        tfile.extractall(newdir)
        folders = os.listdir(newdir)
        folder = None
        for folder in folders:
            if folder == '._.DS_Store':
                continue
        directory = os.path.join(newdir, folder)
        assert os.path.exists(directory)
        files = os.listdir(directory)
        if resource == 'stddata':
            for filename in files:
                if filename.endswith('txt') and filename != 'MANIFEST.txt':
                    result.append(os.path.join(directory,filename))
        elif resource == 'analyses':
            for filename in files:
                # BAD.  Hard-coded for CNV.
                if filename =='all_data_by_genes.txt':
                    result.append(os.path.join(directory,filename))
        else:
            raise ValueError('not recoginzed resource %s' % resource)

    if datatype == "rppa":
        newname = "rppa.txt"
        merge_rppa_files(result, newname)
        return newname
    elif datatype == "mutation":
        newname = "mutation.txt"
        merge_mutation_files(result, newname)
        return newname
    # If only one file found, then return it.
    elif len(result) == 1:
        return result[0]
    # If two files found, then merge them and return.
    elif len(result) == 2:
        x = os.path.split(result[0])[-1]
        # Should not hard-code this.
        s = "Level_3__unc_lowess_normalization_gene_level__data"
        x = x.replace(
            'agilentg4502a_07_1__unc_edu__%s' % s,
            'agilentg4502a_07__unc_edu__%s' % s)
        newname = x
        merge_files(result, newname)
        return newname
    raise ValueError('extract_and_merge_files can only handle two files')


def format_mutation_packager(filename, outfile):
    import shutil
    shutil.copy2(filename, outfile)


def format_firehose_rsem(filename, output):
    import arrayio

    HYB_REF = "Hybridization REF"
    GENE_ID = "gene_id"
    DATA = arrayio.read(filename)
    assert DATA._row_order == [HYB_REF]
    assert DATA._col_order == ["_SAMPLE_NAME", GENE_ID]
    genes = DATA.row_names(HYB_REF)
    gene_symbols = [None] * len(genes)
    gene_ids = [None] * len(genes)
    for i in range(len(genes)):
        x = genes[i].split("|")
        assert len(x) == 2
        gene_symbol, gene_id = x
        if gene_symbol == "?":
            gene_symbol = ""
        gene_ids[i] = gene_id
        gene_symbols[i] = gene_symbol
    f = file(output,'w')
    header = ["Gene ID", "Gene Symbol"] + DATA.col_names("_SAMPLE_NAME")
    f.write("\t".join(header)+'\n')
    for i in range(DATA.nrow()):
        x = [gene_ids[i], gene_symbols[i]] + DATA._X[i]
        assert len(x) == len(header)
        f.write("\t".join(map(str, x))+'\n')
    f.close()


def format_firehose_exonexp(filename, output):
    import arrayio
    
    HYB_REF = "Hybridization REF"
    GENE_ID = "exon"
    DATA = arrayio.read(filename)
    assert DATA._row_order == [HYB_REF]
    assert DATA._col_order == ["_SAMPLE_NAME", GENE_ID]
    # Line 2 of file is:
    # exon raw_counts median_length_normalized RPKM [...]
    # Only want RPKM columns.
    col_headers = DATA.col_names(GENE_ID)
    I_sig = []
    for i in range(len(col_headers)):
        if i % 3 == 2:
            assert col_headers[i] == "RPKM"
            I_sig.append(i)
        else:
            assert col_headers[i] != "RPKM"
    genes = DATA.row_names(HYB_REF)
    chroms = [None] * len(genes)
    starts = [None] * len(genes)
    ends = [None] * len(genes)
    strands = [None] * len(genes)
    for i in range(len(genes)):
        x = genes[i].split(":")
        assert len(x) == 3
        chrom, x, strand = x
        x = x.split("-")
        assert len(x) == 2
        start, end = x
        chroms[i] = chrom
        starts[i] = start
        ends[i] = end
        strands[i] = strand
    x = DATA.col_names("_SAMPLE_NAME")
    sample_names = [x[i] for i in I_sig]
    f = file(output,'w')
    header = ["Chrom", "Start", "End", "Strand"] + sample_names
    f.write("\t".join(header)+'\n')
    prev_line = None
    for i in range(DATA.nrow()):
        sig = [DATA._X[i][j] for j in I_sig]
        x = [chroms[i], starts[i], ends[i], strands[i]] + sig
        assert len(x) == len(header)
        if x == prev_line:
            continue
        prev_line = x
        f.write("\t".join(map(str, x))+'\n')
    f.close()

def annotate_firehose_methylation(filename, output):
    f=file(filename,'r')
    text = f.readlines(2)
    f.close()
    handle = text[1].split('\t')
    assert handle[:5] == ['Composite Element REF','Beta_value',
                           'Gene_Symbol','Chromosome','Genomic_Coordinate']
    f=file(filename,'r')
    all_symbols = {}
    symbols=[]
    for i in f.readlines():
        words = i.split('\t')
        symbol = words[2]
        symbols = symbol.split(";")
        for x in symbols:
            all_symbols[x] = 1
    f.close()
    all_symbols = sorted(all_symbols)
    #Look up all the symbols in the genefinder.
    symbol2id = {}
    genes = genefinder.find_many_genes(all_symbols, tax_id=9606)
    for (symbol, gene) in zip(all_symbols, genes):
        gene_id = gene[1]
        if gene_id is None:
            gene_id = ""
        symbol2id[symbol] = gene_id
    handle = filelib.read_row(filename, header=1)
    samples_names = [handle._header[i] for i in range(len(handle._header)) if not (i-1)%4]
    header = ["Probe.ID", "Gene.ID", "Gene.Symbol", "Chromosome",
              "Genomic.Coordinate"] + samples_names
    f = file(output, 'w')
    f.write("\t".join(header)+'\n')
    with open(filename) as FileObj:
        for lines in FileObj:
            if lines.startswith('Hybridization REF') or lines.startswith('Composite Element REF'):
                continue
            items = lines.split('\t')
            probe_id = items[0]
            Gene_symbols = items[2]
            Chormosome = items[3]
            Genomic_coordinate = items[4]
            values = [items[i] for i in range(len(items)) if not (i-1)%4]
            symbols = Gene_symbols.split(";")
            ids = [symbol2id.get(x, "") for x in symbols]
            #gene_symbol = ";".join(symbols)
            gene_id = ";".join(map(str, ids))
            row = [probe_id,gene_id,Gene_symbols,Chormosome,Genomic_coordinate]+values
            assert len(row) == len(header)
            f.write("\t".join(map(str, row))+'\n')
    f.close()


def format_firehose_mirna(filename, output):
    matrix = [x for x in filelib.read_cols(filename)]
    HYB_REF = "Hybridization REF"
    GENE_ID = "miRNA_ID"
    assert matrix
    assert matrix[0][0] == HYB_REF
    assert matrix[1][0] == GENE_ID
    header0 = matrix[0]
    header1 = matrix[1]
    for i in range(1, len(header1), 3):
        assert header1[i] == "read_count"
        assert header1[i+1] == "reads_per_million_miRNA_mapped"
        assert header1[i+2] == "cross-mapped"

    sample_name = [header0[i] for i in range(2, len(header0), 3)]
    header = ["miRNA ID"] + sample_name
    f = file(output, 'w')
    f.write("\t".join(header)+'\n')
    for i in range(2, len(matrix)):
        x = [matrix[i][j] for j in range(2, len(matrix[i]), 3)]
        x = [matrix[i][0]] + x
        assert len(x) == len(header)
        f.write("\t".join(x)+'\n')
    f.close()

def format_firehose_rppa(filename, output):
    COMP_REF = "Composite.Element.REF"
    COMP_REF_H = "Composite_Element_REF"
    iter = filelib.read_row(filename, header=1)
    #assert iter._header[0] == COMP_REF
    f = file(output, 'w')
    header = ["Gene Symbol", "Gene ID", "Antibody"] + iter._header[1:]
    f.write("\t".join(header)+'\n')
    for d in iter:
        assert hasattr(d, COMP_REF_H)
        x = getattr(d, COMP_REF_H)
        x = x.split("|")
        assert len(x) == 2
        x, antibody = x
        gene_symbols = [x.strip() for x in x.split()]
        x = genefinder.find_many_genes(gene_symbols, tax_id="9606")
        x = [x[1] for x in x]
        x = [x for x in x if x]
        gene_ids = x
        if gene_symbols == ["CDC2"]:
            gene_ids = ["983"]
        
        gene_symbol_str = ";".join(gene_symbols)
        gene_id_str = ";".join(map(str, gene_ids))
        x = [gene_symbol_str, gene_id_str, antibody] + d._cols[1:]
        assert len(x) == len(header)
        f.write("\t".join(map(str, x))+'\n')
    f.close()
    
def format_firehose_gistic(filename, output):
    f = file(output, 'w')
    iter = filelib.read_row(filename, header=1)
    header = ["Gene ID", "Gene Symbol"] + iter._header[2:]
    print >>f, "\t".join(header)
    for d in iter:
        gene_symbol = d.Gene_Symbol
        gene_id = d.Locus_ID
        x = [gene_id, gene_symbol] + d._cols[2:]
        assert len(x) == len(header)
        print >>f, "\t".join(map(str, x))
    f.close()

    
def format_rsem_isoforms(txt_file, outfile):
    import arrayio
    from genomicode import arrayplatformlib
    
    M = arrayio.read(txt_file)
    # detect platform
    x = arrayplatformlib.score_matrix(M, min_score=0.8)
    assert x, "Cannot identify platform."
    header, platform = x.header, x.platform_name
    
    probe_ids = M.row_names(header)
    #if kg5, convert to kg7
    if platform == 'UCSC_human_hg19_kg5':
        new_platform = 'UCSC_human_hg19_kg7'
        kg7_ids = arrayannot.convert_probe_ids(probe_ids,new_platform)
        kg7_header = 'Hybridization REF kg7'
        M = make_matrix_new_ids(M, kg7_ids,kg7_header,1) 
        probe_ids = M.row_names(kg7_header)   
    # add LocusLink ids
    LocusLink_ids = arrayannot.convert_probe_ids(probe_ids,'Entrez_ID_human')