// Train command: labels.jsonl → model.json + report.json.
//
// Time-split validation (train on older labels, evaluate on the newest
// window — what serve actually faces). Headline metric is precision@15
// (rank-based: of the 15 highest-scored validation entries, what fraction
// are positives). Youden's J is computed and reported for information only;
// it is NOT used as a cutoff — serve ranks by raw score.
//
// Reproducible: a fixed seed makes the downsampling deterministic, so the
// same labels produce the same weights/vocab. IDF is fit on the full
// (unbalanced) train set to preserve document frequencies, then the
// majority class is downsampled to 1:1 before LR fit (matches the vendored
// scholscan approach).
package cmds

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"math"
	"math/rand"
	"os"
	"sort"
	"time"

	"fluxrec/core"
)

// trainSeed fixes the downsampling RNG: same labels + same seed → same model.
const trainSeed int64 = 20240301

type TrainCommand struct {
	Labels    string
	ModelOut  string
	ReportOut string
	ValFrac   float64
	Lambda    float64
	MinDF     int
	MaxDF     float64
	NgramMax  int
}

func (c *TrainCommand) Name() string { return "train" }

func (c *TrainCommand) Init(args []string) error {
	fs := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprint(fs.Output(), `Usage: fluxrec train LABELS_JSONL [flags]

Train a TF-IDF + logistic regression model from labels.jsonl. Writes the
model to --model-out and a validation report (precision@15 etc.) to
--report-out. Time-split: the newest --val-frac of labels (by captured_at)
are held out for validation; the rest train the model.
`)
		fs.PrintDefaults()
	}
	fs.StringVar(&c.ModelOut, "model-out", "model.json", "Path to write the trained model JSON")
	fs.StringVar(&c.ReportOut, "report-out", "report.json", "Path to write the validation report JSON")
	fs.Float64Var(&c.ValFrac, "val-frac", 0.2, "Fraction of newest labels held out for validation")
	fs.Float64Var(&c.Lambda, "lambda", 0.001, "L2 regularization strength")
	fs.IntVar(&c.MinDF, "min-df", 2, "Minimum document frequency (absolute)")
	fs.Float64Var(&c.MaxDF, "max-df", 0.8, "Maximum document frequency (ratio)")
	fs.IntVar(&c.NgramMax, "ngram-max", 2, "Max n-gram size (1=unigrams, 2=+bigrams)")

	for _, a := range args {
		if a == "--help" || a == "-h" {
			fs.Usage()
			return flag.ErrHelp
		}
	}
	if len(args) == 0 {
		return fmt.Errorf("labels.jsonl path argument is required")
	}
	c.Labels = args[0]
	if err := fs.Parse(args[1:]); err != nil {
		return err
	}
	if c.ValFrac <= 0 || c.ValFrac >= 1 {
		return fmt.Errorf("--val-frac must be between 0 and 1, got %v", c.ValFrac)
	}
	return nil
}

func (c *TrainCommand) Run(stdin io.Reader, stdout io.Writer) error {
	rows, err := core.ReadLabels(c.Labels)
	if err != nil {
		return err
	}
	if len(rows) == 0 {
		return fmt.Errorf("no labels in %s", c.Labels)
	}

	trainRows, valRows := core.SplitByTime(rows, c.ValFrac)
	if len(trainRows) == 0 || len(valRows) == 0 {
		return fmt.Errorf("time split produced an empty side (train=%d val=%d); adjust --val-frac or accumulate more labels",
			len(trainRows), len(valRows))
	}
	valPos := 0
	for _, r := range valRows {
		if r.Label == core.LabelPos {
			valPos++
		}
	}
	if valPos < 5 {
		return fmt.Errorf("only %d positives in validation (need >=5 for signal, not noise); lower --val-frac or accumulate more labels",
			valPos)
	}

	model, report, err := c.trainAndEvaluate(trainRows, valRows, time.Now().UTC())
	if err != nil {
		return err
	}

	if err := writeJSON(c.ModelOut, model); err != nil {
		return err
	}
	if err := writeJSON(c.ReportOut, report); err != nil {
		return err
	}

	fmt.Fprintf(stdout, "train: pos=%d neg=%d | val: pos=%d neg=%d | precision@15=%.3f youden_j=%.3f | vocab=%d → %s\n",
		report.TrainPos, report.TrainNeg, report.ValPos, report.ValNeg,
		report.PrecisionAt15, report.YoudensJ, report.VocabSize, c.ModelOut)
	return nil
}

