"""Pure validation contracts shared by numerical analysis routines."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch


@dataclass(frozen=True)
class ProjectionResult:
    """Pure rank-1 projection result and metadata."""

    weight: torch.Tensor
    projected: bool
    coefficient_norm_sq: float
    layout: str | None


def _stable_float_dtype(*tensors: torch.Tensor) -> torch.dtype:
    """Select a supported compute dtype without downcasting float64 inputs."""
    if any(tensor.dtype == torch.float64 for tensor in tensors):
        return torch.float64
    if any(tensor.dtype in {torch.float16, torch.bfloat16} for tensor in tensors):
        return torch.float32
    if all(tensor.is_floating_point() for tensor in tensors):
        return tensors[0].dtype
    return torch.float32


def validate_whitened_parameters(
    regularization_eps: object,
    min_variance_ratio: object,
) -> tuple[float, float]:
    """Validate and normalize whitened-SVD tuning parameters."""
    if (
        isinstance(regularization_eps, bool)
        or not isinstance(regularization_eps, (int, float))
        or not math.isfinite(regularization_eps)
        or regularization_eps <= 0
    ):
        raise ValueError("regularization_eps must be a finite positive number")
    if (
        isinstance(min_variance_ratio, bool)
        or not isinstance(min_variance_ratio, (int, float))
        or not math.isfinite(min_variance_ratio)
        or not 0 <= min_variance_ratio < 1
    ):
        raise ValueError("min_variance_ratio must be in the interval [0, 1)")
    return float(regularization_eps), float(min_variance_ratio)


def validate_whitened_request(
    harmful_count: int,
    harmless_count: int,
    n_directions: object,
) -> int:
    """Validate paired sample counts and requested direction count."""
    if harmful_count <= 0 or harmless_count <= 0:
        raise ValueError("harmful and harmless activations must both be non-empty")
    if harmful_count != harmless_count:
        raise ValueError(
            "harmful and harmless activations must have equal sample counts, got "
            f"{harmful_count} and {harmless_count}",
        )
    if isinstance(n_directions, bool) or not isinstance(n_directions, int):
        raise ValueError("n_directions must be a positive integer")
    if n_directions <= 0:
        raise ValueError("n_directions must be a positive integer")
    return n_directions


def orthogonalize_subspace_rows(subspace: torch.Tensor) -> torch.Tensor:
    """Orthogonalize rows of a subspace matrix with QR while preserving dtype/device."""
    if subspace.shape[0] <= 1 or subspace.numel() == 0:
        return subspace
    if not torch.isfinite(subspace).all():
        return subspace
    compute_dtype = _stable_float_dtype(subspace)
    work = subspace.to(dtype=compute_dtype)
    if work.norm() < 1e-8:
        return torch.zeros_like(subspace)

    q, _ = torch.linalg.qr(work.T)
    result = q[:, : subspace.shape[0]].T
    if (result[0] @ work[0]) < 0:
        result[0] = -result[0]
    return result.to(dtype=subspace.dtype, device=subspace.device)


def remove_harmless_principal_components(
    subspace: torch.Tensor,
    harmless_stack: torch.Tensor,
    pc_count: int,
) -> torch.Tensor:
    """Subtract dominant benign activation PCs from refusal directions."""
    if pc_count <= 0 or harmless_stack.shape[0] < 3 or subspace.numel() == 0:
        return subspace

    compute_dtype = _stable_float_dtype(subspace, harmless_stack)
    harmless_work = harmless_stack.to(dtype=compute_dtype)
    centered = harmless_work - harmless_work.mean(dim=0, keepdim=True)
    try:
        _, _, vh = torch.linalg.svd(centered, full_matrices=False)
    except Exception:
        return subspace

    k = min(int(pc_count), vh.shape[0], subspace.shape[1])
    if k <= 0:
        return subspace

    original = subspace.to(dtype=compute_dtype)
    pcs = vh[:k]
    residual = original - (original @ pcs.T) @ pcs
    row_norms = residual.norm(dim=-1, keepdim=True)
    near_zero = row_norms.squeeze(-1) < 1e-8
    if near_zero.any():
        residual[near_zero] = original[near_zero]
        row_norms = residual.norm(dim=-1, keepdim=True)

    residual = residual / row_norms.clamp(min=1e-8)
    if residual.shape[0] > 1:
        residual = orthogonalize_subspace_rows(residual)
    return residual.to(dtype=subspace.dtype, device=subspace.device)


def residualize_against_shield_atoms(
    subspace: torch.Tensor,
    atoms: torch.Tensor,
    ridge: float,
) -> torch.Tensor:
    """Remove protected concept atoms with ridge-regularized projection."""
    if atoms.numel() == 0 or subspace.numel() == 0:
        return subspace

    compute_dtype = _stable_float_dtype(subspace, atoms)
    original = subspace.to(dtype=compute_dtype)
    normalized_atoms = atoms.to(dtype=compute_dtype)
    normalized_atoms = normalized_atoms / normalized_atoms.norm(dim=-1, keepdim=True).clamp(min=1e-8)
    gram = normalized_atoms @ normalized_atoms.T
    eye = torch.eye(gram.shape[0], dtype=gram.dtype, device=gram.device)
    try:
        coeff = torch.linalg.solve(gram + float(ridge) * eye, normalized_atoms @ original.T)
    except Exception:
        return subspace

    residual = original - coeff.T @ normalized_atoms
    row_norms = residual.norm(dim=-1, keepdim=True)
    near_zero = row_norms.squeeze(-1) < 1e-8
    if near_zero.any():
        residual[near_zero] = original[near_zero]
        row_norms = residual.norm(dim=-1, keepdim=True)

    residual = residual / row_norms.clamp(min=1e-8)
    if residual.shape[0] > 1:
        residual = orthogonalize_subspace_rows(residual)
    return residual.to(dtype=subspace.dtype, device=subspace.device)


def select_projection_coefficients(
    coeff: torch.Tensor,
    projection_row_fraction: float,
) -> torch.Tensor:
    """Keep only the strongest projection coefficients when selective projection is requested."""
    if not 0.0 < projection_row_fraction <= 1.0:
        raise ValueError("projection_row_fraction must be in (0.0, 1.0]")
    if projection_row_fraction >= 1.0:
        return coeff

    flat = coeff.detach().abs().reshape(-1).float().cpu()
    n_coeffs = flat.numel()
    if n_coeffs == 0:
        return coeff

    keep = max(1, min(n_coeffs, math.ceil(n_coeffs * projection_row_fraction)))
    if keep >= n_coeffs:
        return coeff

    idx = torch.topk(flat, keep, sorted=False).indices
    mask = torch.zeros(n_coeffs, dtype=torch.bool)
    mask[idx] = True
    mask = mask.reshape(coeff.shape).to(device=coeff.device)
    return coeff * mask.to(dtype=coeff.dtype)


def project_weight_against_direction(
    weight: torch.Tensor,
    direction: torch.Tensor,
    *,
    norm_preserve: bool = False,
    regularization: float = 0.0,
    projection_row_fraction: float = 1.0,
    max_norm_ratio: float = 1.10,
) -> ProjectionResult:
    """Return a pure rank-1 projection update for standard or transposed weight layouts."""
    compute_dtype = _stable_float_dtype(weight, direction)
    work = weight.to(dtype=compute_dtype)
    d = direction.to(device=weight.device, dtype=compute_dtype).reshape(-1, 1)
    if not torch.isfinite(work).all() or not torch.isfinite(d).all():
        return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)
    d_norm = d.norm()
    if d_norm < 1e-8:
        return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)
    d = d / d_norm

    scale = 1.0 - regularization
    original_norm_sq = work.pow(2).sum().item() if norm_preserve else 0.0

    if work.shape[-1] == d.shape[0]:
        layout = "standard"
        coeff = work @ d
        if not torch.isfinite(coeff).all():
            return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=layout)
        coeff_to_remove = select_projection_coefficients(coeff, projection_row_fraction)
        coeff_norm_sq = coeff_to_remove.pow(2).sum().item() if norm_preserve else 0.0
        projected = work - d.T * (scale * coeff_to_remove)
    elif work.shape[0] == d.shape[0]:
        layout = "transposed"
        coeff = d.T @ work
        if not torch.isfinite(coeff).all():
            return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=layout)
        coeff_to_remove = select_projection_coefficients(coeff, projection_row_fraction)
        coeff_norm_sq = coeff_to_remove.pow(2).sum().item() if norm_preserve else 0.0
        projected = work - (scale * d) * coeff_to_remove
    else:
        return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)

    if norm_preserve and original_norm_sq > 0:
        new_norm_sq = max(0.0, original_norm_sq - scale * (2 - scale) * coeff_norm_sq)
        if new_norm_sq > 0:
            ratio = math.sqrt(original_norm_sq / new_norm_sq)
            if ratio > max_norm_ratio:
                ratio = max_norm_ratio
            projected = projected * ratio

    return ProjectionResult(
        weight=projected.to(dtype=weight.dtype, device=weight.device),
        projected=True,
        coefficient_norm_sq=coeff_norm_sq,
        layout=layout,
    )
