import numpy as np
from models.solvers.ortools.ortools_base import ORToolsBase

class ORToolsPCTSP(ORToolsBase):
    def __init__(self, large_value=1e+6, scaling=False):
        super().__init__(large_value, scaling)

    def scaling_feats(self, node_feats):
        return {
            key: (node_feat * self.large_value + 0.5).astype(np.int64)
            if key in ("coords", "prizes", "penalties", "max_length", "min_prize") else
            node_feat
            for key, node_feat in node_feats.items()
        }

    def add_constraints(self, routing, transit_callback_index, manager, data, node_feats):
        # Add penalties to nodes except for the depot
        # ORTools can ignore the nodes with taking the penalties
        penalties = node_feats["penalties"]
        for i in range(1, len(data['distance_matrix'])):
            index = manager.NodeToIndex(i)
            routing.AddDisjunction([index], penalties[i].item())

        # Add other constraints
        self.add_prize_constraints(routing, data, node_feats)
        # self.add_distance_constraints(routing, transit_callback_index, node_feats)
    
    def add_distance_constraints(self, routing, transit_callback_index, node_feats):
        # Add distance dimension
        dim_name = "Distance"
        routing.AddDimension(
            transit_callback_index,
            0,  # Null capacity slack
            node_feats["max_length"].item(), # Maximum distance constraints
            True,  # Start cumul to zero
            dim_name)

    def add_prize_constraints(self, routing, data, node_feats):
        # Add prize dimension
        dim_name = "Prize"
        prizes = node_feats["prizes"]
        def prize_callback(from_node, to_node):
            return prizes[from_node].item()
        prize_callback_index = routing.RegisterTransitCallback(prize_callback)
        routing.AddDimension(
            prize_callback_index,
            0,  # Null capacity slack
            np.sum(prizes).item(),  # Upper bound
            True,  # Start cumul to zero
            dim_name)

        # Minimum prize constraints
        capacity_dimension = routing.GetDimensionOrDie(dim_name)
        for vehicle_id in range(data["num_vehicles"]):  # Only single vehicle
            capacity_dimension.CumulVar(routing.End(vehicle_id)).RemoveInterval(0, node_feats["min_prize"].item())