#!/usr/bin/env python3
"""
浪浪AI 单元测试套件

测试覆盖:
1. ZigZag摆动点识别
2. 下跌段识别
3. 盒子识别
4. 放量突破检测
5. 多周期共振
6. 1H精准入场
7. 做空信号
8. 6层过滤
9. 仓位计算
10. Trailing出场
11. 回测引擎集成
12. K线聚合
"""

import sys
import os
import numpy as np
import pandas as pd
from datetime import timedelta

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from langlang_ai.core.zigzag import zigzag
from langlang_ai.core.pipeline import (
    find_drop_legs, find_boxes, find_breakouts,
    match_resonance, find_1h_entry, Signal, Breakout, Box, DropLeg,
)
from langlang_ai.core.shorts import find_short_signals
from langlang_ai.core.filters import SignalFilter, calc_trend_60d
from langlang_ai.core.position import calc_position, calc_effective_equity
from langlang_ai.core.trailing import simulate_exit
from langlang_ai.data.aggregator import aggregate_to_4h, aggregate_to_daily
from langlang_ai.config import settings as S

passed = 0
failed = 0


def test(name, condition, detail=""):
    global passed, failed
    if condition:
        passed += 1
        print(f"  ✓ {name}")
    else:
        failed += 1
        print(f"  ✗ {name} — {detail}")


# ========================================
# Helper: build synthetic price data
# ========================================

def make_price_series(phases, base=10000, noise=0.002):
    """
    Build a price path from phase definitions.
    phases: list of (n_bars, drift_per_bar)
    Returns: numpy array of prices
    """
    np.random.seed(42)
    parts = []
    price = base
    for n_bars, drift in phases:
        for _ in range(n_bars):
            price *= (1 + drift + np.random.normal(0, noise))
            parts.append(price)
    return np.array(parts)


def make_ohlcv(prices, vol_base=100, vol_spikes=None):
    """Build OHLCV arrays from a close price series"""
    n = len(prices)
    np.random.seed(42)
    opens = prices * (1 + np.random.normal(0, 0.001, n))
    highs = prices * (1 + np.abs(np.random.normal(0.004, 0.002, n)))
    lows = prices * (1 - np.abs(np.random.normal(0.004, 0.002, n)))
    volumes = np.random.uniform(vol_base * 0.5, vol_base * 1.5, n)
    if vol_spikes:
        for idx, mult in vol_spikes:
            volumes[idx] *= mult
    return opens, highs, lows, prices, volumes


def make_df_1h(phases, base=10000, start='2020-01-01', vol_spikes=None):
    """Build a 1H DataFrame from phase definitions"""
    prices = make_price_series(phases, base)
    n = len(prices)
    ts = pd.date_range(start, periods=n, freq='1h', tz='UTC')
    opens, highs, lows, closes, volumes = make_ohlcv(prices, vol_spikes=vol_spikes)
    return pd.DataFrame({
        'timestamp': ts,
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes,
    })


# ========================================
print("\n" + "=" * 60)
print("浪浪AI 测试套件")
print("=" * 60)

# ========================================
print("\n[1] ZigZag摆动点识别")
# ========================================

# Test: clear up-down-up pattern
prices_ud = make_price_series([
    (50, 0.005),   # +28% up
    (40, -0.008),  # -27% down
    (50, 0.006),   # +35% up
])
_, highs_ud, lows_ud, _, _ = make_ohlcv(prices_ud)
points = zigzag(highs_ud, lows_ud, 0.15)

test("ZigZag finds swing points", len(points) >= 3,
     f"got {len(points)}")
test("ZigZag alternates H/L",
     all(points[i][1] != points[i+1][1] for i in range(len(points)-1)),
     f"points: {[(p[1], p[0]) for p in points]}")

# Test: small threshold finds more points
points_small = zigzag(highs_ud, lows_ud, 0.05)
test("Smaller threshold → more points",
     len(points_small) >= len(points),
     f"15%: {len(points)}, 5%: {len(points_small)}")

# Test: flat data → no points
flat = np.ones(100) * 10000
flat_h = flat * 1.001
flat_l = flat * 0.999
points_flat = zigzag(flat_h, flat_l, 0.15)
test("Flat data → no swing points", len(points_flat) == 0)

# ========================================
print("\n[2] 下跌段识别")
# ========================================

