model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import hashlib

# --- Future Entropy Predictor (FEP) ---
# (No changes from V4)
class FutureEntropyPredictor(nn.Module):
    def __init__(self, input_dim=2, hidden_dim=16, output_dim=1, name=""):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.name = name
        self.debug_prints_enabled = False

    def forward(self, current_block_entropy, current_static_target_diff):
        if not torch.is_tensor(current_block_entropy):
            current_block_entropy = torch.tensor([current_block_entropy], device=self.fc1.weight.device, dtype=torch.float32)
        if not torch.is_tensor(current_static_target_diff):
            current_static_target_diff = torch.tensor([current_static_target_diff], device=self.fc1.weight.device, dtype=torch.float32)
        current_block_entropy = current_block_entropy.view(-1, 1)
        current_static_target_diff = current_static_target_diff.view(-1, 1)
        x_in = torch.cat((current_block_entropy, current_static_target_diff), dim=1)
        h = F.relu(self.fc1(x_in))
        predicted_delta_factor_raw = self.fc2(h)
        return predicted_delta_factor_raw.squeeze(-1)

# --- Helper: Entropy Estimator ---
# (No changes from V4)
class EntropyEstimator(nn.Module):
    def __init__(self, d_model, hidden_dim=32, name=""):
        super().__init__()
        self.fc1 = nn.Linear(d_model, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, 1)
        self.name = name
        self.debug_prints_enabled = False
    def forward(self, x, active_mask=None):
        if x.numel() == 0: return torch.tensor(0.0, device=x.device)
        if active_mask is not None:
            if active_mask.dtype != torch.bool: active_mask = active_mask.bool()
            if x.dim() == 3 and active_mask.dim() == 2 and x.shape[:2] == active_mask.shape: x_masked = x[active_mask]
            elif x.dim() == 2 and active_mask.dim() == 1 and x.shape[0] == active_mask.shape[0]: x_masked = x[active_mask]
            else: x_masked = x.reshape(-1, x.size(-1))
        else: x_masked = x.reshape(-1, x.size(-1))
        if x_masked.numel() == 0: return torch.tensor(0.0, device=x.device)
        h = F.relu(self.fc1(x_masked)); return torch.sigmoid(self.fc2(h)).mean()

