"""
The diffusionnet architecture and code is based on the following Github repository:
https://github.com/nmwsharp/diffusion-net/
"""

import sys
import os
import time

import torch
import hashlib
import numpy as np
import scipy

# == Pytorch things


def toNP(x):
    """
    Really, definitely convert a torch tensor to a numpy array
    """
    x = x.detach().to(torch.device("cpu")).numpy()
    return np.ascontiguousarray(x)


def label_smoothing_log_loss(pred, labels, smoothing=0.0):
    n_class = pred.shape[-1]
    one_hot = torch.zeros_like(pred)
    one_hot[labels] = 1.0
    one_hot = one_hot * (1 - smoothing) + (1 - one_hot) * smoothing / (n_class - 1)
    loss = -(one_hot * pred).sum(dim=-1).mean()
    return loss


# Randomly rotate points.
# Torch in, torch out
# Note fornow, builds rotation matrix on CPU.
def random_rotate_points(pts, randgen=None):
    R = random_rotation_matrix(randgen)
    R = torch.from_numpy(R).to(device=pts.device, dtype=pts.dtype)
    return torch.matmul(pts, R)


def random_rotate_points_y(pts):
    angles = torch.rand(1, device=pts.device, dtype=pts.dtype) * (2.0 * np.pi)
    rot_mats = torch.zeros(3, 3, device=pts.device, dtype=pts.dtype)
    rot_mats[0, 0] = torch.cos(angles)
    rot_mats[0, 2] = torch.sin(angles)
    rot_mats[2, 0] = -torch.sin(angles)
    rot_mats[2, 2] = torch.cos(angles)
    rot_mats[1, 1] = 1.0

    pts = torch.matmul(pts, rot_mats)
    return pts


# Numpy things


def sparse_np_to_torch(A):
    """
    Converts a numpy sparse matrix to a PyTorch sparse tensor.

    Args:
        A: A scipy sparse matrix (e.g., COO, CSR).

    Returns:
        PyTorch sparse tensor.
    """
    Acoo = A.tocoo()  # Convert to COO format if not already
    values = torch.tensor(Acoo.data, dtype=torch.float32)
    indices = torch.tensor(np.vstack((Acoo.row, Acoo.col)), dtype=torch.int64)
    shape = Acoo.shape
    return torch.sparse_coo_tensor(indices, values, torch.Size(shape)).coalesce()


# Pytorch sparse to numpy csc matrix
def sparse_torch_to_np(A):
    if len(A.shape) != 2:
        raise RuntimeError("should be a matrix-shaped type; dim is : " + str(A.shape))

    indices = toNP(A.indices())
    values = toNP(A.values())

    mat = scipy.sparse.coo_matrix((values, indices), shape=A.shape).tocsc()

    return mat


# Hash a list of numpy arrays
def hash_arrays(arrs):
    running_hash = hashlib.sha1()
    for arr in arrs:
        binarr = arr.view(np.uint8)
        running_hash.update(binarr)
    return running_hash.hexdigest()


def random_rotation_matrix(randgen=None):
    """
    Creates a random rotation matrix.
    randgen: if given, a np.random.RandomState instance used for random numbers (for reproducibility)
    """
    # adapted from http://www.realtimerendering.com/resources/GraphicsGems/gemsiii/rand_rotation.c

    if randgen is None:
        randgen = np.random.RandomState()

    theta, phi, z = tuple(randgen.rand(3).tolist())

    theta = theta * 2.0 * np.pi  # Rotation about the pole (Z).
    phi = phi * 2.0 * np.pi  # For direction of pole deflection.
    z = z * 2.0  # For magnitude of pole deflection.

    # Compute a vector V used for distributing points over the sphere
    # via the reflection I - V Transpose(V).  This formulation of V
    # will guarantee that if x[1] and x[2] are uniformly distributed,
    # the reflected points will be uniform on the sphere.  Note that V
    # has length sqrt(2) to eliminate the 2 in the Householder matrix.

    r = np.sqrt(z)
    Vx, Vy, Vz = V = (np.sin(phi) * r, np.cos(phi) * r, np.sqrt(2.0 - z))

    st = np.sin(theta)
    ct = np.cos(theta)

    R = np.array(((ct, st, 0), (-st, ct, 0), (0, 0, 1)))
    # Construct the rotation matrix  ( V Transpose(V) - I ) R.

    M = (np.outer(V, V) - np.eye(3)).dot(R)
    return M


# Python string/file utilities
def ensure_dir_exists(d):
    if not os.path.exists(d):
        os.makedirs(d)
