#!/usr/bin/env python3
"""
Extract one example of each of the 56 conversation types from the 56variations dataset
Creates both JSON and Markdown documentation
"""

import json
from collections import defaultdict

def extract_samples_from_file(filepath):
    """Extract one sample per conversation type from the dataset"""
    print(f"Processing {filepath}...")

    # Read file in chunks
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    # Parse JSON objects
    objects = content.split('\n}\n{')

    # Collect one example per type
    examples = {}
    type_shlokas = defaultdict(set)

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

        try:
            data = json.loads(obj)
            conv_type = data.get('conversation_type', '')

            # Skip if we already have this base type
            base_type = conv_type.rsplit('_v', 1)[0] if '_v' in conv_type else conv_type

            # Store first v1 example for each type
            if conv_type.endswith('_v1') or '_v' not in conv_type:
                if base_type not in examples:
                    examples[base_type] = data
                    type_shlokas[base_type].add(data.get('shloka_id', ''))

        except Exception as e:
            continue

        if i % 50000 == 0:
            print(f"  Processed {i} objects, found {len(examples)} types...")

    print(f"\nTotal unique types found: {len(examples)}")
    return examples

def create_markdown_document(examples, output_file="56_QUESTION_TYPES_EXAMPLES.md"):
    """Create comprehensive markdown documentation"""

    # Get metadata from first example
    first = list(examples.values())[0]
    metadata = first.get('metadata', {})
    shloka_id = first.get('shloka_id', 'N/A')

    md = f"""# 56 Question Types - Complete Examples

Generated from: `{shloka_id}`

---

## Source Shloka Information

**Sanskrit:**
```
{metadata.get('sanskrit', 'N/A')}
```

**Transliteration:**
```
{metadata.get('transliteration', 'N/A')}
```

**Section:** {metadata.get('section', 'N/A')}
**Section (Hindi):** {metadata.get('section_hi', 'N/A')}
**Section (Sanskrit):** {metadata.get('section_sa', 'N/A')}
**Chapter:** {metadata.get('chapter', 'N/A')}
**Verse:** {metadata.get('verse', 'N/A')}
**Translation (EN):** {metadata.get('translation_en', 'N/A')[:100]}...

---

## Table of Contents

1. [Category 1: Natural Questions (Conceptual) - 15 types](#category-1-natural-questions-conceptual)
2. [Category 2: Natural Questions (Analytical) - 15 types](#category-2-natural-questions-analytical)
3. [Category 3: Natural Questions (Application) - 15 types](#category-3-natural-questions-application)
4. [Category 4: Reference Questions - 40 types](#category-4-reference-questions)
5. [Category 5: Advanced Questions - 20 types](#category-5-advanced-questions)
6. [Category 6: Cross-Reference Questions - 10 types](#category-6-cross-reference-questions)
7. [Category 7: QnA Variations - 45 types](#category-7-qna-variations)

---

## Category 1: Natural Questions (Conceptual)

These questions treat the content as knowledge, NOT as written text.

### 1.1 English Variations (5 types)

"""

    # Natural Conceptual EN
    for i in range(1, 6):
        key = f"natural_conceptual_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 1.2 Hindi Variations (5 types)

"""

    # Natural Conceptual HI
    for i in range(1, 6):
        key = f"natural_conceptual_hi_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 1.3 Sanskrit Variations (5 types)

"""

    # Natural Conceptual SA
    for i in range(1, 6):
        key = f"natural_conceptual_sa_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """---

## Category 2: Natural Questions (Analytical)

### 2.1 English Variations (5 types)

"""

    # Natural Analytical EN
    for i in range(1, 6):
        key = f"natural_analytical_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 2.2 Hindi Variations (5 types)

"""

    # Natural Analytical HI
    for i in range(1, 6):
        key = f"natural_analytical_hi_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 2.3 Sanskrit Variations (5 types)

"""

    # Natural Analytical SA
    for i in range(1, 6):
        key = f"natural_analytical_sa_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """---

## Category 3: Natural Questions (Application)

### 3.1 English Variations (5 types)

"""

    # Natural Application EN
    for i in range(1, 6):
        key = f"natural_application_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 3.2 Hindi Variations (5 types)

"""

    # Natural Application HI
    for i in range(1, 6):
        key = f"natural_application_hi_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 3.3 Sanskrit Variations (5 types)

"""

    # Natural Application SA
    for i in range(1, 6):
        key = f"natural_application_sa_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """---

## Category 4: Reference Questions

These questions include text references for scholarly/academic use.

