import argparse
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter, MinuteLocator
import pandas as pd
from datetime import datetime

def parse_args():
    parser = argparse.ArgumentParser(description='OHLCV可视化工具')
    parser.add_argument('file_path', type=str, help='OHLCV数据文件路径')
    return parser.parse_args()

def load_data(file_path):
    """加载CSV格式的OHLCV数据"""
    df = pd.read_csv(file_path, header=None, 
                   names=['timestamp', 'open', 'high', 'low', 'close', 'volume'])
    
    # 转换时间戳为datetime
    df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
    return df

def plot_candlestick(df):
    """绘制K线图"""
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), gridspec_kw={'height_ratios': [3,1]})
    
    # 绘制K线
    for idx, row in df.iterrows():
        color = 'g' if row['close'] > row['open'] else 'r'
        ax1.vlines(row['datetime'], row['low'], row['high'], color=color, linewidth=1)
        ax1.vlines(row['datetime'], row['open'], row['close'], color=color, linewidth=3)

    # 设置时间格式
    date_format = DateFormatter('%m-%d %H:%M')
    ax1.xaxis.set_major_formatter(date_format)
    ax1.xaxis.set_major_locator(MinuteLocator(interval=30))
    ax1.set_title('OKX BTC/USDT 1小时K线图')
    ax1.grid(True)
    
    # 绘制成交量
    ax2.bar(df['datetime'], df['volume'], color=[('g' if c > o else 'r') for c, o in zip(df['close'], df['open'])], width=0.002)
    ax2.xaxis.set_major_formatter(date_format)
    ax2.set_title('成交量')
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    args = parse_args()
    try:
        df = load_data(args.file_path)
        plot_candlestick(df)
    except Exception as e:
        print(f"数据处理失败: {str(e)}")