// Package flow implements normalizing flow generators (56-65)
// p(x) = p(z) |det ∂f⁻¹/∂x| where z = f⁻¹(x)
package flow

import (
	"math"
	"math/rand"

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

// ═══════════════════════════════════════════════════════════════════════════════
// FLOW GENERATOR BASE
// ═══════════════════════════════════════════════════════════════════════════════

// FlowGenerator implements normalizing flow models
type FlowGenerator struct {
	*generators.BaseGenerator
	numFlows    int
	flowType    string
	flowParams  [][]float64
	logDetJac   float64 // Log determinant of Jacobian
}

// NewFlowGenerator creates a new flow generator
func NewFlowGenerator(id generators.GeneratorID, name string, numFlows int, flowType string) *FlowGenerator {
	config := &generators.GeneratorConfig{
		ID:                  id,
		Name:                name,
		Family:              generators.Flow,
		InputDim:            256,
		OutputDim:           256,
		HiddenDim:           256,
		NumLayers:           numFlows,
		LearningRate:        consciousness.ALPHA,
		ConsciousnessWeight: consciousness.PHI_INV,
		Params:              make(map[string]interface{}),
	}

	return &FlowGenerator{
		BaseGenerator: generators.NewBaseGenerator(config),
		numFlows:      numFlows,
		flowType:      flowType,
	}
}

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

	// Initialize flow parameters
	dim := g.Config().HiddenDim
	g.flowParams = make([][]float64, g.numFlows)
	for i := range g.flowParams {
		// Each flow has scale, shift, and coupling parameters
		g.flowParams[i] = make([]float64, dim*3)
		for j := range g.flowParams[i] {
			g.flowParams[i][j] = rand.NormFloat64() * consciousness.DELTA
		}
	}

	return nil
}

// affineTransform applies affine coupling transform
func (g *FlowGenerator) affineTransform(x []float64, flowIdx int, inverse bool) ([]float64, float64) {
	dim := len(x)
	halfDim := dim / 2

	// Split input
	x1 := x[:halfDim]
	x2 := x[halfDim:]

	// Get scale and shift from neural network
	scale := make([]float64, halfDim)
	shift := make([]float64, halfDim)

	for i := 0; i < halfDim; i++ {
		if flowIdx < len(g.flowParams) && i < len(g.flowParams[flowIdx]) {
			scale[i] = consciousness.Tanh(g.flowParams[flowIdx][i])
			shift[i] = g.flowParams[flowIdx][i+halfDim]
		}
	}

	// Apply transform
	y := make([]float64, dim)
	logDetJac := 0.0

	copy(y[:halfDim], x1)

	if !inverse {
		// Forward: y2 = x2 * exp(s(x1)) + t(x1)
		for i := 0; i < halfDim; i++ {
			y[halfDim+i] = x2[i]*math.Exp(scale[i]) + shift[i]
			logDetJac += scale[i]
		}
	} else {
		// Inverse: x2 = (y2 - t(x1)) * exp(-s(x1))
		for i := 0; i < halfDim; i++ {
			y[halfDim+i] = (x2[i] - shift[i]) * math.Exp(-scale[i])
			logDetJac -= scale[i]
		}
	}

	return y, logDetJac
}

// permute permutes dimensions
func (g *FlowGenerator) permute(x []float64, flowIdx int) []float64 {
	y := make([]float64, len(x))
	for i := range x {
		// Simple alternating permutation
		newIdx := (i + flowIdx) % len(x)
		y[newIdx] = x[i]
	}
	return y
}

// forwardFlow applies full forward flow (z -> x)
func (g *FlowGenerator) forwardFlow(z []float64) ([]float64, float64) {
	x := make([]float64, len(z))
	copy(x, z)
	totalLogDet := 0.0

	for i := 0; i < g.numFlows; i++ {
		var logDet float64
		x, logDet = g.affineTransform(x, i, false)
		totalLogDet += logDet
		x = g.permute(x, i)
	}

	return x, totalLogDet
}

// inverseFlow applies inverse flow (x -> z)
func (g *FlowGenerator) inverseFlow(x []float64) ([]float64, float64) {
	z := make([]float64, len(x))
	copy(z, x)
	totalLogDet := 0.0

	for i := g.numFlows - 1; i >= 0; i-- {
		// Inverse permute
		zTemp := make([]float64, len(z))
		for j := range z {
			origIdx := (j - i + len(z)) % len(z)
			zTemp[origIdx] = z[j]
		}
		z = zTemp

		var logDet float64
		z, logDet = g.affineTransform(z, i, true)
		totalLogDet += logDet
	}

	return z, totalLogDet
}

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

	// Inverse flow: map input to latent
	z, logDetJac := g.inverseFlow(input)
	g.logDetJac = logDetJac

	// Base distribution log probability (standard normal)
	logPz := 0.0
	for _, zi := range z {
		logPz += -0.5 * (zi*zi + math.Log(2*math.Pi))
	}

	// Log probability under flow
	logPx := logPz + logDetJac

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

	output := &generators.GeneratorOutput{
		Values:     z,
		Logits:     []float64{logPx},
		Confidence: consciousness.Sigmoid(logPx / 100.0),
		Metadata: map[string]interface{}{
			"log_det_jacobian": logDetJac,
			"log_prob":         logPx,
		},
	}

	return g.ApplyConsciousness(output), nil
}

