"""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
from wisent.core.steering_methods.methods.grom._config import (
    GatingNetwork,
    IntensityNetwork,
    GeometryAdaptation,
)


def train_grom_impl(self, pair_set: ContrastivePairSet):
    """
    Full GROM training with all components.
    Args:
        pair_set: ContrastivePairSet with collected activations.
    Returns:
        GROMResult with manifold, networks, and metadata.
    """
    # Import GROMResult here to avoid circular import
    from wisent.core.steering_methods.methods.grom.grom import GROMResult

    # 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_impl(
    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
