from __future__ import annotations

import argparse
import math
import random
import sys
import time
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Sequence, Tuple

import numpy as np
from scipy import stats as sp_stats

from .simulator import PROFILES, BehaviorProfile, run_simulation

Shuffler = Callable[[List[int], random.Random], List[int]]


def _fisher_yates(items: List[int], rng: random.Random) -> List[int]:
    a = items[:]
    for i in range(len(a) - 1, 0, -1):
        j = rng.randint(0, i)
        a[i], a[j] = a[j], a[i]
    return a


def _sattolo(items: List[int], rng: random.Random) -> List[int]:
    a = items[:]
    for i in range(len(a) - 1, 0, -1):
        j = rng.randint(0, i - 1)
        a[i], a[j] = a[j], a[i]
    return a


def _knuth_buffer(items: List[int], rng: random.Random, k: int = 8) -> List[int]:
    buf = items[:k]
    out: List[int] = []
    for item in items[k:]:
        j = rng.randrange(k)
        out.append(buf[j])
        buf[j] = item
    rng.shuffle(buf)
    out.extend(buf)
    return out


def _riffle_gsr(items: List[int], rng: random.Random, riffles: int = 7) -> List[int]:
    a = items[:]
    n = len(a)
    for _ in range(riffles):
        cut = sum(1 for _ in range(n) if rng.random() < 0.5)
        left, right = a[:cut], a[cut:]
        merged: List[int] = []
        li = ri = 0
        while li < len(left) and ri < len(right):
            total = len(left) - li + len(right) - ri
            if rng.random() < (len(left) - li) / total:
                merged.append(left[li])
                li += 1
            else:
                merged.append(right[ri])
                ri += 1
        merged.extend(left[li:])
        merged.extend(right[ri:])
        a = merged
    return a


def _identity(items: List[int], _rng: random.Random) -> List[int]:
    return items[:]


def _make_roundabout_shuffler(profile: BehaviorProfile) -> Shuffler:
    def _shuffler(items: List[int], rng: random.Random) -> List[int]:
        seed = rng.randint(0, 2**31)
        result = run_simulation(items, seed, profile, record_history=False)
        return result.shuffled

    _shuffler.__name__ = f"roundabout_{profile.name}"
    return _shuffler


def _kendall_tau_normalized(perm: List[int]) -> float:
    n = len(perm)
    if n < 2:
        return 0.0
    inv = 0
    for i in range(n):
        for j in range(i + 1, n):
            if perm[i] > perm[j]:
                inv += 1
    return inv / (n * (n - 1) / 2)


def _spearman_footrule_normalized(perm: List[int]) -> float:
    n = len(perm)
    if n < 2:
        return 0.0
    total = sum(abs(perm[i] - i) for i in range(n))
    if n % 2 == 0:
        max_d = n * n / 2
    else:
        max_d = (n * n - 1) / 2
    return total / max_d if max_d else 0.0


def _cayley_distance(perm: List[int]) -> int:
    n = len(perm)
    visited = [False] * n
    cycles = 0
    for i in range(n):
        if visited[i]:
            continue
        cycles += 1
        j = i
        while not visited[j]:
            visited[j] = True
            j = perm[j]
    return n - cycles


def _fixed_points(perm: List[int]) -> int:
    return sum(1 for i, v in enumerate(perm) if v == i)


def _cycle_lengths(perm: List[int]) -> List[int]:
    n = len(perm)
    visited = [False] * n
    lengths: List[int] = []
    for i in range(n):
        if visited[i]:
            continue
        length = 0
        j = i
        while not visited[j]:
            visited[j] = True
            j = perm[j]
            length += 1
        lengths.append(length)
    return lengths


def _lag1_autocorrelation(seq: List[int]) -> float:
    if len(seq) < 3:
        return 0.0
    a = np.array(seq, dtype=float)
    var = a.var()
    if var == 0:
        return 0.0
    return float(np.corrcoef(a[:-1], a[1:])[0, 1])


def _runs_test_pvalue(seq: List[int]) -> float:
    if len(seq) < 10:
        return 1.0
    med = float(np.median(seq))
    binary = [1 if x > med else 0 for x in seq]
    binary = [b for b, x in zip(binary, seq) if x != med]
    if len(binary) < 10:
        return 1.0
    n1 = sum(binary)
    n0 = len(binary) - n1
    if n0 == 0 or n1 == 0:
        return 1.0
    runs = 1 + sum(1 for i in range(1, len(binary)) if binary[i] != binary[i - 1])
    n = n0 + n1
    mu_r = 1 + 2 * n0 * n1 / n
    var_r = 2 * n0 * n1 * (2 * n0 * n1 - n) / (n * n * (n - 1))
    if var_r <= 0:
        return 1.0
    z = (runs - mu_r) / math.sqrt(var_r)
    return float(2 * (1 - sp_stats.norm.cdf(abs(z))))


