package knobloch

import (
	"encoding/json"
	"fmt"
	"os"
	"sync"

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

// SharedVocabs lists the vocabularies intersected for the *_shared plots, which
// restrict every figure to tokens all of these models have in common. Set it to
// the family being swept so a model is only compared against its own kind; an
// empty list skips those plots.
var SharedVocabs = []shelf.Item{
	"models/mbpe/minipile/gpt2_50256_m000_minipile/vocab.json",
	"models/mbpe/minipile/gpt2_50256_m030_minipile/vocab.json",
	"models/mbpe/minipile/gpt2_50256_m050_minipile/vocab.json",
	"models/mbpe/minipile/gpt2_50256_m100_minipile/vocab.json",
}

var sharedCache struct {
	mu sync.Mutex
	m  map[string]map[string]struct{}
}

// sharedTokens caches the intersection per vocabulary list, so a sweep that
// switches SharedVocabs between families gets each family's own intersection
// while still reading the vocabularies only once per family.
func sharedTokens() (map[string]struct{}, error) {
	key := ""

	for _, v := range SharedVocabs {
		key += string(v) + "\n"
	}

	sharedCache.mu.Lock()

	defer sharedCache.mu.Unlock()

	if tokens, ok := sharedCache.m[key]; ok {
		return tokens, nil
	}

	tokens, err := intersectVocabs(SharedVocabs)

	if err != nil {
		return nil, err
	}

	if sharedCache.m == nil {
		sharedCache.m = make(map[string]map[string]struct{})
	}

	sharedCache.m[key] = tokens

	return tokens, nil
}

func intersectVocabs(vocabs []shelf.Item) (map[string]struct{}, error) {
	if len(vocabs) == 0 {
		return nil, fmt.Errorf("no vocabularies to intersect")
	}

	var keep map[string]struct{}

	for _, v := range vocabs {
		tokens, err := loadVocab(shelf.Abs(v))

		if err != nil {
			return nil, err
		}

		if keep == nil {
			keep = tokens

			continue
		}

		for token := range keep {
			if _, ok := tokens[token]; !ok {
				delete(keep, token)
			}
		}
	}

	return keep, nil
}

func loadVocab(name string) (map[string]struct{}, error) {
	file, err := os.Open(name)

	if err != nil {
		return nil, err
	}

	defer file.Close()

	var m map[string]int64

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

	tokens := make(map[string]struct{}, len(m))

	for token := range m {
		tokens[token] = struct{}{}
	}

	return tokens, nil
}

// sharedMask selects the vocab ids on one side of the intersection: shared marks
// the tokens every vocabulary in SharedVocabs has, and its complement isolates
// what a listed model added on its own.
func sharedMask(n int, t *bpe.Tokenizer, shared bool) ([]bool, error) {
	keep, err := sharedTokens()

	if err != nil {
		return nil, err
	}

	itoa := bpe.Itoa(t)

	mask := make([]bool, n)

	for id := range mask {
		_, ok := keep[itoa[int64(id)]]

		mask[id] = ok == shared
	}

	return mask, nil
}

// maskCounts zeroes the counts outside the mask. Every plot skips zero counts
// already, so the existing set of figures works unchanged while the axes stay
// pinned to the full vocabulary.
func maskCounts(r []int, keep []bool) []int {
	masked := make([]int, len(r))

	for id, v := range r {
		if keep[id] {
			masked[id] = v
		}
	}

	return masked
}