"""
测试 cpdtab2 方法结合 putxl 导出
重点测试：
1. 使用 ___ 三下划线分割创建的多维表格（多个 index 和多个 columns）
2. auto_format 功能是否正确合并 index 和 column 的非最后一级单元格
3. Total 列和 Total 行的位置是否正确
"""

import pandas as pd
import numpy as np
import sys
import os

# 添加 pandaspro 到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

import pandaspro as cpd
from pandaspro.io.excel.putexcel import PutxlSet

# 创建测试数据
np.random.seed(42)

# 生成数据
regions = ['华东', '华北', '华南', '西南']
categories = ['电子产品', '服装', '食品']
quarters = ['Q1', 'Q2', 'Q3', 'Q4']
products = ['产品A', '产品B', '产品C']

data = []
for region in regions:
    for category in categories:
        for quarter in quarters:
            for product in products:
                data.append({
                    'Region': region,
                    'Category': category,
                    'Quarter': quarter,
                    'Product': product,
                    'Sales': np.random.randint(100, 1000),
                    'Quantity': np.random.randint(10, 100)
                })

df = pd.DataFrame(data)
df = cpd.FramePro(df)

print("="*80)
print("测试数据创建完成")
print("="*80)
print(df.head(10))
print(f"\n数据形状: {df.shape}")
print(f"列名: {list(df.columns)}")

# 创建 Excel 文件
output_file = os.path.join(os.path.dirname(__file__), 'test_cpdtab2_triple_underscore.xlsx')
ps = PutxlSet(output_file, noisily=True)

# ============================================================================
# 测试 1: 双 index, 双 columns (使用 ___ 分割)
# ============================================================================
print("\n" + "="*80)
print("测试 1: cpdtab2_Region__Category___Quarter__Product (双 index, 双 columns)")
print("="*80)

pivot1 = df.cpdtab2_Region__Category___Quarter__Product
print(pivot1.head(20))
print(f"\nIndex names: {pivot1.index.names}")
print(f"Columns names: {pivot1.columns.names}")
print(f"Shape: {pivot1.shape}")

ps.putxl(
    pivot1,
    sheet_name='01_DoubleIndex_DoubleCol',
    cell='A1',
    index=True,
    auto_format=True,
    index_auto_merge=True,
    sheetreplace=True
)

print("\n✓ 测试 1 完成 - 请检查:")
print("  - Region index 应该被合并（同一 Region 的多个 Category）")
print("  - Quarter columns 应该被合并（同一 Quarter 的多个 Product）")
print("  - 最后一级 index (Category) 不应该被合并")
print("  - 最后一级 column (Product) 不应该被合并")

# ============================================================================
# 测试 2: 双 index, 双 columns, 带 Subtotal
# ============================================================================
print("\n" + "="*80)
print("测试 2: cpdtab2s_Region__Category___Quarter__Product (双 index, 双 columns, 带 Subtotal)")
print("="*80)

pivot2 = df.cpdtab2s_Region__Category___Quarter__Product
print(pivot2.head(20))
print(f"\nIndex names: {pivot2.index.names}")
print(f"Columns names: {pivot2.columns.names}")
print(f"Shape: {pivot2.shape}")

# 检查是否有 Subtotal
print("\n检查 Subtotal 行:")
if isinstance(pivot2.index, pd.MultiIndex):
    subtotal_rows = [idx for idx in pivot2.index if 'Subtotal' in str(idx)]
    print(f"找到 {len(subtotal_rows)} 个 Subtotal 行")
    if subtotal_rows:
        print("前 5 个 Subtotal 行:")
        for idx in subtotal_rows[:5]:
            print(f"  {idx}")

print("\n检查 Subtotal 列:")
if isinstance(pivot2.columns, pd.MultiIndex):
    subtotal_cols = [col for col in pivot2.columns if 'Subtotal' in str(col)]
    print(f"找到 {len(subtotal_cols)} 个 Subtotal 列")
    if subtotal_cols:
        print("前 5 个 Subtotal 列:")
        for col in subtotal_cols[:5]:
            print(f"  {col}")

# 检查 Total 行和列的位置
print("\n检查 Total 位置:")
if isinstance(pivot2.index, pd.MultiIndex):
    total_rows = [idx for idx in pivot2.index if idx[0] in ['Total', 'All']]
    if total_rows:
        print(f"✓ Total 行位于: {pivot2.index.get_loc(total_rows[0])}")
        print(f"  总行数: {len(pivot2)}, Total 应该在最后")

if isinstance(pivot2.columns, pd.MultiIndex):
    total_cols = [col for col in pivot2.columns if col[0] in ['Total', 'All']]
    if total_cols:
        print(f"✓ Total 列位于: {pivot2.columns.get_loc(total_cols[0])}")
        print(f"  总列数: {len(pivot2.columns)}, Total 应该在最后")

