package main

import (
	"bufio"
	"compress/bzip2"
	"compress/gzip"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"io"
	"math"
	"os"
	"path"
	"strconv"
	"time"

	//"github.com/pkg/profile"

	xmlparser "github.com/tamerh/xml-stream-parser"
	// Dá se získat pomocí "go get -u github.com/tamerh/xml-stream-parser"
)

// Funkce na kontrolu jestli má element daný tag
func hasTag(el *xmlparser.XMLElement, name string) bool {
	for _, tag := range el.Childs["tag"] {
		if tag.Attrs["k"] == name {
			return true
		}
	}
	return false
}

// Pomocná funkce vykonávající pravidelně jinou funkci
func doEvery(d time.Duration, f func()) {
	for range time.Tick(d) {
		f()
	}
}

func die(exitcode int, format string, a ...interface{}) {
	fmt.Fprintf(os.Stderr, format, a...)
	os.Exit(1)
}

func main() {
	if len(os.Args) != 4 {
		die(1, "Usage: %s <file.xml.gz> <width_in_px> <output_image_name.png>\n", os.Args[0])
	}
	heatmapWidth, _ := strconv.Atoi(os.Args[2])
	outputImageName := os.Args[3]

	// 1. Otevřeme soubor
	file, err := os.Open(os.Args[1])
	if err != nil {
		die(1, "Cannot open file: %v\n", err)
	}

	// 2. Připravíme si funkci na resetování readeru, protože budeme muset XML soubor projít vícekrát
	var archiveReader io.Reader
	var parser *xmlparser.XMLParser

	resetReader := func(elements ...string) {
		if _, err := file.Seek(0, 0); err != nil {
			die(1, "Cannot seek to beginning of the file")
		}
		// Podle typu dat (podle přípony) si pořídíme buď gzip nebo bzip2 reader
		switch path.Ext(os.Args[1]) {
		case ".gz", "gzip":
			if archiveReader, err = gzip.NewReader(file); err != nil {
				die(1, "Cannot open the gzip reader: %v", err)
			}
		case ".bz2":
			archiveReader = bzip2.NewReader(file)
		default:
			die(1, "Unknown typ of the archive: %s\n", os.Args[1])
		}

		// Dopustíme se triku - odzipování dat spustíme v samostatném vlákně
		// (u .gz souborů to má vliv na zrychlení tak 10%, protože gzip knihovna je velmi
		//  rychlá a "brzdou" je parsování XML, u .bz2 je zrychlení o více než třetinu,
		//  protože bzip knihovna v aktuálním Go ještě není tolik odladěná a brzdí ona)
		r, w := io.Pipe()
		go func() {
			// write everything into the pipe. Decompression happens in this goroutine.
			io.Copy(w, archiveReader)
			w.Close()
		}()
		parser = xmlparser.NewXMLParser(bufio.NewReaderSize(r, 1024*1024*256), elements...)
	}

	// 3. Vyrobíme si pole na pamatování si bodů budov.
	// Musíme si dát pozor na velikost a ukládat je co nejúspornějším způsobem,
	// zároveň v nich však budeme chtít vyhledávat. Takže pole int64 by sice
	// bylo nejlepší, ale map[int64]void (kde void si definujeme jako prázdnou
	// strukturu) bude sice mít nějaký overhead, ale pořád by to mělo být
	// v průměru do 10B na jedno 8B int64 číslo.
	type void struct{}
	buildingsToDisplay := map[int]void{}
	placedBuildings := 0 // počítadlo kolik budov jsme již umístili do heatmapy
	// Proměnné na spočítání obdélníku pro kreslení heatmapy
	lonMin, lonMax := math.Inf(1), math.Inf(-1)
	latMin, latMax := math.Inf(1), math.Inf(-1)

	start := time.Now()
	phaseStart := start
	// Funkce na vypisování statistiky, spustíme ji tentokrát na pozadí
	printStat := func() {
		megabytes := float64(parser.TotalReadSize) / 1024 / 1024
		elapsed := time.Now().Sub(phaseStart)
		fmt.Fprintf(
			os.Stderr,
			"Elapsed: %-20s Read: %10.2fMB (%.2fMB/s)\tFound building: %-10d Placed buildings: %-10d\r",
			elapsed, megabytes, megabytes/elapsed.Seconds(), len(buildingsToDisplay), placedBuildings,
		)
	}
	go doEvery(500*time.Millisecond, printStat)

	////////////////////////////////////////////////////////////////////////
	// PRVNÍ PRŮCHOD - získání prvního bodu pro každou budovu //////////////

	fmt.Fprintf(os.Stderr, "[Phase 1] Getting buildings\n")

	// 1. Pořídíme si reader od začátku souboru a budeme chtít teď jen tagy cest
	resetReader("node", "way")
	parser.SkipElements([]string{"node"}).SkipOuterElements()

	// 2. Najdeme všechny budovy a uložíme si z nich vždy první bod + zjistíme
	// minimální a maximální lon/lat
	for xml := range parser.Stream() {
		switch xml.Name {
		case "node":
			lat, _ := strconv.ParseFloat(xml.Attrs["lat"], 64)
			lon, _ := strconv.ParseFloat(xml.Attrs["lon"], 64)
			lonMin = math.Min(lonMin, lon)
			lonMax = math.Max(lonMax, lon)
			latMin = math.Min(latMin, lat)
			latMax = math.Max(latMax, lat)
		case "way":
			if hasTag(xml, "building") {
				if len(xml.Childs["nd"]) > 0 {
					ID, _ := strconv.Atoi(xml.Childs["nd"][0].Attrs["ref"])
					buildingsToDisplay[ID] = void{}
				}
			}
		}
	}

	// 3. Spočítáme velikost heatmapy
	lonSpan := lonMax - lonMin
	latSpan := latMax - latMin
	cellSize := lonSpan / float64(heatmapWidth)
	heatmapHeight := int(math.Ceil(latSpan / cellSize))

	printStat()
	fmt.Fprintf(os.Stderr, "\nFound %d buildings\n", len(buildingsToDisplay))

	////////////////////////////////////////////////////////////////////////
	// DRUHÝ PRŮCHOD ///////////////////////////////////////////////////////

	fmt.Fprintf(os.Stderr, "[Phase 2] Getting locations of all buildings and placing on the heatmap\n")

	// 1. Resetujeme se na začátek souboru a připravíme se na další průchod
	resetReader("node")
	parser.SkipElements([]string{"way", "relation"}).SkipOuterElements()

	// 2. Vytvoříme si 2D pole do kterého budeme počítat počty budov
	heatmapMax := 0
	heatmapData := make([][]int, heatmapWidth)
	for i := range heatmapData {
		heatmapData[i] = make([]int, heatmapHeight)
	}

	// 3. Projdeme přes všechny body a když potkáme nějaký, který chceme,
	// přidáme ho do heatmapy
	phaseStart = time.Now()
	for xml := range parser.Stream() {
		switch xml.Name {
		case "node":
			ID, _ := strconv.Atoi(xml.Attrs["id"])
			if _, found := buildingsToDisplay[ID]; found {
				lat, _ := strconv.ParseFloat(xml.Attrs["lat"], 64)
				lon, _ := strconv.ParseFloat(xml.Attrs["lon"], 64)
				x := int(math.Floor((lon - lonMin) / cellSize))
				y := int(math.Floor((lat - latMin) / cellSize))
				heatmapData[x][y]++
				if heatmapData[x][y] > heatmapMax {
					heatmapMax = heatmapData[x][y]
				}
				placedBuildings++
			}
		}
	}

	// 4. Vygenerujeme heatmapu a uložíme
	// A) Černobílá varianta
	//heatmap := image.NewGray(image.Rectangle{image.Point{0, 0}, image.Point{heatmapWidth, heatmapHeight}})

	// B) Plnobarevná varianta
	heatmap := image.NewRGBA(image.Rectangle{image.Point{0, 0}, image.Point{heatmapWidth, heatmapHeight}})
	backgroundColor := color.RGBA{R: 0, G: 0, B: 0}
	halfHeatColor := color.RGBA{R: 255, G: 40, B: 0}
	heatColor := color.RGBA{R: 80, G: 255, B: 0}

	for x := range heatmapData {
		for y := range heatmapData[x] {
			// Spočítáme si číslo mezi 0 a 1 udávající hustotu, pro lepší znázornění umocníme hodnotu na 1/3
			value := math.Pow(float64(heatmapData[x][y])/float64(heatmapMax), 1.0/3)
			// A) Černobílá varianta
			//heatmap.SetGray(x, heatmapHeight-y-1, color.Gray{uint8(256 * value)})

			// B) Plnobarevná varianta
			fromColor, toColor := backgroundColor, halfHeatColor
			if value > 0.5 {
				fromColor, toColor = halfHeatColor, heatColor
			}
			computeValue := func(from, to uint8, value float64) uint8 {
				return from + uint8(math.Round((float64(to)-float64(from))*value))
			}
			color := color.RGBA{
				R: computeValue(fromColor.R, toColor.R, value),
				G: computeValue(fromColor.G, toColor.G, value),
				B: computeValue(fromColor.B, toColor.B, value),
				A: 255,
			}
			heatmap.SetRGBA(x, heatmapHeight-y-1, color)
		}
	}

	if f, err := os.Create(outputImageName); err == nil {
		png.Encode(f, heatmap)
		f.Close()
	}

	// Uložíme si i data, kdyby se nám hodily
	if f, err := os.Create(outputImageName + ".data"); err == nil {
		fmt.Fprintf(f, "%#v\n", heatmapData)
		f.Close()
	}

	printStat()
	fmt.Fprintf(os.Stderr, "\nPlaced %d buildings, max heatmap value: %d\n", placedBuildings, heatmapMax)
}