package knobloch

import (
	"fmt"
	"image/color"
	"log"
	"slices"
	"sync/atomic"
	"time"

	"github.com/jonasknobloch/mbpe"
	"go.jknobloc.com/x/tokenizer/bpe"
	"go.jknobloc.com/x/tui"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
	"gonum.org/v1/plot/vg/draw"
)

// TODO we could color code intermediate / terminal tokens
// TODO some mere layer like color coding would be even cooler
// TODO we could even color code morpheme approximations?
// TODO why not simply count runes? + mark terminal?

// NRGBA; color.RGBA is alpha-premultiplied and mangles the blend
var (
	colorOther    = color.NRGBA{R: 108, G: 126, B: 179, A: 96}
	colorMorpheme = color.NRGBA{R: 214, G: 96, B: 77, A: 96}
)

func opaque(c color.NRGBA) color.NRGBA {
	c.A = 255

	return c
}

type NoPreTok struct{}

func (p *NoPreTok) PreTokenize(s string) []string {
	return []string{s}
}

func TokenFrequencies(t *bpe.Tokenizer, name, suffix string) error {
	dict := mbpe.NewDict()

	if err := dict.Load(name); err != nil {
		return err
	}

	m := bpe.MBPE(t)

	m.SetPreTokenizer(&NoPreTok{})

	r := make([]int, len(bpe.Vocab(t)))

	items := dict.Items()

	pb := tui.NewProgressBar("Chunks", 20, len(items), time.Now())

	var processed atomic.Int64

	pb.Start(1*time.Second, func() int {
		return int(processed.Load())
	})

	for _, v := range items {
		n := v.N()

		for _, id := range t.Encode(v.Src()) {
			r[id] += n
		}

		processed.Add(1)
	}

	if err := plotAll(r, nil, t, suffix); err != nil {
		return err
	}

	if len(SharedVocabs) > 0 {
		shared, err := sharedMask(len(r), t, true)

		if err != nil {
			return err
		}

		if err := plotAll(r, shared, t, suffix+"_shared"); err != nil {
			return err
		}

		unshared, err := sharedMask(len(r), t, false)

		if err != nil {
			return err
		}

		if err := plotAll(r, unshared, t, suffix+"_unshared"); err != nil {
			return err
		}

		if err := plotMorphemeFractionOverview(r, shared, unshared, t, plotName("morpheme_fraction_overview", suffix)); err != nil {
			return err
		}

		if err := plotMorphemeContribution(r, shared, unshared, t, plotName("morpheme_contribution", suffix)); err != nil {
			return err
		}

		if err := plotMorphemeFractionLesci(r, shared, unshared, t, plotName("morpheme_fraction_lesci", suffix)); err != nil {
			return err
		}

		if err := plotMorphemeFractionByIDOverview(r, shared, unshared, t, plotName("morpheme_fraction_by_id_overview", suffix)); err != nil {
			return err
		}
	}

	stats, err := NewFrequencyStats(suffix, r, t)

	if err != nil {
		return err
	}

	log.Printf("%s: %d/%d used tokens are morphemes (%.1f%% of types, %.1f%% of tokens), median %d vs %d, %d of top 1000, rank-id rho %.4f",
		stats.Model, stats.MorphemeTypes, stats.Used, stats.TypeShare()*100, stats.TokenShare()*100,
		stats.MorphemeMedian, stats.OtherMedian, stats.MorphemeTop1000, stats.RankID)

	return AppendFrequencyStats("frequency_stats.csv", stats)
}

// plotAll draws every figure for one tokenizer. A nil keep covers the whole
// vocabulary; otherwise the subset it selects is plotted. Note that the morpheme
// fraction gets the unmasked counts, since it ranks over the full vocabulary
// before hiding anything.
func plotAll(r []int, keep []bool, t *bpe.Tokenizer, suffix string) error {
	masked := r

	if keep != nil {
		masked = maskCounts(r, keep)
	}

	if err := plotFrequenciesByIDSplit(masked, t, plotName("frequencies_by_id_split", suffix)); err != nil {
		return err
	}

	// two takes on the rank axis: mixed ranks the whole vocabulary and masks
	// after, so subsets stay on the full rank space, while split ranks each
	// group alone so the two separate
	if err := plotFrequenciesMixed(r, keep, t, plotName("frequencies_mixed", suffix)); err != nil {
		return err
	}

	if err := plotFrequenciesSplit(masked, t, plotName("frequencies_split", suffix)); err != nil {
		return err
	}

	return plotMorphemeFraction(r, keep, t, plotName("morpheme_fraction", suffix))
}

