from __future__ import annotations

from collections.abc import Callable

import torch
import torch.nn.functional as F

from .targets import TargetSpec, parse_target


def indices_mean(predictions: torch.Tensor, target) -> torch.Tensor:
    spec = parse_target(target)
    if spec.type != "indices":
        raise ValueError("indices_mean objective requires an indices target.")
    idx = spec.value.to(predictions.device)
    return predictions.index_select(dim=1, index=idx).mean(dim=1)


def vector_dot(predictions: torch.Tensor, target) -> torch.Tensor:
    spec = parse_target(target)
    weights = _target_vector(spec, predictions).to(predictions.device)
    return predictions @ weights


def vector_cosine(predictions: torch.Tensor, target) -> torch.Tensor:
    spec = parse_target(target)
    vector = _target_vector(spec, predictions).to(predictions.device)
    return F.cosine_similarity(predictions, vector.unsqueeze(0), dim=1)


def weighted_mean(predictions: torch.Tensor, target) -> torch.Tensor:
    spec = parse_target(target)
    weights = _target_vector(spec, predictions).to(predictions.device)
    denom = weights.abs().sum().clamp_min(1e-8)
    return (predictions * weights.unsqueeze(0)).sum(dim=1) / denom


def build_objective(name: str | Callable) -> Callable:
    if callable(name):
        return name
    objectives = {
        "indices_mean": indices_mean,
        "target_vector_dot": vector_dot,
        "vector_dot": vector_dot,
        "target_vector_cosine": vector_cosine,
        "vector_cosine": vector_cosine,
        "weighted_mean": weighted_mean,
    }
    if name not in objectives:
        raise ValueError(f"Unknown objective: {name}")
    return objectives[name]


def _target_vector(spec: TargetSpec, predictions: torch.Tensor) -> torch.Tensor:
    if spec.type in {"vector", "weights"}:
        vector = spec.value.float()
        if vector.numel() != predictions.shape[1]:
            raise ValueError(f"Target vector has {vector.numel()} values, expected {predictions.shape[1]}.")
        return vector.reshape(-1)
    if spec.type == "indices":
        vector = torch.zeros(predictions.shape[1], dtype=predictions.dtype)
        vector[spec.value.long()] = 1.0
        return vector
    raise ValueError(f"Unsupported target type for vector objective: {spec.type}")