#!/usr/bin/env python3

import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from datetime import date, time, datetime
from dataclasses import dataclass
from glob import glob
from typing import Sequence
from functools import reduce

VOLICU_CELKEM = 8_245_962
OKRSKU_CELKEM = 14_857
kandidati_jmena = {
	1: 'Pavel Fischer',
	2: 'Jaroslav Bašta',
	4: 'Petr Pavel',
	5: 'Tomáš Zima',
	6: 'Danuše Nerudová',
	7: 'Andrej Babiš',
	8: 'Karel Diviš',
	9: 'Marek Hilšer',
	}

@dataclass
class VysledekOkrsku:
	poradi: int
	timestamp: datetime
	kandidati: dict[str, int]
	ucast: dict

	def dopln_vysledek(self):
		assert self.poradi != 0, "Dávka má mít kladné pořadí"
		self.kandidati['NEPLATNÉ'] = int(self.ucast['ODEVZDANE_OBALKY']) - int(self.ucast['PLATNE_HLASY'])
		self.kandidati['NEVOLILI'] = int(self.ucast['ZAPSANI_VOLICI']) - int(self.ucast['ODEVZDANE_OBALKY'])
		# Ztracené (vydané a neodevzdané) obálky neřešíme
		# Nezapočítané hlasy neexistují na úrovni okrsku

def nacti_davku(fn) -> dict[tuple[int, int], VysledekOkrsku]:
	okrsky: dict[tuple[int, int], VysledekOkrsku] = {}

	with open(fn, 'r') as f:
		root = ET.parse(f).getroot()
	NS = {'': r'http://www.volby.cz/prezident/'}

	def zpracuj_okrsek(elem) -> tuple[VysledekOkrsku, tuple[int, int]]:
		idOkrsku = tuple(map(int, (elem.get(x) for x in ['CIS_OBEC', 'CIS_OKRSEK'])))
		kandidati = {}
		for kandidat in elem.findall('./HLASY_OKRSEK', NS):
			cislo = int(kandidat.attrib['PORADOVE_CISLO'])
			hlasy = int(kandidat.attrib['HLASY'])
			jmeno = kandidati_jmena[cislo]
			kandidati[jmeno] = hlasy
		return VysledekOkrsku(
			poradi = int(elem.attrib['PORADI_ZPRAC']),
			timestamp = datetime.fromisoformat(elem.attrib['DATUM_CAS_ZPRAC']),
			ucast = elem.find('./UCAST_OKRSEK', NS).attrib,
			kandidati = kandidati,
			), idOkrsku
	
	for okrsek in root.findall('./OKRSEK', NS):
		vysl, id = zpracuj_okrsek(okrsek)
		vysl.dopln_vysledek()
		okrsky[id] = vysl
	
	return okrsky

def zpracuj_davky(fns: Sequence[str]):
	# Zpracovávací fáze: načteme dávku, nahradíme předchozí instance pro okrsky
	# Počítací fáze: sečteme hlasy za všechny okrsky, vyrobíme dump nebo koláč ke každému času
	# Uložíme pod timestampem do nějaké složky.
	# Tyto dvě fáze ale musí probíhat zároveň, pokud chceme mít vykreslený graf za každý okamžik předání
	# Jiný program: vezme dumpy/koláče a naskládá je správně do videa (případně provede jinou analýzu)

	okrsky: dict[tuple[int, int], VysledekOkrsku] = {}
	celkovy_vysledek: dict[str, int] = {}
	for fn in fns:
		print(f'Zpracovávám dávku {fn}')
		nove = nacti_davku(fn)
		podle_casu = sorted(nove.items(), key=lambda x: x[1].timestamp)
		for okrsek, vysledek_okrsku in podle_casu:
			# Přidat okrsek do dat, přepočítat a exportovat výsledek
			if okrsek in okrsky:
				assert vysledek_okrsku.poradi > okrsky[okrsek].poradi
				# Odečteme předchozí výsledek tohohle okrsku
				for k in okrsky[okrsek].kandidati:
					celkovy_vysledek[k] -= okrsky[okrsek].kandidati[k]
			okrsky[okrsek] = vysledek_okrsku
			for k in vysledek_okrsku.kandidati:
				if k not in celkovy_vysledek: celkovy_vysledek[k] = 0
				celkovy_vysledek[k] += vysledek_okrsku.kandidati[k]
			# Počítání dat:
			ts = vysledek_okrsku.timestamp
			celkovy_vysledek['NEZAPOČÍTANÉ'] = 0 # del který funguje i v první iteraci :-)
			celkovy_vysledek['NEZAPOČÍTANÉ'] = VOLICU_CELKEM - sum(celkovy_vysledek.values())
			# Vizualizace / export
			# FIXME: Zkopírováno, negenerické, fuj. Má používat nějaký dedikovaný spoolečný kód.
			plt.clf()
			# Pro srovnání: zpracování okrsků
			# TODO: WTF: když tohle nakreslím až po kandidátech, tak se
			# kandidáti vyrenderují mimo bbox.
			okrsku_secteno = len(okrsky.keys())
			okrsku_zbyva = OKRSKU_CELKEM - okrsku_secteno
			plt.pie(
				[okrsku_secteno, okrsku_zbyva],
				colors=['#000088', '#cccccc'],
				center=(1, -1), radius=0.3,
				)

			# Zastoupení kandidátů
			order = (
				'Petr Pavel',
				'Danuše Nerudová',
				'Marek Hilšer',
				'Pavel Fischer',
				'Karel Diviš',
				'Tomáš Zima',
				'Jaroslav Bašta',
				'Andrej Babiš',
				'NEPLATNÉ',
				'NEVOLILI',
				'NEZAPOČÍTANÉ',
				)
			# FIXME: průběžné výsledky nemají některé kandidáty. Tohle bychom neměli opravovat tady.
			x = [celkovy_vysledek.get(i, 0) for i in order]
			print(ts, x)
			labels = order
			label = str(ts)
			colors = {
				'Petr Pavel': '#627210',
				'Danuše Nerudová': '#811367',
				'Andrej Babiš': '#262161',
				'Jaroslav Bašta': '#B51119',
				'Marek Hilšer': '#CA834E',
				'Pavel Fischer': '#244C76',
				'Karel Diviš': '#3B6E5D',
				'Tomáš Zima': '#E5DE1A',
				'NEPLATNÉ': '#000000',
				'NEVOLILI': '#666666',
				'NEZAPOČÍTANÉ': '#CCCCCC',
				}
			plt.pie(x, labels=labels, colors=[colors[i] for i in order], autopct='%1.3f %%')
			plt.text(0, -1.2, label, ha='center')

			# save
			fn = f'progress_png/{ts}.png'
			plt.savefig(fn)


zpracuj_davky(
	sorted(glob('davky/*.xml'))
	)