import os
import streamlit as st
import json
import tarfile
from base64 import b64encode

st.set_page_config(layout="wide")

PARENT_DIR: str = os.path.join(os.path.dirname(os.path.abspath(__file__)))
EVAL_DIR: str = os.path.join(PARENT_DIR, "eval-results")
st.title("K2 Evaluation Gallery")
st.markdown("""The K2 gallery allows one to browse the output of various evaluations on intermediate K2 checkpoints, which provides an intuitive understanding on how the model develops and improves over time.""")


with st.sidebar:
    with open(os.path.join(PARENT_DIR, "k2-logo.svg"), 'r') as f:
        b64 = b64encode(f.read().encode('utf-8')).decode("utf-8")
        html = f"<img src='https://www.llm360.ai/images/logo-highres.png' width='100' /><img src='data:image/svg+xml;base64,{b64}' width='100' />"
        st.markdown(html, unsafe_allow_html=True)

    metric = st.radio(
        "Choose a metric", options=os.listdir(os.path.join(EVAL_DIR))
    )

    n_shot = st.radio(
        "Selece an n-shot number", os.listdir(os.path.join(EVAL_DIR, metric))
    )

col1, col2 = st.columns(2)

with col1:
    st.header("Checkpoint A")
    ckpt = st.selectbox('Select a checkpoint', sorted(os.listdir(os.path.join(EVAL_DIR, metric, n_shot))), key="A1")
    st.write(f'Veiwing Evaluation Results for Checkpoint: `{ckpt}`')
    file = st.selectbox("Select a file", sorted(os.listdir(os.path.join(EVAL_DIR, metric, n_shot, ckpt))), key="A2")
    with tarfile.open(os.path.join(EVAL_DIR, metric, n_shot, ckpt, file), "r:gz") as f:
        st.json(json.load(f.extractfile(f.next())))

with col2:
    st.header("Checkpoint B")
    ckpt = st.selectbox('Select a checkpoint', sorted(os.listdir(os.path.join(EVAL_DIR, metric, n_shot))), key="B1")
    st.write(f'Veiwing Evaluation Results for Checkpoint: `{ckpt}`')
    file = st.selectbox("Select a file", sorted(os.listdir(os.path.join(EVAL_DIR, metric, n_shot, ckpt))), key="B2")
    with tarfile.open(os.path.join(EVAL_DIR, metric, n_shot, ckpt, file), "r:gz") as f:
        st.json(json.load(f.extractfile(f.next())))