#!/usr/bin/env python3
"""
QUIXZOOM Custom YOLO Training Pipeline
"""

import argparse
import os
import yaml
from pathlib import Path

def create_dataset_config(data_dir, output_file='dataset.yaml'):
    config = {
        'path': os.path.abspath(data_dir),
        'train': 'images/train',
        'val': 'images/val',
        'test': 'images/test',
        'names': {
            0: 'street_lamp',
            1: 'traffic_sign',
            2: 'tree',
            3: 'manhole',
            4: 'utility_box',
            5: 'bench',
        },
        'nc': 6,
    }
    
    with open(output_file, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    
    print(f"[TRAIN] Dataset config saved to {output_file}")
    return output_file

def prepare_dataset(data_dir):
    dirs = [
        'images/train', 'images/val', 'images/test',
        'labels/train', 'labels/val', 'labels/test',
    ]
    
    for d in dirs:
        os.makedirs(os.path.join(data_dir, d), exist_ok=True)
    
    print(f"[TRAIN] Dataset structure prepared in {data_dir}")

def main():
    parser = argparse.ArgumentParser(description='QUIXZOOM YOLO Training')
    parser.add_argument('--data', type=str, required=True, help='Dataset directory')
    parser.add_argument('--epochs', type=int, default=100, help='Training epochs')
    
    args = parser.parse_args()
    
    prepare_dataset(args.data)
    config_file = create_dataset_config(args.data)
    
    print(f"[TRAIN] Ready to train with config: {config_file}")
    print(f"[TRAIN] Run: yolo detect train data={config_file} epochs={args.epochs}")

if __name__ == '__main__':
    main()