"""
BAWIM Mutation Engine — explore the coupling/encoding design space

Maps the Bulk Acoustic Wave Ising Machine (arXiv 2607.02112) equations
onto the SilverSight Q16 fraction system and provides a mutation loop
that systematically varies parameters to improve solution quality.

Papers equations:
  Eq 1 — Ising Hamiltonian     H = -∑ J_ij s_i s_j - ∑ h_i s_i
  Eq 2 — Feedback coupling     c_i = ∑ J_ij s_j
  Eq 3 — MAX-CUT score         Cuts = -½∑ J_ij - ½H
  Eq 4 — NPP (set form)        E(A,B) = |∑A a_i - ∑B a_i|
  Eq 5 — NPP (spin form)       E(s) = |∑ a_i s_i|

Architecture:
  Mutation round → vary one parameter → evaluate → accept/reject → next round
  All arithmetic is exact Fraction (q16_fraction).  Quantize only at output.

Usage:
    python3 python/bawim_mutation_engine.py --rounds 100 --problem maxcut
    python3 python/bawim_mutation_engine.py --rounds 100 --problem npp
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import random
import sys
import time
from dataclasses import dataclass, field
from fractions import Fraction
from itertools import combinations
from typing import Any, Callable, Optional

# ── Import Q16 fraction system (exact rational, no float in compute path) ──

sys.path.insert(0, "python")
from q16_fraction import Q16, q16_mul, q16_div, q16_add, q16_sub, to_raw, from_raw
from q16_fraction import from_float as q16_from_float


# ═══════════════════════════════════════════════════════════════════════════
# §1  PARAMETER BUDGET (mutable)
# ═══════════════════════════════════════════════════════════════════════════

# BAWIM hardware constraints (Eq.2 context: 15-bit J_ij, 5-30% amplitude)
BAWIM_DEFAULTS = {
    "n_spins": 20,                  # N
    "j_resolution_bits": 15,        # J_ij bit depth (paper: 15-bit)
    "j_max": Fraction(2**15, 1),    # max coupling value
    "feedback_min_pct": Fraction(5, 100),    # 5% of RF carrier
    "feedback_max_pct": Fraction(30, 100),   # 30% of RF carrier
    "h_bias_max": Fraction(1, 10),           # max local bias
    "thermal_stability": Fraction(780, 1),   # BAWIM deg/°C (paper: 780)
    "encoding": "one_hot",          # one_hot | binary | hybrid
}


@dataclass
class BawimParams:
    """Mutable parameter set for the BAWIM system.

    Each field can be mutated independently.  The mutation engine
    varies one field per round and evaluates the impact.
    """
    n_spins: int = BAWIM_DEFAULTS["n_spins"]
    j_resolution_bits: int = BAWIM_DEFAULTS["j_resolution_bits"]
    j_max: Fraction = BAWIM_DEFAULTS["j_max"]
    feedback_min_pct: Fraction = BAWIM_DEFAULTS["feedback_min_pct"]
    feedback_max_pct: Fraction = BAWIM_DEFAULTS["feedback_max_pct"]
    h_bias_max: Fraction = BAWIM_DEFAULTS["h_bias_max"]
    thermal_stability: Fraction = BAWIM_DEFAULTS["thermal_stability"]
    encoding: str = BAWIM_DEFAULTS["encoding"]
    seed: int = 42

    def clone(self) -> BawimParams:
        return BawimParams(
            n_spins=self.n_spins,
            j_resolution_bits=self.j_resolution_bits,
            j_max=self.j_max,
            feedback_min_pct=self.feedback_min_pct,
            feedback_max_pct=self.feedback_max_pct,
            h_bias_max=self.h_bias_max,
            thermal_stability=self.thermal_stability,
            encoding=self.encoding,
            seed=self.seed,
        )

    def fingerprint(self) -> str:
        """Deterministic hash for reproducibility tracking."""
        raw = json.dumps({
            "n_spins": self.n_spins,
            "j_res_bits": self.j_resolution_bits,
            "j_max": str(self.j_max),
            "fb_min": str(self.feedback_min_pct),
            "fb_max": str(self.feedback_max_pct),
            "h_max": str(self.h_bias_max),
            "thermal": str(self.thermal_stability),
            "encoding": self.encoding,
            "seed": self.seed,
        }, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(raw.encode()).hexdigest()[:16]


# ═══════════════════════════════════════════════════════════════════════════
# §2  CORE BAWIM EQUATIONS (exact Fraction arithmetic)
# ═══════════════════════════════════════════════════════════════════════════

def ising_hamiltonian(
    spins: list[int],
    J: list[list[Fraction]],
    h: list[Fraction],
) -> Fraction:
    """Eq 1: H = -∑ J_ij s_i s_j - ∑ h_i s_i

    All arithmetic in exact Fraction.  No float, no truncation.
    """
    H = Fraction(0, 1)
    n = len(spins)
    for i in range(n):
        for j in range(i + 1, n):
            H -= J[i][j] * spins[i] * spins[j]
    for i in range(n):
        H -= h[i] * spins[i]
    return H


def feedback_coupling(
    spins: list[int],
    J: list[list[Fraction]],
    i: int,
) -> Fraction:
    """Eq 2: c_i = ∑ J_ij s_j

    The coupling pulse injected back into spin i.
    Amplitude should stay within feedback_min_pct–feedback_max_pct
    of the circulating RF signal to avoid chaotization.
    """
    return sum(J[i][j] * spins[j] for j in range(len(spins)))


def maxcut_score(
    spins: list[int],
    J: list[list[Fraction]],
    H: Optional[Fraction] = None,
) -> Fraction:
    """Eq 3: Cuts = -½∑ J_ij - ½H

    Lower Hamiltonian → more cuts.
    """
    if H is None:
        H = ising_hamiltonian(spins, J, [Fraction(0, 1)] * len(spins))
    sum_J = sum(J[i][j] for i in range(len(spins)) for j in range(i + 1, len(spins)))
    return -sum_J / 2 - H / 2


def npp_spin_form(
    spins: list[int],
    values: list[Fraction],
) -> Fraction:
    """Eq 5: E(s) = |∑ a_i s_i|

    s_i = +1 → subset A; s_i = -1 → subset B.
    Perfect partition when E = 0 (total even) or E = 1 (total odd).
    """
    total = sum(values[i] * spins[i] for i in range(len(spins)))
    return abs(total)


# ═══════════════════════════════════════════════════════════════════════════
# §3  SPIN CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════

def random_spins(n: int, seed: int) -> list[int]:
    rng = random.Random(seed)
    return [1 if rng.random() < 0.5 else -1 for _ in range(n)]


def flip_spin(spins: list[int], idx: int) -> list[int]:
    s = list(spins)
    s[idx] = -s[idx]
    return s


# ═══════════════════════════════════════════════════════════════════════════
# §4  COUPLING MATRIX GENERATORS (mutation targets)
# ═══════════════════════════════════════════════════════════════════════════

def build_maxcut_graph(
    n: int,
    density: Fraction,
    seed: int,
    j_max: Fraction,
) -> list[list[Fraction]]:
    """Generate a MAX-CUT graph coupling matrix.

    density is the fraction of edges present.
    Each edge gets a random coupling in [0, j_max].
    """
    rng = random.Random(seed)
    J = [[Fraction(0, 1) for _ in range(n)] for _ in range(n)]
    threshold = float(density)
    for i in range(n):
        for j in range(i + 1, n):
            if rng.random() < threshold:
                val = rng.randint(1, int(j_max))
                J[i][j] = Fraction(val, 1)
                J[j][i] = J[i][j]
    return J


def build_npp_matrix(
    values: list[Fraction],
) -> tuple[list[list[Fraction]], list[Fraction]]:
    """Build coupling matrix for number partitioning.

    J_ij = a_i * a_j  (Mattis spin glass form)
    h = 0
    """
    n = len(values)
    J = [[Fraction(0, 1) for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(i + 1, n):
            val = values[i] * values[j]
            J[i][j] = val
            J[j][i] = val
    h = [Fraction(0, 1)] * n
    return J, h


# ═══════════════════════════════════════════════════════════════════════════
# §5  SOLVER (greedy descent — mutation target)
# ═══════════════════════════════════════════════════════════════════════════

def greedy_descent(
    spins: list[int],
    J: list[list[Fraction]],
    h: list[Fraction],
    max_flips: int = 100,
) -> tuple[list[int], Fraction]:
    """Greedy energy minimization by single-spin flips.

    Each flip that lowers H is accepted.  Stops when no flip improves.
    — MUTATION TARGET: replace with simulated annealing, QAOA, etc.
    """
    current = list(spins)
    H = ising_hamiltonian(current, J, h)
    improved = True
    n = len(current)
    while improved and max_flips > 0:
        improved = False
        for i in range(n):
            candidate = flip_spin(current, i)
            H_candidate = ising_hamiltonian(candidate, J, h)
            if H_candidate < H:
                current = candidate
                H = H_candidate
                improved = True
                max_flips -= 1
                break
    return current, H


# ═══════════════════════════════════════════════════════════════════════════
# §6  MUTATION OPERATORS
# ═══════════════════════════════════════════════════════════════════════════

@dataclass
class Mutation:
    """A single mutation: which parameter changed, from what to what."""
    target: str            # parameter name
    old_value: Any
    new_value: Any
    delta_energy: Optional[Fraction] = None  # improvement (negative = better)


MUTATION_TARGETS = [
    # n_spins is excluded: changing N changes the graph, making
    # cross-round energy comparisons meaningless.
    "j_resolution_bits",
    "j_max",
    "feedback_min_pct",
    "feedback_max_pct",
    "h_bias_max",
    "encoding",
    "solver",            # swap solver strategy
]


def mutate_param(params: BawimParams, rng: random.Random) -> tuple[BawimParams, Mutation]:
    """Apply one random mutation to a parameter.

    Returns (new_params, description_of_mutation).
    """
    p = params.clone()
    target = rng.choice(MUTATION_TARGETS)
    old_val = getattr(p, target, None)

    if target == "n_spins":
        delta = rng.choice([-5, -3, -1, 1, 3, 5])
        new = max(4, min(200, p.n_spins + delta))
        p.n_spins = new

    elif target == "j_resolution_bits":
        delta = rng.choice([-2, -1, 1, 2])
        new = max(4, min(20, p.j_resolution_bits + delta))
        p.j_resolution_bits = new
        p.j_max = Fraction(2**new, 1)

    elif target == "j_max":
        scale = rng.choice([Fraction(1, 2), Fraction(3, 4), Fraction(1, 1),
                           Fraction(5, 4), Fraction(3, 2), Fraction(2, 1)])
        p.j_max = (p.j_max * scale).limit_denominator(2**p.j_resolution_bits)

    elif target == "feedback_min_pct":
        step = Fraction(1, 100)
        p.feedback_min_pct = max(Fraction(1, 100), min(
            Fraction(49, 100),
            p.feedback_min_pct + rng.choice([-step, step]) * rng.randint(1, 5)
        ))

    elif target == "feedback_max_pct":
        step = Fraction(1, 100)
        p.feedback_max_pct = max(p.feedback_min_pct + step, min(
            Fraction(95, 100),
            p.feedback_max_pct + rng.choice([-step, step]) * rng.randint(1, 5)
        ))

    elif target == "h_bias_max":
        scale = rng.choice([Fraction(1, 2), Fraction(3, 4), Fraction(1, 1),
                           Fraction(3, 2), Fraction(2, 1)])
        p.h_bias_max = max(Fraction(1, 100), min(Fraction(10, 1),
                           (p.h_bias_max * scale).limit_denominator(100)))

    elif target == "encoding":
        p.encoding = rng.choice(["one_hot", "binary", "hybrid"])

    elif target == "solver":
        # Future: swap between greedy_descent, simulated_annealing, QAOA
        pass

    new_val = getattr(p, target, None)
    return p, Mutation(target=target, old_value=old_val, new_value=new_val)


# ═══════════════════════════════════════════════════════════════════════════
# §7  EVALUATION
# ═══════════════════════════════════════════════════════════════════════════

def evaluate_maxcut(params: BawimParams) -> dict[str, Any]:
    """Build a MAX-CUT problem, solve it, return metrics."""
    rng = random.Random(params.seed)
    J = build_maxcut_graph(
        params.n_spins,
        density=Fraction(1, 2),
        seed=params.seed,
        j_max=params.j_max,
    )
    h = [Fraction(0, 1)] * params.n_spins
    spins = random_spins(params.n_spins, params.seed)
    final_spins, H = greedy_descent(spins, J, h)
    cuts = maxcut_score(final_spins, J, H)
    return {
        "problem": "maxcut",
        "n": params.n_spins,
        "H": H,
        "cuts": cuts,
        "cuts_float": float(cuts),
        "spins": final_spins,
    }


def evaluate_npp(params: BawimParams) -> dict[str, Any]:
    """Build a number partitioning problem, solve it, return metrics."""
    rng = random.Random(params.seed)
    values = [Fraction(rng.randint(1, 1000), 1) for _ in range(params.n_spins)]
    J, h = build_npp_matrix(values)
    spins = random_spins(params.n_spins, params.seed + 1)
    final_spins, H = greedy_descent(spins, J, h)
    E = npp_spin_form(final_spins, values)
    return {
        "problem": "npp",
        "n": params.n_spins,
        "E": E,
        "E_float": float(E),
        "H": H,
        "spins": final_spins,
    }


# ═══════════════════════════════════════════════════════════════════════════
# §8  MUTATION ENGINE
# ═══════════════════════════════════════════════════════════════════════════

@dataclass
class MutationRound:
    """Record of one mutation + evaluation."""
    round: int
    params: BawimParams
    mutation: Mutation
    result: dict[str, Any]
    accepted: bool


class BawimMutationEngine:
    """Run mutation rounds to improve BAWIM solution quality.

    Each round:
      1. Clone current params
      2. Mutate one parameter
      3. Evaluate on the problem
      4. Accept if energy improved (or with probability for exploration)
      5. Record the round
    """

    def __init__(
        self,
        problem: str = "maxcut",
        initial_params: Optional[BawimParams] = None,
        temperature: Fraction = Fraction(1, 10),
    ):
        self.problem = problem
        self.params = initial_params or BawimParams()
        self.temperature = temperature
        self.history: list[MutationRound] = []
        self.best_result: Optional[dict[str, Any]] = None
        self.best_energy: Optional[Fraction] = None
        self.rng = random.Random(self.params.seed)

    def evaluate(self, params: BawimParams) -> dict[str, Any]:
        if self.problem == "maxcut":
            return evaluate_maxcut(params)
        elif self.problem == "npp":
            return evaluate_npp(params)
        else:
            raise ValueError(f"Unknown problem: {self.problem}")

    def energy_key(self, result: dict[str, Any]) -> Fraction:
        if self.problem == "maxcut":
            return -result["cuts"]  # more cuts = better (lower energy)
        else:
            return result["E"]  # lower E = better

    def step(self) -> MutationRound:
        new_params, mutation = mutate_param(self.params, self.rng)
        result = self.evaluate(new_params)
        new_energy = self.energy_key(result)

        accepted = False
        if self.best_energy is None:
            accepted = True
        elif new_energy < self.best_energy:
            accepted = True  # strict improvement
        else:
            # Probabilistic acceptance (simulated annealing style)
            delta = new_energy - self.best_energy
            if self.temperature > 0:
                # Convert Fraction to float for exp (I/O boundary — display only)
                p = math.exp(-float(delta) / float(self.temperature))
                accepted = self.rng.random() < p

        mutation.delta_energy = (
            new_energy - self.best_energy if self.best_energy is not None else Fraction(0, 1)
        )

        if accepted:
            self.params = new_params
            self.best_result = result
            self.best_energy = new_energy

        record = MutationRound(
            round=len(self.history),
            params=new_params,
            mutation=mutation,
            result=result,
            accepted=accepted,
        )
        self.history.append(record)
        return record

    def run(self, rounds: int, verbose: bool = True) -> list[MutationRound]:
        for r in range(rounds):
            record = self.step()
            if verbose and (r < 10 or r % 10 == 0 or record.accepted):
                delta = record.mutation.delta_energy
                sign = "✓" if record.accepted else "✗"
                energy = float(self.energy_key(record.result))
                print(
                    f"  [{r:4d}] {sign} "
                    f"{record.mutation.target:20s} "
                    f"Δ={float(delta):+.6f}  "
                    f"E={energy:.6f}"
                )
        return self.history

    def summary(self) -> dict[str, Any]:
        best = self.best_result or {}
        accepted = sum(1 for h in self.history if h.accepted)
        return {
            "problem": self.problem,
            "rounds": len(self.history),
            "accepted": accepted,
            "accept_rate": accepted / max(1, len(self.history)),
            "best_energy": float(self.best_energy) if self.best_energy else None,
            "best_params": {
                "n_spins": self.params.n_spins,
                "j_resolution_bits": self.params.j_resolution_bits,
                "j_max": str(self.params.j_max),
                "feedback_range": f"{self.params.feedback_min_pct}–{self.params.feedback_max_pct}",
                "encoding": self.params.encoding,
            },
            "best_result": {
                k: v for k, v in best.items() if k != "spins"
            } if best else {},
            "fingerprint": self.params.fingerprint(),
        }


# ═══════════════════════════════════════════════════════════════════════════
# §9  CLI
# ═══════════════════════════════════════════════════════════════════════════

def main() -> int:
    parser = argparse.ArgumentParser(
        description="BAWIM Mutation Engine — explore coupling/encoding design space"
    )
    parser.add_argument("--rounds", type=int, default=50, help="mutation rounds")
    parser.add_argument("--problem", choices=["maxcut", "npp"], default="maxcut")
    parser.add_argument("--n-spins", type=int, default=20)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--temperature", type=float, default=0.1,
                        help="SA acceptance temperature (0 = greedy)")
    parser.add_argument("--json", action="store_true", help="output final summary as JSON")
    args = parser.parse_args()

    params = BawimParams(n_spins=args.n_spins, seed=args.seed)
    engine = BawimMutationEngine(
        problem=args.problem,
        initial_params=params,
        temperature=Fraction(args.temperature).limit_denominator(100),
    )

    print(f"BAWIM Mutation Engine — {args.problem.upper()}")
    print(f"  Initial params: n={params.n_spins}, J_bits={params.j_resolution_bits}")
    print(f"  Temperature: {args.temperature}")
    print(f"  Running {args.rounds} rounds...\n")

    start = time.time()
    engine.run(args.rounds)
    elapsed = time.time() - start

    summary = engine.summary()
    print(f"\n--- Summary ({elapsed:.1f}s) ---")
    print(f"  Problem:        {summary['problem']}")
    print(f"  Rounds:         {summary['rounds']} ({summary['accepted']} accepted)")
    print(f"  Accept rate:    {summary['accept_rate']:.2f}")
    if summary['best_energy'] is not None:
        print(f"  Best energy:    {summary['best_energy']:.6f}")
    print(f"  Best params:    {summary['best_params']}")
    print(f"  Fingerprint:    {summary['fingerprint']}")

    if args.json:
        print(json.dumps(summary, indent=2))

    return 0


if __name__ == "__main__":
    sys.exit(main())