"""Helper functions for cone structure analysis."""

from __future__ import annotations

from typing import Dict, Any, List, Tuple

import torch
import torch.nn.functional as F
from wisent.core.utils.config_tools.constants import COMPARE_TOL, PROGRESS_LOG_INTERVAL_20


def compute_pca_directions(
    pos_tensor: torch.Tensor,
    neg_tensor: torch.Tensor,
    n_components: int,
) -> Tuple[torch.Tensor, float]:
    """Compute PCA directions and explained variance ratio."""
    all_activations = torch.cat([pos_tensor, neg_tensor], dim=0)
    mean = all_activations.mean(dim=0, keepdim=True)
    centered = all_activations - mean

    try:
        U, S, Vh = torch.linalg.svd(centered, full_matrices=False)
        k = min(n_components, Vh.shape[0])
        pca_directions = Vh[:k]
        total_var = (S ** 2).sum()
        explained_var = (S[:k] ** 2).sum() / total_var if total_var > 0 else 0.0
        return pca_directions, float(explained_var)
    except Exception:
        return torch.zeros(n_components, pos_tensor.shape[1]), 0.0


def discover_cone_directions(
    pos_tensor: torch.Tensor, neg_tensor: torch.Tensor,
    num_directions: int, optimization_steps: int, learning_rate: float,
    min_cos_sim: float, max_cos_sim: float,
    cone_noise_scale: float = None, cone_loss_weight_dissimilar: float = None,
    cone_loss_weight_cone: float = None, cone_loss_weight_diversity: float = None,
) -> Tuple[torch.Tensor, Dict[str, Any]]:
    """Discover cone directions via gradient optimization."""
    for _n, _v in [("cone_noise_scale", cone_noise_scale), ("cone_loss_weight_dissimilar", cone_loss_weight_dissimilar), ("cone_loss_weight_cone", cone_loss_weight_cone), ("cone_loss_weight_diversity", cone_loss_weight_diversity)]:
        if _v is None: raise ValueError(f"{_n} is required")
    hidden_dim = pos_tensor.shape[1]

    caa_dir = pos_tensor.mean(dim=0) - neg_tensor.mean(dim=0)
    caa_dir = F.normalize(caa_dir, p=2, dim=0)

    directions = torch.randn(num_directions, hidden_dim)
    directions[0] = caa_dir

    for i in range(1, num_directions):
        noise = torch.randn(hidden_dim) * cone_noise_scale
        directions[i] = F.normalize(caa_dir + noise, p=2, dim=0)

    directions = F.normalize(directions, p=2, dim=1)
    directions.requires_grad_(True)

    optimizer = torch.optim.Adam([directions], lr=learning_rate)
    training_losses = []
    total_loss = torch.tensor(0.0)

    for step in range(optimization_steps):
        optimizer.zero_grad()
        dirs_norm = F.normalize(directions, p=2, dim=1)

        pos_proj = pos_tensor @ dirs_norm.T
        neg_proj = neg_tensor @ dirs_norm.T
        separation_loss = -((pos_proj.mean(dim=0) - neg_proj.mean(dim=0)).abs().mean())

        cos_sim = dirs_norm @ dirs_norm.T
        off_diag_mask = 1 - torch.eye(num_directions)
        off_diag = cos_sim * off_diag_mask

        negative_penalty = F.relu(-off_diag).sum()
        too_similar = F.relu(off_diag - max_cos_sim).sum()
        too_dissimilar = F.relu(min_cos_sim - off_diag).sum()

        cone_loss = (negative_penalty + too_similar
                     + cone_loss_weight_dissimilar * too_dissimilar)
        diversity_loss = -off_diag.var()
        total_loss = (separation_loss
                      + cone_loss_weight_cone * cone_loss
                      + cone_loss_weight_diversity * diversity_loss)

        total_loss.backward()
        optimizer.step()

        with torch.no_grad():
            directions.data = F.normalize(directions.data, p=2, dim=1)
            if directions.shape[0] > 1:
                primary = directions[0:1]
                for i in range(1, directions.shape[0]):
                    if (directions[i:i+1] @ primary.T).item() < 0:
                        directions.data[i] = -directions.data[i]

        if step % PROGRESS_LOG_INTERVAL_20 == 0:
            training_losses.append(float(total_loss.item()))

    return directions.detach(), {"training_losses": training_losses,
                                  "final_loss": float(total_loss.item())}


def compute_cone_explained_variance(
    pos_tensor: torch.Tensor,
    neg_tensor: torch.Tensor,
    cone_directions: torch.Tensor,
) -> float:
    """Compute how much variance the cone directions explain."""
    all_activations = torch.cat([pos_tensor, neg_tensor], dim=0)
    mean = all_activations.mean(dim=0, keepdim=True)
    centered = all_activations - mean

    total_var = (centered ** 2).sum()
    if total_var == 0:
        return 0.0

    dirs_norm = F.normalize(cone_directions, p=2, dim=1)
    projections = centered @ dirs_norm.T
    reconstructed = projections @ dirs_norm

    explained_var = (reconstructed ** 2).sum() / total_var
    return float(min(explained_var.item(), 1.0))


