from utils import graph_add_high_breadth, graph_add_high_breadth2, graph_add_high_depth, graph_add_high_depth2, add_chain_of_length, prufer_decode, export_test
from collections import deque
import random
import math


def list_pick_raport(v, l, r):
    return v[random.randint(int(l * len(v)), min(len(v), 1 + int(r * len(v))))]


def simple_tests1(maxn, better, cnt_children, switch_pick_from):
    g = [[] for _ in range(maxn)]
    switch_nodes, final_nodes = [], []

    if better == "dfs":
        g, smallest_unused, switch_nodes = graph_add_high_breadth(
            g,
            smallest_unused = 1,
            tmp_root = 0,
            cnt_new_nodes = maxn//2,
            cnt_children = 3
        )
        switch_node = list_pick_raport(switch_nodes, *switch_pick_from)

        g, smallest_unused, switch_nodes = graph_add_high_depth(
            g,
            smallest_unused = smallest_unused,
            tmp_root = switch_node,
            cnt_new_nodes = maxn - smallest_unused,
            cnt_children = 3
        )
        final_node = list_pick_raport(switch_nodes, *switch_pick_from)
    else:
        g, smallest_unused, switch_nodes = graph_add_high_depth(
            g,
            smallest_unused = 1,
            tmp_root = 0,
            cnt_new_nodes = maxn//2,
            cnt_children = 3
        )
        switch_node = list_pick_raport(switch_nodes, *switch_pick_from)

        g, smallest_unused, switch_nodes = graph_add_high_breadth(
            g,
            smallest_unused = smallest_unused,
            tmp_root = switch_node,
            cnt_new_nodes = maxn - smallest_unused,
            cnt_children = 3
        )
        final_node = list_pick_raport(switch_nodes, *switch_pick_from)

    export_test(g, final_node, message = f"{maxn}_{better}_first_depth_breadth")


def simple_tests2(maxn, better, rap, switch_pick_from):
    g = [[] for _ in range(maxn)]
    switch_nodes, final_nodes = [], []

    if better == "dfs":
        g, smallest_unused, switch_nodes = graph_add_high_breadth2(
            g,
            smallest_unused = 1,
            tmp_root = 0,
            cnt_new_nodes = maxn//2,
            rap_immediate_new_nodes = rap
        )
        switch_node = list_pick_raport(switch_nodes, *switch_pick_from)

        g, smallest_unused, switch_nodes = graph_add_high_depth2(
            g,
            smallest_unused = smallest_unused,
            tmp_root = switch_node,
            cnt_new_nodes = maxn - smallest_unused,
            rap_left_chain = rap
        )
        final_node = list_pick_raport(switch_nodes, *switch_pick_from)
    else:
        g, smallest_unused, switch_nodes = graph_add_high_depth2(
            g,
            smallest_unused = 1,
            tmp_root = 0,
            cnt_new_nodes = maxn//2,
            rap_left_chain = rap
        )
        switch_node = list_pick_raport(switch_nodes, *switch_pick_from)

        g, smallest_unused, switch_nodes = graph_add_high_breadth2(
            g,
            smallest_unused = smallest_unused,
            tmp_root = switch_node,
            cnt_new_nodes = maxn - smallest_unused,
            rap_immediate_new_nodes = rap
        )
        final_node = list_pick_raport(switch_nodes, *switch_pick_from)

    export_test(g, final_node, message = f"{maxn}_{better}_rap{str(rap).replace('.', '')}_spf{''.join([str(y).replace('.', '') for y in switch_pick_from])}_first_depth2_breadth2")


