#!/usr/bin/env python3
"""Benchmark: Finsler QUBO — QAOA vs classical solvers across problem sizes.

Compares QAOA (Cirq), Simulated Annealing (SA), and Levy flights on
TransportQUBOBridge Finsler-Randers QUBOs at n=8, 12, 24, 48 directions.
"""

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

import json
import math
import time
import numpy as np

from qaoa_adapter import (
    FinslerMetric, finsler_metric_to_qubo, qubo_to_ising,
    ising_to_pauli, pauli_to_cirq, qaoa_solve_qubo,
    stochastic_abuse_qubo, qaoa_vs_stochastic_comparison,
)

def generate_finsler_qubo(n_dirs: int, dimension: int = 3, seed: int = 42) -> tuple:
    rng = np.random.RandomState(seed)
    alpha_mass = rng.uniform(0.5, 2.0, size=dimension).tolist()
    beta_wind = rng.uniform(-0.3, 0.3, size=dimension).tolist()
    metric = FinslerMetric(alpha_mass=alpha_mass, beta_wind=beta_wind, dimension=dimension)
    directions = []
    for k in range(n_dirs):
        vec = rng.randn(dimension)
        norm = math.sqrt(sum(v*v for v in vec))
        directions.append([v / norm for v in vec])
    qubo = finsler_metric_to_qubo(metric, directions, normalize=True)
    # Verify anisotropy
    raw = [[0.0]*n_dirs for _ in range(n_dirs)]
    for i in range(n_dirs):
        for j in range(n_dirs):
            if i != j:
                raw[i][j] = metric.crossing_cost(directions[i], directions[j])
    max_aniso = max(abs(raw[i][j] - raw[j][i]) for i in range(n_dirs) for j in range(n_dirs))
    return qubo, metric, directions, max_aniso


def run_benchmark(sizes: list[int] = None, p_layers: int = 1, shots: int = 2000, time_limit: float = 5.0, seed: int = 42) -> dict:
    if sizes is None:
        sizes = [8, 12, 24, 48]
    results = {}
    for n in sizes:
        print(f"\n=== Finsler QUBO n={n} ===")
        qubo, metric, directions, max_aniso = generate_finsler_qubo(n, seed=seed)
        print(f"  max |Q_ij - Q_ji| = {max_aniso:.6f}  {'anisotropic' if max_aniso > 1e-9 else 'symmetric'}")
        row = {"n": n, "max_anisotropy": max_aniso}
        # QAOA
        try:
            t0 = time.time()
            qaoa_result = qaoa_solve_qubo(
                qubo, p_layers=p_layers, shots=shots, backend="cirq"
            )
            t_qaoa = time.time() - t0
            row["qaoa"] = {
                "energy": round(qaoa_result["energy"], 6),
                "time_s": round(t_qaoa, 3),
                "solution": qaoa_result["solution"],
                "p_layers": p_layers,
            }
            print(f"  QAOA(p={p_layers}):  energy={row['qaoa']['energy']},  time={t_qaoa:.3f}s")
        except Exception as e:
            row["qaoa"] = {"error": str(e)}
            print(f"  QAOA: FAILED — {e}")
        # SA
        try:
            t0 = time.time()
            sa_result = stochastic_abuse_qubo(qubo, method="sa", time_limit=time_limit, seed=seed)
            t_sa = time.time() - t0
            row["sa"] = {
                "energy": round(sa_result["energy"], 6),
                "time_s": round(t_sa, 3),
                "solution": sa_result["solution"],
            }
            print(f"  SA:               energy={row['sa']['energy']},  time={t_sa:.3f}s")
        except Exception as e:
            row["sa"] = {"error": str(e)}
            print(f"  SA: FAILED — {e}")
        # Levy flights
        try:
            t0 = time.time()
            levy_result = stochastic_abuse_qubo(qubo, method="levy", time_limit=time_limit, seed=seed)
            t_levy = time.time() - t0
            row["levy"] = {
                "energy": round(levy_result["energy"], 6),
                "time_s": round(t_levy, 3),
                "solution": levy_result["solution"],
            }
            print(f"  Levy flights:     energy={row['levy']['energy']},  time={t_levy:.3f}s")
        except Exception as e:
            row["levy"] = {"error": str(e)}
            print(f"  Levy: FAILED — {e}")
        # Geodesic (ground truth)
        geodesic = geodesic_assignment(metric, directions)
        geo_energy = qubo.energy([1 if g else 0 for g in geodesic])
        all_false_energy = qubo.energy([0]*n)
        row["geodesic"] = {
            "energy": round(geo_energy, 6),
            "all_false_energy": round(all_false_energy, 6),
            "assignment": [int(g) for g in geodesic],
        }
        print(f"  Geodesic truth:   energy={geo_energy:.6f}  (all-false: {all_false_energy:.6f})")
        # Determine winner
        energies = {}
        if "energy" in row.get("qaoa", {}):
            energies["qaoa"] = row["qaoa"]["energy"]
        if "energy" in row.get("sa", {}):
            energies["sa"] = row["sa"]["energy"]
        if "energy" in row.get("levy", {}):
            energies["levy"] = row["levy"]["energy"]
        if energies:
            best_name = min(energies, key=energies.get)
            row["winner"] = {"solver": best_name, "energy": energies[best_name]}
            row["gap_vs_geodesic"] = round(energies[best_name] - geo_energy, 6)
            print(f"  Winner: {best_name}  (gap vs geodesic: {row['gap_vs_geodesic']})")
        results[f"n_{n}"] = row
    return results


def geodesic_assignment(metric, directions, tol=1e-9):
    costs = [metric.finsler_cost(v) for v in directions]
    min_cost = min(costs)
    return [abs(c - min_cost) <= tol for c in costs]


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Benchmark Finsler QUBO solvers")
    parser.add_argument("--sizes", type=int, nargs="+", default=[8, 12, 24, 48])
    parser.add_argument("--p-layers", type=int, default=1)
    parser.add_argument("--shots", type=int, default=2000)
    parser.add_argument("--time-limit", type=float, default=5.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", "-o", help="JSON output path")
    args = parser.parse_args()
    results = run_benchmark(
        sizes=args.sizes,
        p_layers=args.p_layers,
        shots=args.shots,
        time_limit=args.time_limit,
        seed=args.seed,
    )
    output = json.dumps(results, indent=2)
    if args.output:
        with open(args.output, "w") as f:
            f.write(output)
        print(f"\nWrote results to {args.output}")
    else:
        print(output)