/*
 * # Licensed to the LF AI & Data foundation under one
 * # or more contributor license agreements. See the NOTICE file
 * # distributed with this work for additional information
 * # regarding copyright ownership. The ASF licenses this file
 * # to you under the Apache License, Version 2.0 (the
 * # "License"); you may not use this file except in compliance
 * # with the License. You may obtain a copy of the License at
 * #
 * #     http://www.apache.org/licenses/LICENSE-2.0
 * #
 * # Unless required by applicable law or agreed to in writing, software
 * # distributed under the License is distributed on an "AS IS" BASIS,
 * # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * # See the License for the specific language governing permissions and
 * # limitations under the License.
 */

package expr

import (
	"math"

	"github.com/apache/arrow/go/v17/arrow"
	"github.com/apache/arrow/go/v17/arrow/array"

	"github.com/milvus-io/milvus/internal/util/function/chain/types"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
)

// =============================================================================
// Constants (use types package constants)
// =============================================================================

const (
	// Decay function types
	GaussFunction  = types.DecayFuncGauss
	LinearFunction = types.DecayFuncLinear
	ExpFunction    = types.DecayFuncExp

	// Parameter keys for DecayExpr
	FunctionKey = types.DecayParamFunction
	OriginKey   = types.DecayParamOrigin
	ScaleKey    = types.DecayParamScale
	OffsetKey   = types.DecayParamOffset
	DecayKey    = types.DecayParamDecay
)

// =============================================================================
// Types
// =============================================================================

// decayReScorer is a function type for decay calculation.
type decayReScorer func(origin, scale, decay, offset, distance float64) float64

// DecayExpr implements FunctionExpr for decay scoring.
// It takes a numeric input column and outputs the pure decay factor (0~1).
// The combination with $score is handled by a subsequent NumCombineExpr node in the chain.
// Column mapping is handled by MapOp.
//
// Expected inputs (passed from MapOp):
//   - inputs[0]: numeric column to calculate distance from origin
//
// Outputs:
//   - outputs[0]: decay factor column (Float32, values in [0, 1])
type DecayExpr struct {
	BaseExpr
	function  string        // "gauss", "exp", or "linear"
	origin    float64       // origin point
	scale     float64       // scale parameter (must > 0)
	offset    float64       // offset (default 0, must >= 0)
	decay     float64       // decay factor (default 0.5, 0 < decay < 1)
	decayFunc decayReScorer // selected decay function
}

// =============================================================================
// Decay Calculation Functions
// =============================================================================

// gaussianDecay calculates Gaussian decay.
func gaussianDecay(origin, scale, decay, offset, distance float64) float64 {
	adjustedDist := math.Max(0, math.Abs(distance-origin)-offset)
	sigmaSquare := math.Pow(scale, 2.0) / math.Log(decay)
	exponent := math.Pow(adjustedDist, 2.0) / sigmaSquare
	return math.Exp(exponent)
}

// expDecay calculates exponential decay.
func expDecay(origin, scale, decay, offset, distance float64) float64 {
	adjustedDist := math.Max(0, math.Abs(distance-origin)-offset)
	lambda := math.Log(decay) / scale
	return math.Exp(lambda * adjustedDist)
}

// linearDecay calculates linear decay.
func linearDecay(origin, scale, decay, offset, distance float64) float64 {
	adjustedDist := math.Max(0, math.Abs(distance-origin)-offset)
	slope := (1 - decay) / scale
	return math.Max(decay, 1-slope*adjustedDist)
}

// =============================================================================
// Constructor Functions
// =============================================================================

// minDecayValue is the minimum allowed decay value to ensure numerical stability.
// When decay is too close to 0, log(decay) approaches -∞ causing numerical instability.
const minDecayValue = 0.001

// maxDecayValue is the maximum allowed decay value to ensure numerical stability.
// When decay is too close to 1, log(decay) approaches 0 causing division issues.
const maxDecayValue = 0.999

