import torch
from tokenizers import Tokenizer
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
import math
import time
import csv
import os
import re
from safetensors.torch import save_file, load_file

# Idea User assistant tokens
D_MODEL = 768
NUM_HEADS = 12
NUM_LAYERS = 16
DROPOUT = 0.2
MAX_SEQ_LENGTH = 96
LEARNING_RATE = 5e-5
NUM_EPOCHS = 0 # No more overfitting on initial training
INTERACTIVE_EPOCHS = 25
BATCH_SIZE = 32

TOP_K = 20
TOP_P = 0.2
REPETITION_PENALTY = 1.2
TEMPERATURE = 1

PENALTIES_FILE = 'Penalties.csv'


def top_k_top_p_sample(logits, top_k=50, top_p=0.9):
    logits = logits.clone()

    if top_k > 0:
        v, _ = torch.topk(logits, top_k)
        logits[logits < v[-1]] = -float("inf")

    probs = torch.softmax(logits, dim=-1)
    sorted_probs, sorted_idx = torch.sort(probs, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    cutoff = cumulative > top_p
    cutoff[1:] = cutoff[:-1].clone()
    cutoff[0] = False

    probs[sorted_idx[cutoff]] = 0.0
    probs /= probs.sum()

    return torch.multinomial(probs, 1).item()




def load_penalties():
    loaded_penalties = []
    if os.path.exists('Penalties.csv'):
        with open('Penalties.csv', 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            for row in reader:
                if len(row) >= 2:
                    loaded_penalties.append([row[0], float(row[1])])
                elif row:
                    loaded_penalties.append([row[0], 3.0])
    return loaded_penalties

def save_single_penalty(penalty_string):
    """Appends a new penalty to the CSV immediately."""
    with open(PENALTIES_FILE, 'a', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow([penalty_string])
        
SETTINGS_FILE = 'settings.csv'

def save_settings(penalty, temp):
    with open('settings.csv', 'a', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow([penalty, temp])
    print(f"[Console] Logged to settings history: Penalty={penalty}, Temp={temp}")

def load_settings():
    if os.path.exists('settings.csv'):
        with open('settings.csv', 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            last_row = None
            for row in reader:
                if row:
                    last_row = row
            if last_row:
                return float(last_row[0]), float(last_row[1])  
    return 1.0, 1.0 

penalties = load_penalties()


DATA_FILE = 'training_data.txt' 
DEFAULT_TRAINING_DATA = [
    "The quick brown fox jumps over the lazy dog.",
    "A glass of water is clear.",
    "The sun is shining bright and the sky is clear.",
    "The dog and the fox are friends forever.",
    "Coding with Pytorch and Transformers is fun and very rewarding.",
    "A computer runs very fast and never stops.",
    "The windows are big and bright.",
    "A green park is a great place to relax.",
    "The sky is clear today, with no clouds.",
    "The cat jumped over the fence.",
    "The plane has many windows.",
    "A big bird flew over the house.",
    "The plane smoothly landed on the concrete runway.",
    "The bird flew above the bustling city.",
    "The plane had an engine failure and had to land in the river.",
    "The Cessna 172 is a low-wing monoplane.",
    "The plane flew by the trees.",
    "The plane, almost out of fuel, finally landed at an airport.",
    "The angry bird flew away furiously.",
    "A plane is a machine that flies.",
    "The fast plane landed at the bright airport.",
    "The plane quickly landed on the runway.",
    "The letter A is part of the alphabet.",
    "The plane landed hardly on a grass runway in the forest.",
    "The clouds were floating above the ground.",
    "The plane was a very bright plane, it's livery glimmered in the night sky.",
    "The GPWS sounds on a plane are like Caution Terrain PULL up PULL up."
]


def load_data_from_csv(filepath):
    if not os.path.exists(filepath) or os.path.getsize(filepath) == 0:
        print("[Console] Text file not found or empty.")
        return list(DEFAULT_TRAINING_DATA)

    print(f"[Console] Loading training data from {filepath}...")

    with open(filepath, "r", encoding="utf-8") as f:
        raw = f.read()

    # Split on double newlines (conversation blocks)
    blocks = [b.strip() for b in raw.split("\n\n") if b.strip()]

    samples = []
    for b in blocks:
        samples.append(b + "\n<|endoftext|>")

    print(f"[Console] Loaded {len(samples)} training samples.")

    print("\n[Console] ===== TRAINING DATA (PREVIEW) =====\n")
    for i, s in enumerate(samples[:10], 1):
        preview = s.replace("\n", "\\n")
        print(f"{i:02d}: {preview[:300]}")
    print("\n[Console] ===================================\n")

    return samples




def save_data_to_csv(filepath, texts):
    print(f"[Console] Saving training data to {filepath}...")
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(texts[0])


class AobanTokenizer:
    def __init__(self, path="aoban_tokenizer.json"):
        self.tokenizer = Tokenizer.from_file(path)

        self.pad_id = self.tokenizer.token_to_id("<PAD>")
        self.eot_id = self.tokenizer.token_to_id("<|endoftext|>")

        if self.pad_id is None or self.eot_id is None:
            raise ValueError("Tokenizer missing special tokens")

    def encode(self, text, max_len=None):
        ids = self.tokenizer.encode(text).ids

        if max_len is not None:
            ids = ids[:max_len]
            ids += [self.pad_id] * (max_len - len(ids))

        return torch.tensor(ids, dtype=torch.long)

    def decode(self, indices):
        if isinstance(indices, torch.Tensor):
            indices = indices.tolist()
        return self.tokenizer.decode(indices)

    @property
    def vocab_size(self):
        return self.tokenizer.get_vocab_size()


def collate_batch(batch, pad_id):
    xs, ys = zip(*batch)

    max_len = max(x.size(0) for x in xs)

    x_padded = []
    y_padded = []

    for x, y in zip(xs, ys):
        pad_len = max_len - x.size(0)

        x_padded.append(
            torch.cat([x, torch.full((pad_len,), pad_id)])
        )
        y_padded.append(
            torch.cat([y, torch.full((pad_len,), pad_id)])
        )

    return torch.stack(x_padded), torch.stack(y_padded)

class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_len):
        self.samples = []

        for text in texts:
            ids = tokenizer.encode(text).tolist()
            ids.append(tokenizer.eot_id)
            if len(ids) < 2:
                continue

            for i in range(0, len(ids) - 1, max_len):
                chunk = ids[i:i + max_len + 1]
                if len(chunk) < 2:
                    continue
                self.samples.append(chunk)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        chunk = self.samples[idx]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y



class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(1), :].transpose(0, 1)

class TransformerLanguageModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers, dropout, max_len):
        super(TransformerLanguageModel, self).__init__()
        self.model_type = 'Transformer'
        self.d_model = d_model
        self.vocab_size = vocab_size
        
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, max_len)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=d_model*4,
            dropout=dropout,
            batch_first=True,
            activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)

        
        self.fc_out = nn.Linear(d_model, vocab_size)
        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()
        self.fc_out.weight.data.uniform_(-initrange, initrange)

    def forward(self, src):
        src = self.embedding(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)

        seq_len = src.size(1)
        causal_mask = torch.triu(
            torch.full((seq_len, seq_len), float('-inf'), device=src.device),
            diagonal=1
        )

        output = self.transformer(src, mask=causal_mask)
        return self.fc_out(output)

