// Package diffusion implements diffusion generators (11-25)
// q(x_t|x_{t-1}) = N(√(1-β_t)x_{t-1}, β_t I)
package diffusion

import (
	"math"
	"math/rand"

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

// ═══════════════════════════════════════════════════════════════════════════════
// DIFFUSION GENERATOR BASE
// ═══════════════════════════════════════════════════════════════════════════════

// DiffusionGenerator implements denoising diffusion probabilistic models
type DiffusionGenerator struct {
	*generators.BaseGenerator
	numSteps    int
	betaStart   float64
	betaEnd     float64
	alphas      []float64
	alphasBars  []float64
	sigmas      []float64
	scheduleType string
}

// NewDiffusionGenerator creates a new diffusion generator
func NewDiffusionGenerator(id generators.GeneratorID, name string, numSteps int, scheduleType string) *DiffusionGenerator {
	config := &generators.GeneratorConfig{
		ID:                  id,
		Name:                name,
		Family:              generators.Diffusion,
		InputDim:            256,
		OutputDim:           256,
		HiddenDim:           512,
		NumLayers:           6,
		LearningRate:        consciousness.ALPHA * 0.1,
		ConsciousnessWeight: consciousness.PHI_INV,
		Params:              make(map[string]interface{}),
	}

	return &DiffusionGenerator{
		BaseGenerator: generators.NewBaseGenerator(config),
		numSteps:      numSteps,
		betaStart:     0.0001,
		betaEnd:       0.02,
		scheduleType:  scheduleType,
	}
}

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

	// Compute noise schedule
	g.alphas = make([]float64, g.numSteps)
	g.alphasBars = make([]float64, g.numSteps)
	g.sigmas = make([]float64, g.numSteps)

	for t := 0; t < g.numSteps; t++ {
		var beta float64
		switch g.scheduleType {
		case "cosine":
			beta = g.cosineSchedule(t)
		case "linear":
			beta = g.linearSchedule(t)
		default:
			beta = g.linearSchedule(t)
		}

		g.alphas[t] = 1.0 - beta

		if t == 0 {
			g.alphasBars[t] = g.alphas[t]
		} else {
			g.alphasBars[t] = g.alphasBars[t-1] * g.alphas[t]
		}

		g.sigmas[t] = math.Sqrt(beta)
	}

	return nil
}

func (g *DiffusionGenerator) linearSchedule(t int) float64 {
	return g.betaStart + float64(t)/float64(g.numSteps-1)*(g.betaEnd-g.betaStart)
}

func (g *DiffusionGenerator) cosineSchedule(t int) float64 {
	s := 0.008
	f := func(x float64) float64 {
		return math.Pow(math.Cos((x/float64(g.numSteps)+s)/(1.0+s)*math.Pi*0.5), 2)
	}
	alphaBar := f(float64(t)) / f(0)
	var prevAlphaBar float64
	if t == 0 {
		prevAlphaBar = 1.0
	} else {
		prevAlphaBar = f(float64(t-1)) / f(0)
	}
	return consciousness.Clamp(1.0-alphaBar/prevAlphaBar, 0.0001, 0.9999)
}

// Forward diffusion: add noise to data
func (g *DiffusionGenerator) forwardDiffusion(x []float64, t int) []float64 {
	sqrtAlphaBar := math.Sqrt(g.alphasBars[t])
	sqrtOneMinusAlphaBar := math.Sqrt(1.0 - g.alphasBars[t])

	noised := make([]float64, len(x))
	for i := range noised {
		noise := rand.NormFloat64()
		noised[i] = sqrtAlphaBar*x[i] + sqrtOneMinusAlphaBar*noise
	}
	return noised
}