def _bit_bias(seq: List[int]) -> float:
    if len(seq) < 2:
        return 0.0
    ones = sum(1 for a, b in zip(seq, seq[1:]) if b > a)
    return ones / (len(seq) - 1) - 0.5


def _block_entropy(seq: List[int], block: int) -> float:
    if len(seq) < block + 1:
        return 0.0
    bits = [1 if b > a else 0 for a, b in zip(seq, seq[1:])]
    counts: Dict[tuple, int] = defaultdict(int)
    for i in range(len(bits) - block + 1):
        counts[tuple(bits[i : i + block])] += 1
    total = sum(counts.values())
    entropy = 0.0
    for c in counts.values():
        p = c / total
        if p > 0:
            entropy -= p * math.log2(p)
    return entropy / block


def _displacement_cdf(perm: List[int]) -> np.ndarray:
    return np.array([perm[i] - i for i in range(len(perm))], dtype=float)


def _build_position_matrix(perms: List[List[int]], n: int) -> np.ndarray:
    mat = np.zeros((n, n), dtype=float)
    for perm in perms:
        for pos, val in enumerate(perm):
            mat[pos][val] += 1
    return mat


def _chi_sq_uniformity(mat: np.ndarray, n_trials: int) -> Tuple[float, float]:
    n = mat.shape[0]
    expected = n_trials / n
    if expected == 0:
        return 0.0, 1.0
    chi2_total = 0.0
    df_total = 0
    for row in mat:
        chi2_total += float(np.sum((row - expected) ** 2 / expected))
        df_total += n - 1
    p = float(sp_stats.chi2.sf(chi2_total, df_total))
    return chi2_total, p


def _total_variation_from_uniform(mat: np.ndarray, n_trials: int) -> float:
    n = mat.shape[0]
    uniform = n_trials / n
    tvs: List[float] = []
    for row in mat:
        tv = float(0.5 * np.sum(np.abs(row - uniform)) / n_trials)
        tvs.append(tv)
    return float(np.mean(tvs))


@dataclass
class TrialMetrics:
    kendall_tau: float
    spearman_footrule: float
    cayley_distance: int
    fixed_points: int
    mean_cycle_length: float
    lag1_autocorr: float
    runs_pvalue: float
    bit_bias: float
    block_entropy_2: float
    block_entropy_4: float
    mean_abs_displacement: float
    wall_seconds: float


@dataclass
class AggregateResult:
    name: str
    n_trials: int
    n: int
    kendall_tau: float
    spearman_footrule: float
    cayley_distance: float
    fixed_point_frac: float
    mean_cycle_length: float
    lag1_autocorr: float
    runs_pvalue_median: float
    bit_bias: float
    block_entropy_2: float
    block_entropy_4: float
    mean_abs_displacement: float
    chi2_stat: float
    chi2_pvalue: float
    tv_distance: float
    ks_stat_vs_fy: float
    ks_pvalue_vs_fy: float
    wall_seconds_mean: float
    kendall_tau_ci: Tuple[float, float] = (0.0, 0.0)
    spearman_footrule_ci: Tuple[float, float] = (0.0, 0.0)
    bit_bias_ci: Tuple[float, float] = (0.0, 0.0)


def _ci95(values: Sequence[float]) -> Tuple[float, float]:
    if len(values) < 2:
        m = float(np.mean(values))
        return (m, m)
    a = np.array(values)
    m = float(a.mean())
    se = float(a.std(ddof=1) / math.sqrt(len(a)))
    return (m - 1.96 * se, m + 1.96 * se)


