# File: src/clean_external_data_task1.py
import json
import os
from datasets import load_dataset
from tqdm import tqdm

def to_float_safe(x):
    """Chuyển đổi string sang float (xử lý cả '9' và '9.0')"""
    try:
        if x is None: return None
        val = float(x)
        if 0 <= val <= 9: return val
        return None
    except ValueError:
        return None

def parse_hai2131_dataset(dataset):
    """
    Parser chuẩn cho 'hai2131/IELTS-essays-task-1'.
    Kết hợp: Prompt + Image Description + Essay
    """
    print("Đang xử lý dataset 'hai2131'...")
    cleaned = []
    bad_examples = 0
    
    for item in tqdm(dataset, desc="Parsing hai2131"):
        try:
            # 1. Lấy thông tin đầu vào
            prompt = item.get("subject") or ""
            # QUAN TRỌNG: Lấy mô tả ảnh
            img_desc = item.get("image_description") or "" 
            essay = item.get("content") or ""
            
            # 2. Tạo input text kết hợp (Prompt + Context + Essay)
            # Model sẽ đọc toàn bộ chuỗi này
            full_prompt_text = f"PROMPT: {prompt}\n\nIMAGE CONTEXT: {img_desc}"
            
            # 3. Lấy điểm số (Dataset này để điểm dạng string "9")
            scores = {
                "task_response": to_float_safe(item.get("task_response_score")),
                "coherence_cohesion": to_float_safe(item.get("coherence_cohesion_score")),
                "lexical_resource": to_float_safe(item.get("lexical_resource_score")),
                "grammatical_range": to_float_safe(item.get("grammatical_range_accuracy_score"))
            }

            # 4. Kiểm tra hợp lệ
            if essay and all(scores.values()):
                cleaned.append({
                    "prompt_text": full_prompt_text, # Input đặc biệt cho Task 1
                    "essay_text": essay,
                    "scores": scores
                })
            else:
                bad_examples += 1
                
        except Exception:
            bad_examples += 1
            
    print(f"hai2131: kept {len(cleaned)} samples, skipped {bad_examples}")
    return cleaned

def main():
    print("🚀 BẮT ĐẦU XỬ LÝ DATASET TASK 1 (hai2131)")
    
    # Cache dir để tránh tải lại nhiều lần
    cache_dir = "./.cache/huggingface_datasets_task1"
    
    try:
        # Tải dataset (lần này sẽ nhanh vì nó nhỏ gọn)
        dataset = load_dataset("hai2131/IELTS-essays-task-1", split="train", cache_dir=cache_dir)
        
        # Xử lý
        final_dataset = parse_hai2131_dataset(dataset)
        
        if not final_dataset:
            print("LỖI: Không có dữ liệu.")
            return

        # Lưu file
        output_dir = "data"
        if not os.path.exists(output_dir): os.makedirs(output_dir)
        output_path = os.path.join(output_dir, "dataset_for_scorer_task1.json")
        
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(final_dataset, f, ensure_ascii=False, indent=2)
            
        print(f"\n✅ HOÀN TẤT! Đã lưu {len(final_dataset)} mẫu vào '{output_path}'.")
        print("💡 Lưu ý: 'prompt_text' bây giờ chứa cả Đề bài VÀ Mô tả ảnh.")
        
    except Exception as e:
        print(f"❌ Lỗi: {e}")

if __name__ == "__main__":
    main()