func writeJSON(path string, v any) error {
	body, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	if err := os.WriteFile(path, body, 0o644); err != nil {
		return fmt.Errorf("writing %s: %w", path, err)
	}
	return nil
}

// trainAndEvaluate is the pure training core: given split rows and a clock
// value, returns the model envelope and validation report. Pure w.r.t. the
// inputs (no I/O, no wall clock) so two calls with the same rows + now are
// byte-identical — the reproducibility guarantee.
func (c *TrainCommand) trainAndEvaluate(trainRows, valRows []core.LabelRow, now time.Time) (*core.ModelEnvelope, *trainReport, error) {
	trainDocs := make([]string, len(trainRows))
	trainLabels := make([]float64, len(trainRows))
	trainPos, trainNeg := 0, 0
	for i, r := range trainRows {
		trainDocs[i] = core.ScoringText(r.Title, r.FeedTitle)
		if r.Label == core.LabelPos {
			trainLabels[i] = 1.0
			trainPos++
		} else {
			trainLabels[i] = 0.0
			trainNeg++
		}
	}
	valDocs := make([]string, len(valRows))
	valLabels := make([]float64, len(valRows))
	valPos, valNeg := 0, 0
	for i, r := range valRows {
		valDocs[i] = core.ScoringText(r.Title, r.FeedTitle)
		if r.Label == core.LabelPos {
			valLabels[i] = 1.0
			valPos++
		} else {
			valLabels[i] = 0.0
			valNeg++
		}
	}

	// Fit IDF on the full (unbalanced) train set to preserve document
	// frequencies, then downsample the majority to 1:1 before LR fit.
	vectorizer := &core.TFIDFVectorizer{
		NgramMin: 1, NgramMax: c.NgramMax,
		MinDF: c.MinDF, MaxDF: c.MaxDF, VocabCap: 50000,
		Vocabulary: make(map[string]float64),
	}
	vectorizer.Fit(trainDocs)

	balDocs, balLabels := downsampleBalance(trainDocs, trainLabels, trainSeed)
	trainVectors := vectorizer.Transform(balDocs)
	valVectors := vectorizer.Transform(valDocs)

	lr := &core.LogisticRegression{
		LearningRate: 0.5, Lambda: c.Lambda, Iterations: 500, Tolerance: 0.000001,
	}
	lr.Validate()
	weights, err := lr.Fit(trainVectors, balLabels, map[float64]float64{1.0: 1.0, 0.0: 1.0})
	if err != nil {
		return nil, nil, fmt.Errorf("training logistic regression: %w", err)
	}

	precisionAt15, youdenJ, posStats, negStats := evaluate(valVectors, valLabels, weights)

	minCap, maxCap := trainRows[0].CapturedAt, trainRows[0].CapturedAt
	for _, r := range trainRows {
		if r.CapturedAt.Before(minCap) {
			minCap = r.CapturedAt
		}
		if r.CapturedAt.After(maxCap) {
			maxCap = r.CapturedAt
		}
	}

	vecParams := map[string]any{
		"ngram_min": vectorizer.NgramMin, "ngram_max": vectorizer.NgramMax,
		"min_df": vectorizer.MinDF, "max_df": vectorizer.MaxDF, "vocab_cap": vectorizer.VocabCap,
	}
	modelParams := map[string]any{
		"learning_rate": lr.LearningRate, "lambda": lr.Lambda,
		"iterations": lr.Iterations, "tolerance": lr.Tolerance,
	}

	model := &core.ModelEnvelope{
		Algorithm: "tfidf-go", Impl: "go", Version: "1",
		CreatedAt: now,
		Meta: map[string]any{
			"scorer":            "tfidf",
			"seed":              trainSeed,
			"train_size":        len(trainRows),
			"val_size":          len(valRows),
			"train_pos":         trainPos,
			"train_neg":         trainNeg,
			"train_window_min":  minCap.UTC().Format(time.RFC3339),
			"train_window_max":  maxCap.UTC().Format(time.RFC3339),
			"val_frac":          c.ValFrac,
			"precision_at_15":   precisionAt15,
			"youdens_j":         youdenJ,
			"vectorizer_params": vecParams,
			"model_params":      modelParams,
		},
		Vectorizer:   vectorizer.Vocabulary,
		OrderedVocab: vectorizer.OrderedVocab,
		Weights:      weights,
	}

	report := &trainReport{
		Scorer: "tfidf", Seed: trainSeed,
		TrainSize: len(trainRows), ValSize: len(valRows),
		TrainPos: trainPos, TrainNeg: trainNeg,
		ValPos: valPos, ValNeg: valNeg,
		TrainWindowMin: minCap.UTC(), TrainWindowMax: maxCap.UTC(),
		ValFrac: c.ValFrac, VocabSize: len(vectorizer.OrderedVocab),
		PrecisionAt15: precisionAt15, YoudensJ: youdenJ,
		PositiveScores: posStats, NegativeScores: negStats,
		VectorizerParams: vecParams, ModelParams: modelParams,
	}
	return model, report, nil
}

