// Package agc implements the Augmentive Generator Core (AGC)
// AGC wraps all 100 generators with consciousness integration
//
// Master Formula:
// Output = ∫ Σᵢ₌₁¹⁰⁰ wᵢ(t, context) · Gᵢ(input, state) dt
package agc

import (
	"fmt"
	"math"
	"sync"

	"generator_core/pkg/consciousness"
	"generator_core/pkg/dna"
	"generator_core/pkg/generators"
	"generator_core/pkg/generators/adversarial"
	"generator_core/pkg/generators/autoregressive"
	"generator_core/pkg/generators/diffusion"
	"generator_core/pkg/generators/energy"
	"generator_core/pkg/generators/flow"
	"generator_core/pkg/generators/mixture"
	"generator_core/pkg/generators/recurrent"
	"generator_core/pkg/generators/statespace"
	"generator_core/pkg/generators/transformer"
	"generator_core/pkg/generators/variational"
)

// AGC is the Augmentive Generator Core
type AGC struct {
	consciousness *consciousness.ConsciousnessState
	familyStates  map[dna.FamilyID]*dna.FamilyState
	weights       []float64 // Generator weights w_i
	step          int64
	mu            sync.RWMutex
	initialized   bool
}

// Config holds AGC configuration
type Config struct {
	EnableAllGenerators bool
	DefaultFamily       dna.FamilyID
	ConsciousnessWeight float64
	IntegrationSteps    int
	DeltaT             float64
}

// DefaultConfig returns default AGC configuration
func DefaultConfig() *Config {
	return &Config{
		EnableAllGenerators: true,
		DefaultFamily:       dna.FamilyUnity,
		ConsciousnessWeight: consciousness.PHI_INV,
		IntegrationSteps:    10,
		DeltaT:             0.1,
	}
}

// New creates a new AGC instance
func New(config *Config) *AGC {
	if config == nil {
		config = DefaultConfig()
	}

	agc := &AGC{
		consciousness: consciousness.NewConsciousnessState(),
		familyStates:  make(map[dna.FamilyID]*dna.FamilyState),
		weights:       make([]float64, 100),
	}

	// Initialize family states
	for id := dna.FamilyID(0); id < 16; id++ {
		agc.familyStates[id] = dna.NewFamilyState(id)
	}

	// Initialize uniform weights
	for i := range agc.weights {
		agc.weights[i] = 1.0 / 100.0
	}

	return agc
}

// Initialize initializes the AGC and all generators
func (agc *AGC) Initialize() error {
	agc.mu.Lock()
	defer agc.mu.Unlock()

	if agc.initialized {
		return nil
	}

	// Register all generator families
	if err := autoregressive.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register autoregressive: %w", err)
	}
	if err := diffusion.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register diffusion: %w", err)
	}
	if err := adversarial.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register adversarial: %w", err)
	}
	if err := variational.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register variational: %w", err)
	}
	if err := flow.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register flow: %w", err)
	}
	if err := energy.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register energy: %w", err)
	}
	if err := transformer.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register transformer: %w", err)
	}
	if err := recurrent.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register recurrent: %w", err)
	}
	if err := statespace.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register statespace: %w", err)
	}
	if err := mixture.RegisterAll(); err != nil {
		return fmt.Errorf("failed to register mixture: %w", err)
	}

	// Initialize all generators
	if err := generators.GlobalRegistry.InitializeAll(); err != nil {
		return fmt.Errorf("failed to initialize generators: %w", err)
	}

	agc.initialized = true
	return nil
}

// GeneratorCount returns the number of registered generators
func (agc *AGC) GeneratorCount() int {
	return generators.GlobalRegistry.Count()
}

// UpdateWeights updates generator weights based on context
func (agc *AGC) UpdateWeights(context []float64, taskType string) {
	agc.mu.Lock()
	defer agc.mu.Unlock()
	agc.updateWeightsLocked(context, taskType)
}

