#!/usr/bin/python
"""
Anselms Oliv-Index: Auswerten und plotten.
"""

import argparse
import calendar
import collections
import io
import time


def compute_index(code):
	"""gibt einen numerischen Index für einen symbolischen code zurück.

	Das ist schlicht das Verhältnis der o-s zu allen Zeichen (bei einem
	gestrippten string.

	>>> compute_index(" .......   ")
	0.0
	>>> compute_index(".o...o..")
	0.25
	"""
	code = code.strip()
	return code.count("o")/len(code)


def get_timeseries(src_f):
	r"""gibt eine numpy-gerechte Zeitreihe [timestamps], [indexe] zurück.

	src_f (ein file) muss whitespace-getrennt ISO-datum und
	compute-index-kompatiblen code enthalten.

	>>> get_timeseries(io.StringIO("#ignore\n2022-05-07 ooooo.....o."))
	([1651881600], [0.5])
	"""
	timestamps, scores = [], []

	for ln in src_f:
		if ln.startswith("#") or not ln.strip():
			continue

		date_str, code = ln.split()
		timestamps.append(calendar.timegm(time.strptime(date_str, "%Y-%m-%d")))
		scores.append(compute_index(code))

	return timestamps, scores


def decimate(a_list, max_items):
	"""gibt eine Liste zurück, in der bis zu max_items gleichmäßig
	gesampelte Elemente aus a_list sind.

	>>> decimate(list(range(6)), 6)
	[0, 1, 2, 3, 4, 5]
	>>> decimate(list(range(6)), 3)
	[0, 2, 4]
	>>> decimate(list(range(6)), 2)
	[0, 3]
	"""
	if len(a_list)<=max_items:
		return a_list
	
	step = len(a_list)//max_items
	return a_list[::step]


def action_plot(args):
	import matplotlib, matplotlib.figure
	figsize, dpi = (5, 0.5), 200
	fig = matplotlib.figure.Figure(figsize, dpi,
		frameon=False,
		subplotpars=matplotlib.figure.SubplotParams(0,0,1,1,0,0))
	ax = fig.add_subplot()
	ax.set_axis_off()

	with open(args.source_name) as f:
		timestamps, scores = get_timeseries(f)
		colormap = matplotlib.colors.LinearSegmentedColormap("olive", dict(
			red=[(0.0, 1, 0.61), (1,1,0.61)],
			green=[(0.0, 0.99, 0.60), (1,1,0.6)],
			blue=[(0.0, 0.63, 0.21), (1,1,0.21)])).reversed()

		mesh = ax.pcolormesh(
			[scores],
			cmap=colormap)

		to_mesh = mesh.get_transform()
		labels = list(enumerate(timestamps))
		for index, ts in decimate(labels, 15):
			ax.text(index, 0.05,
				time.strftime("%Y-%m-%d", time.gmtime(ts)),
				rotation="vertical",
				verticalalignment="bottom",
				horizontalalignment="left",
				fontsize=5)

		fig.savefig("oliv-index.png")


def smooth_gauss(samples, window_width):
	from scipy import signal
	# pad with the first/last value so we don't rely on zero-padding
	# (yeah, it's not *much* better as it still adds fake signals
	# at the borders, but it's still less misleading in most cases.

	padded = ([samples[0]]*window_width
		+ samples + [samples[-1]]*window_width)
	window = signal.windows.gaussian(window_width, window_width/5)
	window /= sum(window)
	return signal.convolve(padded, window, "same")[window_width:-window_width]


def action_lineplot(args):
	from matplotlib import figure, dates
	import numpy

	fig = figure.Figure(figsize=(7, 3))
	ax = fig.add_subplot()
	with open(args.source_name) as f:
		timestamps, scores = get_timeseries(f)

	timestamps = dates.num2date(numpy.array(timestamps)/86400.)
	scores = smooth_gauss(scores, 40)

	ax.plot(timestamps, scores)
	ax.grid(True)
	fig.autofmt_xdate()
	fig.savefig("oliv-line.png")


def action_heute(args):
	with open(args.source_name) as f:
		parts = f.read()[-1000:].split(" ")
		while not parts[-1]:
			parts.pop()
	
	if not parts:
		raise Exception("Keine erkennbaren Codes")

	print("{:.2f}".format(compute_index(parts[-1])))


def parse_commandline():
	parser  = argparse.ArgumentParser(
		description="Operationen mit dem Oliv-Index")
	parser.add_argument("action", action="store", type=str,
		choices=["test", "plot", "heute", "lineplot"])
	parser.add_argument("--datafile", action="store", type=str,
		metavar="NAME", dest="source_name", default="codes.txt",
		help="Lies codes aus NAME")

	return parser.parse_args()


def action_test(args):
	import doctest
	doctest.testmod()


def main():
	args = parse_commandline()
	globals()["action_"+args.action](args)


if __name__=="__main__":
	main()
