package holdem

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"math/rand/v2"
	"sync"

	"pete/internal/games/cards"
)

// The trainer.
//
// This is counterfactual regret minimisation, and what it produces is policy.gob
// — the table the bots read at the table. It is not on any request path; it runs
// from cmd/holdem-train, for half an hour, and then it is a file.
//
// The one thing worth understanding about it: **it plays the real game.** Every
// move it explores goes through Step, which is the same reducer the felt calls,
// so the blinds, the min-raise, the street completion and the money are the ones
// a player will actually meet. Its info-set key comes out of State.spot, which is
// the same function the bots look themselves up with.
//
// That is not tidiness, it is the whole lesson of the policy this replaces. That
// one was trained against a hand-written model of poker sitting beside the real
// engine — a model where a call always ended the street, the big blind had no
// option, and the payoff was half the pot no matter who had put what in. Then it
// was looked up under a key the trainer never wrote. The result was a 3.4MB file
// that had never once been read, and nobody could tell, because a policy miss is
// not an error. It just quietly isn't there.
//
// So: one engine, one key function, and a test that fails if the bots stop
// finding themselves in the table.

// How much of the game tree to explore. Two raises a street keeps the tree small
// enough to converge; a third barely changes how anybody plays and multiplies the
// nodes.
const (
	maxRaisesPerStreet = 2
	maxDepth           = 40
	trainMCIters       = 60 // noisy, but it is only picking a bucket
)

// regrets is what CFR accumulates: how much better each action would have been.
type regrets map[string]*[numActions]float64

// Trained is the file the bots read.
type Trained struct {
	Strategy map[string][numActions]float64
	Meta     TrainMeta
}

// TrainMeta is what the policy can say about itself. Worth having: a policy is
// otherwise an opaque three megabytes and there is no way to tell a good one from
// a stale one by looking.
type TrainMeta struct {
	Iterations int
	Stakes     string
	Depths     string
	Nodes      int
}

// ---- preflop, measured once ------------------------------------------------

var (
	preflopOnce  sync.Once
	preflopTable [13][13]Equity // [hi][lo] offsuit, [lo][hi] suited, diagonal pairs
)

// preflopEquity is the equity of a starting hand heads-up. There are only 169
// hands that differ from each other, so they are measured properly, once, and
// then it is a lookup — which matters twice: it takes the noise out of a bucket
// boundary, and the trainer visits preflop on every single iteration.
func preflopEquity(hole [2]cards.Card) Equity {
	preflopOnce.Do(func() {
		rng := cards.NewRNG(20260714, 1)
		for a := cards.Ace; a <= cards.King; a++ {
			for b := a; b <= cards.King; b++ {
				lo, hi := rankIdx(a), rankIdx(b)

				// Suited, and the pairs (which can only be offsuit) on the diagonal.
				s1 := cards.Card{Rank: a, Suit: cards.Spades}
				s2 := cards.Card{Rank: b, Suit: cards.Spades}
				if a == b {
					s2.Suit = cards.Hearts
				}
				preflopTable[lo][hi] = equityOf([2]cards.Card{s1, s2}, nil, 1, 10000, rng)

				if a != b {
					o2 := cards.Card{Rank: b, Suit: cards.Hearts}
					preflopTable[hi][lo] = equityOf([2]cards.Card{s1, o2}, nil, 1, 10000, rng)
				}
			}
		}
	})

	lo, hi := rankIdx(hole[0].Rank), rankIdx(hole[1].Rank)
	if lo > hi {
		lo, hi = hi, lo
	}
	if hole[0].Suit == hole[1].Suit {
		return preflopTable[lo][hi] // suited, and the pairs sit here too
	}
	if lo == hi {
		return preflopTable[lo][hi]
	}
	return preflopTable[hi][lo] // offsuit
}

// rankIdx maps a rank to 0–12, with the ace high — which is what it is, before
// the flop.
func rankIdx(r cards.Rank) int {
	if r == cards.Ace {
		return 12
	}
	return int(r) - 2
}

// ---- the traversal ---------------------------------------------------------

