<commit_before>package gencorpus

import (
	"context"
	"encoding/csv"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/google/subcommands"
	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/prove"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
	"golang.org/x/sync/errgroup"
)

type Command struct {
	seed int64
	size int

	games int

	epsilon float64
	depth   int
	threads int

	limit    time.Duration
	analysis string

	stats  bool
	output string
}

func (*Command) Name() string     { return "gencorpus" }
func (*Command) Synopsis() string { return "Generate a corpus of 3x3 positions" }
func (*Command) Usage() string {
	return `gencorpus [flags]
`
}

func (c *Command) SetFlags(flags *flag.FlagSet) {
	flags.IntVar(&c.size, "size", 3, "what size to analyze")
	flags.IntVar(&c.games, "games", 100, "games to generate")
	flags.Int64Var(&c.seed, "seed", 0, "Random seed")
	flags.IntVar(&c.threads, "threads", runtime.NumCPU(), "Number of threads")

	flags.StringVar(&c.analysis, "analysis", "minimax", "Analyze to use: minimax,dfpn,none")
	flags.DurationVar(&c.limit, "limit", 5*time.Second, "Minimax time limit when scoring")

	flags.BoolVar(&c.stats, "stats", false, "compute and print stats")
	flags.IntVar(&c.depth, "depth", 2, "minimax depth")
	flags.Float64Var(&c.epsilon, "epsilon", 0.95, "epsilon for epsilon-greedy generation")

	flags.StringVar(&c.output, "output", "positions.txt", "output file")

}

type game struct {
	positions []*tak.Position
	moves     []tak.Move
}

func growslice[T any](sl []T, newlen int) []T {
	if len(sl) >= newlen {
		return sl
	}
	newsl := make([]T, newlen)
	copy(newsl, sl)
	return newsl
}

type entry struct {
	pos   *tak.Position
	move  tak.Move
	value float64
}

func (c *Command) Execute(ctx context.Context, flag *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	var byLength []int
	var posCount []map[uint64]int

	grp, ctx := errgroup.WithContext(ctx)

	games := make(chan *game)
	positions := make(chan *tak.Position)

	grp.Go(func() error {
		c.generateGames(ctx, games)
		return nil
	})
	grp.Go(func() error {
		defer close(positions)
		rng := rand.New(rand.NewSource(c.seed))

		seen := make(map[uint64]struct{})

		for g := range games {
			if c.stats {
				moves := len(g.positions)
				byLength = growslice(byLength, moves)
				byLength[moves-1] += 1
				posCount = growslice(posCount, moves+1)
				for i, p := range g.positions {
					if posCount[i] == nil {
						posCount[i] = make(map[uint64]int)
					}
					posCount[i][p.Hash()] += 1
				}
			}
			// select position
			var idx int
			r := rng.Float64()
			if r < 0.01 {
				idx = int(rng.Int31n(4))
			} else if r < 0.25 {
				idx = 4 + int(rng.Int31n(5))
			} else if r < 0.95 {
				npos := len(g.positions)
				if npos <= 9 {
					continue
				}
				idx = 9 + int(rng.Int31n(int32(npos)-9))
			}
			if idx >= len(g.positions)-1 {
				continue
			}
			pos := g.positions[idx]
			if _, ok := seen[pos.Hash()]; ok {
				continue
			}
			positions <- pos
		}

		if c.stats {
			for i := range byLength {<commit_msg>we lost the dedup, oops<commit_after>package gencorpus

import (
	"context"
	"encoding/csv"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/google/subcommands"
	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/prove"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
	"golang.org/x/sync/errgroup"
)

type Command struct {
	seed int64
	size int

	games int

	epsilon float64
	depth   int
	threads int

	limit    time.Duration
	analysis string

	stats  bool
	output string
}

func (*Command) Name() string     { return "gencorpus" }
func (*Command) Synopsis() string { return "Generate a corpus of 3x3 positions" }
func (*Command) Usage() string {
	return `gencorpus [flags]
`
}

func (c *Command) SetFlags(flags *flag.FlagSet) {
	flags.IntVar(&c.size, "size", 3, "what size to analyze")
	flags.IntVar(&c.games, "games", 100, "games to generate")
	flags.Int64Var(&c.seed, "seed", 0, "Random seed")
	flags.IntVar(&c.threads, "threads", runtime.NumCPU(), "Number of threads")

	flags.StringVar(&c.analysis, "analysis", "minimax", "Analyze to use: minimax,dfpn,none")
	flags.DurationVar(&c.limit, "limit", 5*time.Second, "Minimax time limit when scoring")

	flags.BoolVar(&c.stats, "stats", false, "compute and print stats")
	flags.IntVar(&c.depth, "depth", 2, "minimax depth")
	flags.Float64Var(&c.epsilon, "epsilon", 0.95, "epsilon for epsilon-greedy generation")

	flags.StringVar(&c.output, "output", "positions.txt", "output file")

}

type game struct {
	positions []*tak.Position
	moves     []tak.Move
}

func growslice[T any](sl []T, newlen int) []T {
	if len(sl) >= newlen {
		return sl
	}
	newsl := make([]T, newlen)
	copy(newsl, sl)
	return newsl
}

type entry struct {
	pos   *tak.Position
	move  tak.Move
	value float64
}

func (c *Command) Execute(ctx context.Context, flag *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	var byLength []int
	var posCount []map[uint64]int

	grp, ctx := errgroup.WithContext(ctx)

	games := make(chan *game)
	positions := make(chan *tak.Position)

	grp.Go(func() error {
		c.generateGames(ctx, games)
		return nil
	})
	grp.Go(func() error {
		defer close(positions)
		rng := rand.New(rand.NewSource(c.seed))

		seen := make(map[uint64]struct{})

		for g := range games {
			if c.stats {
				moves := len(g.positions)
				byLength = growslice(byLength, moves)
				byLength[moves-1] += 1
				posCount = growslice(posCount, moves+1)
				for i, p := range g.positions {
					if posCount[i] == nil {
						posCount[i] = make(map[uint64]int)
					}
					posCount[i][p.Hash()] += 1
				}
			}
			// select position
			var idx int
			r := rng.Float64()
			if r < 0.01 {
				idx = int(rng.Int31n(4))
			} else if r < 0.25 {
				idx = 4 + int(rng.Int31n(5))
			} else if r < 0.95 {
				npos := len(g.positions)
				if npos <= 9 {
					continue
				}
				idx = 9 + int(rng.Int31n(int32(npos)-9))
			}
			if idx >= len(g.positions)-1 {
				continue
			}
			pos := g.positions[idx]
			if _, ok := seen[pos.Hash()]; ok {
				continue
			}
			seen[pos.Hash()] = struct{}{}
			positions <- pos
		}

		if c.stats {
			for i := range byLength {
{'fraction_non_alphanumeric': 0.15834795088502632, 'fraction_numerical': 0.018497847233296123, 'mean_word_length': 3.4231311706629053, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2213723', 'n_tokens_mistral': 2736, 'n_tokens_neox': 2398, 'n_words': 578}