data/prepare_data.py

"""Generate synthetic training data for FSI_Edge.

Creates code + NLP mixtures using rule-based generation and
structural annotations for the novel architecture training.
"""
import os
import json
import random
import math
from tqdm import tqdm


# ============================================================================
# Code templates covering Python, JavaScript, Java, C++, Go, Rust, SQL
# ============================================================================

PYTHON_TEMPLATES = [
    # Data structures
    lambda: f"def {random.choice(['merge_sort', 'quick_sort', 'binary_search', 'bfs', 'dfs', 'dijkstra', 'kruskal', 'prim', 'topological_sort', 'lcs', 'lis', 'knapsack', 'edit_distance', 'levenshtein', 'rabin_karp', 'kmp', 'z_algorithm', 'manacher', 'suffix_array', 'segment_tree', 'fenwick_tree', 'trie', 'union_find', 'heap'])}" + f"({', '.join(random.choice(['arr', 'nums', 'data', 'graph', 'tree', 'values', 'items', 'seq', 's', 't', 'n', 'm', 'k']) for _ in range(random.randint(1, 3)))}):\n" + '    """' + random.choice([
        'Efficient implementation with O(n log n) complexity.',
        'Standard algorithm with optimized memory usage.',
        'Iterative approach with explicit stack.',
        'Recursive with memoization for optimal performance.',
        'Two-pointer technique for linear time.',
    ]) + '"""\n' + '    n = len(' + random.choice(['arr', 'nums', 'data']) + ')\n' + _gen_body(5, 15),
    
    # System design / API
    lambda: f"class {random.choice(['RateLimiter', 'Cache', 'Logger', 'Database', 'Queue', 'Stack', 'Graph', 'Trie', 'Session', 'Router', 'Handler', 'Middleware', 'Parser', 'Tokenizer', 'Validator'])}" + ":\n    " + f'def __init__(self, {", ".join(random.choice(["config", "options", "params", "settings", "timeout", "max_size", "capacity", "rate"]) for _ in range(random.randint(1, 3)))}):\n' + '        self.config = config\n        ' + _gen_body(3, 8),
    
    # Async
    lambda: f"async def {random.choice(['fetch_data', 'process_batch', 'handle_request', 'stream_results', 'poll_status', 'aggregate_metrics', 'sync_replica', 'ingest_events'])}" + f"({', '.join(random.choice(['url', 'session', 'data', 'queue', 'task_id', 'batch', 'stream', 'params']) for _ in range(random.randint(1, 3)))}):\n" + '    """Async version with proper error handling."""\n    try:\n        ' + _gen_body(3, 6) + '\n    except Exception as e:\n        logger.error(f"Failed: {e}")\n        raise',
    
    # Type-annotated functions
    lambda: f"def {random.choice(['process_data', 'transform_input', 'validate_schema', 'compute_metrics', 'serialize_output', 'parse_query', 'normalize_text', 'extract_features'])}" + f"(input_data: {random.choice(['str', 'int', 'float', 'List[str]', 'Dict[str, Any]', 'Optional[Dict]', 'Sequence[int]', 'Iterable[str]', 'Callable', 'Path'])}," + f" {random.choice(['config', 'options', 'params'])}: {random.choice(['Config', 'Options', 'Optional[Dict]', 'dict', 'Any'])} = None) -> {random.choice(['str', 'int', 'bool', 'List[str]', 'Dict[str, Any]', 'Optional[str]', 'Result', 'None'])}:\n" + '    """' + random.choice([
        'Type-safe implementation with input validation.',
        'Processes input and returns formatted result.',
        'Transforms data according to configuration.',
    ]) + '"""\n    ' + _gen_body(5, 10),
]

