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

// Conv2D.hpp — 2D 畳み込み層
//
// im2col + GEMM 方式。Layer<T> を継承し、Matrix<T> のフラット表現を使用。
// 入出力: Matrix<T> (batch × C*H*W)。空間サイズはレイヤー内部で管理。

#ifndef CALX_ML_NN_CONV2D_HPP
#define CALX_ML_NN_CONV2D_HPP

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

namespace calx::ml::nn {

    /// 2D 畳み込み層
    /// 入力: Matrix (batch × inC*inH*inW), 出力: Matrix (batch × outC*outH*outW)
    template<typename T>
    class Conv2D : public Layer<T> {
    public:
        /// @param inChannels   入力チャネル数
        /// @param outChannels  出力チャネル数 (フィルタ数)
        /// @param kernelSize   カーネルサイズ (正方形)
        /// @param inH          入力高さ
        /// @param inW          入力幅
        /// @param stride       ストライド
        /// @param padding      パディング
        /// @param seed         乱数シード
        Conv2D(size_t inChannels, size_t outChannels, size_t kernelSize,
               size_t inH, size_t inW, size_t stride = 1, size_t padding = 0,
               unsigned seed = 42)
            : inC_(inChannels), outC_(outChannels), kSize_(kernelSize)
            , inH_(inH), inW_(inW), stride_(stride), pad_(padding)
        {
            outH_ = (inH_ + 2 * pad_ - kSize_) / stride_ + 1;
            outW_ = (inW_ + 2 * pad_ - kSize_) / stride_ + 1;

            // Kaiming 初期化
            size_t fanIn = inC_ * kSize_ * kSize_;
            T scale = std::sqrt(T{2} / static_cast<T>(fanIn));
            std::mt19937 rng(seed);
            std::normal_distribution<T> dist(T{0}, scale);

            // weight: (outC, inC * kSize * kSize)
            size_t colSize = inC_ * kSize_ * kSize_;
            weight_ = Matrix<T>(outC_, colSize);
            gradWeight_ = Matrix<T>(outC_, colSize, T{0});
            for (size_t i = 0; i < outC_; ++i)
                for (size_t j = 0; j < colSize; ++j)
                    weight_(i, j) = dist(rng);

            // bias: (1, outC)
            bias_ = Matrix<T>(1, outC_, T{0});
            gradBias_ = Matrix<T>(1, outC_, T{0});
        }

        /// 出力空間サイズ
        size_t outH() const { return outH_; }
        size_t outW() const { return outW_; }
        size_t outChannels() const { return outC_; }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            size_t colSize = inC_ * kSize_ * kSize_;
            size_t spatialOut = outH_ * outW_;

            // im2col 変換
            col_ = Matrix<T>(batch * spatialOut, colSize, T{0});
            im2col(input, col_, batch);

            // GEMM: output = col * weight^T + bias
            Matrix<T> output(batch, outC_ * spatialOut);
            for (size_t b = 0; b < batch; ++b) {
                for (size_t oh = 0; oh < spatialOut; ++oh) {
                    size_t colRow = b * spatialOut + oh;
                    for (size_t oc = 0; oc < outC_; ++oc) {
                        T sum = bias_(0, oc);
                        for (size_t k = 0; k < colSize; ++k)
                            sum += col_(colRow, k) * weight_(oc, k);
                        output(b, oc * spatialOut + oh) = sum;
                    }
                }
            }

            inputRows_ = batch;
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = inputRows_;
            size_t colSize = inC_ * kSize_ * kSize_;
            size_t spatialOut = outH_ * outW_;

            // gradWeight = gradOutput^T * col (累積)
            for (size_t b = 0; b < batch; ++b) {
                for (size_t oc = 0; oc < outC_; ++oc) {
                    for (size_t k = 0; k < colSize; ++k) {
                        T sum = T{0};
                        for (size_t oh = 0; oh < spatialOut; ++oh)
                            sum += gradOutput(b, oc * spatialOut + oh)
                                   * col_(b * spatialOut + oh, k);
                        gradWeight_(oc, k) += sum;
                    }
                    // gradBias
                    T bsum = T{0};
                    for (size_t oh = 0; oh < spatialOut; ++oh)
                        bsum += gradOutput(b, oc * spatialOut + oh);
                    gradBias_(0, oc) += bsum;
                }
            }

