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

// Map.hpp
// 外部メモリの zero-copy ラッピング (Eigen::Map 相当)
//
// MatrixMap<T> — 外部 T* を行列として操作 (stride 対応)
// VectorMap<T> — 外部 T* をベクトルとして操作 (stride 対応)

#ifndef CALX_MAP_HPP
#define CALX_MAP_HPP

#include <math/core/common.hpp>
#include <cstddef>
#include <stdexcept>

namespace calx {

    // ── VectorMap<T> ──────────────────────────────────────────────
    // 外部メモリ上のベクトルビュー。所有権を持たない。

    template<typename T>
    class VectorMap {
    public:
        using value_type = T;
        using size_type = std::size_t;
        using difference_type = std::ptrdiff_t;
        using reference = T&;
        using const_reference = const T&;
        using pointer = T*;
        using const_pointer = const T*;

        // 連続メモリ (stride=1)
        VectorMap(T* data, size_type size)
            : data_(data), size_(size), stride_(1) {
            if (!data && size > 0)
                throw std::invalid_argument("VectorMap: null pointer with non-zero size");
        }

        // stride 指定
        VectorMap(T* data, size_type size, size_type stride)
            : data_(data), size_(size), stride_(stride) {
            if (!data && size > 0)
                throw std::invalid_argument("VectorMap: null pointer with non-zero size");
            if (stride == 0)
                throw std::invalid_argument("VectorMap: stride must be > 0");
        }

        // コピー/ムーブ (ポインタのコピーのみ)
        VectorMap(const VectorMap&) = default;
        VectorMap& operator=(const VectorMap&) = default;

        // 要素アクセス
        reference operator[](size_type i) { return data_[i * stride_]; }
        const_reference operator[](size_type i) const { return data_[i * stride_]; }

        reference at(size_type i) {
            if (i >= size_) throw IndexError("VectorMap::at: index out of range");
            return data_[i * stride_];
        }
        const_reference at(size_type i) const {
            if (i >= size_) throw IndexError("VectorMap::at: index out of range");
            return data_[i * stride_];
        }

        // 容量
        size_type size() const noexcept { return size_; }
        size_type stride() const noexcept { return stride_; }
        bool empty() const noexcept { return size_ == 0; }
        pointer data() noexcept { return data_; }
        const_pointer data() const noexcept { return data_; }

        // fill
        void fill(const T& value) {
            for (size_type i = 0; i < size_; ++i)
                data_[i * stride_] = value;
        }

        void zero() { fill(T{}); }

        // ドット積
        T dot(const VectorMap& other) const {
            if (size_ != other.size_)
                throw DimensionError("VectorMap::dot: size mismatch");
            T sum{};
            for (size_type i = 0; i < size_; ++i)
                sum += (*this)[i] * other[i];
            return sum;
        }

        // ノルム二乗
        T squaredNorm() const {
            T sum{};
            for (size_type i = 0; i < size_; ++i) {
                T v = (*this)[i];
                sum += v * v;
            }
            return sum;
        }

        // Matrix<T> / Vector<T> へのコピー出力
        template<typename VecType>
        void copyTo(VecType& dest) const {
            if (dest.size() != size_) dest.resize(size_);
            for (size_type i = 0; i < size_; ++i)
                dest[i] = (*this)[i];
        }

        // 他コンテナからのコピー入力
        template<typename VecType>
        void copyFrom(const VecType& src) {
            if (src.size() != size_)
                throw DimensionError("VectorMap::copyFrom: size mismatch");
            for (size_type i = 0; i < size_; ++i)
                (*this)[i] = src[i];
        }

        // サブビュー
        VectorMap segment(size_type offset, size_type count) const {
            if (offset + count > size_)
                throw IndexError("VectorMap::segment: range exceeds bounds");
            return VectorMap(data_ + offset * stride_, count, stride_);
        }

        // stride 付きイテレータ
        class iterator {
        public:
            using iterator_category = std::random_access_iterator_tag;
            using value_type = T;
            using difference_type = std::ptrdiff_t;
            using pointer = T*;
            using reference = T&;

            iterator() : ptr_(nullptr), stride_(1) {}
            iterator(T* p, size_type s) : ptr_(p), stride_(s) {}

            reference operator*() const { return *ptr_; }
            pointer operator->() const { return ptr_; }
            reference operator[](difference_type n) const { return ptr_[n * static_cast<difference_type>(stride_)]; }

            iterator& operator++() { ptr_ += stride_; return *this; }
            iterator operator++(int) { auto t = *this; ptr_ += stride_; return t; }
            iterator& operator--() { ptr_ -= stride_; return *this; }
            iterator operator--(int) { auto t = *this; ptr_ -= stride_; return t; }

