"""
classical_solver.py -- Classical Fallback Solvers for QUBO

Provides classical optimization baselines for comparison with QAOA:
    - HiGHS MIP solver (exact for small problems)
    - Simulated Annealing (SA) heuristic
    - Both return solutions in the same format as qaoa_solve() for comparison.

Reference: qubo_highs.py -- solve_qubo_highs, _sa_solve_qubo
"""

from __future__ import annotations

import math
import random
import time
from typing import Any, Optional

import numpy as np

from qubo_builder import QUBO, extract_dominant_state


# =========================================================================
# HiGHS MIP Solver
# =========================================================================

def solve_highs(qubo: QUBO, time_limit: float = 60.0) -> dict:
    """Solve QUBO using HiGHS MIP solver.

    Converts QUBO to MIP via linearization of bilinear terms:
        y_{ij} = x_i · x_j  with McCormick inequalities:
        y_{ij} ≤ x_i, y_{ij} ≤ x_j, y_{ij} ≥ x_i + x_j - 1

    Returns solution matching qaoa_solve() format:
        {
            'optimal_state': str,      # Hachimoji state name
            'energy': float,
            'solution': list[int],
            'method': 'highs',
            'status': str,
            'runtime_s': float,
        }
    """
    t0 = time.time()

    try:
        import highspy
        HAS_HIGHS = True
    except ImportError:
        HAS_HIGHS = False

    if not HAS_HIGHS:
        # Fall back to simulated annealing
        return solve_sa(qubo, time_limit=time_limit)

    n = qubo.n

    # Separate linear and quadratic terms
    linear: dict[int, float] = {}
    quadratic: dict[tuple[int, int], float] = {}
    for (i, j), qij in qubo.matrix.items():
        if i == j:
            linear[i] = linear.get(i, 0.0) + qij
        else:
            key = (min(i, j), max(i, j))
            quadratic[key] = quadratic.get(key, 0.0) + qij

    # Variables: x_0..x_{n-1} (binary), y_n.. (continuous for bilinear)
    y_map: dict[tuple[int, int], int] = {}
    y_idx = n
    for (i, j) in quadratic:
        y_map[(i, j)] = y_idx
        y_idx += 1

    num_vars = y_idx
    num_rows = len(quadratic) * 3  # 3 McCormick constraints per pair

    # Build model
    model = highspy.HighsModel()
    lp = model.lp_
    lp.num_col_ = num_vars
    lp.num_row_ = num_rows

    # Objective: min Σ a_i x_i + Σ b_{ij} y_{ij}
    obj = np.zeros(num_vars)
    for i, coeff in linear.items():
        obj[i] = coeff
    for (i, j), coeff in quadratic.items():
        obj[y_map[(i, j)]] = coeff
    lp.col_cost_ = obj

    # Bounds: x binary [0,1], y [0,1]
    lp.col_lower_ = np.zeros(num_vars)
    lp.col_upper_ = np.ones(num_vars)

    # Integrality: x binary, y continuous
    lp.integrality_ = [highspy.HighsVarType.kInteger] * n + \
                       [highspy.HighsVarType.kContinuous] * (num_vars - n)

    # Build constraint matrix (CSC format)
    col_entries: dict[int, list[tuple[int, float]]] = {}
    row_idx = 0

    for (i, j), yi in y_map.items():
        # y_{ij} ≤ x_i
        if yi not in col_entries:
            col_entries[yi] = []
        if i not in col_entries:
            col_entries[i] = []
        col_entries[yi].append((row_idx, 1.0))
        col_entries[i].append((row_idx, -1.0))
        row_idx += 1

        # y_{ij} ≤ x_j
        if j not in col_entries:
            col_entries[j] = []
        col_entries[yi].append((row_idx, 1.0))
        col_entries[j].append((row_idx, -1.0))
        row_idx += 1

        # y_{ij} ≥ x_i + x_j - 1  =>  -y_{ij} + x_i + x_j ≤ 1
        col_entries[yi].append((row_idx, -1.0))
        col_entries[i].append((row_idx, 1.0))
        col_entries[j].append((row_idx, 1.0))
        row_idx += 1

    starts = []
    indices_list = []
    values_list = []
    nnz = 0
    for col in range(num_vars):
        starts.append(nnz)
        if col in col_entries:
            for ridx, val in col_entries[col]:
                indices_list.append(ridx)
                values_list.append(val)
                nnz += 1
    starts.append(nnz)

    lp.a_matrix_.format_ = highspy.MatrixFormat.kColwise
    lp.a_matrix_.start_ = np.array(starts, dtype=np.int32)
    lp.a_matrix_.index_ = np.array(indices_list, dtype=np.int32) if indices_list else np.array([], dtype=np.int32)
    lp.a_matrix_.value_ = np.array(values_list) if values_list else np.array([])

    # Row bounds: all ≤ 0 or ≤ 1
    row_lower = []
    row_upper = []
    for _ in range(len(quadratic)):
        row_lower.append(-1e30)
        row_upper.append(0.0)   # y - x ≤ 0
        row_lower.append(-1e30)
        row_upper.append(0.0)   # y - x' ≤ 0
        row_lower.append(-1e30)
        row_upper.append(1.0)   # -y + x + x' ≤ 1

    lp.row_lower_ = np.array(row_lower)
    lp.row_upper_ = np.array(row_upper)

    # Solve
    h = highspy.Highs()
    h.setOptionValue("time_limit", time_limit)
    h.setOptionValue("output_flag", False)
    h.passModel(model)
    h.run()

    sol = h.getSolution()
    x_vals = sol.col_value

    solution = [int(round(max(0, min(1, x_vals[i])))) for i in range(n)]
    energy = qubo.energy(solution)
    runtime = time.time() - t0

    # Get status
    status_val = h.getInfoValue("primal_solution_status")[1]
    status_map = {0: "unknown", 1: "infeasible", 2: "feasible", 3: "optimal"}
    status_str = status_map.get(status_val, f"status_{status_val}")

    return {
        "optimal_state": extract_dominant_state(solution),
        "energy": energy,
        "solution": solution,
        "method": "highs",
        "status": status_str,
        "runtime_s": round(runtime, 4),
    }


