"""Model Soup: average weights of multiple seed models → single model.
Zero inference cost vs single model, near-ensemble performance.
"""
import sys, os
import copy
import numpy as np

PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)

import torch
from ultralytics import YOLO
from tqdm import tqdm


def soup_models(model_paths, output_path):
    """Average weights of multiple YOLO models."""
    models = [YOLO(p) for p in model_paths]
    base = models[0]

    # Collect state dicts
    state_dicts = []
    for m in models:
        sd = m.model.state_dict()
        state_dicts.append(sd)

    # Average
    avg_sd = {}
    for key in state_dicts[0].keys():
        if any(kw in key.lower() for kw in ['num_batches_tracked', 'running_mean', 'running_var']):
            # Use first model's BN stats
            avg_sd[key] = state_dicts[0][key].clone()
        elif 'bn' in key.lower() or 'batch_norm' in key.lower():
            # Average BN stats
            avg_sd[key] = sum(sd[key] for sd in state_dicts) / len(state_dicts)
        else:
            avg_sd[key] = sum(sd[key] for sd in state_dicts) / len(state_dicts)

    base.model.load_state_dict(avg_sd)
    base.model.eval()
    base.save(output_path)
    print(f'Model Soup saved to {output_path}')
    return output_path


def evaluate_model(model_path, name):
    """Quick eval on val set."""
    model = YOLO(model_path)
    results = model.val(data='Data/Detection_dataset/dataset.yaml', imgsz=1536, batch=2, verbose=False)
    return float(results.box.map50), float(results.box.map)


def main():
    model_paths = [
        'runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt',
        'runs/detect/Detection_experiments/v12_seed_42/weights/best.pt',
        'runs/detect/Detection_experiments/v12_seed_123/weights/best.pt',
    ]

    # Verify all exist
    valid = [p for p in model_paths if os.path.exists(p)]
    print(f'Models available: {len(valid)}/{len(model_paths)}')
    for p in valid:
        print(f'  {p}')

    # Eval individual models
    print(f'\n{"="*50}')
    print('Individual models:')
    singles = {}
    for p in valid:
        name = p.split('/')[-3]
        m50, m5095 = evaluate_model(p, name)
        singles[name] = (m50, m5095)
        print(f'  {name}: mAP50={m50:.4f}, mAP50-95={m5095:.4f}')

    # Model Soup
    print(f'\n{"="*50}')
    print('Model Soup experiments:')

    # Try all combinations
    import itertools
    for r in range(2, len(valid) + 1):
        for combo in itertools.combinations(range(len(valid)), r):
            paths = [valid[i] for i in combo]
            names = [p.split('/')[-3] for p in paths]
            label = ' + '.join(n.split('_')[0] + '_' + n.split('_')[-1][:4] for n in names)
            output = f'runs/soup/soup_{len(paths)}models.pt'
            os.makedirs('runs/soup', exist_ok=True)
            soup_models(paths, output)
            m50, m5095 = evaluate_model(output, label)
            delta = m5095 - singles[names[0]][1]
            print(f'  {label}: mAP50={m50:.4f}, mAP50-95={m5095:.4f} (d={delta:+.4f} vs {names[0]})')

    # WBF baseline for comparison
    print(f'\n{"="*50}')
    print('WBF Ensemble (for comparison):')
    print('  (Run eval_ensemble.py for full WBF results)')


if __name__ == '__main__':
    main()