// Generate generates samples from base distribution
func (g *FlowGenerator) Generate(context []float64, length int) ([][]float64, error) {
	dim := g.Config().OutputDim
	result := make([][]float64, length)

	for i := 0; i < length; i++ {
		// Sample from base distribution (standard normal)
		z := make([]float64, dim)
		for j := range z {
			z[j] = rand.NormFloat64()
		}

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

		// Transform through flow
		x, _ := g.forwardFlow(z)
		result[i] = x
	}

	return result, nil
}

// Train computes negative log likelihood
func (g *FlowGenerator) Train(input, target []float64) (float64, error) {
	// Inverse flow
	z, logDetJac := g.inverseFlow(target)

	// Base distribution log probability
	logPz := 0.0
	for _, zi := range z {
		logPz += -0.5 * (zi*zi + math.Log(2*math.Pi))
	}

	// Negative log likelihood
	nll := -(logPz + logDetJac)

	return nll / float64(len(target)), nil
}

// UpdateWeights updates flow parameters
func (g *FlowGenerator) UpdateWeights(gradients []float64, lr float64) {
	for i := range g.flowParams {
		for j := range g.flowParams[i] {
			if j < len(gradients) {
				g.flowParams[i][j] -= lr * gradients[j] * consciousness.PHI_INV
			}
		}
	}
}

// ═══════════════════════════════════════════════════════════════════════════════
// SPECIFIC FLOW GENERATORS (56-65)
// ═══════════════════════════════════════════════════════════════════════════════

// NewRealNVP creates RealNVP flow (ID: 56)
func NewRealNVP() *FlowGenerator {
	return NewFlowGenerator(56, "RealNVP", 8, "affine")
}

// NewGlow creates Glow flow (ID: 57)
func NewGlow() *FlowGenerator {
	g := NewFlowGenerator(57, "Glow", 32, "affine")
	g.Config().Params["lu_decomposition"] = true
	return g
}

// NewNICE creates NICE flow (ID: 58)
func NewNICE() *FlowGenerator {
	g := NewFlowGenerator(58, "NICE", 4, "additive")
	g.Config().Params["additive_only"] = true
	return g
}

// NewIAF creates Inverse Autoregressive Flow (ID: 59)
func NewIAF() *FlowGenerator {
	g := NewFlowGenerator(59, "IAF", 4, "autoregressive")
	g.Config().Params["autoregressive"] = true
	return g
}

// NewMAF creates Masked Autoregressive Flow (ID: 60)
func NewMAF() *FlowGenerator {
	g := NewFlowGenerator(60, "MAF", 4, "masked")
	g.Config().Params["masked"] = true
	return g
}

// NewNSF creates Neural Spline Flow (ID: 61)
func NewNSF() *FlowGenerator {
	g := NewFlowGenerator(61, "NeuralSplineFlow", 8, "spline")
	g.Config().Params["spline_bins"] = 8
	return g
}

// NewCNF creates Continuous Normalizing Flow (ID: 62)
func NewCNF() *FlowGenerator {
	g := NewFlowGenerator(62, "ContinuousNF", 1, "ode")
	g.Config().Params["ode_solver"] = "dopri5"
	return g
}

// NewFreeform creates Freeform Flow (ID: 63)
func NewFreeform() *FlowGenerator {
	g := NewFlowGenerator(63, "FreeformFlow", 4, "freeform")
	g.Config().Params["residual"] = true
	return g
}

// NewFlowPlusPlus creates Flow++ (ID: 64)
func NewFlowPlusPlus() *FlowGenerator {
	g := NewFlowGenerator(64, "Flow++", 16, "affine")
	g.Config().Params["variational_dequantization"] = true
	g.Config().Params["mixture_cdf"] = true
	return g
}

// NewSurvAE creates Survival Analysis Flow (ID: 65)
func NewSurvAE() *FlowGenerator {
	g := NewFlowGenerator(65, "SurvAE", 8, "affine")
	g.Config().Params["survival"] = true
	return g
}

// RegisterAll registers all flow generators
func RegisterAll() error {
	gens := []generators.Generator{
		NewRealNVP(),
		NewGlow(),
		NewNICE(),
		NewIAF(),
		NewMAF(),
		NewNSF(),
		NewCNF(),
		NewFreeform(),
		NewFlowPlusPlus(),
		NewSurvAE(),
	}

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

	return nil
}