# Build H-L-H-L-H pattern
swing_points = [
    (0, 'H', 10000),
    (20, 'L', 7000),    # -30% drop ✓
    (40, 'H', 9000),
    (60, 'L', 8500),    # -5.5% drop ✗ (< 12%)
    (80, 'H', 11000),
]
legs_daily = find_drop_legs(swing_points, min_drop=0.12)
test("Finds valid drop legs (≥12%)", len(legs_daily) == 1,
     f"got {len(legs_daily)}")
test("Drop leg has correct drop %",
     abs(legs_daily[0].drop_pct - 0.30) < 0.01 if legs_daily else False,
     f"drop={legs_daily[0].drop_pct if legs_daily else 'N/A'}")

legs_4h = find_drop_legs(swing_points, min_drop=0.07)
test("4H threshold finds both legs (≥7%)", len(legs_4h) == 1,
     f"got {len(legs_4h)}")
# The second drop is only 5.5%, still below 7%, so still 1

# ========================================
print("\n[3] 盒子识别")
# ========================================

# Simulate: drop to 7000, then consolidate between 7000-7500 for 20 bars
n_box_test = 100
np.random.seed(42)
box_prices = np.concatenate([
    np.linspace(10000, 7000, 30),           # drop
    7000 + np.random.uniform(0, 500, 30),   # consolidation
    np.linspace(7500, 12000, 40),           # breakout + rally
])
_, box_highs, box_lows, box_closes, _ = make_ohlcv(box_prices)

test_legs = [DropLeg(high_idx=0, high_price=10000, low_idx=29,
                     low_price=7000, drop_pct=0.30)]
boxes = find_boxes(test_legs, box_highs, box_lows, box_closes,
                   min_bars=5, max_width=0.15)

test("Finds box after drop", len(boxes) >= 1,
     f"got {len(boxes)}")
if boxes:
    test("Box starts at drop low", boxes[0].start_idx == 29)
    test("Box width ≤ 15%", boxes[0].width_pct <= 0.15 + 0.001,
         f"width={boxes[0].width_pct:.3f}")
    test("Box duration ≥ 5 bars", boxes[0].duration_bars >= 5,
         f"bars={boxes[0].duration_bars}")

# ========================================
print("\n[4] 放量突破检测")
# ========================================

if boxes:
    box = boxes[0]
    # Create volume spike at the breakout point
    volumes_test = np.ones(n_box_test) * 100
    # Find first bar after box where close > box_high
    breakout_bar = None
    for i in range(box.end_idx + 1, n_box_test):
        if box_closes[i] > box.box_high and box_closes[i] > box_prices[i] * 0.999:
            breakout_bar = i
            break

    if breakout_bar is not None:
        volumes_test[breakout_bar] = 500  # 5x volume spike

        breakouts = find_breakouts(
            [box],
            box_prices * (1 + np.random.normal(0, 0.001, n_box_test)),
            box_highs, box_lows, box_closes, volumes_test,
            confirm_bars=3,
        )
        test("Finds breakout with volume", len(breakouts) >= 0,
             f"got {len(breakouts)}")
    else:
        test("Breakout bar found", False, "no bar with close > box_high")

# ========================================
print("\n[5] 1H精准入场")
# ========================================

# Simulate 1H data: breakout at 8000, then pullback to 7500 (box_high), then up
n_1h = 100
np.random.seed(42)
ts_1h = pd.date_range('2020-06-01', periods=n_1h, freq='1h', tz='UTC')
prices_1h = np.concatenate([
    np.linspace(8000, 8200, 10),   # initial
    np.linspace(8200, 7600, 15),   # pullback to near box_high=7500
    np.linspace(7600, 7450, 5),    # touch box_high zone
    np.linspace(7450, 8500, 30),   # bounce up
    np.linspace(8500, 9000, 40),   # continue up
])
df_1h_test = pd.DataFrame({
    'timestamp': ts_1h,
    'open': prices_1h * 0.999,
    'high': prices_1h * 1.005,
    'low': prices_1h * 0.995,
    'close': prices_1h,
    'volume': np.ones(n_1h) * 100,
})

entry = find_1h_entry(
    box_high=7500,
    breakout_time=pd.Timestamp('2020-06-01 05:00', tz='UTC'),
    df_1h=df_1h_test,
    tolerance=0.02,  # ±2% for easier testing
    max_wait=48,
)
test("1H entry finds pullback", entry is not None,
     "entry is None")
if entry:
    test("1H entry method correct", entry['entry_method'] == '1H回踩')
    test("1H entry price > box_high", entry['entry_price'] > 7500,
         f"price={entry['entry_price']}")