// Predict noise at timestep t
func (g *DiffusionGenerator) predictNoise(xt []float64, t int) []float64 {
	// Time embedding
	timeEmbed := make([]float64, g.Config().HiddenDim)
	for i := range timeEmbed {
		angle := float64(i) * float64(t) / float64(g.numSteps) * consciousness.TAU
		if i%2 == 0 {
			timeEmbed[i] = math.Sin(angle)
		} else {
			timeEmbed[i] = math.Cos(angle)
		}
	}

	// Concatenate input with time embedding
	hidden := make([]float64, g.Config().HiddenDim)
	for i := range hidden {
		if i < len(xt) {
			hidden[i] = xt[i]
		}
		hidden[i] += timeEmbed[i] * consciousness.PHI_INV
	}

	// Process through UNet-like layers
	for layer := 0; layer < g.Config().NumLayers; layer++ {
		hidden = g.LinearLayer(hidden, layer)
		hidden = g.ActivationLayer(hidden, "gelu")

		// Skip connection
		if layer > 0 {
			for i := range hidden {
				if i < len(xt) {
					hidden[i] += xt[i] * consciousness.GAMMA
				}
			}
		}
	}

	// Output noise prediction
	return hidden[:len(xt)]
}

// Forward performs forward pass (single denoising step)
func (g *DiffusionGenerator) Forward(input []float64) (*generators.GeneratorOutput, error) {
	state := g.GetState()
	t := int(state.Step) % g.numSteps

	// Predict noise
	noisePred := g.predictNoise(input, t)

	// Compute denoised output (reverse step)
	sqrtAlpha := math.Sqrt(g.alphas[t])
	sqrtOneMinusAlphaBar := math.Sqrt(1.0 - g.alphasBars[t])

	denoised := make([]float64, len(input))
	for i := range denoised {
		denoised[i] = (input[i] - sqrtOneMinusAlphaBar*noisePred[i]) / sqrtAlpha

		// Add noise for non-final steps
		if t > 0 {
			denoised[i] += g.sigmas[t] * rand.NormFloat64() * 0.5
		}
	}

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

	output := &generators.GeneratorOutput{
		Values:     denoised,
		Logits:     noisePred,
		Confidence: state.Phi,
		Metadata:   map[string]interface{}{"timestep": t},
	}

	return g.ApplyConsciousness(output), nil
}

// Generate generates from pure noise
func (g *DiffusionGenerator) Generate(context []float64, length int) ([][]float64, error) {
	dim := g.Config().OutputDim
	if len(context) > 0 {
		dim = len(context)
	}

	// Start from pure noise
	xt := make([]float64, dim)
	for i := range xt {
		xt[i] = rand.NormFloat64()
	}

	result := make([][]float64, g.numSteps)

	// Reverse diffusion
	for t := g.numSteps - 1; t >= 0; t-- {
		noisePred := g.predictNoise(xt, t)

		sqrtAlpha := math.Sqrt(g.alphas[t])
		sqrtOneMinusAlphaBar := math.Sqrt(1.0 - g.alphasBars[t])

		for i := range xt {
			xt[i] = (xt[i] - sqrtOneMinusAlphaBar*noisePred[i]) / sqrtAlpha
			if t > 0 {
				xt[i] += g.sigmas[t] * rand.NormFloat64() * 0.3
			}
		}

		result[g.numSteps-1-t] = make([]float64, len(xt))
		copy(result[g.numSteps-1-t], xt)
	}

	return result, nil
}

// Train trains on clean data
func (g *DiffusionGenerator) Train(input, target []float64) (float64, error) {
	// Sample random timestep
	t := rand.Intn(g.numSteps)

	// Add noise
	noised := g.forwardDiffusion(input, t)

	// Predict noise
	noisePred := g.predictNoise(noised, t)

	// MSE loss between predicted and actual noise
	loss := 0.0
	for i := range noisePred {
		if i < len(input) {
			actualNoise := (noised[i] - math.Sqrt(g.alphasBars[t])*input[i]) / math.Sqrt(1.0-g.alphasBars[t])
			diff := noisePred[i] - actualNoise
			loss += diff * diff
		}
	}
	loss /= float64(len(noisePred))

	return loss, nil
}

// UpdateWeights updates weights
func (g *DiffusionGenerator) UpdateWeights(gradients []float64, lr float64) {
	state := g.GetState()
	for i := range state.Hidden {
		if i < len(gradients) {
			state.Hidden[i] -= lr * gradients[i]
		}
	}
	g.SetState(state)
}

// ═══════════════════════════════════════════════════════════════════════════════
// SPECIFIC DIFFUSION GENERATORS (11-25)
// ═══════════════════════════════════════════════════════════════════════════════