// NewDecayExpr creates a new DecayExpr with the given parameters.
// Note: Column mapping (which columns to use as input/output) is handled by MapOp,
// not by the function itself.
func NewDecayExpr(function string, origin, scale, offset, decay float64) (*DecayExpr, error) {
	if scale <= 0 {
		return nil, merr.WrapErrParameterInvalidMsg("decay: scale must be > 0, got %f", scale)
	}

	if offset < 0 {
		return nil, merr.WrapErrParameterInvalidMsg("decay: offset must be >= 0, got %f", offset)
	}

	if decay <= 0 || decay >= 1 {
		return nil, merr.WrapErrParameterInvalidMsg("decay: decay must be 0 < decay < 1, got %f", decay)
	}

	// Additional check for numerical stability
	if decay < minDecayValue || decay > maxDecayValue {
		return nil, merr.WrapErrParameterInvalidMsg("decay: decay must be between %f and %f for numerical stability, got %f",
			minDecayValue, maxDecayValue, decay)
	}

	expr := &DecayExpr{
		BaseExpr: *NewBaseExpr("decay", types.AllStages),
		function: function,
		origin:   origin,
		scale:    scale,
		offset:   offset,
		decay:    decay,
	}

	// Select decay function
	switch function {
	case GaussFunction:
		expr.decayFunc = gaussianDecay
	case ExpFunction:
		expr.decayFunc = expDecay
	case LinearFunction:
		expr.decayFunc = linearDecay
	default:
		return nil, merr.WrapErrParameterInvalidMsg("decay: invalid function %q, must be one of [%s, %s, %s]",
			function, GaussFunction, ExpFunction, LinearFunction)
	}

	return expr, nil
}

// NewDecayExprFromParams creates a DecayExpr from a parameter map.
// This is the factory function for the function registry.
// All parameter parsing is handled here, keeping it close to the expr definition.
func NewDecayExprFromParams(_ types.FunctionBuildContext, cfg types.FunctionConfig) (types.FunctionExpr, error) {
	const funcName = "decay"
	reader := types.NewParamReader(funcName, cfg.Params)

	function, err := reader.String(FunctionKey, true)
	if err != nil {
		return nil, err
	}

	origin, err := reader.Float64(OriginKey, true, 0)
	if err != nil {
		return nil, err
	}

	scale, err := reader.Float64(ScaleKey, true, 0)
	if err != nil {
		return nil, err
	}

	offset, err := reader.Float64(OffsetKey, false, 0)
	if err != nil {
		return nil, err
	}

	decayVal, err := reader.Float64(DecayKey, false, 0.5)
	if err != nil {
		return nil, err
	}

	return NewDecayExpr(function, origin, scale, offset, decayVal)
}

// =============================================================================
// FunctionExpr Interface Implementation
// =============================================================================

// Name() and IsRunnable() are inherited from BaseExpr

// OutputDataTypes returns the data types of output columns.
// DecayExpr outputs a single Float32 column (the decay factor).
func (d *DecayExpr) OutputDataTypes() []arrow.DataType {
	return []arrow.DataType{arrow.PrimitiveTypes.Float32}
}

// Execute executes the decay function on input columns and returns the decay factor.
// inputs[0]: the numeric column to calculate decay from
// returns: decay factor column (Float32, values in [0, 1])
func (d *DecayExpr) Execute(ctx *types.FuncContext, inputs []*arrow.Chunked) ([]*arrow.Chunked, error) {
	if len(inputs) != 1 {
		return nil, merr.WrapErrServiceInternalMsg("decay: expected 1 input column, got %d", len(inputs))
	}

	inputCol := inputs[0]

	numChunks := len(inputCol.Chunks())
	decayChunks := make([]arrow.Array, numChunks)

	for chunkIdx := 0; chunkIdx < numChunks; chunkIdx++ {
		inputChunk := inputCol.Chunk(chunkIdx)

		decayChunk, err := d.processChunk(ctx, inputChunk)
		if err != nil {
			for i := 0; i < chunkIdx; i++ {
				decayChunks[i].Release()
			}
			return nil, err
		}
		decayChunks[chunkIdx] = decayChunk
	}

	result := arrow.NewChunked(arrow.PrimitiveTypes.Float32, decayChunks)

	for _, chunk := range decayChunks {
		chunk.Release()
	}

	return []*arrow.Chunked{result}, nil
}

// =============================================================================
// Internal Processing Methods
// =============================================================================

// processChunk processes a single chunk, calculating decay factors.
func (d *DecayExpr) processChunk(ctx *types.FuncContext, inputChunk arrow.Array) (arrow.Array, error) {
	builder := array.NewFloat32Builder(ctx.Pool())
	defer builder.Release()

	for i := range inputChunk.Len() {
		if inputChunk.IsNull(i) {
			builder.AppendNull()
			continue
		}

		distance, err := GetNumericValue(inputChunk, i)
		if err != nil {
			return nil, merr.WrapErrServiceInternalMsg("decay: %v", err)
		}

		decayScore := d.decayFunc(d.origin, d.scale, d.decay, d.offset, distance)
		builder.Append(float32(decayScore))
	}

	return builder.NewArray(), nil
}

// =============================================================================
// Registration
// =============================================================================

func init() {
	types.MustRegisterFunction("decay", NewDecayExprFromParams)
}