#!/usr/bin/env python3

import argparse
import base64
import io
import operator
import statistics

import matplotlib.pyplot as plt

from grimoire.genome import Reader
from grimoire.io import FASTA_stream, GFF_stream

## Command line stuff ##

extended_help = """
%(prog)s provides an overview of a genome annotation as an HTML page. There
are graphs showing the characteristics of chromosomes and genes as well as
an overview of some of the errors/non-canonical features.
"""

parser = argparse.ArgumentParser(
	description='Genome annotation reporter',
	formatter_class=argparse.RawDescriptionHelpFormatter,
	epilog=extended_help)
parser.add_argument('--fasta', required=True, type=str,
	metavar='<path>', help='path to fasta file')
parser.add_argument('--gff', required=True, type=str,
	metavar='<path>', help='path to GFF (or other) annotation file')
parser.add_argument('--title', required=True, type=str,
	metavar='<path>', help='title of report')
parser.add_argument('--html', required=True, type=str,
	metavar='<path>', help='path to html output file')
arg = parser.parse_args()

class CalfoError(Exception):
	pass

####################
# Global Variables #
####################

FigureNumber = 0 # used for enumerating each figure in the document

########################
## Graphical Routines ##
########################

def stats(values):
	if values:
		return {
			'count' : str(len(values)),
			'min' : str(min(values)),
			'max' : str(max(values)),
			'med' : str(round(statistics.median(values))),
			'mean' : str(round(statistics.mean(values), 3)),
			'stdev' : str(round(statistics.stdev(values), 3)),
		}
	else:
		return {'count':0, 'min':'0', 'max':'0', 'med':'0', 'mean':'0', 'stdev':'0'}

def encode_fig(fig):
	buff = io.BytesIO()
	fig.savefig(buff, format='png', bbox_inches='tight')
	buff.seek(0)
	img = base64.b64encode(buff.read()).decode('utf-8')
	return '<br><img src="data:image/png;base64,{0}"><br>'.format(img)

def histogram(data, title, bins):
	global FigureNumber
	FigureNumber += 1
	
	fig, ax = plt.subplots()
	ax.hist(data, bins=bins)
	ax.set(title='Figure ' + str(FigureNumber) + ': ' + title)
		
	text = encode_fig(fig)
	text += '<br><b>Figure ' + str(FigureNumber) + '</b>: '
	text += str(stats(data))
	
	return text

def bar_chart(data, title):
	global FigureNumber
	FigureNumber += 1
	
	fig, ax = plt.subplots()
	
	overflow = 0
	names = []
	values = []
	for i in sorted(data.items(), key=operator.itemgetter(1), reverse=True):
		if len(names) > 30:
			overflow += i[1]
		else:
			names.append(i[0])
			values.append(i[1])
	
	if overflow:
		names.append('others')
		values.append(overflow)
	names.reverse()
	values.reverse()
	
	ax.barh(names, values)
	ax.set(title='Figure ' + str(FigureNumber) + ': ' + title)
	
	text = encode_fig(fig)
	text += '<b>Figure ' + str(FigureNumber) + '</b>: '
	text += 'Total: ' + str(sum(values)) + ': ' + str(data) + '.<p>'

	return text
	
###############################################################################
# Report Init
###############################################################################

doc = open(arg.html, 'w+')
doc.write('<html>')
doc.write('<head>')
doc.write('<meta charset="utf-8">')
doc.write('<title>' + arg.title + '</title>')
doc.write('</head>')
doc.write('<body>')
doc.write('<h1>' + arg.title + '</h1>')

###############################################################################
# Chromosome Report
###############################################################################

doc.write('<h2>Chromosome Report</h2>')
sizes = {}
fasta = FASTA_stream(arg.fasta)
for entry in fasta:
	sizes[entry.id] = len(entry.seq)
doc.write('Origin: ' + arg.fasta + '<p>')
doc.write(bar_chart(sizes, 'Chromosome Sizes'))

###############################################################################
# Generic Annotation Report
###############################################################################

doc.write('<h2>Genome Annotation Report</h2>')
gff = GFF_stream(arg.gff)
types = {}
for entry in gff:
	if entry.type not in types: types[entry.type] = 0
	types[entry.type] += 1
doc.write(bar_chart(types, 'Feature Types'))

###############################################################################
# Protein-Coding Genes Report
###############################################################################