# --- Helper: Seed Parser ---
# (No changes from V4)
class SeedParser:
    def __init__(self, seed_phrase, seed_number_str, d_model, num_adaptive_blocks, num_sub_modules_per_block):
        self.seed_phrase = seed_phrase; self.seed_number_str = seed_number_str; self.d_model = d_model
        self.num_adaptive_blocks = num_adaptive_blocks; self.num_sub_modules_per_block = num_sub_modules_per_block
        self.debug_prints_enabled = True
        if self.debug_prints_enabled: print(f"--- SeedParser Initialization ---\n  Seed Phrase (start): '{self.seed_phrase[:50]}...'\n  Seed Number: {self.seed_number_str}")
        phrase_hash = hashlib.sha256(seed_phrase.encode()).hexdigest(); self.phrase_base_val = int(phrase_hash[:16], 16)
        if self.debug_prints_enabled: print(f"  Phrase Base Value (from hash): {self.phrase_base_val}")
        self.num_sequence = [int(d) for d in seed_number_str if d.isdigit()]
        if not self.num_sequence: self.num_sequence = [sum(bytearray(seed_number_str.encode())) % 10]
        if self.debug_prints_enabled: print(f"  Numerical Sequence (from seed number): {self.num_sequence}")
        self.init_map = self._generate_init_map()
        if self.debug_prints_enabled:
            print(f"  SeedParser: Generated InitMap:")
            for i, block_config in enumerate(self.init_map["block_configs"]):
                gate_inits_str = [f'{g:.3f}' for g in block_config['initial_gate_proportions']]
                raw_gate_scores_str = [f'{g:.3f}' for g in block_config['raw_gate_scores_for_param_init']]
                print(f"    Block {i}: Target Entropy: {block_config['target_entropy']:.4f}, RawGateScores: {raw_gate_scores_str}, InitialGateProps (softmax): {gate_inits_str}")
        if self.debug_prints_enabled: print(f"--- SeedParser Initialized ---")
    def _get_deterministic_value(self, key_name, min_val, max_val, sequence_idx_offset=0): # ... (same as V4)
        key_specific_hash = int(hashlib.sha256(key_name.encode() + self.seed_phrase.encode()).hexdigest()[:8], 16); num_seq_val = 0
        if self.num_sequence:
            for i, digit in enumerate(self.num_sequence): num_seq_val = (num_seq_val * 10 + digit) % 1000003
        combined_seed_val = self.phrase_base_val + key_specific_hash + num_seq_val + sequence_idx_offset
        if max_val == min_val: return min_val
        val_range = max_val - min_val + 1
        return min_val + int(abs(math.sin(float(combined_seed_val)) * 1e5)) % int(val_range)
    def _get_deterministic_float(self, key_name, min_val=0.0, max_val=1.0, sequence_idx_offset=0): # ... (same as V4)
        key_specific_hash = int(hashlib.sha256(key_name.encode() + self.seed_phrase.encode()).hexdigest()[:8], 16); num_seq_val = 0
        if self.num_sequence:
            for i, digit in enumerate(self.num_sequence): num_seq_val = (num_seq_val * 10 + digit) % 1000003
        combined_seed_val = self.phrase_base_val + key_specific_hash + num_seq_val + sequence_idx_offset
        norm_float = (math.sin(float(combined_seed_val) * 0.1) + 1.0) / 2.0
        return min_val + norm_float * (max_val - min_val)
    def _generate_init_map(self): # ... (same as V4, but remember initial_gate_proportions are softmax based)
        init_map = {"block_configs": []}
        for i in range(self.num_adaptive_blocks):
            gate_raw_scores = [self._get_deterministic_float(f"block_{i}_gate_{j}_raw_score", -1.5, 1.5, sequence_idx_offset=i*10 + j) for j in range(self.num_sub_modules_per_block)]
            gate_initial_proportions = F.softmax(torch.tensor(gate_raw_scores), dim=0).tolist() if self.num_sub_modules_per_block > 0 else []
            target_entropy = self._get_deterministic_float(f"block_{i}_target_entropy", 0.15, 0.45, sequence_idx_offset=i)
            init_map["block_configs"].append({"initial_gate_proportions": gate_initial_proportions, "raw_gate_scores_for_param_init": gate_raw_scores, "target_entropy": target_entropy})
        return init_map
    def get_block_config(self, block_idx): # ... (same as V4)
        if 0 <= block_idx < len(self.init_map["block_configs"]): return self.init_map["block_configs"][block_idx]
        return None