JAVASCRIPT_TEMPLATES = [
    lambda: f"function {random.choice(['quickSort', 'mergeSort', 'binarySearch', 'debounce', 'throttle', 'deepClone', 'flatten', 'groupBy', 'chunk', 'pick', 'omit', 'pipe', 'compose', 'memoize'])}" + f"({', '.join(random.choice(['arr', 'obj', 'fn', 'items', 'data', 'key', 'value', 'n', 'delay', 'ms']) for _ in range(random.randint(1, 3)))}" + ") {\n  " + _gen_body(5, 12, indent='  ') + "\n}",
    lambda: f"class {random.choice(['EventEmitter', 'StreamHandler', 'CacheManager', 'RateLimiter', 'ConnectionPool', 'Router', 'Store', 'Controller', 'Service', 'Provider'])}" + " {\n  constructor(" + ", ".join(random.choice(['options', 'config', 'params', 'initialState']) for _ in range(random.randint(1, 2))) + ") {\n    " + _gen_body(3, 6, indent='    ') + "\n  }\n\n  " + _gen_method(random.randint(2, 5)) + "\n}",
    lambda: f"async function {random.choice(['fetchData', 'processStream', 'handleRequest', 'batchProcess', 'pollStatus', 'syncData', 'aggregateResults', 'streamResponse'])}" + f"({', '.join(random.choice(['url', 'data', 'session', 'queue', 'taskId', 'batch', 'stream', 'params']) for _ in range(random.randint(1, 3)))}" + ") {\n  try {\n    " + _gen_body(3, 6, indent='    ') + "\n  } catch (error) {\n    console.error(`Failed: ${error.message}`);\n    throw error;\n  }\n}",
]

JAVA_TEMPLATES = [
    lambda: f"public class {random.choice(['Sorter', 'Parser', 'Validator', 'Processor', 'Handler', 'Service', 'Repository', 'Controller', 'Mapper', 'Util'])}" + " {\n    " + _gen_java_method(random.randint(1, 3)) + "\n}",
    lambda: f"@{random.choice(['Override', 'Autowired', 'Value', 'Bean', 'Component'])}\npublic " + random.choice(['void', 'int', 'String', 'boolean', 'List<String>', 'Optional<Object>']) + f" {random.choice(['process', 'handle', 'execute', 'validate', 'transform', 'compute', 'parse', 'map', 'filter', 'reduce'])}" + f"({random.choice(['String input', 'int value', 'Object data', 'List<String> items', 'Map<String, Object> params'])}) throws {random.choice(['Exception', 'IOException', 'RuntimeException'])}" + " {\n        " + _gen_body(3, 7, indent='        ') + "\n    }",
]

CPP_TEMPLATES = [
    lambda: f"template<typename {random.choice(['T', 'K', 'V', 'Iter', 'Container'])}>\nauto {random.choice(['sort', 'filter', 'transform', 'reduce', 'accumulate', 'partition', 'merge', 'search', 'index_of', 'contains'])}" + f"(const {random.choice(['std::vector<T>&', 'std::list<T>&', 'std::deque<T>&', 'T* begin', 'T* end'])}" + f", {random.choice(['const T& value', 'std::function<bool(T)> pred', 'auto&& func', 'T init'])}) -> {random.choice(['void', 'bool', 'size_t', 'std::vector<T>', 'T'])}" + " {\n    " + _gen_body(5, 10, indent='    ') + "\n}",
    lambda: f"class {random.choice(['Graph', 'Tree', 'Heap', 'Trie', 'SegmentTree', 'FenwickTree', 'UnionFind', 'HashMap', 'ConcurrentQueue', 'ThreadPool'])}" + " {\npublic:\n    " + f"{random.choice(['Graph', 'Tree', 'Heap'])}" + "() { " + _gen_body(2, 4, indent='        ') + " }\n    ~" + "{}();\n    " + _gen_cpp_method(random.randint(2, 4)) + "\nprivate:\n    " + _gen_body(2, 5, indent='    ') + "\n};",
]

RUST_TEMPLATES = [
    lambda: f"fn {random.choice(['quick_sort', 'merge_sort', 'binary_search', 'parse_input', 'validate_data', 'transform_json', 'compute_hash', 'encrypt_data', 'decode_base64', 'compress_data'])}" + f"<{random.choice(['T: Ord', 'T: Clone + Debug', 'T: Hash', 'T: Serialize'])}>" + f"(input: &[{random.choice(['T', 'u8', 'String'])}]" + f", {random.choice(['key: &str', 'config: &Config', 'threshold: usize', 'buffer: &mut [u8]', 'callback: impl Fn(T) -> bool'])}) -> {random.choice(['Vec<T>', 'Result<String, Error>', 'Option<usize>', 'HashMap<String, T>', 'bool', 'usize'])}" + " {\n    " + _gen_body(5, 10, indent='    ') + "\n}",
    lambda: f"pub struct {random.choice(['Cache', 'Pool', 'Router', 'Session', 'Stream', 'Buffer', 'Parser', 'Encoder', 'Decoder', 'Index'])}" + "<" + random.choice(['K', 'K: Hash + Eq', 'T', 'T: Clone']) + "> {\n    " + random.choice(['inner: HashMap<K, V>', 'data: Vec<T>', 'items: VecDeque<T>', 'store: Arc<RwLock<HashMap<K, V>>>']) + ",\n    " + _gen_body(1, 3, indent='    ') + "\n}\n\nimpl" + "<" + random.choice(['K: Hash + Eq', 'T: Clone', 'K, V']) + "> " + random.choice(['Cache', 'Pool', 'Router']) + "<" + random.choice(['K, V', 'K', 'T']) + "> {\n    " + _gen_rust_method(random.randint(2, 4)) + "\n}",
]