def benchmark_shuffler(
    name: str,
    shuffler: Shuffler,
    n: int,
    n_trials: int,
    base_seed: int,
    fy_displacements: Optional[np.ndarray] = None,
) -> AggregateResult:
    items = list(range(n))
    trials: List[TrialMetrics] = []
    perms: List[List[int]] = []
    all_displacements: List[float] = []

    for t in range(n_trials):
        rng = random.Random(base_seed + t)
        t0 = time.perf_counter()
        perm = shuffler(items, rng)
        elapsed = time.perf_counter() - t0

        cycles = _cycle_lengths(perm)
        disp = _displacement_cdf(perm)
        all_displacements.extend(disp.tolist())

        trials.append(
            TrialMetrics(
                kendall_tau=_kendall_tau_normalized(perm),
                spearman_footrule=_spearman_footrule_normalized(perm),
                cayley_distance=_cayley_distance(perm),
                fixed_points=_fixed_points(perm),
                mean_cycle_length=float(np.mean(cycles)),
                lag1_autocorr=_lag1_autocorrelation(perm),
                runs_pvalue=_runs_test_pvalue(perm),
                bit_bias=_bit_bias(perm),
                block_entropy_2=_block_entropy(perm, 2),
                block_entropy_4=_block_entropy(perm, 4),
                mean_abs_displacement=float(np.mean(np.abs(disp))),
                wall_seconds=elapsed,
            )
        )
        perms.append(perm)

    pos_mat = _build_position_matrix(perms, n)
    chi2_stat, chi2_p = _chi_sq_uniformity(pos_mat, n_trials)
    tv = _total_variation_from_uniform(pos_mat, n_trials)

    disp_arr = np.array(all_displacements)
    if fy_displacements is not None and len(fy_displacements) > 0:
        ks_result = sp_stats.ks_2samp(disp_arr, fy_displacements)
        ks_stat = float(ks_result.statistic)
        ks_p = float(ks_result.pvalue)
    else:
        ks_stat = 0.0
        ks_p = 1.0

    kt_vals = [t.kendall_tau for t in trials]
    sf_vals = [t.spearman_footrule for t in trials]
    bb_vals = [t.bit_bias for t in trials]

    return AggregateResult(
        name=name,
        n_trials=n_trials,
        n=n,
        kendall_tau=float(np.mean(kt_vals)),
        spearman_footrule=float(np.mean(sf_vals)),
        cayley_distance=float(np.mean([t.cayley_distance for t in trials])),
        fixed_point_frac=float(np.mean([t.fixed_points for t in trials])) / n,
        mean_cycle_length=float(np.mean([t.mean_cycle_length for t in trials])),
        lag1_autocorr=float(np.mean([t.lag1_autocorr for t in trials])),
        runs_pvalue_median=float(np.median([t.runs_pvalue for t in trials])),
        bit_bias=float(np.mean(bb_vals)),
        block_entropy_2=float(np.mean([t.block_entropy_2 for t in trials])),
        block_entropy_4=float(np.mean([t.block_entropy_4 for t in trials])),
        mean_abs_displacement=float(np.mean([t.mean_abs_displacement for t in trials])),
        chi2_stat=chi2_stat,
        chi2_pvalue=chi2_p,
        tv_distance=tv,
        ks_stat_vs_fy=ks_stat,
        ks_pvalue_vs_fy=ks_p,
        wall_seconds_mean=float(np.mean([t.wall_seconds for t in trials])),
        kendall_tau_ci=_ci95(kt_vals),
        spearman_footrule_ci=_ci95(sf_vals),
        bit_bias_ci=_ci95(bb_vals),
    )


_HEADER_FMT = (
    "{name:<24s} {kt:>8s} {sf:>8s} {cay:>6s} {fp:>6s} {mcl:>6s} "
    "{ac:>7s} {runs:>7s} {bb:>8s} {be2:>6s} {be4:>6s} "
    "{mad:>6s} {chi2p:>8s} {tv:>7s} {ks:>7s} {wt:>9s}"
)
_ROW_FMT = (
    "{name:<24s} {kt:>8.4f} {sf:>8.4f} {cay:>6.2f} {fp:>6.3f} {mcl:>6.2f} "
    "{ac:>+7.3f} {runs:>7.3f} {bb:>+8.4f} {be2:>6.3f} {be4:>6.3f} "
    "{mad:>6.2f} {chi2p:>8.4f} {tv:>7.4f} {ks:>7.4f} {wt:>9.5f}"
)


def _print_header() -> None:
    print(
        _HEADER_FMT.format(
            name="Algorithm",
            kt="Kend-t",
            sf="Spear-F",
            cay="Cayl",
            fp="FP%",
            mcl="MCycL",
            ac="Lag1AC",
            runs="RunsP",
            bb="BitBias",
            be2="BE2",
            be4="BE4",
            mad="MAD",
            chi2p="Chi2-P",
            tv="TV",
            ks="KS-FY",
            wt="Wall(s)",
        )
    )
    print("-" * 155)