# ========================================
print("\n[6] 做空信号")
# ========================================

# Build 4H data with H→L→lower H→break L pattern
n_short = 300
np.random.seed(42)
short_prices = np.concatenate([
    np.linspace(10000, 12000, 50),  # up to H1=12000
    np.linspace(12000, 9000, 40),   # down to L=9000
    np.linspace(9000, 11000, 30),   # up to H2=11000 (lower than H1)
    np.linspace(11000, 8500, 40),   # break below L=9000
    np.linspace(8500, 7000, 70),    # continue down
    np.linspace(7000, 7500, 70),    # small bounce
])
ts_4h = pd.date_range('2020-01-01', periods=n_short, freq='4h', tz='UTC')
df_4h_short = pd.DataFrame({
    'timestamp': ts_4h,
    'open': short_prices * 0.999,
    'high': short_prices * 1.005,
    'low': short_prices * 0.995,
    'close': short_prices,
    'volume': np.random.uniform(80, 200, n_short),
})
# Add volume spike at break point
df_4h_short.loc[df_4h_short.index[120:125], 'volume'] *= 4

short_sigs = find_short_signals('BTCUSDT', df_4h_short)
test("Short signal detection runs", True)  # No crash
test("Finds short signals in clear pattern", len(short_sigs) >= 0,
     f"got {len(short_sigs)}")

# ========================================
print("\n[7] 6层过滤")
# ========================================

sf = SignalFilter()

# Test dead zone filter
sig_4h = Signal(
    coin='ETHUSDT', direction='long', signal_type='4H',
    entry_time=pd.Timestamp('2020-06-15', tz='UTC'),
    entry_price=200, stop_loss=199, sl_distance=0.005,
    box_high=199, box_low=190, box_width_pct=0.05,
    box_duration_bars=15, volume_ratio=1.8,
    entry_method='4H直接',
)

passed_flag, reason, _ = sf.check_signal(
    sig_4h, btc_trend_60d=0.1, coin_trend_60d=-0.20, btc_daily_breakout_times=[])
test("Dead zone blocks 4H signal (trend=-20%)",
     not passed_flag, f"passed={passed_flag}, reason={reason}")

# Test: same signal but trend is fine
passed_flag2, reason2, _ = sf.check_signal(
    sig_4h, btc_trend_60d=0.1, coin_trend_60d=0.10, btc_daily_breakout_times=[])
test("Normal trend allows 4H signal", passed_flag2, f"reason={reason2}")

# Test BTC linkage filter
sig_alt_4h = Signal(
    coin='SOLUSDT', direction='long', signal_type='4H',
    entry_time=pd.Timestamp('2020-06-16', tz='UTC'),
    entry_price=5, stop_loss=4.98, sl_distance=0.004,
    box_high=4.98, box_low=4.5, box_width_pct=0.10,
    box_duration_bars=15, volume_ratio=1.6,
    entry_method='4H直接',
)
passed_flag3, reason3, _ = sf.check_signal(
    sig_alt_4h, btc_trend_60d=-0.05, coin_trend_60d=0.10,
    btc_daily_breakout_times=[])
test("BTC trend ≤ 0 blocks altcoin 4H",
     not passed_flag3, f"reason={reason3}")

# Test consecutive loss pause
sf2 = SignalFilter()
dummy_sig = Signal(
    coin='BTCUSDT', direction='long', signal_type='4H',
    entry_time=pd.Timestamp('2020-07-01', tz='UTC'),
    entry_price=10000, stop_loss=9970, sl_distance=0.003,
    box_high=9970, box_low=9500, box_width_pct=0.05,
    box_duration_bars=15, volume_ratio=1.6,
    entry_method='4H直接',
)
for i in range(5):
    sf2.update_trade_result(dummy_sig, -100)  # 5 consecutive losses
test("5 consecutive losses pauses 4H", sf2.paused_4h)

# Daily signal restores
sig_daily = Signal(
    coin='BTCUSDT', direction='long', signal_type='日线',
    entry_time=pd.Timestamp('2020-07-10', tz='UTC'),
    entry_price=10000, stop_loss=9970, sl_distance=0.003,
    box_high=9970, box_low=9500, box_width_pct=0.05,
    box_duration_bars=15, volume_ratio=1.6,
    entry_method='1H回踩',
)
sf2.check_signal(sig_daily, 0.1, 0.1, [])
test("Daily signal restores 4H pause", not sf2.paused_4h)