            // gradCol = gradOutput * weight
            Matrix<T> gradCol(batch * spatialOut, colSize, T{0});
            for (size_t b = 0; b < batch; ++b) {
                for (size_t oh = 0; oh < spatialOut; ++oh) {
                    size_t colRow = b * spatialOut + oh;
                    for (size_t k = 0; k < colSize; ++k) {
                        T sum = T{0};
                        for (size_t oc = 0; oc < outC_; ++oc)
                            sum += gradOutput(b, oc * spatialOut + oh) * weight_(oc, k);
                        gradCol(colRow, k) = sum;
                    }
                }
            }

            // col2im: gradCol → gradInput
            Matrix<T> gradInput(batch, inC_ * inH_ * inW_, T{0});
            col2im(gradCol, gradInput, batch);
            return gradInput;
        }

        std::vector<ParamGrad<T>> parameters() override {
            return { {&weight_, &gradWeight_}, {&bias_, &gradBias_} };
        }

    private:
        size_t inC_, outC_, kSize_, inH_, inW_, stride_, pad_;
        size_t outH_, outW_;
        size_t inputRows_ = 0;
        Matrix<T> weight_, bias_;
        Matrix<T> gradWeight_, gradBias_;
        Matrix<T> col_;  // im2col キャッシュ

        // im2col: 入力を列行列に変換
        void im2col(const Matrix<T>& input, Matrix<T>& col, size_t batch) {
            size_t spatialOut = outH_ * outW_;
            for (size_t b = 0; b < batch; ++b) {
                for (size_t oh = 0; oh < outH_; ++oh) {
                    for (size_t ow = 0; ow < outW_; ++ow) {
                        size_t colRow = b * spatialOut + oh * outW_ + ow;
                        size_t colIdx = 0;
                        for (size_t ic = 0; ic < inC_; ++ic) {
                            for (size_t kh = 0; kh < kSize_; ++kh) {
                                for (size_t kw = 0; kw < kSize_; ++kw) {
                                    int ih = static_cast<int>(oh * stride_ + kh) - static_cast<int>(pad_);
                                    int iw = static_cast<int>(ow * stride_ + kw) - static_cast<int>(pad_);
                                    if (ih >= 0 && ih < static_cast<int>(inH_) &&
                                        iw >= 0 && iw < static_cast<int>(inW_)) {
                                        size_t inIdx = ic * inH_ * inW_ + ih * inW_ + iw;
                                        col(colRow, colIdx) = input(b, inIdx);
                                    }
                                    ++colIdx;
                                }
                            }
                        }
                    }
                }
            }
        }

        // col2im: 列行列を入力形状に逆変換
        void col2im(const Matrix<T>& col, Matrix<T>& output, size_t batch) {
            size_t spatialOut = outH_ * outW_;
            for (size_t b = 0; b < batch; ++b) {
                for (size_t oh = 0; oh < outH_; ++oh) {
                    for (size_t ow = 0; ow < outW_; ++ow) {
                        size_t colRow = b * spatialOut + oh * outW_ + ow;
                        size_t colIdx = 0;
                        for (size_t ic = 0; ic < inC_; ++ic) {
                            for (size_t kh = 0; kh < kSize_; ++kh) {
                                for (size_t kw = 0; kw < kSize_; ++kw) {
                                    int ih = static_cast<int>(oh * stride_ + kh) - static_cast<int>(pad_);
                                    int iw = static_cast<int>(ow * stride_ + kw) - static_cast<int>(pad_);
                                    if (ih >= 0 && ih < static_cast<int>(inH_) &&
                                        iw >= 0 && iw < static_cast<int>(inW_)) {
                                        size_t inIdx = ic * inH_ * inW_ + ih * inW_ + iw;
                                        output(b, inIdx) += col(colRow, colIdx);
                                    }
                                    ++colIdx;
                                }
                            }
                        }
                    }
                }
            }
        }
    };

    /// Flatten 層 (Conv 出力 → Linear 入力のブリッジ)
    /// 入力がすでにフラット (batch × features) なので実質パススルー
    template<typename T>
    class Flatten : public Layer<T> {
    public:
        Matrix<T> forward(const Matrix<T>& input) override { return input; }
        Matrix<T> backward(const Matrix<T>& gradOutput) override { return gradOutput; }
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_CONV2D_HPP