def _print_row(r: AggregateResult) -> None:
    print(
        _ROW_FMT.format(
            name=r.name,
            kt=r.kendall_tau,
            sf=r.spearman_footrule,
            cay=r.cayley_distance,
            fp=r.fixed_point_frac,
            mcl=r.mean_cycle_length,
            ac=r.lag1_autocorr,
            runs=r.runs_pvalue_median,
            bb=r.bit_bias,
            be2=r.block_entropy_2,
            be4=r.block_entropy_4,
            mad=r.mean_abs_displacement,
            chi2p=r.chi2_pvalue,
            tv=r.tv_distance,
            ks=r.ks_stat_vs_fy,
            wt=r.wall_seconds_mean,
        )
    )


def _print_interpretation(results: List[AggregateResult]) -> None:
    fy = next((r for r in results if r.name == "Fisher-Yates"), None)
    if fy is None:
        return

    print("\n--- Interpretation Guide ---")
    print("An ideal uniform shuffle matches Fisher-Yates on all metrics:")
    print(f"  Kendall-tau ~ {fy.kendall_tau:.4f}  (fraction of max inversions)")
    print(f"  Spearman footrule ~ {fy.spearman_footrule:.4f}")
    print(
        f"  Fixed-point fraction ~ {fy.fixed_point_frac:.3f}  (expect ~1/n = {1 / fy.n:.3f})"
    )
    print("  Bit bias ~ 0.0000  (unbiased successive comparison)")
    print("  Block entropy (k=2) ~ 1.000  (1 bit per symbol = maximum)")
    print("  Chi2-P > 0.05  (position frequencies indistinguishable from uniform)")
    print("  TV distance ~ 0  (total variation from uniform)")
    print("  KS-FY ~ 0  (displacement distribution matches Fisher-Yates)")
    print("  Runs-P > 0.05  (no serial dependence)")
    print("  Lag-1 autocorrelation ~ 0")

    print("\nIdentity (no shuffle) is the degenerate baseline:")
    ident = next((r for r in results if r.name == "Identity"), None)
    if ident:
        print(f"  Kendall-tau = {ident.kendall_tau:.4f}, TV = {ident.tv_distance:.4f}")

    print("\nRoundabout profiles fall on a spectrum between these extremes.")
    for r in results:
        if not r.name.startswith("roundabout_"):
            continue
        quality = r.kendall_tau / fy.kendall_tau if fy.kendall_tau else 0
        uniformity = "PASS" if r.chi2_pvalue > 0.05 else "FAIL"
        serial = "PASS" if r.runs_pvalue_median > 0.05 else "FAIL"
        print(
            f"  {r.name}: mixing={quality:.1%} of F-Y | "
            f"uniformity={uniformity} (p={r.chi2_pvalue:.3f}) | "
            f"independence={serial} (runs-p={r.runs_pvalue_median:.3f}) | "
            f"bit-bias={r.bit_bias:+.4f}"
        )


