import cvxpy as cp
import numpy as np

def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=False):
    """
        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
            - cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
            - loss: when return_loss=True, sum of squared differences loss
    """
    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
    # Use Big M method to ensure unique grid indices
    M = max(grid_shape) * 10
    for i in range(N):
        for j in range(i+1, N):
            # At least one of the two constraints below must be true
            y1 = cp.Variable(boolean=True)
            y2 = cp.Variable(boolean=True)
            constraints.append(y1 + y2 >= 1)
            
            # Enforces separation by at least 1 in the x direction
            if robot_locations[i, 0] >= robot_locations[j, 0]:
                constraints.append(grid_indices[i, 0] - grid_indices[j, 0] + M * (1 - y1) >= 1)
            else:
                constraints.append(grid_indices[j, 0] - grid_indices[i, 0] + M * (1 - y1) >= 1)
            
            # Enforces separation by at least 1 in the y direction
            if robot_locations[i, 1] >= robot_locations[j, 1]:
                constraints.append(grid_indices[i, 1] - grid_indices[j, 1] + M * (1 - y2) >= 1)
            else:
                constraints.append(grid_indices[j, 1] - grid_indices[i, 1] + M * (1 - y2) >= 1)
    
    # Solve the optimization problem
    prob = cp.Problem(cp.Minimize(cost), constraints)
    prob.solve(solver=cp.SCIP)

    if prob.status not in ["optimal", "optimal_inaccurate"]:
        print("Problem could not be solved to optimality.")
        return None
    
    if return_loss:
        return origin.value, cell_centers.value, prob.value
    return origin.value, cell_centers.value


def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), num_angles=18, return_angle_info=False):
    """
        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
            - sample_size (int): number of evenly spaced angles to sample between 0 and pi/2 radians
        outputs:
            - grid_center (tuple): center of the grid in continuous space
            - cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
            - angle_info (dict): when return_angle_info=True, dict of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle
    """ 
    
    # Dictionary of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle }
    angle_info = {}
    
    # The rotation angle from the angles sampled which minimizes loss
    min_loss_angle = 0
    min_loss = None 
    
    from tqdm import tqdm
    
    for angle in tqdm(np.linspace(start=0, stop=np.pi/2, num=num_angles, endpoint=False), desc="Processing..."):
        rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)],
                                    [np.sin(angle), np.cos(angle)]])
        
        # Rotate robot locations about the origin in continuous space
        rotated_robot_locations = robot_locations @ rotation_matrix
        
        # Run place_grid on transformed robot locations
        rotated_origin, rotated_cell_centers, loss = place_grid(rotated_robot_locations, cell_size, grid_shape, return_loss=True)
        rotated_grid_center = rotated_origin + np.array(grid_shape) * cell_size / 2
        
        # Undo the rotation transformation for the origin and cell centers 
        grid_center = rotated_grid_center @ np.linalg.inv(rotation_matrix)
        cell_centers = rotated_cell_centers @ np.linalg.inv(rotation_matrix)
        
        # Check if the loss is the smallest observed so far
        if min_loss is None or loss < min_loss:
            min_loss_angle = angle
            min_loss = loss
            
        angle_info[angle] = grid_center, cell_centers, loss
    
    if return_angle_info:
        return *angle_info[min_loss_angle][:2], min_loss_angle, angle_info
    return *angle_info[min_loss_angle][:2], min_loss_angle


def plot_grid(origin, cell_size, grid_shape):
    import matplotlib.pyplot as plt
    
    # 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-')
        
        
def plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle):
    rotation_matrix = np.array([[np.cos(rotation_angle), -np.sin(rotation_angle)],
                                [np.sin(rotation_angle), np.cos(rotation_angle)]])
    grid_center = np.array(grid_center)
    
    import matplotlib.pyplot as plt
    
    # Draw the grid
    for i in range(grid_shape[1] + 1):
        # Construct vertical line endpoints
        line_points = np.array([[(-grid_shape[1] / 2 + i) * cell_size, (-grid_shape[1] / 2 + i) * cell_size],
                                [-grid_shape[0] / 2, grid_shape[0] / 2]])
        # Rotate vertical lines and draw
        rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
        plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')
    for i in range(grid_shape[1] + 1):
        # Construct horizontal line endpoints
        line_points = np.array([[-grid_shape[1] / 2, grid_shape[1] / 2],
                                [(-grid_shape[0] / 2 + i) * cell_size, (-grid_shape[0] / 2 + i) * cell_size]])
        # Rotate horizontal lines and draw
        rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
        plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')


def main(allow_rotations, num_robots, num_angles, seed):
    if seed is not None:
        np.random.seed(seed)
    
    robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
    cell_size = 1
    grid_shape = (5, 5)
    
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(6, 6))
    
    if allow_rotations:
        grid_center, cell_centers, rotation_angle = place_grid_with_rotation(robot_locations, cell_size, grid_shape, num_angles=num_angles)
        print("Grid Center:", grid_center)
        # angles = angle_info.keys()
        # losses = [loss for (_, _, loss) in angle_info.values()]
        # plt.scatter(angles, losses, c='r')
        
        plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle)
    else:
        origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
        print("Grid Origin (Bottom-Left Corner):", origin)
        
        plot_grid(origin, cell_size, grid_shape)
    
    # 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(loc='upper left')

    plt.show()

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--allow_rotations", 
        type=bool, 
        default=False
    )
    parser.add_argument(
        "--num_robots", 
        type=int, 
        default=2
    )
    parser.add_argument(
        "--num_angles", 
        type=int, 
        default=18
    )
    parser.add_argument(
        "--seed", 
        type=int, 
        default=None
    )
    args = parser.parse_args()

    main(args.allow_rotations, args.num_robots, args.num_angles, args.seed)