#%%
import json
from collections import Counter
import pandas as pd
from IPython.display import display
from functools import partial

def load_counts(path):
    with open(path, 'r') as f:
        return json.load(f)

# Load the counts
counts_a = load_counts('cache/cultura-dotnum.json')
counts_b = load_counts('cache/lmsys-dotnum.json')

# --- Generalized Tokenization Functions (Robust to `first_char=None`) ---

def count_tokens_pure_single_digit(first_char, digits_part):
    """Baseline: Each digit is a token. A prefix, if it exists, is also a token."""
    return (1 if first_char is not None else 0) + len(digits_part)

def count_tokens_3_digit_l2r_r2l(first_char, digits_part):
    """3-digit chunking. A prefix, if it exists, is a separate token."""
    return (1 if first_char is not None else 0) + (len(digits_part) + 2) // 3

def count_tokens_prefix_aware_single_digit(first_char, digits_part, mergeable_chars):
    """Single-digit tokenization where a prefix might merge."""
    prefix_cost = 0
    if first_char is not None and first_char not in mergeable_chars:
        prefix_cost = 1
    return prefix_cost + len(digits_part)

def count_tokens_hybrid_no_merge(first_char, digits_part):
    """Hybrid Variant 1: 1-2 digits atomic (no merge); 3+ digits split (but first merges with prefix)."""
    return 1 if first_char is None else 2 if len(digits_part) <= 2 else len(digits_part)

def count_tokens_hybrid_space_merge(first_char, digits_part):
    """Hybrid Variant 2: 1-2 digits atomic (merges with space only); 3+ digits split."""
    return 1 if first_char is None or first_char == ' ' else 2 if len(digits_part) <= 2 else len(digits_part)

def count_tokens_hybrid_space_punct_merge(first_char, digits_part, mergeable_chars):
    """Hybrid Variant 3: 1-2 digits atomic (merges with special chars); 3+ digits split."""
    return 1 if first_char is None or first_char in mergeable_chars else 2 if len(digits_part) <= 2 else len(digits_part)

# --- Main Calculation Function with Restored Individual Analysis ---
def calculate_compressibility(counts, label):
    """Tests the full suite of strategies including individual character merges."""
    
    special_chars_to_test = [' ', '-', '(', ',', '.']
    
    strategies = {
        # Baselines and Standards
        "Pure Single-Digit (Baseline)": count_tokens_pure_single_digit,
        "3-Digit L2R/R2L": count_tokens_3_digit_l2r_r2l,
        
        # Refined Hybrid Strategies
        "Hybrid (1-2 Digits Atomic, No Merge)": count_tokens_hybrid_no_merge,
        "Hybrid (1-2 Digits Atomic, Space-Merge)": count_tokens_hybrid_space_merge,
        "Hybrid (1-2 Digits Atomic, Space+Punct Merge)": partial(count_tokens_hybrid_space_punct_merge, mergeable_chars=set(special_chars_to_test)),
    }
    
    # --- RESTORED: INDIVIDUAL SINGLE-DIGIT MERGE ANALYSIS ---
    for char in special_chars_to_test:
        strategy_name = f"Single-Digit (Merge '{repr(char)}' Only)"
        strategies[strategy_name] = partial(count_tokens_prefix_aware_single_digit, mergeable_chars={char})
    
    # Add the aggregate "merge all" strategy for single-digit for a summary view
    strategies["Single-Digit (Merge ALL Special)"] = partial(count_tokens_prefix_aware_single_digit, mergeable_chars=set(special_chars_to_test))
    # --- END RESTORED SECTION ---

    totals = {name: 0 for name in strategies}
    
    for text, freq in counts.items():
        if not text: continue
        
        # Corrected Parsing Logic
        if text[0].isdigit():
            first_char = None
            digits_part = text
        else:
            first_char = text[0]
            digits_part = text[1:]

        for name, func in strategies.items():
            totals[name] += func(first_char, digits_part) * freq
            
    df = pd.DataFrame(list(totals.items()), columns=['Strategy', 'Total Tokens'])
    df = df.sort_values('Total Tokens', ascending=True)
    
    baseline_tokens = df.loc[df['Strategy'] == 'Pure Single-Digit (Baseline)', 'Total Tokens'].iloc[0]
    df['Compression Ratio (vs Baseline)'] = baseline_tokens / df['Total Tokens']
    df['% Reduction (vs Baseline)'] = (1 - 1 / df['Compression Ratio (vs Baseline)']) * 100
    
    display(df.style.set_caption(f"Predicted Compressibility on '{label}' Dataset (Full Analysis)")
              .background_gradient(subset=['Compression Ratio (vs Baseline)'], cmap='Greens')
              .format({'Compression Ratio (vs Baseline)': '{:.3f}x', '% Reduction (vs Baseline)': '{:.2f}%'}))
    return df