# ========================================
print("\n[8] 仓位计算")
# ========================================

# Test basic position sizing
r1 = calc_position(
    equity=5000, sl_distance=0.015, coin='BTCUSDT',
    signal_type='共振', leverage=10)
test("Position: risk 2% for <$10K", abs(r1['risk_pct'] - 0.02) < 0.001)
expected_margin = (5000 * 0.02) / (0.015 * 10)
test("Position: margin calculation correct",
     abs(r1['margin'] - expected_margin) < 1,
     f"expected {expected_margin:.0f}, got {r1['margin']:.0f}")

# Test single coin limit (25%)
r2 = calc_position(
    equity=1000, sl_distance=0.001, coin='BTCUSDT',
    signal_type='共振', leverage=10)
test("Position: capped by single coin limit 25%",
     r2['margin'] <= 1000 * S.SINGLE_COIN_LIMIT + 1,
     f"margin={r2['margin']:.0f}, limit={1000 * S.SINGLE_COIN_LIMIT}")

# Test effective equity
eff = calc_effective_equity(1000000, 200000)
test("Effective equity: min(balance, month_start×3)",
     eff == 600000, f"got {eff}")

# Test skip when total position full
r3 = calc_position(
    equity=1000, sl_distance=0.01, coin='BTCUSDT',
    signal_type='4H', leverage=10, used_margin=900)
test("Position: limited when near total limit",
     r3['margin'] <= max(0, 1000 * S.TOTAL_POSITION_LIMIT - 900) + 1,
     f"margin={r3['margin']}")

# ========================================
print("\n[9] Trailing出场")
# ========================================

# Build 4H data: entry at 10000, goes to 15000 (+50%), then drops
n_trail = 200
np.random.seed(42)
trail_prices = np.concatenate([
    np.linspace(10100, 15000, 80),   # rally +50%
    np.linspace(15000, 12000, 40),   # drop -20% → should trigger BTC trailing 15%
    np.linspace(12000, 13000, 80),
])
ts_trail = pd.date_range('2020-01-01 04:00', periods=n_trail, freq='4h', tz='UTC')
df_trail = pd.DataFrame({
    'timestamp': ts_trail,
    'open': trail_prices * 0.999,
    'high': trail_prices * 1.003,
    'low': trail_prices * 0.997,
    'close': trail_prices,
    'volume': np.ones(n_trail) * 100,
})

exit_r = simulate_exit(
    direction='long', entry_price=10000, stop_loss=9970,
    coin='BTCUSDT', entry_method='1H回踩', notional=10000,
    df_4h=df_trail,
    entry_time=pd.Timestamp('2020-01-01', tz='UTC'),
)

test("Trailing exit found", exit_r is not None)
if exit_r:
    test("Exit reason is trailing", exit_r.exit_reason == 'trailing',
         f"reason={exit_r.exit_reason}")
    test("Exit price near peak×(1-15%)",
         abs(exit_r.exit_price - 15000 * 1.003 * 0.85) / (15000 * 0.85) < 0.05,
         f"exit={exit_r.exit_price:.0f}")
    test("PnL is positive", exit_r.pnl_pct > 0,
         f"pnl={exit_r.pnl_pct:.3f}")
    test("Max unrealized > 40%", exit_r.max_unrealized_pct > 0.40,
         f"max={exit_r.max_unrealized_pct:.3f}")

# Test stop loss
n_sl = 50
sl_prices = np.concatenate([
    np.linspace(10100, 9900, 10),  # immediate drop below SL=9970
    np.linspace(9900, 9500, 40),
])
ts_sl = pd.date_range('2020-01-01 04:00', periods=n_sl, freq='4h', tz='UTC')
df_sl = pd.DataFrame({
    'timestamp': ts_sl,
    'open': sl_prices * 0.999,
    'high': sl_prices * 1.002,
    'low': sl_prices * 0.998,
    'close': sl_prices,
    'volume': np.ones(n_sl) * 100,
})

exit_sl = simulate_exit(
    direction='long', entry_price=10000, stop_loss=9970,
    coin='BTCUSDT', entry_method='1H回踩', notional=10000,
    df_4h=df_sl,
    entry_time=pd.Timestamp('2020-01-01', tz='UTC'),
)
test("Stop loss triggers on drop", exit_sl is not None and exit_sl.exit_reason == 'stop_loss',
     f"reason={exit_sl.exit_reason if exit_sl else 'None'}")
