import pandas as pd
import streamlit as st
import re
# Define color maps for both light and dark modes
COLOR_MAP = {
    "light": {
        "yellow": "background-color: rgba(255, 255, 128, 0.3)",  # Reasoning models
        "green": "background-color: rgba(192, 255, 192, 0.3)",   # Linear attention hybrid
        "blue": "background-color: rgba(192, 192, 255, 0.3)"     # SSM hybrid models
    },
}

def get_color_map():
    """Returns the appropriate color map based on the current theme."""
    return COLOR_MAP["light"]

def style_zero_context(df):
    """
    Similar approach to style_long_context:
    1) color rows based on model name
    2) numeric formatting
    """
    import pandas as pd

    # Example color dict, tweak as needed:
    color_mapping = {
        "minimax-text-01": get_color_map()["green"],
        "jamba-1.5-large": get_color_map()["blue"],
        "deepseek-r1": get_color_map()["yellow"],
        "o1-mini": get_color_map()["yellow"],
        "o3-mini": get_color_map()["yellow"],
        "qwq-32b-preview": get_color_map()["yellow"],
        # Add any other special-cased models here
        # "o1-mini": COLOR_MAP["yellow"], etc.
    }
    # Add links to model names
    df["Model"] = df.apply(lambda row: f'<a href="{row["Link"]}" target="_blank">{row["Model"]}</a>', axis=1)
    df.drop(columns=["Link"], inplace=True)

    styler = df.style.apply(
        lambda row: [color_mapping.get(re.sub(r'<[^>]+>', '', row["Model"]), "")]*len(row),
        axis=1
    )
    
    # # Attach custom tooltips (optional)
    # tooltips = pd.DataFrame("", index=df.index, columns=df.columns)
    # if "1st<50% op" in df.columns:
    #     tooltips["1st<50% op"] = "First operation number with accuracy <50%"
    # if "1st<10% op" in df.columns:
    #     tooltips["1st<10% op"] = "First operation number with accuracy <10%"
    # if "Avg. Acc op≤30" in df.columns:
    #     tooltips["Avg. Acc op≤30"] = "Average accuracy of first 30 operations"
    # styler = styler.set_tooltips(tooltips)

    # Apply numeric formatting
    styler = styler.format({
        "Symbolic": "{:,.2f}",       # Format as number with thousands separator and 1 decimal place
        "Medium": "{:,.2f}",        # Format as number with thousands separator and 2 decimal places
        "Hard": "{:,.2f}",          # Format as number with thousands separator and 2 decimal places
        "1st<50% op": "{:,.0f}",      # Format as plain integer (no decimal places)
        "1st<10% op": "{:,.0f}",      # Format as plain integer (no decimal places)
        "Avg. Acc op≤30": "{:.4f}",  # Format with 4 decimal places
        "Average↑": "{:,.2f}"      # Format as number with thousands separator and 2 decimal places
    })

    html = styler.to_html(escape=False, index=False)

    # Updated regex: target model name *before* link replacement and use word boundary
    for model, style in color_mapping.items():
        html = re.sub(rf'<tr[^>]*>\s*<td[^>]*>{re.escape(model)}<', rf'<tr style="{style}"><td>', html, re.M)

    html = re.sub(
        r'<table(.*?)>', 
        r'<table\1 style="width:100%; border-collapse:collapse;">', 
        html
    )
    return html # Return the modified HTML

    
    return styler
# Add styling for model types
def style_long_context(df):
    color_mapping = {
        "minimax-text-01": get_color_map()["green"],
        "jamba-1.5-large": get_color_map()["blue"]
    }

    df["Model"] = df.apply(lambda row: f'<a href="{row["Link"]}" target="_blank">{row["Model"]}</a>', axis=1)
    df.drop(columns=["Link"], inplace=True)
    
    styled_df = df.style.apply(
        lambda row: [color_mapping.get(re.sub(r'<[^>]+>', '', row["Model"]), "")]*len(row),
        axis=1
    ).format({
        "8K": "{:,.2f}",
        "16K": "{:,.2f}", 
        "32K": "{:,.2f}",
        "Average↑": "{:,.2f}",
    })

    # Convert to HTML and add <a> tags
    html = styled_df.to_html(escape=False, index=False)

    # Updated regex: target model name *before* link replacement and use word boundary
    for model, style in color_mapping.items():
        html = re.sub(rf'<tr[^>]*>\s*<td[^>]*>{re.escape(model)}<', rf'<tr style="{style}"><td>', html, re.M)

    html = re.sub(
        r'<table(.*?)>', 
        r'<table\1 style="width:100%; border-collapse:collapse;">', 
        html
    )

    return html # Return the modified HTML