def simple_tests2_combo(maxn, cnt_components, better, rap, rap_expand):
    g = [[] for _ in range(maxn)]
    switch_nodes, final_nodes = [], []
    component_size = (maxn - 1) // cnt_components

    if better == "dfs":
        g, new_smallest_unused, switch_nodes = graph_add_high_breadth2(g, smallest_unused = 1, tmp_root = 0, cnt_new_nodes = component_size, rap_immediate_new_nodes = rap)    
        expand_from = [(i, 'd') for i in range(1, new_smallest_unused)]
    else:
        g, new_smallest_unused, switch_nodes = graph_add_high_depth2(g, smallest_unused = 1, tmp_root = 0, cnt_new_nodes = component_size, rap_left_chain = rap)
        expand_from = [(i, 'b') for i in range(1, new_smallest_unused)]
    smallest_unused = new_smallest_unused

    random.shuffle(expand_from)
    expand_from = expand_from[:int(rap_expand * len(expand_from))]
    qu = deque(expand_from)

    for i in range(1, cnt_components):
        nod, type = qu.popleft()

        if type == 'd':
            g, new_smallest_unused, switch_nodes = graph_add_high_depth2(g, smallest_unused = smallest_unused, tmp_root = nod, cnt_new_nodes = component_size, rap_left_chain = rap)
            expand_from = [(i, 'b') for i in range(smallest_unused, new_smallest_unused)]
        else:
            g, new_smallest_unused, switch_nodes = graph_add_high_breadth2(g, smallest_unused = smallest_unused, tmp_root = nod, cnt_new_nodes = component_size, rap_immediate_new_nodes = rap)
            expand_from = [(i, 'd') for i in range(smallest_unused, new_smallest_unused)]
        smallest_unused = new_smallest_unused

        random.shuffle(expand_from)
        expand_from = expand_from[:int(rap_expand * len(expand_from))]
        qu.extend(expand_from)

    g = g[:smallest_unused]
    final_node = random.randint(len(g)-1 - component_size, len(g))
    export_test(g, final_node, message = f"{len(g)}_{better}_first_combo")


def anti_brute_force_dfs_bfs(maxn):
    g = [[] for _ in range(maxn)]
    num_breadth_layers = int(random.uniform(0.8, 1.2) * maxn/6)

    tmp_root = 0
    smallest_unused = 1
    for _ in range(num_breadth_layers):
        for i in range(smallest_unused, smallest_unused + 3):
            g[tmp_root].append(i)
            g[i].append(tmp_root)
        tmp_root = smallest_unused + 1
        smallest_unused += 3

    g, smallest_unused, switch_nodes = graph_add_high_depth2(g, smallest_unused, tmp_root, maxn - smallest_unused, 0.8)
    final_node = list_pick_raport(switch_nodes, 1/2, 3/4)

    export_test(g, final_node, message = f"{maxn}_anti_brute_force_dfs_bfs")


def anti_brute_force_dfs_bfs_sqrt(maxn):
    g = [[] for _ in range(maxn)]
    smallest_unused = 1

    tmp_root = 0
    for b_new in range(int((1.5 * maxn) ** 0.5), 10, -1):
        g, smallest_unused, switch_nodes = graph_add_high_breadth2(g, smallest_unused, tmp_root, b_new, (b_new - 2) / b_new)
        tmp_root = switch_nodes[-1]

    g, smallest_unused, switch_nodes = graph_add_high_depth2(g, smallest_unused, tmp_root, maxn - smallest_unused, 0.8)
    final_node = list_pick_raport(switch_nodes, 1/2, 3/4)

    export_test(g, final_node, message = f"{maxn}_anti_brute_force_dfs_bfs_more_sqrt")


def anti_brute_force_dfs_bfs_increasing_subtree_sizes(maxn):
    g = [[] for _ in range(maxn)]
    smallest_unused = 1

    lo, hi = map(int, [maxn ** 0.5, (3 * maxn) ** 0.5 - 3])
    swap_at = int(random.uniform(0.8, 1.2) * (lo + hi) / 2)

    tmp_root = 0
    for l in range(hi, lo-1, -1):
        for _ in range(3):
            g[tmp_root].append(smallest_unused)
            g[smallest_unused].append(tmp_root)
            smallest_unused += 1
        
        g, smallest_unused = add_chain_of_length(g, g[tmp_root][-1] if l >= swap_at else g[tmp_root][-3], l, smallest_unused)
        tmp_root = g[tmp_root][-2]

    g = g[:smallest_unused]
    final_node = tmp_root
    
    export_test(g, final_node, message = f"{len(g)}_anti_brute_force_dfs_bfs_increasing_subtree_sizes_same_child")


