clear; clc; % 清除所有变量，防止 length 等变量名被占用

%% 1. 读取图像
img_rgb = imread('test.jpg'); 
% 扔掉第四个通道
if size(img_rgb, 3) == 4, img_rgb = img_rgb(:,:,1:3); end
img = double(img_rgb); % 转为 double 进行计算，但保持整数值

%% 2. RCT 可逆颜色变换 (对应图 10.1 左下角)
% 提取 R, G, B 分量
R = img(:,:,1);
G = img(:,:,2);
B = img(:,:,3);

% 执行 RCT 公式 裁切出分量 这里已经完成分解步骤 P293
Y_prime = floor((R + 2*G + B) / 4);
Db = B - G;
Dr = R - G;

%% 3. 分块 
% Tiling 返回各分量独立的结构体数组，三个分量坐标相同
Y_tiles  = Tiling(Y_prime);
Db_tiles = Tiling(Db);
Dr_tiles = Tiling(Dr);

num_tiles = length(Y_tiles);

% 统一的 tiles 结构体数组
tiles = struct();
for i = 1:num_tiles
    % 坐标元信息（三个分量相同，只存一次）
    tiles(i).x0     = Y_tiles(i).x0;
    tiles(i).y0     = Y_tiles(i).y0;
    tiles(i).width   = Y_tiles(i).width;
    tiles(i).height  = Y_tiles(i).height;

    % 各分隔块的分量的原始 tile 数据
    tiles(i).Y.raw   = Y_tiles(i).data;
    tiles(i).Db.raw  = Db_tiles(i).data;
    tiles(i).Dr.raw  = Dr_tiles(i).data;
end

%% 4. 5/3 DWT 小波变换 + 构建三层金字塔
num_levels = 3; % 定义金字塔层数

for i = 1:num_tiles
    % 遍历三个分量 (Y, Db, Dr)
    comp_names = {'Y', 'Db', 'Dr'};
    
    for c = 1 : 3
        comp = comp_names{c};
        
        % 1. 初始化：当前要处理的数据是原始 Tile 数据
        current_data = tiles(i).(comp).raw;
        
        % 2. 创建当前分量的层级容器
        tiles(i).(comp).levels = struct();

        % 3. 循环 3 次，构建 3 级金字塔
        for lvl = 1 : num_levels
            % 进行 5/3 DWT 变换，此处获得的是分量切割成四个子带，但是这里是交织起来的
            interleaved = dwt53_inplace(current_data);
            
            % 把交织子带提取出四个子带
            LL = interleaved(1:2:end, 1:2:end);
            HL = interleaved(1:2:end, 2:2:end);
            LH = interleaved(2:2:end, 1:2:end);
            HH = interleaved(2:2:end, 2:2:end);
            
            % 把高频子带存入当前层级 (Level 1 是最外层，Level 3 是最内层)
            tiles(i).(comp).levels(lvl).HL = HL;
            tiles(i).(comp).levels(lvl).LH = LH;
            tiles(i).(comp).levels(lvl).HH = HH;
            
            % 把新的 LL 作为下一次循环的输入
            current_data = LL;
        end
        
        % 4. 循环结束后，current_data 就是最顶层的塔尖 (LL3)
        tiles(i).(comp).levels(num_levels).LL = current_data;
    end
end

%% 5. 子带划分
comp_names = {'Y', 'Db', 'Dr'};
% 对分量循环取出
for c = 1 : 3
    comp = comp_names{c};
    %对分割块循环取出
    for i = 1:num_tiles
        %对每个分割块的子带取出
        for lvl = 1:num_levels
            % 对每一层的 HL, LH, HH 切块
            tiles(i).(comp).levels(lvl).HL = Child_Bound_Tiling(tiles(i).(comp).levels(lvl).HL);
            tiles(i).(comp).levels(lvl).LH = Child_Bound_Tiling(tiles(i).(comp).levels(lvl).LH);
            tiles(i).(comp).levels(lvl).HH = Child_Bound_Tiling(tiles(i).(comp).levels(lvl).HH);
        end
            % 对最顶层的塔尖 LL 切块
            tiles(i).(comp).levels(num_levels).LL = Child_Bound_Tiling(tiles(i).(comp).levels(num_levels).LL);
    end
end

%% Tier-1: 嵌入式块编码 (EBCOT)
% 对每个 code-block 执行位平面编码 + MQ 算术编码
tiles = embedded_block_encoder(tiles);

