"""
Test all Python code snippets from docs/hrp-denoising.mdx
"""
import traceback

import horizon as hz

results = []


def run_snippet(name, fn):
    """Run a snippet function, print PASS/FAIL, record result."""
    try:
        fn()
        print(f"PASS: {name}")
        results.append((name, True, None))
    except Exception as e:
        tb = traceback.format_exc()
        print(f"FAIL: {name}")
        print(f"  Error: {e}")
        print(f"  {tb}")
        results.append((name, False, str(e)))


# ---------------------------------------------------------------------------
# Snippet 1: Basic hz.hrp_weights (3x3 covariance)
# docs/hrp-denoising.mdx lines 46-58
# ---------------------------------------------------------------------------
def snippet_1_basic_hrp():
    cov = [
        [0.04, 0.02, 0.01],
        [0.02, 0.09, 0.03],
        [0.01, 0.03, 0.16],
    ]

    result = hz.hrp_weights(cov)
    print(result.weights)         # e.g., [0.52, 0.30, 0.18]
    print(result.sorted_indices)  # Quasi-diagonal ordering, e.g., [0, 1, 2]
    print(result.linkage)         # Clustering linkage: [(cluster_a, cluster_b, distance), ...]

    # Validate structure
    assert len(result.weights) == 3, f"Expected 3 weights, got {len(result.weights)}"
    assert abs(sum(result.weights) - 1.0) < 1e-9, f"Weights don't sum to 1: {sum(result.weights)}"
    assert len(result.sorted_indices) == 3, f"Expected 3 sorted_indices"
    assert len(result.linkage) == 2, f"Expected 2 linkage entries for 3 assets, got {len(result.linkage)}"
    assert all(w > 0 for w in result.weights), "All weights should be positive"


run_snippet("Snippet 1: Basic hz.hrp_weights (3x3)", snippet_1_basic_hrp)


# ---------------------------------------------------------------------------
# Snippet 2: Block structure HRP (4x4 covariance)
# docs/hrp-denoising.mdx lines 83-101
# ---------------------------------------------------------------------------
def snippet_2_block_hrp():
    # Two low-variance assets and two high-variance assets
    # with block correlation structure
    cov = [
        [0.04, 0.03, 0.001, 0.001],  # asset 0: low var, correlated with 1
        [0.03, 0.04, 0.001, 0.001],  # asset 1: low var, correlated with 0
        [0.001, 0.001, 0.09, 0.06],  # asset 2: high var, correlated with 3
        [0.001, 0.001, 0.06, 0.09],  # asset 3: high var, correlated with 2
    ]

    result = hz.hrp_weights(cov)
    print(result.weights)
    # Low-var block (assets 0,1) gets more total weight than high-var block (assets 2,3)
    low_var_total = result.weights[0] + result.weights[1]
    high_var_total = result.weights[2] + result.weights[3]
    print(f"Low-var block: {low_var_total:.2%}")
    print(f"High-var block: {high_var_total:.2%}")

    # Validate
    assert abs(sum(result.weights) - 1.0) < 1e-9, "Weights don't sum to 1"
    assert low_var_total > high_var_total, (
        f"Low-var block ({low_var_total:.4f}) should have more weight than "
        f"high-var block ({high_var_total:.4f})"
    )


run_snippet("Snippet 2: Block structure HRP (4x4)", snippet_2_block_hrp)


# ---------------------------------------------------------------------------
# Snippet 3: hz.denoise_covariance
# docs/hrp-denoising.mdx lines 114-127
# ---------------------------------------------------------------------------
def snippet_3_denoise():
    cov = [
        [0.04, 0.02, 0.01],
        [0.02, 0.09, 0.03],
        [0.01, 0.03, 0.16],
    ]

    denoised = hz.denoise_covariance(cov, n_observations=100)
    print(denoised.covariance)     # Cleaned N x N covariance matrix
    print(denoised.eigenvalues)    # Eigenvalues (sorted descending)
    print(denoised.n_signals)      # Number of signal eigenvalues (above MP threshold)
    print(denoised.n_noise)        # Number of noise eigenvalues (below MP threshold)

    # Validate
    assert len(denoised.covariance) == 3, "Denoised cov should be 3x3"
    assert all(len(row) == 3 for row in denoised.covariance), "Each row should have 3 elements"
    assert len(denoised.eigenvalues) == 3, "Should have 3 eigenvalues"
    assert denoised.n_signals + denoised.n_noise == 3, "Signal + noise should equal N"
    assert denoised.n_signals >= 0, "n_signals should be non-negative"
    assert denoised.n_noise >= 0, "n_noise should be non-negative"
    # Eigenvalues should be sorted descending
    for i in range(len(denoised.eigenvalues) - 1):
        assert denoised.eigenvalues[i] >= denoised.eigenvalues[i + 1], "Eigenvalues not sorted descending"


