import pandas as pd  
import numpy as np  
import vectorbt as vbt  
from numba import njit  
import vectorbt as vbt
import pandas as pd
import numpy as np

import pandas as pd
from loguru import logger
from chart import resample_data, QuantChart
from db_config import DatabaseManager, DatabaseConfig
from vectorbt.base.reshape_fns import to_2d_array  


def get_kline(code: str) -> list:
    """
    获取所有指数代码
    :return:
    """
    db_config = DatabaseConfig()
    logger.info(f"数据库连接: {db_config.connection_string}")

    db_manager = DatabaseManager(db_config)
    sql = f"SELECT date as time, open, high, low, close, volume FROM public.index_kline where code='{code}' order by date;"
    res = db_manager.execute_query(sql)
    data_list = [dict(item) for item in res]
    df = pd.DataFrame(data_list)
    df["time"] = pd.to_datetime(df["time"])
    num_cols = ["open", "high", "low", "close", "volume"]
    for col in num_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce").astype(float)
    return df

symbol = "399998"
timeframe = "1D"

df = get_kline(code=symbol)
df = resample_data(df, timeframe)
df.rename(columns={'time': 'date'}, inplace=True)
print(df.head())
if 'date' in df.columns:  
    df = df.set_index('date')  
price = df['close']  
  
# 2. 计算90天滚动波动率(年化)  
returns = price.pct_change()  
volatility_90d = returns.rolling(window=90, min_periods=90).std() * np.sqrt(365)  
  
# 3. 计算波动率倒数作为权重  
inv_vol = 1 / volatility_90d  
# 标准化权重(可选,使其更易解释)  
inv_vol_normalized = inv_vol / inv_vol.rolling(window=252).mean()  
  
# 4. 创建每周重新平衡的信号  
# 获取每周最后一个交易日  
weekly_rebalance = pd.Series(False, index=price.index)  
weekly_last_days = price.resample('W').last().index  
for date in weekly_last_days:  
    # 找到最接近的交易日  
    idx = price.index.get_indexer([date], method='ffill')[0]  
    if idx >= 0:  
        weekly_rebalance.iloc[idx] = True  
  
# 5. 定义订单函数  
@njit  
def order_func_nb(c, inv_vol_arr, rebalance_arr):  
    # 获取当前的波动率倒数权重  
    inv_vol_now = vbt.nb.flex_select_auto_nb(inv_vol_arr, c.i, c.col, False)  
    rebalance_now = vbt.nb.flex_select_auto_nb(rebalance_arr, c.i, c.col, False)  
      
    # 只在重新平衡日调整仓位  
    if not rebalance_now or np.isnan(inv_vol_now):  
        return vbt.nb.order_nothing_nb()  
      
    # 目标仓位 = 总价值 * 波动率倒数权重  
    # 这里使用 TargetPercent 类型,权重越高仓位越大  
    target_percent = min(inv_vol_now, 1.0)  # 限制最大100%仓位  
      
    return vbt.nb.order_nb(  
        size=target_percent,  
        size_type=vbt.SizeType.TargetPercent,  
        direction=vbt.Direction.LongOnly  
    )  
  
# 6. 运行回测  
pf = vbt.Portfolio.from_order_func(  
    price,  
    order_func_nb,  
    to_2d_array(inv_vol_normalized),  
    to_2d_array(weekly_rebalance),  
    init_cash=100,  
    freq='1D'  
)  
  
# 7. 查看结果  
print(pf.stats())  
print(f"\n波动率倒数策略 vs 买入持有:")  
print(f"总收益率: {pf.total_return():.2%}")