// downsampleBalance downsamples the majority class to a 1:1 ratio with the
// minority, using a seeded RNG for reproducibility. If already balanced,
// returns the inputs unchanged.
func downsampleBalance(docs []string, labels []float64, seed int64) ([]string, []float64) {
	var posDocs, negDocs []string
	var posLabels, negLabels []float64
	for i, label := range labels {
		if label == 1.0 {
			posDocs = append(posDocs, docs[i])
			posLabels = append(posLabels, label)
		} else {
			negDocs = append(negDocs, docs[i])
			negLabels = append(negLabels, label)
		}
	}
	if len(posDocs) == len(negDocs) {
		return docs, labels
	}
	var majorityDocs, minorityDocs []string
	var majorityLabels, minorityLabels []float64
	if len(negDocs) > len(posDocs) {
		majorityDocs, minorityDocs = negDocs, posDocs
		majorityLabels, minorityLabels = negLabels, posLabels
	} else {
		majorityDocs, minorityDocs = posDocs, negDocs
		majorityLabels, minorityLabels = posLabels, negLabels
	}
	minoritySize := len(minorityDocs)
	rng := rand.New(rand.NewSource(seed))
	indices := make([]int, len(majorityDocs))
	for i := range indices {
		indices[i] = i
	}
	rng.Shuffle(len(indices), func(i, j int) {
		indices[i], indices[j] = indices[j], indices[i]
	})
	outDocs := make([]string, 0, minoritySize*2)
	outLabels := make([]float64, 0, minoritySize*2)
	outDocs = append(outDocs, minorityDocs...)
	outLabels = append(outLabels, minorityLabels...)
	for i := 0; i < minoritySize; i++ {
		outDocs = append(outDocs, majorityDocs[indices[i]])
		outLabels = append(outLabels, majorityLabels[indices[i]])
	}
	return outDocs, outLabels
}

