// Command morphbands reports morpheme content per band of vocabulary-id space,
// counted two ways: by type and weighted by corpus frequency.
//
// The two disagree sharply. Morpheme types concentrate early -- around two
// thirds of the first 2,048 sub-words -- and thin out into the tail, but the
// tail is where most of the vocabulary is, so a type count says little about how
// much text is actually covered by morphemic sub-words. Weighting each sub-word
// by how often the tokenizer emits it over the corpus answers that instead.
//
// Frequencies come from encoding the dictionary, the same way the frequency
// plots do, so this takes a few minutes per tokenizer and is the slow part.
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/jonasknobloch/mbpe"
	"go.jknobloc.com/x/research/knobloch"
	"go.jknobloc.com/x/shelf"
	"go.jknobloc.com/x/tokenizer/bpe"
)

// frequencies counts how often each vocabulary id is emitted over the
// dictionary, each word weighted by its corpus count.
func frequencies(t *bpe.Tokenizer, dictPath string) ([]int, error) {
	dict := mbpe.NewDict()

	if err := dict.Load(dictPath); err != nil {
		return nil, err
	}

	// the dictionary is already pre-tokenized, so re-splitting would corrupt it
	m := bpe.MBPE(t)
	m.SetPreTokenizer(&knobloch.NoPreTok{})

	r := make([]int, len(bpe.Vocab(t)))

	for _, v := range dict.Items() {
		n := v.N()

		for _, id := range t.Encode(v.Src()) {
			r[id] += n
		}
	}

	return r, nil
}

func parseBands(s string, size int) []int {
	var edges []int

	for _, f := range strings.Split(s, ",") {
		v, err := strconv.Atoi(strings.TrimSpace(f))

		if err != nil {
			log.Fatalf("bad band edge %q: %v", f, err)
		}

		edges = append(edges, v)
	}

	if edges[0] != 0 {
		edges = append([]int{0}, edges...)
	}

	if last := edges[len(edges)-1]; last < size {
		edges = append(edges, size)
	}

	return edges
}

func share(part, whole int) string {
	if whole == 0 {
		return ""
	}

	return fmt.Sprintf("%.2f", 100*float64(part)/float64(whole))
}

func main() {
	dirs := flag.String("dirs", "results/knobloch/minipile_19_ctrl/%s_minipile",
		"tokenizer directory, %s is the alignment name")
	dict := flag.String("dict", "results/knobloch/minipile_19_ctrl/dict.txt",
		"pre-tokenized corpus dictionary")
	list := flag.String("alignments", "m000,m050,m100,mi050,mi100",
		"comma-separated alignments to report")
	bands := flag.String("bands", "2048,8192,32768,100512,131072,262144",
		"comma-separated upper edges of the id bands")
	out := flag.String("out", "morpheme_bands.csv", "output CSV")

	flag.Parse()

	morph, err := knobloch.Morphemes(shelf.Abs(knobloch.SegmentsPath))

	if err != nil {
		log.Fatal(err)
	}

	log.Printf("%d known morphemes", len(morph))

	f, err := os.Create(*out)

	if err != nil {
		log.Fatal(err)
	}

	defer f.Close()

	w := csv.NewWriter(f)
	defer w.Flush()

	if err := w.Write([]string{"tokenizer", "band_lo", "band_hi", "size",
		"morpheme_types", "type_share", "occurrences", "morpheme_occurrences",
		"token_share"}); err != nil {
		log.Fatal(err)
	}

	for _, n := range strings.Split(*list, ",") {
		n = strings.TrimSpace(n)
		dir := shelf.Item(fmt.Sprintf(*dirs, n))

		tok, err := bpe.NewTokenizerFromFiles(
			shelf.Abs(dir+"/vocab.json"), shelf.Abs(dir+"/merges.txt"),
			bpe.Config{Recover: false})

		if err != nil {
			log.Fatal(err)
		}

		vocab := bpe.Vocab(tok)

		log.Printf("%s: encoding dictionary over %d sub-words", n, len(vocab))

		r, err := frequencies(tok, shelf.Abs(shelf.Item(*dict)))

		if err != nil {
			log.Fatal(err)
		}

		edges := parseBands(*bands, len(vocab))

		log.Printf("%-6s %10s %9s %9s %11s %9s", n, "band", "types%", "tokens%", "occurrences", "morphs")

		for i := 0; i+1 < len(edges); i++ {
			lo, hi := edges[i], edges[i+1]

			var types, morphTypes, occ, morphOcc int

			for id := lo; id < hi && id < len(vocab); id++ {
				types++
				occ += r[id]

				if _, ok := morph[vocab[id]]; ok {
					morphTypes++
					morphOcc += r[id]
				}
			}

			row := []string{n, strconv.Itoa(lo), strconv.Itoa(hi), strconv.Itoa(types),
				strconv.Itoa(morphTypes), share(morphTypes, types),
				strconv.Itoa(occ), strconv.Itoa(morphOcc), share(morphOcc, occ)}

			if err := w.Write(row); err != nil {
				log.Fatal(err)
			}

			log.Printf("%-6s %10s %8s%% %8s%% %11d %9d", "",
				fmt.Sprintf("%d-%d", lo, hi), share(morphTypes, types),
				share(morphOcc, occ), occ, morphTypes)
		}

		w.Flush()
	}

	log.Printf("wrote %s", *out)
}