from typing import Literal

import torch

from mipcandy.common import Normalize


def convert_ids_to_logits(ids: torch.Tensor, d: Literal[1, 2, 3], num_classes: int) -> torch.Tensor:
    if ids.dtype != torch.int or ids.min() < 0:
        raise TypeError("`ids` should be positive integers")
    d += 1
    if ids.ndim != d:
        if ids.ndim == d + 1 and ids.shape[1] == 1:
            ids = ids.squeeze(1)
        else:
            raise ValueError(f"`ids` should be {d} dimensional or {d + 1} dimensional with single channel")
    logits = torch.zeros((ids.shape[0], num_classes, *ids.shape[1:]), device=ids.device, dtype=torch.float32)
    logits.scatter_(1, ids.unsqueeze(1).long(), 1)
    return logits


def convert_logits_to_ids(logits: torch.Tensor, *, channel_dim: int = 1) -> torch.Tensor:
    return logits.round().int() if logits.shape[channel_dim] < 2 else logits.argmax(channel_dim, keepdim=True)


def auto_convert(image: torch.Tensor) -> torch.Tensor:
    return (image * 255 if 0 <= image.min() <= image.max() <= 1 else Normalize(domain=(0, 255))(image)).int()
