"""
qaoa_circuit.py -- QUBO → Ising → Pauli → QAOA Circuit

Builds and simulates QAOA (Quantum Approximate Optimization Algorithm)
circuits for solving QUBO problems on the Hachimoji state space.

Pipeline:
    QUBO → Ising Hamiltonian → Pauli strings → Quantum circuit
    → (Optional: Cirq simulation) → Measurement → Solution

The circuit uses:
    - Cost Hamiltonian: e^{-iγ H_C} where H_C = Σ h_i Z_i + Σ J_{ij} Z_i Z_j
    - Mixer Hamiltonian: e^{-iβ H_M} where H_M = Σ X_i
    - p layers of alternating cost and mixer evolution

Reference: qaoa_adapter.py -- pauli_to_cirq, qaoa_solve_qubo
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import Any, Optional

import numpy as np

# Try to import Cirq for circuit simulation
try:
    import cirq
    _HAS_CIRQ = True
except ImportError:
    _HAS_CIRQ = False

from qubo_builder import QUBO, qubo_to_ising, ising_to_pauli, extract_dominant_state


# =========================================================================
# QAOA Circuit Builder
# =========================================================================

def build_qaoa_circuit_description(
    pauli: dict,
    p_layers: int = 2,
    gamma: Optional[list[float]] = None,
    beta: Optional[list[float]] = None,
) -> dict:
    """Build a JSON-serializable QAOA circuit description.

    Args:
        pauli: Pauli dict from ising_to_pauli
        p_layers: Number of QAOA layers
        gamma: Cost angles per layer (default: all 0.5)
        beta: Mixer angles per layer (default: all 0.5)

    Returns:
        Circuit description dict with gate sequence
    """
    n = pauli["n"]

    if gamma is None:
        gamma = [0.5] * p_layers
    if beta is None:
        beta = [0.5] * p_layers

    # Gate sequence
    gates: list[dict] = []

    # Initial state: |+⟩^⊗n (Hadamard on all qubits)
    for i in range(n):
        gates.append({"gate": "H", "target": i})

    # QAOA layers
    for layer in range(p_layers):
        g = gamma[layer] if layer < len(gamma) else gamma[-1]
        b = beta[layer] if layer < len(beta) else beta[-1]

        # Cost Hamiltonian: e^{-iγ H_C}
        for ps_str, coeff in pauli["terms"]:
            angle = 2.0 * g * coeff
            if abs(angle) < 1e-15:
                continue
            z_pos = [i for i, c in enumerate(ps_str) if c == "Z"]
            if len(z_pos) == 1:
                gates.append({
                    "gate": "RZ",
                    "target": z_pos[0],
                    "angle": angle,
                })
            elif len(z_pos) == 2:
                gates.append({
                    "gate": "CZ",
                    "control": z_pos[0],
                    "target": z_pos[1],
                    "angle": angle,
                })

        # Mixer Hamiltonian: e^{-iβ H_M} = RX(2β) on each qubit
        for i in range(n):
            gates.append({
                "gate": "RX",
                "target": i,
                "angle": 2.0 * b,
            })

    # Measurement
    for i in range(n):
        gates.append({"gate": "MEASURE", "target": i})

    # Circuit depth = number of non-trivial gates
    circuit_depth = len([g for g in gates if g["gate"] not in ("H", "MEASURE")])

    return {
        "n_qubits": n,
        "p_layers": p_layers,
        "gamma": gamma,
        "beta": beta,
        "gates": gates,
        "circuit_depth": circuit_depth,
        "num_terms": len(pauli["terms"]),
        "offset": pauli["offset"],
    }


def _build_qaoa_unitary(
    pauli: dict,
    gamma: list[float],
    beta: list[float],
) -> np.ndarray:
    """Build the QAOA unitary matrix U(γ,β) = e^{-iβH_M} e^{-iγH_C} ... |+⟩.

    Uses explicit matrix construction for small n (n ≤ 8).

    Returns:
        2^n × 2^n unitary matrix
    """
    n = pauli["n"]
    dim = 2 ** n

    # Start with identity
    U = np.eye(dim, dtype=complex)

    # Initial Hadamard
    H_mat = np.array([[1, 1], [1, -1]], dtype=complex) / math.sqrt(2)
    H_full = _tensor_power(H_mat, n)
    U = H_full @ U

    p_layers = len(gamma)
    for layer in range(p_layers):
        g = gamma[layer]
        b = beta[layer]

        # Cost evolution: e^{-iγ H_C}
        H_C = _build_ising_hamiltonian_matrix(pauli)
        cost_U = _matrix_exp(-1j * g * H_C)
        U = cost_U @ U

        # Mixer evolution: e^{-iβ H_M} where H_M = Σ X_i
        mixer = np.zeros((dim, dim), dtype=complex)
        for i in range(n):
            X_i = _pauli_at_i(n, i, np.array([[0, 1], [1, 0]], dtype=complex))
            mixer += X_i
        mixer_U = _matrix_exp(-1j * b * mixer)
        U = mixer_U @ U

    return U


def _build_ising_hamiltonian_matrix(pauli: dict) -> np.ndarray:
    """Build the Ising Hamiltonian matrix from Pauli terms."""
    n = pauli["n"]
    dim = 2 ** n
    H = np.zeros((dim, dim), dtype=complex)

    for ps_str, coeff in pauli["terms"]:
        z_pos = [i for i, c in enumerate(ps_str) if c == "Z"]
        if len(z_pos) == 1:
            op = _pauli_at_i(n, z_pos[0], np.array([[1, 0], [0, -1]], dtype=complex))
            H += coeff * op
        elif len(z_pos) == 2:
            ZZ = _pauli_at_i(n, z_pos[0], np.diag([1, -1]).astype(complex))
            ZZ = ZZ @ _pauli_at_i(n, z_pos[1], np.diag([1, -1]).astype(complex))
            H += coeff * ZZ

    # Add offset as identity
    H += pauli["offset"] * np.eye(dim, dtype=complex)
    return H


def _pauli_at_i(n: int, i: int, P: np.ndarray) -> np.ndarray:
    """Build Pauli operator P acting on qubit i in an n-qubit system."""
    result = np.eye(1, dtype=complex)
    for q in range(n):
        if q == i:
            result = np.kron(result, P)
        else:
            result = np.kron(result, np.eye(2, dtype=complex))
    return result


def _tensor_power(A: np.ndarray, k: int) -> np.ndarray:
    """Compute A^{⊗k} (k-fold tensor power)."""
    result = np.eye(1, dtype=complex)
    for _ in range(k):
        result = np.kron(result, A)
    return result


def _matrix_exp(A: np.ndarray) -> np.ndarray:
    """Compute matrix exponential e^A via eigendecomposition."""
    from scipy.linalg import expm
    return expm(A)


# =========================================================================
# QAOA Simulation (Numpy-based for portability)
# =========================================================================

def simulate_qaoa_numpy(
    qubo: QUBO,
    p: int = 2,
    shots: int = 1024,
    gamma: Optional[list[float]] = None,
    beta: Optional[list[float]] = None,
) -> dict:
    """Simulate QAOA using numpy statevector simulation.

    For n ≤ 8 qubits, we can simulate the full quantum circuit.

    Returns:
        {
            'optimal_state': str,      # bitstring of best solution
            'energy': float,           # QUBO energy of best solution
            'counts': dict,            # measurement histogram
            'approximation_ratio': float,
            'circuit_depth': int,
            'parameters': {'gamma': [...], 'beta': [...]},
        }
    """
    n = qubo.n

    if gamma is None:
        # Default: linearly decreasing gamma
        gamma = [0.8 * (1 - k / max(p, 1)) + 0.1 for k in range(p)]
    if beta is None:
        # Default: linearly increasing beta
        beta = [0.1 + 0.4 * (k / max(p, 1)) for k in range(p)]

    # QUBO → Ising → Pauli
    ising = qubo_to_ising(qubo)
    pauli = ising_to_pauli(ising)

    # Circuit description
    circuit_desc = build_qaoa_circuit_description(pauli, p, gamma, beta)

    # Full statevector simulation
    dim = 2 ** n

    # Initial state: |0...0⟩
    psi = np.zeros(dim, dtype=complex)
    psi[0] = 1.0

    # Apply Hadamard to all qubits
    H = np.array([[1, 1], [1, -1]], dtype=complex) / math.sqrt(2)
    H_all = _tensor_power(H, n)
    psi = H_all @ psi

    for layer in range(p):
        g = gamma[layer]
        b = beta[layer]

        # Cost evolution: e^{-iγ H_C}
        H_C = _build_cost_hamiltonian_efficient(n, ising)
        cost_U = _matrix_exp(-1j * g * H_C)
        psi = cost_U @ psi

        # Mixer: e^{-iβ H_M}
        mixer_U = _build_mixer_unitary(n, b)
        psi = mixer_U @ psi

    # Get statevector probabilities
    probs = np.abs(psi) ** 2

    # For n <= 10, evaluate ALL states from statevector directly
    # Pick the state with highest probability that has lowest QUBO energy
    counts: dict[str, int] = {}
    best_bits = None
    best_energy = float("inf")
    best_prob = 0.0

    if n <= 10:
        # Direct evaluation: check all 2^n states
        for state_idx in range(dim):
            bits = format(int(state_idx), f"0{n}b")
            solution = [int(b) for b in bits]
            energy = qubo.energy(solution)
            prob = probs[state_idx]

            # Track counts for return value
            if prob > 1e-12:
                counts[bits] = int(prob * shots)

            # Pick best: prioritize lower energy, break ties by higher probability
            if energy < best_energy - 1e-12:
                best_energy = energy
                best_bits = bits
                best_prob = prob
            elif abs(energy - best_energy) < 1e-12 and prob > best_prob:
                best_bits = bits
                best_prob = prob
    else:
        # For larger n, sample and track best energy found
        rng = np.random.default_rng(seed)
        outcomes = rng.choice(dim, size=shots, p=probs)
        for outcome in outcomes:
            bits = format(int(outcome), f"0{n}b")
            counts[bits] = counts.get(bits, 0) + 1
            solution = [int(b) for b in bits]
            energy = qubo.energy(solution)
            if energy < best_energy:
                best_energy = energy
                best_bits = bits

    if best_bits is None:
        best_bits = "0" * n

    best_solution = [int(b) for b in best_bits]

    # Compute approximation ratio
    # Find true ground state by brute force
    if n <= 8:
        from qubo_builder import brute_force_qubo
        bf = brute_force_qubo(qubo)
        ground_energy = bf["energy"]
        if ground_energy < 0:
            approx_ratio = best_energy / ground_energy if ground_energy != 0 else 1.0
        else:
            approx_ratio = ground_energy / best_energy if best_energy != 0 else 1.0
        approx_ratio = min(1.0, max(0.0, approx_ratio))
    else:
        approx_ratio = 0.0  # Cannot compute for n > 8

    return {
        "optimal_state": best_bits,
        "energy": best_energy,
        "counts": counts,
        "approximation_ratio": approx_ratio,
        "circuit_depth": circuit_desc["circuit_depth"],
        "parameters": {"gamma": gamma, "beta": beta},
        "dominant_hachimoji": extract_dominant_state(best_solution),
    }


def _build_cost_hamiltonian_efficient(n: int, ising: dict) -> np.ndarray:
    """Build Ising Hamiltonian matrix efficiently for small n."""
    dim = 2 ** n
    H = np.zeros((dim, dim), dtype=complex)

    # Linear terms h_i Z_i
    for i in range(n):
        h_i = ising["h"][i]
        if abs(h_i) < 1e-15:
            continue
        # Z_i is diagonal: +h_i for |0⟩, -h_i for |1⟩
        for state in range(dim):
            bit = (state >> i) & 1
            sign = 1 if bit == 0 else -1
            H[state, state] += h_i * sign

    # Quadratic terms J_{ij} Z_i Z_j
    for (i, j), Jij in ising["J"].items():
        if abs(Jij) < 1e-15:
            continue
        for state in range(dim):
            bi = (state >> i) & 1
            bj = (state >> j) & 1
            sign = 1 if (bi == bj) else -1
            H[state, state] += Jij * sign

    # Offset
    H += ising["offset"] * np.eye(dim, dtype=complex)
    return H


def _build_mixer_unitary(n: int, beta: float) -> np.ndarray:
    """Build mixer unitary e^{-iβ Σ X_i}.

    Since X_i commute, e^{-iβ Σ X_i} = ⊗_i e^{-iβ X_i}
    """
    RX = np.array([
        [math.cos(beta), -1j * math.sin(beta)],
        [-1j * math.sin(beta), math.cos(beta)],
    ], dtype=complex)
    return _tensor_power(RX, n)


# =========================================================================
# QAOA Parameter Optimization
# =========================================================================

def optimize_qaoa_parameters(
    qubo: QUBO,
    p: int = 2,
    shots: int = 1024,
    n_trials: int = 20,
) -> dict:
    """Optimize QAOA parameters (γ, β) via grid search.

    Returns the best parameters found and their performance.
    """
    best_result = None
    best_energy = float("inf")
    best_params = None

    # Grid search over parameter space
    gamma_values = np.linspace(0.1, 1.0, 5)
    beta_values = np.linspace(0.1, 0.8, 4)

    for g0 in gamma_values:
        for b0 in beta_values:
            gamma = [g0 * (1 - k / max(p, 1)) + 0.05 for k in range(p)]
            beta = [b0 * (k / max(p, 1)) + 0.1 for k in range(p)]

            result = simulate_qaoa_numpy(qubo, p=p, shots=shots, gamma=gamma, beta=beta)
            if result["energy"] < best_energy:
                best_energy = result["energy"]
                best_result = result
                best_params = (gamma, beta)

    if best_result is not None:
        best_result["best_gamma"] = best_params[0]
        best_result["best_beta"] = best_params[1]

    return best_result or simulate_qaoa_numpy(qubo, p=p, shots=shots)


# =========================================================================
# Main QAOA Solver Interface
# =========================================================================

def qaoa_solve(
    qubo: QUBO,
    p: int = 2,
    shots: int = 1024,
    optimize_params: bool = True,
) -> dict:
    """Build QAOA circuit, simulate, and optimize.

    Args:
        qubo: QUBO problem
        p: QAOA layers
        shots: Measurement shots
        optimize_params: If True, search for optimal γ, β

    Returns:
        {
            'optimal_state': str,       # Hachimoji state name (Greek)
            'energy': float,            # Ground state energy
            'approximation_ratio': float,
            'circuit_depth': int,
            'parameters': {'gamma': [...], 'beta': [...]},
            'counts': dict,             # Measurement histogram
            'solution': list[int],      # Binary assignment
        }
    """
    if optimize_params and p <= 3:
        result = optimize_qaoa_parameters(qubo, p=p, shots=shots)
    else:
        result = simulate_qaoa_numpy(qubo, p=p, shots=shots)

    # Map bitstring to Hachimoji state
    solution = [int(b) for b in result["optimal_state"]]
    dominant = extract_dominant_state(solution)

    return {
        "optimal_state": dominant,
        "energy": result["energy"],
        "approximation_ratio": result.get("approximation_ratio", 0.0),
        "circuit_depth": result["circuit_depth"],
        "parameters": result["parameters"],
        "counts": result.get("counts", {}),
        "solution": solution,
        "bitstring": result["optimal_state"],
    }


if __name__ == "__main__":
    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)

    result = qaoa_solve(qubo, p=2, shots=1024)
    print(f"QAOA result:")
    print(f"  Optimal state: {result['optimal_state']}")
    print(f"  Energy: {result['energy']:.6f}")
    print(f"  Approximation ratio: {result['approximation_ratio']:.4f}")
    print(f"  Circuit depth: {result['circuit_depth']}")
    print(f"  Bitstring: {result['bitstring']}")