def _plot_results(results: List[AggregateResult], output: str) -> str:
    import matplotlib.pyplot as plt

    names = [r.name for r in results]
    n_alg = len(names)
    x = np.arange(n_alg)

    fig, axes = plt.subplots(3, 3, figsize=(18, 13))
    fig.patch.set_facecolor("#fafaf8")
    fig.suptitle(
        f"Roundabout Shuffler Benchmark  (n={results[0].n}, trials={results[0].n_trials})",
        fontsize=15,
        weight="bold",
        y=0.98,
    )

    bar_kw = dict(width=0.52, edgecolor="#1f2937", linewidth=0.6)
    colors = []
    for name in names:
        if name == "Fisher-Yates":
            colors.append("#16a34a")
        elif name == "Identity":
            colors.append("#9ca3af")
        elif name.startswith("roundabout_"):
            colors.append("#2563eb")
        else:
            colors.append("#d97706")

    def _bar(ax, values, title, ylabel, ideal=None):
        ax.bar(x, values, color=colors, **bar_kw)
        ax.set_xticks(x)
        ax.set_xticklabels(names, rotation=38, ha="right", fontsize=8)
        ax.set_title(title, fontsize=11, weight="bold")
        ax.set_ylabel(ylabel, fontsize=9)
        if ideal is not None:
            ax.axhline(
                ideal,
                color="#dc2626",
                ls="--",
                lw=1.0,
                alpha=0.7,
                label=f"ideal={ideal}",
            )
            ax.legend(fontsize=8)
        ax.grid(axis="y", alpha=0.3)

    _bar(
        axes[0, 0],
        [r.kendall_tau for r in results],
        "Kendall-tau (normalised)",
        "fraction of max inversions",
    )
    _bar(
        axes[0, 1],
        [r.spearman_footrule for r in results],
        "Spearman Footrule (normalised)",
        "fraction of max displacement",
    )
    ideal_fp = 1.0 / results[0].n if results[0].n else 0
    _bar(
        axes[0, 2],
        [r.fixed_point_frac for r in results],
        "Fixed-Point Fraction",
        "fraction",
        ideal=ideal_fp,
    )
    _bar(
        axes[1, 0],
        [r.bit_bias for r in results],
        "Bit Bias (successive comparison)",
        "bias from 0.5",
        ideal=0.0,
    )
    _bar(
        axes[1, 1],
        [r.block_entropy_2 for r in results],
        "Block Entropy (k=2)",
        "bits/symbol",
        ideal=1.0,
    )
    _bar(
        axes[1, 2],
        [r.tv_distance for r in results],
        "Total Variation from Uniform",
        "TV distance",
        ideal=0.0,
    )
    _bar(
        axes[2, 0],
        [r.lag1_autocorr for r in results],
        "Lag-1 Autocorrelation",
        "correlation",
        ideal=0.0,
    )
    _bar(
        axes[2, 1],
        [r.ks_stat_vs_fy for r in results],
        "KS Distance vs Fisher-Yates",
        "KS statistic",
        ideal=0.0,
    )
    _bar(
        axes[2, 2], [r.wall_seconds_mean for r in results], "Mean Wall Time", "seconds"
    )

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.savefig(output, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return output


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Benchmark roundabout shuffler vs reference algorithms"
    )
    parser.add_argument("--size", type=int, default=16)
    parser.add_argument("--trials", type=int, default=300)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--no-plot", action="store_true")
    parser.add_argument("--output", type=str, default="benchmark.png")
    parser.add_argument(
        "--profiles",
        nargs="*",
        default=None,
        help="Roundabout profiles to test (default: all)",
    )
    args = parser.parse_args()

    n = args.size
    n_trials = args.trials
    base_seed = args.seed

    shufflers: Dict[str, Shuffler] = {
        "Fisher-Yates": _fisher_yates,
        "Sattolo": _sattolo,
        "Knuth-Buffer(k=8)": lambda items, rng: _knuth_buffer(items, rng, k=8),
        "Riffle(7)": lambda items, rng: _riffle_gsr(items, rng, riffles=7),
        "Riffle(3)": lambda items, rng: _riffle_gsr(items, rng, riffles=3),
        "Identity": _identity,
    }

    profile_names = args.profiles if args.profiles else list(PROFILES.keys())
    for pname in profile_names:
        if pname not in PROFILES:
            print(f"WARNING: unknown profile '{pname}', skipping", file=sys.stderr)
            continue
        shufflers[f"roundabout_{pname}"] = _make_roundabout_shuffler(PROFILES[pname])

    print(f"Benchmarking {len(shufflers)} algorithms, n={n}, trials={n_trials}")
    print()

    fy_displacements: Optional[np.ndarray] = None
    results: List[AggregateResult] = []
    for name, shuffler in shufflers.items():
        sys.stdout.write(f"  {name:<28s} ... ")
        sys.stdout.flush()

        r = benchmark_shuffler(
            name=name,
            shuffler=shuffler,
            n=n,
            n_trials=n_trials,
            base_seed=base_seed,
            fy_displacements=fy_displacements,
        )
        results.append(r)
        sys.stdout.write(
            f"Kt={r.kendall_tau:.4f}  TV={r.tv_distance:.4f}  BB={r.bit_bias:+.4f}  {r.wall_seconds_mean:.5f}s\n"
        )

        if name == "Fisher-Yates":
            items = list(range(n))
            disps: List[float] = []
            for t in range(n_trials):
                rng = random.Random(base_seed + t)
                perm = shuffler(items, rng)
                disps.extend(float(perm[i] - i) for i in range(n))
            fy_displacements = np.array(disps)

    print()
    _print_header()
    for r in results:
        _print_row(r)
    _print_interpretation(results)

    if not args.no_plot:
        out = _plot_results(results, args.output)
        print(f"\nFigure saved: {out}")
    print()


if __name__ == "__main__":
    main()