func plotName(base, suffix string) string {
	if suffix == "" {
		return base + ".png"
	}

	return base + "_" + suffix + ".png"
}

func plotFrequenciesByIDSplit(m []int, t *bpe.Tokenizer, out string) error {
	a, b, err := splitScatter(pointsByID(m), t)

	if err != nil {
		return err
	}

	return renderFrequenciesSplit(a, b, "Token ID", len(m), out)
}

// plotFrequenciesSplit ranks tokens by frequency rather than id. Each group is
// ranked within itself: a shared ranking puts both on one curve, where the
// groups overplot each other and no composition is visible.
func plotFrequenciesSplit(m []int, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	var a, b []int

	for id, v := range m {
		if v <= 0 {
			continue
		}

		if morph(id) {
			a = append(a, v)
		} else {
			b = append(b, v)
		}
	}

	return renderFrequenciesSplit(rankPoints(a), rankPoints(b), "Token frequency rank", len(m), out)
}

// plotFrequenciesMixed ranks every token on one shared scale and colors each
// point by group, so composition is read along a single curve rather than from
// the two separate curves plotFrequenciesSplit draws.
//
// Ranking covers the whole vocabulary and keep only hides points afterwards, so
// a subset keeps the rank positions it holds in the full figure. Ranking within
// the subset instead would stretch it back across the axis; that is what
// plotFrequenciesSplit does, deliberately, as the other half of the pair.
func plotFrequenciesMixed(m []int, keep []bool, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	ids := rankedIDs(m)

	pts := make(plotter.XYs, 0, len(ids))
	flags := make([]bool, 0, len(ids))

	for i, id := range ids {
		if keep != nil && !keep[id] {
			continue
		}

		pts = append(pts, plotter.XY{X: float64(i + 1), Y: float64(m[id])})
		flags = append(flags, morph(id))
	}

	if len(pts) == 0 {
		return fmt.Errorf("no tokens selected")
	}

	return renderFrequenciesMixed(pts, flags, "Token frequency rank", len(m), out)
}

// rankWindow is wide enough to smooth the per-token noise but still leaves a few
// hundred points across a typical vocabulary.
const rankWindow = 250

// plotMorphemeFraction plots the morpheme share of a window over the frequency
// ranking. Encoding composition as position rather than hue keeps it readable,
// which renderFrequenciesMixed cannot manage on a single curve.
//
// Ranking covers the whole vocabulary and keep only hides points afterwards, so
// a subset lands on the same x positions it occupies in the full figure. Ranking
// within the subset instead would stretch it back across the axis and make the
// subsets impossible to compare.
func plotMorphemeFraction(m []int, keep []bool, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	ids := rankedIDs(m)

	if len(ids) == 0 {
		return fmt.Errorf("no used tokens")
	}

	segments, overall, shown := morphemeFractionCurve(ids, keep, morph)

	if shown == 0 {
		return fmt.Errorf("no tokens selected")
	}

	return renderMorphemeFraction([]fractionSeries{{
		segments: segments,
		overall:  overall,
		ref:      true,
		color:    opaque(colorMorpheme),
		label:    "Window share",
	}}, "Token frequency rank", len(m), out)
}

// plotMorphemeFractionOverview overlays the full vocabulary and both sides of the
// intersection. The shared ranking makes the three curves directly comparable,
// so their overall shares go in the legend instead of as reference lines.
func plotMorphemeFractionOverview(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	ids := rankedIDs(m)

	if len(ids) == 0 {
		return fmt.Errorf("no used tokens")
	}

	var series []fractionSeries

	for _, c := range []struct {
		keep  []bool
		color color.NRGBA
		label string
	}{
		{nil, color.NRGBA{R: 90, G: 90, B: 90, A: 255}, "All"},
		{shared, opaque(colorOther), "Shared"},
		{unshared, opaque(colorMorpheme), "Unshared"},
	} {
		segments, overall, shown := morphemeFractionCurve(ids, c.keep, morph)

		if shown == 0 {
			continue
		}

		series = append(series, fractionSeries{
			segments: segments,
			overall:  overall,
			color:    c.color,
			label:    fmt.Sprintf("%s (%.1f%%)", c.label, overall),
		})
	}

	if len(series) == 0 {
		return fmt.Errorf("no tokens selected")
	}

	return renderMorphemeFraction(series, "Token frequency rank", len(m), out)
}