# --- Run the Analysis ---
print("Analyzing compressibility for the 'cultura' dataset...")
df_comp_a = calculate_compressibility(counts_a, 'cultura')

print("\nAnalyzing compressibility for the 'lmsys' dataset...")
df_comp_b = calculate_compressibility(counts_b, 'lmsys')

Analyzing compressibility for the 'cultura' dataset...
Predicted Compressibility on 'cultura' Dataset (Full Analysis)
 	Strategy	Total Tokens	Compression Ratio (vs Baseline)	% Reduction (vs Baseline)
4	Hybrid (1-2 Digits Atomic, Space+Punct Merge)	225417007	2.421x	58.69%
3	Hybrid (1-2 Digits Atomic, Space-Merge)	269107998	2.028x	50.69%
1	3-Digit L2R/R2L	366724895	1.488x	32.80%
10	Single-Digit (Merge ALL Special)	421483074	1.295x	22.76%
2	Hybrid (1-2 Digits Atomic, No Merge)	422609237	1.291x	22.56%
5	Single-Digit (Merge '' '' Only)	450811058	1.210x	17.39%
9	Single-Digit (Merge ''.'' Only)	534918089	1.020x	1.98%
6	Single-Digit (Merge ''-'' Only)	536801464	1.017x	1.63%
7	Single-Digit (Merge ''('' Only)	540591582	1.009x	0.94%
8	Single-Digit (Merge '','' Only)	541146205	1.008x	0.83%
0	Pure Single-Digit (Baseline)	545696331	1.000x	0.00%

Analyzing compressibility for the 'lmsys' dataset...
Predicted Compressibility on 'lmsys' Dataset (Full Analysis)
 	Strategy	Total Tokens	Compression Ratio (vs Baseline)	% Reduction (vs Baseline)
4	Hybrid (1-2 Digits Atomic, Space+Punct Merge)	13302468	2.038x	50.93%
3	Hybrid (1-2 Digits Atomic, Space-Merge)	16103631	1.683x	40.60%
1	3-Digit L2R/R2L	20198663	1.342x	25.49%
10	Single-Digit (Merge ALL Special)	20551954	1.319x	24.19%
2	Hybrid (1-2 Digits Atomic, No Merge)	22437988	1.208x	17.23%
5	Single-Digit (Merge '' '' Only)	22570078	1.201x	16.74%
9	Single-Digit (Merge ''.'' Only)	26458510	1.025x	2.40%
6	Single-Digit (Merge ''-'' Only)	26501468	1.023x	2.24%
7	Single-Digit (Merge ''('' Only)	26697561	1.015x	1.52%
8	Single-Digit (Merge '','' Only)	26757485	1.013x	1.29%
0	Pure Single-Digit (Baseline)	27108287	1.000x	0.00%


analyze and propose a set of core strategies to test and present. Do we need to add others?

Note data ingestion was done with
pattern = r'.[0-9]+'

total_size = 0

for row in tqdm(dataset):
    texts = []
    if 'text' in row:
        texts.append(row['text'])
    if 'conversation' in row:
        content = "\n".join(turn['content'] for turn in row['conversation'])
        content = re.sub(r'NAME_\d','NAME', content)
        texts.append(content)
    for text in texts:
        text = str(text)
        total_size += len(text)
        for m in re.finditer(r'(.){50} \d{3}\b.{50}',text): 
            print(m)
        for pat in re.findall(pattern, text):
            num_counts[pat] += 1

print(f"Total unique numbers: {len(num_counts):,d}")
print(f"Total size of numbers: {sum(len(k) * v for k, v in num_counts.items()):,d} characters")
print(f"Total size of dataset: {total_size:,d} characters")