import streamlit as st
from utils import QNLP
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

st.write("### QNLP demo")
lang = st.selectbox("Select Language", ("EN", "ZH"))
input = st.text_input("Text Input")

def plot_data(array:np.ndarray):
    fig, ax = plt.subplots()

    ax.set_xlabel("Value")
    ax.set_ylabel("% Total")
    value, count = np.unique(array,return_counts=True)
    count = count * 100 / len(array)
    ax.set_xlim([-5,260])
    ax.bar(value, count, edgecolor="blue", align="edge")

    return fig

if input.strip():
    results = QNLP(lang.lower()).process_sentence(input)
    subtabs = []
    if len(results) > 1:
        subtabs = [f"Sentence {n+1}" for n in range(len(results))]

    tabs = st.tabs(["Overall"] + subtabs)

    overall_tab = tabs[0]
    detail_tabs = tabs[1:]

    with overall_tab:
        st.header("Overall")
        
        whole_array = [np.sum(res.array, axis=-1) for res in results if res.job.done()]
        if len(whole_array):                
            whole_array = np.concatenate(whole_array)

            total = len(whole_array)
            value, count = np.unique(whole_array,return_counts=True)
            count = count/total*100

            df = pd.DataFrame({    
                "Bit Value" : value, 
                "Percentage" : count,
            })

            col1, col2 = st.columns([1,2])
            with col1:
                st.dataframe(df, hide_index=True)
            with col2:
                st.pyplot(plot_data(whole_array))
        else:
            st.write("Input sentences are too big to be processed, try breaking them into smaller sentences")

    for idx, (tab, result) in enumerate(zip(detail_tabs, results)):
        with tab:
            st.header(f"Sentence {idx}")
            st.write(' '.join(result.tokens))
            
            if result.valid:
                value, count = np.unique(result.array,return_counts=True)
                count = count/total*100

                df = pd.DataFrame({    
                    "Bit Value" : value, 
                    "Percentage" : count,
                })

                col1, col2 = st.columns([1,2])
                with col1:
                    st.dataframe(df, hide_index=True)
                with col2:
                    st.pyplot(plot_data(result.array))
            else:
                st.write(f"Sentence Discarded due to lack of qubits to process")
else:
    st.write("Choose a langauge and input some sentences to start !")