// plotMorphemeContribution decomposes the overall morpheme share into the part
// each side of the intersection supplies. Every curve is divided by all tokens
// in the bin rather than by its own subset, so the two stack to exactly the
// combined curve, which the rate view in plotMorphemeFractionOverview does not.
//
// The denominator also being the full bin removes the small-sample noise that
// makes the rate view ragged wherever a subset is thin, and leaves no empty bins
// to break the line. The cost is that a low value here mixes "morpheme-poor"
// with "barely present", which the rate view keeps apart.
func plotMorphemeContribution(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	ids := rankedIDs(m)

	if len(ids) == 0 {
		return fmt.Errorf("no used tokens")
	}

	sharedPts, sharedTotal := contributionCurve(ids, shared, morph)
	_, unsharedTotal := contributionCurve(ids, unshared, morph)
	combined, combinedTotal := contributionCurve(ids, nil, morph)

	// the combined curve is filled to the baseline first, then the shared part is
	// filled over it, so the salmon left showing is the unshared band. The same
	// combined values are drawn again unfilled on top, since the total is a curve
	// worth reading in its own right and not just the top edge of a stack.
	return renderMorphemeContribution([]fractionSeries{
		{
			segments: []plotter.XYs{combined},
			color:    opaque(colorMorpheme),
			fill:     true,
			label:    fmt.Sprintf("Unshared (%.1f%%)", unsharedTotal),
		},
		{
			segments: []plotter.XYs{sharedPts},
			color:    opaque(colorOther),
			fill:     true,
			label:    fmt.Sprintf("Shared (%.1f%%)", sharedTotal),
		},
		{
			segments: []plotter.XYs{combined},
			color:    color.NRGBA{R: 90, G: 90, B: 90, A: 255},
			label:    fmt.Sprintf("Combined (%.1f%%)", combinedTotal),
		},
	}, combinedTotal, len(m), out)
}

// contributionCurve reports the morphemes selected by keep as a percentage of
// every token in the bin, so curves for disjoint keeps add up.
func contributionCurve(ids []int, keep []bool, morph func(int) bool) (plotter.XYs, float64) {
	pts := make(plotter.XYs, 0, len(ids)/rankWindow+1)

	morphemes := 0

	for start := 0; start < len(ids); start += rankWindow {
		end := min(start+rankWindow, len(ids))

		count := 0

		for _, id := range ids[start:end] {
			if keep != nil && !keep[id] {
				continue
			}

			if morph(id) {
				count++
			}
		}

		morphemes += count

		pts = append(pts, plotter.XY{
			X: float64(start + (end-start)/2),
			Y: 100 * float64(count) / float64(end-start),
		})
	}

	if len(ids) == 0 {
		return pts, 0
	}

	return pts, 100 * float64(morphemes) / float64(len(ids))
}

func renderMorphemeContribution(series []fractionSeries, overall float64, vocab int, out string) error {
	p := plot.New()

	p.X.Label.Text = "Token frequency rank"
	p.Y.Label.Text = "Share of all tokens in window (%)"

	// pin to the vocabulary so figures stay comparable between tokenizers
	p.X.Min = 0
	p.X.Max = float64(vocab)

	p.Y.Min = 0
	p.Y.Max = 100

	p.Add(plotter.NewGrid())

	ref := plotter.NewFunction(func(float64) float64 {
		return overall
	})

	ref.Color = color.NRGBA{R: 131, G: 131, B: 131, A: 255}
	ref.Width = vg.Points(1)
	ref.Dashes = []vg.Length{vg.Points(4), vg.Points(3)}

	p.Add(ref)
	p.Legend.Add("Combined mean", ref)

	for _, s := range series {
		for _, pts := range s.segments {
			line, err := plotter.NewLine(pts)

			if err != nil {
				return err
			}

			line.Color = s.color
			line.Width = vg.Points(1)

			if s.fill {
				// opaque: the shared fill has to occlude the combined fill
				// underneath it, otherwise the combined area shows through and
				// the visible red is the total rather than the unshared band
				line.FillColor = opaque(s.color)
			} else {
				line.Width = vg.Points(1.5)
			}

			p.Add(line)
			p.Legend.Add(s.label, line)
		}
	}

	p.Legend.Top = true
	p.Legend.Padding = vg.Points(4)

	return p.Save(12*vg.Inch, 6*vg.Inch, out)
}