# --- Adaptive Block (V5 changes) ---
class AdaptiveBlock(nn.Module):
    MAX_DYNAMIC_ENTROPY_ADJUSTMENT_RANGE = 0.05
    INITIAL_HEURISTIC_STRENGTH = 0.025 # V5: Start strength for heuristic
    FINAL_HEURISTIC_STRENGTH = 0.005   # V5: End strength for heuristic

    def __init__(self, d_model, n_heads, d_ff, dropout, seed_parser_config_for_block, block_idx, num_sub_modules=3):
        super().__init__()
        self.d_model = d_model; self.block_idx = block_idx; self.num_sub_modules = num_sub_modules
        self.config_from_seed = seed_parser_config_for_block; self.debug_prints_enabled = True

        raw_gate_param_inits_list = self.config_from_seed.get("raw_gate_scores_for_param_init", [0.0] * self.num_sub_modules)
        if len(raw_gate_param_inits_list) != self.num_sub_modules:
            raw_gate_param_inits_list = [0.0] * self.num_sub_modules
        self.gates_params = nn.Parameter(torch.tensor(raw_gate_param_inits_list, dtype=torch.float32))
        # V5: Store initial raw scores as a buffer for alignment loss
        self.register_buffer('initial_raw_gate_scores_buffer', torch.tensor(raw_gate_param_inits_list, dtype=torch.float32))

        if self.debug_prints_enabled:
            raw_gate_scores_str = [f'{g:.3f}' for g in raw_gate_param_inits_list]
            print(f"  Initializing AdaptiveBlock {self.block_idx} with seed config: StaticSeedTgtEnt={self.config_from_seed['target_entropy']:.3f}, InitialRawGateScores={raw_gate_scores_str}")

        self.sub_module_0 = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.sub_module_1 = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model))
        self.sub_module_2 = nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(), nn.Dropout(dropout))
        self.sub_modules = nn.ModuleList([self.sub_module_0, self.sub_module_1, self.sub_module_2])
        if self.num_sub_modules > len(self.sub_modules): self.num_sub_modules = len(self.sub_modules)
        elif self.num_sub_modules <= 0: raise ValueError(f"AdaptiveBlock {self.block_idx} must have at least one sub_module.")

        self.norm1 = nn.LayerNorm(d_model); self.norm2 = nn.LayerNorm(d_model)
        self.dropout_layer = nn.Dropout(dropout) # V5 Renamed from self.dropout to avoid conflict
        self.output_entropy_estimator = EntropyEstimator(d_model, name=f"Block{block_idx}_OutEntropy")
        self.fep = FutureEntropyPredictor(input_dim=2, hidden_dim=16, output_dim=1, name=f"Block{block_idx}_FEP")

        self.wiring_phase_active = False
        self.static_seed_target_entropy = self.config_from_seed.get("target_entropy", 0.25)
        self.current_epoch_in_wiring = 0 # V5
        self.total_wiring_epochs = 1     # V5: Default to 1 to prevent division by zero if not set

    # V5: set_wiring_phase now takes epoch info for decaying strength
    def set_wiring_phase(self, active, current_epoch_num=0, total_wiring_epochs=1):
        self.wiring_phase_active = active
        if active:
            self.current_epoch_in_wiring = current_epoch_num
            self.total_wiring_epochs = total_wiring_epochs if total_wiring_epochs > 0 else 1

    def _get_current_heuristic_strength(self):
        if not self.wiring_phase_active or self.total_wiring_epochs <= 1:
            return self.INITIAL_HEURISTIC_STRENGTH # Or some default if not wiring

        # Linear decay from INITIAL to FINAL strength over total_wiring_epochs
        progress = min(self.current_epoch_in_wiring / (self.total_wiring_epochs -1 ), 1.0) if self.total_wiring_epochs >1 else 1.0

        decayed_strength = self.INITIAL_HEURISTIC_STRENGTH - progress * (self.INITIAL_HEURISTIC_STRENGTH - self.FINAL_HEURISTIC_STRENGTH)
        return decayed_strength

    def forward(self, x, key_padding_mask=None, attn_mask=None):
        # V5: Sigmoid activations
        current_gates_activations = torch.sigmoid(self.gates_params)

        if self.debug_prints_enabled and self.wiring_phase_active:
            print(f"    AdaptiveBlock {self.block_idx} (Wiring ON, Epoch {self.current_epoch_in_wiring+1}/{self.total_wiring_epochs}) Input x: {x.shape}, RawG: {[f'{g.item():.3f}' for g in self.gates_params.data]}, SigmoidG: {[f'{s.item():.3f}' for s in current_gates_activations.data]}")

        x_norm_submodules = self.norm1(x)
        outputs = []
        for i, module_instance in enumerate(self.sub_modules):
            if i >= self.num_sub_modules: break
            if i == 0: module_out, _ = module_instance(x_norm_submodules, x_norm_submodules, x_norm_submodules, key_padding_mask=key_padding_mask, attn_mask=attn_mask, need_weights=False)
            else: module_out = module_instance(x_norm_submodules)
            outputs.append(module_out * current_gates_activations[i]) # V5: Apply sigmoid activation here

        if not outputs: final_out_unnorm = x
        else:
            # V5: Summing activated outputs (no further multiplication by gates needed here as it's done above)
            weighted_sum = torch.sum(torch.stack(outputs, dim=0), dim=0)
            final_out_unnorm = x + self.dropout_layer(weighted_sum)

        final_out_norm = self.norm2(final_out_unnorm)
        current_output_entropy = self.output_entropy_estimator(final_out_norm, active_mask=~key_padding_mask if key_padding_mask is not None else None)
        current_static_target_diff = current_output_entropy - self.static_seed_target_entropy
        dynamic_target_entropy_for_heuristic = self.static_seed_target_entropy
        predicted_delta_factor_for_report = torch.tensor(0.0, device=x.device)

        if self.wiring_phase_active and self.training:
            predicted_delta_factor_raw = self.fep(current_output_entropy.detach(), current_static_target_diff.detach())
            predicted_delta_factor_tanh = torch.tanh(predicted_delta_factor_raw)
            dynamic_adjustment = predicted_delta_factor_tanh * self.MAX_DYNAMIC_ENTROPY_ADJUSTMENT_RANGE
            dynamic_target_entropy_for_heuristic = self.static_seed_target_entropy + dynamic_adjustment.item()
            dynamic_target_entropy_for_heuristic = max(0.01, min(0.99, dynamic_target_entropy_for_heuristic))
            predicted_delta_factor_for_report = predicted_delta_factor_tanh

            with torch.no_grad():
                entropy_diff_for_heuristic = current_output_entropy - dynamic_target_entropy_for_heuristic
                # V5: Decaying heuristic strength
                base_adjustment_strength = self._get_current_heuristic_strength()
                adaptive_strength_factor = min(max(abs(entropy_diff_for_heuristic.item()) * 7.0, 0.3), 2.5)
                adjustment_strength = base_adjustment_strength * adaptive_strength_factor

                if self.debug_prints_enabled:
                    print(f"    AdaptiveBlock {self.block_idx} WIRING PRE-ADJUST: RawG={[f'{g.item():.3f}' for g in self.gates_params.data]}, SigmoidG={[f'{s.item():.3f}' for s in current_gates_activations.data]}")
                    print(f"      OutEnt={current_output_entropy.item():.4f}, StaticTgtEnt={self.static_seed_target_entropy:.4f}, FEPΔFactor={predicted_delta_factor_tanh.item():.4f}, DynTgtEnt={dynamic_target_entropy_for_heuristic:.4f}, ED_Dyn={entropy_diff_for_heuristic.item():.4f}, BaseHeurStr={base_adjustment_strength:.4f} AdjStr={adjustment_strength:.4f}")

                if entropy_diff_for_heuristic.item() > 1e-4:
                    self.gates_params.data[0] -= adjustment_strength
                    self.gates_params.data[1] += adjustment_strength * 0.6
                    if self.num_sub_modules > 2: self.gates_params.data[2] += adjustment_strength * 0.4
                elif entropy_diff_for_heuristic.item() < -1e-4:
                    self.gates_params.data[0] += adjustment_strength
                    self.gates_params.data[1] -= adjustment_strength * 0.6
                    if self.num_sub_modules > 2: self.gates_params.data[2] -= adjustment_strength * 0.4
                self.gates_params.data.clamp_(-3.5, 3.5)
                if self.debug_prints_enabled:
                    print(f"    AdaptiveBlock {self.block_idx} WIRING POST-ADJUST: RawG={[f'{g.item():.3f}' for g in self.gates_params.data]}, SigmoidG={[f'{s.item():.3f}' for s in torch.sigmoid(self.gates_params.data)]}")

        # V5: Return sigmoid activations
        return final_out_norm, current_output_entropy, current_gates_activations, self.gates_params.data.clone(), predicted_delta_factor_for_report, torch.tensor(dynamic_target_entropy_for_heuristic, device=x.device)

