to_dimod

@dataclass
class DimodBuilder:
    """Dimod model builder.

    Attributes:
        compiled_model (CompiledInstance): Compiled model.
        _model (dimod.ConstrainedQuadraticModel): Dimod model.
        _objective (dimod.QuadraticModel): Objective of dimod model.
        custome_penalties (dict[str, dimod.BinaryQuadraticModel]): Custom penalties.
    """

    compiled_model: CompiledInstance
    _model: dimod.ConstrainedQuadraticModel
    _objective: dimod.QuadraticModel
    custome_penalties: dict[str, dimod.BinaryQuadraticModel]

    def get_model(
        self, penalty_weight: typ.Optional[dict[str, float]] = None
    ) -> dimod.BinaryQuadraticModel:
        """Get dimod model.

        Args:
            penalty_weight (dict[str, float], optional): Penalty weight. Defaults to None.

        Returns:
            dimod.BinaryQuadraticModel: Dimod model.
        """
        if penalty_weight is None:
            penalty_weight = {}
        if len(self.custome_penalties) == 0:
            self._model.set_objective(self._objective)
            return self._model

        objective = self._objective.copy()
        for label, weight in penalty_weight.items():
            if label not in self.custome_penalties:
                objective = self.custome_penalties[label]
            else:
                objective += weight * self.custome_penalties[label]
        self._model.set_objective(objective)
        return self._model

    def decode_from_cqm_result(self, sampleset: dimod.SampleSet) -> jm.SampleSet:
        """Decode sampleset from dimod model.

        Args:
            sampleset (dimod.SampleSet): Dimod sampleset.

        Returns:
            jm.SampleSet: Decoded sampleset.
        """

        # Parse dimod variables
        # Assumption: sampleset.variables -> [(label, tuple[int, ...]), ...]
        # See detail in transpile_to_dimod
        subscripts: dict[str, list[list[int]]] = {}
        for var_label, var_subscripts in sampleset.variables:  # type: ignore
            if var_label not in subscripts:
                subscripts[var_label] = []
            subscripts[var_label].append(list(var_subscripts))

        solution_set = {
            label: [] for label in self.compiled_model.deci_var_shape.keys()
        }
        for sample in sampleset.samples():
            sparse_solutions: dict = {
                label: (tuple([[] for _ in range(len(shape))]), [], shape)
                for label, shape in self.compiled_model.deci_var_shape.items()
            }
            for key, value in sample.items():
                key_label = key[0]
                key_subs = key[1]
                if value == 0.0:
                    continue
                if key_label in sparse_solutions:
                    # set coordinate
                    for index, coo in enumerate(key_subs):
                        sparse_solutions[key_label][0][index].append(coo)
                    # set value
                    sparse_solutions[key_label][1].append(value)

            for label, sparse_solution in sparse_solutions.items():
                solution_set[label].append(sparse_solution)

        jm_record = jm.Record(
            solution=solution_set,  # type: ignore
            num_occurrences=sampleset.record["num_occurrences"].tolist(),
        )

        # Check feasibility and evaluate violation
        const_violation = {label: [] for label in self.compiled_model.constraint.keys()}
        for index, is_feasible in enumerate(sampleset.record["is_feasible"]):
            if is_feasible:
                for label in self.compiled_model.constraint.keys():
                    const_violation[label].append(0.0)
            else:
                formatted_solution = {}
                sparse_solution: dict = jm_record.solution  # type: ignore
                for label, _solutions in sparse_solution.items():
                    subs, _values, _shape = _solutions[index]
                    _subs_array = np.array(list(subs)).T
                    for _ind, value in enumerate(_values):
                        _s = tuple(_subs_array[_ind].tolist())
                        var_index = self.compiled_model.var_map.var_map[label][_s]
                        formatted_solution[var_index] = value
                violations = eval_constraint(formatted_solution, self.compiled_model)
                for label, violation in violations.items():
                    const_obj = self.compiled_model.problem.constraints[label]
                    cond = const_obj.condition
                    if cond.kind == subs_expr.ConstraintKind.EQUAL:
                        violation_sum = np.sum(np.abs(list(violation.values())))
                    else:
                        abs_value = np.array(list(violation.values()))
                        violation_sum = np.sum(abs_value[abs_value > 0])
                    const_violation[label].append(violation_sum)

        # In CQM solver, energy == objective
        energy = sampleset.record["energy"]
        if self.compiled_model.problem.sense == ProblemSense.MINIMIZE:
            objective = sampleset.record["energy"]
        else:
            objective = -1 * sampleset.record["energy"]
        evaluation = jm.Evaluation(
            energy=energy.tolist(),  # type: ignore
            objective=objective.tolist(),
            constraint_violations=const_violation,
        )

        return jm.SampleSet(
            record=jm_record,  # type: ignore
            evaluation=evaluation,
            measuring_time=jm.MeasuringTime(),
        )