GO_TEMPLATES = [
    lambda: f"func {random.choice(['QuickSort', 'MergeSort', 'BinarySearch', 'ParseJSON', 'ValidateInput', 'TransformData', 'ComputeHash', 'CompressData', 'EncryptPayload', 'DecodeBase64'])}" + f"({', '.join(random.choice(['data []byte', 'input string', 'items []int', 'ctx context.Context', 'w io.Writer', 'r io.Reader', 'key string', 'config Config']) for _ in range(random.randint(1, 3)))}" + f") ({random.choice(['[]byte', 'string', 'int', 'error', 'Result', '(*Result, error)', 'bool'])}" + ") {\n\t" + _gen_body(5, 10, indent='\t') + "\n}",
    lambda: f"type {random.choice(['Cache', 'Pool', 'Router', 'Session', 'Store', 'Service', 'Handler', 'Middleware', 'Parser', 'Index'])}" + " struct {\n\t" + random.choice(['mu sync.RWMutex', 'data map[string]interface{}', 'items []Item', 'store *sql.DB', 'client *http.Client', 'ctx context.Context', 'cancel context.CancelFunc']) + "\n\t" + _gen_body(1, 3, indent='\t') + "\n}\n\nfunc (c *" + random.choice(['Cache', 'Pool', 'Router']) + ") " + _gen_go_method(2, 4) + "\n",
]

SQL_TEMPLATES = [
    lambda: f"WITH {random.choice(['cte_filtered', 'ranked_items', 'grouped_data', 'recent_entries', 'aggregated_metrics'])} AS (\n    SELECT\n        " + random.choice(['id', 'name', 'category', 'status']) + ",\n        " + random.choice(['COUNT(*)', 'SUM(amount)', 'AVG(score)', 'MAX(created_at)', 'MIN(updated_at)']) + " as " + random.choice(['total', 'avg_score', 'max_date', 'min_date']) + ",\n        " + random.choice(['ROW_NUMBER() OVER (PARTITION BY category ORDER BY score DESC)', 'RANK() OVER (ORDER BY total DESC)', 'DENSE_RANK() OVER (ORDER BY created_at)']) + " as rn\n    FROM " + random.choice(['users', 'orders', 'products', 'sessions', 'events', 'transactions']) + "\n    WHERE " + random.choice(['status = \'active\'', 'created_at > NOW() - INTERVAL \'30 days\'', 'amount > 100', 'deleted_at IS NULL']) + "\n    GROUP BY " + random.choice(['id', 'name', 'category', 'status']) + "\n    HAVING " + random.choice(['COUNT(*) > 1', 'SUM(amount) > 1000', 'AVG(score) > 50']) + "\n)\nSELECT * FROM cte_filtered WHERE rn = 1 ORDER BY total DESC LIMIT 10;",
    lambda: "CREATE TABLE " + random.choice(['users', 'orders', 'products', 'sessions', 'audit_log', 'inventory', 'payments', 'shipping', 'reviews', 'categories']) + " (\n    id BIGSERIAL PRIMARY KEY,\n    " + random.choice(['uuid UUID DEFAULT gen_random_uuid() UNIQUE NOT NULL', 'name VARCHAR(255) NOT NULL', 'email VARCHAR(255) UNIQUE NOT NULL', 'status VARCHAR(50) DEFAULT \'active\'', 'created_at TIMESTAMPTZ DEFAULT NOW()', 'updated_at TIMESTAMPTZ DEFAULT NOW()']) + ",\n    " + random.choice(['metadata JSONB DEFAULT \'{}\'', 'parent_id BIGINT REFERENCES ' + random.choice(['users', 'orders', 'products']) + '(id)', 'index UNIQUE', 'CHECK (amount > 0)', 'EXCLUDE USING gist (period WITH &&)']) + "\n);\n\nCREATE INDEX idx_" + random.choice(['users', 'orders', 'products']) + "_" + random.choice(['created_at', 'status', 'email', 'uuid']) + " ON " + random.choice(['users', 'orders', 'products']) + "(" + random.choice(['created_at', 'status', 'email', 'uuid']) + ");\n",
]

