"""
缠论量化引擎 v1.0
按标准定义实现: 包含处理→分型→笔→线段→中枢→买卖点
"""
import pandas as pd
import numpy as np

# ── 1. 包含关系处理 ──
def merge_candles(df):
    """
    缠论K线包含处理:
    - 向上趋势: 取高高(高点取max, 低点取max)
    - 向下趋势: 取低低(低点取min, 高点取min)
    方向由前两根非包含K线决定
    """
    klines = df[["open","high","low","close"]].copy().values
    n = len(klines)
    merged = []
    direction = 1  # 1=向上, -1=向下
    
    i = 0
    while i < n:
        if i == 0:
            merged.append(klines[i])
            i += 1
            continue
        
        prev = merged[-1]
        curr = klines[i]
        
        # 判断包含关系: curr的高低点都在prev范围内
        if (curr[2] >= prev[2] and curr[1] <= prev[1]) or \
           (curr[2] <= prev[2] and curr[1] >= prev[1]):
            # 包含 → 合并
            if direction == 1:  # 向上: 取高高
                new_high = max(prev[1], curr[1])
                new_low = max(prev[2], curr[2])
            else:  # 向下: 取低低
                new_high = min(prev[1], curr[1])
                new_low = min(prev[2], curr[2])
            
            # 保持open/close的相对位置
            new_open = prev[0]
            new_close = curr[3]
            merged[-1] = np.array([new_open, new_high, new_low, new_close])
        else:
            # 不包含 → 确定方向
            if curr[1] > prev[1]:
                direction = 1
            else:
                direction = -1
            merged.append(curr)
        
        i += 1
    
    return pd.DataFrame(merged, columns=["open","high","low","close"])

# ── 2. 分型识别 ──
def identify_fenxing(df):
    """
    顶分型: 中间K线高点最高, 低点最高
    底分型: 中间K线低点最低, 高点最低
    """
    H, L = df["high"].values, df["low"].values
    n = len(df)
    fenxing = []  # (index, type)  type=1顶分型, -1=底分型
    
    for i in range(1, n-1):
        # 顶分型
        if H[i] > H[i-1] and H[i] > H[i+1] and L[i] > L[i-1] and L[i] > L[i+1]:
            fenxing.append((i, 1, H[i], L[i]))  # (idx, type, top, bottom)
        # 底分型
        elif L[i] < L[i-1] and L[i] < L[i+1] and H[i] < H[i-1] and H[i] < H[i+1]:
            fenxing.append((i, -1, H[i], L[i]))
    
    return fenxing

# ── 3. 笔 (Bi) ──
def identify_bi(fenxing, df_merged):
    """
    笔: 相邻的底分型和顶分型之间，至少间隔1根K线
    上升笔: 底分型→顶分型
    下降笔: 顶分型→底分型
    """
    if len(fenxing) < 2:
        return []
    
    bi_list = []
    i = 0
    H, L = df_merged["high"].values, df_merged["low"].values
    
    while i < len(fenxing) - 1:
        fx1 = fenxing[i]
        
        # 找下一个相反类型且位置合理的分型
        j = i + 1
        best_j = -1
        while j < len(fenxing):
            fx2 = fenxing[j]
            if fx2[1] != fx1[1]:  # 类型相反
                if fx2[0] - fx1[0] >= 3:  # 至少间隔1根(含分型自身)
                    # 上升笔: 底→顶, 顶必须高于底
                    if fx1[1] == -1 and fx2[2] > fx1[3]:
                        # 检查中间有没有更低的底
                        mid_low = min(L[fx1[0]:fx2[0]+1])
                        if mid_low >= fx1[3]:
                            best_j = j
                            break
                    # 下降笔: 顶→底, 底必须低于顶
                    elif fx1[1] == 1 and fx2[3] < fx1[2]:
                        mid_high = max(H[fx1[0]:fx2[0]+1])
                        if mid_high <= fx1[2]:
                            best_j = j
                            break
            j += 1
        
        if best_j > 0:
            fx2 = fenxing[best_j]
            bi_list.append({
                "start_idx": fx1[0],
                "end_idx": fx2[0],
                "start_type": "底" if fx1[1] == -1 else "顶",
                "end_type": "顶" if fx2[1] == 1 else "底",
                "direction": "up" if fx1[1] == -1 else "down",
                "start_price": df_merged["close"].iloc[fx1[0]],
                "end_price": df_merged["close"].iloc[fx2[0]],
            })
            i = best_j
        else:
            i += 1
    
    return bi_list

