#!python
import argparse
from pymeiji import meiji_solver
import networkx as nx

parser = argparse.ArgumentParser(description="produces the tree decomposition, i.e. the dot td file")

parser.add_argument('pkcm_file', type=str, help='.pkcm file, describing a pseudoknotted covarance model.')

args = parser.parse_args()

# BUILING HELIX GRAPH
partner = {}
half_helices = []

for line in open(args.pkcm_file).readlines():
    if line.startswith("Helix"):
        name = line.split(' ')[0]
        i = int(line.split(' ')[-4])-1
        j = int(line.split(' ')[-3])-1
        k = int(line.split(' ')[-2])-1
        l = int(line.split(' ')[-1].rstrip('\n'))-1

        half_helices.append((i,j))
        half_helices.append((k,l))

        partner[(k,l)] = (i,j)

    if line.startswith("## END HELIX ARRANGEMENT DESCRIPTION"):
        break

half_helices = sorted(half_helices, key=lambda x : x[0])

cnt1 = 0
cnt2 = 1

import copy

helix_graph = nx.Graph()
pos = {}

for k in range(len(half_helices)):

    helix_graph.add_edge(cnt1, cnt2)

    i,j = half_helices[k]
    pos[(i,j)] = (copy.copy(cnt1), copy.copy(cnt2))

    if (i,j) in partner.keys():
        l,m = partner[(i,j)]
        p1,p2 = pos[(l,m)]
        helix_graph.add_edge(cnt1, p1)
        helix_graph.add_edge(cnt2, p2)
        helix_graph.add_edge(cnt2, p1)
        helix_graph.add_edge(cnt1, p2)
    

    cnt1 += 1
    cnt2 += 1
# HELIX GRAPH BUILT

tw, bags, td_edges = meiji_solver(helix_graph.edges)

print(tw)
print(bags)
print(td_edges)

for u, content in bags.items():
    for k in range(len(content)-1):
        assert(content[k]<content[k+1])

# MAKING IT NICE
queue = [(-1,min(bags.keys()))]
root = min(bags.keys())
N = max(bags.keys())
while len(queue) > 0:
    p,u = queue.pop()
    print(p,u)

    # here, if there is in bag[u]
    # more than one element which is not
    # in bag[p], introduce a new bag
    # x such that bag[x] = bag[p]\bag[u] + one extra
    # element in bag[u] but not in bag[p].
    # in td_edges[u], p should be replaced by x
    # and in td_edges[p], u should be replaced by x/
    # (the bag bags[x] is inserted between bags p and u)

    set_u = set(bags[u])
    if p!=-1:
        set_p = set(bags[p])
    else:
        set_p = set([])

    if len(set_u - set_p) > 1:
        elem = (set_u-set_p).pop()
        new_bag = [e for e in bags[u] if e!=elem]
        bags[N+1] = new_bag

        td_edges[u] = [ngbh for ngbh in td_edges[u] if ngbh != p]
        td_edges[u].append(N+1)

        if p!=-1:
            td_edges[p] = [ngbh for ngbh in td_edges[p] if ngbh != u]
            td_edges[p].append(N+1)

            td_edges[N+1] = [p,u]
        else:
            td_edges[N+1] = [u]
            root = N+1

        queue.append((p,N+1))

        N+= 1

        continue
        


    for v in td_edges[u]:
        if v!=p:
            queue.append((u,v))

f = open(args.pkcm_file+'.td','w')

print('root',root,file=f)

for k, bag in bags.items():
    print('bag',k,[u+1 for u in bag],file=f)

for k, ngbh in td_edges.items():
    print("bag_adj",k,ngbh,file=f)
