from __future__ import annotations

import numpy as np
import ot
import torch
from sklearn.cluster import MiniBatchKMeans

from .common import pairwise_cosine_distance


def gw_pseudo_targets(
    vision: torch.Tensor,
    text: torch.Tensor,
    clusters: int,
    seed: int,
    max_iter: int = 100,
) -> dict[str, torch.Tensor | float | int]:
    """Build structure-only pseudo-correspondences between modality prototypes."""
    x = vision.float().numpy()
    y = text.float().numpy()
    k = min(clusters, len(x), len(y))
    vx = MiniBatchKMeans(
        n_clusters=k,
        random_state=seed,
        batch_size=min(4096, len(x)),
        n_init=3,
        max_iter=200,
    ).fit(x)
    ty = MiniBatchKMeans(
        n_clusters=k,
        random_state=seed + 1,
        batch_size=min(4096, len(y)),
        n_init=3,
        max_iter=200,
    ).fit(y)

    cx = pairwise_cosine_distance(vx.cluster_centers_)
    cy = pairwise_cosine_distance(ty.cluster_centers_)
    p = np.bincount(vx.labels_, minlength=k).astype(np.float64)
    q = np.bincount(ty.labels_, minlength=k).astype(np.float64)
    p /= p.sum()
    q /= q.sum()

    coupling, log = ot.gromov.gromov_wasserstein(
        cx,
        cy,
        p,
        q,
        loss_fun="square_loss",
        armijo=False,
        log=True,
        max_iter=max_iter,
        tol_rel=1e-8,
        tol_abs=1e-8,
    )
    target_centers = coupling @ ty.cluster_centers_
    target_centers /= p[:, None].clip(min=1e-12)
    target_centers /= np.linalg.norm(target_centers, axis=1, keepdims=True).clip(
        min=1e-12
    )
    row_entropy = -np.sum(
        (coupling / p[:, None].clip(min=1e-12))
        * np.log(
            (coupling / p[:, None].clip(min=1e-12)).clip(min=1e-12)
        ),
        axis=1,
    ).mean()

    return {
        "vision_centers": torch.from_numpy(vx.cluster_centers_).float(),
        "text_centers": torch.from_numpy(ty.cluster_centers_).float(),
        "target_centers": torch.from_numpy(target_centers).float(),
        "vision_assignments": torch.from_numpy(vx.labels_).long(),
        "coupling": torch.from_numpy(coupling).float(),
        "gw_distance": float(log["gw_dist"]),
        "coupling_row_entropy": float(row_entropy),
        "clusters": k,
    }