#!/usr/bin/env python3
"""Test that checkpoint loading properly reinitializes SASS core."""

import sys
sys.path.insert(0, 'C:\\MONIKA')

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

print("=== Testing Checkpoint Load Fix ===")

# Load model from checkpoint
config = TrainingConfig()
config.checkpoint_path = "storage/proto_lm/checkpoint.pt"
model = ProtoLanguageModel(config)

print(f"\nInitial state:")
print(f"  Step: {model.step}")
print(f"  Vocab size: {model.vocab.size()}")
print(f"  Embed capacity: {model.embed.num_embeddings}")
print(f"  Output capacity: {model.output.out_features}")

# Try a training step
try:
    loss = model.training_step("Testing checkpoint fix")
    print(f"\n✓ Training step succeeded, loss: {loss:.4f}")
except Exception as e:
    print(f"\n✗ Training step failed: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Try another training step to be sure
try:
    loss = model.training_step("Second test sentence")
    print(f"✓ Second training step succeeded, loss: {loss:.4f}")
except Exception as e:
    print(f"✗ Second training step failed: {e}")
    sys.exit(1)

# Try generation
try:
    response = model.sample("Hello", max_tokens=20)
    print(f"\n✓ Generation succeeded: {response[:50]}...")
except Exception as e:
    print(f"\n✗ Generation failed: {e}")
    sys.exit(1)

print("\n=== All tests passed! ===")