// Compare vocabulary in equal-sized random samples of token occurrences.
package main

import (
	"encoding/csv"
	"fmt"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

const sampleSize, trials = 1000, 1000

type result struct {
	author   string
	total    int
	mean, sd float64
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run() error {
	input, err := os.Open("outputs/author-unique-words.tsv")
	if err != nil {
		return err
	}
	defer input.Close()
	reader := csv.NewReader(input)
	reader.Comma = '\t'
	authors, err := reader.ReadAll()
	if err != nil {
		return err
	}
	slugPattern := regexp.MustCompile(`[^a-zA-Z0-9]+`)
	var results []result
	var excluded [][]string
	for _, author := range authors[1:] {
		slug := strings.ToLower(strings.Trim(slugPattern.ReplaceAllString(author[0], "-"), "-"))
		data, err := os.ReadFile(filepath.Join("outputs", slug+".txt"))
		if err != nil {
			return err
		}
		var cumulative []int
		total := 0
		for _, row := range strings.Split(strings.TrimSpace(string(data)), "\n") {
			fields := strings.Split(strings.TrimSuffix(row, "\r"), "\t")
			if len(fields) != 2 {
				return fmt.Errorf("invalid row for %s", author[0])
			}
			count, err := strconv.Atoi(fields[1])
			if err != nil || count < 1 {
				return fmt.Errorf("invalid frequency for %s", author[0])
			}
			total += count
			cumulative = append(cumulative, total)
		}
		if total < sampleSize {
			excluded = append(excluded, []string{author[0], strconv.Itoa(total), "Fewer than 1000 tokens"})
			continue
		}
		// Reset the seed for reproducibility independent of author ordering.
		rng := rand.New(rand.NewSource(20261007))
		mean, m2 := 0.0, 0.0
		for trial := 1; trial <= trials; trial++ {
			// Floyd's algorithm draws distinct occurrence indices without replacement.
			positions := make(map[int]bool, sampleSize)
			words := make(map[int]bool, sampleSize)
			for j := total - sampleSize; j < total; j++ {
				position := rng.Intn(j + 1)
				if positions[position] {
					position = j
				}
				positions[position] = true
				word := sort.Search(len(cumulative), func(i int) bool { return cumulative[i] > position })
				words[word] = true
			}
			value := float64(len(words))
			delta := value - mean
			mean += delta / float64(trial)
			m2 += delta * (value - mean)
		}
		results = append(results, result{author[0], total, mean, math.Sqrt(m2 / float64(trials-1))})
	}
	sort.Slice(results, func(i, j int) bool { return results[i].mean > results[j].mean })
	rows := [][]string{{"Author", "Total words", "Sample words", "Trials", "Mean unique words", "Mean unique / sample words", "Sample standard deviation"}}
	for _, r := range results {
		rows = append(rows, []string{r.author, strconv.Itoa(r.total), strconv.Itoa(sampleSize), strconv.Itoa(trials), fmt.Sprintf("%.4f", r.mean), fmt.Sprintf("%.8f", r.mean/sampleSize), fmt.Sprintf("%.4f", r.sd)})
	}
	if err := writeTSV("outputs/author-vocabulary-sampled.tsv", rows); err != nil {
		return err
	}
	if err := writeTSV("outputs/author-vocabulary-sampled-excluded.tsv", append([][]string{{"Author", "Total words", "Reason"}}, excluded...)); err != nil {
		return err
	}
	fmt.Printf("Compared %d authors; excluded %d; sample=%d, trials=%d\n", len(results), len(excluded), sampleSize, trials)
	for _, r := range results {
		fmt.Printf("%s\t%.4f\n", r.author, r.mean)
	}
	return nil
}

func writeTSV(path string, rows [][]string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	w := csv.NewWriter(file)
	w.Comma = '\t'
	err = w.WriteAll(rows)
	closeErr := file.Close()
	if err != nil {
		return err
	}
	return closeErr
}