### 4.1 Translation (5 × 3 languages = 15 types)

"""

    # Reference Translation
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"reference_translation_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 4.2 ID Lookup (5 × 3 languages = 15 types)

"""

    # Reference ID Lookup
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"reference_id_lookup_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 4.3 Complete Details (5 × 3 languages = 15 types)

"""

    # Reference Complete Details
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"reference_complete_details_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """---

## Category 5: Advanced Questions

### 5.1 Clinical Case (5 types)

"""

    # Advanced Clinical Case
    for i in range(1, 6):
        key = f"advanced_clinical_case_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 5.2 Word Analysis (5 types)

"""

    # Advanced Word Analysis
    for i in range(1, 6):
        key = f"advanced_word_analysis_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 5.3 Historical Context (5 types)

"""

    # Advanced Historical Context
    for i in range(1, 6):
        key = f"advanced_historical_context_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 5.4 Comparative (5 types)

"""

    # Advanced Comparative
    for i in range(1, 6):
        key = f"advanced_comparative_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """---

## Category 6: Cross-Reference Questions

### 6.1 Intra-chapter (5 types)

"""

    # Cross-reference Intra-chapter
    for i in range(1, 6):
        key = f"cross_reference_intra_chapter_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 6.2 Section Relationship (5 types)

"""

    # Cross Section Relationship
    for i in range(1, 6):
        key = f"cross_section_relationship_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """---

## Category 7: QnA Variations

### 7.1 Meaning (5 × 3 languages = 15 types)

"""

    # QnA Meaning
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_meaning_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.2 Concept (5 × 3 languages = 15 types)

"""

    # QnA Concept
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_concept_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.3 Location (5 × 3 languages = 15 types)

"""

    # QnA Location
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_location_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.4 Context (5 × 3 languages = 15 types)

"""

    # QnA Context
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_context_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.5 Short (5 × 3 languages = 15 types)

"""

    # QnA Short
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_short_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.6 Section Info (5 × 3 languages = 15 types)

"""

    # QnA Section Info
    for lang in ['en', 'hi', 'sa']:
        md += f"**Language: {lang.upper()}**\n\n"
        for i in range(1, 6):
            key = f"qna_section_info_{lang}_v{i}"
            if key in examples:
                ex = examples[key]
                user = ex['messages'][1]['content']
                md += f"- **`{key}`**\n\n  Q: {user}\n\n"
        md += "\n"

    md += """### 7.7 Analytical (5 types)

"""

    # QnA Analytical
    for i in range(1, 6):
        key = f"qna_analytical_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 7.8 Word-by-word (5 types)

"""

    # QnA Word-by-word
    for i in range(1, 6):
        key = f"qna_word_by_word_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    md += """### 7.9 Application (5 types)

"""

    # QnA Application
    for i in range(1, 6):
        key = f"qna_application_en_v{i}"
        if key in examples:
            ex = examples[key]
            user = ex['messages'][1]['content']
            model = ex['messages'][2]['content'][:400]
            md += f"**`{key}`**\n\n"
            md += f"**Q:** {user}\n\n"
            md += f"**A:** {model}...\n\n---\n\n"

    # Summary
    md += f"""
---

## Summary

Total Question Types Documented: **{len(examples)}**

All examples extracted from the 56variations dataset.
Format: ShareGPT with system/user/model messages.
Each type has 5 variations (v1-v5) with different phrasings.

"""

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(md)

    print(f"Markdown saved to: {output_file}")
    return md

def create_json_output(examples, output_file="56_QUESTION_TYPES_EXAMPLES.json"):
    """Create JSON output with all examples"""

    # Simplify structure for JSON
    simple = {}
    for key, data in examples.items():
        simple[key] = {
            "shloka_id": data.get('shloka_id'),
            "question": data['messages'][1]['content'],
            "answer": data['messages'][2]['content'][:500]  # Truncate for size
        }

    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(simple, f, ensure_ascii=False, indent=2)

    print(f"JSON saved to: {output_file}")

def main():
    print("="*60)
    print("EXTRACTING 56 QUESTION TYPES EXAMPLES")
    print("="*60)

    examples = extract_samples_from_file("gemma4_ayurveda_56variations.jsonl")

    print("\nCreating documentation...")
    create_markdown_document(examples)
    create_json_output(examples)

    print("\n" + "="*60)
    print("DONE")
    print("="*60)

    # List all found types
    print("\nAll found question types:")
    for key in sorted(examples.keys()):
        print(f"  - {key}")

if __name__ == "__main__":
    main()