import pandas as pd
import streamlit as st


@st.cache_data
def load_comparisons(comparison_file, billtext_file, modelseg_file):
    df = pd.read_json(comparison_file)
    
    most_recent_ver = (
        df.groupby('target__bill_id')
        ['target__doc_rank'].max()
        .reset_index()
        .assign(target__is_most_recent_doc=True)
    )
    
    df = (
        df.merge(most_recent_ver, how='left')
        .fillna({'target__is_most_recent_doc':False})
        .astype({'target__is_most_recent_doc':'bool'})
    )
    
    text_df = pd.read_json(billtext_file)
    seg_df = pd.read_json(modelseg_file)
    
    return df, text_df, seg_df


def render(comparison_file, billtext_file, modelseg_file):
    if st.button('Load data'):
        (
            st.session_state.df, 
            st.session_state.text_df, 
            st.session_state.seg_df
        ) = load_comparisons(comparison_file, billtext_file, modelseg_file)
    if 'df' not in st.session_state:
        (
            st.session_state.df, 
            st.session_state.text_df, 
            st.session_state.seg_df
        ) = load_comparisons(comparison_file, billtext_file, modelseg_file)
        
    df = st.session_state.df
    text_df = st.session_state.text_df
    seg_df = st.session_state.seg_df

    return df, text_df, seg_df