#!/usr/bin/env python3
"""Initialize MONIKA with aggressive scaling for RTX 5060 16GB."""

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

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

# Aggressive scaling for RTX 5060 16GB VRAM
config = TrainingConfig()
config.embed_dim = 768  # 6x baseline
config.sequence_length = 512  # 4x baseline  
config.checkpoint_path = "storage/proto_lm/massive_model.pt"
config.device = "cuda"
config.learning_rate = 1e-4
config.vocab_growth_chunk = 128  # Larger chunks
config.vocab_growth_headroom = 64

print("=== Initializing MASSIVE MONIKA Model ===")
print(f"Embed dim: {config.embed_dim}")
print(f"Sequence length: {config.sequence_length}")
print(f"SASS layers: 16 (hardcoded)")
print(f"Device: {config.device}")

model = ProtoLanguageModel(config)

# Calculate parameters
embed_params = model.embed.weight.numel()
output_params = model.output.weight.numel() + model.output.bias.numel()
core_params = sum(p.numel() for n, p in model.named_parameters() if 'core' in n)
total_params = sum(p.numel() for p in model.parameters())

print(f"\n=== Model Architecture ===")
print(f"Embedding: {embed_params:,} params")
print(f"SASS Core: {core_params:,} params")
print(f"Output: {output_params:,} params")
print(f"TOTAL: {total_params:,} params ({total_params/1e6:.1f}M)")

# Test first step
print("\nTesting training...")
loss = model.training_step("The quick brown fox jumps over the lazy dog.")
print(f"First step loss: {loss:.4f}")

# Save
checkpoint = model.save_checkpoint(reason="massive-model-init", tags=["768-dim", "16-layers", "aggressive-scale"])
print(f"\nSaved: {checkpoint}")
print("\n=== Ready for bulk training ===")