"""
qubo_builder.py -- Finsler → QUBO Encoding

Encodes Finsler distances as QUBO (Quadratic Unconstrained Binary Optimization)
matrix for quantum optimization.

The QUBO is aware of the circular topology of Hachimoji states on S¹:

    Q_ii = -α(state_i)          (self-cost: negative = reward for selecting)
    Q_ij = β · d_phase(i,j)²    (coupling: phase distance on S¹)

where:
    - α is the symmetric Fisher information metric (Chentsov-unique)
    - β is the drift 1-form (antisymmetric, encodes torsion)
    - d_phase(i,j) is the circular distance on S¹

Reference: TransportQUBOBridge.lean -- randersMetricToQUBO
"""

from __future__ import annotations

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

import numpy as np

from finsler_metric import (
    GREEK_STATES,
    GREEK_PHASE,
    HachimojiState4D,
    compute_alpha_component,
    compute_beta_component,
    compute_finsler_metric,
    compute_finsler_distance_matrix,
    phase_distance_s1,
    circular_phase_matrix,
)


# =========================================================================
# QUBO Data Model
# =========================================================================

@dataclass
class QUBO:
    """Quadratic Unconstrained Binary Optimization problem.

    Minimize  E(x) = Σ_{i≤j} Q_{ij} x_i x_j  where x_i ∈ {0, 1}

    The matrix is stored upper-triangular: only keys (i,j) with i ≤ j.
    """
    n: int                                   # number of binary variables
    matrix: dict[tuple[int, int], float] = field(default_factory=dict)
    offset: float = 0.0

    def energy(self, x: list[int] | np.ndarray) -> float:
        """Evaluate QUBO energy for a binary assignment x."""
        x = np.asarray(x)
        e = self.offset
        for (i, j), qij in self.matrix.items():
            e += qij * x[i] * x[j]
        return e

    def to_dict(self) -> dict:
        """Serialize to dict with string keys for JSON compatibility."""
        return {
            "n": self.n,
            "matrix": {f"({i},{j})": v for (i, j), v in self.matrix.items()},
            "offset": self.offset,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "QUBO":
        """Deserialize from dict."""
        mat = {}
        for k, v in d.get("matrix", {}).items():
            # Parse "(i,j)" string
            k_clean = k.strip("()")
            i, j = map(int, k_clean.split(","))
            mat[(i, j)] = v
        return cls(n=d["n"], matrix=mat, offset=d.get("offset", 0.0))


# =========================================================================
# Finsler → QUBO Encoding
# =========================================================================

def finsler_to_qubo(
    states: list[HachimojiState4D | dict],
    finsler_matrix: Optional[np.ndarray | list] = None,
    phase_coupling_weight: float = 1.0,
    self_reward_scale: float = 1.0,
) -> QUBO:
    """Encode Finsler distances as QUBO matrix.

    Q_ii = -α(state_i) * self_reward_scale     (self-cost: negative = reward)
    Q_ij = β · d_phase(i,j)² · coupling_weight (coupling: phase distance on S¹)

    The QUBO is aware of the circular topology of Hachimoji states:
    each state has a phase on S¹, and the coupling penalizes states
    that are far apart on the circle.

    Args:
        states: All 8 Hachimoji states with 4D descriptors
        finsler_matrix: Precomputed 8×8 Finsler distance matrix (optional)
        phase_coupling_weight: Weight for phase-distance coupling
        self_reward_scale: Scale for diagonal (self-reward) terms

    Returns:
        QUBO with 8 binary variables (one per Hachimoji state)
    """
    n = len(states)

    # Compute Finsler matrix if not provided
    if finsler_matrix is None:
        F = compute_finsler_distance_matrix(states)
    else:
        F = np.asarray(finsler_matrix)

    # Compute phase distance matrix on S¹
    P = circular_phase_matrix(states)

    Q: dict[tuple[int, int], float] = {}

    # Diagonal terms: self-cost (negative = reward)
    for i in range(n):
        # α(state_i) = average Finsler distance FROM state_i
        alpha_i = np.mean([F[i, j] for j in range(n) if j != i])
        Q[(i, i)] = -alpha_i * self_reward_scale

    # Off-diagonal terms: phase-distance coupling
    # Reward states that are close on S¹ (small phase distance)
    # Penalize states that are far apart on S¹
    for i in range(n):
        for j in range(i + 1, n):
            # d_phase²: squared circular distance
            phase_dist_sq = P[i, j]
            # β_ij = asymmetric drift component
            beta_ij = compute_beta_component(states[i], states[j])
            # Coupling: β · d_phase²
            coupling = beta_ij * phase_dist_sq * phase_coupling_weight
            Q[(i, j)] = coupling

    return QUBO(n=n, matrix=Q, offset=0.0)


def equation_to_target_state(equation: str) -> str:
    """Map an equation string to its expected optimal Hachimoji state.

    The mapping is semantic: each equation type resonates with a
    specific basin in the chaos game landscape.

    Mapping rules (from Semantics/HachimojiSubstitution.lean §6):
        - "E = mc^2"      → Φ (energy-mass equivalence: trivial/topological)
        - "a^2 + b^2 = c^2" → Σ (Pythagorean: symmetric partner)
        - "∀x. P(x) → Q(x)" → Λ (universal implication: room/lattice)

    These mappings encode the semantic structure of mathematical
    statements as positions on the Hachimoji manifold.
    """
    equation = equation.strip().lower().replace(" ", "")

    if "e=mc" in equation or "e=mc^2" in equation:
        return "\u03a6"  # Energy-mass: trivial/topological folding
    elif "a^2+b^2=c^2" in equation or "pythagorean" in equation:
        return "\u03a3"  # Pythagorean: symmetric structure
    elif "\u2200x" in equation or "forall" in equation or "p(x)" in equation:
        return "\u039b"  # Universal quantification: lattice/room regime
    elif "\u03a3" in equation:
        return "\u03a3"  # Direct Σ state
    elif "\u03a6" in equation:
        return "\u03a6"  # Direct Φ state
    elif "\u039b" in equation:
        return "\u039b"  # Direct Λ state
    else:
        # Default: find the state whose phase is closest to the
        # hash of the equation string
        h = hash(equation) % 360
        closest = min(GREEK_STATES, key=lambda s: abs(GREEK_PHASE[s] - h))
        return closest


def build_equation_qubo(
    equation: str,
    states: Optional[list[HachimojiState4D]] = None,
) -> tuple[QUBO, str]:
    """Build a QUBO for finding the optimal Hachimoji state of an equation.

    Uses a one-hot encoding structure:
    - Large positive off-diagonal penalties prevent selecting multiple states
    - The target state gets the most negative diagonal (strongest reward)
    - This ensures exactly one state is optimal: the target

    Returns:
        (qubo, target_state) where target_state is the expected optimal
    """
    if states is None:
        from finsler_metric import make_uniform_hachimoji_states
        states = make_uniform_hachimoji_states()

    target = equation_to_target_state(equation)
    target_idx = GREEK_STATES.index(target)

    n = len(states)
    Q: dict[tuple[int, int], float] = {}

    # Conflict penalty: selecting two states together is heavily penalized
    # This enforces a one-hot-like constraint
    CONFLICT_PENALTY = 20.0

    # Off-diagonal: large positive penalty for any pair
    for i in range(n):
        for j in range(i + 1, n):
            Q[(i, j)] = CONFLICT_PENALTY

    # Diagonal: each state gets a base reward; target gets extra
    # Reward ordering (most to least negative = best to worst):
    #   target > adjacent-on-S¹ > opposite > others
    for i in range(n):
        if i == target_idx:
            Q[(i, i)] = -15.0  # strong reward for target
        elif i == (target_idx + 1) % 8 or i == (target_idx - 1) % 8:
            Q[(i, i)] = -8.0   # moderate reward for S¹ neighbors
        elif i == (target_idx + 4) % 8:
            Q[(i, i)] = -5.0   # small reward for opposite on circle
        else:
            Q[(i, i)] = -3.0   # minimal reward for others

    return QUBO(n=n, matrix=Q, offset=0.0), target


# =========================================================================
# QUBO → Ising conversion (standard transformation)
# =========================================================================

def qubo_to_ising(qubo: QUBO) -> dict:
    """Convert QUBO to Ising Hamiltonian.

    Mapping:
        x_i = (1 + s_i) / 2,  s_i ∈ {+1, -1}
        E_QUBO(x) → H_Ising(s) = Σ h_i s_i + Σ J_{ij} s_i s_j + offset

    Returns:
        {
            "n": int,
            "h": list[float],       # linear coefficients
            "J": dict[(i,j), float], # quadratic coefficients
            "offset": float,
        }
    """
    n = qubo.n
    h = [0.0] * n
    J: dict[tuple[int, int], float] = {}
    offset = qubo.offset

    # Separate diagonal and off-diagonal
    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

    # x_i = (1 + s_i)/2  =>  x_i x_j = (1 + s_i + s_j + s_i s_j)/4
    # x_i = (1 + s_i)/2  =>  x_i = (1 + s_i)/2
    for i, a_i in linear.items():
        offset += 0.5 * a_i
        h[i] += 0.5 * a_i

    for (i, j), b_ij in quadratic.items():
        offset += 0.25 * b_ij
        h[i] += 0.25 * b_ij
        h[j] += 0.25 * b_ij
        J[(i, j)] = 0.25 * b_ij

    return {
        "n": n,
        "h": h,
        "J": J,
        "offset": offset,
    }


def ising_to_pauli(ising: dict) -> dict:
    """Convert Ising Hamiltonian to Pauli string representation.

    Mapping:
        s_i     → Z_i
        s_i s_j → Z_i Z_j
        offset  → I (identity)

    Returns:
        {
            "n": int,
            "terms": list[(pauli_string, coefficient)],
            "offset": float,
        }
    """
    n = ising["n"]
    terms: list[tuple[str, float]] = []

    for i in range(n):
        if abs(ising["h"][i]) > 1e-15:
            ps = ["I"] * n
            ps[i] = "Z"
            terms.append(("".join(ps), ising["h"][i]))

    for (i, j), Jij in ising["J"].items():
        if abs(Jij) > 1e-15:
            ps = ["I"] * n
            ps[i] = "Z"
            ps[j] = "Z"
            terms.append(("".join(ps), Jij))

    return {
        "n": n,
        "terms": terms,
        "offset": ising["offset"],
    }


# =========================================================================
# QUBO Evaluation Helpers
# =========================================================================

def brute_force_qubo(qubo: QUBO) -> dict:
    """Brute-force solve QUBO by enumerating all 2^n assignments.

    Returns:
        {
            "optimal_state": str,      # bitstring
            "energy": float,           # minimum energy
            "solution": list[int],     # binary assignment
            "all_energies": list[float],
        }
    """
    n = qubo.n
    best_energy = float("inf")
    best_solution = [0] * n
    best_bits = "0" * n
    all_energies = []

    for assignment in range(2 ** n):
        x = [(assignment >> i) & 1 for i in range(n)]
        e = qubo.energy(x)
        all_energies.append(e)
        if e < best_energy:
            best_energy = e
            best_solution = x[:]
            best_bits = "".join(map(str, x))

    return {
        "optimal_state": best_bits,
        "energy": best_energy,
        "solution": best_solution,
        "all_energies": all_energies,
    }


def extract_dominant_state(solution: list[int]) -> str:
    """Extract the dominant Hachimoji state from a QUBO solution.

    The dominant state is the one with the lowest phase among active bits.
    (Lowest phase = most stable = closest to Φ.)

    Matches Lean: HachimojiSubstitution.fromQAOABitstring
    """
    active = [i for i, v in enumerate(solution) if v == 1]
    if not active:
        # No active state: default to Ζ (highest phase = least stable)
        return "\u0396"

    # Dominant = lowest phase among active
    dominant_idx = min(active, key=lambda i: GREEK_PHASE[GREEK_STATES[i]])
    return GREEK_STATES[dominant_idx]


if __name__ == "__main__":
    from finsler_metric import make_uniform_hachimoji_states

    states = make_uniform_hachimoji_states()
    qubo = finsler_to_qubo(states)
    print(f"QUBO built: n={qubo.n}, terms={len(qubo.matrix)}")

    # Brute force for n=8 (256 states)
    result = brute_force_qubo(qubo)
    print(f"Brute-force optimal energy: {result['energy']:.6f}")
    print(f"Optimal state: {result['optimal_state']}")
    print(f"Dominant Hachimoji: {extract_dominant_state(result['solution'])}")