import gradio as gr
import stanza



import torch
import numpy
# Allowlist numpy.ndarray for weights-only loading
torch.serialization.add_safe_globals([numpy.ndarray])
# Allowlist the required global for weights-only loading
# torch.serialization.add_safe_globals([numpy.core.multiarray._reconstruct])

# Download and initialize the Stanza pipeline
stanza.download('en')
nlp = stanza.Pipeline(lang='en', processors='tokenize,pos,constituency')


def generate_reordering_rule(english, reordered):
    tree = nlp(english).sentences[0].constituency
    reordered_tokens = reordered.split()
    rules = []

    def extract(node):
        if not hasattr(node, 'children') or all(isinstance(c, str) for c in node.children):
            return

        child_labels = tuple(c.label if hasattr(c, 'label') else c for c in node.children)
        rule_key = (node.label, child_labels)

        # Get leaf tokens for each child
        child_tokens = []
        for c in node.children:
            if hasattr(c, 'leaf_labels'):
                child_tokens.append(' '.join(c.leaf_labels()))
            else:
                child_tokens.append(c)

        # Try to find token positions in reordered sentence
        positions = []
        for tok in child_tokens:
            first_word = tok.split()[0]
            try:
                pos = reordered_tokens.index(first_word)
            except ValueError:
                pos = -1
            positions.append(pos)

        # Infer reordering function
        reordered_indices = sorted(range(len(child_tokens)), key=lambda i: positions[i])
        rule_func = f"lambda {', '.join(f'c{i}' for i in range(len(child_tokens)))}: " \
                    f"{' + '.join(f'c{i}' for i in reordered_indices)}"

        rules.append((rule_key, rule_func))

        for c in node.children:
            if hasattr(c, 'children'):
                extract(c)

    extract(tree)

    # Format rules for display
    if not rules:
        return "No reordering rules could be inferred."
    return "\n".join([f"{key}: {func}" for key, func in rules])

# Gradio interface
demo = gr.Interface(
    fn=generate_reordering_rule,
    inputs=[
        gr.Textbox(lines=2, label="English Sentence", placeholder="e.g. I want to eat the cake."),
        gr.Textbox(lines=2, label="Reordered Sentence", placeholder="e.g. I cake the eat to want.")
    ],
    outputs="text",
    title="Reordering Rule Generator",
    description="Enter an English sentence and its reordered version. This app extracts syntactic transformation rules using Stanza's constituency parser."
)

if __name__ == "__main__":
    demo.launch()