            iterator& operator+=(difference_type n) { ptr_ += n * static_cast<difference_type>(stride_); return *this; }
            iterator& operator-=(difference_type n) { ptr_ -= n * static_cast<difference_type>(stride_); return *this; }

            friend iterator operator+(iterator it, difference_type n) { it += n; return it; }
            friend iterator operator+(difference_type n, iterator it) { it += n; return it; }
            friend iterator operator-(iterator it, difference_type n) { it -= n; return it; }
            friend difference_type operator-(const iterator& a, const iterator& b) {
                return (a.ptr_ - b.ptr_) / static_cast<difference_type>(a.stride_);
            }

            bool operator==(const iterator& o) const { return ptr_ == o.ptr_; }
            bool operator!=(const iterator& o) const { return ptr_ != o.ptr_; }
            bool operator<(const iterator& o) const { return ptr_ < o.ptr_; }
            bool operator<=(const iterator& o) const { return ptr_ <= o.ptr_; }
            bool operator>(const iterator& o) const { return ptr_ > o.ptr_; }
            bool operator>=(const iterator& o) const { return ptr_ >= o.ptr_; }

        private:
            T* ptr_;
            size_type stride_;
        };

        using const_iterator = const iterator;

        iterator begin() { return iterator(data_, stride_); }
        iterator end() { return iterator(data_ + size_ * stride_, stride_); }

    private:
        T* data_;
        size_type size_;
        size_type stride_;
    };

    // ── ConstVectorMap<T> ─────────────────────────────────────────

    template<typename T>
    class ConstVectorMap {
    public:
        using value_type = T;
        using size_type = std::size_t;
        using const_reference = const T&;
        using const_pointer = const T*;

        ConstVectorMap(const T* data, size_type size)
            : data_(data), size_(size), stride_(1) {}

        ConstVectorMap(const T* data, size_type size, size_type stride)
            : data_(data), size_(size), stride_(stride) {
            if (stride == 0)
                throw std::invalid_argument("ConstVectorMap: stride must be > 0");
        }

        // VectorMap からの暗黙変換
        ConstVectorMap(const VectorMap<T>& vm)
            : data_(vm.data()), size_(vm.size()), stride_(vm.stride()) {}

        const_reference operator[](size_type i) const { return data_[i * stride_]; }
        const_reference at(size_type i) const {
            if (i >= size_) throw IndexError("ConstVectorMap::at: index out of range");
            return data_[i * stride_];
        }

        size_type size() const noexcept { return size_; }
        size_type stride() const noexcept { return stride_; }
        bool empty() const noexcept { return size_ == 0; }
        const_pointer data() const noexcept { return data_; }

        T dot(const ConstVectorMap& other) const {
            if (size_ != other.size_)
                throw DimensionError("ConstVectorMap::dot: size mismatch");
            T sum{};
            for (size_type i = 0; i < size_; ++i)
                sum += (*this)[i] * other[i];
            return sum;
        }

        T squaredNorm() const {
            T sum{};
            for (size_type i = 0; i < size_; ++i) {
                T v = (*this)[i];
                sum += v * v;
            }
            return sum;
        }

        template<typename VecType>
        void copyTo(VecType& dest) const {
            if (dest.size() != size_) dest.resize(size_);
            for (size_type i = 0; i < size_; ++i)
                dest[i] = (*this)[i];
        }

        ConstVectorMap segment(size_type offset, size_type count) const {
            if (offset + count > size_)
                throw IndexError("ConstVectorMap::segment: range exceeds bounds");
            return ConstVectorMap(data_ + offset * stride_, count, stride_);
        }

    private:
        const T* data_;
        size_type size_;
        size_type stride_;
    };

    // ── MatrixMap<T> ──────────────────────────────────────────────
    // 外部メモリ上の行列ビュー。Row-major。所有権を持たない。

    template<typename T>
    class MatrixMap {
    public:
        using value_type = T;
        using size_type = std::size_t;
        using reference = T&;
        using const_reference = const T&;
        using pointer = T*;
        using const_pointer = const T*;

        // 連続 row-major (outerStride = cols)
        MatrixMap(T* data, size_type rows, size_type cols)
            : data_(data), rows_(rows), cols_(cols), outerStride_(cols) {
            if (!data && rows * cols > 0)
                throw std::invalid_argument("MatrixMap: null pointer with non-zero size");
        }

        // outerStride 指定 (行間のメモリ距離)
        MatrixMap(T* data, size_type rows, size_type cols, size_type outerStride)
            : data_(data), rows_(rows), cols_(cols), outerStride_(outerStride) {
            if (!data && rows * cols > 0)
                throw std::invalid_argument("MatrixMap: null pointer with non-zero size");
            if (outerStride < cols)
                throw std::invalid_argument("MatrixMap: outerStride must be >= cols");
        }

