package mathkit

import (
	"fmt"
	"math"
)

// Sigmoid is the logistic function 1 / (1 + e^-x), computed in the
// numerically stable branch-on-sign form so it doesn't overflow
// math.Exp for very negative x (a plain 1/(1+exp(-x)) blows up
// exp(-x) to +Inf for x below about -745 and silently returns 0
// instead of propagating the overflow visibly).
func Sigmoid(x float64) float64 {
	if x >= 0 {
		z := math.Exp(-x)
		return 1 / (1 + z)
	}
	z := math.Exp(x)
	return z / (1 + z)
}

// ReLU is max(0, x).
func ReLU(x float64) float64 {
	if x > 0 {
		return x
	}
	return 0
}

// Softmax turns a vector of raw scores ("logits") into a probability
// distribution: non-negative entries that sum to 1. Subtracts the max
// before exponentiating (the standard numerically-stable formulation)
// so large logits don't overflow math.Exp.
func Softmax(x []float64) []float64 {
	if len(x) == 0 {
		return nil
	}
	max := x[0]
	for _, v := range x[1:] {
		if v > max {
			max = v
		}
	}
	out := make([]float64, len(x))
	var sum float64
	for i, v := range x {
		e := math.Exp(v - max)
		out[i] = e
		sum += e
	}
	for i := range out {
		out[i] /= sum
	}
	return out
}

// MSELoss returns the mean squared error between predictions and
// targets: mean((pred - target)^2), the default regression loss.
func MSELoss(predictions, targets []float64) (float64, error) {
	if len(predictions) != len(targets) {
		return 0, fmt.Errorf("mse_loss: predictions and targets have different lengths (%d vs %d)", len(predictions), len(targets))
	}
	if len(predictions) == 0 {
		return 0, fmt.Errorf("mse_loss: empty input")
	}
	var sum float64
	for i := range predictions {
		d := predictions[i] - targets[i]
		sum += d * d
	}
	return sum / float64(len(predictions)), nil
}

// crossEntropyEps clamps predicted probabilities away from exactly 0
// or 1 before taking a log — ln(0) is -Inf, which would make a single
// confidently-wrong prediction blow up the whole averaged loss to
// +Inf instead of just contributing a large-but-finite penalty. Every
// production cross-entropy implementation (e.g. scikit-learn's
// log_loss) clips for exactly this reason.
const crossEntropyEps = 1e-12

// CrossEntropyLoss returns the mean binary cross-entropy between
// predicted probabilities and 0/1 (or soft, in [0,1]) targets:
// mean(-(y*ln(p) + (1-y)*ln(1-p))), the default classification loss.
func CrossEntropyLoss(predictions, targets []float64) (float64, error) {
	if len(predictions) != len(targets) {
		return 0, fmt.Errorf("cross_entropy_loss: predictions and targets have different lengths (%d vs %d)", len(predictions), len(targets))
	}
	if len(predictions) == 0 {
		return 0, fmt.Errorf("cross_entropy_loss: empty input")
	}
	var sum float64
	for i, p := range predictions {
		if p < 0 || p > 1 {
			return 0, fmt.Errorf("cross_entropy_loss: predictions must be probabilities in [0, 1], got %g at index %d", p, i)
		}
		p = math.Min(math.Max(p, crossEntropyEps), 1-crossEntropyEps)
		y := targets[i]
		sum += -(y*math.Log(p) + (1-y)*math.Log(1-p))
	}
	return sum / float64(len(predictions)), nil
}

// LinearRegressionResult is a fitted ordinary-least-squares linear
// model.
type LinearRegressionResult struct {
	// Coefficients[0] is the intercept, Coefficients[1:] line up with
	// the feature columns of X in order.
	Coefficients []float64
	Predictions  []float64
	Residuals    []float64
	RSquared     float64
}

// LinearRegression fits y ≈ X·β + intercept by ordinary least squares
// via the normal equations β = (XᵀX)^-1 Xᵀy, with an intercept column
// of 1s prepended automatically (the caller supplies only the real
// feature columns in X). The normal-equation approach is the direct,
// closed-form solution to minimizing squared error — the textbook
// ML-101 regression fit, as opposed to reaching it iteratively via
// GradientDescent on the MSE loss.
func LinearRegression(X Matrix, y []float64) (*LinearRegressionResult, error) {
	rows, cols, err := X.Dims()
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}
	if rows != len(y) {
		return nil, fmt.Errorf("linear_regression: X has %d rows but y has %d entries", rows, len(y))
	}
	if rows <= cols {
		return nil, fmt.Errorf("linear_regression: need more observations (%d) than features+1 (%d) to fit uniquely", rows, cols+1)
	}

	// Prepend the intercept column of 1s.
	design := make(Matrix, rows)
	for i := 0; i < rows; i++ {
		row := make([]float64, cols+1)
		row[0] = 1
		copy(row[1:], X[i])
		design[i] = row
	}

	designT, err := Transpose(design)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}
	xtx, err := MatMul(designT, design)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}
	xtxInv, err := Inverse(xtx)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: features are collinear (XᵀX is singular): %w", err)
	}
	xty, err := MatVecMul(designT, y)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}
	beta, err := MatVecMul(xtxInv, xty)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}

	preds, err := MatVecMul(design, beta)
	if err != nil {
		return nil, fmt.Errorf("linear_regression: %w", err)
	}
	residuals := make([]float64, rows)
	var ssRes float64
	for i := range preds {
		residuals[i] = y[i] - preds[i]
		ssRes += residuals[i] * residuals[i]
	}
	meanY, _ := Mean(y)
	var ssTot float64
	for _, v := range y {
		d := v - meanY
		ssTot += d * d
	}
	rSquared := 1.0
	if ssTot > 0 {
		rSquared = 1 - ssRes/ssTot
	}

	return &LinearRegressionResult{
		Coefficients: beta,
		Predictions:  preds,
		Residuals:    residuals,
		RSquared:     rSquared,
	}, nil
}