"""GROM train_grom method and joint optimization loop."""
from __future__ import annotations
from typing import List, Dict, Any, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from wisent.core.activations.core.atoms import LayerActivations, RawActivationMap, LayerName
from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
from wisent.core.errors import InsufficientDataError
    def train_grom(self, pair_set: ContrastivePairSet) -> GROMResult:
        """
        Full GROM training with all components.
        Args:
            pair_set: ContrastivePairSet with collected activations.
        Returns:
            GROMResult with manifold, networks, and metadata.
        """
        # Collect activations
        buckets = self._collect_from_set(pair_set)
        if not buckets:
            raise InsufficientDataError(reason="No valid activation pairs found")
        # Detect num_layers from available data if not set
        # Find max layer index to determine model size
        max_layer_idx = 0
        for layer_name in buckets.keys():
            try:
                layer_idx = int(str(layer_name).split("_")[-1])
                max_layer_idx = max(max_layer_idx, layer_idx)
            except (ValueError, IndexError):
                pass
        # Resolve steering_layers and sensor_layer based on detected num_layers
        detected_num_layers = max_layer_idx + 1  # layers are 0-indexed
        if self.config.steering_layers is None or self.config.sensor_layer is None:
            self.config.resolve_layers(detected_num_layers)
        # Filter to steering layers and determine hidden dim
        layer_names = []
        hidden_dim = None
        for layer_name in sorted(buckets.keys()):
            pos_list, neg_list = buckets[layer_name]
            if not pos_list or not neg_list:
                continue
            # Check if layer matches steering_layers config
            try:
                layer_idx = int(str(layer_name).split("_")[-1])
                if layer_idx not in self.config.steering_layers:
                    continue
            except (ValueError, IndexError):
                pass  # Include if can't parse
            layer_names.append(layer_name)
            if hidden_dim is None:
                hidden_dim = pos_list[0].reshape(-1).shape[0]
        if not layer_names or hidden_dim is None:
            raise InsufficientDataError(reason="No valid steering layers found")
        # Resolve network dimensions based on actual hidden_dim
        if self.config.gate_hidden_dim is None or self.config.intensity_hidden_dim is None:
            self.config.resolve_network_dims(hidden_dim)
        num_layers = len(layer_names)
        # Geometry analysis and adaptation
        geometry_adaptation = None
        effective_num_directions = self.config.num_directions
        enable_gating = True
        if self.config.adapt_to_geometry:
            geometry_adaptation = self._analyze_and_adapt_geometry(
                buckets, layer_names, hidden_dim
            )
            effective_num_directions = geometry_adaptation.adapted_num_directions
            enable_gating = geometry_adaptation.gating_enabled
        # Initialize components with adapted configuration
        directions = self._initialize_directions(
            buckets, layer_names, hidden_dim, 
            num_directions=effective_num_directions
        )
        gate_network: Optional[GatingNetwork] = None
        if enable_gating:
            gate_network = GatingNetwork(hidden_dim, self.config.gate_hidden_dim)
        intensity_network = IntensityNetwork(
            hidden_dim, num_layers, 
            self.config.intensity_hidden_dim, 
            self.config.max_alpha
        )
        direction_weights = {
            layer: torch.ones(effective_num_directions) / effective_num_directions
            for layer in layer_names
        }
        # Make direction weights trainable
        direction_weight_params = {
            layer: nn.Parameter(torch.zeros(effective_num_directions))
            for layer in layer_names
        }
        # Prepare data tensors
        data = self._prepare_data_tensors(buckets, layer_names)
        # Find sensor layer
        sensor_layer = self._find_sensor_layer(layer_names)
        # Joint optimization
        directions, gate_network, intensity_network, direction_weights = self._joint_optimization(
            directions=directions,
            gate_network=gate_network,
            intensity_network=intensity_network,
            direction_weight_params=direction_weight_params,
            data=data,
            layer_names=layer_names,
            sensor_layer=sensor_layer,
            enable_gating=enable_gating,
        )
        return GROMResult(
            directions=directions,
            gate_network=gate_network,
            intensity_network=intensity_network,
            direction_weights=direction_weights,
            layer_order=layer_names,
            metadata={
                "config": self.config.__dict__,
                "num_layers": num_layers,
                "hidden_dim": hidden_dim,
                "sensor_layer": sensor_layer,
                "training_logs": self._training_logs,
                "effective_num_directions": effective_num_directions,
                "gating_enabled": enable_gating,
            },
            geometry_adaptation=geometry_adaptation,
        )
    def _joint_optimization(
        self,
        directions: Dict[LayerName, torch.Tensor],
        gate_network: Optional[GatingNetwork],
        intensity_network: IntensityNetwork,
        direction_weight_params: Dict[LayerName, nn.Parameter],
        data: Dict[str, Dict[LayerName, torch.Tensor]],
        layer_names: List[LayerName],
        sensor_layer: LayerName,
        enable_gating: bool = True,
    ) -> Tuple[Dict[LayerName, torch.Tensor], Optional[GatingNetwork], IntensityNetwork, Dict[LayerName, torch.Tensor]]:
        """
        Joint end-to-end optimization of all GROM components.
        """
        # Make directions trainable
        direction_params = {layer: nn.Parameter(dirs.clone()) for layer, dirs in directions.items()}
        # Collect all parameters
        all_params = []
        all_params.extend(direction_params.values())
        if gate_network is not None:
            all_params.extend(gate_network.parameters())
        all_params.extend(intensity_network.parameters())
        all_params.extend(direction_weight_params.values())
        optimizer = torch.optim.AdamW(all_params, lr=self.config.learning_rate, weight_decay=0.01)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=self.config.optimization_steps, eta_min=self.config.learning_rate * 0.1
        )
        best_loss = float('inf')
        best_state = None
        for step in range(self.config.optimization_steps):
            optimizer.zero_grad()
            # Compute effective directions (weighted sum)
            effective_dirs = {}
            for layer in layer_names:
                weights = F.softmax(direction_weight_params[layer], dim=0)
                dirs = direction_params[layer]
                dirs_norm = F.normalize(dirs, p=2, dim=1)
                effective_dirs[layer] = (weights.unsqueeze(-1) * dirs_norm).sum(dim=0)
            # Get sensor layer data
            pos_sensor = data["pos"][sensor_layer]
            neg_sensor = data["neg"][sensor_layer]
            # Predict gates (or use constant 1.0 if gating disabled)
            if gate_network is not None:
                pos_gate = gate_network(pos_sensor, self.config.gate_temperature)
                neg_gate = gate_network(neg_sensor, self.config.gate_temperature)
            else:
                pos_gate = torch.ones(pos_sensor.shape[0], device=pos_sensor.device)
                neg_gate = torch.ones(neg_sensor.shape[0], device=neg_sensor.device)
            # Predict intensities
            pos_intensity = intensity_network(pos_sensor)  # [N_pos, num_layers]
            neg_intensity = intensity_network(neg_sensor)  # [N_neg, num_layers]
            # Compute losses
            loss, loss_components = self._compute_grom_loss(
                direction_params=direction_params,
                effective_dirs=effective_dirs,
                pos_gate=pos_gate,
                neg_gate=neg_gate,
                pos_intensity=pos_intensity,
                neg_intensity=neg_intensity,
                data=data,
                layer_names=layer_names,
                step=step,
                direction_weight_params=direction_weight_params,
            )
            loss.backward()
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(all_params, max_norm=1.0)
            optimizer.step()
            scheduler.step()
            # Apply constraints to directions
            with torch.no_grad():
                for layer in layer_names:
                    direction_params[layer].data = self._apply_direction_constraints(
                        direction_params[layer].data
                    )
            # Track best
            if loss.item() < best_loss:
                best_loss = loss.item()
                best_state = {
                    "directions": {l: p.detach().clone() for l, p in direction_params.items()},
                    "gate_network": {k: v.detach().clone() for k, v in gate_network.state_dict().items()} if gate_network is not None else None,
                    "intensity_network": {k: v.detach().clone() for k, v in intensity_network.state_dict().items()},
                    "direction_weights": {l: F.softmax(p.detach().clone(), dim=0) for l, p in direction_weight_params.items()},
                }
            # Log
            if step % 20 == 0 or step == self.config.optimization_steps - 1:
                # Compute direction weight statistics
                weight_stds = []
                weight_maxes = []
                for layer in layer_names:
                    weights = F.softmax(direction_weight_params[layer], dim=0)
                    weight_stds.append(weights.std().item())
                    weight_maxes.append(weights.max().item())
                self._training_logs.append({
                    "step": step,
                    "total_loss": loss.item(),
                    "lr": scheduler.get_last_lr()[0],
                    **{k: v.item() for k, v in loss_components.items()},
                    "pos_gate_mean": pos_gate.mean().item(),
                    "neg_gate_mean": neg_gate.mean().item(),
                    "pos_intensity_mean": pos_intensity.mean().item(),
                    "neg_intensity_mean": neg_intensity.mean().item(),
                    "direction_weight_std_mean": sum(weight_stds) / len(weight_stds),
                    "direction_weight_max_mean": sum(weight_maxes) / len(weight_maxes),
                })
        # Restore best state
        if best_state is not None:
            final_directions = best_state["directions"]
            if gate_network is not None and best_state["gate_network"] is not None:
                gate_network.load_state_dict(best_state["gate_network"])
            intensity_network.load_state_dict(best_state["intensity_network"])
            final_weights = best_state["direction_weights"]
        else:
            final_directions = {l: p.detach() for l, p in direction_params.items()}
            final_weights = {l: F.softmax(p.detach(), dim=0) for l, p in direction_weight_params.items()}
        # Final normalization
        if self.config.normalize:
            final_directions = {l: F.normalize(d, p=2, dim=1) for l, d in final_directions.items()}
        # POLARITY CORRECTION: Ensure directions point from neg to pos
        # After optimization, directions may have flipped. We check if pos samples
        # have higher projection than neg samples. If not, flip the direction.
        for layer in layer_names:
            pos_data = data["pos"][layer]
            neg_data = data["neg"][layer]
            # Compute effective direction for this layer
            weights = final_weights[layer]
            dirs = final_directions[layer]
            eff_dir = (weights.unsqueeze(-1) * dirs).sum(dim=0)
            eff_dir = F.normalize(eff_dir, p=2, dim=0)
            # Compute mean projections
            pos_proj = (pos_data * eff_dir).sum(dim=1).mean()
            neg_proj = (neg_data * eff_dir).sum(dim=1).mean()
            # If neg > pos, flip all directions for this layer
            if neg_proj > pos_proj:
                final_directions[layer] = -final_directions[layer]
        return final_directions, gate_network, intensity_network, final_weights