// idFractionCurve bins by token id instead of by rank, over the half-open id
// range [lo, hi). Empty bins end the current segment rather than being
// interpolated. A nil keep covers every token.
func idFractionCurve(m []int, keep []bool, morph func(int) bool, lo, hi int) (segments []plotter.XYs, overall float64, shown int) {
	var current plotter.XYs

	morphemes := 0

	for start := lo; start < hi; start += rankWindow {
		end := min(start+rankWindow, hi)

		count, total := 0, 0

		for id := start; id < end; id++ {
			if id < 0 || id >= len(m) || m[id] <= 0 || (keep != nil && !keep[id]) {
				continue
			}

			total++

			if morph(id) {
				count++
			}
		}

		if total == 0 {
			if len(current) > 0 {
				segments = append(segments, current)

				current = nil
			}

			continue
		}

		shown += total
		morphemes += count

		current = append(current, plotter.XY{
			X: float64(start + (end-start)/2),
			Y: 100 * float64(count) / float64(total),
		})
	}

	if len(current) > 0 {
		segments = append(segments, current)
	}

	if shown > 0 {
		overall = 100 * float64(morphemes) / float64(shown)
	}

	return segments, overall, shown
}

// plotMorphemeFractionByIDOverview is the token id counterpart of
// plotMorphemeFractionOverview: same three series, binned by vocab position over
// the whole vocabulary rather than by frequency rank, and with no lesci filter.
func plotMorphemeFractionByIDOverview(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
	morph, err := isMorpheme(t)

	if err != nil {
		return err
	}

	var series []fractionSeries

	for _, c := range []struct {
		keep  []bool
		color color.NRGBA
		label string
	}{
		{nil, color.NRGBA{R: 90, G: 90, B: 90, A: 255}, "All"},
		{shared, opaque(colorOther), "Shared"},
		{unshared, opaque(colorMorpheme), "Unshared"},
	} {
		segments, overall, shown := idFractionCurve(m, c.keep, morph, 0, len(m))

		if shown == 0 {
			continue
		}

		series = append(series, fractionSeries{
			segments: segments,
			overall:  overall,
			color:    c.color,
			label:    fmt.Sprintf("%s (%.1f%%)", c.label, overall),
		})
	}

	if len(series) == 0 {
		return fmt.Errorf("no tokens selected")
	}

	return renderMorphemeFraction(series, "Token ID", len(m), out)
}

// rankedIDs lists the used vocab ids most frequent first.
func rankedIDs(m []int) []int {
	ids := make([]int, 0, len(m))

	for id, v := range m {
		if v > 0 {
			ids = append(ids, id)
		}
	}

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

	return ids
}

// morphemeFractionCurve bins the ranking and reports the morpheme share of the
// selected tokens per bin. Bins are fixed on the ranking it is given, so a bin
// covers the same ranks for every keep; empty bins end the current segment
// rather than letting the line interpolate across them.
func morphemeFractionCurve(ids []int, keep []bool, morph func(int) bool) (segments []plotter.XYs, overall float64, shown int) {
	var current plotter.XYs

	morphemes := 0

	for start := 0; start < len(ids); start += rankWindow {
		end := min(start+rankWindow, len(ids))

		count, total := 0, 0

		for _, id := range ids[start:end] {
			if keep != nil && !keep[id] {
				continue
			}

			total++

			if morph(id) {
				count++
			}
		}

		if total == 0 {
			if len(current) > 0 {
				segments = append(segments, current)

				current = nil
			}

			continue
		}

		shown += total
		morphemes += count

		current = append(current, plotter.XY{
			X: float64(start + (end-start)/2),
			Y: 100 * float64(count) / float64(total),
		})
	}

	if len(current) > 0 {
		segments = append(segments, current)
	}

	if shown > 0 {
		overall = 100 * float64(morphemes) / float64(shown)
	}

	return segments, overall, shown
}

type fractionSeries struct {
	segments []plotter.XYs
	overall  float64
	ref      bool
	fill     bool
	color    color.NRGBA
	label    string
}

func renderMorphemeFraction(series []fractionSeries, labelX string, vocab int, out string) error {
	p := plot.New()

	p.X.Label.Text = labelX
	p.Y.Label.Text = "Morphemes in window (%)"

	// pin to the vocabulary so figures stay comparable between tokenizers
	p.X.Min = 0
	p.X.Max = float64(vocab)

	p.Y.Min = 0
	p.Y.Max = 100

	p.Add(plotter.NewGrid())

	for _, s := range series {
		if s.ref {
			ref := plotter.NewFunction(func(float64) float64 {
				return s.overall
			})

			ref.Color = color.NRGBA{R: 131, G: 131, B: 131, A: 255}
			ref.Width = vg.Points(1)
			ref.Dashes = []vg.Length{vg.Points(4), vg.Points(3)}

			p.Add(ref)
			p.Legend.Add("Overall share", ref)
		}

		// one line per segment; only the first carries the legend entry
		for i, pts := range s.segments {
			line, err := plotter.NewLine(pts)

			if err != nil {
				return err
			}

			line.Color = s.color
			line.Width = vg.Points(1.5)

			p.Add(line)

			if i == 0 {
				p.Legend.Add(s.label, line)
			}
		}
	}

	p.Legend.Top = true
	p.Legend.Padding = vg.Points(4)

	return p.Save(12*vg.Inch, 6*vg.Inch, out)
}

