import requests
import json
import sys
import re
import argparse
from __future__ import print_function


# The endpoint that CravatQuerys are submitted to
endpoint = 'http://www.cravat.us/CRAVAT/rest/service/query'


# newline and delimiter values used in the output file
delimiter = "\t"
newline = "\n"


# Defualt indices for intepretting a cravat file's row of data in to a CravatQuery
cr_mapping = {
	'chromosome': 1,
	'position': 2,
	'strand': 3,
	'reference': 4,
	'alternate': 5
}


# The neccessary attributes neeeded to submit a query.
query_keys = [
	'chromosome', 'position', 'strand', 'reference', 'alternate'
]


# Expected response keys from server. Ordered in list so that galaxy output has uniform column ordering run-to-run.
# If cravat server returns additional keys, they are appended to and included in output.
ordered_keys = [
	"Chromosome", "Position", "Strand", "Reference base(s)", "Alternate base(s)",
	"HUGO symbol", "S.O. transcript", "Sequence ontology protein change", "Sequence ontology",
	"S.O. all transcripts", "gnomAD AF", "gnomAD AF (African)", "gnomAD AF (Amrican)",
	"gnomAD AF (Ashkenazi Jewish)", "gnomAD AF (East Asian)", "gnomAD AF (Finnish)",
	"gnomAD AF (Non-Finnish European)", "gnomAD AF (Other)", "gnomAD AF (South Asian)",
	"1000 Genomes AF", "ESP6500 AF (average)", "ESP6500 AF (European American)",
	"ESP6500 AF (African American)", "COSMIC transcript", "COSMIC protein change", 
	"COSMIC variant count [exact nucleotide change]", "cosmic_site_nt", "CGL driver class",
	"TARGET", "dbSNP", "cgc_role", "cgc_inheritance", "cgc_tumor_type_somatic",
	"cgc_tumor_type_germline", "ClinVar", "ClinVar disease identifier", "ClinVar XRef",
	"GWAS Phenotype (GRASP)", "GWAS PMID (GRASP)", "Protein 3D variant"
]


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input',
                            '-i',
                            required = True,
                            help='Input path to a cravat file for querying',)
    parser.add_argument('--output',
                            '-o',
                            default = None,
                            help = 'Output path to write results from query')
    return parser.parse_args()


def format_chromosome(chrom):
	""" : Ensure chromosome entry is propely formatted for use as querying attribute. """
	if chrom[0:3] == 'chr':
		return chrom
	return 'chr' + str(chrom)


def get_query_string(row):
	""" : From a row dict, return a query string for the Cravat server.
		: The row dict is cravat headeres associated to their values of that row.
	"""
	return '_'.join([ row['chromosome'], row['position'], row['strand'], row['reference'], row['alternate'] ])


def query(in_path, out_path):
	""" : From a Cravat the file at in_path, query each line on the Cravat server.
		: Write the response values to file at out_path.
	"""
	with open(in_path, 'r') as in_file, \
	open(out_path, 'w') as out_file:
		for line in in_file:
			try:
				line = line.strip().split('\t')
				# row is dict of cravat col headers assioted values in this line
				row = { header: line[index] for header, index in cr_mapping.items() }
				row['chromosome'] = format_chromosome(row['chromosome'])
				query_string = get_query_string(row)
				call = requests.get(endpoint, params={ 'mutation': query_string })
				if call.status_code != 200 or call.text == "":
					raise requests.RequestException('Bad server response for query="{}". Respone code: "{}", Response Text: "{}"'
						.format(query_string, call.status_code, call.text))
				json_response = json.loads(call.text)
				# See if server returned additional json key-vals not expected in ordered_keys
				for key in json_response:
					if key not in ordered_keys:
						ordered_keys.append(key)
				# Write key in order of ordered_keys to standardize order of output columns
				wrote = False
				for key in ordered_keys:
					if key in json_response:
						val = json_response[key]
					else:
						val = None
					# Standardize format for  numeric values
					try:
						val = float(val)
						val = format(val, ".4f")
					except:
						pass
					if wrote:
						out_file.write(delimiter)
					out_file.write(str(val))
					wrote = True
				out_file.write(newline)		
			except Exception as e:
				print(e, file=sys.stderr)
				continue


if __name__ == "__main__":
	cli_args = get_args()
	if cli_args.output == None:
		base, _ = os.path.split(cli_args.input)
		cli_args.output = os.path.join(base, "cravat_converted.txt") 
	query(cli_args.input, cli_args.output)