package domain

// Symbols is the checkpoint's trained vocabulary, ported rune-for-rune from
// text/symbols.py (`[_pad] + list(_special) + list(_punctuation) +
// list(_letters) + _dummy`, including the literal backslash inside
// _punctuation's raw Python triple-quoted string, and the 84 "wrong tokens"
// placeholder entries the original training pipeline never sorted out - see
// its own comment: "I trained with wrong tokens... these thing is for that").
var Symbols = buildSymbols()

// SymbolToID mirrors `_symbol_to_id = {s: i for i, s in enumerate(symbols)}`:
// later duplicate runes (e.g. '-' appears in both _special and _punctuation)
// overwrite earlier ones, matching Python dict-construction semantics exactly.
var SymbolToID = buildSymbolToID()

func buildSymbols() []rune {
	var symbols []rune
	symbols = append(symbols, '_')                                                                         // _pad
	symbols = append(symbols, []rune("-~%#@&*$")...)                                                       // _special
	symbols = append(symbols, []rune("!\"'(),.:;?{}<>\\^[]/+- ")...)                                       // _punctuation
	symbols = append(symbols, []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890")...) // _letters
	for range 84 {
		symbols = append(symbols, '=') // _dummy
	}
	return symbols
}

func buildSymbolToID() map[rune]int {
	m := make(map[rune]int, len(Symbols))
	for i, s := range Symbols {
		m[s] = i
	}
	return m
}

// Intersperse ports commons.py::intersperse: pads seq with item between every
// element and at both ends (used when hps.data.add_blank is true, which it
// is for this checkpoint's configs/uma.json).
func Intersperse(seq []int64, item int64) []int64 {
	result := make([]int64, len(seq)*2+1)
	for i := range result {
		result[i] = item
	}
	for i, v := range seq {
		result[i*2+1] = v
	}
	return result
}