# --- Positional Encoding ---
# (No changes from V4)
class PositionalEncoding(nn.Module): # ... (same as V4)
    def __init__(self,d_model,dropout=0.1,max_len=512): super().__init__(); self.dropout=nn.Dropout(p=dropout); pe=torch.zeros(max_len,d_model); pos=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1); div=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model)); pe[:,0::2]=torch.sin(pos*div); pe[:,1::2]=torch.cos(pos*div); self.register_buffer('pe',pe.unsqueeze(0))
    def forward(self,x): x=x+self.pe[:,:x.size(1),:]; return self.dropout(x)

# --- Main SWCK Model (V5 changes) ---
class SWCKModel(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, d_ff, num_adaptive_blocks,
                 dropout, seed_phrase, seed_number_str, num_sub_modules_per_block=3):
        super().__init__()
        self.d_model = d_model; self.seed_phrase = seed_phrase; self.seed_number_str = seed_number_str
        self.debug_prints_enabled = True
        if self.debug_prints_enabled: print(f"--- Initializing SWCKModel (V5) ---")
        self.seed_parser = SeedParser(seed_phrase, seed_number_str, d_model, num_adaptive_blocks, num_sub_modules_per_block)
        self.seed_parser.debug_prints_enabled = self.debug_prints_enabled
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        self.adaptive_blocks = nn.ModuleList()
        for i in range(num_adaptive_blocks):
            block_config = self.seed_parser.get_block_config(i)
            if block_config is None: raise ValueError(f"SWCKModel Error: Could not get seed config for block {i}")
            new_block = AdaptiveBlock(d_model, n_heads, d_ff, dropout, block_config, block_idx=i, num_sub_modules=num_sub_modules_per_block)
            new_block.debug_prints_enabled = self.debug_prints_enabled
            self.adaptive_blocks.append(new_block)
            if self.debug_prints_enabled: print(f"  SWCKModel: Added AdaptiveBlock {i} (V5 with Sigmoid Gates, Decaying Heuristic)")
        self.fc_out = nn.Linear(d_model, vocab_size)
        self.overall_output_entropy_estimator = EntropyEstimator(d_model, name="OverallOutEntropy")
        self.overall_output_entropy_estimator.debug_prints_enabled = False
        self._init_weights()
        if self.debug_prints_enabled: print(f"--- SWCKModel V5 Initialized (Vocab: {vocab_size}, d_model: {d_model}, Blocks: {num_adaptive_blocks}x{num_sub_modules_per_block}sub) ---")

    def _init_weights(self): # ... (same as V4)
        initrange = 0.1; self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_(); self.fc_out.weight.data.uniform_(-initrange, initrange)

    # V5: set_wiring_phase now takes epoch info
    def set_wiring_phase(self, active, current_epoch_num=0, total_wiring_epochs=1):
        if self.debug_prints_enabled:
            print(f"SWCKModel: Setting wiring phase to {active} for all blocks (Epoch {current_epoch_num+1}/{total_wiring_epochs} of wiring if active).")
        for block in self.adaptive_blocks:
            block.set_wiring_phase(active, current_epoch_num, total_wiring_epochs)

    def forward(self, src_tokens, src_key_padding_mask=None):
        if self.debug_prints_enabled:
            print(f"\n--- SWCKModel Forward Pass (Training: {self.training}) ---")
            print(f"  Input src_tokens: {src_tokens.shape}")
            if src_key_padding_mask is not None: print(f"  Input src_key_padding_mask: {src_key_padding_mask.shape} (True means pad)")
        x = self.embedding(src_tokens) * math.sqrt(self.d_model)
        x = self.pos_encoder(x)
        if self.debug_prints_enabled: print(f"  After Embedding & PosEnc, x: {x.shape}")

        block_output_entropies = []
        current_block_gate_activations = [] # V5: Changed from softmaxes
        current_block_gate_raw_params = []
        fep_predicted_delta_factors = []
        dynamic_target_entropies_used = []

        for i, block in enumerate(self.adaptive_blocks):
            if self.debug_prints_enabled: print(f"  Processing AdaptiveBlock {i}...")
            # V5 AdaptiveBlock returns sigmoid activations
            x, block_entropy, current_gate_acts, raw_gate_params, fep_delta, dyn_target_ent = block(x, key_padding_mask=src_key_padding_mask, attn_mask=None)

            block_output_entropies.append(block_entropy)
            current_block_gate_activations.append(current_gate_acts) # V5
            current_block_gate_raw_params.append(raw_gate_params)
            fep_predicted_delta_factors.append(fep_delta)
            dynamic_target_entropies_used.append(dyn_target_ent)

            if self.debug_prints_enabled:
                acts_str = [f'{act.item():.3f}' for act in current_gate_acts] # V5
                raw_str = [f'{rp.item():.3f}' for rp in raw_gate_params]
                fep_delta_str = f"{fep_delta.item():.3f}" if torch.is_tensor(fep_delta) else "N/A"
                dyn_target_str = f"{dyn_target_ent.item():.3f}" if torch.is_tensor(dyn_target_ent) else "N/A"
                print(f"  Output x from Block {i}: {x.shape}, MeasEnt: {block_entropy.item():.4f}, FEPΔFactor: {fep_delta_str}, DynTgtUsed: {dyn_target_str}, SigmoidG: {acts_str}, RawG: {raw_str}") # V5

        logits = self.fc_out(x)
        if self.debug_prints_enabled: print(f"  Output logits: {logits.shape}")
        final_active_mask = ~src_key_padding_mask if src_key_padding_mask is not None else None
        overall_entropy = self.overall_output_entropy_estimator(x, active_mask=final_active_mask)
        if self.debug_prints_enabled: print(f"  Overall Final Representation Entropy: {overall_entropy.item():.4f}")

        entropy_report = {
            "block_output_entropies": block_output_entropies,
            "overall_output_entropy": overall_entropy,
            "current_block_gate_activations": current_block_gate_activations, # V5
            "current_block_gate_params": current_block_gate_raw_params,
            # "initial_block_gate_targets" (softmax based) is removed from report as it's less relevant with sigmoid gates
            # The alignment loss will use the initial_raw_gate_scores_buffer directly from the block.
            "fep_predicted_delta_factors": fep_predicted_delta_factors,
            "dynamic_target_entropies_used": dynamic_target_entropies_used
        }
        if self.debug_prints_enabled: print(f"--- SWCKModel Forward Pass Complete ---")
        return logits, entropy_report