#!/usr/bin/env python3
"""
TinyGSM Dataset Initialization - Replace Test with GSM8K
按正常流程生成 train/val/test，最后把 test.bin 替换为 GSM8K 测试集
"""

import os
import re
import json
from pathlib import Path
from datasets import load_dataset
from tqdm import tqdm

from lmr.tokenizer import Tokenizer
from lmr.data.disk_dataset import DiskDataset

# Dataset sources
TINYGSM_HF = "TinyGSM/TinyGSM"
GSM8K_HF = "openai/gsm8k"


def initialize_dataset(dataset_config, dataset_dir):
    """
    Initialize TinyGSM dataset, then replace test with GSM8K.
    
    Steps:
    1. Load TinyGSM train split
    2. Split into train/val/test (normal process)
    3. Generate train.bin and validation.bin
    4. Load GSM8K test set
    5. Replace test.bin with GSM8K data
    """
    from pathlib import Path
    from lmr.tokenizer import Tokenizer
    from lmr.data.disk_dataset import DiskDataset

    dataset_dir = Path(dataset_dir)
    print(f"\n{'='*60}")
    print(f"Initializing TinyGSM dataset: {dataset_config.dataset_name}")
    print(f"Process: TinyGSM train/val + GSM8K test")
    print(f"{'='*60}\n")

    tokenizer = Tokenizer.get_instance()
    component_name = getattr(dataset_config, "component_name", "tinygsm")
    
    # Optional whitelist
    component_whitelist = getattr(dataset_config, "component_whitelist", None)
    if component_whitelist is not None and component_name not in component_whitelist:
        print(f"Component {component_name} not on whitelist, skipping...")
        return

    # Token limits
    token_limits = {
        "train": _parse_token_limit(getattr(dataset_config, "max_tokens_train", None)),
        "validation": _parse_token_limit(getattr(dataset_config, "max_tokens_validation", None)),
        "test": _parse_token_limit(getattr(dataset_config, "max_tokens_test", None)),
    }
    tokens_buffer = _parse_token_limit(getattr(dataset_config, "tokens_buffer", "10k"))

    # Create output directory
    component_dir = dataset_dir / dataset_config.dataset_name
    component_dir.mkdir(parents=True, exist_ok=True)

    print(f"Output directory: {component_dir}")
    print(f"Token limits: {token_limits}")
    print(f"Tokens buffer: {tokens_buffer}\n")

    # ===== Step 1-3: Process TinyGSM train and validation =====
    for split_name in ("train", "validation"):
        
        token_limit = token_limits[split_name]
        if token_limit is not None and tokens_buffer is not None:
            token_limit += tokens_buffer

        print(f"{'='*60}")
        print(f"Processing {split_name} split from TinyGSM...")
        print(f"Token limit (with buffer): {token_limit}")
        print(f"{'='*60}")

        try:
            ds_stream = _load_tinygsm_split(split_name, dataset_config)
        except Exception as e:
            print(f"Warning: Could not load {split_name} split: {e}")
            continue

        output_path = component_dir / f"{split_name}.bin"
        metadata_path = component_dir / f"metadata_{split_name}.json"

        # Transform to text format
        ds_transformed = _transform_tinygsm_to_text(ds_stream)

        # Generate binary file
        print(f"Generating binary file: {output_path}")
        
        DiskDataset.generate_bin(
            ds_transformed,
            tokenizer,
            output_path,
            token_limit=token_limit,
            metadata_path=metadata_path,
            pad_to_length = 1024
        )

        print(f"✓ Completed {split_name} split")
        print(f"  Binary file: {output_path}")
        print(f"  Metadata: {metadata_path}\n")

    # ===== Step 4-5: Replace test with GSM8K =====
    split_name = "test"
    token_limit = token_limits[split_name]
    if token_limit is not None and tokens_buffer is not None:
        token_limit += tokens_buffer

    print(f"{'='*60}")
    print(f"Replacing {split_name} split with GSM8K test set...")
    print(f"Token limit (with buffer): {token_limit}")
    print(f"{'='*60}")

    try:
        # Load GSM8K test
        print("Loading GSM8K test set from HuggingFace...")
        ds_gsm8k_test = load_dataset(GSM8K_HF, "main", split="test", streaming=False)
        print(f"✓ Loaded GSM8K test set: {len(ds_gsm8k_test)} samples\n")
    except Exception as e:
        print(f"❌ Error loading GSM8K test set: {e}")
        print("Keeping TinyGSM test split...")
        return

    output_path = component_dir / "test.bin"
    metadata_path = component_dir / "metadata_test.json"

    # Transform GSM8K to text format
    ds_transformed = _transform_gsm8k_to_text(ds_gsm8k_test)

    # Generate binary file (replaces any existing test.bin)
    print(f"Generating binary file: {output_path}")
    DiskDataset.generate_bin(
        ds_transformed,
        tokenizer,
        output_path,
        token_limit=token_limit,
        metadata_path=metadata_path
    )

    print(f"✓ Completed {split_name} split (GSM8K)")
    print(f"  Binary file: {output_path}")
    print(f"  Metadata: {metadata_path}\n")

    print(f"{'='*60}")
    print("✓ Dataset initialization complete!")
    print(f"  Train: TinyGSM")
    print(f"  Validation: TinyGSM")
    print(f"  Test: GSM8K (1319 samples)")
    print(f"{'='*60}\n")