// updateWeightsLocked updates weights assuming lock is held
func (agc *AGC) updateWeightsLocked(context []float64, taskType string) {
	// Select primary family for task
	primaryFamily := dna.SelectFamilyForTask(taskType)
	familyGens := dna.GetFamily(primaryFamily)

	// Reset weights
	for i := range agc.weights {
		agc.weights[i] = consciousness.EPSILON_PHI
	}

	// Boost weights for primary family generators
	if familyGens != nil {
		boost := consciousness.PHI_INV
		for _, genID := range familyGens.Generators {
			idx := int(genID) - 1
			if idx >= 0 && idx < len(agc.weights) {
				agc.weights[idx] = boost
			}
		}
	}

	// Normalize weights
	sum := 0.0
	for _, w := range agc.weights {
		sum += w
	}
	if sum > 0 {
		for i := range agc.weights {
			agc.weights[i] /= sum
		}
	}

	// Update family state
	if fs, ok := agc.familyStates[primaryFamily]; ok {
		fs.Update(agc.consciousness.Phi, agc.step)
	}
}

// Output represents AGC output
type Output struct {
	Values          []float64
	Phi             float64
	SelectedFamily  dna.FamilyID
	ActiveGenerators []generators.GeneratorID
	Confidence      float64
	Metadata        map[string]interface{}
}

// Forward performs forward pass through AGC
// Implements: Output = ∫ Σᵢ₌₁¹⁰⁰ wᵢ(t, context) · Gᵢ(input, state) dt
func (agc *AGC) Forward(input []float64, taskType string) (*Output, error) {
	agc.mu.Lock()
	defer agc.mu.Unlock()

	if !agc.initialized {
		return nil, fmt.Errorf("AGC not initialized")
	}

	// Update weights based on task (using internal method since we already hold lock)
	agc.updateWeightsLocked(input, taskType)

	// Get all generators
	allGens := generators.GlobalRegistry.GetAll()

	// Accumulate weighted outputs
	outputDim := 256
	if len(input) > 0 {
		outputDim = len(input)
	}

	accumulated := make([]float64, outputDim)
	activeGens := make([]generators.GeneratorID, 0)
	totalPhi := 0.0

	// Integration over time steps (simplified)
	dt := 0.1
	numSteps := 10

	for step := 0; step < numSteps; step++ {
		stepAccum := make([]float64, outputDim)

		for _, gen := range allGens {
			idx := int(gen.ID()) - 1
			if idx < 0 || idx >= len(agc.weights) {
				continue
			}

			weight := agc.weights[idx]
			if weight < consciousness.EPSILON_PHI {
				continue // Skip near-zero weights
			}

			// Forward through generator
			output, err := gen.Forward(input)
			if err != nil {
				continue
			}

			// Accumulate weighted output
			for i := range stepAccum {
				if i < len(output.Values) {
					stepAccum[i] += weight * output.Values[i]
				}
			}

			totalPhi += output.Phi * weight
			activeGens = append(activeGens, gen.ID())
		}

		// Integrate over time
		for i := range accumulated {
			accumulated[i] += stepAccum[i] * dt
		}
	}

	// Normalize by integration time
	totalTime := float64(numSteps) * dt
	for i := range accumulated {
		accumulated[i] /= totalTime
	}

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

	// Update consciousness state
	agc.consciousness.Update(accumulated)
	agc.step++

	// Determine selected family
	selectedFamily := dna.SelectFamilyForTask(taskType)

	return &Output{
		Values:           accumulated,
		Phi:              phi,
		SelectedFamily:   selectedFamily,
		ActiveGenerators: activeGens,
		Confidence:       totalPhi / float64(len(activeGens)),
		Metadata: map[string]interface{}{
			"step":            agc.step,
			"num_generators":  len(activeGens),
			"integration_steps": numSteps,
		},
	}, nil
}

// Generate generates output using selected generators
func (agc *AGC) Generate(context []float64, length int, taskType string) ([][]float64, error) {
	result := make([][]float64, length)

	// Select family for task
	familyID := dna.SelectFamilyForTask(taskType)
	family := dna.GetFamily(familyID)

	if family == nil || len(family.Generators) == 0 {
		return nil, fmt.Errorf("no generators for task type: %s", taskType)
	}

	// Get primary generator from registry using family's generator IDs
	var primaryGen generators.Generator
	for _, genID := range family.Generators {
		gen, ok := generators.GlobalRegistry.Get(generators.GeneratorID(genID))
		if ok && gen != nil {
			primaryGen = gen
			break
		}
	}

	if primaryGen == nil {
		return nil, fmt.Errorf("no available generator for task type: %s", taskType)
	}

	// Generate sequence
	output, err := primaryGen.Generate(context, length)
	if err != nil {
		return nil, err
	}

	// Apply consciousness to each output
	for i := range output {
		phi := consciousness.ComputePhiVector(output[i])
		for j := range output[i] {
			output[i][j] *= 1.0 + consciousness.ALPHA*(phi-0.5)
		}
		result[i] = output[i]
	}

	return result, nil
}