# =========================================================================
# Simulated Annealing
# =========================================================================

def solve_sa(
    qubo: QUBO,
    time_limit: float = 5.0,
    initial_temp: float = 10.0,
    cooling_rate: float = 0.9995,
    seed: int = 42,
) -> dict:
    """Solve QUBO with Simulated Annealing.

    Standard SA: flip random bits, accept if energy decreases or
    with probability exp(-ΔE/T).

    Returns solution matching qaoa_solve() format:
        {
            'optimal_state': str,      # Hachimoji state name
            'energy': float,
            'solution': list[int],
            'method': 'sa',
            'iterations': int,
            'runtime_s': float,
        }
    """
    t0 = time.time()
    rng = random.Random(seed)

    n = qubo.n
    Q_dict = dict(qubo.matrix)

    def _energy(x):
        e = qubo.offset
        for (i, j), qij in Q_dict.items():
            e += qij * x[i] * x[j]
        return e

    # Initialize random solution
    x = [rng.randint(0, 1) for _ in range(n)]
    current_energy = _energy(x)
    best_x = x[:]
    best_energy = current_energy

    T = initial_temp
    iterations = 0
    n_vals = list(range(n))

    while (time.time() - t0) < time_limit:
        i = rng.choice(n_vals)
        x[i] = 1 - x[i]  # flip bit
        new_energy = _energy(x)
        delta = new_energy - current_energy

        if delta < 0 or rng.random() < math.exp(-delta / max(T, 1e-10)):
            current_energy = new_energy
            if current_energy < best_energy:
                best_x = x[:]
                best_energy = current_energy
        else:
            x[i] = 1 - x[i]  # revert

        T *= cooling_rate
        iterations += 1

    runtime = time.time() - t0

    return {
        "optimal_state": extract_dominant_state(best_x),
        "energy": best_energy,
        "solution": best_x,
        "method": "sa",
        "iterations": iterations,
        "runtime_s": round(runtime, 4),
    }


# =========================================================================
# Unified Solver Interface
# =========================================================================

def solve_classical(qubo: QUBO, method: str = "highs", **kwargs) -> dict:
    """Solve QUBO with classical methods for comparison.

    Methods:
        - "highs": HiGHS MIP solver (exact, falls back to SA)
        - "sa": Simulated annealing heuristic

    Returns solution matching qaoa_solve() format for comparison.
    """
    if method == "highs":
        return solve_highs(qubo, **{k: v for k, v in kwargs.items() if k in ["time_limit"]})
    elif method == "sa":
        return solve_sa(qubo, **{k: v for k, v in kwargs.items() if k in ["time_limit", "initial_temp", "cooling_rate", "seed"]})
    else:
        raise ValueError(f"Unknown method: {method}. Use 'highs' or 'sa'.")


def compare_solvers(qubo: QUBO, time_limit: float = 2.0) -> dict:
    """Run all classical solvers and compare results.

    Returns:
        {
            "highs": {...},
            "sa": {...},
            "best": {"method": str, "energy": float},
            "agreement": bool,  # whether all solvers agree on state
        }
    """
    highs_result = solve_highs(qubo, time_limit=time_limit)
    sa_result = solve_sa(qubo, time_limit=time_limit)

    # Find best
    results = {"highs": highs_result, "sa": sa_result}
    best_method = min(results, key=lambda m: results[m]["energy"])

    # Check agreement
    states = [r["optimal_state"] for r in results.values()]
    agreement = len(set(states)) == 1

    return {
        "highs": highs_result,
        "sa": sa_result,
        "best": {"method": best_method, "energy": results[best_method]["energy"]},
        "agreement": agreement,
        "all_states": {m: r["optimal_state"] for m, r in results.items()},
    }


if __name__ == "__main__":
    import sys
    sys.path.insert(0, "/mnt/agents/output/rebuild/stage4-optimize")
    from finsler_metric import make_uniform_hachimoji_states
    from qubo_builder import finsler_to_qubo

    states = make_uniform_hachimoji_states()
    qubo = finsler_to_qubo(states)

    print("Testing classical solvers on 8-state Hachimoji QUBO...")
    comparison = compare_solvers(qubo, time_limit=1.0)

    print(f"\nHiGHS: state={comparison['highs']['optimal_state']}, "
          f"energy={comparison['highs']['energy']:.6f}, "
          f"status={comparison['highs']['status']}")
    print(f"SA:    state={comparison['sa']['optimal_state']}, "
          f"energy={comparison['sa']['energy']:.6f}, "
          f"iterations={comparison['sa']['iterations']}")
    print(f"\nAgreement: {comparison['agreement']}")
    print(f"Best: {comparison['best']['method']} with energy {comparison['best']['energy']:.6f}")