def _load_tinygsm_split(split_name, dataset_config):
    """Load TinyGSM split, creating val from train if needed."""
    val_ratio = float(getattr(dataset_config, "val_ratio", 0.01))

    # Try to get available splits
    available_splits = []
    try:
        hf_all = load_dataset(TINYGSM_HF, split=None)
        if isinstance(hf_all, dict):
            available_splits = list(hf_all.keys())
    except Exception:
        available_splits = []

    hf_name = "train" if split_name == "train" else split_name

    # If split exists, use streaming
    if hf_name in available_splits:
        print(f"Using existing '{hf_name}' split (streaming mode)")
        return load_dataset(TINYGSM_HF, split=hf_name, streaming=True)

    # Otherwise, load train and split locally
    print(f"Split '{hf_name}' not found, loading 'train' and splitting locally")
    ds_full = load_dataset(TINYGSM_HF, split="train", streaming=False)

    if val_ratio <= 0:
        if split_name == "train":
            return ds_full.shuffle(seed=42)
        else:
            return iter([])

    # Compute validation count
    total = len(ds_full)
    val_count = max(1, int(round(total * val_ratio)))
    test_count = max(1, int(round(total * val_ratio)))  # Same size as val

    # Split
    split1 = ds_full.train_test_split(test_size=test_count, seed=42)
    ds_remain = split1["train"]
    ds_test = split1["test"]

    split2 = ds_remain.train_test_split(test_size=val_count, seed=43)
    ds_train = split2["train"]
    ds_val = split2["test"]

    # Return requested split
    if split_name == "train":
        return ds_train.shuffle(seed=44)
    elif split_name == "validation":
        return ds_val.shuffle(seed=45)
    else:  # test (will be replaced by GSM8K later)
        return ds_test.shuffle(seed=46)


def _transform_tinygsm_to_text(dataset_stream):
    """Transform TinyGSM samples to training format."""
    current = 0
    for sample in dataset_stream:
        try:
            current += 1

            if isinstance(sample, dict):
                getf = sample.get
            else:
                getf = lambda k, default=None: getattr(sample, k, default)

            question = (getf("question", None) or "").strip()
            code = (getf("code", None) or "").strip()
            answer = (getf("answer", None) or getf("final_answer", None) or "").strip()

            if not question:
                continue

            text = _format_training_text(question, code, answer)

            if current % 100000 == 0:
                print(text[:200])

            yield {"text": text}

        except Exception as e:
            print(f"Warning: Error processing sample #{current}: {e}")
            continue


def _transform_gsm8k_to_text(dataset):
    """
    Transform GSM8K test set to TinyGSM training format.
    
    GSM8K format: 
        question: "Natalia sold clips..."
        answer: "Step 1...\n#### 42"
    
    Convert to:
        <|bos|>Question: ...
        Solution:
        # Step 1...
        result = 42
        Answer: 42<|eos|>
    """
    current = 0
    for sample in dataset:
        try:
            current += 1

            question = (sample.get("question", "") or "").strip()
            answer_text = (sample.get("answer", "") or "").strip()

            if not question or not answer_text:
                continue

            # Extract numerical answer from GSM8K format
            final_answer = _extract_gsm8k_answer(answer_text)

            # Convert to code format
            code = _convert_gsm8k_to_code(question, answer_text, final_answer)

            # Format in TinyGSM style
            text = _format_training_text(question, code, final_answer)

            if current % 100 == 0:
                print(f"GSM8K sample {current}:")
                print(text[:300])
                print("---")

            yield {"text": text}

        except Exception as e:
            print(f"Warning: Error processing GSM8K sample #{current}: {e}")
            continue


