import itertools
import math
import random
import numpy as np
from .architecture import vertical_neighbors, horizontal_neighbors
import rustworkx as rx


def route_gate(indexed_gate, grid_len, grid_height, msf_faces, mapping, to_remove, take_first_ms):
    device_graph = rx.generators.grid_graph(rows=grid_height, cols=grid_len)
    device_graph.remove_nodes_from(list(to_remove))
    shortest_path_len =  2**31 - 1
    shortest_pair = None

    id,gate = indexed_gate 
    if len(gate) == 2:  
        pairs = [(vn, hn)  for vn in vertical_neighbors(mapping[gate[0]],grid_len, grid_height, omitted_edges=[])
                    for hn in horizontal_neighbors(mapping[gate[1]],grid_len, grid_height, omitted_edges=[])]
    else:
        sorted_msf = sorted(msf_faces, key=lambda m: abs(list(reversed(divmod(m, grid_len)))[0] - list(reversed(divmod(mapping[gate[0]], grid_len)))[0]) +abs(list(reversed(divmod(m, grid_len)))[1] - list(reversed(divmod(mapping[gate[0]], grid_len)))[1]))
        pairs =  [(vn, hn)
                for magic_state in sorted_msf
                for vn in vertical_neighbors(mapping[gate[0]],grid_len, grid_height, omitted_edges=[])
                for hn in horizontal_neighbors(magic_state,grid_len, grid_height, omitted_edges=[])]
    pairs = filter(lambda p : device_graph.has_node(p[0]) and device_graph.has_node(p[1]), pairs )       
    for s,t in pairs:
        const_1 = lambda _ : 1
        dist_dict = (rx.dijkstra_shortest_path_lengths(device_graph, edge_cost_fn=const_1, node=s, goal=t ))
        if t in dist_dict.keys():
            dist = dist_dict[t]
        else: dist = math.inf
        if dist < shortest_path_len:
            shortest_path_len = dist
            shortest_pair = s,t
            if take_first_ms and len(gate) == 1:
                break
    if shortest_pair != None:
        s,t = shortest_pair 
        path = list(rx.dijkstra_shortest_paths(device_graph, source=s, target = t)[t])
        if s not in path:
            path = [s]+path
        if t not in path:
            path.append(t)
        route = [(id, gate, path)]
        for v in path:
            to_remove.add(v)
    else:
        route = []
    return route, to_remove
    
def try_order(order, executable, grid_len, grid_height,msf_faces, mapping, take_first_ms):
    step = []
    to_remove=initialize_to_remove(msf_faces, mapping)
    for i in range(len(executable)):
        gate = list(executable.items())[order[i]]
        route, to_remove = route_gate(gate, grid_len, grid_height, msf_faces, mapping, to_remove,take_first_ms)
        step.extend(route)
    return step

def initialize_to_remove(msf_faces, mapping):
    to_remove = set()
    for q in mapping.keys():
             to_remove.add(mapping[q])
    for f in msf_faces:
            to_remove.add(f)
    return to_remove

def shortest_path(gate, mapping, grid_len, grid_height, msf_faces):
    shortest_path_len =  2**31 - 1
    if len(gate) == 1:
        return shortest_path_len
    if len(gate) == 2:
        pairs = [(vn, hn)  for vn in vertical_neighbors(mapping[gate[0]],grid_len, grid_height, omitted_edges=[])
                    for hn in horizontal_neighbors(mapping[gate[1]],grid_len, grid_height, omitted_edges=[])]
    else:
        pairs =  [(vn, hn)
                for magic_state in msf_faces
                for vn in vertical_neighbors(mapping[gate[0]],grid_len, grid_height, omitted_edges=[])
                for hn in horizontal_neighbors(magic_state,grid_len, grid_height, omitted_edges=[])]
    for s,t in pairs:
        dist = rx.dijkstra_shortest_path_lengths(rx.generators.grid_graph(rows=grid_height, cols=grid_len), node=s, goal=t )[t]
        if dist < shortest_path_len:
            shortest_path_len = dist
    return shortest_path_len

def gates_routed(step, remaining_gates, crit_dict):
    return len(step)

def criticality(step, remaining_gates, crit_dict):
    paths = 0
    for id,qubits,path in step:
        dependent = get_dependent_gates((id,qubits), remaining_gates)
        depths = get_depth_by_qubit(dependent)
        crit_path = max(depths.get(q,0) for q in depths.keys())
        paths += 1+crit_path
    return paths

