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

import os
import re
import math
import pandas as pd

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

# 1. 设置包含 MV-HEVC 日志 (.log) 文件的文件夹路径
LOG_FOLDER = 'I:\\#132\\post132\\PSNR\\logs\\MV-HEVC'  # <--- 请务必修改为您的文件夹路径

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


# --- 核心功能代码 ---

def psnr_to_mse(psnr, max_val=MAX_SIGNAL_VALUE):
    """将 PSNR (dB) 值转换回 MSE。"""
    if psnr >= 100:
        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: # 处理零或极小的MSE，避免log(inf)错误
        return float('inf')
    return 10 * math.log10((max_val ** 2) / mse)

def extract_data_from_log(file_path):
    """
    从 MV-HEVC 日志的 .log 文件中提取 Layer 0/1 的 Y-PSNR 和总码率。
    """
    try:
        # 使用标准文件读取方式，不再需要 python-docx
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            full_text = f.read()

        # --- 正则表达式定义 (与之前相同) ---
        psnr_pattern_template = r"SUMMARY.*?LayerId\s+{layer_id}.*?\n\s*\d+\s+a\s+\S+\s+([\d\.]+)"
        total_bitrate_pattern = re.compile(r"Bytes written to file:.*?\((\d+\.?\d*)\s*kbps\)")

        # --- 数据提取 ---
        match_l0 = re.search(psnr_pattern_template.format(layer_id=0), full_text, re.DOTALL)
        match_l1 = re.search(psnr_pattern_template.format(layer_id=1), full_text, re.DOTALL)
        match_bitrate = total_bitrate_pattern.search(full_text)

        # --- 检查所有数据是否都已找到 ---
        if not match_l0 or not match_l1 or not match_bitrate:
            missing = []
            if not match_l0: missing.append("Layer 0 Summary Y-PSNR")
            if not match_l1: missing.append("Layer 1 Summary Y-PSNR")
            if not match_bitrate: missing.append("Total Bitrate line")
            print(f"警告: 在文件 '{os.path.basename(file_path)}' 中未找到: {', '.join(missing)}")
            return None

        # --- 返回提取到的所有数据 ---
        return {
            'psnr_l0': float(match_l0.group(1)),
            'psnr_l1': float(match_l1.group(1)),
            'total_bitrate_kbps': float(match_bitrate.group(1))
        }

    except Exception as e:
        print(f"错误: 无法读取或解析文件 '{os.path.basename(file_path)}'。错误: {e}")
        return None

def main():
    """主执行函数"""
    print("--- 开始处理 MV-HEVC .log 文件 ---")

    if not os.path.isdir(LOG_FOLDER):
        print(f"错误: 文件夹 '{LOG_FOLDER}' 不存在。请检查路径配置。")
        return

    results = []
    # 用于从文件名 (如 '..._QP17.log') 中提取QP值的正则表达式
    qp_pattern = re.compile(r'_QP(\d+)\.log$', re.IGNORECASE)

    print(f"\n正在扫描文件夹: {LOG_FOLDER}")
    # 按文件名排序后遍历
    for filename in sorted(os.listdir(LOG_FOLDER)):
        # 检查文件后缀是否为 .log
        if filename.lower().endswith(".log"):
            file_path = os.path.join(LOG_FOLDER, filename)
            
            print(f"处理文件: {filename}")
            
            # 提取QP值
            qp_match = qp_pattern.search(filename)
            qp = int(qp_match.group(1)) if qp_match else -1

            # 从文件中提取数据
            data = extract_data_from_log(file_path)

            if data:
                # 提取各层PSNR和总码率
                psnr_l0 = data['psnr_l0']
                psnr_l1 = data['psnr_l1']
                total_bitrate = data['total_bitrate_kbps']
                
                # --- 【核心计算】: 通过MSE方式计算平均PSNR ---
                # 1. 将各层PSNR转换为MSE
                mse_l0 = psnr_to_mse(psnr_l0)
                mse_l1 = psnr_to_mse(psnr_l1)
                
                # 2. 计算MSE的平均值
                avg_mse = (mse_l0 + mse_l1) / 2.0
                
                # 3. 将平均MSE转换回PSNR
                avg_psnr = mse_to_psnr(avg_mse)
                
                # 将所有结果存入列表
                results.append({
                    'Filename': filename,
                    'QP': qp,
                    'Total Bitrate (kbps)': total_bitrate,
                    'Avg Y-PSNR (dB)': avg_psnr,
                    'Y-PSNR L0 (dB)': psnr_l0, # 保留原始值以供参考
                    'Y-PSNR L1 (dB)': psnr_l1  # 保留原始值以供参考
                })

    if not results:
        print("\n错误: 在指定文件夹中没有找到任何有效的.log日志文件或未能成功提取数据。")
        return

    # --- 使用 pandas 创建并打印结果表格 ---
    df = pd.DataFrame(results)
    df = df.sort_values(by='QP').reset_index(drop=True)
    
    # 设置Pandas的浮点数显示格式
    pd.options.display.float_format = '{:,.4f}'.format
    
    print("\n" + "="*80)
    print(f"{'计算结果汇总':^80}")
    print("="*80)
    
    print(df.to_string(index=False))
    
    print("\n处理完成。")


if __name__ == "__main__":
    main()