// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// Diffusion.hpp — DDPM (Denoising Diffusion Probabilistic Model)
//
// 拡散過程: x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε
// 逆拡散: ノイズ予測ネットワークで ε を推定し段階的にデノイズ

#ifndef CALX_ML_GENERATIVE_DIFFUSION_HPP
#define CALX_ML_GENERATIVE_DIFFUSION_HPP

#include <ml/nn/Layer.hpp>
#include <ml/nn/Activation.hpp>
#include <ml/nn/Optimizer.hpp>
#include <cmath>
#include <random>

namespace calx::ml {

    /// DDPM (Denoising Diffusion Probabilistic Model)
    template<typename T>
    class DDPM {
    public:
        /// @param dataDim    データ次元
        /// @param hiddenDim  ノイズ予測ネットワークの隠れ層次元
        /// @param numSteps   拡散ステップ数
        /// @param betaStart  ノイズスケジュール開始値
        /// @param betaEnd    ノイズスケジュール終了値
        /// @param seed       乱数シード
        DDPM(size_t dataDim, size_t hiddenDim, size_t numSteps = 100,
             T betaStart = T{0.0001}, T betaEnd = T{0.02}, unsigned seed = 42)
            : dataDim_(dataDim), numSteps_(numSteps), rng_(seed)
            // ノイズ予測ネット: (dataDim + 1) → hidden → dataDim
            // +1 はタイムステップ埋め込み
            , net1_(dataDim + 1, hiddenDim, true, seed)
            , net2_(hiddenDim, hiddenDim, true, seed + 1)
            , net3_(hiddenDim, dataDim, true, seed + 2)
            , relu1_(), relu2_()
        {
            // ノイズスケジュール (線形)
            beta_.resize(numSteps);
            alpha_.resize(numSteps);
            alphaBar_.resize(numSteps);
            for (size_t t = 0; t < numSteps; ++t) {
                beta_[t] = betaStart + (betaEnd - betaStart) * static_cast<T>(t) / static_cast<T>(numSteps - 1);
                alpha_[t] = T{1} - beta_[t];
                alphaBar_[t] = (t == 0) ? alpha_[t] : alphaBar_[t - 1] * alpha_[t];
            }
        }

        /// 拡散: x_0 にノイズを加えて x_t を生成
        /// @return {x_t, ε} のペア
        std::pair<Matrix<T>, Matrix<T>> diffuse(const Matrix<T>& x0, size_t t) {
            size_t batch = x0.rows();
            std::normal_distribution<T> normal(T{0}, T{1});

            Matrix<T> eps(batch, dataDim_);
            Matrix<T> xt(batch, dataDim_);
            T sqrtAbar = std::sqrt(alphaBar_[t]);
            T sqrtOneMinusAbar = std::sqrt(T{1} - alphaBar_[t]);

            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < dataDim_; ++d) {
                    eps(b, d) = normal(rng_);
                    xt(b, d) = sqrtAbar * x0(b, d) + sqrtOneMinusAbar * eps(b, d);
                }
            return { std::move(xt), std::move(eps) };
        }

        /// ノイズ予測
        Matrix<T> predictNoise(const Matrix<T>& xt, size_t t) {
            size_t batch = xt.rows();
            // タイムステップを正規化して入力に結合
            T tNorm = static_cast<T>(t) / static_cast<T>(numSteps_);
            Matrix<T> input(batch, dataDim_ + 1);
            for (size_t b = 0; b < batch; ++b) {
                for (size_t d = 0; d < dataDim_; ++d)
                    input(b, d) = xt(b, d);
                input(b, dataDim_) = tNorm;
            }

            auto h1 = relu1_.forward(net1_.forward(input));
            auto h2 = relu2_.forward(net2_.forward(h1));
            return net3_.forward(h2);
        }

        /// 1ステップ学習 (ノイズ予測の MSE)
        T trainStep(const Matrix<T>& x0, nn::OptimizerBase<T>& opt) {
            auto params = parameters();

            // ランダムなタイムステップ
            std::uniform_int_distribution<size_t> tDist(0, numSteps_ - 1);
            size_t t = tDist(rng_);

            auto [xt, eps] = diffuse(x0, t);
            auto epsPred = predictNoise(xt, t);

            // MSE 損失
            Matrix<T> grad;
            T loss = nn::mseLoss(epsPred, eps, grad);

            nn::OptimizerBase<T>::zeroGrad(params);
            auto d3 = net3_.backward(grad);
            auto d2 = net2_.backward(relu2_.backward(d3));
            net1_.backward(relu1_.backward(d2));

            opt.step(params);
            return loss;
        }

        /// サンプリング (逆拡散)
        Matrix<T> sample(size_t batch) {
            std::normal_distribution<T> normal(T{0}, T{1});

            // x_T ~ N(0, I)
            Matrix<T> x(batch, dataDim_);
            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < dataDim_; ++d)
                    x(b, d) = normal(rng_);

            // 逆拡散
            for (int t = static_cast<int>(numSteps_) - 1; t >= 0; --t) {
                auto epsPred = predictNoise(x, t);

                T sqrtAlpha = std::sqrt(alpha_[t]);
                T sqrtOneMinusAbar = std::sqrt(T{1} - alphaBar_[t]);
                T coeff = (T{1} - alpha_[t]) / sqrtOneMinusAbar;

                for (size_t b = 0; b < batch; ++b)
                    for (size_t d = 0; d < dataDim_; ++d) {
                        x(b, d) = (x(b, d) - coeff * epsPred(b, d)) / sqrtAlpha;
                        if (t > 0) {
                            T sigma = std::sqrt(beta_[t]);
                            x(b, d) += sigma * normal(rng_);
                        }
                    }
            }
            return x;
        }

        std::vector<nn::ParamGrad<T>> parameters() {
            auto p = net1_.parameters();
            auto p2 = net2_.parameters();
            auto p3 = net3_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            p.insert(p.end(), p3.begin(), p3.end());
            return p;
        }

    private:
        size_t dataDim_, numSteps_;
        std::mt19937 rng_;

        nn::Linear<T> net1_, net2_, net3_;
        nn::ReLU<T> relu1_, relu2_;

        std::vector<T> beta_, alpha_, alphaBar_;
    };

} // namespace calx::ml

#endif // CALX_ML_GENERATIVE_DIFFUSION_HPP
