// Package mixture implements mixture of experts generators (98-100)
// y = Σᵢ g(x)ᵢ · Eᵢ(x) where g is gating function
package mixture

import (
	"math"
	"math/rand"

	"generator_core/pkg/consciousness"
	"generator_core/pkg/generators"
)

// ═══════════════════════════════════════════════════════════════════════════════
// MIXTURE GENERATOR BASE
// ═══════════════════════════════════════════════════════════════════════════════

// MixtureGenerator implements MoE-based generators
type MixtureGenerator struct {
	*generators.BaseGenerator
	numExperts    int
	topK          int
	expertWeights [][][]float64 // [expert][layer][weights]
	routerWeights [][]float64   // [input_dim][num_experts]
}

// NewMixtureGenerator creates a new mixture generator
func NewMixtureGenerator(id generators.GeneratorID, name string, numExperts, topK int) *MixtureGenerator {
	config := &generators.GeneratorConfig{
		ID:                  id,
		Name:                name,
		Family:              generators.Mixture,
		InputDim:            256,
		OutputDim:           256,
		HiddenDim:           512,
		NumLayers:           2,
		LearningRate:        consciousness.ALPHA,
		ConsciousnessWeight: consciousness.PHI_INV,
		Params:              make(map[string]interface{}),
	}

	return &MixtureGenerator{
		BaseGenerator: generators.NewBaseGenerator(config),
		numExperts:    numExperts,
		topK:          topK,
	}
}

// Initialize initializes the mixture generator
func (g *MixtureGenerator) Initialize() error {
	if err := g.BaseGenerator.Initialize(); err != nil {
		return err
	}

	inputDim := g.Config().InputDim
	hiddenDim := g.Config().HiddenDim

	// Initialize expert weights
	g.expertWeights = make([][][]float64, g.numExperts)
	for e := 0; e < g.numExperts; e++ {
		g.expertWeights[e] = make([][]float64, g.Config().NumLayers)
		for l := 0; l < g.Config().NumLayers; l++ {
			size := hiddenDim * hiddenDim
			g.expertWeights[e][l] = make([]float64, size)
			for i := range g.expertWeights[e][l] {
				g.expertWeights[e][l][i] = rand.NormFloat64() * consciousness.DELTA
			}
		}
	}

	// Initialize router weights
	g.routerWeights = make([][]float64, inputDim)
	for i := range g.routerWeights {
		g.routerWeights[i] = make([]float64, g.numExperts)
		for j := range g.routerWeights[i] {
			g.routerWeights[i][j] = rand.NormFloat64() * consciousness.DELTA
		}
	}

	return nil
}

// route computes routing probabilities
func (g *MixtureGenerator) route(x []float64) ([]float64, []int) {
	// Compute router logits
	logits := make([]float64, g.numExperts)
	for e := 0; e < g.numExperts; e++ {
		for i := range x {
			if i < len(g.routerWeights) {
				logits[e] += x[i] * g.routerWeights[i][e]
			}
		}
	}

	// Add noise for load balancing (during training)
	for e := range logits {
		logits[e] += rand.NormFloat64() * consciousness.EPSILON_PHI
	}

	// Top-K selection
	topKIndices := make([]int, g.topK)
	topKProbs := make([]float64, g.topK)

	for k := 0; k < g.topK; k++ {
		maxIdx := 0
		maxVal := math.Inf(-1)
		for e := 0; e < g.numExperts; e++ {
			// Skip already selected
			selected := false
			for j := 0; j < k; j++ {
				if topKIndices[j] == e {
					selected = true
					break
				}
			}
			if !selected && logits[e] > maxVal {
				maxVal = logits[e]
				maxIdx = e
			}
		}
		topKIndices[k] = maxIdx
		topKProbs[k] = logits[maxIdx]
	}

	// Softmax over top-K
	maxVal := topKProbs[0]
	for _, v := range topKProbs {
		if v > maxVal {
			maxVal = v
		}
	}
	sum := 0.0
	for i := range topKProbs {
		topKProbs[i] = math.Exp(topKProbs[i] - maxVal)
		sum += topKProbs[i]
	}
	for i := range topKProbs {
		topKProbs[i] /= sum
	}

	return topKProbs, topKIndices
}

