#!/usr/bin/env python

import os
import re
import sys
import csv
import errno
import requests
import argparse


## Example ids supported by this script
SRA_IDS = ['PRJNA63463', 'SAMN00765663', 'SRA023522', 'SRP003255', 'SRR390278', 'SRS282569', 'SRX111814']
ENA_IDS = ['ERA2421642', 'ERP120836', 'ERR674736', 'ERS4399631', 'ERX629702', 'PRJEB7743', 'SAMEA3121481']
GEO_IDS = ['GSE18729', 'GSM465244']
ID_REGEX = r'^[A-Z]+'
PREFIX_LIST = sorted(list(set([re.search(ID_REGEX,x).group() for x in SRA_IDS + ENA_IDS + GEO_IDS])))


def parse_args(args=None):
    Description = 'Download and create a run information metadata file from SRA/ENA/GEO identifiers.'
    Epilog = """Example usage: python fetch_sra_runinfo.py <FILE_IN> <FILE_OUT>"""

    parser = argparse.ArgumentParser(description=Description, epilog=Epilog)
    parser.add_argument('FILE_IN', help="File containing database identifiers, one per line.")
    parser.add_argument('FILE_OUT', help="Output file in tab-delimited format.")
    parser.add_argument('-pl', '--platform', type=str, dest="PLATFORM", default='', help="Comma-separated list of platforms to use for filtering. Accepted values = 'ILLUMINA', 'OXFORD_NANOPORE' (default: '').")
    parser.add_argument('-ll', '--library_layout', type=str, dest="LIBRARY_LAYOUT", default='', help="Comma-separated list of library layouts to use for filtering. Accepted values = 'SINGLE', 'PAIRED' (default: '').")
    return parser.parse_args(args)


def validate_csv_param(param,validVals,param_desc):
    validList = []
    if param:
        userVals = param.split(',')
        intersect = list(set(userVals) & set(validVals))
        if len(intersect) == len(userVals):
            validList = intersect
        else:
            print("ERROR: Please provide a valid {} parameter!\nProvided values = {}\nAccepted values = {}".format(param_desc,param,','.join(validVals)))
            sys.exit(1)
    return validList


def make_dir(path):
    if not len(path) == 0:
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise


def fetch_url(url,encoding='utf-8'):
    try:
        r = requests.get(url)
    except requests.exceptions.RequestException as e:
        raise SystemExit(e)
    if r.status_code != 200:
        print("ERROR: Connection failed\nError code '{}'".format(r.status_code))
        sys.exit(1)
    return r.content.decode(encoding).splitlines()


def id_to_srx(db_id):
    ids = []
    url = 'https://trace.ncbi.nlm.nih.gov/Traces/sra/sra.cgi?save=efetch&db=sra&rettype=runinfo&term={}'.format(db_id)
    for row in csv.DictReader(fetch_url(url), delimiter=','):
        ids.append(row['Experiment'])
    return ids


def id_to_erx(db_id):
    ids = []
    fields = ['run_accession', 'experiment_accession']
    url = 'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession={}&result=read_run&fields={}'.format(db_id,','.join(fields))
    for row in csv.DictReader(fetch_url(url), delimiter='\t'):
        ids.append(row['experiment_accession'])
    return ids


def gse_to_srx(db_id):
    ids = []
    url = 'https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc={}&targ=gsm&view=data&form=text'.format(db_id)
    gsm_ids = [x.split('=')[1].strip() for x in fetch_url(url) if x.find('GSM') != -1]
    for gsm_id in gsm_ids:
        ids += id_to_srx(gsm_id)
    return ids


def get_ena_fields():
    fields = []
    url = 'https://www.ebi.ac.uk/ena/portal/api/returnFields?dataPortal=ena&format=tsv&result=read_run'
    for row in csv.DictReader(fetch_url(url), delimiter='\t'):
        fields.append(row['columnId'])
    return fields


def fetch_sra_runinfo(FileIn,FileOut,platformList=[],libraryLayoutList=[]):
    total_out = 0
    seen_ids = []; run_ids = []
    header = []
    make_dir(os.path.dirname(FileOut))
    ena_fields = get_ena_fields()
    fin = open(FileIn,'r')
    fout = open(FileOut,'w')
    while True:
        line = fin.readline()
        if line:
            db_id = line.strip()
            match = re.search(ID_REGEX, db_id)
            if match:
                prefix = match.group()
                if prefix in PREFIX_LIST:
                    if not db_id in seen_ids:

                        ids = [db_id]
                        ## Resolve/expand these ids against GEO URL
                        if prefix in ['GSE']:
                            ids = gse_to_srx(db_id)

                        ## Resolve/expand these ids against SRA URL
                        elif prefix in ['GSM', 'PRJNA', 'SAMN', 'SRR']:
                            ids = id_to_srx(db_id)

                        ## Resolve/expand these ids against ENA URL
                        elif prefix in ['ERR']:
                            ids = id_to_erx(db_id)

                        ## Resolve/expand to get run identifier from ENA and write to file
                        for id in ids:
                            url = 'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession={}&result=read_run&fields={}'.format(id,','.join(ena_fields))
                            csv_dict = csv.DictReader(fetch_url(url), delimiter='\t')
                            for row in csv_dict:
                                run_id = row['run_accession']
                                if not run_id in run_ids:

                                    writeID = True
                                    if platformList:
                                        if row['instrument_platform'] not in platformList:
                                            writeID = False
                                    if libraryLayoutList:
                                        if row['library_layout'] not in libraryLayoutList:
                                            writeID = False

                                    if writeID:
                                        if total_out == 0:
                                            header = sorted(row.keys())
                                            fout.write('{}\n'.format('\t'.join(sorted(header))))
                                        else:
                                            if header != sorted(row.keys()):
                                                print("ERROR: Metadata columns do not match for id {}!\nLine: '{}'".format(run_id,line.strip()))
                                                sys.exit(1)
                                        fout.write('{}\n'.format('\t'.join([row[x] for x in header])))
                                        total_out += 1
                                    run_ids.append(run_id)
                        seen_ids.append(db_id)
                else:
                    id_str = ', '.join([x + "*" for x in PREFIX_LIST])
                    print("ERROR: Please provide a valid database id starting with {}!\nLine: '{}'".format(id_str,line.strip()))
                    sys.exit(1)
            else:
                id_str = ', '.join([x + "*" for x in PREFIX_LIST])
                print("ERROR: Please provide a valid database id starting with {}!\nLine: '{}'".format(id_str,line.strip()))
                sys.exit(1)
        else:
            break
            fin.close()
            fout.close()


def main(args=None):
    args = parse_args(args)
    platformList = validate_csv_param(args.PLATFORM,validVals=['ILLUMINA', 'OXFORD_NANOPORE'],param_desc='--platform')
    libraryLayoutList = validate_csv_param(args.LIBRARY_LAYOUT,validVals=['SINGLE', 'PAIRED'],param_desc='--library_layout')
    fetch_sra_runinfo(args.FILE_IN,args.FILE_OUT,platformList,libraryLayoutList)

if __name__ == '__main__':
    sys.exit(main())
