#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
需求：
1. 生成 1 张「跨-Sheet 5 指标汇总」柱状图（保留）  
2. 额外生成 5 张「单指标跨-Sheet 对比柱状图」  
3. 全部输出到指定文件夹
python six_plots.py your.xlsx -o output_folder
"""

import argparse, pandas as pd, matplotlib, os, re
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# ---------- 工具 ----------
def find_avg_row(df):
    for idx, row in df.iterrows():
        if any(isinstance(cell, str) and re.search(r'平均值|Average', str(cell)) for cell in row):
            return idx
    return df.index[-1]

def extract_values(df):
    idx = find_avg_row(df)
    vals = [float(v) for v in df.iloc[idx].dropna().iloc[-5:] if str(v).replace('.', '').isdigit()]
    if len(vals) != 5:
        vals = [float(v) for v in df.iloc[-1].dropna().iloc[-5:] if str(v).replace('.', '').isdigit()]
    return {
        "total_throughput": vals[4],
        "generate_throughput": vals[3],
        "singel_road_generate_throughput": vals[1],
        "generate_throughput_without_ttft": vals[0],
        "ttft": vals[2]
    }

# ---------- 绘图 ----------
def draw_bar(data, labels, colors, title, out_path, xlabels):
    fig, ax = plt.subplots(figsize=(max(len(xlabels) * 0.8, 12), 6))
    fig.patch.set_facecolor('#ffffff')
    ax.set_facecolor('#ffffff')
    bars = ax.bar(xlabels, data, color=colors, width=0.6)
    for bar, v in zip(bars, data):
        height = bar.get_height()
        ax.annotate(f"{v:.1f}%",
                    xy=(bar.get_x() + bar.get_width() / 2, height),
                    xytext=(0, 3), textcoords='offset points',
                    ha='center', va='bottom', fontsize=7, rotation=60,
                    bbox=dict(boxstyle="round,pad=0.15", facecolor='white', edgecolor='none'))
    ax.set_ylabel("Percentage")
    ax.set_title(title, fontsize=13, weight='bold')
    ax.grid(axis='y', linestyle='-', linewidth=0.25, alpha=0.2)
    ax.tick_params(axis='x', rotation=30)
    plt.tight_layout()
    for ext in ['png', 'pdf']:
        plt.savefig(f"{out_path}.{ext}", dpi=300 if ext == 'png' else None, bbox_inches='tight')
    plt.close()

# ---------- 主流程 ----------
def main():
    parser = argparse.ArgumentParser(description='生成 6 张柱状图：1 张汇总 + 5 张单指标')
    parser.add_argument("xlsx_path", help="输入 Excel 文件")
    parser.add_argument("-o", "--output", default="output", help="输出文件夹")
    args = parser.parse_args()

    xls = pd.ExcelFile(args.xlsx_path)
    summary = {sheet: extract_values(pd.read_excel(xls, sheet_name=sheet, header=None))
               for sheet in xls.sheet_names}

    labels = [
        "total_throughput",
        "generate_throughput",
        "singel_road_generate_throughput",
        "generate_throughput_without_ttft",
        "ttft"
    ]
    colors = ['#005F73', '#0A9396', '#94D2BD', '#E9D8A6', '#d62728']
    os.makedirs(args.output, exist_ok=True)

    # 1. 跨-Sheet 5 指标汇总（原图保持不变）
    models = list(summary.keys())
    x = range(len(models))
    bar_width = 0.14
    fig, ax = plt.subplots(figsize=(max(len(models) * 1.1, 14), 8))
    fig.patch.set_facecolor('#F8F9FA')
    ax.set_facecolor('#FFFFFF')
    for i, lab in enumerate(labels):
        vals = [summary[m][lab] for m in models]
        offset = (i - 1.5) * bar_width
        bars = ax.bar([p + offset for p in x], vals, width=bar_width,
                      label=lab.replace('_', ' ').title(), color=colors[i])
        for bar, v in zip(bars, vals):
            height = bar.get_height()
            line_start = (bar.get_x() + bar.get_width() / 2, height)
            text_pos   = (line_start[0], height + max(vals) * 0.05)
            ax.plot([line_start[0], text_pos[0]], [line_start[1], text_pos[1]],
                    color='#555555', linestyle='--', linewidth=0.8, alpha=0.7)
            ax.text(text_pos[0], text_pos[1], f"{v:.1f}%",
                    ha='center', va='bottom', fontsize=7, rotation=60,
                    color='#000000', weight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(models, rotation=25, ha='right', fontsize=10, weight='bold')
    ax.set_ylabel("Percentage", fontsize=12, weight='bold')
    ax.set_title("Throughput Comparison", fontsize=16, weight='bold', pad=20)
    ax.grid(axis='y', linestyle='-', linewidth=0.3, alpha=0.3, color='#DEE2E6')
    ax.legend(frameon=False, loc='upper left', bbox_to_anchor=(1.02, 1), fontsize=9)
    plt.tight_layout()
    for ext in ['png', 'pdf']:
        plt.savefig(os.path.join(args.output, f"summary.{ext}"), dpi=300 if ext == 'png' else None,
                    bbox_inches='tight')
    plt.close()

    # 2-6. 5 张单指标跨-Sheet 对比图
    for lab, color in zip(labels, colors):
        vals = [summary[s][lab] for s in models]
        draw_bar(vals, [lab], [color],
                 f"{lab.replace('_', ' ').title()} - All Sheets",
                 os.path.join(args.output, lab),
                 models)

    print(f"✅ 6 张图已生成：{args.output}")

if __name__ == '__main__':
    main()