ps.putxl(
    pivot2,
    sheet_name='02_DoubleIndex_DoubleCol_Sub',
    cell='A1',
    index=True,
    auto_format=True,
    index_auto_merge=True,
    sheetreplace=True
)

print("\n✓ 测试 2 完成 - 请检查:")
print("  - Region index 应该被合并且居中")
print("  - Quarter columns 应该被合并且居中")
print("  - Subtotal 行应该在每个 Region 的最后")
print("  - Subtotal 列应该在每个 Quarter 的最后")
print("  - Total 行应该在表格最下方")
print("  - Total 列应该在表格最右边")

# ============================================================================
# 测试 3: 三 index, 双 columns, 带 Subtotal 和聚合函数 (sum)
# ============================================================================
print("\n" + "="*80)
print("测试 3: cpdtab2ssum_Region__Category__Product___Quarter__Sales (三 index, 单 column, 带 sum)")
print("="*80)

pivot3 = df.cpdtab2ssum_Region__Category__Product___Quarter__Sales
print(pivot3.head(20))
print(f"\nIndex names: {pivot3.index.names}")
print(f"Columns names: {pivot3.columns.names}")
print(f"Shape: {pivot3.shape}")

ps.putxl(
    pivot3,
    sheet_name='03_TripleIndex_SingleCol_Sum',
    cell='A1',
    index=True,
    auto_format=True,
    index_auto_merge=True,
    sheetreplace=True
)

print("\n✓ 测试 3 完成 - 请检查:")
print("  - Region 和 Category 应该被合并（前两级 index）")
print("  - Product 不应该被合并（最后一级 index）")
print("  - Quarter column header 应该正常显示")
print("  - Subtotal 行应该出现在每个 Region 的最后")

# ============================================================================
# 测试 4: 单 index, 三 columns, 带 Subtotal
# ============================================================================
print("\n" + "="*80)
print("测试 4: cpdtab2s_Region___Quarter__Category__Product (单 index, 三 columns)")
print("="*80)

pivot4 = df.cpdtab2s_Region___Quarter__Category__Product
print(pivot4.head(10))
print(f"\nIndex names: {pivot4.index.names}")
print(f"Columns names: {pivot4.columns.names}")
print(f"Shape: {pivot4.shape}")

ps.putxl(
    pivot4,
    sheet_name='04_SingleIndex_TripleCol',
    cell='A1',
    index=True,
    auto_format=True,
    index_auto_merge=True,
    sheetreplace=True
)

print("\n✓ 测试 4 完成 - 请检查:")
print("  - Quarter 和 Category 应该被合并（前两级 columns）")
print("  - Product 不应该被合并（最后一级 column）")
print("  - Region index 正常显示（单级 index 不需要合并）")

# ============================================================================
# 测试 5: 不带 Subtotal 的双 index, 双 columns (对比测试)
# ============================================================================
print("\n" + "="*80)
print("测试 5: cpdtab2_Region__Category___Quarter__Product (无 Subtotal, 用于对比)")
print("="*80)

pivot5 = df.cpdtab2_Region__Category___Quarter__Product
print(pivot5.head(20))

ps.putxl(
    pivot5,
    sheet_name='05_DoubleIndex_DoubleCol_NoSub',
    cell='A1',
    index=True,
    auto_format=True,
    index_auto_merge=True,
    sheetreplace=True
)

print("\n✓ 测试 5 完成 - 请检查:")
print("  - 应该没有 Subtotal 行和列")
print("  - 合并单元格的逻辑应该与测试 2 相同")

# ============================================================================
# 保存并关闭
# ============================================================================
ps.close()

print("\n" + "="*80)
print("所有测试完成！")
print("="*80)
print(f"\n输出文件: {output_file}")
print("\n请打开 Excel 文件检查以下内容:")
print("\n1. 自动格式化检查:")
print("   ✓ Header 是否为蓝色背景、白色字体、粗体、居中")
print("   ✓ 是否有网格线")
print("   ✓ 整体外框是否为粗黑边框")
print("\n2. 合并单元格检查:")
print("   ✓ MultiIndex index 的前 N-1 级是否被合并且居中")
print("   ✓ MultiIndex columns 的前 N-1 级是否被合并且居中")
print("   ✓ 最后一级的 index 和 columns 是否保持不合并")
print("\n3. Subtotal 位置检查 (测试 2, 3, 4):")
print("   ✓ Subtotal 行是否在每个第一级 index 分组的最后")
print("   ✓ Subtotal 列是否在每个第一级 column 分组的最后")
print("   ✓ Total 行是否在表格最下方")
print("   ✓ Total 列是否在表格最右边")
print("\n4. 边框检查:")
print("   ✓ 每个第一级 index 分组是否有分区边框")
print("   ✓ 每个第一级 column 分组的首列是否有粗左边框")
print("\n如果发现问题，请报告具体的工作表名称和问题描述。")
print("="*80)


