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

// jpeg.hpp — Baseline JPEG エンコーダ/デコーダ
//
// 外部ライブラリ不要の自前実装。
// Baseline JPEG: 8×8 DCT, Huffman 符号化, YCbCr 色空間。
// 圧縮品質は 1-100 で指定。

#ifndef CALX_IMAGE_IO_JPEG_HPP
#define CALX_IMAGE_IO_JPEG_HPP

#include <image/Image.hpp>
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <vector>

namespace calx::image {

    namespace detail_jpeg {

        // 標準輝度量子化テーブル (品質 50)
        constexpr int lumaQuant50[64] = {
            16,11,10,16,24,40,51,61, 12,12,14,19,26,58,60,55,
            14,13,16,24,40,57,69,56, 14,17,22,29,51,87,80,62,
            18,22,37,56,68,109,103,77, 24,35,55,64,81,104,113,92,
            49,64,78,87,103,121,120,101, 72,92,95,98,112,100,103,99
        };

        // ジグザグ順序
        constexpr int zigzag[64] = {
             0, 1, 8,16, 9, 2, 3,10,
            17,24,32,25,18,11, 4, 5,
            12,19,26,33,40,48,41,34,
            27,20,13, 6, 7,14,21,28,
            35,42,49,56,57,50,43,36,
            29,22,15,23,30,37,44,51,
            58,59,52,45,38,31,39,46,
            53,60,61,54,47,55,62,63
        };

        // 品質から量子化テーブルを計算
        inline std::array<int, 64> makeQuantTable(int quality) {
            quality = std::clamp(quality, 1, 100);
            int scale = quality < 50 ? (5000 / quality) : (200 - 2 * quality);
            std::array<int, 64> table;
            for (int i = 0; i < 64; ++i) {
                int val = (lumaQuant50[i] * scale + 50) / 100;
                table[i] = std::clamp(val, 1, 255);
            }
            return table;
        }

        // 8×8 DCT (Type-II)
        inline void dct8x8(double block[64]) {
            double tmp[64];
            // 行方向
            for (int i = 0; i < 8; ++i) {
                for (int u = 0; u < 8; ++u) {
                    double sum = 0;
                    for (int x = 0; x < 8; ++x)
                        sum += block[i*8+x] * std::cos((2*x+1)*u*3.14159265358979/(16.0));
                    double cu = (u == 0) ? 1.0/std::sqrt(2.0) : 1.0;
                    tmp[i*8+u] = 0.5 * cu * sum;
                }
            }
            // 列方向
            for (int u = 0; u < 8; ++u) {
                for (int v = 0; v < 8; ++v) {
                    double sum = 0;
                    for (int y = 0; y < 8; ++y)
                        sum += tmp[y*8+u] * std::cos((2*y+1)*v*3.14159265358979/(16.0));
                    double cv = (v == 0) ? 1.0/std::sqrt(2.0) : 1.0;
                    block[v*8+u] = 0.5 * cv * sum;
                }
            }
        }

        // 8×8 IDCT (Type-III)
        inline void idct8x8(double block[64]) {
            double tmp[64];
            for (int v = 0; v < 8; ++v) {
                for (int y = 0; y < 8; ++y) {
                    double sum = 0;
                    for (int u = 0; u < 8; ++u) {
                        double cu = (u == 0) ? 1.0/std::sqrt(2.0) : 1.0;
                        sum += cu * block[v*8+u] * std::cos((2*y+1)*u*3.14159265358979/(16.0));
                    }
                    tmp[v*8+y] = 0.5 * sum;
                }
            }
            for (int x = 0; x < 8; ++x) {
                for (int y = 0; y < 8; ++y) {
                    double sum = 0;
                    for (int v = 0; v < 8; ++v) {
                        double cv = (v == 0) ? 1.0/std::sqrt(2.0) : 1.0;
                        sum += cv * tmp[v*8+x] * std::cos((2*y+1)*v*3.14159265358979/(16.0));
                    }
                    block[y*8+x] = 0.5 * sum;
                }
            }
        }

    } // namespace detail_jpeg

    /// JPEG エンコード → DCT 係数列 (デバッグ・テスト用)
    /// 実際のファイル出力は Huffman 符号化が必要で大規模実装になるため、
    /// ここでは DCT + 量子化のコア部分のみ提供。
    template<typename T>
    struct JpegBlock {
        std::array<int, 64> coefficients;  // 量子化後の DCT 係数 (ジグザグ順)
    };

    /// 画像を 8×8 ブロックに分割して DCT + 量子化
    template<typename T>
    std::vector<JpegBlock<T>> jpegEncode(const Image<T>& img, int quality = 75) {
        assert(img.numChannels() == 1);  // グレースケールのみ
        size_t H = img.height(), W = img.width();
        auto qtable = detail_jpeg::makeQuantTable(quality);

        size_t bH = (H + 7) / 8, bW = (W + 7) / 8;
        std::vector<JpegBlock<T>> blocks(bH * bW);

        for (size_t by = 0; by < bH; ++by) {
            for (size_t bx = 0; bx < bW; ++bx) {
                double block[64];
                for (int dy = 0; dy < 8; ++dy)
                    for (int dx = 0; dx < 8; ++dx) {
                        size_t y = by * 8 + dy, x = bx * 8 + dx;
                        if (y < H && x < W)
                            block[dy*8+dx] = static_cast<double>(img.at(y, x, 0)) * 255.0 - 128.0;
                        else
                            block[dy*8+dx] = 0;
                    }

                detail_jpeg::dct8x8(block);

                auto& jb = blocks[by * bW + bx];
                for (int i = 0; i < 64; ++i) {
                    int zi = detail_jpeg::zigzag[i];
                    jb.coefficients[i] = static_cast<int>(std::round(block[zi] / qtable[zi]));
                }
            }
        }
        return blocks;
    }

    /// DCT 係数からグレースケール画像を復元
    template<typename T>
    Image<T> jpegDecode(const std::vector<JpegBlock<T>>& blocks,
                         size_t H, size_t W, int quality = 75)
    {
        auto qtable = detail_jpeg::makeQuantTable(quality);
        size_t bH = (H + 7) / 8, bW = (W + 7) / 8;
        Image<T> img(H, W, 1);

        for (size_t by = 0; by < bH; ++by) {
            for (size_t bx = 0; bx < bW; ++bx) {
                auto& jb = blocks[by * bW + bx];
                double block[64] = {};

                // 逆量子化 + 逆ジグザグ
                for (int i = 0; i < 64; ++i) {
                    int zi = detail_jpeg::zigzag[i];
                    block[zi] = static_cast<double>(jb.coefficients[i]) * qtable[zi];
                }

                detail_jpeg::idct8x8(block);

                for (int dy = 0; dy < 8; ++dy)
                    for (int dx = 0; dx < 8; ++dx) {
                        size_t y = by * 8 + dy, x = bx * 8 + dx;
                        if (y < H && x < W) {
                            double val = (block[dy*8+dx] + 128.0) / 255.0;
                            img.at(y, x, 0) = static_cast<T>(std::clamp(val, 0.0, 1.0));
                        }
                    }
            }
        }
        return img;
    }

} // namespace calx::image

#endif // CALX_IMAGE_IO_JPEG_HPP
