package sdsengine

import (
    "context"
    "fmt"
    "log"
    "math"
    "sync"
    "time"

    "github.com/yourorg/3drenderer"  // 3D渲染器接口
    "github.com/yourorg/diffusion"   // 扩散模型接口
    "github.com/yourorg/geometry"    // 几何处理库
)

// SDSConfig 定义SDS优化器的配置参数
type SDSConfig struct {
    // 学习率，控制3D表示更新的步长
    LearningRate float64 `json:"learning_rate" default:"0.01"`
    
    // 噪声调度参数，控制扩散过程的时间步分布
    NoiseSchedule string `json:"noise_schedule" default:"cosine"`
    
    // 视角采样数量，每次迭代渲染的视角数
    ViewCount int `json:"view_count" default:"4"`
    
    // 光照控制参数，是否启用动态光照调整
    EnableLighting bool `json:"enable_lighting" default:"true"`
    
    // 正则化权重，平衡SDS损失与几何约束
    RegularizationWeight float64 `json:"reg_weight" default:"0.1"`
}

// SceneRepresentation 表示3D场景的抽象接口
type SceneRepresentation interface {
    // Render 从指定视角和光照条件渲染2D图像
    Render(camera *geometry.Camera, light *geometry.Lighting) (*geometry.Image, error)
    
    // Update 根据梯度更新3D表示
    Update(gradient []float64) error
    
    // GetParameters 获取当前3D表示的可学习参数
    GetParameters() []float64
    
    // Save 保存当前场景到文件
    Save(path string) error
}

// SDSOptimizer 核心SDS优化器结构体
type SDSOptimizer struct {
    config      SDSConfig
    scene       SceneRepresentation
    diffusionModel *diffusion.Model
    renderer    *3drenderer.DifferentiableRenderer
    
    // 优化状态
    iteration    int
    bestLoss     float64
    lossHistory  []float64
    
    // 并发控制
    mu           sync.RWMutex
    ctx          context.Context
    cancel       context.CancelFunc
}

// NewSDSOptimizer 创建SDS优化器实例
func NewSDSOptimizer(config SDSConfig, scene SceneRepresentation, diffusionModel *diffusion.Model) *SDSOptimizer {
    ctx, cancel := context.WithCancel(context.Background())
    
    return &SDSOptimizer{
        config:         config,
        scene:          scene,
        diffusionModel: diffusionModel,
        renderer:       3drenderer.NewDifferentiableRenderer(),
        iteration:      0,
        bestLoss:       math.MaxFloat64,
        lossHistory:    make([]float64, 0),
        ctx:            ctx,
        cancel:         cancel,
    }
}

// OptimizeStep 执行单步SDS优化
func (opt *SDSOptimizer) OptimizeStep(camera *geometry.Camera, light *geometry.Lighting, textCondition string) (float64, error) {
    opt.mu.Lock()
    defer opt.mu.Unlock()
    
    // 1. 从当前3D场景渲染2D图像
    renderedImage, err := opt.scene.Render(camera, light)
    if err != nil {
        return 0, fmt.Errorf("渲染失败: %w", err)
    }
    
    // 2. 对渲染图像添加噪声
    // 随机采样时间步t，控制噪声强度
    t := opt.sampleTimestep()
    noisedImage := addNoise(renderedImage, t)
    
    // 3. 使用扩散模型预测噪声
    predictedNoise, err := opt.diffusionModel.PredictNoise(noisedImage, t, textCondition)
    if err != nil {
        return 0, fmt.Errorf("噪声预测失败: %w", err)
    }
    
    // 4. 计算SDS损失梯度
    // 公式：∇_θ L_SDS = w(t) * (ε_ϕ - ε) * (∂x/∂θ)
    weight := computeWeight(t, opt.config.NoiseSchedule)
    noiseGradient := subtractNoise(predictedNoise, noisedImage.Noise)
    
    // 通过可微分渲染器反向传播梯度到3D表示
    renderGradient := opt.renderer.Backward(renderedImage, camera, light)
    sdsGradient := multiplyGradients(weight, noiseGradient, renderGradient)
    
    // 5. 应用正则化
    if opt.config.RegularizationWeight > 0 {
        regGradient := opt.computeRegularization()
        sdsGradient = addGradients(sdsGradient, regGradient)
    }
    
    // 6. 更新3D表示
    err = opt.scene.Update(sdsGradient)
    if err != nil {
        return 0, fmt.Errorf("场景更新失败: %w", err)
    }
    
    // 7. 计算当前损失值
    loss := computeLoss(predictedNoise, noisedImage.Noise)
    opt.lossHistory = append(opt.lossHistory, loss)
    if loss < opt.bestLoss {
        opt.bestLoss = loss
    }
    
    opt.iteration++
    
    return loss, nil
}

// sampleTimestep 根据噪声调度采样时间步
func (opt *SDSOptimizer) sampleTimestep() float64 {
    // 使用余弦调度，在早期和晚期时间步均匀采样
    // 这样能同时优化低频和高频细节
    u := randFloat64()
    switch opt.config.NoiseSchedule {
    case "cosine":
        return math.Acos(u) / (math.Pi / 2)
    case "linear":
        return u
    default:
        return u
    }
}

// computeWeight 计算时间步对应的权重
func computeWeight(t float64, schedule string) float64 {
    // 权重函数设计：早期时间步（小t）权重低，晚期时间步（大t）权重高
    // 这使得优化过程先关注低频结构，后优化高频细节
    switch schedule {
    case "cosine":
        return 1.0 - math.Cos(t*math.Pi/2)
    case "linear":
        return t
    default:
        return t
    }
}

