#%%

import polars as pl
import numpy as np
import time
import os
from pathlib import Path

# 测试配置
STOCK_COUNTS = [10, 100, 1000]
COLUMN_CONFIGS = [
    {"int_cols": 1, "float_cols": 1},
    {"int_cols": 10, "float_cols": 10},
    {"int_cols": 200, "float_cols": 200}
]
TARGET_ELEMENTS = 100_000_000  # 1亿个元素

# 存储格式和压缩方式
FORMATS = {
    'feather': ['uncompressed', 'lz4', 'zstd'],
    'parquet': ['uncompressed', 'snappy', 'gzip', 'brotli', 'lz4', 'zstd'],
    'avro': ['uncompressed', 'snappy', 'deflate']
}

# 排序方式
SORT_METHODS = ['time_only', 'stock_time']

def generate_data(stock_count, int_cols, float_cols, sort_method):
    """生成测试数据"""
    # 计算需要的行数
    total_cols = 2 + int_cols + float_cols  # stock_id + timestamp + 其他列
    rows = TARGET_ELEMENTS // total_cols

    print(f"生成数据: {stock_count}只股票, {rows}行, {total_cols}列")

    # 生成基础数据
    np.random.seed(42)
    stock_ids = np.random.randint(0, stock_count, rows)
    timestamps = np.random.randint(1640995200, 1640995200 + 86400 * 10, rows)  # 2022年时间戳

    data = {
        'stock_id': stock_ids,
        'timestamp': timestamps
    }

    # 添加整数列 (订单数量)
    for i in range(int_cols):
        col_name = f'quantity_{i}' if int_cols > 1 else 'quantity'
        # data[col_name] = np.random.randint(1, 10000, rows)
        data[col_name] = np.random.poisson(1000, rows)

    # 添加浮点数列 (价格)
    for i in range(float_cols):
        col_name = f'price_{i}' if float_cols > 1 else 'price'
        base_price = np.random.uniform(10, 1000, rows)
        volatility = np.random.normal(0, 0.02, rows)
        data[col_name] = base_price * (1 + volatility)

    df = pl.DataFrame(data)

    # 排序
    if sort_method == 'time_only':
        df = df.sort('timestamp')
    elif sort_method == 'stock_time':
        df = df.sort(['stock_id', 'timestamp'])

    return df

def test_write_read(df, filename, format_type, compression):
    """测试写入和读取性能"""

    # 写入测试
    if format_type == 'feather':
        filepath = f"{filename}.feather"
        start_time = time.time()
        if compression == 'uncompressed':
            df.write_ipc(filepath)
        else:
            df.write_ipc(filepath, compression=compression)
        write_time = time.time() - start_time

    elif format_type == 'parquet':
        filepath = f"{filename}.parquet"
        start_time = time.time()
        df.write_parquet(filepath, compression=compression)
        write_time = time.time() - start_time

    elif format_type == 'avro':
        filepath = f"{filename}.avro"
        start_time = time.time()
        df.write_avro(filepath, compression=compression)
        write_time = time.time() - start_time

    # 读取测试
    start_time = time.time()
    if format_type == 'feather':
        test_df = pl.read_ipc(filepath)
    elif format_type == 'parquet':
        test_df = pl.read_parquet(filepath)
    elif format_type == 'avro':
        test_df = pl.read_avro(filepath)
    read_time = time.time() - start_time

    # 获取文件大小
    file_size = os.path.getsize(filepath)

    # 删除文件
    os.remove(filepath)

    return write_time, read_time, file_size

