import json
import math

import numpy as np

from fxhnt.application.book_allocator import (
    align_edges,
    apply_drawdown_killswitch,
    combine_book,
    evaluate_book,
)


def test_drawdown_killswitch_flattens_and_reenters():
    # a sharp drawdown then recovery: killswitch flattens during the crash, re-enters after recovery
    book = {}
    d = 0
    for r in [0.01] * 20:           # ramp up (peak)
        book[d] = r; d += 1
    for r in [-0.05] * 5:           # -25% crash -> breach kill_dd=0.15
        book[d] = r; d += 1
    for r in [0.02] * 30:           # recovery
        book[d] = r; d += 1
    out = apply_drawdown_killswitch(book, kill_dd=0.15, reenter_dd=0.07)
    days = sorted(book)
    # no lookahead: identical until the kill triggers
    assert out[days[0]] == book[days[0]]
    # at least one day is flattened to 0 during/after the crash
    assert any(out[k] == 0.0 and book[k] != 0.0 for k in days)
    # after full recovery the switch re-enters (last day passes through)
    assert out[days[-1]] == book[days[-1]]


def _synth_edge(seed: int, n: int = 2000, drift: float = 0.0008, vol: float = 0.008,
                start: int = 1000) -> dict[int, float]:
    """A synthetic daily net-return series (positive drift + iid noise) keyed by epoch-day.

    Long sample + a drift comfortably above the noise floor so each edge has a similar,
    clearly-positive standalone Sharpe (the precondition for a diversification uplift — three
    independent same-Sharpe sleeves combine to a higher Sharpe by ~sqrt(n_edges))."""
    rng = np.random.default_rng(seed)
    rets = drift + vol * rng.standard_normal(n)
    return {start + i: float(rets[i]) for i in range(n)}


def test_align_edges_union_and_zero_fill():
    a = {1: 0.01, 2: 0.02, 3: 0.03}        # days 1-3
    b = {3: 0.10, 4: 0.20, 5: 0.30}        # days 3-5 (overlap on day 3, disjoint otherwise)
    dates, arrs = align_edges({"a": a, "b": b})
    assert dates == [1, 2, 3, 4, 5]        # sorted union
    assert np.allclose(arrs["a"], [0.01, 0.02, 0.03, 0.0, 0.0])   # 0.0-fill where absent
    assert np.allclose(arrs["b"], [0.0, 0.0, 0.10, 0.20, 0.30])
    assert set(arrs) == {"a", "b"}


def test_combine_book_hits_vol_target_and_diversification_uplift():
    edges = {
        "e1": _synth_edge(1),
        "e2": _synth_edge(2),
        "e3": _synth_edge(3),
    }
    target_vol = 0.12
    kelly = 0.5
    ppy = 365
    book = combine_book(edges, target_vol=target_vol, periods_per_year=ppy,
                        kelly_fraction=kelly, vol_lookback=60, rebalance_every=21)
    assert isinstance(book, dict)
    rets = np.array([book[d] for d in sorted(book)])
    # (2) book scaled toward the vol target — fractional-Kelly scales the vol-targeted exposure,
    # so the EFFECTIVE target the book is driven toward is kelly_fraction * target_vol.
    effective_target = kelly * target_vol
    realized_ann_vol = float(rets.std() * math.sqrt(ppy))
    assert abs(realized_ann_vol - effective_target) / effective_target < 0.40

    # diversification uplift: combined Sharpe >= max standalone Sharpe.
    def ann_sharpe(x: np.ndarray) -> float:
        x = np.asarray(x, float)
        return float(x.mean() / x.std() * math.sqrt(ppy)) if x.std() > 0 else 0.0

    combined_sr = ann_sharpe(rets)
    standalone = [ann_sharpe(np.array(list(e.values()))) for e in edges.values()]
    assert combined_sr >= max(standalone) - 1e-9


def test_combine_book_deterministic():
    edges = {"e1": _synth_edge(1), "e2": _synth_edge(2), "e3": _synth_edge(3)}
    b1 = combine_book(edges)
    b2 = combine_book(edges)
    assert b1.keys() == b2.keys()
    assert all(b1[d] == b2[d] for d in b1)


def test_anticorrelated_edges_reduce_vol():
    rng = np.random.default_rng(7)
    n = 800
    base = 0.0003 + 0.012 * rng.standard_normal(n)
    a = {1000 + i: float(0.0003 + base[i]) for i in range(n)}
    b = {1000 + i: float(0.0003 - base[i]) for i in range(n)}   # anti-correlated to a
    book = combine_book({"a": a, "b": b}, target_vol=10.0, vol_lookback=60, rebalance_every=21)
    # With a hard leverage cap (no vol-target inflation possible at this absurd target), the
    # combined book of two anti-correlated sleeves has far lower vol than either standalone.
    rets = np.array([book[d] for d in sorted(book)])
    sa = np.array(list(a.values()))
    sb = np.array(list(b.values()))
    assert rets.std() < 0.5 * min(sa.std(), sb.std())


def test_evaluate_book_report_is_json_safe():
    edges = {"e1": _synth_edge(1), "e2": _synth_edge(2), "e3": _synth_edge(3)}
    report = evaluate_book(edges, periods_per_year=365, oos_fraction=0.40)
    assert {"stats", "verdict", "correlation_matrix", "standalone_sharpe"} <= set(report)
    assert {"cagr", "ann_vol", "sharpe", "max_drawdown", "n_obs"} <= set(report["stats"])
    assert {"passed", "dsr", "is_sharpe", "oos_sharpe", "n_trials"} <= set(report["verdict"])
    # correlation matrix is pairwise over the edges
    cm = report["correlation_matrix"]
    assert set(cm) == {"e1", "e2", "e3"}
    for k in cm:
        assert set(cm[k]) == {"e1", "e2", "e3"}
        assert abs(cm[k][k] - 1.0) < 1e-9     # self-correlation = 1
    assert set(report["standalone_sharpe"]) == {"e1", "e2", "e3"}
    json.dumps(report)                          # must be JSON-serializable (no NaN/np types)