package main

// TODO: see TODO in hocr package
//
// TODO: Simplify things into functions more; this works well, but is a bit of a rush job

import (
	"flag"
	"fmt"
	"image/png"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"git.rescribe.xyz/testingtools/hocr"
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: line-conf-buckets hocr1 [hocr2] [...]\n")
		fmt.Fprintf(os.Stderr, "Copies image-text line pairs into different directories according\n")
		fmt.Fprintf(os.Stderr, "to the average character probability for the line.\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	lines := make(hocr.LineDetails, 0)

	for _, f := range flag.Args() {
		file, err := ioutil.ReadFile(f)
		if err != nil {
			log.Fatal(err)
		}

		h, err := hocr.Parse(file)
		if err != nil {
			log.Fatal(err)
		}

		pngfn := strings.Replace(f, ".hocr", ".png", 1)
		pngf, err := os.Open(pngfn)
		if err != nil {
			log.Fatal(err)
		}
		defer pngf.Close()
		img, err := png.Decode(pngf)
		if err != nil {
			log.Fatal(err)
		}

		n := strings.Replace(filepath.Base(f), ".hocr", "", 1)
		newlines, err := hocr.GetLineDetails(h, img, n)
		if err != nil {
			log.Fatal(err)
		}
		for _, l := range newlines {
			lines = append(lines, l)
		}
	}

	sort.Sort(lines)

	worstnum := 0
	mediumnum := 0
	bestnum := 0

	outdir := "buckets" // TODO: set this from cmdline
	todir := ""

	for _, l := range lines {
		switch {
		case l.Avgconf < 0.95:
			todir = "bad"
			worstnum++
		case l.Avgconf < 0.98:
			todir = "95to98"
			mediumnum++
		default:
			todir = "98plus"
			bestnum++
		}

		avgstr := strconv.FormatFloat(l.Avgconf, 'f', 5, 64)
		avgstr = avgstr[2:]
		outname := filepath.Join(outdir, todir, l.Hocrname + "_" + l.Name + "_" + avgstr + ".png")

		err := os.MkdirAll(filepath.Join(outdir, todir), 0700)
		if err != nil {
			log.Fatal(err)
		}

		outfile, err := os.Create(outname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create %s\n", outname)
			log.Fatal(err)
		}
		defer outfile.Close()

		err = l.Img.CopyLineTo(outfile)
		if err != nil {
			log.Fatal(err)
		}

		outname = filepath.Join(outdir, todir, l.Hocrname + "_" + l.Name + "_" + avgstr + ".txt")
		outfile, err = os.Create(outname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create %s\n", outname)
			log.Fatal(err)
		}
		defer outfile.Close()

		_, err = io.WriteString(outfile, l.Text)
		if err != nil {
			log.Fatal(err)
		}

		// TODO: test whether the line.img works properly with multiple hocrs, as it could be that as it's a pointer, it always points to the latest image (don't think so, but not sure)
	}

	total := worstnum + mediumnum + bestnum

	fmt.Printf("Copied lines to %s\n", outdir)
	fmt.Printf("---------------------------------\n")
	fmt.Printf("Lines 98%%+ quality:     %d%%\n", 100 * bestnum / total)
	fmt.Printf("Lines 95-98%% quality:   %d%%\n", 100 * mediumnum / total)
	fmt.Printf("Lines <95%% quality:     %d%%\n", 100 * worstnum / total)
}