import tiktoken
import torch
import biggerbrain as biggerbrain
import ai_extras as A_E
from safetensors.torch import save_file, load_file

model = biggerbrain.initmodel("cpu")
model = model._orig_mod if hasattr(model, '_orig_mod') else model

while True:
    user_input = input("You: ")
    user_input = user_input.lower()
    if user_input.lower() in {"exit", "quit"}:
        print("Exiting the app.")
        break
    
    elif user_input == "print model":
        model.print_parameter_breakdown(model)
    elif user_input == "cpu":
        model.to("cpu")
        print("Model moved to CPU.")
    elif user_input == "train":
        if 'pretrain_ds' not in locals():
            pretrain_ds = A_E.StreamDataset(bin_file="C:\\AIs\\biggerbrain2_135m\\total_dataset.bin", seq_len=model.sequencelength)
            print(f"Dataset loaded: {len(pretrain_ds)} samples") 
        model.trainingloop(data=pretrain_ds, epochs=10, lr=3e-4, batchsize=4, accumulation_steps=32, warmup_steps=5000)#train
    elif user_input.lower() == "load":
        weights = load_file("model.safetensors")
        model.load_state_dict(weights, strict=False)
        print("Weights loaded!")
    elif user_input.lower() == "check1":
        print("alpha pre:")
        print(model.get_parameter("alpha_pre").item())
        print("alpha loop:")
        print(model.get_parameter("alpha_loop").item())
        print("alpha post:")
        print(model.get_parameter("alpha_post").item())
        print("alpha mem:")
        print(model.get_parameter("alpha_mem").item())
        model.debugprints = True
        model.forward_training(1)
    elif user_input.lower() == "check2":
        print("\n--- Model Internal Stats ---")
        # Print all alpha parameters dynamically
        for name, param in model.named_parameters():
            if 'alpha' in name:
                # Using .item() to get the actual number instead of the tensor object
                print(f"{name}: {param.item():.6f}")
        
        # Check the Engram Gate (mem_gate)
        if hasattr(model, 'mem_gate'):
            # We look at the bias because that's what controls the initial "openness"
            gate_bias = model.mem_gate.bias.item()
            # Calculate the actual sigmoid value to see the % it's open
            gate_open_pct = torch.sigmoid(torch.tensor(gate_bias)).item() * 100
            print(f"mem_gate bias: {gate_bias:.6f} ({gate_open_pct:.2f}% open)")
        
        print("----------------------------\n")
    elif user_input.lower() == "check3":
        for name, param in model.named_parameters():
            if 'alpha' in name:
                print(f"{name} | Requires Grad: {param.requires_grad} | Device: {param.device}")
    elif user_input == "debug":
        model.debugprints = True
    else:
        biggerbrain.think(prompt=user_input, model=model, max_length=10, iter=3, top_k=10, temperature=1.0)