import streamlit as st
import spacy
from spacy_streamlit import visualize_ner
from support_functions import HealthseaPipe
import operator

def visualize_pipeline():
    healthsea_pipe = HealthseaPipe()

    color_code = {
        "POSITIVE": ("#3C9E58", "#1B7735"),
        "NEGATIVE": ("#FF166A", "#C0094B"),
        "NEUTRAL": ("#7E7E7E", "#4E4747"),
        "ANAMNESIS": ("#E49A55", "#AD6B2D"),
    }

    example_reviews = [
        "This is great for joint pain.",
        "This help joint pain but causes rashes",
        "I'm diagnosed with gastritis. This product helped!",
        "Made my insomnia worse",
        "Didn't help my energy levels",
    ]

    # Functions
    def kpi(n, text):
        html = f"""
        <div class='kpi'>
            <h1>{n}</h1>
            <span>{text}</span>
        </div>
        """
        return html


    def central_text(text):
        html = f"""<h2 class='central_text'>{text}</h2>"""
        return html


    def format_clause(text, meta, pred):
        html = f"""
        <div>
            <div class="clause" style="background-color:{color_code[pred][0]} ; box-shadow: 0px 5px {color_code[pred][1]}; border-color:{color_code[pred][1]};">
                <div class="clause_text">{text}</div>
            </div>
            <div class="clause_meta">
                <div>{meta}</div>
            </div>
        </div>"""
        return html


    def format_effect(text, pred):
        html = f"""
        <div>
            <div class="clause" style="background-color:{color_code[pred][0]} ; box-shadow: 0px 5px {color_code[pred][1]}; border-color:{color_code[pred][1]};">
                <div class="clause_text">{text}</div>
            </div>
        </div>"""
        return html

    load_state = st.markdown ("#### Loading...")
    # Load model
    try:
        load_state.markdown ("#### Loading model...")
        nlp = spacy.load("en_healthsea")

    # Download model
    except LookupError:
        import nltk
        import benepar
        load_state.markdown ("#### Downloading model...")
        benepar.download('benepar_en3')
    load_state.markdown ("#### Loading done!")

    # Pipeline
    st.markdown("""---""")

    st.markdown(central_text("⚙️ Pipeline"), unsafe_allow_html=True)

    check = st.checkbox("Use predefined examples")

    if not check:
        text = st.text_input(label="Write a review", value="This is great for joint pain!")
    else:
        text = st.selectbox("Predefined example reviews", example_reviews)
    doc = nlp(text)

    # NER
    visualize_ner(
        doc,
        labels=nlp.get_pipe("ner").labels,
        show_table=False,
        title="✨ Named Entity Recognition",
        colors={"CONDITION": "#FF4B76", "BENEFIT": "#629B68"},
    )

    st.markdown("""---""")

    # Segmentation, Blinding, Classification
    st.markdown("## 🔮 Segmentation, Blinding, Classification")

    clauses = healthsea_pipe.get_clauses(doc)
    for doc_clause, clause in zip(clauses, doc._.clauses):
        classification = max(clause["cats"].items(), key=operator.itemgetter(1))[0]
        percentage = round(float(clause["cats"][classification]) * 100, 2)
        meta = f"{clause['ent_name']} ({classification} {percentage}%)"

        st.markdown(
            format_clause(doc_clause.text, meta, classification), unsafe_allow_html=True
        )
        st.markdown("\n")

    st.markdown("""---""")

    # Aggregation
    st.markdown("## 🔗 Aggregation")

    for effect in doc._.health_effects:
        st.markdown(
            format_effect(
                f"{doc._.health_effects[effect]['effect']} effect on {effect}",
                doc._.health_effects[effect]["effect"],
            ),
            unsafe_allow_html=True,
        )
        st.markdown("\n")

    st.markdown("""---""")
    # Indepth
    st.markdown("## 🔧 Pipeline attributes")
    clauses_col, effect_col = st.columns(2)

    clauses_col.markdown("### doc._.clauses")
    for clause in doc._.clauses:
        clauses_col.json(clause)
    effect_col.markdown("### doc._.health_effects")
    effect_col.json(doc._.health_effects)