def criticality_fast(step, remaining_gates, crit_dict):
    paths = 0
    for id,qubits,path in step:
        paths += crit_dict[id]
    return paths 

def build_crit_dict(gates):
    crit_dict ={}
    for id,qubits in gates.items():
        dependent = get_dependent_gates((id,qubits), gates)
        depths = get_depth_by_qubit(dependent)
        crit_path = max(depths.get(q,0) for q in qubits)
        crit_dict[id] = crit_path
    return crit_dict
def build_crit_dict_fast(gates):
    crit_dict ={}
    for id,qubits in gates.items():
        depths = get_depth_by_qubit_p(id, gates)
        crit_path = max(depths.get(q,0) for q in qubits)
        crit_dict[id] = crit_path
    return crit_dict

def dependent(step, remaining_gates):
    deps = 0
    for id,qubits,path in step:
        dependent = get_dependent_gates((id,qubits), remaining_gates)
        deps += len(dependent)
    return deps    

def best_realizable_set_found(gates, executable, arch, mapping, initial_order, reward_name, crit_dict, order_fraction, temperature=10, cooling_rate=0.1, termination_temp=0.1, take_first_ms=False):
    grid_len = arch['width']
    grid_height = arch['height']
    msf_faces = arch['magic_states']
    t_indices = [i  for (i,(id, gate)) in enumerate(executable.items()) if len(gate) == 1]
    cnot_indices = [i for (i,(id, gate)) in enumerate(executable.items()) if len(gate) == 2]
    if initial_order  == 'naive':
        best_order = cnot_indices + t_indices
        best_step = (try_order(best_order, executable, grid_len, grid_height, msf_faces,mapping, take_first_ms))
        current_order = best_order
        current_step = best_step
    elif initial_order == 'random':
        best_order = cnot_indices + t_indices
        random.shuffle(best_order)
        best_step = (try_order(best_order, executable, grid_len, grid_height, msf_faces, mapping, take_first_ms))
        current_order = best_order
        current_step = best_step
    elif initial_order == "shortest_first":
        shortest_cnot = sorted(tuple(range(len(cnot_indices))), key= lambda x : shortest_path(list(executable.items())[x][1], mapping, grid_len, grid_height, msf_faces))
        shortest_t = sorted(tuple(range(len(t_indices))), key= lambda x : shortest_path(list(executable.items())[x][1], mapping, grid_len, grid_height, msf_faces))
        shortest_first = shortest_cnot+shortest_t
        best_order = shortest_first
        best_step = (try_order(best_order, executable, grid_len, grid_height, msf_faces, mapping, take_first_ms))
        current_order = best_order
        current_step = best_step
    routed_ids = [x[0] for x in  best_step]
    best_remaining_gates  = {k:v for k,v in gates.items() if k not in routed_ids}
    name_to_func = { "gates_routed" : gates_routed,
                     "criticality" : criticality_fast,
                     "dependent" : dependent
    }

    reward_func = name_to_func[reward_name]
    orders_tried_count = 1 
    if len(executable) < 2:
        return best_step, 1
    
    elif (len(cnot_indices) < 5 and  len(t_indices) < 5) and cooling_rate != 1:
        #print("exhaustive step")
        all_cnot_orders = itertools.permutations(cnot_indices)
        all_t_orders = itertools.permutations(t_indices)
        orders = list(itertools.product(all_cnot_orders, all_t_orders))
        random.shuffle(orders)
        sample_size = int(len(orders)*order_fraction)
        #print(sample_size, len(orders))

        orders_to_explore = orders[:sample_size]
        for cnot_order, t_order in orders_to_explore:
            order = list(cnot_order) + list(t_order)
            new_step = try_order(order, executable, grid_len, grid_height, msf_faces, mapping, take_first_ms)
            orders_tried_count += 1
            routed_ids = [x[0] for x in  new_step]
            new_remaining_gates = {k:v for k,v in gates.items() if k not in routed_ids}
            #print(f"considering step {new_step}", f"reward: {reward_func(new_step, new_remaining_gates)}")
            if reward_func(new_step, new_remaining_gates, crit_dict)> reward_func(best_step, best_remaining_gates, crit_dict):
                best_step = new_step
        return best_step, orders_tried_count
    
    else:
        while temperature > termination_temp:
            new_order = current_order.copy()
            cnots, ts = new_order[:len(cnot_indices)], new_order[len(cnot_indices):]
            if len(cnots) > 1:
                ind1, ind2 = np.random.choice(range(len(cnot_indices)), size=2, replace=False)
                cnots[ind1], cnots[ind2] = cnots[ind2], cnots[ind1]
            if len(ts) > 1:
                ind1, ind2 = np.random.choice(range(len(t_indices)), size=2, replace=False)
                ts[ind1], ts[ind2] = ts[ind2], ts[ind1]
            new_order = cnots + ts 
            new_step = try_order(new_order, executable, grid_len, grid_height, msf_faces, mapping, take_first_ms)
            orders_tried_count += 1
            routed_ids = [x[0] for x in  new_step]
            new_remaining_gates = {k:v for k,v in gates.items() if k not in routed_ids}
            delta_curr = reward_func(current_step, best_remaining_gates, crit_dict) - reward_func(new_step, new_remaining_gates, crit_dict)
            delta_best = reward_func(best_step, best_remaining_gates, crit_dict) - reward_func(new_step, new_remaining_gates, crit_dict)
            if delta_curr < 0 or np.random.rand() < np.exp(-delta_curr / temperature):
                current_order = new_order
                current_step = new_step
            if delta_best < 0:
                #print(len(best_step))
                best_order = new_order
                best_step = new_step
            temperature *= 1 - cooling_rate
        return best_step, orders_tried_count

