import torch
import numpy as np
from scipy.sparse import csc_array, csr_matrix, issparse, spmatrix


def torch_csc_to_scipy_csr(A: torch.Tensor) -> csr_matrix:
    if A.layout != torch.sparse_csc:
        raise TypeError(f"Expected torch.sparse_csc, got {A.layout}")

    A_csc = csc_array((
        A.values().cpu().numpy(),
        A.row_indices().cpu().numpy().astype(np.int32),
        A.ccol_indices().cpu().numpy().astype(np.int32),
    ), shape=A.shape)

    return csr_matrix(A_csc)


def torch_csc_to_scipy_csc(A: torch.Tensor) -> csc_array:
    if A.layout != torch.sparse_csc:
        raise TypeError(f"Expected torch.sparse_csc, got {A.layout}")

    return csc_array((
        A.values().cpu().numpy(),
        A.row_indices().cpu().numpy().astype(np.int32),
        A.ccol_indices().cpu().numpy().astype(np.int32),
    ), shape=A.shape)


def scipy_to_torch_csc(A: spmatrix, device: torch.device = torch.device("cpu")) -> torch.Tensor:
    A_csc = A.tocsc()
    return torch.sparse_csc_tensor(
        torch.tensor(A_csc.indptr, dtype=torch.int64, device=device),
        torch.tensor(A_csc.indices, dtype=torch.int64, device=device),
        torch.tensor(A_csc.data, dtype=torch.float64, device=device),
        size=A_csc.shape,
    )


def to_numpy(r: torch.Tensor) -> np.ndarray:
    return r.detach().cpu().numpy()


def from_numpy(x: np.ndarray, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
    return torch.from_numpy(np.ascontiguousarray(x)).to(dtype=dtype, device=device)
