package knobloch

import (
	"encoding/gob"
	"io"
	"os"
	"sync"

	"go.jknobloc.com/x/shelf"
	"go.jknobloc.com/x/tokenizer/bpe"
)

type segmentsEntry struct {
	Value []string
	OK    bool
}

// SegmentsPath holds the morphological segmentation that decides which vocab
// tokens count as morphemes; resolved against the shelf root on first use.
var SegmentsPath shelf.Item = "results/knobloch/minipile_19_fix/segments.gob"

var morphemesCache struct {
	once sync.Once
	m    map[string]int
	err  error
}

// morphemes caches the parsed segmentation; the gob is several hundred MB and
// every plot and stats pass needs the same map.
func morphemes() (map[string]int, error) {
	morphemesCache.once.Do(func() {
		morphemesCache.m, morphemesCache.err = Morphemes(shelf.Abs(SegmentsPath))
	})

	return morphemesCache.m, morphemesCache.err
}

// isMorpheme reports whether a vocab id spells out a known morpheme; the
// returned closure holds the id-to-token map so callers build it once.
func isMorpheme(t *bpe.Tokenizer) (func(id int) bool, error) {
	m, err := morphemes()

	if err != nil {
		return nil, err
	}

	itoa := bpe.Itoa(t)

	return func(id int) bool {
		token, ok := itoa[int64(id)]

		if !ok {
			panic("unexpected token ID")
		}

		_, isMorph := m[token]

		return isMorph
	}, nil
}

func Morphemes(name string) (map[string]int, error) {
	var file *os.File

	if f, err := os.Open(name); err != nil {
		return nil, err
	} else {
		file = f

		defer file.Close()
	}

	var segments map[string]segmentsEntry

	if m, err := loadSegments(file); err != nil {
		return nil, err
	} else {
		segments = m
	}

	morphemes := make(map[string]int)

	for _, v := range segments {
		// the segmenter returns the unsegmented compound when it fails
		if !v.OK || len(v.Value) == 1 {
			continue
		}

		for _, s := range v.Value {
			morphemes[s]++
		}
	}

	return morphemes, nil

	// r := make([]int64, 0)

	// atoi := bpe.Atoi(t)
	//
	// for k, v := range morphemes {
	// 	if v < 0 {
	// 		continue
	// 	}
	//
	// 	i, ok := atoi[k]
	//
	// 	if !ok {
	// 		continue
	// 	}
	//
	// 	r = append(r, i)
	// }
	//
	// return r, nil
}

func loadSegments(r io.Reader) (map[string]segmentsEntry, error) {
	m := make(map[string]segmentsEntry)

	dec := gob.NewDecoder(r)

	if err := dec.Decode(&m); err != nil {
		return nil, err
	}

	return m, nil
}