#!/usr/bin/env python3
"""
Script để random pick 50% samples từ label_data.json với 3 seeds khác nhau
và lưu thành 3 files riêng biệt.
"""

import json
import random
import os
from typing import List, Dict, Any

def load_json_data(file_path: str) -> List[Dict[str, Any]]:
    """Load dữ liệu từ file JSON"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"✅ Đã load {len(data)} samples từ {file_path}")
        return data
    except FileNotFoundError:
        print(f"❌ Không tìm thấy file: {file_path}")
        return []
    except json.JSONDecodeError as e:
        print(f"❌ Lỗi decode JSON: {e}")
        return []

def random_sample_data(data: List[Dict[str, Any]], seed: int, sample_ratio: float = 0.5, sample_size = None) -> List[Dict[str, Any]]:
    """Random pick samples từ data với seed và tỷ lệ sampling cụ thể"""
    random.seed(seed)
    if sample_size is None:
        sample_size = int(len(data) * sample_ratio)
    sampled_data = random.sample(data, sample_size)
    print(f"🎲 Seed {seed}: Đã sample {len(sampled_data)}/{len(data)} samples ({sample_ratio*100:.1f}%)")
    return sampled_data

def save_json_data(data: List[Dict[str, Any]], file_path: str) -> bool:
    """Lưu dữ liệu vào file JSON"""
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"💾 Đã lưu {len(data)} samples vào {file_path}")
        return True
    except Exception as e:
        print(f"❌ Lỗi khi lưu file {file_path}: {e}")
        return False

def main():
    """Hàm chính"""
    # Đường dẫn file input
    input_file = "label_data.json"
    
    # Các seeds cần sử dụng
    seeds = [1997, 42, 2025]
    # seeds = [42]
    
    # Tỷ lệ sampling (50%)
    
    # 8, 9, 10
    sample_ratio = 0.5
    
    # 11,12,13
    sample_ratio = 0.3
    
    # 16
    # sample_ratio = 0.01
    
    # 15, 17, 18
    sample_size = 200
    
    print("🔄 Bắt đầu quá trình random sampling...")
    print("=" * 50)
    
    # Load dữ liệu gốc
    original_data = load_json_data(input_file)
    if not original_data:
        print("❌ Không thể load dữ liệu. Thoát chương trình.")
        return
    
    print(f"📊 Tổng số samples trong file gốc: {len(original_data)}")
    print(f"🎯 Sẽ sample {sample_ratio*100:.1f}% = {int(len(original_data) * sample_ratio)} samples")
    print("=" * 50)
    
    # Thực hiện sampling với từng seed
    for seed in seeds:
        print(f"\n🌱 Xử lý với seed: {seed}")
        
        # Random sample data
        sampled_data = random_sample_data(original_data, seed, sample_ratio, sample_size = sample_size)
        
        # Tạo tên file output
        output_file = f"label_data_sample_seed_{seed}.json"
        
        # Lưu file
        success = save_json_data(sampled_data, output_file)
        
        if success:
            print(f"✅ Hoàn thành sampling với seed {seed}")
        else:
            print(f"❌ Lỗi khi xử lý seed {seed}")
    
    print("\n" + "=" * 50)
    print("🎉 Hoàn thành toàn bộ quá trình!")
    print("\n📁 Các file đã được tạo:")
    for seed in seeds:
        output_file = f"016_no_fine_tune_label_data_sample_seed_{seed}.json"
        if os.path.exists(output_file):
            print(f"   ✅ {output_file}")
        else:
            print(f"   ❌ {output_file} (không tồn tại)")

if __name__ == "__main__":
    main()
