from pulp import *
import numpy as np

def vrp_mip_model(distances, demands, vehicle_capacity, num_vehicles):
    """
    Basic VRP using Mixed-Integer Programming

    Parameters:
    - distances: NxN distance matrix
    - demands: list of demands at each location
    - vehicle_capacity: max capacity per vehicle
    - num_vehicles: number of available vehicles
    """

    n = len(distances)  # includes depot at index 0
    customers = range(1, n)
    all_nodes = range(n)
    vehicles = range(num_vehicles)

    # Create problem
    prob = LpProblem("VRP", LpMinimize)

    # Decision variables
    # x[i,j,k] = 1 if vehicle k travels from i to j
    x = LpVariable.dicts("route",
                         [(i,j,k) for i in all_nodes
                          for j in all_nodes
                          for k in vehicles if i != j],
                         cat='Binary')

    # u[i,k] = position of node i in route of vehicle k (for subtour elimination)
    u = LpVariable.dicts("order",
                         [(i,k) for i in customers for k in vehicles],
                         lowBound=1,
                         upBound=n-1,
                         cat='Integer')

    # Objective: minimize total distance
    prob += lpSum([distances[i][j] * x[i,j,k]
                   for i in all_nodes
                   for j in all_nodes
                   for k in vehicles
                   if i != j])

    # Constraints

    # 1. Each customer visited exactly once
    for j in customers:
        prob += lpSum([x[i,j,k]
                      for i in all_nodes
                      for k in vehicles
                      if i != j]) == 1

    # 2. Flow conservation: enter and leave each node
    for k in vehicles:
        for j in all_nodes:
            prob += (lpSum([x[i,j,k] for i in all_nodes if i != j]) ==
                    lpSum([x[j,i,k] for i in all_nodes if i != j]))

    # 3. Each vehicle starts from depot
    for k in vehicles:
        prob += lpSum([x[0,j,k] for j in customers]) <= 1

    # 4. Each vehicle returns to depot
    for k in vehicles:
        prob += lpSum([x[i,0,k] for i in customers]) <= 1

    # 5. Vehicle capacity constraint
    for k in vehicles:
        prob += lpSum([demands[j] * x[i,j,k]
                      for i in all_nodes
                      for j in customers
                      if i != j]) <= vehicle_capacity

    # 6. Subtour elimination (Miller-Tucker-Zemlin)
    for i in customers:
        for j in customers:
            for k in vehicles:
                if i != j:
                    prob += u[i,k] - u[j,k] + n * x[i,j,k] <= n - 1

    # Solve
    prob.solve(PULP_CBC_CMD(msg=1))

    # Extract routes
    routes = extract_routes(x, vehicles, n)

    return {
        'status': LpStatus[prob.status],
        'total_distance': value(prob.objective),
        'routes': routes
    }

def extract_routes(x, vehicles, n):
    """Extract route sequences from solution"""
    routes = []
    for k in vehicles:
        route = [0]  # start at depot
        current = 0
        while True:
            for j in range(n):
                if j != current and (current,j,k) in x:
                    if x[current,j,k].varValue > 0.5:
                        if j == 0:  # back to depot
                            route.append(0)
                            if len(route) > 2:  # non-empty route
                                routes.append(route)
                            current = -1
                            break
                        else:
                            route.append(j)
                            current = j
                            break
            if current == -1 or current == 0:
                break
    return routes

import numpy as np

def nearest_neighbor_vrp(distance_matrix, demands, capacity, depot=0):
    """
    Nearest neighbor heuristic for VRP

    Simple construction heuristic - builds routes by always
    selecting the nearest unvisited customer
    """

    n = len(distance_matrix)
    unvisited = set(range(1, n))  # exclude depot
    routes = []

    while unvisited:
        route = [depot]
        current_load = 0
        current_location = depot

        while unvisited:
            # Find nearest feasible customer
            nearest = None
            nearest_dist = float('inf')

            for customer in unvisited:
                if current_load + demands[customer] <= capacity:
                    dist = distance_matrix[current_location][customer]
                    if dist < nearest_dist:
                        nearest = customer
                        nearest_dist = dist

            if nearest is None:
                break  # No feasible customer, return to depot

            route.append(nearest)
            current_load += demands[nearest]
            current_location = nearest
            unvisited.remove(nearest)

        route.append(depot)  # return to depot
        routes.append(route)

    total_distance = sum(calculate_route_distance(route, distance_matrix)
                        for route in routes)

    return {
        'routes': routes,
        'total_distance': total_distance,
        'num_vehicles': len(routes)
    }

def calculate_route_distance(route, distance_matrix):
    """Calculate total distance of a route"""
    distance = 0
    for i in range(len(route) - 1):
        distance += distance_matrix[route[i]][route[i+1]]
    return distance