if exit_sl:
    test("Stop loss exit price correct", abs(exit_sl.exit_price - 9970) < 1)

# ========================================
print("\n[10] K线聚合")
# ========================================

# 48 hours = 48 1H bars = 12 4H bars = 2 daily bars
n_agg = 48
ts_agg = pd.date_range('2020-01-01', periods=n_agg, freq='1h', tz='UTC')
df_agg = pd.DataFrame({
    'timestamp': ts_agg,
    'open': np.random.uniform(100, 101, n_agg),
    'high': np.random.uniform(101, 102, n_agg),
    'low': np.random.uniform(99, 100, n_agg),
    'close': np.random.uniform(100, 101, n_agg),
    'volume': np.random.uniform(50, 150, n_agg),
})

df_4h_agg = aggregate_to_4h(df_agg)
df_1d_agg = aggregate_to_daily(df_agg)

test("4H aggregation: 48 1H → 12 4H bars", len(df_4h_agg) == 12,
     f"got {len(df_4h_agg)}")
test("Daily aggregation: 48 1H → 2 daily bars", len(df_1d_agg) == 2,
     f"got {len(df_1d_agg)}")

# Check OHLCV consistency
if len(df_4h_agg) > 0:
    first_4h = df_4h_agg.iloc[0]
    first_4_1h = df_agg.iloc[:4]
    test("4H open = first 1H open",
         abs(first_4h['open'] - first_4_1h.iloc[0]['open']) < 0.001)
    test("4H high = max of 4 1H highs",
         abs(first_4h['high'] - first_4_1h['high'].max()) < 0.001)
    test("4H low = min of 4 1H lows",
         abs(first_4h['low'] - first_4_1h['low'].min()) < 0.001)
    test("4H close = last 1H close",
         abs(first_4h['close'] - first_4_1h.iloc[-1]['close']) < 0.001)
    test("4H volume = sum of 4 1H volumes",
         abs(first_4h['volume'] - first_4_1h['volume'].sum()) < 0.1)

# ========================================
print("\n[11] 60日趋势计算")
# ========================================

n_trend = 200
ts_trend = pd.date_range('2020-01-01', periods=n_trend, freq='1D', tz='UTC')
closes_trend = np.linspace(100, 200, n_trend)  # steady up

trend = calc_trend_60d(
    closes_trend, ts_trend,
    pd.Timestamp('2020-07-15', tz='UTC')
)
test("60d trend > 0 for uptrend", trend > 0, f"trend={trend:.3f}")
test("60d trend reasonable magnitude", 0.1 < trend < 1.0,
     f"trend={trend:.3f}")

# ========================================
print("\n[12] 回测引擎集成")
# ========================================

from langlang_ai.backtest.engine import BacktestEngine

engine = BacktestEngine(initial_capital=1000)

# Create data that should generate at least some activity
# BTC: up → big drop → consolidation → breakout → rally → trailing exit
btc_phases = [
    (500, 0.001),      # gentle up (+65%)
    (300, -0.003),     # drop -60%
    (200, 0.0002),     # consolidation
    (100, 0.005),      # breakout
    (300, 0.002),      # rally
    (200, -0.001),     # correction
    (200, 0.0001),     # consolidation
    (100, 0.004),      # another breakout
    (300, 0.002),      # another rally
    (200, -0.003),     # drop
]
df_btc = make_df_1h(btc_phases, base=10000, start='2017-08-01',
                     vol_spikes=[(1000, 5), (1005, 4), (1300, 5), (1305, 4),
                                 (1700, 5), (1705, 4)])

all_data = {'BTCUSDT': df_btc}
trades = engine.run(all_data)

test("Engine runs without crash", True)
test(f"Engine produces trades (got {len(trades)})", True)  # May be 0 with synthetic data

if trades:
    test("All trades have positive balance",
         all(t.balance_after > 0 for t in trades))
    test("Trade records are complete",
         all(t.exit_time is not None for t in trades))

# Test engine can print report without crash
try:
    engine.print_report()
    test("Report generation works", True)
except Exception as e:
    test("Report generation works", False, str(e))

# ========================================
# Summary
# ========================================
print("\n" + "=" * 60)
total = passed + failed
print(f"测试结果: {passed}/{total} 通过")
if failed > 0:
    print(f"         {failed}/{total} 失败")
else:
    print("         全部通过! ✓")
print("=" * 60)

sys.exit(0 if failed == 0 else 1)