# ============================================================================
# NLP Templates for general language understanding
# ============================================================================

NLP_TEMPLATES = [
    lambda: f"Question: {random.choice(['Explain the concept of', 'Describe how', 'What is the difference between', 'Compare and contrast', 'Provide an example of', 'Why is', 'How does', 'What are the advantages of', 'Summarize', 'Analyze the'])} {random.choice(['recursion', 'dynamic programming', 'object-oriented programming', 'functional programming', 'machine learning', 'neural networks', 'time complexity', 'space complexity', 'distributed systems', 'microservices', 'REST APIs', 'graph databases', 'blockchain', 'quantum computing', 'compiler design', 'operating systems', 'network protocols', 'database indexing', 'caching strategies', 'load balancing'])}?\n\nAnswer: {random.choice(['Let me explain this step by step.', 'Here is a comprehensive overview.', 'I will break this down into key concepts.', 'The fundamental idea is that', 'At its core, this involves'])} " + _gen_nlp_explanation(50, 150) + "[/ANSWER]",
    lambda: f"<|code|> Write a function that {random.choice(['sorts an array using quicksort', 'finds the longest palindromic substring', 'computes the nth fibonacci number', 'checks if a string is a valid parentheses sequence', 'finds the shortest path in a graph', 'serializes a binary tree', 'implements a LRU cache', 'finds all anagrams in a string', 'solves the N-Queens problem', 'implements a trie for autocomplete'])}.\n\n<|thought|> Let me think about this carefully. {_gen_nlp_explanation(30, 80)} Therefore, I will implement it as follows:\n\n<|answer|> def " + random.choice(['solve', 'find', 'compute', 'check', 'implement', 'serialize', 'deserialize', 'search', 'sort', 'count']) + "(..." + ")",
    lambda: f"<|explain|> Explain this code:\n```python\n{random.choice(PYTHON_TEMPLATES)()}\n```\n\n<|answer|> This code " + _gen_nlp_explanation(40, 120),
    lambda: f"<|debug|> Find the bug:\n```python\ndef buggy_function(x):\n    result = []\n    for i in range(len(x)):\n        if x[i] % 2 == 0:\n            result.append(x[i])\n    return result\n```\nThe function filters even numbers but has a performance issue.\n\n<|answer|> The bug is " + _gen_nlp_explanation(30, 80),
    lambda: f"<|test|> Generate unit tests for:\n```python\n{random.choice(PYTHON_TEMPLATES)()}\n```\n\n<|answer|> "
]


def _gen_body(min_lines, max_lines, indent='    '):
    """Generate random function body."""
    lines = []
    statements = [
        f'{indent}if n == 0:',
        f'{indent}    return 0',
        f'{indent}result = []',
        f'{indent}for i in range(n):',
        f'{indent}    result.append(i * 2)',
        f'{indent}while left <= right:',
        f'{indent}    mid = (left + right) // 2',
        f'{indent}    if arr[mid] == target:',
        f'{indent}        return mid',
        f'{indent}    elif arr[mid] < target:',
        f'{indent}        left = mid + 1',
        f'{indent}    else:',
        f'{indent}        right = mid - 1',
        f'{indent}return ",".join(map(str, result))',
        f'{indent}temp = [0] * (right - left + 1)',
        f'{indent}i = left',
        f'{indent}j = mid + 1',
        f'{indent}hash_map = {{}}',
        f'{indent}if key in hash_map:',
        f'{indent}    return hash_map[key]',
        f'{indent}cache[key] = value',
        f'{indent}if cache.get(key) is not None:',
        f'{indent}    return cache[key]',
        f'{indent}seen = set()',
        f'{indent}if item not in seen:',
        f'{indent}    seen.add(item)',
        f'{indent}    queue.append(item)',
        f'{indent}visited[node] = True',
        f'{indent}for neighbor in graph[node]:',
        f'{indent}    if not visited[neighbor]:',
        f'{indent}        dfs(neighbor)',
        f'{indent}result = sum(data) / len(data) if data else 0',
        f'{indent}return [x for x in items if x > threshold]',
    ]
    
    n = random.randint(min_lines, max_lines)
    for _ in range(n):
        lines.append(random.choice(statements))
    return '\n'.join(lines)