def check_half_space_consistency(directions: torch.Tensor) -> float:
    """Check what fraction of directions are in the same half-space."""
    if directions.shape[0] <= 1:
        return 1.0

    dirs_norm = F.normalize(directions, p=2, dim=1)
    primary = dirs_norm[0:1]
    cos_with_primary = (dirs_norm @ primary.T).squeeze()
    positive_count = (cos_with_primary > 0).sum().item()
    return positive_count / directions.shape[0]


def test_positive_combinations(
    pos_tensor: torch.Tensor, neg_tensor: torch.Tensor, directions: torch.Tensor,
    cone_significance_threshold: float = None, cone_pos_ratio_weight: float = None, cone_sig_ratio_weight: float = None,
) -> float:
    """Test if difference vectors can be represented as positive combinations."""
    for _n, _v in [("cone_significance_threshold", cone_significance_threshold), ("cone_pos_ratio_weight", cone_pos_ratio_weight), ("cone_sig_ratio_weight", cone_sig_ratio_weight)]:
        if _v is None: raise ValueError(f"{_n} is required")
    diff = pos_tensor.mean(dim=0) - neg_tensor.mean(dim=0)
    diff_norm = F.normalize(diff.unsqueeze(0), p=2, dim=1)
    dirs_norm = F.normalize(directions, p=2, dim=1)

    projections = (diff_norm @ dirs_norm.T).squeeze()
    positive_projections = (projections >= 0).sum().item()
    significant_projections = (
        projections > cone_significance_threshold
    ).sum().item()

    pos_ratio = positive_projections / directions.shape[0]
    sig_ratio = significant_projections / directions.shape[0]

    return cone_pos_ratio_weight * pos_ratio + cone_sig_ratio_weight * sig_ratio


def compute_cosine_similarity_matrix(directions: torch.Tensor) -> torch.Tensor:
    """Compute pairwise cosine similarity matrix."""
    dirs_norm = F.normalize(directions, p=2, dim=1)
    return dirs_norm @ dirs_norm.T


def compute_avg_off_diagonal(matrix: torch.Tensor) -> float:
    """Compute average of off-diagonal elements."""
    n = matrix.shape[0]
    if n <= 1:
        return 1.0
    mask = 1 - torch.eye(n)
    return float((matrix * mask).sum() / (n * (n - 1)))


def compute_separation_scores(
    pos_tensor: torch.Tensor,
    neg_tensor: torch.Tensor,
    directions: torch.Tensor,
) -> List[float]:
    """Compute separation score for each direction."""
    dirs_norm = F.normalize(directions, p=2, dim=1)
    pos_proj = pos_tensor @ dirs_norm.T
    neg_proj = neg_tensor @ dirs_norm.T
    separation = pos_proj.mean(dim=0) - neg_proj.mean(dim=0)
    return separation.tolist()


def compute_cone_score(
    pca_explained: float, cone_explained: float, half_space_score: float,
    avg_cos_sim: float, pos_combo_score: float, separation_scores: List[float],
    cone_cos_threshold_low: float = None, cone_cos_score_low: float = None,
    cone_cos_threshold_high: float = None, cone_significance_threshold: float = None,
    cone_weight_cos: float = None, cone_weight_sig: float = None,
    cone_weight_effect: float = None, cone_weight_sample: float = None, cone_weight_ratio: float = None,
) -> float:
    """Compute overall cone score combining all metrics."""
    for _n, _v in [("cone_cos_threshold_low", cone_cos_threshold_low), ("cone_cos_score_low", cone_cos_score_low), ("cone_cos_threshold_high", cone_cos_threshold_high), ("cone_significance_threshold", cone_significance_threshold), ("cone_weight_cos", cone_weight_cos), ("cone_weight_sig", cone_weight_sig), ("cone_weight_effect", cone_weight_effect), ("cone_weight_sample", cone_weight_sample), ("cone_weight_ratio", cone_weight_ratio)]:
        if _v is None: raise ValueError(f"{_n} is required")
    var_ratio = cone_explained / max(pca_explained, COMPARE_TOL)
    var_score = min(var_ratio, 1.0)

    half_space_component = half_space_score

    if avg_cos_sim < 0:
        cos_score = 0.0
    elif avg_cos_sim < cone_cos_threshold_low:
        cos_score = (avg_cos_sim / cone_cos_threshold_low
                     * cone_cos_score_low)
    elif avg_cos_sim <= cone_cos_threshold_high:
        cos_score = 1.0
    else:
        cos_score = max(
            cone_cos_score_low,
            1.0 - (avg_cos_sim - cone_cos_threshold_high)
            / cone_cos_threshold_low,
        )

    combo_component = pos_combo_score

    significant_directions = sum(
        1 for s in separation_scores
        if abs(s) > cone_significance_threshold
    )
    multi_dir_score = min(
        significant_directions / max(len(separation_scores), 1), 1.0
    )

    cone_score = (
        cone_weight_cos * var_score
        + cone_weight_sig * half_space_component
        + cone_weight_effect * cos_score
        + cone_weight_sample * combo_component
        + cone_weight_ratio * multi_dir_score
    )

    return float(cone_score)