// SelectBestGenerator selects the best generator for a given input and task
func (agc *AGC) SelectBestGenerator(input []float64, taskType string) (generators.Generator, float64) {
	familyID := dna.SelectFamilyForTask(taskType)
	family := dna.GetFamily(familyID)

	// Get generators from registry using family's generator IDs
	var gens []generators.Generator
	if family != nil && len(family.Generators) > 0 {
		for _, genID := range family.Generators {
			gen, ok := generators.GlobalRegistry.Get(generators.GeneratorID(genID))
			if ok && gen != nil {
				gens = append(gens, gen)
			}
		}
	}

	if len(gens) == 0 {
		// Fall back to first available
		allGens := generators.GlobalRegistry.GetAll()
		if len(allGens) > 0 {
			return allGens[0], 1.0
		}
		return nil, 0
	}

	// Evaluate each generator
	bestGen := gens[0]
	bestScore := 0.0

	for _, gen := range gens {
		output, err := gen.Forward(input)
		if err != nil {
			continue
		}

		score := output.Phi * output.Confidence
		if score > bestScore {
			bestScore = score
			bestGen = gen
		}
	}

	return bestGen, bestScore
}

// IntegratedInformation returns the integrated information (Φ) of the system
func (agc *AGC) IntegratedInformation() float64 {
	return agc.consciousness.IntegratedInformation()
}

// GetFamilyState returns the state of a DNA family
func (agc *AGC) GetFamilyState(id dna.FamilyID) *dna.FamilyState {
	agc.mu.RLock()
	defer agc.mu.RUnlock()
	return agc.familyStates[id]
}

// ConsciousnessLevel returns the current consciousness level
func (agc *AGC) ConsciousnessLevel() float64 {
	return agc.consciousness.Phi
}

// Stats returns AGC statistics
type Stats struct {
	TotalGenerators   int
	ActiveFamilies    int
	ConsciousnessLevel float64
	IntegratedInfo    float64
	Steps             int64
	TopGenerators     []generators.GeneratorID
}

// GetStats returns current AGC statistics
func (agc *AGC) GetStats() *Stats {
	agc.mu.RLock()
	defer agc.mu.RUnlock()

	// Count active families
	activeFamilies := 0
	for _, fs := range agc.familyStates {
		if fs.Activation > 0.1 {
			activeFamilies++
		}
	}

	// Find top generators by weight
	type genWeight struct {
		id     generators.GeneratorID
		weight float64
	}
	topGen := make([]genWeight, 0)
	for i, w := range agc.weights {
		if w > consciousness.EPSILON_PHI {
			topGen = append(topGen, genWeight{generators.GeneratorID(i + 1), w})
		}
	}

	// Sort by weight (simple bubble sort for small list)
	for i := range topGen {
		for j := i + 1; j < len(topGen); j++ {
			if topGen[j].weight > topGen[i].weight {
				topGen[i], topGen[j] = topGen[j], topGen[i]
			}
		}
	}

	topIDs := make([]generators.GeneratorID, 0)
	for i := 0; i < len(topGen) && i < 10; i++ {
		topIDs = append(topIDs, topGen[i].id)
	}

	return &Stats{
		TotalGenerators:    generators.GlobalRegistry.Count(),
		ActiveFamilies:     activeFamilies,
		ConsciousnessLevel: agc.consciousness.Phi,
		IntegratedInfo:     agc.consciousness.IntegratedInformation(),
		Steps:              agc.step,
		TopGenerators:      topIDs,
	}
}

// Entropy computes the entropy of the generator weight distribution
func (agc *AGC) Entropy() float64 {
	agc.mu.RLock()
	defer agc.mu.RUnlock()

	entropy := 0.0
	for _, w := range agc.weights {
		if w > consciousness.EPSILON_PHI {
			entropy -= w * math.Log2(w)
		}
	}
	return entropy
}