#!/usr/bin/env python3
"""
Save Models for Hugging Face Deployment

Run this cell AFTER training TCN, ResNet-1D, and the stacking ensemble.
This will save all necessary files for deployment.

Required variables in memory:
- tcn_results (with 'model' key or separate tcn model)
- resnet1d_results (with 'model' key or separate resnet model)  
- meta_model (trained sklearn meta-learner)
- data (with 'sequential' containing max_patches and hidden_dim)
"""

import torch
import pickle
import os
import shutil

# ==================== CONFIGURATION ====================

OUTPUT_DIR = '/kaggle/working/deploy_models'
os.makedirs(OUTPUT_DIR, exist_ok=True)

print("="*80)
print("💾 SAVING MODELS FOR HUGGING FACE DEPLOYMENT")
print("="*80)

# ==================== SAVE CONFIG ====================

print("\n📝 Saving configuration...")

config = {
    'max_patches': data['sequential']['max_patches'] if 'max_patches' in data['sequential'] else 256,
    'hidden_dim': data['sequential']['hidden_dim'] if 'hidden_dim' in data['sequential'] else 2048,
}

config_path = os.path.join(OUTPUT_DIR, 'config.pkl')
with open(config_path, 'wb') as f:
    pickle.dump(config, f)
print(f"✅ Config saved to {config_path}")
print(f"   • max_patches: {config['max_patches']}")
print(f"   • hidden_dim: {config['hidden_dim']}")

# ==================== SAVE TCN MODEL ====================

print("\n📦 Saving TCN model...")

# Check if TCN model state dict already exists
if os.path.exists('tcn_best.pth'):
    shutil.copy('tcn_best.pth', os.path.join(OUTPUT_DIR, 'tcn_best.pth'))
    print(f"✅ TCN model copied from tcn_best.pth")
else:
    print("⚠️  tcn_best.pth not found. Please save your TCN model:")
    print("   torch.save(tcn_model.state_dict(), 'tcn_best.pth')")

# ==================== SAVE RESNET-1D MODEL ====================

print("\n📦 Saving ResNet-1D model...")

# Check if ResNet model state dict already exists
if os.path.exists('resnet1d_best.pth'):
    shutil.copy('resnet1d_best.pth', os.path.join(OUTPUT_DIR, 'resnet1d_best.pth'))
    print(f"✅ ResNet-1D model copied from resnet1d_best.pth")
else:
    print("⚠️  resnet1d_best.pth not found. Please save your ResNet model:")
    print("   torch.save(resnet_model.state_dict(), 'resnet1d_best.pth')")

# ==================== SAVE META-LEARNER ====================

print("\n📦 Saving meta-learner...")

try:
    meta_path = os.path.join(OUTPUT_DIR, 'meta_model.pkl')
    with open(meta_path, 'wb') as f:
        pickle.dump(meta_model, f)
    print(f"✅ Meta-learner saved to {meta_path}")
except NameError:
    print("⚠️  meta_model not found. Please ensure the stacking ensemble has been trained.")

# ==================== CREATE ZIP ====================

print("\n📦 Creating deployment package...")

import zipfile

zip_path = '/kaggle/working/huggingface_deploy.zip'
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
    for root, dirs, files in os.walk(OUTPUT_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            arcname = os.path.relpath(file_path, OUTPUT_DIR)
            zipf.write(file_path, arcname)

zip_size_mb = os.path.getsize(zip_path) / (1024**2)
print(f"✅ Deployment package created: {zip_path}")
print(f"   Size: {zip_size_mb:.2f} MB")

# ==================== SUMMARY ====================

print("\n" + "="*80)
print("📋 DEPLOYMENT CHECKLIST")
print("="*80)

print("\n✅ Files saved:")
for f in os.listdir(OUTPUT_DIR):
    size_kb = os.path.getsize(os.path.join(OUTPUT_DIR, f)) / 1024
    print(f"   • {f} ({size_kb:.1f} KB)")

print("\n📝 Next Steps:")
print("1. Download huggingface_deploy.zip from Kaggle")
print("2. Create a new Hugging Face Space:")
print("   huggingface-cli repo create your-username/ai-image-detector --type space --space_sdk gradio")
print("3. Clone and add files:")
print("   git clone https://huggingface.co/spaces/your-username/ai-image-detector")
print("   cd ai-image-detector")
print("   # Extract model files to models/ folder")
print("   # Copy app.py, requirements.txt, README.md")
print("4. Push to Hugging Face:")
print("   git add .")
print("   git commit -m 'Add AI image detector'")
print("   git push")

print("\n💡 Alternative: Upload directly via Hugging Face web interface")
print("   Go to huggingface.co/new-space and upload files")

print("="*80)
print("✅ SAVE COMPLETE!")
print("="*80)

# Display download link
from IPython.display import FileLink
FileLink(zip_path)