# -*- coding: utf-8 -*-

import os
import re
import numpy as np
from docx import Document
import math

# --- 用户配置区 ---

# 1. 设置包含左视图 DOCX 文件的文件夹路径
LEFT_VIEW_FOLDER = 'I:\\#132\\post132\\PSNR\\logs\\HEVC\\streetView_captured\\left'

# 2. 设置包含右视图 DOCX 文件的文件夹路径
RIGHT_VIEW_FOLDER = 'I:\\#132\\post132\\PSNR\\logs\\HEVC\\streetView_captured\\right'

# 3. 您的实验中使用的 QP 值列表 (用于生成测试数据，主逻辑会自动发现)
QP_VALUES = [17, 22, 27, 32, 37]

# 4. 设置 PSNR 计算中的最大可能信号值 (8-bit 为 255)
MAX_SIGNAL_VALUE = 255.0

# --- 脚本核心代码 ---

def psnr_to_mse(psnr, max_val=MAX_SIGNAL_VALUE):
    """将 PSNR (dB) 值转换回 MSE。"""
    if psnr == float('inf') or psnr >= 100: # Handle perfect or near-perfect PSNR
        return 0.0
    return (max_val ** 2) / (10 ** (psnr / 10))

def mse_to_psnr(mse, max_val=MAX_SIGNAL_VALUE):
    """将 MSE 值转换为 PSNR (dB)。"""
    if mse <= 1e-10: # Handle zero or very small MSE
        return float('inf')
    return 10 * math.log10((max_val ** 2) / mse)

def extract_hm_data_from_docx(file_path):
    """
    从 HM 编码器日志的 DOCX 文件中提取 Bitrate 和 SUMMARY 中的 Y-PSNR。
    """
    try:
        document = Document(file_path)
        full_text = "\n".join([para.text for para in document.paragraphs])
        
        # 提取 Bitrate (保持原有逻辑)
        bitrate_match = re.search(r"Bytes written to file:.*?\((\d+\.?\d*)\s*kbps\)", full_text)
        if not bitrate_match:
            bitrate_match = re.search(r"SUMMARY.*?\n\s*\d+\s+a\s+(\d+\.\d+)", full_text, re.DOTALL)
        
        if not bitrate_match:
            print(f"警告: 在文件 '{os.path.basename(file_path)}' 中未找到 Bitrate。")
            return None
        bitrate = float(bitrate_match.group(1))

        # --- 【核心修改】 ---
        # 使用新的正则表达式提取 SUMMARY 部分的 Y-PSNR
        # 正则表达式解释:
        # SUMMARY.*?         - 匹配 "SUMMARY" 和之后任意字符 (非贪婪模式)
        # \n\s*\d+\s+a\s+    - 匹配换行符、可选空格、帧数、空格、字母'a'、空格
        # \S+\s+             - 匹配码率 (非空格字符) 和后面的空格
        # ([\d\.]+)          - 捕获我们想要的PSNR值 (数字和点号)
        summary_psnr_match = re.search(
            r"SUMMARY.*?\n\s*\d+\s+a\s+\S+\s+([\d\.]+)", 
            full_text, 
            re.DOTALL
        )
        
        if not summary_psnr_match:
            print(f"警告: 在文件 '{os.path.basename(file_path)}' 的 SUMMARY 部分未找到 Y-PSNR。")
            return None
        
        summary_psnr = float(summary_psnr_match.group(1))
        
        # 返回包含码率和单个summary_psnr的字典
        return {'bitrate': bitrate, 'summary_psnr': summary_psnr}
    
    except Exception as e:
        print(f"错误: 无法读取或解析文件 '{os.path.basename(file_path)}'。错误信息: {e}")
        return None

def find_and_group_files_by_sequence(left_folder, right_folder):
    """
    扫描文件夹，根据 'basename_view_QP.docx' 格式找到文件对，并按 'basename' 分组。
    返回一个字典，键是序列名，值是包含各QP点信息的列表。
    e.g., {'dog_captured': [{'qp': 17, 'left_path': ..., 'right_path': ...}, ...]}
    """
    sequences = {}
    print(f"正在扫描左视图文件夹: {left_folder}")
    try:
        left_files = os.listdir(left_folder)
    except FileNotFoundError:
        print(f"错误: 左视图文件夹 '{left_folder}' 不存在。")
        return {}

    left_file_pattern = re.compile(r"^(.*?)_left_(\d+)\.docx$")

    for left_filename in left_files:
        match = left_file_pattern.match(left_filename)
        if match:
            base_name, qp_str = match.groups()
            qp = int(qp_str)
            
            right_filename = f"{base_name}_right_{qp}.docx"
            right_file_path = os.path.join(right_folder, right_filename)
            
            if os.path.exists(right_file_path):
                if base_name not in sequences:
                    sequences[base_name] = []
                
                sequences[base_name].append({
                    'qp': qp,
                    'left_path': os.path.join(left_folder, left_filename),
                    'right_path': right_file_path
                })
            else:
                print(f"警告: 找到 '{left_filename}'，但未找到对应右视图 '{right_filename}'。")
    
    # 对每个序列内部的QP点进行排序
    for base_name in sequences:
        sequences[base_name].sort(key=lambda x: x['qp'])
        
    return sequences