def _gen_method(n):
    """Generate random class methods."""
    methods = []
    for _ in range(n):
        m = f"  {random.choice(['get', 'set', 'add', 'remove', 'find', 'update', 'clear', 'has', 'size', 'isEmpty'])}" + f"({', '.join(random.choice(['key', 'value', 'item', 'id', 'name', 'data', 'config']) for _ in range(random.randint(0, 2)))})" + " {\n    " + _gen_body(2, 5, indent='    ') + "\n  }"
        methods.append(m)
    return '\n\n  '.join(methods)


def _gen_java_method(n):
    """Generate random Java methods."""
    methods = []
    for _ in range(n):
        visibility = random.choice(['public', 'private', 'protected'])
        is_static = random.choice(['', ' static'])
        return_type = random.choice(['void', 'int', 'String', 'boolean', 'List<String>', 'Optional<Object>', 'Map<String, Object>'])
        m = f"{visibility}{is_static} {return_type} {random.choice(['process', 'handle', 'execute', 'validate', 'transform', 'compute'])}" + f"({random.choice(['String input', 'int value', 'Object data', 'List<String> items'])}" + ") {\n        " + _gen_body(3, 7, indent='        ') + "\n    }"
        methods.append(m)
    return '\n\n    '.join(methods)


def _gen_cpp_method(n):
    """Generate random C++ methods."""
    methods = []
    for _ in range(n):
        m = f"{random.choice(['void', 'int', 'bool', 'std::string', 'size_t', 'auto'])} {random.choice(['insert', 'erase', 'find', 'contains', 'size', 'empty', 'clear', 'merge', 'split', 'traverse'])}" + f"({', '.join(random.choice(['const T& value', 'size_t index', 'const K& key', 'auto&& callback', 'const std::string& path']) for _ in range(random.randint(0, 2)))})" + " {\n        " + _gen_body(3, 7, indent='        ') + "\n    }"
        methods.append(m)
    return '\n\n    '.join(methods)


def _gen_rust_method(n):
    """Generate random Rust methods."""
    methods = []
    for _ in range(n):
        m = f"    pub fn {random.choice(['new', 'insert', 'get', 'remove', 'contains', 'len', 'is_empty', 'clear', 'iter', 'into_iter'])}" + f"({', '.join(random.choice(['&self', '&mut self', 'key: K', 'value: V', 'other: &Self', 'index: usize']) for _ in range(random.randint(1, 3)))})" + f" -> {random.choice(['Self', 'Option<&V>', 'bool', 'usize', 'Result<(), Error>', 'Vec<&K>', 'Iterator<Item = &K>'])}" + " {\n        " + _gen_body(3, 7, indent='        ') + "\n    }"
        methods.append(m)
    return '\n    '.join(methods)


def _gen_go_method(min_n, max_n):
    """Generate random Go methods."""
    methods = []
    for _ in range(random.randint(min_n, max_n)):
        m = f"func (c *" + random.choice(['Cache', 'Pool', 'Router']) + ") " + f"{random.choice(['Get', 'Set', 'Add', 'Remove', 'Find', 'Update', 'Clear', 'Has', 'Size'])}" + f"({', '.join(random.choice(['key string', 'value interface{}', 'item Item', 'ctx context.Context', 'id int64']) for _ in range(random.randint(1, 2)))}) " + f"({random.choice(['interface{}', 'error', 'bool', '(*Item, error)', 'int', 'string'])}"+ ") {\n\t" + _gen_body(3, 7, indent='\t') + "\n}"
        methods.append(m)
    return '\n\n\t'.join(methods)