        MatrixMap(const MatrixMap&) = default;
        MatrixMap& operator=(const MatrixMap&) = default;

        // 要素アクセス
        reference operator()(size_type row, size_type col) {
            return data_[row * outerStride_ + col];
        }
        const_reference operator()(size_type row, size_type col) const {
            return data_[row * outerStride_ + col];
        }

        reference at(size_type row, size_type col) {
            checkIndex(row, col);
            return data_[row * outerStride_ + col];
        }
        const_reference at(size_type row, size_type col) const {
            checkIndex(row, col);
            return data_[row * outerStride_ + col];
        }

        // 容量
        size_type rows() const noexcept { return rows_; }
        size_type cols() const noexcept { return cols_; }
        size_type size() const noexcept { return rows_ * cols_; }
        size_type outerStride() const noexcept { return outerStride_; }
        bool empty() const noexcept { return rows_ == 0 || cols_ == 0; }
        bool isSquare() const noexcept { return rows_ == cols_; }
        pointer data() noexcept { return data_; }
        const_pointer data() const noexcept { return data_; }

        // 連続メモリかどうか
        bool isContiguous() const noexcept { return outerStride_ == cols_; }

        // fill / zero
        void fill(const T& value) {
            for (size_type i = 0; i < rows_; ++i)
                for (size_type j = 0; j < cols_; ++j)
                    data_[i * outerStride_ + j] = value;
        }
        void zero() { fill(T{}); }

        // 行ビュー / 列ビュー
        VectorMap<T> row(size_type r) {
            if (r >= rows_) throw IndexError("MatrixMap::row: index out of range");
            return VectorMap<T>(data_ + r * outerStride_, cols_, 1);
        }
        ConstVectorMap<T> row(size_type r) const {
            if (r >= rows_) throw IndexError("MatrixMap::row: index out of range");
            return ConstVectorMap<T>(data_ + r * outerStride_, cols_, 1);
        }

        VectorMap<T> col(size_type c) {
            if (c >= cols_) throw IndexError("MatrixMap::col: index out of range");
            return VectorMap<T>(data_ + c, rows_, outerStride_);
        }
        ConstVectorMap<T> col(size_type c) const {
            if (c >= cols_) throw IndexError("MatrixMap::col: index out of range");
            return ConstVectorMap<T>(data_ + c, rows_, outerStride_);
        }

        // 対角ビュー
        VectorMap<T> diagonal() {
            size_type n = std::min(rows_, cols_);
            return VectorMap<T>(data_, n, outerStride_ + 1);
        }
        ConstVectorMap<T> diagonal() const {
            size_type n = std::min(rows_, cols_);
            return ConstVectorMap<T>(data_, n, outerStride_ + 1);
        }

        // ブロック (サブ行列)
        MatrixMap block(size_type startRow, size_type startCol,
                        size_type blockRows, size_type blockCols) {
            if (startRow + blockRows > rows_ || startCol + blockCols > cols_)
                throw IndexError("MatrixMap::block: range exceeds bounds");
            return MatrixMap(data_ + startRow * outerStride_ + startCol,
                             blockRows, blockCols, outerStride_);
        }

        // Matrix<T> へのコピー出力
        template<typename MatType>
        void copyTo(MatType& dest) const {
            if (dest.rows() != rows_ || dest.cols() != cols_)
                dest.resize(rows_, cols_);
            for (size_type i = 0; i < rows_; ++i)
                for (size_type j = 0; j < cols_; ++j)
                    dest(i, j) = (*this)(i, j);
        }

        // 他行列からのコピー入力
        template<typename MatType>
        void copyFrom(const MatType& src) {
            if (src.rows() != rows_ || src.cols() != cols_)
                throw DimensionError("MatrixMap::copyFrom: dimension mismatch");
            for (size_type i = 0; i < rows_; ++i)
                for (size_type j = 0; j < cols_; ++j)
                    (*this)(i, j) = src(i, j);
        }

        // 転置コピー
        template<typename MatType>
        void transposeTo(MatType& dest) const {
            if (dest.rows() != cols_ || dest.cols() != rows_)
                dest.resize(cols_, rows_);
            for (size_type i = 0; i < rows_; ++i)
                for (size_type j = 0; j < cols_; ++j)
                    dest(j, i) = (*this)(i, j);
        }

        // トレース
        T trace() const {
            if (rows_ != cols_)
                throw DimensionError("MatrixMap::trace: matrix must be square");
            T sum{};
            for (size_type i = 0; i < rows_; ++i)
                sum += (*this)(i, i);
            return sum;
        }