// Train runs external-sampling MCCFR for n hands and returns the average
// strategy. Each worker keeps its own tables and they are summed at the end,
// which is what makes this embarrassingly parallel and is the only reason it
// finishes in half an hour.
func Train(n, workers int, t Tier, minBB, maxBB int64, seed uint64, progress func(done int)) *Trained {
	if workers < 1 {
		workers = 1
	}

	type table struct {
		reg regrets
		avg regrets
	}
	out := make([]table, workers)

	var wg sync.WaitGroup
	var done sync.Mutex
	completed := 0

	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func(w int) {
			defer wg.Done()
			tr := &trainer{
				reg:   regrets{},
				avg:   regrets{},
				tier:  t,
				minBB: minBB,
				maxBB: maxBB,
				rng:   cards.NewRNG(seed, uint64(w)+1),
			}

			share := n / workers
			if w < n%workers {
				share++
			}
			for i := 0; i < share; i++ {
				tr.iterate(uint64(w)<<40 | uint64(i))
				if progress != nil && i%2000 == 0 {
					done.Lock()
					completed += 2000
					c := completed
					done.Unlock()
					progress(c)
				}
			}
			out[w] = table{tr.reg, tr.avg}
		}(w)
	}
	wg.Wait()

	// Sum the workers' average-strategy tables, then normalise each node into the
	// probabilities a bot will actually play.
	total := regrets{}
	for _, tab := range out {
		for key, v := range tab.avg {
			acc, ok := total[key]
			if !ok {
				acc = &[numActions]float64{}
				total[key] = acc
			}
			for i, x := range v {
				acc[i] += x
			}
		}
	}

	strategy := make(map[string][numActions]float64, len(total))
	for key, v := range total {
		var sum float64
		for _, x := range v {
			sum += x
		}
		var probs [numActions]float64
		if sum > 0 {
			for i, x := range v {
				probs[i] = x / sum
			}
		} else {
			for i := range probs {
				probs[i] = 1.0 / numActions
			}
		}
		strategy[key] = probs
	}

	return &Trained{
		Strategy: strategy,
		Meta: TrainMeta{
			Iterations: n,
			Stakes:     fmt.Sprintf("%d/%d", t.SB, t.BB),
			Depths:     fmt.Sprintf("%d–%d BB", minBB, maxBB),
			Nodes:      len(strategy),
		},
	}
}

type trainer struct {
	reg   regrets
	avg   regrets
	tier  Tier
	minBB int64
	maxBB int64
	rng   *rand.Rand

	// A hand's equity on a given street depends on the cards and nothing else —
	// not on how the betting went to get there. The deck is fixed for the whole
	// iteration, so the flop is the same flop down every branch, and this is
	// measured once per seat per street instead of once per node.
	eq   [2][4]Equity
	have [2][4]bool
}

// iterate deals one hand and walks it once for each player.
//
// The stack depth is drawn fresh every hand, across the whole range the table
// allows. This is the fix for the policy that came before: it was trained at ten
// big blinds and nothing else, so four out of five spots in a real cash game fell
// outside anything it had ever seen. A hand of poker is a different game at 20
// big blinds than at 100 — that is most of what makes it a game — and the bots
// have to have played both.
func (tr *trainer) iterate(id uint64) {
	depth := tr.minBB
	if tr.maxBB > tr.minBB {
		depth += tr.rng.Int64N(tr.maxBB - tr.minBB + 1)
	}
	stack := depth * tr.tier.BB

	// No rake while learning. The bots should learn to play poker, not to beat a
	// fee, and the fee is the house's business.
	t := tr.tier
	t.RakePct = 0

	s, err := Open(t, stack, stack, id, tr.rng.Uint64())
	if err != nil {
		return
	}
	start := [2]int64{s.Seats[0].Stack + s.Seats[0].Bet, s.Seats[1].Stack + s.Seats[1].Bet}

	tr.have = [2][4]bool{} // one deal, one set of boards, one set of equities

	for me := 0; me < 2; me++ {
		tr.walk(s.Clone(), me, start, 0)
	}
}

// equity is the cached measurement for this seat on this street.
func (tr *trainer) equity(s State, seat int) Equity {
	st := s.Street
	if st > River {
		st = River
	}
	if !tr.have[seat][st] {
		tr.eq[seat][st] = s.equityFor(seat, trainMCIters, tr.rng)
		tr.have[seat][st] = true
	}
	return tr.eq[seat][st]
}

