# Source code for eqc_models.graph.coloring

    # (C) Quantum Computing Inc., 2026.
    from typing import Dict, Tuple
    import numpy as np
    import networkx as nx
    from eqc_models.base.quadratic import ConstrainedQuadraticModel

    class GraphColoringModel(ConstrainedQuadraticModel):
        """
        Graph k-coloring as a constrained quadratic model.

        Binary decision variables x_{i,j} indicate that node i is assigned color
        j, for i over the m graph nodes and j over k colors. The one-hot
        assignment constraint sum_j x_{i,j} = 1 forces each node to take exactly
        one color. The objective counts monochromatic edges (edges whose
        endpoints share a color):

            min  sum_{(u,v) in E} sum_{j=0}^{k-1} x_{u,j} * x_{v,j}

        A proper k-coloring exists iff the optimal objective value is zero.
        """

        _objective = None

        def __init__(self, G: nx.Graph, k: int):
            self.G = G
            self.k = k
            if k < 2:
                raise ValueError("k must be greater than 1")
            A, b = self._build_constraints()
            c, J = self.costFunction()
            ConstrainedQuadraticModel.__init__(self, c, J, A, b)
            n = len(G.nodes) * k
            self.upper_bound = np.ones((n,))

        def decode(self, solution: np.ndarray) -> np.ndarray:
            """ Decode a raw solution to a one-hot color assignment per node. """
            decoded_solution = np.zeros_like(solution, dtype=np.int32)
            k = self.k
            G = self.G
            for i, u in enumerate(G.nodes):
                idx = slice(k * i, k * (i + 1))
                spins = solution[idx]
                mx = np.max(spins)
                for j in range(k):
                    if spins[j] == mx:
                        decoded_solution[k * i + j] = 1
                        break
            return decoded_solution

        def coloring(self, solution: np.ndarray) -> Dict:
            """ Return a dictionary mapping each node to its assigned color (1..k). """
            k = self.k
            G = self.G
            color = {}
            for i, u in enumerate(G.nodes):
                for j in range(k):
                    if solution[i * k + j] == 1:
                        color[u] = j + 1
            return color

        def getConflictCount(self, coloring: Dict) -> int:
            """ Count edges whose endpoints share a color (monochromatic edges). """
            conflicts = 0
            for u, v in self.G.edges:
                if coloring[u] == coloring[v]:
                    conflicts += 1
            return conflicts

        def isProperColoring(self, coloring: Dict) -> bool:
            """ True iff no edge is monochromatic (a valid k-coloring of G). """
            return self.getConflictCount(coloring) == 0

        def costFunction(self) -> Tuple:
            G = self.G
            node_map = list(G.nodes)
            m = len(G.nodes)
            n = self.k * m
            # quadratic objective penalizes same-color assignments on every edge;
            # linear portion is zero
            objective = np.zeros((n, n), dtype=np.float32)
            for u, v in G.edges:
                i = node_map.index(u)
                j = node_map.index(v)
                ibase = i * self.k
                jbase = j * self.k
                for c in range(self.k):
                    objective[ibase + c, jbase + c] += 1
            return (np.zeros((n, 1)), objective)

        def _build_constraints(self):
            G = self.G
            node_map = list(G.nodes)
            m = len(G.nodes)
            n = self.k * m

            # one-hot: exactly one color per node
            A = np.zeros((m, n))
            b = np.ones((m,))
            for u in G.nodes:
                i = node_map.index(u)
                ibase = i * self.k
                A[i, ibase:ibase + self.k] = 1
            return A, b

        @property
        def constraints(self):
            """ Return LHS, RHS in numpy matrix format """

            return self.lhs, self.rhs

        @property
        def objective(self):
            """ Return the quadratic objective as NxN+1 matrix """

            return self._objective