doc.write('<h2>Protein-coding Genes Report</h2>')

tpg = [] # transcript per gene
ept = [] # exons per transcript
ipt = [] # introns per transcript
u5pt = [] # 5'utrs per transcript
u3pt = [] # 3'utrs per transcript
ilen = [] # intron lengths
elen = [] # exon lengths
u5len = [] # 5'utr lengths
u3len = [] # 3'utr lengths
tlen = [] # tx lengths
clen = [] # cds lengths

# error reporting stuff
gene_ok, gene_err = 0, 0
tx_ok, tx_err = 0, 0
f_ok, f_err = 0, 0
gene_issues = {}
tx_issues = {}
f_issues = {}

genome = Reader(fasta=arg.fasta, gff=arg.gff)
for chr in genome:
	genes = chr.ftable.build_genes()
	
	# stats section
	for gene in genes:
		tpg.append(len(gene.children))
		for tx in gene.transcripts():
				
			## count distributions
			ept.append(len(tx.exons))
			ipt.append(len(tx.introns))
			u5pt.append(len(tx.utr5s))
			u3pt.append(len(tx.utr3s))
								
			#  lengths
			for e in tx.exons: elen.append(e.end - e.beg + 1)
			for u in tx.utr5s: u5len.append(u.end - u.beg + 1)
			for u in tx.utr3s: u3len.append(u.end - u.beg + 1)
			for i in tx.introns: ilen.append(i.end - i.beg + 1)
			clen.append(len(tx.cds_str()))
			tlen.append(tx.end - tx.beg + 1)

	# error section
	for gene in genes:
		if not gene.issues: 
			gene_ok += 1
		else:
			gene_err += 1
			for issue in gene.issues:
				if issue not in gene_issues: gene_issues[issue] = 0
				gene_issues[issue] += 1
		
		# transcript level
		for tx in gene.transcripts():
			if not tx.issues:
				tx_ok += 1
			else:
				tx_err += 1
				for issue in tx.issues:
					if issue not in tx_issues: tx_issues[issue] = 0
					tx_issues[issue] += 1
			
			# feature level
			for f in tx.children:
				if not f.issues:
					f_ok += 1
				else:
					f_err += 1
					for issue in f.issues:
						if issue not in f.issues: f_issues[issue] = 0
						f_issues[issue] += 1

doc.write(histogram(tpg, 'mRNAs per Gene', 10))
doc.write(histogram(ept, 'Exons per mRNA', 10))
doc.write(histogram(ipt, 'Introns per mRNA', 10))
doc.write(histogram(u5pt, 'UTR5s per mRNA', 10))
doc.write(histogram(u3pt, 'UTR3s per mRNA', 10))
doc.write(histogram(tlen, 'mRNA Lengths', 10))
doc.write(histogram(clen, 'CDS Lengths', 10))
doc.write(histogram(elen, 'Exon Lengths', 10))
doc.write(histogram(ilen, 'Intron Lengths', 10))
doc.write(histogram(u5len, 'UTR5 Lengths', 10))
doc.write(histogram(u3len, 'UTR3 Lengths', 10))

###############################################################################
# Protein-Coding Gene Errors : General
###############################################################################

doc.write('<h2>Protein-coding Gene Annotation Errors</h2>')

doc.write('Gene errors: {} / {} = {}'.format(gene_err, gene_ok,
	gene_err /(gene_err + gene_ok)))
doc.write(bar_chart(gene_issues, 'Gene Errors'))

doc.write('mRNA errors: {} / {} = {}'.format(tx_err, tx_ok,
	tx_err /(tx_err + tx_ok)))
doc.write(bar_chart(tx_issues, 'mRNA Errors'))

doc.write('Component errors: {} / {} = {}'.format(f_err, f_ok,
	f_err /(f_err + f_ok)))
doc.write(bar_chart(f_issues, 'Component Errors'))

###############################################################################
# Footer
###############################################################################

FOOTER = """
This document was prepared using <em>fathom</em>,
part of the <em>grimoire</em> tools from the KorfLab.
For more information, see
<a href="https://github.com/KorfLab/grimoire">grimoire</a> at GitHub.
"""

# probably add the date of the report to the footer

doc.write('</body>')
doc.write('<footer>' + FOOTER + '</footer>')
doc.write('</html>')
