"""
测试 putxl 的三个新功能：
1. MultiIndex columns header 自动合并
2. 自动默认格式（蓝色header、白色字体、网格线、thick_outer边框）
3. 基于 index 的条件格式化（如给Subtotal行添加底纹）
"""
import pandas as pd
import numpy as np
from pandaspro.io.excel.putexcel import PutxlSet

# 创建模拟的 pivot table 数据
# 场景：销售数据，包含多个维度
np.random.seed(42)

# 创建原始数据
data = {
    'Region': ['North', 'North', 'North', 'North', 'South', 'South', 'South', 'South',
               'East', 'East', 'East', 'East'],
    'Product': ['A', 'B', 'A', 'B', 'A', 'B', 'A', 'B', 'A', 'B', 'A', 'B'],
    'Quarter': ['Q1', 'Q1', 'Q2', 'Q2', 'Q1', 'Q1', 'Q2', 'Q2', 'Q1', 'Q1', 'Q2', 'Q2'],
    'Sales': [100, 150, 120, 180, 90, 140, 110, 170, 130, 160, 140, 190],
    'Profit': [20, 30, 25, 35, 18, 28, 22, 34, 26, 32, 28, 38]
}

df = pd.DataFrame(data)

# 创建 pivot table 1: 两层 columns (值 x Quarter)
print("="*80)
print("创建测试数据...")
print("="*80)

pivot1 = df.pivot_table(
    values=['Sales', 'Profit'],
    index=['Region', 'Product'],
    columns='Quarter',
    aggfunc='sum',
    margins=True,  # 添加总计行
    margins_name='Total'  # 总计行名称
)

print("\nPivot Table 1 (MultiIndex columns):")
print(pivot1.head())
print(f"\nColumns: {pivot1.columns.tolist()}")
print(f"Index names: {pivot1.index.names}")

# 创建 pivot table 2: 带 Subtotal 的三层 index
# 先创建分组数据
grouped_data = []
for region in ['North', 'South', 'East']:
    for product in ['A', 'B']:
        sales = df[(df['Region'] == region) & (df['Product'] == product)]['Sales'].sum()
        profit = df[(df['Region'] == region) & (df['Product'] == product)]['Profit'].sum()
        grouped_data.append({
            'Region': region,
            'Category': product,
            'Type': 'Regular',
            'Sales': sales,
            'Profit': profit
        })
    # 添加 Subtotal
    sales_subtotal = df[df['Region'] == region]['Sales'].sum()
    profit_subtotal = df[df['Region'] == region]['Profit'].sum()
    grouped_data.append({
        'Region': region,
        'Category': 'Subtotal',
        'Type': 'Subtotal',
        'Sales': sales_subtotal,
        'Profit': profit_subtotal
    })

pivot2 = pd.DataFrame(grouped_data).set_index(['Region', 'Category', 'Type'])

print("\n" + "="*80)
print("Pivot Table 2 (带 Subtotal 的三层 index):")
print("="*80)
print(pivot2)

# 导出到 Excel
print("\n" + "="*80)
print("导出到 Excel...")
print("="*80)

output_file = 'test_pivot_improvements.xlsx'
ps = PutxlSet(output_file)

# 测试 1: 基本的 MultiIndex columns - 启用 auto_format
print("\n1️⃣  测试 1: MultiIndex columns 自动合并 + 默认格式")
ps.putxl(
    pivot1,
    sheet_name='01_AutoFormat',
    cell='A1',
    index=True,
    auto_format=True,  # 启用自动格式
    tab_color='blue'
)
print("   ✅ 完成")

# 测试 2: 关闭 auto_format 查看对比
print("\n2️⃣  测试 2: 不使用 auto_format（对比）")
ps.putxl(
    pivot1,
    sheet_name='02_NoAutoFormat',
    cell='A1',
    index=True,
    auto_format=False,  # 关闭自动格式
    tab_color='gray'
)
print("   ✅ 完成")

# 测试 3: 基于 index 的条件格式化 - 高亮 Subtotal 行
print("\n3️⃣  测试 3: 基于 index 的条件格式化（Subtotal 行）")
ps.putxl(
    pivot2,
    sheet_name='03_IndexCondFormat',
    cell='A1',
    index=True,
    auto_format=True,
    cd_format={
        'index_level': 'Type',  # 基于 Type 这个 index level
        'rules': {
            'Subtotal': 'fill=lightyellow; bold=True'  # Subtotal 行黄色底纹+加粗
        }
    },
    tab_color='orange'
)
print("   ✅ 完成")

# 测试 4: 综合测试 - 使用通配符匹配 Subtotal
print("\n4️⃣  测试 4: 使用通配符匹配 *Subtotal")
# 修改一些行名包含 Subtotal
pivot3 = pivot2.copy()
pivot3.index = pd.MultiIndex.from_tuples([
    ('North', 'A', 'Regular'),
    ('North', 'B', 'Regular'),
    ('North', 'Subtotal', 'Region Subtotal'),  # 修改为包含 Subtotal
    ('South', 'A', 'Regular'),
    ('South', 'B', 'Regular'),
    ('South', 'Subtotal', 'Region Subtotal'),
    ('East', 'A', 'Regular'),
    ('East', 'B', 'Regular'),
    ('East', 'Subtotal', 'Region Subtotal'),
], names=['Region', 'Category', 'Type'])

ps.putxl(
    pivot3,
    sheet_name='04_WildcardMatch',
    cell='A1',
    index=True,
    auto_format=True,
    cd_format={
        'index_level': 'Type',
        'rules': {
            '*Subtotal': 'fill=lightgreen; bold=True; italic=True'  # 匹配所有包含 Subtotal 的
        }
    },
    tab_color='green'
)
print("   ✅ 完成")

# 测试 5: 完整示例 - 所有功能组合
print("\n5️⃣  测试 5: 完整示例（所有功能组合）")
ps.putxl(
    pivot1,
    sheet_name='05_Complete',
    cell='A1',
    index=True,
    auto_format=True,
    cd_format={
        'index_level': 'Region',
        'rules': {
            'Total': 'fill=lightblue; bold=True'  # Total 行蓝色底纹
        }
    },
    df_format={
        'number_format=#,##0': 'data'  # 数据区域添加千位分隔符
    },
    tab_color='purple'
)
print("   ✅ 完成")

print("\n" + "="*80)
print(f"✅ 所有测试完成！文件已导出: {output_file}")
print("="*80)
print("\n包含以下工作表:")
print("  1. 01_AutoFormat - 启用 auto_format（蓝色header、合并、边框）")
print("  2. 02_NoAutoFormat - 不启用 auto_format（对比）")
print("  3. 03_IndexCondFormat - 基于 index 条件格式化（Subtotal行黄色）")
print("  4. 04_WildcardMatch - 通配符匹配 *Subtotal（绿色）")
print("  5. 05_Complete - 完整示例（所有功能组合）")
print("\n请打开 Excel 查看效果！")
print("="*80)

# 读取并验证
print("\n" + "="*80)
print("验证导出结果...")
print("="*80)

xl_file = pd.ExcelFile(output_file)
print(f"\n工作表列表: {xl_file.sheet_names}")
print("\n✅ 验证完成！")

