package knobloch

import (
	"encoding/csv"
	"os"
	"slices"
	"strconv"

	"go.jknobloc.com/x/tokenizer/bpe"
	"gonum.org/v1/gonum/stat"
)

// FrequencyStats summarizes how much of a tokenizer's vocabulary, and how much
// of its token mass, lines up with the morpheme inventory.
type FrequencyStats struct {
	Model string

	// types are distinct vocabulary entries, occurrences are corpus counts
	Vocab          int
	Used           int
	MorphemeTypes  int
	MorphemeTokens int64
	OtherTokens    int64

	MorphemeMedian int
	OtherMedian    int

	// morphemes among the 1000 most frequent tokens
	MorphemeTop1000 int

	// RankID is Spearman's correlation between a token's vocabulary position and
	// its corpus frequency rank. A strict frequency objective merges the most
	// frequent pair first, so position and frequency rank move together and this
	// sits near 1. The further a tokenizer departs from that objective, the more
	// the two come apart, so it measures deviation from the frequency criterion
	// independently of what the competing criterion happens to be.
	RankID float64
}

func (s FrequencyStats) TypeShare() float64 {
	if s.Used == 0 {
		return 0
	}

	return float64(s.MorphemeTypes) / float64(s.Used)
}

func (s FrequencyStats) TokenShare() float64 {
	total := s.MorphemeTokens + s.OtherTokens

	if total == 0 {
		return 0
	}

	return float64(s.MorphemeTokens) / float64(total)
}

func NewFrequencyStats(model string, r []int, t *bpe.Tokenizer) (FrequencyStats, error) {
	morph, err := morphemes()

	if err != nil {
		return FrequencyStats{}, err
	}

	itoa := bpe.Itoa(t)

	s := FrequencyStats{
		Model: model,
		Vocab: len(r),
	}

	var mFreq, oFreq []int

	for id, f := range r {
		if f <= 0 {
			continue
		}

		s.Used++

		if _, ok := morph[itoa[int64(id)]]; ok {
			s.MorphemeTypes++
			s.MorphemeTokens += int64(f)

			mFreq = append(mFreq, f)
		} else {
			s.OtherTokens += int64(f)

			oFreq = append(oFreq, f)
		}
	}

	slices.Sort(mFreq)
	slices.Sort(oFreq)

	s.MorphemeMedian = median(mFreq)
	s.OtherMedian = median(oFreq)

	ids := make([]int, 0, s.Used)

	for id, f := range r {
		if f > 0 {
			ids = append(ids, id)
		}
	}

	slices.SortFunc(ids, func(a, b int) int {
		return r[b] - r[a]
	})

	for _, id := range ids[:min(1000, len(ids))] {
		if _, ok := morph[itoa[int64(id)]]; ok {
			s.MorphemeTop1000++
		}
	}

	s.RankID = rankIDCorrelation(r)

	return s, nil
}

// rankIDCorrelation is Spearman's rho between vocabulary position and corpus
// frequency rank over the used tokens. gonum has no Spearman, so the ranks are
// built here and fed to Pearson, which is the same thing by definition.
func rankIDCorrelation(r []int) float64 {
	ids := make([]int, 0, len(r))

	for id, f := range r {
		if f > 0 {
			ids = append(ids, id)
		}
	}

	if len(ids) < 2 {
		return 0
	}

	// ids is ascending and has no duplicates, so a token's id rank is its index
	position := make([]float64, len(ids))

	for i := range ids {
		position[i] = float64(i + 1)
	}

	order := slices.Clone(ids)

	slices.SortFunc(order, func(a, b int) int {
		return r[b] - r[a]
	})

	// tied counts share the average of the ranks they span
	ranks := make(map[int]float64, len(order))

	for i := 0; i < len(order); {
		j := i

		for j+1 < len(order) && r[order[j+1]] == r[order[i]] {
			j++
		}

		avg := float64(i+j+2) / 2

		for k := i; k <= j; k++ {
			ranks[order[k]] = avg
		}

		i = j + 1
	}

	frequency := make([]float64, len(ids))

	for i, id := range ids {
		frequency[i] = ranks[id]
	}

	return stat.Correlation(position, frequency, nil)
}

func median(v []int) int {
	if len(v) == 0 {
		return 0
	}

	return v[len(v)/2]
}

var frequencyStatsHeader = []string{
	"model", "vocab", "used", "morpheme_types", "type_share",
	"morpheme_tokens", "other_tokens", "token_share",
	"morpheme_median", "other_median", "morpheme_top_1000", "rank_id_spearman",
}

func (s FrequencyStats) row() []string {
	return []string{
		s.Model,
		strconv.Itoa(s.Vocab),
		strconv.Itoa(s.Used),
		strconv.Itoa(s.MorphemeTypes),
		strconv.FormatFloat(s.TypeShare(), 'f', 4, 64),
		strconv.FormatInt(s.MorphemeTokens, 10),
		strconv.FormatInt(s.OtherTokens, 10),
		strconv.FormatFloat(s.TokenShare(), 'f', 4, 64),
		strconv.Itoa(s.MorphemeMedian),
		strconv.Itoa(s.OtherMedian),
		strconv.Itoa(s.MorphemeTop1000),
		strconv.FormatFloat(s.RankID, 'f', 4, 64),
	}
}

// AppendFrequencyStats adds one row per run so numbers for different tokenizers
// accumulate in a single table.
func AppendFrequencyStats(name string, s FrequencyStats) error {
	_, err := os.Stat(name)

	if err != nil && !os.IsNotExist(err) {
		return err
	}

	fresh := os.IsNotExist(err)

	file, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)

	if err != nil {
		return err
	}

	defer file.Close()

	w := csv.NewWriter(file)

	if fresh {
		if err := w.Write(frequencyStatsHeader); err != nil {
			return err
		}
	}

	if err := w.Write(s.row()); err != nil {
		return err
	}

	w.Flush()

	return w.Error()
}