import pandas as pd
import numpy as np
from loguru import logger

def generate(output_dir):
    logger.info("开始生成市场分析图")
    
    indices = ['上证指数', '深证成指', '创业板指', '沪深300']
    current = [3150, 10200, 2050, 3800]
    change_pct = [
        np.random.uniform(-2, 2),
        np.random.uniform(-2, 2),
        np.random.uniform(-2, 2),
        np.random.uniform(-2, 2)
    ]
    
    data = pd.DataFrame({
        '指数': indices,
        '点位': current,
        '涨跌幅(%)': [round(x, 2) for x in change_pct]
    })
    
    import matplotlib.pyplot as plt
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    colors = ['green' if x < 0 else 'red' for x in change_pct]
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))
    
    bars = ax1.bar(data['指数'], data['点位'], color=colors, alpha=0.7)
    ax1.set_title('主要指数点位', fontsize=14, fontweight='bold')
    ax1.set_ylabel('点位')
    ax1.grid(axis='y', linestyle='--', alpha=0.3)
    
    for bar in bars:
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2., height,
                 f'{int(height)}', ha='center', va='bottom')
    
    bars2 = ax2.bar(data['指数'], data['涨跌幅(%)'], color=colors, alpha=0.7)
    ax2.set_title('涨跌幅', fontsize=14, fontweight='bold')
    ax2.set_ylabel('涨跌幅 (%)')
    ax2.axhline(y=0, color='black', linestyle='-', linewidth=0.5)
    ax2.grid(axis='y', linestyle='--', alpha=0.3)
    
    for bar in bars2:
        height = bar.get_height()
        va = 'bottom' if height > 0 else 'top'
        ax2.text(bar.get_x() + bar.get_width()/2., height,
                 f'{height:+.2f}%', ha='center', va=va, fontweight='bold')
    
    plt.suptitle('A股市场概况', fontsize=18, fontweight='bold')
    plt.tight_layout()
    
    from pathlib import Path
    output_path = Path(output_dir) / 'market_analysis.png'
    plt.savefig(str(output_path), dpi=150, bbox_inches='tight')
    plt.close()
    
    logger.info(f"市场分析图已生成: {output_path}")
    return str(output_path)

if __name__ == "__main__":
    from pathlib import Path
    generate(Path.cwd())