    private:
        void checkIndex(size_type row, size_type col) const {
            if (row >= rows_)
                throw IndexError("MatrixMap::at: row index out of range");
            if (col >= cols_)
                throw IndexError("MatrixMap::at: column index out of range");
        }

        T* data_;
        size_type rows_;
        size_type cols_;
        size_type outerStride_;  // 行間の要素数 (row-major)
    };

    // ── ConstMatrixMap<T> ─────────────────────────────────────────

    template<typename T>
    class ConstMatrixMap {
    public:
        using value_type = T;
        using size_type = std::size_t;
        using const_reference = const T&;
        using const_pointer = const T*;

        ConstMatrixMap(const T* data, size_type rows, size_type cols)
            : data_(data), rows_(rows), cols_(cols), outerStride_(cols) {}

        ConstMatrixMap(const T* data, size_type rows, size_type cols, size_type outerStride)
            : data_(data), rows_(rows), cols_(cols), outerStride_(outerStride) {
            if (outerStride < cols)
                throw std::invalid_argument("ConstMatrixMap: outerStride must be >= cols");
        }

        // MatrixMap からの暗黙変換
        ConstMatrixMap(const MatrixMap<T>& mm)
            : data_(mm.data()), rows_(mm.rows()), cols_(mm.cols()),
              outerStride_(mm.outerStride()) {}

        const_reference operator()(size_type row, size_type col) const {
            return data_[row * outerStride_ + col];
        }
        const_reference at(size_type row, size_type col) const {
            if (row >= rows_ || col >= cols_)
                throw IndexError("ConstMatrixMap::at: index out of range");
            return data_[row * outerStride_ + col];
        }

        size_type rows() const noexcept { return rows_; }
        size_type cols() const noexcept { return cols_; }
        size_type size() const noexcept { return rows_ * cols_; }
        size_type outerStride() const noexcept { return outerStride_; }
        bool empty() const noexcept { return rows_ == 0 || cols_ == 0; }
        bool isSquare() const noexcept { return rows_ == cols_; }
        const_pointer data() const noexcept { return data_; }
        bool isContiguous() const noexcept { return outerStride_ == cols_; }

        ConstVectorMap<T> row(size_type r) const {
            if (r >= rows_) throw IndexError("ConstMatrixMap::row: index out of range");
            return ConstVectorMap<T>(data_ + r * outerStride_, cols_, 1);
        }
        ConstVectorMap<T> col(size_type c) const {
            if (c >= cols_) throw IndexError("ConstMatrixMap::col: index out of range");
            return ConstVectorMap<T>(data_ + c, rows_, outerStride_);
        }
        ConstVectorMap<T> diagonal() const {
            size_type n = std::min(rows_, cols_);
            return ConstVectorMap<T>(data_, n, outerStride_ + 1);
        }

        template<typename MatType>
        void copyTo(MatType& dest) const {
            if (dest.rows() != rows_ || dest.cols() != cols_)
                dest.resize(rows_, cols_);
            for (size_type i = 0; i < rows_; ++i)
                for (size_type j = 0; j < cols_; ++j)
                    dest(i, j) = (*this)(i, j);
        }

        T trace() const {
            if (rows_ != cols_)
                throw DimensionError("ConstMatrixMap::trace: matrix must be square");
            T sum{};
            for (size_type i = 0; i < rows_; ++i)
                sum += (*this)(i, i);
            return sum;
        }

    private:
        const T* data_;
        size_type rows_;
        size_type cols_;
        size_type outerStride_;
    };

    // ── ヘルパー関数 ──────────────────────────────────────────────

    // Eigen 風のファクトリ関数
    template<typename T>
    MatrixMap<T> mapMatrix(T* data, std::size_t rows, std::size_t cols) {
        return MatrixMap<T>(data, rows, cols);
    }

    template<typename T>
    MatrixMap<T> mapMatrix(T* data, std::size_t rows, std::size_t cols, std::size_t outerStride) {
        return MatrixMap<T>(data, rows, cols, outerStride);
    }

    template<typename T>
    ConstMatrixMap<T> mapMatrix(const T* data, std::size_t rows, std::size_t cols) {
        return ConstMatrixMap<T>(data, rows, cols);
    }

    template<typename T>
    VectorMap<T> mapVector(T* data, std::size_t size) {
        return VectorMap<T>(data, size);
    }

    template<typename T>
    VectorMap<T> mapVector(T* data, std::size_t size, std::size_t stride) {
        return VectorMap<T>(data, size, stride);
    }

    template<typename T>
    ConstVectorMap<T> mapVector(const T* data, std::size_t size) {
        return ConstVectorMap<T>(data, size);
    }

} // namespace calx

#endif // CALX_MAP_HPP
