#!/usr/bin/env python3
"""SWA: son N intermediate ckpt'in ağırlık ortalaması.

Usage: python swa_average.py runs/detect/.../yolo_fold0/weights/  [N=5]
"""
import sys, glob, copy
from pathlib import Path
import torch

def main():
    if len(sys.argv) < 2:
        print("usage: swa_average.py <weights_dir> [N=5]"); sys.exit(1)
    wdir = Path(sys.argv[1])
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 5

    # Collect epoch ckpts (epoch{N}.pt or last.pt)
    ckpts = sorted(wdir.glob('epoch*.pt'))
    if len(ckpts) < n:
        print(f"only {len(ckpts)} epoch ckpts; using all + last.pt");
        last = wdir / 'last.pt'
        if last.exists() and last not in ckpts:
            ckpts.append(last)
    ckpts = ckpts[-n:]
    print(f"Averaging {len(ckpts)}: {[c.name for c in ckpts]}")

    avg_state = None; avg_meta = None
    for i, c in enumerate(ckpts):
        d = torch.load(c, map_location='cpu')
        sd = d['model'].state_dict() if hasattr(d['model'], 'state_dict') else d['model']
        if i == 0:
            avg_state = {k: v.float().clone() for k, v in sd.items()}
            avg_meta = d
        else:
            for k in avg_state:
                if k in sd and avg_state[k].shape == sd[k].shape:
                    avg_state[k] += sd[k].float()
    for k in avg_state:
        avg_state[k] /= len(ckpts)

    # Replace model state, save
    if hasattr(avg_meta['model'], 'load_state_dict'):
        avg_meta['model'].load_state_dict({k: v.to(next(avg_meta['model'].parameters()).dtype) for k,v in avg_state.items()})
    else:
        avg_meta['model'] = avg_state

    out = wdir / 'swa.pt'
    torch.save(avg_meta, out)
    print(f"SWA -> {out}")

if __name__ == '__main__':
    main()