// Package adversarial implements adversarial generators (26-40)
// min_G max_D V(D,G) = E[log D(x)] + E[log(1-D(G(z)))]
package adversarial

import (
	"math"
	"math/rand"

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

// ═══════════════════════════════════════════════════════════════════════════════
// ADVERSARIAL GENERATOR BASE
// ═══════════════════════════════════════════════════════════════════════════════

// AdversarialGenerator implements GAN-style generators
type AdversarialGenerator struct {
	*generators.BaseGenerator
	latentDim       int
	generatorLayers int
	discriminatorLayers int
	genWeights      [][]float64
	discWeights     [][]float64
}

// NewAdversarialGenerator creates a new adversarial generator
func NewAdversarialGenerator(id generators.GeneratorID, name string, latentDim int) *AdversarialGenerator {
	config := &generators.GeneratorConfig{
		ID:                  id,
		Name:                name,
		Family:              generators.Adversarial,
		InputDim:            latentDim,
		OutputDim:           256,
		HiddenDim:           512,
		NumLayers:           4,
		LearningRate:        consciousness.ALPHA * 0.5,
		ConsciousnessWeight: consciousness.PHI_INV,
		Params:              make(map[string]interface{}),
	}

	return &AdversarialGenerator{
		BaseGenerator:       generators.NewBaseGenerator(config),
		latentDim:           latentDim,
		generatorLayers:     4,
		discriminatorLayers: 3,
	}
}

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

	// Initialize generator weights
	g.genWeights = make([][]float64, g.generatorLayers)
	for i := range g.genWeights {
		size := g.Config().HiddenDim * g.Config().HiddenDim
		g.genWeights[i] = make([]float64, size)
		for j := range g.genWeights[i] {
			g.genWeights[i][j] = rand.NormFloat64() * consciousness.PHI_INV / math.Sqrt(float64(g.Config().HiddenDim))
		}
	}

	// Initialize discriminator weights
	g.discWeights = make([][]float64, g.discriminatorLayers)
	for i := range g.discWeights {
		size := g.Config().HiddenDim * g.Config().HiddenDim
		g.discWeights[i] = make([]float64, size)
		for j := range g.discWeights[i] {
			g.discWeights[i][j] = rand.NormFloat64() * consciousness.PHI_INV / math.Sqrt(float64(g.Config().HiddenDim))
		}
	}

	return nil
}

// generateFromLatent generates output from latent vector
func (g *AdversarialGenerator) generateFromLatent(z []float64) []float64 {
	hidden := make([]float64, g.Config().HiddenDim)
	for i := range hidden {
		if i < len(z) {
			hidden[i] = z[i]
		}
	}

	// Generator forward pass
	for layer := 0; layer < g.generatorLayers; 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.genWeights[layer]) {
					sum += g.genWeights[layer][idx] * hidden[j]
				}
			}
			newHidden[i] = sum
		}
		hidden = g.ActivationLayer(newHidden, "relu")

		// Batch norm approximation
		mean := 0.0
		for _, v := range hidden {
			mean += v
		}
		mean /= float64(len(hidden))
		for i := range hidden {
			hidden[i] -= mean
		}
	}

	// Final activation (tanh for [-1, 1] output)
	return g.ActivationLayer(hidden, "tanh")
}

// discriminate computes discriminator output
func (g *AdversarialGenerator) discriminate(x []float64) float64 {
	hidden := make([]float64, g.Config().HiddenDim)
	for i := range hidden {
		if i < len(x) {
			hidden[i] = x[i]
		}
	}

	// Discriminator forward pass
	for layer := 0; layer < g.discriminatorLayers; 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.discWeights[layer]) {
					sum += g.discWeights[layer][idx] * hidden[j]
				}
			}
			newHidden[i] = sum
		}
		hidden = g.ActivationLayer(newHidden, "relu")

		// Leaky ReLU (negative slope 0.2)
		for i := range hidden {
			if hidden[i] < 0 {
				hidden[i] *= 0.2
			}
		}
	}

	// Final sigmoid
	sum := 0.0
	for _, v := range hidden {
		sum += v
	}
	return consciousness.Sigmoid(sum / float64(len(hidden)))
}

// Forward performs generation from input latent
func (g *AdversarialGenerator) Forward(input []float64) (*generators.GeneratorOutput, error) {
	state := g.GetState()

	// Generate from latent
	generated := g.generateFromLatent(input)

	// Get discriminator score
	discScore := g.discriminate(generated)

	state.Hidden = generated
	state.Phi = consciousness.ComputePhiVector(generated)
	state.Step++
	g.SetState(state)

	output := &generators.GeneratorOutput{
		Values:     generated,
		Logits:     []float64{discScore},
		Confidence: discScore,
		Metadata:   map[string]interface{}{"disc_score": discScore},
	}

	return g.ApplyConsciousness(output), nil
}

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

	for i := 0; i < length; i++ {
		// Sample latent
		z := make([]float64, g.latentDim)
		for j := range z {
			z[j] = rand.NormFloat64()
		}

		// Condition on context if provided
		for j := range context {
			if j < len(z) {
				z[j] += context[j] * consciousness.ALPHA
			}
		}

		output, err := g.Forward(z)
		if err != nil {
			return nil, err
		}
		result[i] = output.Values
	}

	return result, nil
}

