from flask import Blueprint
from itertools import permutations
from app.extensions import db
from app.models.vehicle import Vehicle
from app.models.shift import Shift, ShiftSchema
from app.models.battery_change import BatteryChange
from math import sqrt, pow

auto_routes = Blueprint('auto_routes', __name__)

@auto_routes.route('/auto/<lat>/<long>')
def generate_auto_shift(lat, long):
    lat = float(lat)
    long = float(long)
    closest_vehicles = get_closest_vehicles(lat, long, 20)

    solution = kruskals(closest_vehicles)

    order = get_final_order(solution, closest_vehicles)  
        
    new_shift = Shift()
    for i, ind in enumerate(order):
        id = closest_vehicles[ind].id
        new_change = BatteryChange(**{
            "order": i,
            "shift_id": new_shift.id,
            "vehicle_id": id
        })
        new_shift.battery_changes.append(new_change)
    db.session.add(new_shift)
    db.session.commit()

    shift_schema = ShiftSchema(many=False)

    return shift_schema.dumps(new_shift)

def get_closest_vehicles(lat, long, num):
    vehicles = Vehicle.query.all()
    distances = []
    for vehicle in vehicles:
        distance = sqrt(pow(abs(lat) - abs(vehicle.location_lat), 2) + pow(abs(long) - abs(vehicle.location_long), 2))
        distances.append((vehicle, distance))

    distances = sorted(distances, key=lambda x: x[1])[:num]
    return [x[0] for x in distances]

def get_distance_between_vehicles(left, right):
    # Assumes all lats/longs of cars to be consistently positive or negatively.
    # Would most likely include data point for city/region and only query cars in that area
    return sqrt(pow(abs(left.location_lat) - abs(right.location_lat), 2) + pow(abs(left.location_long) - abs(right.location_long), 2))

def get_sorted_edges(vehicles):
    '''
    Returns list of tuples representing edges (<source_node>, <other_node>, <distance>) 
    where node is index of vehicle in sorted vehicles by distance from starting point.
    The return list is sorted by edge distance.
    '''
    if not vehicles or len(vehicles) == 1: return []

    edge_list = []
    for source_ind, vehicle in enumerate(vehicles):
        for other_ind, other_vehicle in enumerate(vehicles):
            # So we don't double count edges
            if source_ind < other_ind:
                distance = get_distance_between_vehicles(vehicle, other_vehicle)
                edge_list.append((source_ind, other_ind, distance))

    return sorted(edge_list, key = lambda edge: edge[2])

def kruskals(vehicles):
    '''
    Gets the optimal edges to reduce the distance travelled between nodes.
    Returns list of lists where each index represents a vehicle in the closest
    vehicle list (indexes match) and each sublist represents the edges
    to other nodes
    '''
    edges = get_sorted_edges(vehicles)

    parents = [i for i in range(len(vehicles))]
    ranks = [0 for _ in range(len(vehicles))]
    solution = [[] for _ in range(len(vehicles))]

    def find(vertex, parents):
        if vertex != parents[vertex]:
            parents[vertex] = find(parents[vertex], parents)

        return parents[vertex]
    
    def union(root1, root2, parents, ranks):
        if ranks[root1] < ranks[root2]:
            parents[root1] = root2
        elif ranks[root1] > ranks[root2]:
            parents[root2] = root1
        else:
            parents[root2] = root1
            ranks[root1] += 1

    for edge in edges:
        root1 = find(edge[0], parents)
        root2 = find(edge[1], parents)
        if root1 != root2:
            solution[edge[0]].append((edge[1], edge[2]))
            solution[edge[1]].append((edge[0], edge[2]))
            union(root1, root2, parents, ranks)

    return solution

def get_final_order(solution, closest_vehicles):
    '''
    Uses the solution from kruskals algorithm to determine the final route.
    '''

    # Find first vertex with only one edge, this will be out starting point
    # i.e. Closest leaf node
    ind = 0
    for i in range(len(solution)):
        if len(solution[i]) == 1:
            ind = i
            break
            
    order = [ind]
    # Keep a stack to go back to previous nodes if we hit a dead end
    stack = []
    while len(order) != len(closest_vehicles):
        while len(solution[ind]) != 0:
            # Iterate through edges of current node, removing once travelled
            edge = solution[ind].pop(0)

            # New node, add to order and check that node's edges
            if edge[0] not in order:
                if solution[ind]:
                    stack.append(ind)
                ind = edge[0]
                order.append(ind)
                break
            # Hit the end of edges for current node, go back
            if len(solution[ind]) == 0 and stack:
                ind = stack.pop()
            
            # If we're out of edges in the current node and nowhere to go back to, 
            # we must have visited all nodes, return out
            elif len(solution[ind]) == 0 and not stack:
                break

    return order