def _gen_nlp_explanation(min_words, max_words):
    """Generate natural language explanation text."""
    concepts = [
        'time complexity', 'space complexity', 'data structure', 'algorithm',
        'recursive solution', 'iterative approach', 'divide and conquer',
        'dynamic programming', 'greedy algorithm', 'backtracking',
        'binary search', 'sorting', 'graph traversal', 'pattern matching',
        'hash table', 'binary tree', 'heap', 'stack', 'queue',
        'linked list', 'array', 'matrix', 'string manipulation',
        'bit manipulation', 'two pointers', 'sliding window',
        'union find', 'trie', 'segment tree', 'fenwick tree',
    ]
    
    verbs = ['processes', 'computes', 'finds', 'returns', 'transforms',
             'validates', 'parses', 'generates', 'extracts', 'merges',
             'splits', 'sorts', 'searches', 'traverses', 'optimizes']
    
    n_words = random.randint(min_words, max_words)
    words = []
    
    templates = [
        f"The {random.choice(concepts)} {random.choice(verbs)} the input",
        f"This {random.choice(concepts)} is {random.choice(['efficient', 'optimal', 'simple', 'robust', 'scalable', 'flexible'])}",
        f"We can {random.choice(verbs)} this by using a {random.choice(concepts)}",
        f"The key insight is that {random.choice(concepts)} allows us to",
        f"By leveraging {random.choice(concepts)}, we achieve {random.choice(['O(n)', 'O(log n)', 'O(n log n)', 'O(1)', 'O(n²)'])}",
        f"This is a {random.choice(['common', 'standard', 'novel', 'elegant', 'fundamental'])} approach in {random.choice(concepts)}",
        f"The {random.choice(['main', 'primary', 'core', 'essential'])} idea behind this is",
        f"For optimal {random.choice(['performance', 'memory usage', 'readability', 'maintainability'])}, we",
        f"This handles edge cases like {random.choice(['empty input', 'duplicates', 'invalid data', 'boundary conditions', 'null values'])}",
        f"The {random.choice(['algorithm', 'function', 'implementation', 'solution'])} {random.choice(['scales well', 'handles large inputs', 'is production-ready', 'is well-tested', 'is easy to understand'])}",
    ]
    
    while len(words) < n_words:
        t = random.choice(templates)
        words.extend(t.split())
    
    return ' '.join(words[:n_words])


# ============================================================================
# Data generation
# ============================================================================

ALL_LANG_TEMPLATES = {
    'python': PYTHON_TEMPLATES,
    'javascript': JAVASCRIPT_TEMPLATES,
    'java': JAVA_TEMPLATES,
    'cpp': CPP_TEMPLATES,
    'rust': RUST_TEMPLATES,
    'go': GO_TEMPLATES,
    'sql': SQL_TEMPLATES,
}

ALL_NLP_TEMPLATES = NLP_TEMPLATES


def generate_sample(template_type='code'):
    """Generate a single training sample."""
    if template_type == 'nlp':
        template = random.choice(ALL_NLP_TEMPLATES)
        text = template()
        return {
            'text': text,
            'type': 'nlp',
            'task': random.choice(['explain', 'debug', 'test', 'qa', 'codegen']),
        }
    else:
        lang = random.choice(list(ALL_LANG_TEMPLATES.keys()))
        template = random.choice(ALL_LANG_TEMPLATES[lang])
        code = template()
        return {
            'code': code,
            'language': lang,
            'type': 'code',
        }


def generate_dataset(output_path, n_samples=100000, code_ratio=0.6, nlp_ratio=0.4):
    """Generate synthetic training dataset."""
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    with open(output_path, 'w') as f:
        for _ in tqdm(range(n_samples), desc="Generating data"):
            if random.random() < code_ratio:
                sample = generate_sample('code')
            else:
                sample = generate_sample('nlp')
            f.write(json.dumps(sample) + '\n')
    
    return output_path


def generate_multi_lang(output_dir, total_samples=500000):
    """Generate multi-language dataset split into shards."""
    os.makedirs(output_dir, exist_ok=True)
    
    shard_size = 10000
    n_shards = total_samples // shard_size
    
    for shard in range(n_shards):
        shard_path = os.path.join(output_dir, f'shard_{shard:04d}.jsonl')
        generate_dataset(shard_path, shard_size)
    
    print(f"Generated {n_shards} shards ({total_samples} samples) in {output_dir}")
    
    # Generate NLP only dataset for the NLP mixture training
    nlp_dir = os.path.join(output_dir, 'nlp')
    os.makedirs(nlp_dir, exist_ok=True)
    for shard in range(n_shards // 2):
        shard_path = os.path.join(nlp_dir, f'nlp_shard_{shard:04d}.jsonl')
        generate_dataset(shard_path, shard_size, code_ratio=0.0, nlp_ratio=1.0)
    
    print(f"Generated NLP dataset in {nlp_dir}")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', type=str, default='/FSI_Edge/data/train')
    parser.add_argument('--samples', type=int, default=500000)
    parser.add_argument('--multi-lang', action='store_true', default=True)
    args = parser.parse_args()
    
    generate_multi_lang(args.output, args.samples)