def train_model(model, data_loader, optimizer, criterion, device, epochs):
    model.train()
    
    for epoch in range(1, epochs + 1):
        total_loss = 0.0
        
        for batch in data_loader:
            x_batch, y_batch = batch
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)
            
            optimizer.zero_grad()

            output = model(x_batch)

            output = output.view(-1, output.size(-1))
            y_batch = y_batch.view(-1)

            current_loss = criterion(output, y_batch)
            current_loss.backward()

            if current_loss.item() > 3:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
            else:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            torch.cuda.empty_cache()

            optimizer.step()
            total_loss += current_loss.item()
        
        avg_loss = total_loss / len(data_loader)
        
        if epochs > 10 and epoch % (epochs // 10) == 0:
            print(f"Epoch {epoch}/{epochs}, Average Loss: {avg_loss:.4f}")
        elif epochs > 0 and epochs <= 50 and epoch % 10 == 0:
            print(f"Epoch {epoch}/{epochs}, Average Loss: {avg_loss:.4f}")
    print(f"[Console] The {NUM_LAYERS} Layers have been updated.")

def generate_text(model, tokenizer, prompt, max_len, device,
                  top_k=40, top_p=0.9, penalty=1.2, temperature=0.8):

    model.eval()

    input_ids = tokenizer.encode(prompt).unsqueeze(0).to(device)
    prompt_len = input_ids.size(1)
    generated = input_ids[0].tolist()
    eos_id = tokenizer.eot_id

    for _ in range(max_len):

        src = input_ids[:, -MAX_SEQ_LENGTH:]

        with torch.no_grad():
            output = model(src)

        logits = output[:, -1, :].squeeze(0)
        logits[tokenizer.pad_id] = -float("inf")
        
        if _ == 0:
            logits[eos_id] = -float("inf")

        for idx in set(generated[-32:]):
            logits[idx] /= penalty

        logits /= temperature

        next_token = top_k_top_p_sample(
            logits,
            top_k=top_k,
            top_p=top_p
        )
        
        if next_token == eos_id:
            break

        generated.append(next_token)
        input_ids = torch.cat(
            [input_ids, torch.tensor([[next_token]], device=device)],
            dim=1
        )
        
    new_tokens = generated[prompt_len:]
    return tokenizer.decode(new_tokens).replace("<pad>", "").strip()




last_generated_text = None
last_user_prompt = None
current_tokenizer = None
current_model = None
device = torch.device("cuda")
live_data_updates = [] 
initial_training_texts = [] 

def initialize_or_retrain(initial_train=True, use_live_data=False, epochs=NUM_EPOCHS):
    global device, current_tokenizer, current_model, live_data_updates, initial_training_texts

    if initial_train:
        initial_training_texts = load_data_from_csv(DATA_FILE)
    
    training_data = list(initial_training_texts)

    if use_live_data:
        print(f"[Console] Retraining on {len(initial_training_texts)} base examples plus {len(live_data_updates)} new examples.")
        training_data.extend(live_data_updates)

    old_vocab_size = current_tokenizer.vocab_size if current_tokenizer else 0
    current_tokenizer = AobanTokenizer("aoban_tokenizer.json")
    new_vocab_size = current_tokenizer.vocab_size
    print(f"[Console] Vocab Size: {new_vocab_size} unique tokens.")

    if new_vocab_size != old_vocab_size or initial_train:
        current_model = TransformerLanguageModel(
            vocab_size=new_vocab_size, 
            d_model=D_MODEL, 
            nhead=NUM_HEADS, 
            num_layers=NUM_LAYERS, 
            dropout=DROPOUT,
            max_len=MAX_SEQ_LENGTH
        ).to(device)
    
    weights_path = "aoban_weights.safetensors"
    if os.path.exists(weights_path):
        checkpoint = load_file(weights_path, device=str(device))
        saved_vocab_size = checkpoint['embedding.weight'].shape[0]
        current_vocab_size = current_model.embedding.weight.shape[0]

        if saved_vocab_size != current_vocab_size:
            print(f"[Console] Updating Aoban-2.7-AST-M architecture from {saved_vocab_size} to {current_vocab_size} tokens...")
            new_state_dict = current_model.state_dict()
            for key, value in checkpoint.items():
                if key in new_state_dict:
                    if value.shape == new_state_dict[key].shape:
                        new_state_dict[key] = value
                    else:
                        new_state_dict[key][:saved_vocab_size] = value[:saved_vocab_size]
            current_model.load_state_dict(new_state_dict)
            current_model.to(device)
            print(f"[Console] Aoban-2.7-AST-M is ready on {device}.")
        else:
            current_model.load_state_dict(checkpoint)

    dataset = TextDataset(training_data, current_tokenizer, MAX_SEQ_LENGTH)
    data_loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        collate_fn=lambda batch: collate_batch(batch, current_tokenizer.pad_id)
    )

    
    optimizer = Adam(current_model.parameters(), lr=LEARNING_RATE)
    criterion = nn.CrossEntropyLoss(ignore_index=current_tokenizer.pad_id)

    print(f"\n[Console] Starting {epochs} epochs with {len(dataset)} examples...")
    train_model(current_model, data_loader, optimizer, criterion, device, epochs)

    if use_live_data:
        initial_training_texts = training_data 
        save_data_to_csv(DATA_FILE, initial_training_texts)
        live_data_updates = [] 

    temp_path = "aoban_weights_temp.safetensors"
    try:
        save_file(current_model.state_dict(), temp_path)
        if os.path.exists(weights_path):
            os.replace(temp_path, weights_path)
        else:
            os.rename(temp_path, weights_path)
        print("[Console] Aoban-2.7-AST-M state permanently saved to disk.")
    except Exception as e:
        print(f"[System Bug] Could not save Aoban-2.7-AST-M state: {e}")