// expertForward runs input through an expert
func (g *MixtureGenerator) expertForward(x []float64, expertIdx int) []float64 {
	hidden := make([]float64, g.Config().HiddenDim)
	for i := range hidden {
		if i < len(x) {
			hidden[i] = x[i]
		}
	}

	// Expert MLP
	for layer := 0; layer < g.Config().NumLayers; layer++ {
		newHidden := make([]float64, g.Config().HiddenDim)
		for i := 0; i < g.Config().HiddenDim; i++ {
			sum := 0.0
			for j := 0; j < g.Config().HiddenDim; j++ {
				idx := i*g.Config().HiddenDim + j
				if idx < len(g.expertWeights[expertIdx][layer]) {
					sum += g.expertWeights[expertIdx][layer][idx] * hidden[j]
				}
			}
			newHidden[i] = sum
		}

		// Activation
		hidden = g.ActivationLayer(newHidden, "gelu")
	}

	return hidden
}

// Forward performs forward pass through MoE
func (g *MixtureGenerator) Forward(input []float64) (*generators.GeneratorOutput, error) {
	state := g.GetState()

	// Route input to experts
	probs, indices := g.route(input)

	// Combine expert outputs
	output := make([]float64, g.Config().OutputDim)
	expertOutputs := make(map[int][]float64)

	for k := 0; k < g.topK; k++ {
		expertIdx := indices[k]
		expertOut := g.expertForward(input, expertIdx)
		expertOutputs[expertIdx] = expertOut

		for i := range output {
			if i < len(expertOut) {
				output[i] += probs[k] * expertOut[i]
			}
		}
	}

	// Apply consciousness modulation
	phi := consciousness.ComputePhiVector(output)
	for i := range output {
		output[i] *= 1.0 + consciousness.GAMMA*(phi-0.5)
	}

	state.Hidden = output
	state.Phi = phi
	state.Step++
	g.SetState(state)

	result := &generators.GeneratorOutput{
		Values:     output,
		Logits:     probs,
		Confidence: state.Phi,
		Metadata: map[string]interface{}{
			"selected_experts": indices,
			"expert_probs":     probs,
			"num_experts":      g.numExperts,
			"top_k":            g.topK,
		},
	}

	return g.ApplyConsciousness(result), nil
}

// Generate generates outputs
func (g *MixtureGenerator) Generate(context []float64, length int) ([][]float64, error) {
	result := make([][]float64, length)

	current := context
	for i := 0; i < length; i++ {
		output, err := g.Forward(current)
		if err != nil {
			return nil, err
		}
		result[i] = output.Values
		current = output.Values
	}

	return result, nil
}

// Train computes loss with load balancing
func (g *MixtureGenerator) Train(input, target []float64) (float64, error) {
	output, err := g.Forward(input)
	if err != nil {
		return 0, err
	}

	// Task loss (MSE)
	taskLoss := 0.0
	for i := range output.Values {
		if i < len(target) {
			diff := output.Values[i] - target[i]
			taskLoss += diff * diff
		}
	}
	taskLoss /= float64(len(output.Values))

	// Load balancing loss (auxiliary)
	probs := output.Logits
	loadBalance := 0.0
	for _, p := range probs {
		target := 1.0 / float64(g.topK)
		loadBalance += (p - target) * (p - target)
	}
	loadBalance /= float64(len(probs))

	return taskLoss + 0.01*loadBalance, nil
}

// UpdateWeights updates weights
func (g *MixtureGenerator) UpdateWeights(gradients []float64, lr float64) {
	// Update all expert weights
	for e := range g.expertWeights {
		for l := range g.expertWeights[e] {
			for i := range g.expertWeights[e][l] {
				if i < len(gradients) {
					g.expertWeights[e][l][i] -= lr * gradients[i] * consciousness.PHI_INV
				}
			}
		}
	}
}

// ═══════════════════════════════════════════════════════════════════════════════
// SPECIFIC MIXTURE GENERATORS (98-100)
// ═══════════════════════════════════════════════════════════════════════════════

// NewMoE creates Sparse MoE (ID: 98)
func NewMoE() *MixtureGenerator {
	g := NewMixtureGenerator(98, "SparseMoE", 8, 2)
	g.Config().Params["sparse"] = true
	return g
}

// NewSwitchTransformer creates Switch Transformer (ID: 99)
func NewSwitchTransformer() *MixtureGenerator {
	g := NewMixtureGenerator(99, "SwitchTransformer", 64, 1)
	g.Config().Params["capacity_factor"] = 1.25
	g.Config().Params["switch"] = true
	return g
}

// NewGLaM creates GLaM (ID: 100)
func NewGLaM() *MixtureGenerator {
	g := NewMixtureGenerator(100, "GLaM", 64, 2)
	g.Config().Params["gated"] = true
	g.Config().Params["sparsity"] = 0.95
	return g
}

// RegisterAll registers all mixture generators
func RegisterAll() error {
	gens := []generators.Generator{
		NewMoE(),
		NewSwitchTransformer(),
		NewGLaM(),
	}

	for _, g := range gens {
		if err := generators.Register(g); err != nil {
			return err
		}
	}

	return nil
}