import subprocess
import os
import json
from datasets import load_dataset

os.makedirs('/tmp/glm-project/multi_data', exist_ok=True)
os.makedirs('/tmp/glm-project/multi_adapters', exist_ok=True)

dataset_names = [
    'TeichAI/glm-4.7-2000x',
    'TeichAI/claude-4.5-opus-high-reasoning-250x',
    'TeichAI/gpt-5.1-high-reasoning-1000x',
    'TeichAI/gemini-3-pro-preview-high-reasoning-1000x',
    'TeichAI/gpt-5.2-high-reasoning-250x'
]

train_data = []

print('Downloading and formatting multi-teacher datasets...')
for name in dataset_names:
    print(f"Loading {name}...")
    try:
        ds = load_dataset(name, split='train')
        count = 0
        for item in ds:
            if 'messages' in item:
                text = ""
                for m in item['messages']:
                    text += f"{m['role']}: {m['content']}\n"
                train_data.append({'text': text})
                count += 1
            elif 'prompt' in item and 'response' in item:
                train_data.append({'text': f"User: {item['prompt']}\nAssistant: {item['response']}"})
                count += 1
            else:
                keys = list(item.keys())
                if len(keys) >= 2:
                    train_data.append({'text': f"User: {item[keys[0]]}\nAssistant: {item[keys[1]]}"})
                    count += 1
        print(f"  -> Added {count} examples.")
    except Exception as e:
        print(f"Failed to load {name}: {e}")

import random
random.seed(42)
random.shuffle(train_data)

print(f"\nTotal mixed training examples: {len(train_data)}")

with open('/tmp/glm-project/multi_data/train.jsonl', 'w') as f:
    for item in train_data:
        f.write(json.dumps(item) + '\n')

with open('/tmp/glm-project/multi_data/valid.jsonl', 'w') as f:
    for i in range(min(10, len(train_data))):
        f.write(json.dumps(train_data[i]) + '\n')

print('\nStarting Multi-Teacher LoRA fine-tuning (Batch Size: 3, Context: 32k)...')
lora_cmd = [
    '/tmp/glm-project/venv/bin/mlx_lm.lora',
    '--model', '/tmp/GLM-4.7-Flash-REAP-23B-A3B',
    '--train',
    '--data', '/tmp/glm-project/multi_data',
    '--iters', '4512', 
    '--batch-size', '1',
    '--num-layers', '8',
    '--max-seq-length', '32768',
    '--grad-checkpoint',
    '--adapter-path', '/tmp/glm-project/multi_adapters'
]
subprocess.run(lora_cmd, check=True)

print('\nFusing multi-teacher adapters...')
fuse_cmd = [
    '/tmp/glm-project/venv/bin/mlx_lm.fuse',
    '--model', '/tmp/GLM-4.7-Flash-REAP-23B-A3B',
    '--adapter-path', '/tmp/glm-project/multi_adapters',
    '--save-path', '/tmp/glm-project/multi_distilled'
]
subprocess.run(fuse_cmd, check=True)

print('\nQuantizing to Q6...')
quant_cmd = [
    '/tmp/glm-project/venv/bin/mlx_lm.convert',
    '--hf-path', '/tmp/glm-project/multi_distilled',
    '--mlx-path', '/tmp/GLM-4.7-Flash-23B-MultiTeacher-MLX-6bit',
    '-q', '--q-bits', '6', '--q-group-size', '64'
]
subprocess.run(quant_cmd, check=True)

print('\nMulti-teacher pipeline complete!')