"""
Full Training Pipeline
Combine synthetic + quiXzoom data for production AI models
"""

import sys
sys.path.insert(0, '/home/bernt/.openclaw/workspace/iom')

from ai_pipeline.data_collection import DataCollector, SyntheticDataGenerator
from ai_pipeline.quixzoom_data_ingestion import QuixzoomDataIngestion
from ai_pipeline.training_pipeline import TrainingPipeline, TrainingConfig

import os
import shutil


def combine_datasets():
    """Combine synthetic and quiXzoom datasets"""
    print("=== Combining Datasets ===\n")
    
    # Create combined dataset
    combined_path = "/tmp/iom_combined_dataset"
    os.makedirs(combined_path, exist_ok=True)
    
    # 1. Generate synthetic data
    print("1. Generating synthetic data...")
    synthetic_collector = DataCollector(f"{combined_path}/synthetic")
    generator = SyntheticDataGenerator()
    
    scenes = ["street_view", "building_facade", "bridge", "road", "sidewalk", "park"]
    
    for i in range(100):
        scene = scenes[i % len(scenes)]
        synthetic = generator.generate_synthetic_image(scene, num_defects=3)
        
        synthetic_collector.add_annotation(
            image_id=f"syn_{i:04d}",
            filename=f"syn_{i:04d}.jpg",
            width=synthetic["width"],
            height=synthetic["height"],
            objects=synthetic["objects"],
            scene_type=scene,
            split="train" if i < 80 else "val" if i < 90 else "test"
        )
    
    # 2. Ingest quiXzoom data
    print("2. Ingesting quiXzoom data...")
    quixzoom_ingestion = QuixzoomDataIngestion(f"{combined_path}/quixzoom")
    quixzoom_results = quixzoom_ingestion.run_full_pipeline()
    
    # 3. Merge datasets
    print("3. Merging datasets...")
    
    merged_path = f"{combined_path}/merged"
    os.makedirs(f"{merged_path}/images/train", exist_ok=True)
    os.makedirs(f"{merged_path}/images/val", exist_ok=True)
    os.makedirs(f"{merged_path}/images/test", exist_ok=True)
    os.makedirs(f"{merged_path}/labels/train", exist_ok=True)
    os.makedirs(f"{merged_path}/labels/val", exist_ok=True)
    os.makedirs(f"{merged_path}/labels/test", exist_ok=True)
    
    # Copy synthetic data
    for split in ["train", "val", "test"]:
        src_img = f"{combined_path}/synthetic/images/{split}"
        src_lbl = f"{combined_path}/synthetic/labels/{split}"
        dst_img = f"{merged_path}/images/{split}"
        dst_lbl = f"{merged_path}/labels/{split}"
        
        if os.path.exists(src_img):
            for f in os.listdir(src_img):
                shutil.copy(f"{src_img}/{f}", dst_img)
        if os.path.exists(src_lbl):
            for f in os.listdir(src_lbl):
                shutil.copy(f"{src_lbl}/{f}", dst_lbl)
    
    # Copy quiXzoom data
    for split in ["train", "val", "test"]:
        src = f"{combined_path}/quixzoom/{split}"
        dst = f"{merged_path}/labels/{split}"
        if os.path.exists(src):
            for f in os.listdir(src):
                shutil.copy(f"{src}/{f}", dst)
    
    # Count merged data
    train_count = len(os.listdir(f"{merged_path}/labels/train"))
    val_count = len(os.listdir(f"{merged_path}/labels/val"))
    test_count = len(os.listdir(f"{merged_path}/labels/test"))
    
    print(f"\nMerged dataset:")
    print(f"  Train: {train_count}")
    print(f"  Val: {val_count}")
    print(f"  Test: {test_count}")
    print(f"  Total: {train_count + val_count + test_count}")
    
    return merged_path


def train_production_model(dataset_path: str):
    """Train production-ready model"""
    print("\n=== Training Production Model ===\n")
    
    from ultralytics import YOLO
    
    # Create data.yaml
    import json
    data_yaml = {
        "path": dataset_path,
        "train": "images/train",
        "val": "images/val",
        "test": "images/test",
        "nc": 10,
        "names": [
            "street_light", "traffic_sign", "bench", "trash_can",
            "sidewalk", "road", "building", "bridge", "tree", "graffiti"
        ]
    }
    
    with open(f"{dataset_path}/data.yaml", "w") as f:
        json.dump(data_yaml, f, indent=2)
    
    # Load model
    model = YOLO("yolov8n.pt")
    
    # Train with better settings for production
    print("Training production model...")
    results = model.train(
        data=f"{dataset_path}/data.yaml",
        epochs=50,  # More epochs for production
        batch=16,
        imgsz=640,
        device="cpu",  # Change to "0" for GPU
        project="/tmp/iom_production_models",
        name="yolo_infrastructure_v1",
        exist_ok=True,
        patience=10,  # Early stopping
        save=True,
        pretrained=True
    )
    
    print(f"Training complete!")
    
    return model


def export_production_model(model):
    """Export model to all production formats"""
    print("\n=== Exporting Production Model ===\n")
    
    formats = ["onnx", "torchscript", "openvino", "engine"]
    
    for fmt in formats:
        try:
            print(f"Exporting to {fmt}...")
            model.export(format=fmt)
            print(f"  ✓ {fmt} export successful")
        except Exception as e:
            print(f"  ✗ {fmt} export failed: {e}")
    
    print("\nExport complete!")


def main():
    """Main production training pipeline"""
    print("=" * 60)
    print("PRODUCTION AI TRAINING PIPELINE")
    print("=" * 60)
    
    # 1. Combine datasets
    dataset_path = combine_datasets()
    
    # 2. Train model
    model = train_production_model(dataset_path)
    
    # 3. Export
    export_production_model(model)
    
    print("\n" + "=" * 60)
    print("PRODUCTION TRAINING COMPLETE")
    print("=" * 60)
    print("\nModels saved to: /tmp/iom_production_models/")
    print("\nNext steps:")
    print("1. Validate model on test set")
    print("2. Deploy to production API")
    print("3. Monitor performance")
    print("4. Collect more data and retrain")


if __name__ == '__main__':
    main()