def main():
    """主函数"""
    print("开始金融数据存储性能测试...")

    # 创建结果存储
    results = []

    # 计算总测试数
    total_tests = 0
    for stock_count in STOCK_COUNTS:
        for col_config in COLUMN_CONFIGS:
            for sort_method in SORT_METHODS:
                for format_type in FORMATS:
                    total_tests += len(FORMATS[format_type])

    current_test = 0

    # 开始测试
    for stock_count in STOCK_COUNTS:
        for col_config in COLUMN_CONFIGS:
            int_cols = col_config['int_cols']
            float_cols = col_config['float_cols']

            for sort_method in SORT_METHODS:
                # 生成数据
                df = generate_data(stock_count, int_cols, float_cols, sort_method)
                rows = len(df)
                total_elements = len(df) * len(df.columns)

                for format_type in FORMATS:
                    for compression in FORMATS[format_type]:
                        current_test += 1

                        filename = f"test_{stock_count}_{int_cols}_{float_cols}_{sort_method}_{format_type}_{compression}"

                        print(f"\n进度: {current_test}/{total_tests}")
                        print(f"测试: {stock_count}只股票, {int_cols}+{float_cols}列, {sort_method}, {format_type} {compression}")

                        try:
                            write_time, read_time, file_size = test_write_read(df, filename, format_type, compression)

                            # 记录结果
                            result = {
                                '股票数': stock_count,
                                '整数列数': int_cols,
                                '浮点列数': float_cols,
                                '排序方式': '时间排序' if sort_method == 'time_only' else '股票+时间排序',
                                '格式': format_type,
                                '压缩方式': compression,
                                '写入时间(秒)': round(write_time, 3),
                                '读取时间(秒)': round(read_time, 3),
                                '文件大小(MB)': round(file_size / 1024 / 1024, 2),
                                '行数': rows,
                                '总元素数': total_elements,
                                '写入速度(MB/s)': round(file_size / 1024 / 1024 / write_time, 2),
                                '读取速度(MB/s)': round(file_size / 1024 / 1024 / read_time, 2)
                            }
                            results.append(result)

                            print(f"  写入: {write_time:.3f}s, 读取: {read_time:.3f}s, 大小: {file_size/1024/1024:.1f}MB")

                        except Exception as e:
                            print(f"  错误: {e}")
                            continue

    # 创建结果DataFrame
    if results:
        results_df = pl.DataFrame(results)

        # 保存结果
        results_df.write_csv("storage_performance_results.csv")
        print(f"\n结果已保存到 storage_performance_results.csv")

        # 显示完整结果
        print("\n=== 完整测试结果 ===")
        print(results_df)

        # 找出最佳表现者
        print("\n=== 最佳表现者 ===")

        # 最快写入
        fastest_write = results_df.sort('写入时间(秒)').row(0)
        print(f"\n最快写入:")
        print(f"  {fastest_write[3]} {fastest_write[4]} {fastest_write[5]} - {fastest_write[6]}秒")

        # 最快读取
        fastest_read = results_df.sort('读取时间(秒)').row(0)
        print(f"\n最快读取:")
        print(f"  {fastest_read[3]} {fastest_read[4]} {fastest_read[5]} - {fastest_read[7]}秒")

        # 最小文件
        smallest_file = results_df.sort('文件大小(MB)').row(0)
        print(f"\n最小文件:")
        print(f"  {smallest_file[3]} {smallest_file[4]} {smallest_file[5]} - {smallest_file[8]}MB")

        # 按格式分组统计
        print("\n=== 按格式分组统计 ===")
        format_stats = results_df.group_by('格式').agg([
            pl.col('写入时间(秒)').mean().alias('平均写入时间'),
            pl.col('读取时间(秒)').mean().alias('平均读取时间'),
            pl.col('文件大小(MB)').mean().alias('平均文件大小'),
            pl.col('写入速度(MB/s)').mean().alias('平均写入速度'),
            pl.col('读取速度(MB/s)').mean().alias('平均读取速度')
        ]).sort('格式')
        print(format_stats)

        # 按压缩方式分组统计
        print("\n=== 按压缩方式分组统计 ===")
        compression_stats = results_df.group_by('压缩方式').agg([
            pl.col('写入时间(秒)').mean().alias('平均写入时间'),
            pl.col('读取时间(秒)').mean().alias('平均读取时间'),
            pl.col('文件大小(MB)').mean().alias('平均文件大小')
        ]).sort('压缩方式')
        print(compression_stats)

        # 按股票数量分组统计
        print("\n=== 按股票数量分组统计 ===")
        stock_stats = results_df.group_by('股票数').agg([
            pl.col('写入时间(秒)').mean().alias('平均写入时间'),
            pl.col('读取时间(秒)').mean().alias('平均读取时间'),
            pl.col('文件大小(MB)').mean().alias('平均文件大小')
        ]).sort('股票数')
        print(stock_stats)

    print("\n测试完成!")
    return results

if __name__ == "__main__":
    results = main()