"""
helix_cdc/routing/qubo_adapters.py
===================================

Format converters between QUBO representations used across the Helix substrate.

Formats:
1. **Routing format**: Dict {(i,j): weight} - used by helix_cdc.routing
2. **Operator format**: {"size": n, "terms": [...], "offset": 0.0} - used by quantum operators
3. **Canonical format**: {"n": n, "linear": {...}, "quadratic": {...}} - helix_qubo:v1 schema
4. **BQM format**: (linear_dict, quadratic_dict) - used by dimod

Usage:
    from helix_cdc.routing.qubo_adapters import (
        routing_qubo_to_operator_format,
        operator_format_to_routing_qubo,
        canonical_to_routing_qubo,
        routing_qubo_to_canonical,
        load_qubo_file,
    )
"""

from __future__ import annotations

import json
import hashlib
from pathlib import Path
from typing import Dict, Tuple, Any, Optional, List
from dataclasses import dataclass
from datetime import datetime, timezone


@dataclass
class QUBOProblem:
    """Unified QUBO problem representation."""

    n: int
    linear: Dict[int, float]
    quadratic: Dict[Tuple[int, int], float]
    offset: float = 0.0
    constraints: Optional[Dict[str, Any]] = None
    metadata: Optional[Dict[str, Any]] = None

    def to_routing_format(self) -> Tuple[Dict[Tuple[int, int], float], int]:
        """Convert to routing format {(i,j): weight}."""
        Q = {}
        # Linear terms as diagonal
        for i, h in self.linear.items():
            Q[(i, i)] = h
        # Quadratic terms
        for (i, j), J in self.quadratic.items():
            Q[(i, j)] = J
        return Q, self.n

    def to_operator_format(self) -> Dict[str, Any]:
        """Convert to quantum operator format."""
        terms = []
        for i, h in self.linear.items():
            terms.append({"i": i, "j": i, "w": h})
        for (i, j), J in self.quadratic.items():
            terms.append({"i": i, "j": j, "w": J})
        return {
            "size": self.n,
            "terms": terms,
            "offset": self.offset,
            "label": self.metadata.get("label", "unlabeled") if self.metadata else "unlabeled"
        }

    def to_canonical_format(self) -> Dict[str, Any]:
        """Convert to canonical helix_qubo:v1 format."""
        result = {
            "n": self.n,
            "linear": {str(i): v for i, v in self.linear.items()},
            "quadratic": {f"{i},{j}": v for (i, j), v in self.quadratic.items()},
            "offset": self.offset,
        }
        if self.constraints:
            result["constraints"] = self.constraints
        if self.metadata:
            result["metadata"] = self.metadata
        return result

    def compute_sha256(self) -> str:
        """Compute deterministic SHA256 hash of QUBO."""
        canonical = {
            "n": self.n,
            "linear": {str(i): float(v) for i, v in sorted(self.linear.items())},
            "quadratic": {f"{i},{j}": float(v) for (i, j), v in sorted(self.quadratic.items())},
            "offset": float(self.offset),
        }
        blob = json.dumps(canonical, sort_keys=True, separators=(",", ":")).encode()
        return hashlib.sha256(blob).hexdigest()


# =============================================================================
# Routing Format Converters
# =============================================================================

def routing_qubo_to_operator_format(
    Q: Dict[Tuple[int, int], float],
    n: int,
    label: str = "unlabeled",
) -> Dict[str, Any]:
    """
    Convert routing QUBO {(i,j): weight} to operator format.

    Args:
        Q: QUBO dict with tuple keys
        n: Number of variables
        label: Optional problem label

    Returns:
        Operator format dict for DWave/IBM operators
    """
    terms = []
    for (i, j), w in Q.items():
        terms.append({"i": int(i), "j": int(j), "w": float(w)})

    return {
        "size": n,
        "terms": terms,
        "offset": 0.0,
        "label": label,
    }


def operator_format_to_routing_qubo(
    op: Dict[str, Any],
) -> Tuple[Dict[Tuple[int, int], float], int]:
    """
    Convert operator format to routing QUBO.

    Args:
        op: Operator format dict with "size" and "terms"

    Returns:
        (Q, n) where Q is {(i,j): weight}
    """
    n = int(op["size"])
    Q = {}

    for term in op["terms"]:
        i, j, w = int(term["i"]), int(term["j"]), float(term["w"])
        key = (min(i, j), max(i, j)) if i != j else (i, j)
        Q[key] = Q.get(key, 0.0) + w

    return Q, n


# =============================================================================
# Canonical Format Converters
# =============================================================================