// Train performs adversarial training step
func (g *AdversarialGenerator) Train(input, target []float64) (float64, error) {
	// Sample latent
	z := make([]float64, g.latentDim)
	for i := range z {
		z[i] = rand.NormFloat64()
	}

	// Generate fake sample
	fake := g.generateFromLatent(z)

	// Discriminator loss
	realScore := g.discriminate(target)
	fakeScore := g.discriminate(fake)

	dLoss := -math.Log(realScore+consciousness.EPSILON_PHI) - math.Log(1.0-fakeScore+consciousness.EPSILON_PHI)

	// Generator loss
	gLoss := -math.Log(fakeScore + consciousness.EPSILON_PHI)

	return (dLoss + gLoss) / 2.0, nil
}

// UpdateWeights updates weights
func (g *AdversarialGenerator) UpdateWeights(gradients []float64, lr float64) {
	// Update generator weights
	for layer := range g.genWeights {
		for i := range g.genWeights[layer] {
			if i < len(gradients) {
				g.genWeights[layer][i] -= lr * gradients[i] * consciousness.PHI_INV
			}
		}
	}
}

// ═══════════════════════════════════════════════════════════════════════════════
// SPECIFIC ADVERSARIAL GENERATORS (26-40)
// ═══════════════════════════════════════════════════════════════════════════════

// NewVanillaGAN creates Vanilla GAN (ID: 26)
func NewVanillaGAN() *AdversarialGenerator {
	return NewAdversarialGenerator(26, "VanillaGAN", 100)
}

// NewDCGAN creates DCGAN (ID: 27)
func NewDCGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(27, "DCGAN", 100)
	g.Config().Params["conv"] = true
	return g
}

// NewWGAN creates Wasserstein GAN (ID: 28)
func NewWGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(28, "WGAN", 100)
	g.Config().Params["wasserstein"] = true
	return g
}

// NewWGANGP creates WGAN with Gradient Penalty (ID: 29)
func NewWGANGP() *AdversarialGenerator {
	g := NewAdversarialGenerator(29, "WGAN-GP", 100)
	g.Config().Params["gradient_penalty"] = 10.0
	return g
}

// NewStyleGAN creates StyleGAN (ID: 30)
func NewStyleGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(30, "StyleGAN", 512)
	g.Config().Params["style_dim"] = 512
	g.Config().Params["mapping_layers"] = 8
	return g
}

// NewStyleGAN2 creates StyleGAN2 (ID: 31)
func NewStyleGAN2() *AdversarialGenerator {
	g := NewAdversarialGenerator(31, "StyleGAN2", 512)
	g.Config().Params["style_dim"] = 512
	g.Config().Params["path_regularization"] = true
	return g
}

// NewStyleGAN3 creates StyleGAN3 (ID: 32)
func NewStyleGAN3() *AdversarialGenerator {
	g := NewAdversarialGenerator(32, "StyleGAN3", 512)
	g.Config().Params["alias_free"] = true
	return g
}

// NewBigGAN creates BigGAN (ID: 33)
func NewBigGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(33, "BigGAN", 128)
	g.Config().Params["class_conditional"] = true
	return g
}

// NewProGAN creates Progressive GAN (ID: 34)
func NewProGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(34, "ProGAN", 512)
	g.Config().Params["progressive"] = true
	return g
}

// NewCycleGAN creates CycleGAN (ID: 35)
func NewCycleGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(35, "CycleGAN", 256)
	g.Config().Params["cycle_consistency"] = true
	return g
}

// NewPix2Pix creates Pix2Pix (ID: 36)
func NewPix2Pix() *AdversarialGenerator {
	g := NewAdversarialGenerator(36, "Pix2Pix", 256)
	g.Config().Params["conditional"] = true
	return g
}

// NewGauGAN creates GauGAN/SPADE (ID: 37)
func NewGauGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(37, "GauGAN", 256)
	g.Config().Params["spade"] = true
	return g
}

// NewStarGAN creates StarGAN (ID: 38)
func NewStarGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(38, "StarGAN", 256)
	g.Config().Params["multi_domain"] = true
	return g
}

// NewInfoGAN creates InfoGAN (ID: 39)
func NewInfoGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(39, "InfoGAN", 128)
	g.Config().Params["info_reg"] = true
	return g
}

// NewACGAN creates AC-GAN (ID: 40)
func NewACGAN() *AdversarialGenerator {
	g := NewAdversarialGenerator(40, "AC-GAN", 100)
	g.Config().Params["auxiliary_classifier"] = true
	return g
}

// RegisterAll registers all adversarial generators
func RegisterAll() error {
	gens := []generators.Generator{
		NewVanillaGAN(),
		NewDCGAN(),
		NewWGAN(),
		NewWGANGP(),
		NewStyleGAN(),
		NewStyleGAN2(),
		NewStyleGAN3(),
		NewBigGAN(),
		NewProGAN(),
		NewCycleGAN(),
		NewPix2Pix(),
		NewGauGAN(),
		NewStarGAN(),
		NewInfoGAN(),
		NewACGAN(),
	}

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

	return nil
}