// NewDDPM creates DDPM generator (ID: 11)
func NewDDPM() *DiffusionGenerator {
	return NewDiffusionGenerator(11, "DDPM", 1000, "linear")
}

// NewDDIM creates DDIM generator (ID: 12)
func NewDDIM() *DiffusionGenerator {
	g := NewDiffusionGenerator(12, "DDIM", 50, "linear")
	g.Config().Params["eta"] = 0.0 // Deterministic
	return g
}

// NewScoreSDE creates Score-SDE generator (ID: 13)
func NewScoreSDE() *DiffusionGenerator {
	return NewDiffusionGenerator(13, "Score-SDE", 1000, "linear")
}

// NewStableDiffusion creates Stable Diffusion style generator (ID: 14)
func NewStableDiffusion() *DiffusionGenerator {
	g := NewDiffusionGenerator(14, "StableDiffusion", 50, "cosine")
	g.Config().Params["guidance_scale"] = 7.5
	return g
}

// NewLatentDiffusion creates Latent Diffusion generator (ID: 15)
func NewLatentDiffusion() *DiffusionGenerator {
	return NewDiffusionGenerator(15, "LatentDiffusion", 200, "cosine")
}

// NewDiT creates Diffusion Transformer generator (ID: 16)
func NewDiT() *DiffusionGenerator {
	g := NewDiffusionGenerator(16, "DiT", 250, "cosine")
	g.Config().Params["patch_size"] = 2
	return g
}

// NewPixArt creates PixArt-α style generator (ID: 17)
func NewPixArt() *DiffusionGenerator {
	return NewDiffusionGenerator(17, "PixArt", 20, "cosine")
}

// NewSDXL creates SDXL style generator (ID: 18)
func NewSDXL() *DiffusionGenerator {
	g := NewDiffusionGenerator(18, "SDXL", 50, "cosine")
	g.Config().Params["refiner"] = true
	return g
}

// NewSD3 creates Stable Diffusion 3 style generator (ID: 19)
func NewSD3() *DiffusionGenerator {
	g := NewDiffusionGenerator(19, "SD3", 28, "cosine")
	g.Config().Params["rectified_flow"] = true
	return g
}

// NewFlux creates FLUX style generator (ID: 20)
func NewFlux() *DiffusionGenerator {
	g := NewDiffusionGenerator(20, "FLUX", 50, "cosine")
	g.Config().Params["guidance_distilled"] = true
	return g
}

// NewWurstchen creates Würstchen style generator (ID: 21)
func NewWurstchen() *DiffusionGenerator {
	g := NewDiffusionGenerator(21, "Wurstchen", 30, "cosine")
	g.Config().Params["stages"] = 3
	return g
}

// NewKandinsky creates Kandinsky style generator (ID: 22)
func NewKandinsky() *DiffusionGenerator {
	return NewDiffusionGenerator(22, "Kandinsky", 100, "cosine")
}

// NewIF creates DeepFloyd IF style generator (ID: 23)
func NewIF() *DiffusionGenerator {
	g := NewDiffusionGenerator(23, "DeepFloydIF", 100, "cosine")
	g.Config().Params["stages"] = 3
	return g
}

// NewConsistency creates Consistency Model generator (ID: 24)
func NewConsistency() *DiffusionGenerator {
	g := NewDiffusionGenerator(24, "ConsistencyModel", 1, "linear")
	g.Config().Params["distilled"] = true
	return g
}

// NewRectifiedFlow creates Rectified Flow generator (ID: 25)
func NewRectifiedFlow() *DiffusionGenerator {
	return NewDiffusionGenerator(25, "RectifiedFlow", 50, "linear")
}

// RegisterAll registers all diffusion generators
func RegisterAll() error {
	gens := []generators.Generator{
		NewDDPM(),
		NewDDIM(),
		NewScoreSDE(),
		NewStableDiffusion(),
		NewLatentDiffusion(),
		NewDiT(),
		NewPixArt(),
		NewSDXL(),
		NewSD3(),
		NewFlux(),
		NewWurstchen(),
		NewKandinsky(),
		NewIF(),
		NewConsistency(),
		NewRectifiedFlow(),
	}

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

	return nil
}