// walk returns what the hand is worth to `me`, in chips, from here.
func (tr *trainer) walk(s State, me int, start [2]int64, depth int) float64 {
	if s.Phase != PhaseBetting || depth > maxDepth {
		// The hand is over (or we have gone far enough to call it over). What it was
		// worth is simply what the player has now against what they sat down with —
		// the real number, out of the real engine, side pots and all.
		return float64(s.Seats[me].Stack - start[me])
	}

	seat := s.ToAct
	key := s.spotKey(seat, tr.equity(s, seat))

	mask := s.mask(seat)
	if raises(s.History) >= maxRaisesPerStreet {
		mask[actRaiseHalf], mask[actRaisePot] = false, false
	}

	reg := tr.reg[key]
	if reg == nil {
		reg = &[numActions]float64{}
		tr.reg[key] = reg
	}
	strat := match(*reg, mask)

	// The opponent's turn: sample one line and follow it. That is the "external
	// sampling" part, and it is what keeps a hand from costing 5^12 traversals.
	if seat != me {
		avg := tr.avg[key]
		if avg == nil {
			avg = &[numActions]float64{}
			tr.avg[key] = avg
		}
		for i, p := range strat {
			avg[i] += p
		}
		return tr.walk(tr.play(s, seat, sample(strat, tr.rng)), me, start, depth+1)
	}

	// Our turn: try everything, and regret what we didn't do.
	var values [numActions]float64
	var node float64
	for a := 0; a < numActions; a++ {
		if !mask[a] {
			continue
		}
		values[a] = tr.walk(tr.play(s, seat, a), me, start, depth+1)
		node += strat[a] * values[a]
	}
	for a := 0; a < numActions; a++ {
		if mask[a] {
			reg[a] += values[a] - node
		}
	}
	return node
}

// play applies one abstract action through the real reducer.
func (tr *trainer) play(s State, seat, action int) State {
	next, _, err := Step(s.Clone(), s.moveFor(action, seat))
	if err != nil {
		// The mask and the rules disagreed, which is a bug in one of them. Fold and
		// carry on rather than poison the whole run.
		next, _, err = Step(s.Clone(), Move{Kind: Fold})
		if err != nil {
			return s
		}
	}
	return next
}

// match is regret matching: play each action in proportion to how much you wish
// you had played it. An action nobody regrets not taking gets played uniformly.
func match(reg [numActions]float64, mask [numActions]bool) [numActions]float64 {
	var strat [numActions]float64
	var sum float64
	for i, r := range reg {
		if mask[i] && r > 0 {
			sum += r
		}
	}
	if sum > 0 {
		for i, r := range reg {
			if mask[i] && r > 0 {
				strat[i] = r / sum
			}
		}
		return strat
	}

	n := 0
	for _, ok := range mask {
		if ok {
			n++
		}
	}
	if n == 0 {
		strat[actCallCheck] = 1
		return strat
	}
	for i, ok := range mask {
		if ok {
			strat[i] = 1 / float64(n)
		}
	}
	return strat
}

func sample(strat [numActions]float64, rng *rand.Rand) int {
	r := rng.Float64()
	var sum float64
	for i, p := range strat {
		sum += p
		if r < sum {
			return i
		}
	}
	return actCallCheck
}

// raises counts the bets and raises on this street, which is what the tree is
// capped on.
func raises(history string) int {
	n := 0
	for _, c := range history {
		if c == 'r' || c == 'R' {
			n++
		}
	}
	return n
}

// ---- the file --------------------------------------------------------------

// Save writes a trained policy.
func Save(w io.Writer, t *Trained) error { return gob.NewEncoder(w).Encode(t) }

// Load reads one. It is only used by the tests — the bots read the embedded copy.
func Load(r io.Reader) (*Trained, error) {
	var t Trained
	if err := gob.NewDecoder(r).Decode(&t); err != nil {
		return nil, err
	}
	return &t, nil
}

// loadTrained decodes the embedded policy in the new format.
func loadTrained(b []byte) (*Trained, error) { return Load(bytes.NewReader(b)) }