from collections import Counter

import graphviz
from optimum.bettertransformer import BetterTransformer
import penman
from penman.models.noop import NoOpModel
from mbart_amr.constraints.constraints import AMRLogitsProcessor
from mbart_amr.data.linearization import linearized2penmanstr
from mbart_amr.data.tokenization import AMRMBartTokenizer
from transformers import MBartForConditionalGeneration, LogitsProcessorList

import streamlit as st

if "logits_processor" not in st.session_state:
    st.session_state["logits_processor"] = None

if "tokenizer" not in st.session_state:
    st.session_state["tokenizer"] = None

if "model" not in st.session_state:
    st.session_state["tokenizer"] = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-to-amr", src_lang="en_XX")
    st.session_state["model"] = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-to-amr")
    st.session_state["model"] = BetterTransformer.transform(st.session_state["model"], keep_original_model=False)
    st.session_state["model"].resize_token_embeddings(len(st.session_state["tokenizer"]))
    st.session_state["logits_processor"] = AMRLogitsProcessor(st.session_state["tokenizer"],
                                                              st.session_state["model"].config.max_length)

st.title("📝 Parse text into AMR")

text = st.text_input(label="Text to transform (en)")

if text and "model" in st.session_state:
    gen_kwargs = {
        "max_length": st.session_state["model"].config.max_length,
        "num_beams": st.session_state["model"].config.num_beams,
        "logits_processor": LogitsProcessorList([st.session_state["logits_processor"]]) if st.session_state[
            "logits_processor"] else None
    }

    encoded = st.session_state["tokenizer"](text, return_tensors="pt")
    generated = st.session_state["model"].generate(**encoded, **gen_kwargs)
    linearized = st.session_state["tokenizer"].decode_and_fix(generated)[0]
    penman_str = linearized2penmanstr(linearized)

    try:
        graph = penman.decode(penman_str, model=NoOpModel())
    except Exception as exc:
        st.write(f"The generated graph is not valid so it cannot be visualized correctly. Below is the closest attempt"
                 f" to a valid graph but note that this is invalid Penman.")
        st.code(penman_str)

        with st.expander("Error trace"):
            st.write(exc)
    else:
        visualized = graphviz.Digraph(node_attr={"color": "#3aafa9", "style": "rounded,filled", "shape": "box",
                                                 "fontcolor": "white"})

        # Count which names occur multiple times, e.g. t/talk-01 t2/talk-01
        nodename_c = Counter([item[2] for item in graph.triples if item[1] == ":instance"])
        # Generated initial nodenames for each variable, e.g. {"t": "talk-01",  "t2": "talk-01"}
        nodenames = {item[0]: item[2] for item in graph.triples if item[1] == ":instance"}

        # Modify nodenames, so that the values are unique, e.g. {"t": "talk-01 (1)",  "t2": "talk-01 (2)"}
        # but only the value occurs more than once
        nodename_str_c = Counter()
        for varname in nodenames:
            nodename = nodenames[varname]
            if nodename_c[nodename] > 1:
                nodename_str_c[nodename] += 1
                nodenames[varname] = f"{nodename} ({nodename_str_c[nodename]})"

        def get_node_name(item: str):
            return nodenames[item] if item in nodenames else item

        try:
            for triple in graph.triples:
                if triple[1] == ":instance":
                    continue
                else:
                    visualized.edge(get_node_name(triple[0]), get_node_name(triple[2]), label=triple[1])
        except Exception as exc:
            st.write("The generated graph is not valid so it cannot be visualized correctly. Below is the closest attempt"
                     " to a valid graph but note that this is probably invalid Penman.")
            st.code(penman_str)
            st.write("The initial linearized output of the model was:")
            st.code(linearized)

            with st.expander("Error trace"):
                st.write(exc)
        else:
            st.subheader("Graph visualization")
            st.graphviz_chart(visualized, use_container_width=True)

            # Download
            img = visualized.pipe(format="png")
            st.download_button("Download graph", img, mime="image/png")

            # Additional info
            st.subheader("Model output and Penman graph")
            st.write("The linearized output of the model (after some post-processing) is:")
            st.code(linearized)
            st.write("When converted into Penman, it looks like this:")
            st.code(penman.encode(graph))


########################
# Information, socials #
########################
st.markdown("## Contact ✒️")

st.markdown("Would you like  additional functionality in the demo? Or just want to get in touch?"
            " Give me a shout on [Twitter](https://twitter.com/BramVanroy)"
            " or add me on [LinkedIn](https://www.linkedin.com/in/bramvanroy/)!")