function all_tiles = Tiling (img_in)
    
    [H, W] = size(img_in);              % 获取图像高度(H)和宽度(W)
    
    %% 在标准中 画布的原点、图像原点、拼接块原点是可以不对齐的 这里为了方便 直接对齐
    Xsiz = W;       Ysiz = H;           % 画布大小
    XTsiz = 512;    YTsiz = 512;        % Tile 的宽和高 
    XOsiz = 0;      YOsiz = 0;          % 图像原点偏移 
    XTOsiz = 0;     YTOsiz = 0;         % Tile 的起始坐标
    
    %% 3. 计算 Tile 的总数 (行数和列数)
    % ceil 向上取整，确保覆盖整个图像
    Num_Tiles_X = ceil((Xsiz - XTOsiz) / XTsiz);
    Num_Tiles_Y = ceil((Ysiz - YTOsiz) / YTsiz);
    
    fprintf('图像大小: %d x %d\n', W, H);
    fprintf('Tile 大小: %d x %d\n', XTsiz, YTsiz);
    fprintf('划分网格: %d (列) x %d (行)\n', Num_Tiles_X, Num_Tiles_Y);
    
    %% 4. 开始遍历并切分矩阵，创建结构体数组传出 
    tile_count = 0;
    all_tiles = struct();  %创建结构体数组，用来存放所有的分割块
    for ty = 0 : Num_Tiles_Y - 1      % 垂直方向遍历 (对应矩阵的行)
        for tx = 0 : Num_Tiles_X - 1  % 水平方向遍历 (对应矩阵的列)
            % Tile理论起始点
            x_start = XTOsiz + tx * XTsiz;
            y_start = YTOsiz + ty * YTsiz;
            % Tile理论终止点
            x_end_theoretical = x_start + XTsiz;
            y_end_theoretical = y_start + YTsiz;
            % 画布分割的终止点，防止越界
            x_end = min(x_end_theoretical, Xsiz);
            y_end = min(y_end_theoretical, Ysiz);
            % 矩阵分块
            tile_matrix = img_in(y_start + 1 : y_end, x_start + 1 : x_end);
            tile_count = tile_count + 1;
            all_tiles(tile_count).data = tile_matrix;      % 存图像数据(应该是用不到，暂时先存储)
            all_tiles(tile_count).x0 = x_start;            % 存左上角 X 坐标
            all_tiles(tile_count).y0 = y_start;            % 存左上角 Y 坐标
            all_tiles(tile_count).width = x_end - x_start; % 存宽度
            all_tiles(tile_count).height = y_end - y_start;% 存高度
        end
    end
end

% 原位计算的 5/3 整数小波变换
% 输入: img_in - 二维矩阵 (RCT变换后的数据)
% 输出: img_out - 交织排列的四个子带 (LL, HL, LH, HH)
function img_out = dwt53_inplace(img_in)

    img_out = double(img_in); % 转为 double 防止加减法溢出
    [rows, cols] = size(img_out); %行：rows 列：cols（亮度......

    %% 1. 行变换 (Horizontal Transform)
    for r = 1:rows
        % === 第一步：预测 (Predict) ===
        % 用相邻偶数预测奇数，计算行高频细节
        row = img_out(r, :);
        if cols >= 3
            row(2:2:end-1) = row(2:2:end-1) - floor((row(1:2:end-2) + row(3:2:end)) / 2);
        end
        if mod(cols, 2) == 0
             % 简单边界处理：直接减去左邻居 
             row(end) = row(end) - row(end-1);
        end
        % === 第二步：更新 (Update) ===
        % 用刚刚算出的行高频，去修正偶数点，计算行低频轮廓
        if cols >= 3
            row(3:2:end-1) = row(3:2:end-1) + floor((row(2:2:end-2) + row(4:2:end)) / 4);
        end
        row(1) = row(1) + floor(row(2) / 2); 
        img_out(r, :) = row;
    end

    %% 2. 列变换 (Vertical Transform)
    for c = 1:cols
        % === 第一步：预测 (Predict) ===
        % 用相邻偶数预测奇数，计算行高频细节
        col = img_out(:, c);
        if rows >= 3
            col(2:2:end-1) = col(2:2:end-1) - floor((col(1:2:end-2) + col(3:2:end)) / 2);
        end
        if mod(rows, 2) == 0
             % 简单边界处理：直接减去左邻居 
             col(end) = col(end) - col(end-1);
        end
        % === 第二步：更新 (Update) ===
        % 用刚刚算出的行高频，去修正偶数点，计算行低频轮廓
        if rows >= 3
            col(3:2:end-1) = col(3:2:end-1) + floor((col(2:2:end-2) + col(4:2:end)) / 4);
        end
        col(1) = col(1) + floor(col(2) / 2); 
        img_out(:, c) = col; 
    end

end

function Bound_Tiling_out = Child_Bound_Tiling(subband_matrix)
% CHILD_BOUND_TILING 对子带矩阵进行两级划分：Precinct  Code-block
%
% 第一级：用 Precinct 尺寸 (PPx × PPy) 把子带切成 Precinct 分区
% 第二级：用 Code-block 尺寸 (XJsiz × YJsiz) 把每个 Precinct 切成编码块
%
% 输出结构:
%   Bound_Tiling_out.precincts(py, px).blocks(by, bx).data
%   Bound_Tiling_out.precincts(py, px).blocks(by, bx).row_start / row_end / col_start / col_end