func rankPoints(values []int) plotter.XYs {
	slices.SortFunc(values, func(a, b int) int {
		return b - a
	})

	pts := make(plotter.XYs, 0, len(values))

	for i, v := range values {
		pts = append(pts, plotter.XY{X: float64(i + 1), Y: float64(v)})
	}

	return pts
}

func pointsByID(m []int) plotter.XYs {
	pts := make(plotter.XYs, 0, len(m))

	for id, v := range m {
		if v <= 0 {
			continue
		}

		pts = append(pts, plotter.XY{X: float64(id), Y: float64(v)})
	}

	return pts
}

func splitScatter(pts plotter.XYs, t *bpe.Tokenizer) (a, b plotter.XYs, err error) {
	morph, err := isMorpheme(t)

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

	for _, pt := range pts {
		if morph(int(pt.X)) {
			a = append(a, pt)
		} else {
			b = append(b, pt)
		}
	}

	return a, b, nil
}

func renderFrequenciesSplit(a, b plotter.XYs, labelX string, vocab int, out string) error {
	p := plot.New()

	p.X.Label.Text = labelX
	p.Y.Label.Text = "Token frequency"

	p.Y.Scale = plot.LogScale{}

	p.Y.Tick.Marker = plot.LogTicks{Prec: -1}

	// pin to the vocabulary so figures stay comparable between tokenizers
	p.X.Min = 0
	p.X.Max = float64(vocab)

	p.Add(plotter.NewGrid())

	addScatter := func(pts plotter.XYs, c color.NRGBA, label string) error {
		s, err := plotter.NewScatter(pts)

		if err != nil {
			return err
		}

		s.Color = c
		s.Radius = vg.Points(1)
		s.Shape = draw.CircleGlyph{}

		p.Add(s)

		return addLegendKey(p, c, label)
	}

	if err := addScatter(b, colorOther, "Other"); err != nil {
		return err
	}

	if err := addScatter(a, colorMorpheme, "Morpheme"); err != nil {
		return err
	}

	p.Legend.Padding = vg.Points(4)

	return p.Save(12*vg.Inch, 7*vg.Inch, out)
}

// renderFrequenciesMixed colors a single scatter per point. Overlaying two
// plotters instead would draw one group entirely on top of the other, which
// hides the composition wherever both share a rank range.
func renderFrequenciesMixed(pts plotter.XYs, morph []bool, labelX string, vocab int, out string) error {
	p := plot.New()

	p.X.Label.Text = labelX
	p.Y.Label.Text = "Token frequency"

	p.Y.Scale = plot.LogScale{}

	p.Y.Tick.Marker = plot.LogTicks{Prec: -1}

	// pin to the vocabulary so figures stay comparable between tokenizers
	p.X.Min = 0
	p.X.Max = float64(vocab)

	p.Add(plotter.NewGrid())

	s, err := plotter.NewScatter(pts)

	if err != nil {
		return err
	}

	s.GlyphStyleFunc = func(i int) draw.GlyphStyle {
		c := colorOther

		if morph[i] {
			c = colorMorpheme
		}

		return draw.GlyphStyle{Color: c, Radius: vg.Points(1), Shape: draw.CircleGlyph{}}
	}

	p.Add(s)

	if err := addLegendKey(p, colorOther, "Other"); err != nil {
		return err
	}

	if err := addLegendKey(p, colorMorpheme, "Morpheme"); err != nil {
		return err
	}

	p.Legend.Padding = vg.Points(4)

	return p.Save(12*vg.Inch, 7*vg.Inch, out)
}

// addLegendKey adds a solid, larger glyph; the plotted points are too small and
// too faint to read as a legend key.
func addLegendKey(p *plot.Plot, c color.NRGBA, label string) error {
	thumb, err := plotter.NewScatter(plotter.XYs{})

	if err != nil {
		return err
	}

	thumb.Color = opaque(c)
	thumb.Radius = vg.Points(3)
	thumb.Shape = draw.CircleGlyph{}

	p.Legend.Add(label, thumb)

	return nil
}