def sim_anneal_route(gates, arch, mapping, temperature, cooling_rate, termination_temp, order_fraction, initial_order='random', reward_name="criticality", take_first_ms=True):
    timesteps = []
    grid_len = arch['width']
    grid_height = arch['height']
    msf_faces = arch['magic_states']
    mapping = {q : p for (q,p) in mapping}
    gates_id_table = { i : gate for i,gate, in enumerate(gates)}
    crit_dict = {}
    if temperature > termination_temp:
        crit_dict = build_crit_dict_fast(gates_id_table)
    tried_steps = 0
    while len(gates_id_table) != 0:
        executable, remaining = executable_subset(gates_id_table)
        step, tried = best_realizable_set_found(gates_id_table, executable, arch, mapping,  order_fraction=order_fraction, crit_dict=crit_dict, temperature=temperature, cooling_rate=cooling_rate, termination_temp=termination_temp, initial_order=initial_order, reward_name=reward_name, take_first_ms=take_first_ms)
        tried_steps += tried
        timesteps.append(step)
        routed_ids = [x[0] for x in  step]
        not_executed  = {id : gates_id_table[id] for id in executable if id not in routed_ids}
        gates_id_table = {**not_executed, **remaining}
    #print(f'routing orders tried {tried_steps}')
    return timesteps, tried_steps

def get_depth_by_qubit(gates):
    depth_by_qubit = {}
    for i in gates:
        qubits = gates[i]
        depths = (depth_by_qubit.get(q,  0) for q in qubits)
        max_depth = max(depths)
        for qubit in qubits:
            depth_by_qubit[qubit] = 1 + max_depth
    return depth_by_qubit

def get_depth_by_qubit_p(start_id,gates):
    depth_by_qubit = {}
    touched_qubits = {q for q in gates[start_id]}
    for i in range(start_id, len(gates)):
        qubits = gates[i]
        if len(touched_qubits.intersection(qubits)) > 0:
            touched_qubits.update(qubits)
            depths = (depth_by_qubit.get(q,  0) for q  in qubits)
            max_depth = max(depths)
            for qubit in qubits:
                depth_by_qubit[qubit] = 1 + max_depth
    return depth_by_qubit

def get_dependent_gates(gate_tuple, remaining):
    id, initial_gate = gate_tuple
    dependent = {}
    dependent[id] = initial_gate
    for id,gate in remaining.items():

        if any(depends_on((id,gate), added) for added in dependent.items()):
            dependent[id] = gate

    return dependent

def executable_subset(gates : dict):
    executable = {}
    remainining = {}
    blocked_qubits = set()
    for i,gate in gates.items():
            not_blocked = all([q not in blocked_qubits for q in gate])
            if not_blocked:
                executable[i] = gates[i]
                blocked_qubits.update(set(q for q in gate))
            else:
                remainining[i] = gates[i]
                blocked_qubits.update(set(q for q in gate))
    return executable, remainining



def depends_on(g1, g2):
    return g1[0] > g2[0] and len(set(g1[1]).intersection(g2[1])) > 0

