from st_aggrid import AgGrid, JsCode, ColumnsAutoSizeMode

import streamlit as st

HEADER_STYLE = {'fontSize': '18px'}
CELL_STYLE = {'fontSize': '18px'}

# Add theme color and grid styles
st.markdown("""
    <style>
        :root {
            --theme-color: rgb(129, 150, 64);
            --theme-color-light: rgba(129, 150, 64, 0.2);
        }
        
        /* AG Grid specific overrides */
        .ag-theme-alpine {
            --ag-selected-row-background-color: var(--theme-color-light) !important;
            --ag-row-hover-color: var(--theme-color-light) !important;
            --ag-selected-tab-color: var(--theme-color) !important;
            --ag-range-selection-border-color: var(--theme-color) !important;
            --ag-range-selection-background-color: var(--theme-color-light) !important;
        }
        
        .ag-row-hover {
            background-color: var(--theme-color-light) !important;
        }
        
        .ag-row-selected {
            background-color: var(--theme-color-light) !important;
        }
        
        .ag-row-focus {
            background-color: var(--theme-color-light) !important;
        }
        
        .ag-cell-focus {
            border-color: var(--theme-color) !important;
        }
        
        /* Keep existing styles */
        .center-text {
            text-align: center;
            color: var(--theme-color);
        }
        .center-image {
            display: block;
            margin-left: auto;
            margin-right: auto;
        }
        h2 {
            color: var(--theme-color) !important;
        }
        .ag-header-cell {
            background-color: var(--theme-color) !important;
            color: white !important;
        }
        a {
            color: var(--theme-color) !important;
        }
        a:hover {
            color: rgba(129, 150, 64, 0.8) !important;
        }
    </style>
""", unsafe_allow_html=True)

#  logo
# st.markdown('<img src="https://www.voyageai.com/logo.svg" class="center-image" width="200">', unsafe_allow_html=True)

# title
st.markdown('<h2 class="center-text">Embedding Benchmark For Retrieval</h2>', unsafe_allow_html=True)

group_name = "$group_name$"

data_engine = st.session_state["data_engine"]

df = data_engine.jsons_to_df()[:]

df = df[df["group_name"] == group_name].sort_values(by="ndcg_at_10", ascending=False)

# setting column config
grid_options = {
    'columnDefs': [
        {
            'headerName': 'Model Name',
            'field': 'model_name',
            'pinned': 'left',
            'sortable': False,
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
            'cellRenderer': JsCode("""class CustomHTML {
                init(params) {
                    const link = params.data.reference;
                    this.eGui = document.createElement('div');
                    this.eGui.innerHTML = link ? 
                        `<a href="${link}" target="_blank">${params.value}</a>` : 
                        params.value;
                }

                getGui() {
                    return this.eGui;
                }
            }"""),
        },
        {
            'headerName': 'NDCG@10',
            'field': 'ndcg_at_10',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
        },
        {
            'headerName': 'Data Type',
            'field': 'embd_dtype',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
        },
        {
            'headerName': 'Embd Dim',
            'field': 'embd_dim',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
        },
        {
            'headerName': 'Model Size (# of Parameters)',
            'field': 'num_params',
            'cellDataType': 'number',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
        },
        {
            'headerName': 'Context Length',
            'field': 'max_tokens',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
        },
        {
            'headerName': 'Query Instruction',
            'field': 'query_instruct',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
            'suppressSizeToFit': True,

        },
        {
            'headerName': 'Corpus Instruction',
            'field': 'corpus_instruct',
            'headerStyle': HEADER_STYLE,
            'cellStyle': CELL_STYLE,
            'suppressSizeToFit': True,

        },
    ],
    'defaultColDef': {
        'filter': True,
        'sortable': True,
        'resizable': True
    },
    'autoSizeStrategy': {
        'type': 'fitCellContents'
    }
}

ag = AgGrid(
    df,
    enable_enterprise_modules=False,
    gridOptions=grid_options,
    allow_unsafe_jscode=True,
    columns_auto_size_mode=ColumnsAutoSizeMode.FIT_CONTENTS,
    theme="streamlit",
)