import json
import random

print("Creating 200k subset...")

# Collect all unique shloka_ids first
print("Step 1: Collecting unique shloka IDs...")
shloka_ids = []
with open('gemma4_ayurveda_ultra_enriched.jsonl', 'r', encoding='utf-8') as f:
    buffer = ""
    for line in f:
        buffer += line
        if line.strip() == "}":
            try:
                conv = json.loads(buffer)
                sid = conv.get('shloka_id', '')
                if sid and sid != 'INVALID' and sid not in shloka_ids:
                    shloka_ids.append(sid)
            except:
                pass
            buffer = ""

print(f"Found {len(shloka_ids)} unique shlokas")

# Calculate how many per shloka
per_shloka = 200000 // len(shloka_ids)
print(f"Will select ~{per_shloka} conversations per shloka")

# Now collect conversations
print("\nStep 2: Collecting conversations...")
selected = []
shloka_counts = {sid: 0 for sid in shloka_ids}

with open('gemma4_ayurveda_ultra_enriched.jsonl', 'r', encoding='utf-8') as f:
    buffer = ""
    count = 0
    for line in f:
        buffer += line
        if line.strip() == "}":
            try:
                conv = json.loads(buffer)
                sid = conv.get('shloka_id', '')
                if sid in shloka_counts and shloka_counts[sid] < per_shloka:
                    selected.append(conv)
                    shloka_counts[sid] += 1
                    count += 1

                    if count % 10000 == 0:
                        print(f"  Collected {count} conversations")

                    if count >= 200000:
                        break
            except:
                pass
            buffer = ""

print(f"\nCollected {len(selected)} conversations")

# Fill up if needed
if len(selected) < 200000:
    print(f"Need {200000 - len(selected)} more, collecting additional...")
    with open('gemma4_ayurveda_ultra_enriched.jsonl', 'r', encoding='utf-8') as f:
        buffer = ""
        for line in f:
            buffer += line
            if line.strip() == "}":
                try:
                    conv = json.loads(buffer)
                    sid = conv.get('shloka_id', '')
                    if sid in shloka_counts:
                        selected.append(conv)
                        if len(selected) >= 200000:
                            break
                except:
                    pass
                buffer = ""

print(f"\nFinal count: {len(selected)}")

# Shuffle
print("Shuffling...")
random.shuffle(selected)
final = selected[:200000]

# Write
print("Writing to file...")
with open('gemma4_ayurveda_200k_subset.jsonl', 'w', encoding='utf-8') as f:
    for conv in final:
        f.write(json.dumps(conv, ensure_ascii=False, indent=2) + '\n')

print(f"\nDone! Created gemma4_ayurveda_200k_subset.jsonl with {len(final)} conversations")

# Stats
type_counts = {}
for conv in final:
    ctype = conv.get('conversation_type', '').split('_')[0]
    type_counts[ctype] = type_counts.get(ctype, 0) + 1

print("\nTop conversation types:")
for t, c in sorted(type_counts.items(), key=lambda x: -x[1])[:10]:
    print(f"  {t}: {c}")