def clarke_wright_savings(distance_matrix, demands, capacity):
    """
    Clarke-Wright Savings Algorithm for VRP

    Starts with individual routes to each customer,
    then merges routes based on savings
    """

    n = len(distance_matrix)
    depot = 0

    # Calculate savings for all customer pairs
    savings = []
    for i in range(1, n):
        for j in range(i+1, n):
            saving = (distance_matrix[depot][i] +
                     distance_matrix[depot][j] -
                     distance_matrix[i][j])
            savings.append((saving, i, j))

    # Sort savings in descending order
    savings.sort(reverse=True, key=lambda x: x[0])

    # Initialize: each customer has its own route
    routes = {i: [depot, i, depot] for i in range(1, n)}
    route_loads = {i: demands[i] for i in range(1, n)}

    # Try to merge routes based on savings
    for saving_value, i, j in savings:
        # Find routes containing i and j
        route_i = None
        route_j = None

        for route_id, route in routes.items():
            if i in route:
                route_i = route_id
            if j in route:
                route_j = route_id

        # Can only merge if i and j are in different routes
        if route_i != route_j:
            route_i_obj = routes[route_i]
            route_j_obj = routes[route_j]

            # Check if merging is feasible
            combined_load = route_loads[route_i] + route_loads[route_j]

            if combined_load <= capacity:
                # Check if i and j are at ends of their routes
                i_at_end = (route_i_obj[1] == i or route_i_obj[-2] == i)
                j_at_end = (route_j_obj[1] == j or route_j_obj[-2] == j)

                if i_at_end and j_at_end:
                    # Merge routes
                    new_route = merge_routes(route_i_obj, route_j_obj, i, j)

                    # Update routes
                    routes[route_i] = new_route
                    route_loads[route_i] = combined_load

                    del routes[route_j]
                    del route_loads[route_j]

    return {
        'routes': list(routes.values()),
        'total_distance': sum(calculate_route_distance(route, distance_matrix)
                            for route in routes.values()),
        'num_vehicles': len(routes)
    }

def merge_routes(route1, route2, i, j):
    """Merge two routes through nodes i and j"""
    # Remove depot endpoints
    r1 = route1[1:-1]
    r2 = route2[1:-1]

    # Orient routes so i and j connect properly
    if r1[-1] == i and r2[0] == j:
        merged = [0] + r1 + r2 + [0]
    elif r1[-1] == i and r2[-1] == j:
        merged = [0] + r1 + r2[::-1] + [0]
    elif r1[0] == i and r2[0] == j:
        merged = [0] + r1[::-1] + r2 + [0]
    elif r1[0] == i and r2[-1] == j:
        merged = [0] + r2 + r1 + [0]
    else:
        merged = [0] + r1 + r2 + [0]  # fallback

    return merged

import math

def sweep_algorithm(coordinates, demands, capacity, depot_idx=0):
    """
    Sweep Algorithm for VRP

    Sorts customers by polar angle from depot,
    then builds routes by sweeping clockwise
    """

    depot = coordinates[depot_idx]
    n = len(coordinates)

    # Calculate angles from depot
    customers = []
    for i in range(n):
        if i != depot_idx:
            dx = coordinates[i][0] - depot[0]
            dy = coordinates[i][1] - depot[1]
            angle = math.atan2(dy, dx)
            customers.append((i, angle, demands[i]))

    # Sort by angle
    customers.sort(key=lambda x: x[1])

    # Build routes by sweeping
    routes = []
    current_route = [depot_idx]
    current_load = 0

    for customer_id, angle, demand in customers:
        if current_load + demand <= capacity:
            current_route.append(customer_id)
            current_load += demand
        else:
            # Start new route
            current_route.append(depot_idx)
            routes.append(current_route)
            current_route = [depot_idx, customer_id]
            current_load = demand

    # Add final route
    if len(current_route) > 1:
        current_route.append(depot_idx)
        routes.append(current_route)

    return {
        'routes': routes,
        'num_vehicles': len(routes)
    }

import random
import math

def simulated_annealing_vrp(initial_solution, distance_matrix,
                           demands, capacity,
                           initial_temp=1000, cooling_rate=0.995,
                           iterations=10000):
    """
    Simulated Annealing for VRP

    Improves initial solution through iterative neighbor search
    with probabilistic acceptance of worse solutions
    """

    current_solution = initial_solution.copy()
    current_cost = calculate_total_distance(current_solution, distance_matrix)

    best_solution = current_solution.copy()
    best_cost = current_cost

    temperature = initial_temp

    for iteration in range(iterations):
        # Generate neighbor solution
        neighbor = generate_neighbor(current_solution, demands, capacity)
        neighbor_cost = calculate_total_distance(neighbor, distance_matrix)

        # Calculate acceptance probability
        delta = neighbor_cost - current_cost

        if delta < 0:
            # Better solution, always accept
            accept = True
        else:
            # Worse solution, accept with probability
            accept_prob = math.exp(-delta / temperature)
            accept = random.random() < accept_prob

        if accept:
            current_solution = neighbor
            current_cost = neighbor_cost

            if current_cost < best_cost:
                best_solution = current_solution.copy()
                best_cost = current_cost

        # Cool down
        temperature *= cooling_rate

    return {
        'routes': best_solution,
        'total_distance': best_cost
    }

