package entropy

import (
	"fmt"

	"go.jknobloc.com/x/dataset"
	"go.jknobloc.com/x/llm"
)

type logProb struct {
	document int
	token    int
	value    float32
	offset   int
}

func Context(model llm.Causal, tokenizer llm.Tokenizer, data dataset.Reader) error {
	evaluatorConfig := llm.EvaluatorConfig{
		BatchSize:  32,
		NumWorkers: 16,
	}

	tokenBufferConfig := llm.TokenBufferConfig{
		Window:     1024,
		Stride:     512,
		PadLeft:    false,
		PadRight:   false,
		PadTokenID: 256,
	}

	eval := llm.NewEvaluator(model, tokenizer, func(job llm.Job, logProbs []float32, tokens []int) []logProb {
		r := make([]logProb, len(tokens))

		for i, token := range tokens {
			r[i] = logProb{
				document: job.Document,
				token:    token,
				value:    logProbs[i],
				offset:   job.Position*tokenBufferConfig.Stride + job.Seen + i,
			}
		}

		return r
	}, evaluatorConfig)

	return eval.RunAndCollect("Context", data, tokenBufferConfig, func(r []logProb) error {
		for _, l := range r {
			fmt.Println(l) // TODO implement
		}

		return nil
	})
}