import kaiwu as kw
import numpy as np
from kaiwu.cim import SimulatedCIMOptimizer

class CIMSolver:
    def __init__(self):
        self.rng = np.random.default_rng(42)
        # kw.license.init(user_id=   , sdk_code=    )
        self.worker = SimulatedCIMOptimizer(
            pump=1.0, noise=0.05, laps=120,
            delta_time=0.2, normalization=1.0,
            iterations=25, size_limit=100, rng=self.rng
        )

    def solve(self, G):
        n = G.numberOfNodes()
        w = np.zeros((n, n))
        for u, v, w_uv in G.iterEdgesWeights():
            w[u, v] = w[v, u] = w_uv

        ising_matrix = -w
        output = self.worker.solve(ising_matrix)

        opt = kw.sampler.optimal_sampler(ising_matrix, output, 0)
        best_solution = opt[0][0]
        best_solution = (best_solution > 0).astype(int)
        max_cut_value = (np.sum(-ising_matrix) - np.dot(-ising_matrix, best_solution).dot(best_solution)) / 4
        # max_cut_value = best_solution.T @ ising_matrix @ best_solution  # x^T W x

        return best_solution, max_cut_value