run_snippet("Snippet 3: hz.denoise_covariance", snippet_3_denoise)


# ---------------------------------------------------------------------------
# Snippet 4: hz.detone_covariance
# docs/hrp-denoising.mdx lines 162-172
# ---------------------------------------------------------------------------
def snippet_4_detone():
    cov = [
        [0.04, 0.02, 0.01],
        [0.02, 0.09, 0.03],
        [0.01, 0.03, 0.16],
    ]

    detoned = hz.detone_covariance(cov)
    # detoned is a list[list[float]] -- the covariance matrix without the market mode

    # Validate structure
    assert isinstance(detoned, list), f"Expected list, got {type(detoned)}"
    assert len(detoned) == 3, f"Expected 3 rows, got {len(detoned)}"
    assert all(len(row) == 3 for row in detoned), "Each row should have 3 elements"

    # Trace of detoned should be smaller than original
    orig_trace = 0.04 + 0.09 + 0.16
    detoned_trace = sum(detoned[i][i] for i in range(3))
    print(f"Original trace: {orig_trace}")
    print(f"Detoned trace: {detoned_trace}")
    assert detoned_trace < orig_trace, (
        f"Detoned trace ({detoned_trace}) should be smaller than original ({orig_trace})"
    )


run_snippet("Snippet 4: hz.detone_covariance", snippet_4_detone)