# ── 4. 线段 ──
def identify_xduan(bi_list):
    """
    线段: 至少由3笔组成，方向交替
    线段被破坏: 出现反向笔突破前一笔的端点
    """
    if len(bi_list) < 3:
        return []
    
    xduan_list = []
    start = 0
    
    while start < len(bi_list) - 2:
        # 取连续3笔
        b1, b2, b3 = bi_list[start], bi_list[start+1], bi_list[start+2]
        
        # 必须方向交替
        if b1["direction"] == b2["direction"]:
            start += 1
            continue
        
        # 线段破坏判断
        if b1["direction"] == "up":
            # 上升线段: b1上, b2下, b3上. 看b2是否跌破b1起点
            if b3["end_price"] > b1["end_price"]:
                # 确认上升线段
                end = start + 3
                while end < len(bi_list) - 1:
                    if bi_list[end]["direction"] == "down":
                        if bi_list[end]["end_price"] < bi_list[start+1]["start_price"]:
                            break
                    end += 1
                
                xduan_list.append({
                    "start_bi": start,
                    "end_bi": end - 1,
                    "direction": "up",
                    "high": max(b["end_price"] for b in bi_list[start:end]),
                    "low": min(b["start_price"] for b in bi_list[start:end]),
                })
                start = end
            else:
                start += 1
        else:
            if b3["end_price"] < b1["end_price"]:
                end = start + 3
                while end < len(bi_list) - 1:
                    if bi_list[end]["direction"] == "up":
                        if bi_list[end]["end_price"] > bi_list[start+1]["start_price"]:
                            break
                    end += 1
                
                xduan_list.append({
                    "start_bi": start,
                    "end_bi": end - 1,
                    "direction": "down",
                    "high": max(b["start_price"] for b in bi_list[start:end]),
                    "low": min(b["end_price"] for b in bi_list[start:end]),
                })
                start = end
            else:
                start += 1
    
    return xduan_list

# ── 5. 中枢 (Zhongshu/Pivot) ──
def identify_zhongshu(xduan_list):
    """
    中枢: 至少3个连续线段的重叠区间
    ZG = min(线段高点), ZD = max(线段低点)
    """
    if len(xduan_list) < 3:
        return []
    
    zs_list = []
    i = 0
    
    while i < len(xduan_list) - 2:
        highs = [x["high"] for x in xduan_list[i:i+3]]
        lows = [x["low"] for x in xduan_list[i:i+3]]
        
        ZG = min(highs)  # 中枢上沿
        ZD = max(lows)   # 中枢下沿
        
        if ZG > ZD:  # 有重叠 → 中枢成立
            # 向后扩展中枢
            end = i + 3
            while end < len(xduan_list):
                h = xduan_list[end]["high"]
                l = xduan_list[end]["low"]
                # 新线段与现有中枢有重叠 → 扩展
                if max(ZG, h) - min(ZD, l) > 0 and h > ZD and l < ZG:
                    ZG = min(ZG, h)
                    ZD = max(ZD, l)
                    end += 1
                else:
                    break
            
            zs_list.append({
                "start_xd": i,
                "end_xd": end - 1,
                "ZG": ZG,  # 中枢高点
                "ZD": ZD,  # 中枢低点
                "ZZ": (ZG + ZD) / 2,  # 中枢中轴
            })
            i = end
        else:
            i += 1
    
    return zs_list

# ── 主函数 ──
def run_chanlun(df_original):
    """运行完整缠论分析"""
    print("="*60)
    print("缠论量化引擎 v1.0")
    print("="*60)
    
    # 1. 包含处理
    df = merge_candles(df_original)
    print(f"  包含处理: {len(df_original)} → {len(df)} K线")
    
    # 2. 分型
    fx = identify_fenxing(df)
    tops = [f for f in fx if f[1] == 1]
    bots = [f for f in fx if f[1] == -1]
    print(f"  分型: {len(fx)} 个 (顶分型{len(tops)}, 底分型{len(bots)})")
    
    # 3. 笔
    bi = identify_bi(fx, df)
    up_bi = [b for b in bi if b["direction"] == "up"]
    dn_bi = [b for b in bi if b["direction"] == "down"]
    print(f"  笔: {len(bi)} 笔 (上升{len(up_bi)}, 下降{len(dn_bi)})")
    
    # 4. 线段
    xd = identify_xduan(bi)
    print(f"  线段: {len(xd)} 段")
    
    # 5. 中枢
    zs = identify_zhongshu(xd)
    print(f"  中枢: {len(zs)} 个")
    
    return {
        "merged_df": df,
        "fenxing": fx,
        "bi": bi,
        "xduan": xd,
        "zhongshu": zs,
    }


# ── 测试 ──
if __name__ == "__main__":
    df_raw = pd.read_parquet("data/btc_multidim.parquet")
    d = df_raw.resample("1h").agg({"open":"first","high":"max","low":"min","close":"last"}).dropna()
    
    # 用最后3个月测试(数据太多会慢)
    test_data = d.iloc[-90*24:]  # 约90天
    result = run_chanlun(test_data)
    
    # 打印最近一个中枢
    if result["zhongshu"]:
        zs = result["zhongshu"][-1]
        print(f"\n  最近中枢: ZG={zs['ZG']:.2f} ZD={zs['ZD']:.2f} ZZ={zs['ZZ']:.2f}")
        print(f"  中枢宽度: {(zs['ZG']-zs['ZD'])/zs['ZZ']*100:.2f}%")
    
    # 笔统计
    if result["bi"]:
        prices_up = [b["end_price"]/b["start_price"]-1 for b in result["bi"] if b["direction"]=="up"]
        prices_dn = [b["start_price"]/b["end_price"]-1 for b in result["bi"] if b["direction"]=="down"]
        if prices_up:
            print(f"\n  上升笔平均涨幅: {np.mean(prices_up)*100:.2f}%")
        if prices_dn:
            print(f"  下降笔平均跌幅: {np.mean(prices_dn)*100:.2f}%")