"""Re-serialize FrankHall models so they can be loaded without __main__ context.

Run this ONCE after downloading models from GDrive:
    python reserialize_models.py

This patches the pickle module to redirect __main__.FrankHallOrdinal
to models.frank_hall.FrankHallOrdinal, then re-saves each model.
"""

import sys
import os
import types
from pathlib import Path

# Add parent to path so models.frank_hall is importable
sys.path.insert(0, str(Path(__file__).parent))

from models.frank_hall import FrankHallOrdinal

# Patch: when unpickling, redirect __main__.FrankHallOrdinal
# to our local class
import __main__
__main__.FrankHallOrdinal = FrankHallOrdinal

import joblib

SAVED_DIR = Path(__file__).parent / "saved_models"

targets = [
    "articulation_ordinal", "pause_freq_ordinal", "pause_dur_ordinal",
    "cognitive_load_ordinal", "utterance_constraints_ordinal",
]

print("Re-serializing FrankHall models...")
for t in targets:
    model_path = SAVED_DIR / f"ordinal_{t}.joblib"
    if not model_path.exists():
        print(f"  SKIP: {model_path.name} (not found)")
        continue

    try:
        model = joblib.load(model_path)
        # Re-save — now pickle will reference models.frank_hall.FrankHallOrdinal
        # instead of __main__.FrankHallOrdinal
        joblib.dump(model, model_path)
        print(f"  OK: {model_path.name} ({type(model).__name__})")
    except Exception as e:
        print(f"  ERROR: {model_path.name}: {e}")

# Also patch inference.py to import FrankHallOrdinal before loading
print("\nDone. Models can now be loaded without the training notebook context.")
print("Test with: python -c \"from models.inference import predict; print('OK')\"")