import streamlit as st
import pandas as pd

covariate_columns = {
    'content_domain': 'Content Domain',
    'language': 'Language',
    'rater_group': 'Rater Group',
}

id_vars = [
    'mean_z', 'text', 'content_domain', 'language', 
    'rater_group', 'study', 'instrument'
]

def load_data():
    st.session_state.df = (
        pd
        .read_feather(path='data.feather')
        .query('partition == "test" | partition == "dev"')
        .melt(
            value_vars=['sentiment_model', 'desirability_model'],
            var_name='x_group',
            value_name='x',
            id_vars=id_vars
            )
        .replace(
            to_replace={
                'en': 'English',
                'de': 'German',
                'other': 'Other',
                'personality': 'Personality',
                'laypeople': 'Laypeople',
                'students': 'Students',
                'sentiment_model': 'Sentiment Model',
                'desirability_model': 'Desirability Model'
            }    
        )
        .rename(columns=covariate_columns)
        .rename(
            columns={
                'mean_z': 'Human-ratings',
                'x': 'Machine-ratings',
            }
        )
    )