from pathlib import Path
import string
import networkx as nx
import numpy as np
import itertools
import ast
from tqdm import tqdm
import os
import io

from dbca.datasets.pcfg_set.set_sample import SETSample
from dbca.sample_set import SampleSet
from dbca.storage import SampleStore

UNARY = ['copy', 'reverse', 'shift', 'echo', 'swap_first_last', 'repeat']
BINARY = ['append', 'prepend', 'remove_first', 'remove_second']
COUNTER = {'copy': 0, 'reverse': 0, 'shift': 0, 'echo': 0, 'swap_first_last': 0, 'repeat': 0, 'append': 0, 'prepend': 0, 'remove_first': 0, 'remove_second': 0, 'LEAF': 0}

# Replace any sequence vars in a given SET with 'LEAF', e.g., 'append A1 B2 copy C3' --> 'append LEAF copy LEAF'
def replace_vars(tokens):
    new_tokens = []
    prev = None
    for token in tokens:
        if token in COUNTER:
            new_tokens.append(token)
            prev = token
        elif token != ',' and prev != 'LEAF':
            new_tokens.append('LEAF')
            prev = 'LEAF'
        elif token == ',':
            prev = ','
    return new_tokens

# Given a list of tokens in SET, convert to a list of NetworkX graphs.
def string_to_tree(token_list, raw):
    samples = []
    for i in range(len(token_list)):
        tree = nx.DiGraph()
        counter = dict(COUNTER)
        _, idx = _string_to_tree(token_list[i], counter, tree)
        assert idx == len(token_list[i])
        samples.append(SETSample(tree, raw[i]))
    return samples


def _string_to_tree(tokens, counter, tree, idx=0):
    """
    Performs a preorder traversal of tokens to build a tree.

    Parameters
    ----------
    tokens
        A list of the tokens in the tree formed from preorder traversal 
    tree
        A networkX DiGraph object for storing the tree
    
    Returns
    -------
    root
        The string of the current node
    idx
        The idx of the next token.
    """
    token = tokens[idx]
    idx += 1
    if token in COUNTER:
        counter[token] += 1
        num = f'_{counter[token]}' if counter[token] > 9 else f'_0{counter[token]}'
        root = token + num
        tree.add_node(root)

        if token != 'LEAF':
            left, idx = _string_to_tree(tokens, counter, tree, idx=idx)         
            tree.add_edge(root, left)

        if token in BINARY:
            right, idx = _string_to_tree(tokens, counter, tree, idx=idx)
            tree.add_edge(root, right)

    else:
        print(tokens)
        print(idx)
        print(token)
        assert False

    return root, idx

def build_SET_samples(path, logger):
    assert os.path.isfile(path)
    logger.info(f"Loading data from {path} ...")
    with io.open(path, mode='r', encoding='utf-8') as f:
        lines = [line.rstrip().replace('( ', '').replace(' )', '') for line in f]
    data = [xy.split('\t') for xy in lines]
    data = [xy[0].split() for xy in data if len(xy) == 2]
    data = [replace_vars(x) for x in data]
    logger.info(f"Loaded {len(data)} strings from the disk.")
    return string_to_tree(data, lines)