# ---------------------------------------------------------------------------
# Snippet 5: Full pipeline (sample covariance -> denoise -> HRP)
# docs/hrp-denoising.mdx lines 191-225
# ---------------------------------------------------------------------------
def snippet_5_full_pipeline():
    # 1. Compute sample covariance from returns
    # returns_matrix: T observations x N assets
    returns_matrix = [
        [0.01, -0.02, 0.005, 0.008],
        [0.02,  0.01, -0.01, 0.003],
        [-0.01, 0.03, 0.02, -0.005],
        [0.005, -0.01, 0.01, 0.012],
        # ... more observations
    ]

    T = len(returns_matrix)
    N = len(returns_matrix[0])

    # Compute sample covariance (or use numpy)
    means = [sum(r[j] for r in returns_matrix) / T for j in range(N)]
    cov = [[0.0] * N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            cov[i][j] = sum(
                (r[i] - means[i]) * (r[j] - means[j])
                for r in returns_matrix
            ) / (T - 1)

    # 2. Denoise the covariance matrix
    denoised = hz.denoise_covariance(cov, n_observations=T)
    print(f"Signal eigenvalues: {denoised.n_signals}")
    print(f"Noise eigenvalues: {denoised.n_noise}")

    # 3. Compute HRP weights on the denoised covariance
    result = hz.hrp_weights(denoised.covariance)
    for i, w in enumerate(result.weights):
        print(f"Asset {i}: {w:.2%}")

    # Validate
    assert len(result.weights) == N, f"Expected {N} weights"
    assert abs(sum(result.weights) - 1.0) < 1e-9, "Weights don't sum to 1"


run_snippet("Snippet 5: Full pipeline (cov -> denoise -> HRP)", snippet_5_full_pipeline)


# ---------------------------------------------------------------------------
# Snippet 6: Pipeline with detoning
# docs/hrp-denoising.mdx lines 230-254
# ---------------------------------------------------------------------------
def snippet_6_pipeline_detoning():
    # For clustering purposes, detone first to remove market factor,
    # then run HRP on the original (or denoised) covariance

    cov = [
        [0.04, 0.03, 0.001, 0.001],
        [0.03, 0.04, 0.001, 0.001],
        [0.001, 0.001, 0.09, 0.06],
        [0.001, 0.001, 0.06, 0.09],
    ]

    # Option A: Standard HRP
    standard = hz.hrp_weights(cov)

    # Option B: Denoise then HRP
    denoised = hz.denoise_covariance(cov, n_observations=200)
    denoised_hrp = hz.hrp_weights(denoised.covariance)

    # Option C: Detone (for analysis)
    detoned = hz.detone_covariance(cov)

    print("Standard HRP weights:", [f"{w:.3f}" for w in standard.weights])
    print("Denoised HRP weights:", [f"{w:.3f}" for w in denoised_hrp.weights])

    # Validate
    assert abs(sum(standard.weights) - 1.0) < 1e-9
    assert abs(sum(denoised_hrp.weights) - 1.0) < 1e-9
    assert isinstance(detoned, list)
    assert len(detoned) == 4


run_snippet("Snippet 6: Pipeline with detoning", snippet_6_pipeline_detoning)


# ---------------------------------------------------------------------------
# Snippet 7: Combining with fractional differentiation
# docs/hrp-denoising.mdx lines 261-302
# Uses synthetic price data since the doc has [...] placeholders
# ---------------------------------------------------------------------------
def snippet_7_frac_diff_pipeline():
    import math
    import random

    random.seed(42)

    # Generate synthetic price series (geometric random walk)
    def make_prices(n=200, start=100.0, drift=0.0001, vol=0.02):
        prices = [start]
        for _ in range(n - 1):
            ret = drift + vol * random.gauss(0, 1)
            prices.append(prices[-1] * math.exp(ret))
        return prices

    # Multiple asset price series
    asset_prices = {
        "btc": make_prices(200, 50000, 0.0002, 0.03),
        "eth": make_prices(200, 3000, 0.0001, 0.04),
        "sol": make_prices(200, 100, 0.0003, 0.05),
    }

    # Fractionally differentiate each series
    stationary = {}
    for name, prices in asset_prices.items():
        d_star, _ = hz.min_frac_diff(prices, max_d=1.0, n_steps=20)
        stationary[name] = hz.frac_diff_ffd(prices, d=d_star)
        print(f"{name}: d*={d_star:.3f}, length={len(stationary[name])}")

    # Trim to common length
    min_len = min(len(s) for s in stationary.values())
    names = list(stationary.keys())
    returns_matrix = []
    for t in range(min_len):
        row = [stationary[name][t] for name in names]
        returns_matrix.append(row)

    # Compute covariance and run HRP
    T = len(returns_matrix)
    N = len(names)
    means = [sum(r[j] for r in returns_matrix) / T for j in range(N)]
    cov = [[0.0] * N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            cov[i][j] = sum(
                (r[i] - means[i]) * (r[j] - means[j])
                for r in returns_matrix
            ) / (T - 1)

    denoised = hz.denoise_covariance(cov, n_observations=T)
    result = hz.hrp_weights(denoised.covariance)

    for name, w in zip(names, result.weights):
        print(f"{name}: {w:.2%}")

    # Validate
    assert len(result.weights) == N, f"Expected {N} weights"
    assert abs(sum(result.weights) - 1.0) < 1e-9, "Weights don't sum to 1"
    assert all(w > 0 for w in result.weights), "All weights should be positive"


run_snippet("Snippet 7: Frac diff + denoise + HRP pipeline", snippet_7_frac_diff_pipeline)


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)

passed = sum(1 for _, ok, _ in results if ok)
failed = sum(1 for _, ok, _ in results if not ok)
total = len(results)

for name, ok, err in results:
    status = "PASS" if ok else "FAIL"
    line = f"  [{status}] {name}"
    if err:
        line += f"  -- {err}"
    print(line)

print(f"\nTotal: {total}  |  Passed: {passed}  |  Failed: {failed}")
if failed == 0:
    print("All snippets passed!")
else:
    print(f"{failed} snippet(s) failed.")
