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

// VAE.hpp — Variational Autoencoder
//
// Encoder → (μ, log σ²) → Reparameterization → Decoder
// 損失: 再構成誤差 (MSE) + KL ダイバージェンス

#ifndef CALX_ML_GENERATIVE_VAE_HPP
#define CALX_ML_GENERATIVE_VAE_HPP

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

namespace calx::ml {

    /// VAE (Variational Autoencoder)
    template<typename T>
    class VAE {
    public:
        /// @param inputDim   入力次元
        /// @param hiddenDim  隠れ層次元
        /// @param latentDim  潜在空間次元
        /// @param seed       乱数シード
        VAE(size_t inputDim, size_t hiddenDim, size_t latentDim, unsigned seed = 42)
            : inputDim_(inputDim), latentDim_(latentDim), rng_(seed)
            , encHidden_(inputDim, hiddenDim, true, seed)
            , encMu_(hiddenDim, latentDim, true, seed + 1)
            , encLogVar_(hiddenDim, latentDim, true, seed + 2)
            , decHidden_(latentDim, hiddenDim, true, seed + 3)
            , decOut_(hiddenDim, inputDim, true, seed + 4)
            , reluEnc_(), reluDec_(), sigmoidOut_()
        {}

        /// 順伝播 (encode + reparameterize + decode)
        Matrix<T> forward(const Matrix<T>& input) {
            size_t batch = input.rows();

            // Encoder
            auto h = reluEnc_.forward(encHidden_.forward(input));
            mu_ = encMu_.forward(h);
            logVar_ = encLogVar_.forward(h);

            // Reparameterization: z = μ + σ * ε
            std::normal_distribution<T> normal(T{0}, T{1});
            z_ = Matrix<T>(batch, latentDim_);
            eps_ = Matrix<T>(batch, latentDim_);
            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < latentDim_; ++d) {
                    eps_(b, d) = normal(rng_);
                    T sigma = std::exp(T{0.5} * logVar_(b, d));
                    z_(b, d) = mu_(b, d) + sigma * eps_(b, d);
                }

            // Decoder
            auto dh = reluDec_.forward(decHidden_.forward(z_));
            return sigmoidOut_.forward(decOut_.forward(dh));
        }

        /// 損失計算 (再構成 + KL)
        T loss(const Matrix<T>& input, const Matrix<T>& recon, Matrix<T>& gradRecon) {
            size_t batch = input.rows();

            // 再構成誤差 (BCE 簡易版: MSE で代替)
            T reconLoss = nn::mseLoss(recon, input, gradRecon);

            // KL ダイバージェンス: -0.5 * Σ(1 + log σ² - μ² - σ²)
            T klLoss = T{0};
            T invN = T{1} / static_cast<T>(batch);
            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < latentDim_; ++d) {
                    T m = mu_(b, d), lv = logVar_(b, d);
                    klLoss += T{-0.5} * (T{1} + lv - m * m - std::exp(lv));
                }
            klLoss *= invN;

            return reconLoss + klLoss;
        }

        /// 逆伝播
        void backward(const Matrix<T>& gradRecon) {
            size_t batch = gradRecon.rows();

            // Decoder backward
            auto dDecOut = sigmoidOut_.backward(gradRecon);
            auto dDecHidden = decOut_.backward(dDecOut);
            auto dz = decHidden_.backward(reluDec_.backward(dDecHidden));

            // KL 勾配を z, μ, logVar に加算
            T invN = T{1} / static_cast<T>(batch);
            Matrix<T> dMu(batch, latentDim_);
            Matrix<T> dLogVar(batch, latentDim_);
            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < latentDim_; ++d) {
                    T m = mu_(b, d), lv = logVar_(b, d);
                    // dz を μ と logVar に分配
                    dMu(b, d) = dz(b, d) + m * invN;  // KL: dKL/dμ = μ
                    T sigma = std::exp(T{0.5} * lv);
                    dLogVar(b, d) = dz(b, d) * eps_(b, d) * T{0.5} * sigma
                                    + T{0.5} * (std::exp(lv) - T{1}) * invN;
                }

            // Encoder backward
            auto dEncMu = encMu_.backward(dMu);
            auto dEncLogVar = encLogVar_.backward(dLogVar);

            // 二つの勾配を加算
            Matrix<T> dh(batch, dEncMu.cols());
            for (size_t b = 0; b < batch; ++b)
                for (size_t j = 0; j < dh.cols(); ++j)
                    dh(b, j) = dEncMu(b, j) + dEncLogVar(b, j);

            encHidden_.backward(reluEnc_.backward(dh));
        }

        /// 全パラメータ
        std::vector<nn::ParamGrad<T>> parameters() {
            auto p = encHidden_.parameters();
            auto p2 = encMu_.parameters();
            auto p3 = encLogVar_.parameters();
            auto p4 = decHidden_.parameters();
            auto p5 = decOut_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            p.insert(p.end(), p3.begin(), p3.end());
            p.insert(p.end(), p4.begin(), p4.end());
            p.insert(p.end(), p5.begin(), p5.end());
            return p;
        }

        /// 潜在空間からサンプリングしてデコード
        Matrix<T> decode(const Matrix<T>& z) {
            auto dh = reluDec_.forward(decHidden_.forward(z));
            return sigmoidOut_.forward(decOut_.forward(dh));
        }

        /// 最新のμを取得
        const Matrix<T>& latentMu() const { return mu_; }
        const Matrix<T>& latentLogVar() const { return logVar_; }

    private:
        size_t inputDim_, latentDim_;
        std::mt19937 rng_;

        nn::Linear<T> encHidden_, encMu_, encLogVar_;
        nn::Linear<T> decHidden_, decOut_;
        nn::ReLU<T> reluEnc_, reluDec_;
        nn::Sigmoid<T> sigmoidOut_;

        Matrix<T> mu_, logVar_, z_, eps_;
    };

} // namespace calx::ml

#endif // CALX_ML_GENERATIVE_VAE_HPP
