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

// GAN.hpp — Generative Adversarial Network
//
// Generator: ノイズ → 偽データ
// Discriminator: データ → 真偽判定 (0/1)

#ifndef CALX_ML_GENERATIVE_GAN_HPP
#define CALX_ML_GENERATIVE_GAN_HPP

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

namespace calx::ml {

    /// GAN (Generative Adversarial Network)
    template<typename T>
    class GAN {
    public:
        /// @param noiseDim   ノイズ入力次元
        /// @param hiddenDim  隠れ層次元
        /// @param dataDim    データ次元
        /// @param seed       乱数シード
        GAN(size_t noiseDim, size_t hiddenDim, size_t dataDim, unsigned seed = 42)
            : noiseDim_(noiseDim), dataDim_(dataDim), rng_(seed)
            // Generator: noise → hidden → data
            , gHidden_(noiseDim, hiddenDim, true, seed)
            , gOut_(hiddenDim, dataDim, true, seed + 1)
            , gRelu_(), gTanh_()
            // Discriminator: data → hidden → 1
            , dHidden_(dataDim, hiddenDim, true, seed + 2)
            , dOut_(hiddenDim, 1, true, seed + 3)
            , dRelu_(), dSigmoid_()
        {}

        /// Generator 順伝播
        Matrix<T> generate(const Matrix<T>& noise) {
            auto h = gRelu_.forward(gHidden_.forward(noise));
            return gTanh_.forward(gOut_.forward(h));
        }

        /// Discriminator 順伝播
        Matrix<T> discriminate(const Matrix<T>& data) {
            auto h = dRelu_.forward(dHidden_.forward(data));
            return dSigmoid_.forward(dOut_.forward(h));
        }

        /// Discriminator を1ステップ学習
        /// @return Discriminator loss
        T trainDiscriminator(const Matrix<T>& realData, const Matrix<T>& noise,
                              nn::OptimizerBase<T>& optD)
        {
            size_t batch = realData.rows();
            auto dParams = discriminatorParameters();

            // 本物: label = 1
            auto realPred = discriminate(realData);
            Matrix<T> onesTarget(batch, 1, T{1});
            Matrix<T> gradReal;
            T lossReal = nn::binaryCrossEntropyLoss(realPred, onesTarget, gradReal);

            nn::OptimizerBase<T>::zeroGrad(dParams);
            dSigmoid_.backward(gradReal);
            dOut_.backward(dSigmoid_.backward(gradReal));
            auto dH = dHidden_.backward(dRelu_.backward(dOut_.backward(dSigmoid_.backward(gradReal))));

            // 偽物: label = 0
            auto fakeData = generate(noise);
            auto fakePred = discriminate(fakeData);
            Matrix<T> zerosTarget(batch, 1, T{0});
            Matrix<T> gradFake;
            T lossFake = nn::binaryCrossEntropyLoss(fakePred, zerosTarget, gradFake);

            // Discriminator backward (偽データに対して)
            auto dSig = dSigmoid_.backward(gradFake);
            auto dO = dOut_.backward(dSig);
            dHidden_.backward(dRelu_.backward(dO));

            optD.step(dParams);
            return (lossReal + lossFake) / T{2};
        }

        /// Generator を1ステップ学習
        /// @return Generator loss
        T trainGenerator(const Matrix<T>& noise, nn::OptimizerBase<T>& optG) {
            size_t batch = noise.rows();
            auto gParams = generatorParameters();

            auto fakeData = generate(noise);
            auto fakePred = discriminate(fakeData);

            // Generator は Discriminator を騙したい → label = 1
            Matrix<T> onesTarget(batch, 1, T{1});
            Matrix<T> gradPred;
            T loss = nn::binaryCrossEntropyLoss(fakePred, onesTarget, gradPred);

            nn::OptimizerBase<T>::zeroGrad(gParams);

            // Discriminator backward (勾配を Generator に伝播)
            auto dSig = dSigmoid_.backward(gradPred);
            auto dO = dOut_.backward(dSig);
            auto dH = dHidden_.backward(dRelu_.backward(dO));

            // dH は Generator 出力への勾配
            auto dGTanh = gTanh_.backward(dH);
            auto dGO = gOut_.backward(dGTanh);
            gHidden_.backward(gRelu_.backward(dGO));

            optG.step(gParams);
            return loss;
        }

        /// ノイズ生成
        Matrix<T> sampleNoise(size_t batch) {
            std::normal_distribution<T> dist(T{0}, T{1});
            Matrix<T> noise(batch, noiseDim_);
            for (size_t b = 0; b < batch; ++b)
                for (size_t d = 0; d < noiseDim_; ++d)
                    noise(b, d) = dist(rng_);
            return noise;
        }

        std::vector<nn::ParamGrad<T>> generatorParameters() {
            auto p = gHidden_.parameters();
            auto p2 = gOut_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            return p;
        }

        std::vector<nn::ParamGrad<T>> discriminatorParameters() {
            auto p = dHidden_.parameters();
            auto p2 = dOut_.parameters();
            p.insert(p.end(), p2.begin(), p2.end());
            return p;
        }

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

        // Generator
        nn::Linear<T> gHidden_, gOut_;
        nn::ReLU<T> gRelu_;
        nn::TanhLayer<T> gTanh_;

        // Discriminator
        nn::Linear<T> dHidden_, dOut_;
        nn::ReLU<T> dRelu_;
        nn::Sigmoid<T> dSigmoid_;
    };

} // namespace calx::ml

#endif // CALX_ML_GENERATIVE_GAN_HPP
