import cvxpy as cp
import numpy as np

def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)):
    """
        Place a grid to cover robot locations with alignment to centers.

        inputs:
            - robot_locations (list): locations of robots involved in conflict [[x,y], [x,y], ...]
            - cell_size (float): the width of each grid cell in continuous space
            - grid_shape (tuple): (# of rows, # of columns) of the grid
        outputs:
            - origin (tuple): bottom-left corner of the grid in continuous space
    """
    robot_locations = np.array(robot_locations)
    N = len(robot_locations)
    
    # Decision variable: Bottom-left corner of the grid in continuous space
    origin = cp.Variable(2, name='origin')
    
    # Decision variable: Integer grid indices for each robot
    grid_indices = cp.Variable((N, 2), integer=True, name='grid_indices')
    
    # Calculate cell centers for each robot based on grid indices
    # Reshape origin to (1, 2) for broadcasting
    cell_centers = cp.reshape(origin, (1, 2), order='C') + grid_indices * cell_size + cell_size / 2
    
    # Objective: Minimize the sum of squared distances
    cost = cp.sum_squares(robot_locations - cell_centers)
    
    # Constraints
    constraints = []
    
    # Grid indices must be non-negative
    constraints.append(grid_indices >= 0)
    
    # Grid indices must fit within grid bounds
    if grid_shape[0] == grid_shape[1]: # Square grid
        constraints.append(grid_indices <= grid_shape[0] - 1)
    else: # Rectangular grid
        constraints.append(grid_indices[:,0] <= grid_shape[1] - 1)
        constraints.append(grid_indices[:,1] <= grid_shape[0] - 1)
    
    # No two robots can share a cell
    # Reformulation of the constraints
    #   abs(grid_indices[i, 0] - grid_indices[j, 0]) >= 1 and
    #   abs(grid_indices[i, 1] - grid_indices[j, 1]) >= 1
    # to be compatible with solver
    for i in range(N):
        for j in range(i+1, N):
            # Auxiliary variable for the distance between cell centers i and j in the x direction
            xdist = cp.Variable(name=f"xdist_{i}_{j}")  
            constraints.append(xdist >= grid_indices[i, 0] - grid_indices[j, 0])
            constraints.append(xdist >= -(grid_indices[i, 0] - grid_indices[j, 0]))
            
            # Auxiliary variable for the distance between cell centers i and j in the y direction
            ydist = cp.Variable(name=f"ydist_{i}_{j}")  
            constraints.append(ydist >= grid_indices[i, 1] - grid_indices[j, 1])
            constraints.append(ydist >= -(grid_indices[i, 1] - grid_indices[j, 1]))
            
            # Enforce that robots must be at least one cell apart
            constraints.append(xdist + ydist >= cell_size) 
    
    # Solve the optimization problem
    prob = cp.Problem(cp.Minimize(cost), constraints)
    prob.solve(solver=cp.SCIP, scip_params={
        "numerics/feastol": 1e-6, 
        "numerics/dualfeastol": 1e-6,
    })

    if prob.status not in ["optimal", "optimal_inaccurate"]:
        print("problem could not be solved to optimality")
        return None
    print(prob.status)
    return origin.value, cell_centers.value
        
def main():
    robot_locations = [(1.2, 1.6), (1.6, 1.2)]
    cell_size = 1
    grid_shape = (5, 5)
    
    origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
    print("Grid Origin (Bottom-Left Corner):", origin)
    print(cell_centers)

    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(4, 4))
    # Draw the grid
    for i in range(grid_shape[1] + 1):
        # Draw vertical lines
        plt.plot([origin[0] + i * cell_size, origin[0] + i * cell_size], 
                    [origin[1], origin[1] + grid_shape[0] * cell_size], 'k-')
    for i in range(grid_shape[0] + 1):
        # Draw horizontal lines
        plt.plot([origin[0], origin[0] + grid_shape[1] * cell_size], 
                    [origin[1] + i * cell_size, origin[1] + i * cell_size], 'k-')

    # Plot robot locations
    robot_locations = np.array(robot_locations)
    plt.scatter(robot_locations[:, 0], robot_locations[:, 1], c='r', label='Robot Locations')

    # Plot cell centers
    cell_centers = np.array(cell_centers)
    plt.scatter(cell_centers[:, 0], cell_centers[:, 1], c='b', label='Cell Centers')
    plt.legend()

    plt.show()

if __name__ == "__main__":
    main()