def _extract_gsm8k_answer(answer_text: str) -> str:
    """
    Extract numerical answer from GSM8K format.
    GSM8K uses "#### number" to mark the final answer.
    
    支持多种格式：
    - #### 42
    - #### 3.14
    - #### -5
    - #### 1,234
    """
    # Method 1: #### pattern (primary)
    match = re.search(r'####\s*([+-]?[\d,]+\.?\d*)', answer_text)
    if match:
        return match.group(1).replace(',', '').strip()
    
    # Method 2: "result =" pattern
    match = re.search(r'result\s*=\s*([+-]?[\d,]+\.?\d*)', answer_text, re.IGNORECASE)
    if match:
        return match.group(1).replace(',', '').strip()
    
    # Method 3: Last number (fallback)
    numbers = re.findall(r'([+-]?[\d,]+\.?\d*)', answer_text)
    if numbers:
        return numbers[-1].replace(',', '').strip()
    
    return "0"


def _convert_gsm8k_to_code(question: str, answer_text: str, final_answer: str) -> str:
    """
    Convert GSM8K natural language solution to Python-like code.
    
    Strategy: Extract reasoning steps as comments, add final calculation.
    """
    # Extract reasoning (text before ####)
    if "####" in answer_text:
        reasoning = answer_text.split("####")[0].strip()
    else:
        reasoning = answer_text.strip()
    
    # Build code
    code_lines = ["# Solution:"]
    
    # Add reasoning steps as comments
    for line in reasoning.split('\n'):
        line = line.strip()
        if line:
            # Remove repeated equals signs (<<...>>)
            line = re.sub(r'<<[^>]+>>', '', line)
            code_lines.append(f"# {line}")
    
    # Add final result
    if final_answer:
        code_lines.append(f"\nresult = {final_answer}")
    
    return "\n".join(code_lines)


def _format_training_text(question, code, answer):
    """
    Format training text in TinyGSM style:
    
    <|bos|>Question: ...
    
    Solution:
    ```python
    ...
    ```
    
    Answer: ...<|eos|>
    """
    q = (question or "").strip()
    c = (code or "").rstrip()
    a = (answer or "").strip()

    parts = []
    parts.append("<|bos|>Question: " + q)
    parts.append("")

    if c:
        parts.append("Solution:")
        parts.append("```python")
        parts.append(c)
        parts.append("```")
        parts.append("")
    else:
        parts.append("Solution: ")
        parts.append("")

    if a:
        parts.append("Answer: " + a)
    else:
        parts.append("Answer: ")

    parts.append("<|eos|>")

    return "\n".join(parts)


def _parse_token_limit(limit_str):
    """Parse token limit string to integer."""
    if limit_str is None:
        return None

    s = str(limit_str).lower().strip()

    if s in ("none", "null", "na", ""):
        return None

    try:
        if s.endswith("b"):
            return int(float(s[:-1]) * 1_000_000_000)
        elif s.endswith("m"):
            return int(float(s[:-1]) * 1_000_000)
        elif s.endswith("k"):
            return int(float(s[:-1]) * 1_000)
        else:
            return int(s)
    except Exception:
        print(f"Warning: Could not parse token limit '{limit_str}' -> treating as None")
        return None


# Standalone execution
if __name__ == "__main__":
    from types import SimpleNamespace

    config = SimpleNamespace(
        dataset_name="tinygsm",
        component_name="tinygsm",
        max_tokens_train="50m",
        max_tokens_validation=None,
        max_tokens_test=None,
        tokens_buffer="10k",
        component_whitelist=None,
        val_ratio=0.01  # 1% for validation
    )

    dataset_dir = Path("./data/processed")

    print("=" * 60)
    print("TinyGSM Dataset Initialization")
    print("=" * 60)

    initialize_dataset(config, dataset_dir)

    print("\n" + "=" * 60)
    print("✓ Initialization Complete!")
    print("=" * 60)