def generate_neighbor(routes, demands, capacity):
    """Generate neighbor solution using local search operators"""

    neighbor = [route.copy() for route in routes]
    operator = random.choice(['swap', 'relocate', '2-opt'])

    if operator == 'swap' and len(neighbor) >= 2:
        # Swap customers between two routes
        route1_idx = random.randint(0, len(neighbor) - 1)
        route2_idx = random.randint(0, len(neighbor) - 1)
        while route2_idx == route1_idx:
            route2_idx = random.randint(0, len(neighbor) - 1)

        route1 = neighbor[route1_idx]
        route2 = neighbor[route2_idx]

        if len(route1) > 2 and len(route2) > 2:
            pos1 = random.randint(1, len(route1) - 2)
            pos2 = random.randint(1, len(route2) - 2)
            route1[pos1], route2[pos2] = route2[pos2], route1[pos1]

    elif operator == 'relocate':
        # Move customer from one route to another
        if len(neighbor) >= 2:
            route1_idx = random.randint(0, len(neighbor) - 1)
            route2_idx = random.randint(0, len(neighbor) - 1)
            while route2_idx == route1_idx:
                route2_idx = random.randint(0, len(neighbor) - 1)

            route1 = neighbor[route1_idx]
            route2 = neighbor[route2_idx]

            if len(route1) > 2:
                pos1 = random.randint(1, len(route1) - 2)
                customer = route1.pop(pos1)
                pos2 = random.randint(1, len(route2) - 1)
                route2.insert(pos2, customer)

    elif operator == '2-opt':
        # 2-opt within a single route
        route_idx = random.randint(0, len(neighbor) - 1)
        route = neighbor[route_idx]

        if len(route) > 4:
            i = random.randint(1, len(route) - 3)
            j = random.randint(i + 1, len(route) - 2)
            route[i:j+1] = reversed(route[i:j+1])

    return neighbor

def calculate_total_distance(routes, distance_matrix):
    """Calculate total distance across all routes"""
    total = 0
    for route in routes:
        for i in range(len(route) - 1):
            total += distance_matrix[route[i]][route[i+1]]
    return total

def vrptw_earliest_start_heuristic(customers, vehicles, time_windows, service_times, travel_times):
    """
    VRPTW heuristic: insert customers based on earliest feasible start time

    Parameters:
    - customers: list of customer IDs
    - vehicles: list of vehicle objects
    - time_windows: dict {customer_id: (earliest, latest)}
    - service_times: dict {customer_id: service_duration}
    - travel_times: function(from_loc, to_loc) -> travel_time
    """

    routes = []
    unserved = set(customers)

    for vehicle in vehicles:
        if not unserved:
            break

        route = [vehicle.depot]
        current_time = 0
        current_location = vehicle.depot
        route_load = 0

        while unserved:
            # Find best customer to insert
            best_customer = None
            best_start_time = float('inf')

            for customer in unserved:
                # Check capacity
                if route_load + customer.demand > vehicle.capacity:
                    continue

                # Calculate arrival time
                travel_time = travel_times(current_location, customer.location)
                arrival_time = current_time + travel_time

                # Check time window feasibility
                earliest, latest = time_windows[customer.id]

                if arrival_time > latest:
                    continue  # Too late

                start_time = max(arrival_time, earliest)

                if start_time < best_start_time:
                    best_start_time = start_time
                    best_customer = customer

            if best_customer is None:
                break  # No feasible customer

            # Add customer to route
            route.append(best_customer)
            unserved.remove(best_customer)

            travel_time = travel_times(current_location, best_customer.location)
            arrival_time = current_time + travel_time
            earliest, latest = time_windows[best_customer.id]
            current_time = max(arrival_time, earliest) + service_times[best_customer.id]
            current_location = best_customer.location
            route_load += best_customer.demand

        route.append(vehicle.depot)
        routes.append(route)

    return routes

def assign_customers_to_depots(customers, depots, distance_matrix):
    """
    Assign customers to nearest depot for MDVRP

    Returns clusters of customers for each depot
    """

    depot_clusters = {depot: [] for depot in depots}

    for customer in customers:
        # Find nearest depot
        nearest_depot = None
        min_distance = float('inf')

        for depot in depots:
            dist = distance_matrix[customer][depot]
            if dist < min_distance:
                min_distance = dist
                nearest_depot = depot

        depot_clusters[nearest_depot].append(customer)

    return depot_clusters

def solve_mdvrp(customers, depots, distance_matrix, demands, capacity):
    """
    Solve Multi-Depot VRP

    1. Assign customers to depots
    2. Solve VRP for each depot independently
    """

    # Assign customers to depots
    clusters = assign_customers_to_depots(customers, depots, distance_matrix)

    all_routes = []
    total_distance = 0

    for depot, depot_customers in clusters.items():
        if not depot_customers:
            continue

        # Solve VRP for this depot
        depot_solution = clarke_wright_savings(
            distance_matrix,
            demands,
            capacity
        )

        all_routes.extend(depot_solution['routes'])
        total_distance += depot_solution['total_distance']

    return {
        'routes': all_routes,
        'total_distance': total_distance,
        'depot_clusters': clusters
    }