from __future__ import annotations

import torch

from staplebridge.training.trajectory import WeightedTrajectory


def assign_lead_local_positive_weights(
    trajectory_groups: list[list[WeightedTrajectory]],
) -> dict[str, float | int]:
    """Assign lead-equal, within-lead Boltzmann weights to positive paths.

    Exact committed-plan completions are normal behavior-cloning positives.
    Off-plan and unfinished paths remain explicit failures with zero positive
    weight; teaching their chosen actions as positives would contradict strict
    committed-plan inference.  A separate failure-avoidance objective may be
    added later without changing this positive-path definition.
    """
    if not trajectory_groups:
        return {
            "n_leads": 0,
            "n_positive": 0,
            "n_failure": 0,
            "n_leads_without_positive": 0,
            "total_positive_weight": 0.0,
        }
    lead_scale = 1.0 / float(len(trajectory_groups))
    n_positive = 0
    n_failure = 0
    n_without_positive = 0
    for group in trajectory_groups:
        for trajectory in group:
            trajectory.weight = 0.0
        positives = [trajectory for trajectory in group if trajectory.is_positive]
        n_positive += len(positives)
        n_failure += len(group) - len(positives)
        if not positives:
            n_without_positive += 1
            continue
        energies = torch.tensor(
            [trajectory.terminal_energy for trajectory in positives],
            dtype=torch.float64,
        )
        weights = torch.softmax(-energies, dim=0).tolist()
        for trajectory, weight in zip(positives, weights):
            trajectory.weight = lead_scale * float(weight)
    return {
        "n_leads": len(trajectory_groups),
        "n_positive": n_positive,
        "n_failure": n_failure,
        "n_leads_without_positive": n_without_positive,
        "total_positive_weight": float(
            sum(trajectory.weight for group in trajectory_groups for trajectory in group)
        ),
    }


def compute_weighted_path_loss(trajs: list[WeightedTrajectory], log_probs: list[torch.Tensor]) -> torch.Tensor:
    device = log_probs[0].device if log_probs else torch.device("cpu")
    weights = torch.tensor([t.weight for t in trajs], dtype=torch.float32, device=device)
    if weights.numel() == 0:
        return torch.tensor(0.0, requires_grad=True, device=device)
    total = torch.tensor(0.0, dtype=torch.float32, device=device)
    offset = 0
    for i, traj in enumerate(trajs):
        n = len(traj.steps)
        if n == 0:
            continue
        lp = torch.stack(log_probs[offset : offset + n]).sum()
        total = total - weights[i] * lp
        offset += n
    return total