// computeRegularization 计算几何正则化梯度
func (opt *SDSOptimizer) computeRegularization() []float64 {
    // 实现视角一致性正则化
    // 惩罚不同视角间颜色差异，确保多视角一致性
    params := opt.scene.GetParameters()
    regularization := make([]float64, len(params))
    
    // 简化的拉普拉斯平滑正则化
    for i := 1; i < len(params)-1; i++ {
        // 二阶差分约束，鼓励参数平滑变化
        regularization[i] = opt.config.RegularizationWeight * (2*params[i] - params[i-1] - params[i+1])
    }
    
    return regularization
}

// InteractiveOptimization 交互式优化循环，支持动态参数调整
func (opt *SDSOptimizer) InteractiveOptimization(
    ctx context.Context,
    cameraStream <-chan *geometry.Camera,
    lightStream <-chan *geometry.Lighting,
    textCondition string,
) error {
    log.Printf("开始交互式SDS优化，初始配置: %+v", opt.config)
    
    ticker := time.NewTicker(100 * time.Millisecond) // 100ms更新一次
    defer ticker.Stop()
    
    for {
        select {
        case <-ctx.Done():
            log.Println("优化被用户中断")
            return ctx.Err()
            
        case <-ticker.C:
            // 从通道获取最新的相机和光照参数
            var camera *geometry.Camera
            var light *geometry.Lighting
            
            select {
            case camera = <-cameraStream:
                // 更新相机视角
            default:
                // 使用默认视角
                camera = geometry.NewDefaultCamera()
            }
            
            select {
            case light = <-lightStream:
                // 更新光照条件
            default:
                light = geometry.NewDefaultLighting()
            }
            
            // 执行一步优化
            loss, err := opt.OptimizeStep(camera, light, textCondition)
            if err != nil {
                log.Printf("优化步骤 %d 失败: %v", opt.iteration, err)
                continue
            }
            
            // 每10步输出一次优化状态
            if opt.iteration%10 == 0 {
                log.Printf("迭代 %d, 损失: %.4f, 最佳损失: %.4f", 
                    opt.iteration, loss, opt.bestLoss)
            }
            
            // 检查收敛条件
            if len(opt.lossHistory) > 100 {
                avgLoss := average(opt.lossHistory[len(opt.lossHistory)-100:])
                if avgLoss < 0.001 {
                    log.Println("优化收敛，停止迭代")
                    return nil
                }
            }
        }
    }
}

// SaveCheckpoint 保存当前优化状态
func (opt *SDSOptimizer) SaveCheckpoint(path string) error {
    opt.mu.RLock()
    defer opt.mu.RUnlock()
    
    // 保存3D场景
    err := opt.scene.Save(path + "/scene.glb")
    if err != nil {
        return fmt.Errorf("保存场景失败: %w", err)
    }
    
    // 保存优化状态
    log.Printf("检查点已保存到 %s, 迭代次数: %d", path, opt.iteration)
    return nil
}

// 辅助函数
func addNoise(img *geometry.Image, t float64) *geometry.Image {
    // 根据时间步t添加高斯噪声
    noise := make([]float64, len(img.Data))
    sigma := math.Sqrt(t) // 噪声标准差随时间步增加
    for i := range noise {
        noise[i] = randNormal(0, sigma)
        img.Data[i] += noise[i]
    }
    img.Noise = noise
    return img
}

func subtractNoise(predicted, actual []float64) []float64 {
    result := make([]float64, len(predicted))
    for i := range predicted {
        result[i] = predicted[i] - actual[i]
    }
    return result
}

func multiplyGradients(weight float64, noiseGrad, renderGrad []float64) []float64 {
    result := make([]float64, len(noiseGrad))
    for i := range noiseGrad {
        result[i] = weight * noiseGrad[i] * renderGrad[i]
    }
    return result
}

func addGradients(a, b []float64) []float64 {
    result := make([]float64, len(a))
    for i := range a {
        result[i] = a[i] + b[i]
    }
    return result
}

func computeLoss(predicted, actual []float64) float64 {
    var loss float64
    for i := range predicted {
        diff := predicted[i] - actual[i]
        loss += diff * diff
    }
    return loss / float64(len(predicted))
}

func average(values []float64) float64 {
    var sum float64
    for _, v := range values {
        sum += v
    }
    return sum / float64(len(values))
}

func randFloat64() float64 {
    // 生产环境应使用crypto/rand
    return float64(time.Now().UnixNano()%10000) / 10000
}

func randNormal(mean, stddev float64) float64 {
    // Box-Muller变换生成正态分布随机数
    u1 := randFloat64()
    u2 := randFloat64()
    return mean + stddev*math.Sqrt(-2*math.Log(u1))*math.Cos(2*math.Pi*u2)
}

func (opt *SDSOptimizer) parallelOptimizeStep(cameras []*geometry.Camera, light *geometry.Lighting, textCondition string) (float64, error) {
    var wg sync.WaitGroup
    gradients := make([][]float64, len(cameras))
    losses := make([]float64, len(cameras))
    errors := make([]error, len(cameras))
    
    for i, cam := range cameras {
        wg.Add(1)
        go func(idx int, camera *geometry.Camera) {
            defer wg.Done()
            // 每个goroutine独立计算一个视角的梯度
            loss, grad, err := opt.computeSingleViewGradient(camera, light, textCondition)
            if err != nil {
                errors[idx] = err
                return
            }
            gradients[idx] = grad
            losses[idx] = loss
        }(i, cam)
    }
    
    wg.Wait()
    
    // 聚合所有视角的梯度
    aggregatedGrad := aggregateGradients(gradients)
    avgLoss := average(losses)
    
    // 更新3D表示
    err := opt.scene.Update(aggregatedGrad)
    if err != nil {
        return 0, fmt.Errorf("更新失败: %w", err)
    }
    
    return avgLoss, nil
}