import json
import random

def parse_jsonl_pretty(filepath):
    """Parse pretty-printed JSONL file"""
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    # Split by closing brace followed by newline and opening brace
    objects = content.split('\n}\n{')
    conversations = []

    for i, obj in enumerate(objects):
        if i == 0:
            obj = obj + '\n}'
        elif i == len(objects) - 1:
            obj = '{\n' + obj
        else:
            obj = '{\n' + obj + '\n}'

        try:
            conv = json.loads(obj)
            conversations.append(conv)
        except:
            pass

    return conversations

def create_200k_subset():
    print("Parsing dataset (this may take a moment)...")
    conversations = parse_jsonl_pretty('gemma4_ayurveda_ultra_enriched.jsonl')

    print(f"Loaded {len(conversations)} total conversations")

    # Group by shloka_id
    shloka_conversations = {}
    for conv in conversations:
        sid = conv.get('shloka_id', '')
        if sid and sid != 'INVALID':
            if sid not in shloka_conversations:
                shloka_conversations[sid] = []
            shloka_conversations[sid].append(conv)

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

    # Priority types - balanced selection
    priority_types = [
        # Core translations
        'direct_translation', 'direct_translation_hi_full', 'direct_translation_sa_full',
        # Reference queries
        'reference_query', 'reference_query_hi_full', 'reference_query_sa_full',
        # Multi-turn
        'multi_turn', 'multi_turn_hi_full', 'multi_turn_sa_full',
        # Analysis
        'complete_analysis', 'complete_analysis_hi_full', 'complete_analysis_sa_full',
        # Sanskrit input
        'sanskrit_input', 'sanskrit_input_hi_full', 'sanskrit_input_sa_full',
        # Rich contextual QnA
        'qna_conceptual_en', 'qna_conceptual_hi', 'qna_conceptual_sa',
        'qna_practical_en', 'qna_practical_hi', 'qna_practical_sa',
        'qna_significance_en', 'qna_significance_hi', 'qna_significance_sa',
        # Meaning QnA
        'qna_meaning', 'qna_meaning_hi', 'qna_meaning_sa',
        'qna_specific_verse', 'qna_specific_verse_hi', 'qna_specific_verse_sa',
        # Entity-based QnA
        'entity_disease_en', 'entity_disease_hi', 'entity_disease_sa',
        'entity_procedure_en', 'entity_procedure_hi', 'entity_procedure_sa',
        'entity_symptom_en', 'entity_symptom_hi', 'entity_symptom_sa',
        'entity_substance_en', 'entity_substance_hi', 'entity_substance_sa',
        # Context QnA
        'chapter_context', 'chapter_context_hi_full', 'chapter_context_sa_full',
        'cross_reference', 'cross_reference_hi_full', 'cross_reference_sa_full',
        'qna_chapter', 'qna_chapter_hi', 'qna_chapter_sa',
        'qna_location', 'qna_location_hi', 'qna_location_sa',
        'qna_context', 'qna_context_hi', 'qna_context_sa',
        # Language specific
        'hindi_translation', 'hindi_translation_hi_full', 'hindi_translation_sa_full',
        'transliteration_request', 'transliteration_request_hi_full', 'transliteration_request_sa_full',
        'citation_identification', 'citation_identification_hi_full', 'citation_identification_sa_full',
    ]

    # Select balanced subset
    selected = []

    for sid, convs in shloka_conversations.items():
        # Group by type
        by_type = {}
        for c in convs:
            ctype = c.get('conversation_type', '')
            if ctype not in by_type:
                by_type[ctype] = []
            by_type[ctype].append(c)

        # Select one from each priority type
        shloka_selected = []
        for ptype in priority_types:
            if ptype in by_type and by_type[ptype]:
                shloka_selected.append(by_type[ptype][0])

        # Fill remaining with any available types to reach ~56 per shloka
        if len(shloka_selected) < 56:
            for ctype, clist in by_type.items():
                if len(shloka_selected) >= 56:
                    break
                # Add if not already in selected
                if ctype not in [c.get('conversation_type') for c in shloka_selected]:
                    shloka_selected.append(clist[0])

        # Limit to exactly 56 if more
        if len(shloka_selected) > 56:
            shloka_selected = shloka_selected[:56]

        selected.extend(shloka_selected)

    print(f"Selected {len(selected)} conversations")

    # Shuffle
    random.shuffle(selected)

    # Limit to exactly 200k
    final = selected[:200000]

    # Write
    print("\nWriting 200k subset...")
    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"\nDataset complete!")
    print(f"Total conversations: {len(final)}")
    print(f"Shlokas covered: {len(shloka_conversations)}")
    print(f"Average per shloka: {len(final) / len(shloka_conversations):.1f}")
    print(f"Output: gemma4_ayurveda_200k_subset.jsonl")

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

    print(f"\nConversation type distribution:")
    for t, c in sorted(type_counts.items(), key=lambda x: -x[1])[:15]:
        print(f"  {t}: {c} ({c/len(final)*100:.1f}%)")

if __name__ == "__main__":
    create_200k_subset()