#!/usr/bin/env python3
"""
Combine all labeled datasets and retrain LoRA models.
"""

import json
import random
from pathlib import Path

# Load all labeled data
all_data = []

for fname in [
    "labeled_verified.jsonl",
    "labeled_expanded.jsonl", 
    "labeled_large.jsonl",
    "labeled_output.jsonl",
    "labeled_real_world.jsonl",
]:
    path = Path(f"data/{fname}")
    if path.exists():
        with open(path) as f:
            for line in f:
                try:
                    item = json.loads(line.strip())
                    # Normalize format
                    labels = item.get("labels", {})
                    text = item.get("text", item.get("raw_text", ""))
                    if text and labels.get("sentiment"):
                        all_data.append({
                            "text": text,
                            "sentiment": labels["sentiment"],
                            "event_type": labels.get("event_type", "unknown"),
                            "entities": labels.get("entities", []),
                        })
                except Exception as e:
                    print(f"Error in {fname}: {e}")

# Deduplicate by text hash
seen = set()
unique = []
for d in all_data:
    h = hash(d["text"][:200])
    if h not in seen:
        seen.add(h)
        unique.append(d)

print(f"Total unique samples: {len(unique)}")

# Split
random.shuffle(unique)
split = int(0.9 * len(unique))
train_data = unique[:split]
val_data = unique[split:]

print(f"Train: {len(train_data)} | Val: {len(val_data)}")

# Save combined dataset
with open("data/combined_train.jsonl", "w") as f:
    for d in train_data:
        f.write(json.dumps(d) + "\n")

with open("data/combined_val.jsonl", "w") as f:
    for d in val_data:
        f.write(json.dumps(d) + "\n")

print("Saved combined datasets")