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

// Pooling.hpp — プーリング層
//
// MaxPool2D, AvgPool2D。
// 入出力: Matrix<T> (batch × C*H*W) のフラット表現。

#ifndef CALX_ML_NN_POOLING_HPP
#define CALX_ML_NN_POOLING_HPP

#include <ml/nn/Layer.hpp>
#include <algorithm>
#include <limits>

namespace calx::ml::nn {

    /// MaxPool2D
    template<typename T>
    class MaxPool2D : public Layer<T> {
    public:
        /// @param channels  チャネル数
        /// @param inH       入力高さ
        /// @param inW       入力幅
        /// @param poolSize  プーリングサイズ (正方形)
        /// @param stride    ストライド (0 = poolSize)
        MaxPool2D(size_t channels, size_t inH, size_t inW,
                  size_t poolSize = 2, size_t stride = 0)
            : C_(channels), inH_(inH), inW_(inW)
            , pool_(poolSize), stride_(stride == 0 ? poolSize : stride)
        {
            outH_ = (inH_ - pool_) / stride_ + 1;
            outW_ = (inW_ - pool_) / stride_ + 1;
        }

        size_t outH() const { return outH_; }
        size_t outW() const { return outW_; }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            size_t outSpatial = outH_ * outW_;
            Matrix<T> output(batch, C_ * outSpatial);
            maxIndices_ = std::vector<size_t>(batch * C_ * outSpatial, 0);

            for (size_t b = 0; b < batch; ++b) {
                for (size_t c = 0; c < C_; ++c) {
                    for (size_t oh = 0; oh < outH_; ++oh) {
                        for (size_t ow = 0; ow < outW_; ++ow) {
                            T maxVal = std::numeric_limits<T>::lowest();
                            size_t maxIdx = 0;
                            for (size_t ph = 0; ph < pool_; ++ph) {
                                for (size_t pw = 0; pw < pool_; ++pw) {
                                    size_t ih = oh * stride_ + ph;
                                    size_t iw = ow * stride_ + pw;
                                    size_t inIdx = c * inH_ * inW_ + ih * inW_ + iw;
                                    if (input(b, inIdx) > maxVal) {
                                        maxVal = input(b, inIdx);
                                        maxIdx = inIdx;
                                    }
                                }
                            }
                            size_t outIdx = c * outSpatial + oh * outW_ + ow;
                            output(b, outIdx) = maxVal;
                            maxIndices_[b * C_ * outSpatial + outIdx] = maxIdx;
                        }
                    }
                }
            }
            inputCols_ = input.cols();
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = gradOutput.rows();
            size_t outSpatial = outH_ * outW_;
            Matrix<T> gradInput(batch, inputCols_, T{0});

            for (size_t b = 0; b < batch; ++b) {
                for (size_t c = 0; c < C_; ++c) {
                    for (size_t oh = 0; oh < outH_; ++oh) {
                        for (size_t ow = 0; ow < outW_; ++ow) {
                            size_t outIdx = c * outSpatial + oh * outW_ + ow;
                            size_t maxIdx = maxIndices_[b * C_ * outSpatial + outIdx];
                            gradInput(b, maxIdx) += gradOutput(b, outIdx);
                        }
                    }
                }
            }
            return gradInput;
        }

    private:
        size_t C_, inH_, inW_, pool_, stride_, outH_, outW_;
        size_t inputCols_ = 0;
        std::vector<size_t> maxIndices_;
    };

    /// AvgPool2D
    template<typename T>
    class AvgPool2D : public Layer<T> {
    public:
        AvgPool2D(size_t channels, size_t inH, size_t inW,
                  size_t poolSize = 2, size_t stride = 0)
            : C_(channels), inH_(inH), inW_(inW)
            , pool_(poolSize), stride_(stride == 0 ? poolSize : stride)
        {
            outH_ = (inH_ - pool_) / stride_ + 1;
            outW_ = (inW_ - pool_) / stride_ + 1;
        }

        size_t outH() const { return outH_; }
        size_t outW() const { return outW_; }

        Matrix<T> forward(const Matrix<T>& input) override {
            size_t batch = input.rows();
            size_t outSpatial = outH_ * outW_;
            T invPool = T{1} / static_cast<T>(pool_ * pool_);
            Matrix<T> output(batch, C_ * outSpatial);

            for (size_t b = 0; b < batch; ++b) {
                for (size_t c = 0; c < C_; ++c) {
                    for (size_t oh = 0; oh < outH_; ++oh) {
                        for (size_t ow = 0; ow < outW_; ++ow) {
                            T sum = T{0};
                            for (size_t ph = 0; ph < pool_; ++ph)
                                for (size_t pw = 0; pw < pool_; ++pw) {
                                    size_t ih = oh * stride_ + ph;
                                    size_t iw = ow * stride_ + pw;
                                    sum += input(b, c * inH_ * inW_ + ih * inW_ + iw);
                                }
                            output(b, c * outSpatial + oh * outW_ + ow) = sum * invPool;
                        }
                    }
                }
            }
            inputCols_ = input.cols();
            return output;
        }

        Matrix<T> backward(const Matrix<T>& gradOutput) override {
            size_t batch = gradOutput.rows();
            size_t outSpatial = outH_ * outW_;
            T invPool = T{1} / static_cast<T>(pool_ * pool_);
            Matrix<T> gradInput(batch, inputCols_, T{0});

            for (size_t b = 0; b < batch; ++b) {
                for (size_t c = 0; c < C_; ++c) {
                    for (size_t oh = 0; oh < outH_; ++oh) {
                        for (size_t ow = 0; ow < outW_; ++ow) {
                            T g = gradOutput(b, c * outSpatial + oh * outW_ + ow) * invPool;
                            for (size_t ph = 0; ph < pool_; ++ph)
                                for (size_t pw = 0; pw < pool_; ++pw) {
                                    size_t ih = oh * stride_ + ph;
                                    size_t iw = ow * stride_ + pw;
                                    gradInput(b, c * inH_ * inW_ + ih * inW_ + iw) += g;
                                }
                        }
                    }
                }
            }
            return gradInput;
        }

    private:
        size_t C_, inH_, inW_, pool_, stride_, outH_, outW_;
        size_t inputCols_ = 0;
    };

} // namespace calx::ml::nn

#endif // CALX_ML_NN_POOLING_HPP