def canonical_to_routing_qubo(
    canonical: Dict[str, Any],
) -> Tuple[Dict[Tuple[int, int], float], int]:
    """
    Convert canonical helix_qubo:v1 format to routing QUBO.

    Args:
        canonical: Dict with "n", "linear", "quadratic"

    Returns:
        (Q, n) where Q is {(i,j): weight}
    """
    n = int(canonical["n"])
    Q = {}

    # Linear terms as diagonal
    for k, v in canonical.get("linear", {}).items():
        i = int(k)
        Q[(i, i)] = float(v)

    # Quadratic terms
    for k, v in canonical.get("quadratic", {}).items():
        i, j = map(int, k.split(","))
        Q[(min(i, j), max(i, j))] = float(v)

    return Q, n


def routing_qubo_to_canonical(
    Q: Dict[Tuple[int, int], float],
    n: int,
    offset: float = 0.0,
    constraints: Optional[Dict[str, Any]] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Convert routing QUBO to canonical helix_qubo:v1 format.

    Args:
        Q: QUBO dict with tuple keys
        n: Number of variables
        offset: Constant offset
        constraints: Optional constraint annotations
        metadata: Optional problem metadata

    Returns:
        Canonical format dict
    """
    linear = {}
    quadratic = {}

    for (i, j), w in Q.items():
        if i == j:
            linear[str(i)] = float(w)
        else:
            # Ensure upper triangular (i < j)
            key = f"{min(i,j)},{max(i,j)}"
            quadratic[key] = float(w)

    result = {
        "n": n,
        "linear": linear,
        "quadratic": quadratic,
        "offset": offset,
    }

    if constraints:
        result["constraints"] = constraints
    if metadata:
        result["metadata"] = metadata

    return result


# =============================================================================
# BQM Format Converters
# =============================================================================

def routing_qubo_to_bqm(
    Q: Dict[Tuple[int, int], float],
) -> Tuple[Dict[int, float], Dict[Tuple[int, int], float]]:
    """
    Convert routing QUBO to BQM format (linear, quadratic).

    This is the format expected by dimod.BinaryQuadraticModel.

    Args:
        Q: QUBO dict with tuple keys

    Returns:
        (linear, quadratic) dicts
    """
    linear = {}
    quadratic = {}

    for (i, j), w in Q.items():
        if i == j:
            linear[i] = linear.get(i, 0.0) + float(w)
        else:
            # Ensure upper triangular
            key = (min(i, j), max(i, j))
            quadratic[key] = quadratic.get(key, 0.0) + float(w)

    return linear, quadratic


def bqm_to_routing_qubo(
    linear: Dict[int, float],
    quadratic: Dict[Tuple[int, int], float],
) -> Dict[Tuple[int, int], float]:
    """
    Convert BQM format to routing QUBO.

    Args:
        linear: Linear coefficients
        quadratic: Quadratic coefficients

    Returns:
        Q dict with tuple keys
    """
    Q = {}

    for i, h in linear.items():
        Q[(i, i)] = float(h)

    for (i, j), J in quadratic.items():
        Q[(min(i, j), max(i, j))] = float(J)

    return Q


# =============================================================================
# File I/O
# =============================================================================

def load_qubo_file(path: Path) -> QUBOProblem:
    """
    Load QUBO from JSON file (auto-detect format).

    Supports:
    - Canonical format (helix_qubo:v1)
    - Operator format (terms list)
    - Legacy scripts/qubo_build.py format

    Args:
        path: Path to QUBO JSON file

    Returns:
        QUBOProblem instance
    """
    data = json.loads(Path(path).read_text())

    # Detect format
    if "terms" in data:
        # Operator format
        Q, n = operator_format_to_routing_qubo(data)
        linear = {i: Q.get((i, i), 0.0) for i in range(n) if (i, i) in Q}
        quadratic = {k: v for k, v in Q.items() if k[0] != k[1]}
        return QUBOProblem(
            n=n,
            linear=linear,
            quadratic=quadratic,
            offset=data.get("offset", 0.0),
            metadata={"source": "operator_format", "label": data.get("label", "")},
        )

    elif "n" in data and "linear" in data:
        # Canonical or legacy format
        n = int(data["n"])
        linear = {int(k): float(v) for k, v in data.get("linear", {}).items()}
        quadratic = {}
        for k, v in data.get("quadratic", {}).items():
            i, j = map(int, k.split(","))
            quadratic[(min(i, j), max(i, j))] = float(v)

        return QUBOProblem(
            n=n,
            linear=linear,
            quadratic=quadratic,
            offset=data.get("offset", 0.0),
            constraints=data.get("constraints"),
            metadata=data.get("metadata", {"source": "canonical"}),
        )

    elif "qubo" in data:
        # Wrapper format (from scripts/qubo_build.py)
        inner = data["qubo"]
        n = int(inner["n"])
        linear = {int(k): float(v) for k, v in inner.get("linear", {}).items()}
        quadratic = {}
        for k, v in inner.get("quadratic", {}).items():
            i, j = map(int, k.split(","))
            quadratic[(min(i, j), max(i, j))] = float(v)

        return QUBOProblem(
            n=n,
            linear=linear,
            quadratic=quadratic,
            offset=inner.get("offset", 0.0),
            constraints={"k_hot": inner.get("k")} if "k" in inner else None,
            metadata={
                "source": "tensor",
                "tensor_ref": data.get("tensor_ref"),
                "label": data.get("tensor_ref", ""),
                "sha256": data.get("qubo_sha256"),
            },
        )

    else:
        raise ValueError(f"Unknown QUBO format in {path}")


def save_qubo_file(
    problem: QUBOProblem,
    path: Path,
    format: str = "canonical",
) -> None:
    """
    Save QUBO to JSON file.

    Args:
        problem: QUBOProblem instance
        path: Output path
        format: "canonical" | "operator"
    """
    if format == "canonical":
        data = problem.to_canonical_format()
        # Add SHA256 if not present
        if "metadata" not in data:
            data["metadata"] = {}
        if "sha256" not in data.get("metadata", {}):
            data["metadata"]["sha256"] = problem.compute_sha256()
        data["metadata"]["created_utc"] = datetime.now(timezone.utc).isoformat()
    elif format == "operator":
        data = problem.to_operator_format()
    else:
        raise ValueError(f"Unknown format: {format}")

    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(data, indent=2))


# =============================================================================
# Validation
# =============================================================================

def validate_qubo(
    data: Dict[str, Any],
    strict: bool = False,
) -> Dict[str, Any]:
    """
    Validate QUBO against helix_qubo:v1 schema.

    Args:
        data: QUBO dict to validate
        strict: If True, fail on warnings

    Returns:
        {"valid": bool, "errors": [...], "warnings": [...]}
    """
    errors = []
    warnings = []

    # Required fields
    if "n" not in data:
        errors.append("Missing required field: n")
    if "linear" not in data:
        errors.append("Missing required field: linear")
    if "quadratic" not in data:
        errors.append("Missing required field: quadratic")

    if errors:
        return {"valid": False, "errors": errors, "warnings": warnings}

    n = data["n"]

    # Validate n
    if not isinstance(n, int) or n < 1:
        errors.append(f"Invalid n: {n} (must be positive integer)")

    # Validate linear
    for k, v in data.get("linear", {}).items():
        try:
            i = int(k)
            if i < 0 or i >= n:
                errors.append(f"Linear index out of range: {k}")
        except ValueError:
            errors.append(f"Invalid linear key: {k}")
        if not isinstance(v, (int, float)):
            errors.append(f"Invalid linear value for {k}: {v}")

    # Validate quadratic
    for k, v in data.get("quadratic", {}).items():
        try:
            parts = k.split(",")
            if len(parts) != 2:
                errors.append(f"Invalid quadratic key format: {k}")
                continue
            i, j = int(parts[0]), int(parts[1])
            if i < 0 or i >= n or j < 0 or j >= n:
                errors.append(f"Quadratic index out of range: {k}")
            if i >= j:
                warnings.append(f"Quadratic key not upper triangular: {k}")
        except ValueError:
            errors.append(f"Invalid quadratic key: {k}")
        if not isinstance(v, (int, float)):
            errors.append(f"Invalid quadratic value for {k}: {v}")

    # Check for NaN/Inf
    for k, v in data.get("linear", {}).items():
        if isinstance(v, float) and (v != v or abs(v) == float("inf")):
            errors.append(f"Linear term {k} has NaN/Inf value")
    for k, v in data.get("quadratic", {}).items():
        if isinstance(v, float) and (v != v or abs(v) == float("inf")):
            errors.append(f"Quadratic term {k} has NaN/Inf value")

    valid = len(errors) == 0 and (not strict or len(warnings) == 0)
    return {"valid": valid, "errors": errors, "warnings": warnings}


# =============================================================================
# Energy Evaluation
# =============================================================================

def evaluate_qubo_energy(
    bits: List[int],
    Q: Dict[Tuple[int, int], float],
    offset: float = 0.0,
) -> float:
    """
    Evaluate QUBO energy for a bitstring.

    Args:
        bits: Binary solution vector
        Q: QUBO dict (routing format)
        offset: Constant offset

    Returns:
        Energy value
    """
    energy = offset
    for (i, j), w in Q.items():
        energy += w * bits[i] * bits[j]
    return energy