def anti_brute_force_dfs_bfs_increasing_subtree_sizes_random_sides(maxn, start_with):
    g = [[] for _ in range(maxn)]
    smallest_unused = 1

    lo, hi = map(int, [maxn ** 0.5, (3 * maxn) ** 0.5 - 3])
    right_proba, decrease_by = (1 if start_with == "dfs" else 0), (1/(hi-lo+1) if start_with == "dfs" else -1/(hi-lo+1))

    tmp_root = 0
    for l in range(hi, lo-1, -1):
        for _ in range(3):
            g[tmp_root].append(smallest_unused)
            g[smallest_unused].append(tmp_root)
            smallest_unused += 1
        
        g, smallest_unused = add_chain_of_length(g, g[tmp_root][-1] if random.random() < right_proba else g[tmp_root][-3], l, smallest_unused)
        tmp_root = g[tmp_root][-2]
        right_proba -= decrease_by

    g = g[:smallest_unused]
    final_node = tmp_root
    
    export_test(g, final_node, message = f"{len(g)}_anti_brute_force_{'dfs_bfs' if start_with == 'dfs' else 'bfs_dfs'}_increasing_subtree_sizes_same_child_random")


def not_anti_nsqrtn(maxn):
    g = [[] for _ in range(maxn)]
    smallest_unused = 1

    a = int((maxn-1) / 4 * random.uniform(0.8, 1.2))
    b = (maxn-1 - 2*a) // 3

    tmp_root = 0
    for chain_ind in range(a):
        for _ in range(2):
            g[tmp_root].append(smallest_unused)
            g[smallest_unused].append(tmp_root)
            smallest_unused += 1
        if chain_ind+1 < a:
            tmp_root = g[tmp_root][-2]

    g, smallest_unused = add_chain_of_length(g, g[tmp_root][-2], b+1, smallest_unused)

    tmp_root = g[tmp_root][-1]
    for chain_ind in range(b):
        for _ in range(2):
            g[tmp_root].append(smallest_unused)
            g[smallest_unused].append(tmp_root)
            smallest_unused += 1
        tmp_root = g[tmp_root][-1]
    
    g = g[:smallest_unused]
    final_node = tmp_root
    
    export_test(g, final_node, message = f"{len(g)}_not_anti_sqrtn")


def random_tree(maxn):
    code = [random.randint(0, maxn-1) for _ in range(maxn - 2)]
    edges = prufer_decode(code)

    g = [[] for _ in range(maxn)]
    for x, y in edges:
        g[x].append(y)
        g[y].append(x)

    depth, stk = [0 for _ in range(maxn)], [0]
    max_depth, final_node = 0, 0
    while stk:
        nod = stk.pop()
        for nn in g[nod]:
            if depth[nn] == 0 and nn != 0:
                depth[nn] = 1 + depth[nod]
                if depth[nn] > max_depth:
                    max_depth = depth[nn]
                    final_node = nn
                stk.append(nn)

    print(f"{max_depth = }")

    export_test(g, final_node, message = f"{maxn}_random_tree")


def main():
    maxn = 10 ** 5
    better = ["bfs", "dfs"][0] # cu ce e mai bine sa incep.
    cnt_children = [2, 3][1]
    switch_pick_from = [(1/8, 1/5), (1/4, 1/2)][0] # cu orice as incepe, am un ultim strat. aleg din el in ce nod vreau sa schimb.
    rap = [0.6, 0.7, 0.75, 0.8][3]

    # simple_tests1(maxn, better, cnt_children, switch_pick_from)
    # simple_tests2(maxn, better, rap, switch_pick_from)
    # simple_tests2_combo(maxn, int(maxn ** 0.5), better, 0.8, 0.1)
    # anti_brute_force_dfs_bfs(maxn)
    # anti_brute_force_dfs_bfs_sqrt(maxn)
    # anti_brute_force_dfs_bfs_increasing_subtree_sizes(maxn)
    # anti_brute_force_dfs_bfs_increasing_subtree_sizes_random_sides(maxn, start_with = "bfs")
    # anti_brute_force_dfs_bfs_same_child_count_same_subtree_increase_ct(maxn, 10**3)
    not_anti_nsqrtn(maxn)
    # random_tree(maxn)


if __name__ == "__main__":
    main()