package magstore

import (
	"context"
	"fmt"
	"math"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/missionworkforce/missionmatch/internal/platform/learn"
)

// DefaultCalibrationBins is the cold start resolution of the Platt calibration map: ten monotone bins of width
// 0.1 over the raw probability range [0, 1]. A later isotonic fit past roughly a thousand labels can write a
// finer map.
const DefaultCalibrationBins = 10

// TrainResult summarizes a training run: whether it was skipped (and why), and, when it ran, the challenger
// version written plus the fitted raw feature space coefficients and the sample composition. It is the value
// both cmd/train-model logs and the integration test asserts on.
type TrainResult struct {
	Skipped      bool
	SkipReason   string
	Version      int
	N            int
	Wins         int
	Losses       int
	FeatureNames []string
	Intercept    float64
	Coeffs       []float64 // raw feature space, aligned to FeatureNames
	Iterations   int
	Bins         int
}

// Train reads the pooled decided bid labels, fits an interpretable win probability logistic model, and writes a
// new CHALLENGER version (never a champion). It refuses to fit (Skipped) below the events per predictor floor or
// on a single class, leaving the prior only champion in force. It converts the solver's standardized
// coefficients into raw feature space (so the serve path applies them directly to gathered raw feature values)
// and materializes a Platt calibration map into monotone bins. When dryRun is set it fits and returns the
// coefficients but writes nothing. minEvents <= 0 auto computes the floor as twenty decided events per fitted
// parameter (the intercept plus one per feature).
func Train(ctx context.Context, pool *pgxpool.Pool, modelKey string, l2 float64, minEvents int, dryRun bool) (TrainResult, error) {
	rows, err := ReadLabeledBids(ctx, pool)
	if err != nil {
		return TrainResult{}, err
	}
	x, y, featureNames := learn.BuildDataset(rows)
	var res TrainResult
	res.FeatureNames = featureNames
	for _, v := range y {
		res.N++
		if v == 1 {
			res.Wins++
		} else {
			res.Losses++
		}
	}

	floor := minEvents
	if floor <= 0 {
		floor = 20 * (len(featureNames) + 1)
	}
	if res.N < floor {
		res.Skipped = true
		res.SkipReason = fmt.Sprintf("%d decided bids is below the events floor of %d", res.N, floor)
		return res, nil
	}
	if res.Wins == 0 || res.Losses == 0 {
		res.Skipped = true
		res.SkipReason = fmt.Sprintf("the labels are a single class (%d win, %d loss)", res.Wins, res.Losses)
		return res, nil
	}

	model, err := learn.FitLogistic(x, y, l2)
	if err != nil {
		return TrainResult{}, err
	}
	res.Iterations = model.Iterations

	// Standardized to raw feature space: rawCoef_j = beta_j / scale_j;
	// rawIntercept = beta0 - SUM_j beta_j * mean_j / scale_j.
	res.Intercept = model.Intercept
	res.Coeffs = make([]float64, len(featureNames))
	for j := range featureNames {
		res.Coeffs[j] = model.Coeffs[j] / model.Scales[j]
		res.Intercept -= model.Coeffs[j] * model.Means[j] / model.Scales[j]
	}

	// The in sample log odds for the Platt calibration fit.
	etas := make([]float64, len(x))
	for i := range x {
		eta := res.Intercept
		for j := range featureNames {
			eta += res.Coeffs[j] * x[i][j]
		}
		etas[i] = eta
	}
	bins := materializeCalibration(etas, y, DefaultCalibrationBins)
	res.Bins = len(bins)

	if dryRun {
		return res, nil
	}

	notes := fmt.Sprintf("Challenger fit by cmd/train-model over %d decided bids (%d win, %d loss) pooled across tenants, ridge l2=%v. OutputKind opportunity_score, never a candidate score. champion=false until cmd/eval-model passes the four fifths gate and a human promotes it.", res.N, res.Wins, res.Losses, l2)
	// The training feature means are stored in the feature_spec so the serve path can impute a feature with no
	// history (a segment with no decided bids) at its mean, keeping a learned model at the base rate for a no
	// history segment rather than extrapolating off a misleading zero.
	means := make([]float64, len(featureNames))
	copy(means, model.Means)
	version, err := WriteChallenger(ctx, pool, modelKey, Challenger{
		FeatureNames: featureNames,
		Intercept:    res.Intercept,
		Coeffs:       res.Coeffs,
		Means:        means,
		Bins:         bins,
		Notes:        notes,
	})
	if err != nil {
		return TrainResult{}, err
	}
	res.Version = version
	return res, nil
}

// materializeCalibration fits Platt scaling on the in sample log odds and materializes it into monotone
// calibration bins over the raw probability range [0, 1], the shape aedt_calibration_bins stores and the serve
// path looks up. For each bin the empirical rate is the Platt calibrated probability at the bin midpoint
// (calibrated = sigmoid(a * logit(mid) + b)). A Platt fit failure returns no bins, the honest identity map.
func materializeCalibration(etas []float64, y []int, binCount int) []Bin {
	a, b, err := learn.FitPlatt(etas, y)
	if err != nil {
		return nil
	}
	bins := make([]Bin, 0, binCount)
	for k := 0; k < binCount; k++ {
		lo := float64(k) / float64(binCount)
		hi := float64(k+1) / float64(binCount)
		mid := (lo + hi) / 2
		score := math.Log(mid / (1 - mid)) // logit of the midpoint; midpoints 0.05..0.95 are always in range
		cal := 1.0 / (1.0 + math.Exp(-(a*score + b)))
		bins = append(bins, Bin{Lo: lo, Hi: hi, EmpiricalRate: cal})
	}
	return bins
}