%% 1. 获取子带的实际尺寸
    [rows, cols] = size(subband_matrix);

%% 2. 定义划分参数
    % --- Precinct 分区参数 ---
    PPx = 128;          % Precinct 宽度 (必须是 2 的幂)
    PPy = 128;          % Precinct 高度 (必须是 2 的幂)

    % --- Code-block 编码块参数 ---
    XJsiz = 32;         % 编码块宽度 (必须是 2 的幂)
    YJsiz = 32;         % 编码块高度 (必须是 2 的幂)

%% 3. 第一级：计算 Precinct 数量并遍历
    num_prec_y = ceil(rows / PPy);
    num_prec_x = ceil(cols / PPx);

    precincts = struct();

    for py = 0 : num_prec_y - 1
        for px = 0 : num_prec_x - 1
            % 计算当前 Precinct 在子带中的行列范围
            pr_start = py * PPy + 1;
            pr_end   = min((py + 1) * PPy, rows);
            pc_start = px * PPx + 1;
            pc_end   = min((px + 1) * PPx, cols);

            % 提取当前 Precinct 的子矩阵
            prec_matrix = subband_matrix(pr_start:pr_end, pc_start:pc_end);
            prec_rows = pr_end - pr_start + 1;
            prec_cols = pc_end - pc_start + 1;

            %% 4. 第二级：在 Precinct 内按 Code-block 切分
            num_cb_y = ceil(prec_rows / YJsiz);
            num_cb_x = ceil(prec_cols / XJsiz);

            blocks = struct();
            block_idx = 1;

            for by = 0 : num_cb_y - 1
                for bx = 0 : num_cb_x - 1
                    % 计算当前 Code-block 在 Precinct 内的行列范围
                    cb_r_start = by * YJsiz + 1;
                    cb_r_end   = min((by + 1) * YJsiz, prec_rows);
                    cb_c_start = bx * XJsiz + 1;
                    cb_c_end   = min((bx + 1) * XJsiz, prec_cols);

                    % 提取 Code-block 数据
                    cb_data = prec_matrix(cb_r_start:cb_r_end, cb_c_start:cb_c_end);

                    % 存入结构体（坐标相对于子带原点）
                    blocks(block_idx).row_start = pr_start + cb_r_start - 1;
                    blocks(block_idx).row_end   = pr_start + cb_r_end - 1;
                    blocks(block_idx).col_start = pc_start + cb_c_start - 1;
                    blocks(block_idx).col_end   = pc_start + cb_c_end - 1;
                    blocks(block_idx).data      = cb_data;

                    block_idx = block_idx + 1;
                end
            end

            % 存储当前 Precinct 的信息
            precincts(py + 1, px + 1).row_start  = pr_start;
            precincts(py + 1, px + 1).row_end    = pr_end;
            precincts(py + 1, px + 1).col_start  = pc_start;
            precincts(py + 1, px + 1).col_end    = pc_end;
            precincts(py + 1, px + 1).num_blocks_y = num_cb_y;
            precincts(py + 1, px + 1).num_blocks_x = num_cb_x;
            precincts(py + 1, px + 1).total_blocks = num_cb_y * num_cb_x;
            precincts(py + 1, px + 1).blocks       = blocks;
        end
    end

%% 5. 输出汇总
    Bound_Tiling_out.subband_size   = [rows, cols];
    Bound_Tiling_out.precinct_size  = [PPy, PPx];
    Bound_Tiling_out.block_size     = [YJsiz, XJsiz];
    Bound_Tiling_out.num_prec_y     = num_prec_y;
    Bound_Tiling_out.num_prec_x     = num_prec_x;
    Bound_Tiling_out.total_precincts = num_prec_y * num_prec_x;
    Bound_Tiling_out.precincts      = precincts;

end