def perform_rlhf_step(model, tokenizer, text, learning_rate):
    """Performs a single, high-intensity update on a specific sentence."""
    model.train()
    optimizer = Adam(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss(ignore_index=current_tokenizer.pad_id)
    
    tokens = tokenizer.encode(text, max_len=MAX_SEQ_LENGTH).to(device)
    x = tokens[:-1].unsqueeze(0)
    y = tokens[1:].unsqueeze(0)
    
    optimizer.zero_grad()
    output = model(x)
    loss = criterion(output.reshape(-1, output.size(-1)), y.reshape(-1))
    loss.backward()
    optimizer.step()
    return loss.item()

def interactive_mode():
    global live_data_updates, last_generated_text, last_user_prompt, penalties
    global REPETITION_PENALTY, TEMPERATURE 
    
    file_existed_before_run = os.path.exists(DATA_FILE)
    global device
    device = torch.device("cuda")
    if not torch.cuda.is_available():
        print("Required GPU not located, running on CPU instead.")
        device = torch.device("cpu")
    print(f"[Console] Using device: {device}")
    initialize_or_retrain(initial_train=True, use_live_data=False, epochs=NUM_EPOCHS)
    if not file_existed_before_run:
        print("\n[SYSTEM] CSV file was empty/missing. Forcing initial save of default knowledge...")
        save_data_to_csv(DATA_FILE, initial_training_texts)
        print("[SYSTEM] Default 27 sentences are now permanently written to training_data.csv.")
    
    print("\n" + "=" * 60)
    print("🤖 Console Information🤖")
    print("1. Type a phrase to generate text (max 10 words).")
    print("2. Use '!add [sentence]' to queue new training data.")
    print("3. Use '!accept' to add the model's last **full** sentence to the training queue.")
    print(f"4. Use '!retrain' to re-train the model on new data (runs for {INTERACTIVE_EPOCHS} epochs) **and save it**.")
    print(f"5. Use '!refine' to re-train on existing data (runs for {INTERACTIVE_EPOCHS} epochs) **without saving.**")
    print("6. Use '!penalty <value>' to regenerate with a different repetition penalty (higher = less repetition).")
    print("7. Type 'quit' or 'exit' to stop.")
    print("8. Type '!help' to see this message again.")
    print("9. Use '!instead [corrected text]' to replace the last output with a corrected version.")
    print("10. Use !endorse to endorse the model on its generated text.")
    print("11. Use !neurons to find active neurons.")
    print("=" * 60)
    find_emergent_neuron(current_model, current_tokenizer, "Hello How can I help you today?", "What can i do for you?")
    
    while True:
        try:
            user_input = input("You: ")
            
            if user_input.lower() in ['quit', 'exit']:
                break
            
            if user_input.lower().startswith('!add '):
                sentence = user_input[5:].strip()
                if sentence:
                    live_data_updates.append(sentence)
                    print(f"[Console] Added sentence to update queue: '{sentence}'")
                    print(f"[Console] Current update queue size: {len(live_data_updates)}. Type '!retrain' to apply and save changes.")
                last_generated_text = None
                last_user_prompt = None
                continue
            
            if user_input.lower().strip() == '!accept':
                if last_generated_text and last_user_prompt:
                    full_sentence_parts = [last_user_prompt.strip(), last_generated_text.strip()]
                    sentence_to_add = " ".join(full_sentence_parts)
                    sentence_to_add = " ".join(sentence_to_add.split())
                    
                    if sentence_to_add and len(sentence_to_add.split()) > 4: 
                        live_data_updates.append(sentence_to_add)
                        print(f"[Console] ACCEPTED: The full sentence '{sentence_to_add}' added to update queue.")
                        print(f"[Console] Current update queue size: {len(live_data_updates)}. Type '!retrain' to apply and save changes.")
                        
                        last_generated_text = None
                        last_user_prompt = None
                    else:
                        print("[Console] Cannot accept: The reconstructed sentence was too short or incomplete. Please use '!add [full sentence]' instead.")
                else:
                    print("[Console] No text generated or prompt found. Generate text first.")
                continue
            
            if user_input.lower() == '!help':
                print("\n🤖 Console Information🤖")
                print("1. Type a phrase to generate text (max 10 words).")
                print("2. Use '!add [sentence]' to queue new training data.")
                print("3. Use '!accept' to add the model's last **full** sentence to the training queue.")
                print(f"4. Use '!retrain' to re-train the model on new data (runs for {INTERACTIVE_EPOCHS} epochs) **and save it**.")
                print(f"5. Use '!refine' to re-train on existing data (runs for {INTERACTIVE_EPOCHS} epochs) **without saving.**")
                print("6. Use '!penalty <value>' to regenerate with a different repetition penalty (higher = less repetition).")
                print("7. Type 'quit' or 'exit' to stop.")
                print("8. Type '!help' to see this message again.")
                print("9. Use '!instead [corrected text]' to replace the last output with a corrected version.")
                print("10. Use !endorse to endorse the model on its generated text.")
                print("11. Use !neurons to find active neurons.")
                print("=" * 60 + "\n")
                continue


            if user_input.lower() == '!retrain':
                if not live_data_updates:
                    print("[Console] No new data to train on. Use '!add [sentence]' first.")
                    continue
                
                print(f"\n[Console] RETRAINING MODEL ON NEW DATA ({INTERACTIVE_EPOCHS} EPOCHS)...")
                initialize_or_retrain(initial_train=False, use_live_data=True, epochs=INTERACTIVE_EPOCHS)
                
                temp_weights = "aoban_weights_temp.safetensors"
                final_weights = "aoban_weights.safetensors"
                
                try:
                    save_file(current_model.state_dict(), temp_weights)
                    os.replace(temp_weights, final_weights)
                    print("[Console] Retraining complete. New knowledge acquired and SAVED.")
                except Exception as e:
                    print(f"[Console Bug] Retrain save failed: {e}")
                
                last_generated_text = None 
                last_user_prompt = None
                continue
            if user_input.lower().startswith('!endorse'):
                if last_generated_text and last_user_prompt:
                    try:
                        parts = user_input.split()
                        multiplier = int(parts[1]) if len(parts) > 1 else 5
                        
                        full_sentence = f"{last_user_prompt.strip()} {last_generated_text.strip()}"
                        
                        print(f"🚀 [RLHF] Injecting reward into {NUM_LAYERS} layers...")
                        rlhf_lr = LEARNING_RATE * 2
                        for i in range(3):
                            l = perform_rlhf_step(current_model, current_tokenizer, full_sentence, rlhf_lr)
                        for _ in range(multiplier):
                            live_data_updates.append(full_sentence)
                        
                        print(f"[SYSTEM] ENDORSED: Weights updated and added {multiplier}x to queue.")
                        print(f"[SYSTEM] Alignment Target: {full_sentence}")
                    except Exception as e:
                        print(f"[SYSTEM] RLHF Update Failed: {e}")
                else:
                    print("[SYSTEM] Nothing to endorse.")
                continue
            
            if user_input.lower() == '!refine':
                print(f"\n[Console] REFINING MODEL ON EXISTING DATA ({INTERACTIVE_EPOCHS} EPOCHS)...")
                initialize_or_retrain(initial_train=False, use_live_data=False, epochs=INTERACTIVE_EPOCHS)
                
                temp_weights = "aoban_weights_temp.safetensors"
                final_weights = "aoban_weights.safetensors"
                
                try:
                    save_file(current_model.state_dict(), temp_weights)
                    os.replace(temp_weights, final_weights)
                    print("[Console] Refinement complete. Knowledge deepened and PERMANENTLY SAVED.")
                except Exception as e:
                    print(f"[Console Bug] Save failed but brain is still active in RAM: {e}")
                continue
            if user_input.lower().startswith('!instead '):
                if last_user_prompt and last_generated_text:
                    corrected_output = user_input[9:].strip()
                    
                    penalty_record = f"{last_user_prompt} {last_generated_text}"
                    penalties.append(penalty_record)
                    with open('Penalties.csv', 'a', newline='', encoding='utf-8') as f:
                        csv.writer(f).writerow([penalty_record, REPETITION_PENALTY])
                    
                    full_correct_sentence = f"{last_user_prompt} {corrected_output}"
                    for _ in range(5):
                        live_data_updates.append(full_correct_sentence)
                    
                    print(f"[SYSTEM] Fixed! '{last_generated_text}' is now penalized.")
                    print(f"[SYSTEM] Added correction: '{full_correct_sentence}' to training queue.")
                    
                    last_generated_text = corrected_output
                    save_file(current_model.state_dict(), "aoban_weights.safetensors")
                    print("[Console] Model weights permanently saved to aoban_weights.safetensors")
                else:
                    print("[SYSTEM] Nothing to replace. Generate text first.")
                continue
            if user_input.lower().startswith('!penalty '): 
                try:
                    new_val = float(user_input[9:].strip())
                    REPETITION_PENALTY = new_val
                    save_settings(REPETITION_PENALTY, TEMPERATURE) 
                    
                    if last_user_prompt and last_generated_text:
                        penalty_record = f"{last_user_prompt} {last_generated_text}"
                        penalties.append(penalty_record)
                        
                        with open('Penalties.csv', 'a', newline='', encoding='utf-8') as f:
                            writer = csv.writer(f)
                            writer.writerow([penalty_record, REPETITION_PENALTY])
                    
                    print(f"[Console] Regenerating with new saved penalty={REPETITION_PENALTY}...")
                    generated_text = generate_text(
                        current_model, 
                        current_tokenizer, 
                        last_user_prompt, 
                        MAX_SEQ_LENGTH, 
                        device, 
                        TOP_K, 
                        REPETITION_PENALTY, 
                        TEMPERATURE
                    )
                    
                    print(f"Model: {generated_text}")
                    last_generated_text = generated_text
                    print("\n[Console] If this full sentence is perfect, type '!accept'.")
                except ValueError:
                    print(f"[Console] Invalid value. Usage: !penalty <number>.")
                continue
            
            if user_input.lower() == "!neurons":
                interpret_aoban_neurons(current_model, current_tokenizer, device)


            if user_input.strip() and not user_input.lower().startswith(('!',)):
                prompt = user_input.strip()
                
                if len(prompt.split()) > MAX_SEQ_LENGTH - 1:
                    print(f"[Console] Prompt too long. Max {MAX_SEQ_LENGTH - 1} words supported.")
                    last_generated_text = None
                    last_user_prompt = None
                    continue
                last_user_prompt = prompt
                
                generated_text = generate_text(
                    current_model, 
                    current_tokenizer, 
                    prompt, 
                    MAX_SEQ_LENGTH, 
                    device, 
                    TOP_K, 
                    REPETITION_PENALTY,
                    TEMPERATURE
                )
                print(f"Model: {generated_text}")
                
                last_generated_text = generated_text
                
                print("\n[Console] If this full sentence is perfect, type '!accept' to add it to the training queue.")
            
        except KeyboardInterrupt:
            print("\nExiting interactive mode.")
            break
        except Exception as e:
            print(f"An error occurred: {e}")
            break


def get_rlhf_reward(generated_text, target_text):
    if generated_text.strip() == target_text.strip():
        return 2.0 
    words = generated_text.split()
    if len(words) == 0: return -1.0
    unique_ratio = len(set(words)) / len(words)
    
    reward = unique_ratio * 1.5
    return reward

def interpret_aoban_neurons(model, tokenizer, device):
    """Scans the 10 layers to identify Sentiment vs. Logic neurons."""
    model.eval()
    
    mad_prompt = "The angry bird flew away furiously."
    chill_prompt = "The sun is shining bright and the sky is clear."
    
    def get_activations(text):
        tokens = tokenizer.encode(text).to(device).unsqueeze(0)
        with torch.no_grad():
            x = model.embedding(tokens) * math.sqrt(D_MODEL)
            x = model.pos_encoder(x)
            for i in range(5):
                x = model.transformer.layers[i](x)
            return x.squeeze(0).mean(dim=0)
        seq_len = x.size(1)
        mask = torch.triu(
            torch.full((seq_len, seq_len), float('-inf'), device=x.device),
            diagonal=1
        )

        for i in range(5):
            x = model.transformer.layers[i](x, src_mask=mask)


    mad_act = get_activations(mad_prompt)
    chill_act = get_activations(chill_prompt)
    diff = (mad_act - chill_act).abs()
    sentiment_indices = torch.topk(diff, 5).indices.tolist()
    
    print(f"🧠 [Aoban Brain Map] Total Neurons per Layer: {D_MODEL}")
    print(f"-> Sentiment Neurons (Mood-Sensitive): {sentiment_indices}")
    print(f"-> Stability Neurons (Low Variance): {[i for i in range(10) if diff[i] < 0.01]}")
    
    return sentiment_indices

def find_emergent_neuron(model, tokenizer, positive_text, negative_text):
    model.eval()
    device = next(model.parameters()).device

    with torch.no_grad():
        pos_idx = tokenizer.encode(positive_text).unsqueeze(0).to(device)
        neg_idx = tokenizer.encode(negative_text).unsqueeze(0).to(device)

        # embeddings + positional encoding
        pos_emb = model.embedding(pos_idx) * math.sqrt(model.d_model)
        pos_emb = model.pos_encoder(pos_emb)
        neg_emb = model.embedding(neg_idx) * math.sqrt(model.d_model)
        neg_emb = model.pos_encoder(neg_emb)

        # causal masks
        pos_len = pos_emb.size(1)
        neg_len = neg_emb.size(1)

        pos_mask = torch.triu(
            torch.full((pos_len, pos_len), float('-inf'), device=device),
            diagonal=1
        )
        neg_mask = torch.triu(
            torch.full((neg_len, neg_len), float('-inf'), device=device),
            diagonal=1
        )

        pos_out = model.transformer(pos_emb, mask=pos_mask)
        neg_out = model.transformer(neg_emb, mask=neg_mask)

        pos_final = pos_out.mean(dim=1)
        neg_final = neg_out.mean(dim=1)

        diff = torch.abs(pos_final - neg_final)

        diff[0, 0] = 0
        diff[0, 1] = 0

        top_neuron = torch.argmax(diff).item()
        strength = diff[0, top_neuron].item()

        print(f"\n[System] Deep Sentiment Neuron found at Index: {top_neuron}")
        print(f"[System] Activation Strength Difference: {strength:.4f}")
        


if __name__ == "__main__":
    print(f"[Console] Global Settings Initialized: Penalty={REPETITION_PENALTY}")
    
    interactive_mode()