// evaluate scores the validation set and returns precision@15 (headline),
// Youden's J (information only), and per-class score stats.
func evaluate(valVectors [][]float64, valLabels []float64, weights []float64) (precisionAt15, youdenJ float64, posStats, negStats map[string]float64) {
	type scored struct {
		score float64
		label float64
		idx   int
	}
	ss := make([]scored, len(valVectors))
	for i, v := range valVectors {
		s, err := core.PredictScore(v, weights)
		if err != nil {
			s = 0
		}
		ss[i] = scored{s, valLabels[i], i}
	}

	// precision@15: rank by score desc (ties by idx), take top 15, count positives.
	ranked := make([]scored, len(ss))
	copy(ranked, ss)
	sort.SliceStable(ranked, func(i, j int) bool {
		if ranked[i].score == ranked[j].score {
			return ranked[i].idx < ranked[j].idx
		}
		return ranked[i].score > ranked[j].score
	})
	k := 15
	if k > len(ranked) {
		k = len(ranked)
	}
	topPos := 0
	for i := 0; i < k; i++ {
		if ranked[i].label == 1.0 {
			topPos++
		}
	}
	if k > 0 {
		precisionAt15 = float64(topPos) / float64(k)
	}

	// Youden's J (information only): best sensitivity + specificity - 1 over
	// a threshold sweep. Not used as a cutoff anywhere.
	bestJ := -1.0
	for t := 5; t <= 95; t++ {
		threshold := float64(t) / 100.0
		var tp, fp, tn, fn int
		for _, s := range ss {
			pred := s.score > threshold
			switch {
			case pred && s.label == 1.0:
				tp++
			case pred && s.label == 0.0:
				fp++
			case !pred && s.label == 1.0:
				fn++
			default:
				tn++
			}
		}
		sens := 0.0
		if tp+fn > 0 {
			sens = float64(tp) / float64(tp+fn)
		}
		spec := 0.0
		if tn+fp > 0 {
			spec = float64(tn) / float64(tn+fp)
		}
		if j := sens + spec - 1.0; j > bestJ {
			bestJ = j
		}
	}
	youdenJ = bestJ

	var pos, neg []float64
	for _, s := range ss {
		if s.label == 1.0 {
			pos = append(pos, s.score)
		} else {
			neg = append(neg, s.score)
		}
	}
	posStats = scoreStats(pos)
	negStats = scoreStats(neg)
	return
}

// scoreStats returns min/max/mean/std for a slice of scores.
func scoreStats(scores []float64) map[string]float64 {
	if len(scores) == 0 {
		return map[string]float64{"min": 0, "max": 0, "mean": 0, "std": 0}
	}
	min, max, sum := scores[0], scores[0], 0.0
	for _, s := range scores {
		if s < min {
			min = s
		}
		if s > max {
			max = s
		}
		sum += s
	}
	mean := sum / float64(len(scores))
	variance := 0.0
	for _, s := range scores {
		variance += (s - mean) * (s - mean)
	}
	variance /= float64(len(scores))
	return map[string]float64{
		"min": min, "max": max, "mean": mean, "std": math.Sqrt(variance),
	}
}

// trainReport is the validation report written to --report-out.
type trainReport struct {
	Scorer           string             `json:"scorer"`
	Seed             int64              `json:"seed"`
	TrainSize        int                `json:"train_size"`
	ValSize          int                `json:"val_size"`
	TrainPos         int                `json:"train_pos"`
	TrainNeg         int                `json:"train_neg"`
	ValPos           int                `json:"val_pos"`
	ValNeg           int                `json:"val_neg"`
	TrainWindowMin   time.Time          `json:"train_window_min"`
	TrainWindowMax   time.Time          `json:"train_window_max"`
	ValFrac          float64            `json:"val_frac"`
	VocabSize        int                `json:"vocab_size"`
	PrecisionAt15    float64            `json:"precision_at_15"`
	YoudensJ         float64            `json:"youdens_j"`
	PositiveScores   map[string]float64 `json:"positive_scores"`
	NegativeScores   map[string]float64 `json:"negative_scores"`
	VectorizerParams map[string]any     `json:"vectorizer_params"`
	ModelParams      map[string]any     `json:"model_params"`
}