#!/usr/bin/env python3
"""
Convert trained .pth files to .safetensors format and place them in the
correct directory structure.

The source .pth files are read from the ZX2 external drive (read-only).
Output .safetensors files are written to the HuggingFace repo structure.

Usage:
    python3 convert_weights.py

Requires:
    pip install safetensors torch
"""

import os
import sys
import json
import torch
from safetensors.torch import save_file


# ── Path configuration ──
ZX2_BASE = '/Volumes/ZX2 1TB/se-alexnet/ClassifiedWithCondition'
REPO_ROOT = os.path.dirname(os.path.abspath(__file__))


# ── Mapping: (model_type, pretrain, squeeze) → ZX2 source path ──
# Format: (config_dest_relpath, source_pth_path)
WEIGHT_MAPPINGS = [
    # ── Raw AlexNet ──
    ('alexnet/facebased',
     f'{ZX2_BASE}/RawAlexNet/FaceBased/AlexNet.pth'),
    ('alexnet/objectbased',
     f'{ZX2_BASE}/RawAlexNet/ObjectBased/AlexNet.pth'),

    # ── VGG16 ──
    ('vgg16/facebased',
     f'{ZX2_BASE}/VGG16/FaceBased/vgg16Net.pth'),
    ('vgg16/objectbased',
     f'{ZX2_BASE}/VGG16/ObjectBased/vgg16Net.pth'),
]

# ── SE Location 1 ──
for pretrain in ['FaceBased', 'ObjectBased']:
    dest_pretrain = pretrain.lower()
    for r in [2, 4, 8, 16, 32]:
        dest = f'se-location1/{dest_pretrain}/squeeze-{r}'
        src = f'{ZX2_BASE}/SELocate-1/{pretrain}/squeeze-{r}/AlexNet.pth'
        WEIGHT_MAPPINGS.append((dest, src))

# ── SE Location 2 ──
for pretrain in ['FaceBased', 'ObjectBased']:
    dest_pretrain = pretrain.lower()
    for r in [2, 4, 8, 16, 32]:
        dest = f'se-location2/{dest_pretrain}/squeeze-{r}'
        src = f'{ZX2_BASE}/SELocate-2/{pretrain}/squeeze-{r}/AlexNet.pth'
        WEIGHT_MAPPINGS.append((dest, src))

# ── SE Location 3 ──
for pretrain in ['FaceBased', 'ObjectBased']:
    dest_pretrain = pretrain.lower()
    for r in [2, 4, 8, 16, 32]:
        dest = f'se-location3/{dest_pretrain}/squeeze-{r}'
        src = f'{ZX2_BASE}/SELocate-3-best/{pretrain}/squeeze-{r}/AlexNet.pth'
        WEIGHT_MAPPINGS.append((dest, src))


def convert_single(src_path, dest_dir, dest_name='model.safetensors'):
    """Convert a single .pth file to .safetensors."""
    dest_path = os.path.join(REPO_ROOT, dest_dir, dest_name)

    if os.path.exists(dest_path):
        print(f'  ⏭  {dest_dir}/model.safetensors already exists, skipping')
        return

    os.makedirs(os.path.join(REPO_ROOT, dest_dir), exist_ok=True)

    print(f'  Loading: {os.path.basename(src_path)} ...', end=' ', flush=True)
    try:
        state_dict = torch.load(src_path, map_location='cpu', weights_only=True)
    except Exception as e:
        print(f'ERROR: {e}')
        return

    # Unwrap checkpoint if needed
    if 'model' in state_dict:
        state_dict = state_dict['model']

    # Check with config to verify compatibility
    config_path = os.path.join(REPO_ROOT, dest_dir, 'config.json')
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = json.load(f)
        # Optional: verify last layer dims match num_classes
        last_linear_keys = [k for k in state_dict if 'weight' in k
                           and len(state_dict[k].shape) == 2]
        if last_linear_keys:
            last_w = state_dict[last_linear_keys[-1]]
            expected = config.get('num_classes')
            if expected and last_w.shape[0] != expected:
                print(f'\n  ⚠️  num_classes mismatch: state_dict={last_w.shape[0]}, '
                      f'config={expected}')

    print(f'saving .safetensors ({len(state_dict)} keys) ...', end=' ', flush=True)
    save_file(state_dict, dest_path)
    size_mb = os.path.getsize(dest_path) / (1024 * 1024)
    print(f'✓ ({size_mb:.0f} MB)')


def main():
    print('=' * 60)
    print('SE-AlexNet: .pth → .safetensors Conversion')
    print('=' * 60)
    print(f'Source: {ZX2_BASE}')
    print(f'Target: {REPO_ROOT}')
    print(f'Total:  {len(WEIGHT_MAPPINGS)} models to convert')
    print()

    success = 0
    skipped = 0
    failed = 0

    for dest, src in WEIGHT_MAPPINGS:
        if not os.path.exists(src):
            print(f'  ✗ MISSING: {src}')
            failed += 1
            continue
        try:
            convert_single(src, dest)
            success += 1
        except Exception as e:
            print(f'  ✗ FAILED: {e}')
            failed += 1

    print()
    print(f'Done. {success} converted, {skipped} skipped, {failed} failed.')


if __name__ == '__main__':
    main()