def main():
    """主执行函数"""
    print("--- 开始处理 HM 编码器日志 ---")
    
    sequences = find_and_group_files_by_sequence(LEFT_VIEW_FOLDER, RIGHT_VIEW_FOLDER)

    if not sequences:
        print("\n错误: 未找到任何匹配的序列文件。请检查文件夹路径和文件名格式。")
        return

    print(f"\n成功找到 {len(sequences)} 个序列，开始计算...\n")

    grand_total_bitrate = 0.0
    grand_all_avg_mses = [] # 存储每个QP对的平均MSE

    # --- 循环处理每个序列 ---
    for base_name, qp_data_list in sequences.items():
        print("="*80)
        print(f"序列: {base_name}")
        print("="*80)
        print(f"{'QP':<5} {'Total Bitrate (kbps)':<25} {'Combined PSNR (dB)':<25}")
        print(f"{'-'*5} {'-'*25} {'-'*25}")

        sequence_total_bitrate = 0.0
        sequence_all_avg_mses = [] # 存储该序列每个QP对的平均MSE
        
        # --- 循环处理该序列的每个QP点 ---
        for qp_data in qp_data_list:
            left_data = extract_hm_data_from_docx(qp_data['left_path'])
            right_data = extract_hm_data_from_docx(qp_data['right_path'])

            if not left_data or not right_data:
                print(f"{qp_data['qp']:<5} {'数据提取失败':<51}")
                continue

            # --- 【核心修改】 ---
            # 计算当前QP对的码率和组合PSNR
            pair_bitrate = left_data['bitrate'] + right_data['bitrate']
            
            # 1. 从左右视图数据中获取 summary_psnr
            left_summary_psnr = left_data['summary_psnr']
            right_summary_psnr = right_data['summary_psnr']
            
            # 2. 将PSNR转换为MSE
            left_mse = psnr_to_mse(left_summary_psnr)
            right_mse = psnr_to_mse(right_summary_psnr)
            
            # 3. 计算两个MSE的平均值
            pair_avg_mse = (left_mse + right_mse) / 2.0
            
            # 4. 将平均MSE转换回PSNR，得到组合PSNR
            pair_psnr = mse_to_psnr(pair_avg_mse)

            # 打印当前QP对的结果
            print(f"{qp_data['qp']:<5} {pair_bitrate:<25.2f} {pair_psnr:<25.4f}")

            # 累加到序列总和
            sequence_total_bitrate += pair_bitrate
            sequence_all_avg_mses.append(pair_avg_mse)
        
        # --- 计算并打印当前序列的汇总信息 ---
        if sequence_all_avg_mses:
            # 计算所有QP点的平均MSE的平均值
            sequence_avg_mse = np.mean(sequence_all_avg_mses)
            sequence_avg_psnr = mse_to_psnr(sequence_avg_mse)
            print("\n--- 序列汇总 ---\n")
            print(f"  总码率 (所有QP点): {sequence_total_bitrate:.2f} kbps")
            print(f"  平均PSNR (所有QP点): {sequence_avg_psnr:.4f} dB\n")

            # 累加到全局总和
            grand_total_bitrate += sequence_total_bitrate
            grand_all_avg_mses.extend(sequence_all_avg_mses)

    # --- 生成最终的全局报告 ---
    print("\n" + "="*80)
    print(f"{'全局最终汇总报告':^80}")
    print("="*80)
    
    if not grand_all_avg_mses:
        print("没有成功处理任何文件。")
        return
        
    grand_avg_mse = np.mean(grand_all_avg_mses)
    grand_avg_psnr = mse_to_psnr(grand_avg_mse)

    print(f"处理的总序列数: {len(sequences)}")
    print(f"全局总码率 (所有序列, 所有QP点): {grand_total_bitrate:.2f} kbps")
    print(f"全局平均PSNR (基于所有QP对的平均MSE): {grand_avg_psnr:.4f} dB")
    
    print("\n" + "="*80)
    print("处理完成。")


if __name__ == "__main__":
    # 自动创建包含所有指定QP值的示例文件
    # (如果文件夹不存在，此部分将提示用户，主逻辑会处理错误)
    if not os.path.exists(LEFT_VIEW_FOLDER) or not os.path.exists(RIGHT_VIEW_FOLDER):
        print(f"警告: 文件夹 '{LEFT_VIEW_FOLDER}' 或 '{RIGHT_VIEW_FOLDER}' 不存在。")
        print("请确保路径正确。脚本将继续尝试执行，但可能会失败。")

    main()