function tiles = embedded_block_encoder(tiles)
%EMBEDDED_BLOCK_ENCODER JPEG2000 Tier-1 嵌入式块编码器
%
%   遍历 tiles 结构体中每个 code-block，执行完整的 EBCOT Tier-1 编码：
%     1. BitPlaneInit  - 初始化位平面分解和状态矩阵
%     2. 对每个位平面 p = P:-1:1 执行三遍扫描:
%        a. 显著性传播扫描 (SignificancePropagationPass)
%        b. 幅度细化扫描   (MagnitudeRefinementPass)
%        c. 清理扫描       (CleanUpPass)
%     3. MQ 编码器将 (D, CX) 对编码为压缩字节流
%
%   输入:
%       tiles - 结构体数组，由 main.m 构建，包含:
%               tiles(i).(comp).levels(lvl).(subband).precincts(py,px).blocks(bi).data
%
%   输出:
%       tiles - 同一结构体数组，在每个 block 上追加:
%               .bitstream  - MQ 编码后的压缩字节流 (uint8 数组)
%               .num_passes - 编码的位平面数

    component_names = {'Y', 'Db', 'Dr'};
    num_tiles = length(tiles);

    %% 初始化 MQ 编码器查找表 (只初始化一次，所有 code-block 共享)
    [CX_Table_init, PET_Table] = mq_inits;

    for tile_index = 1:num_tiles
        for comp_idx = 1:length(component_names)
            comp = component_names{comp_idx};
            levels = tiles(tile_index).(comp).levels;
            num_levels = length(levels);

            for lvl = 1:num_levels
                % 确定当前层级的子带列表
                if lvl == num_levels
                    subband_names = {'LL', 'HL', 'LH', 'HH'};
                else
                    subband_names = {'HL', 'LH', 'HH'};
                end

                for sb_idx = 1:length(subband_names)
                    sb_name = subband_names{sb_idx};
                    subband = levels(lvl).(sb_name);
                    precincts = subband.precincts;

                    for py = 1:size(precincts, 1)
                        for px = 1:size(precincts, 2)
                            blocks = precincts(py, px).blocks;

                            for bi = 1:length(blocks)
                                data = blocks(bi).data;

                                % ===== 编码单个 code-block =====
                                [bitstream, num_passes] = encode_single_codeblock(...
                                    data, sb_name, CX_Table_init, PET_Table);

                                % 写回结果
                                tiles(tile_index).(comp).levels(lvl).(sb_name) ...
                                    .precincts(py, px).blocks(bi).bitstream = bitstream;
                                tiles(tile_index).(comp).levels(lvl).(sb_name) ...
                                    .precincts(py, px).blocks(bi).num_passes = num_passes;

                                fprintf('  Tile %d, %s, level %d, %s, precinct(%d,%d), block %d: %d bytes\n', ...
                                    tile_index, comp, lvl, sb_name, py, px, bi, length(bitstream));
                            end
                        end
                    end
                end
            end
        end
    end
end


function [bitstream, P] = encode_single_codeblock(data, band, CX_Table_init, PET_Table)
%ENCODE_SINGLE_CODEBLOCK 对单个 code-block 执行完整的 Tier-1 编码
%
%   输入:
%       data            - code-block 数据矩阵 (height × width)
%       band            - 子带名称 ('LL','HL','LH','HH')
%       CX_Table_init   - CX 表初始值
%       PET_Table       - 概率估计表
%
%   输出:
%       bitstream - 压缩字节流 (uint8 数组)
%       P         - 位平面数

    %% 1. 位平面初始化
    [Sign_Array, S1, S2, S3, V, P] = BitPlaneInit(data);

    if P == 0
        % code-block 全为零，输出空字节流
        bitstream = uint8([]);
        return
    end

    [height, width] = size(data);

    % 预分配 (D, CX) 数组 (估计最大长度)
    max_pairs = height * width * P * 4;
    D = uint8(zeros(max_pairs, 1));
    CX = uint8(zeros(max_pairs, 1));
    enc_ptr = 1;

    %% 2. 对每个位平面执行三遍扫描 (从 MSB 到 LSB)
    for bp = 1:P
        % 2a. 显著性传播扫描
        [D, CX, S1, S2, S3, enc_ptr] = SignificancePropagationPass(...
            D, CX, V, bp, S1, S2, S3, Sign_Array, enc_ptr, band);

        % 2b. 幅度细化扫描
        [D, CX, S1, S2, S3, enc_ptr] = MagnitudeRefinementPass(...
            D, CX, V, bp, S1, S2, S3, Sign_Array, enc_ptr, band);

        % 2c. 清理扫描
        [D, CX, enc_ptr, S1, S2, S3] = CleanUpPass(...
            CX, D, V, bp, S1, S2, S3, Sign_Array, enc_ptr, band);

        % 重置 S3 (每个位平面编码完成后清零)
        S3 = zeros(height, width);
    end

    % 裁剪到实际长度
    num_pairs = enc_ptr - 1;
    D  = D(1:num_pairs);
    CX = CX(1:num_pairs);

    %% 3. MQ 编码
    CX_Table = CX_Table_init;  % 每个 code-block 使用独立的 CX 表副本
    reg_enc = mq_encoder_init;

    % CX 值加 1 偏移 (MATLAB 1-based 索引)
    CX = CX + 1;

    for i = 1:num_pairs
        [reg_enc, CX_Table] = mq_encoder(D(i), CX(i), reg_enc, CX_Table, PET_Table);
    end

    %% 4. 终止编码器，获取字节流
    bitstream = mq_encoder_end(reg_enc);

end