"""Tests the 'basic' hadamard transform based operations (SORF, SRHT, FHT)
for CPU and -- if GPU is available -- for GPU as well. We must test
for both single and double precision."""
import sys
import unittest
import numpy as np
from scipy.linalg import hadamard
import cupy as cp

from xGPR.xgpr_cpu_rfgen_cpp_ext import cpuFastHadamardTransform as cFHT
from xGPR.xgpr_cpu_rfgen_cpp_ext import cpuFastHadamardTransform2D as cFHT2D
from xGPR.xgpr_cpu_rfgen_cpp_ext import cpuSRHT as cSRHT

from xGPR.xgpr_cuda_rfgen_cpp_ext import cudaSRHT
from xGPR.xgpr_cuda_rfgen_cpp_ext import cudaFastHadamardTransform2D as cudaFHT2D


#To test the C extension for the fast hadamard transform, we compare the
#results to matrix multiplication with Hadamard matrices generated by
#Scipy (result should be the same).
class TestFastHadamardTransform(unittest.TestCase):
    """Runs tests for basic functionality (i.e. FHT and SORF)
    for float and double precision for CPU and (if available) GPU."""


    def test_3d_array_transform(self):
        """Tests the 3d FHT transform. This is for CPU only;
        the cuda module does not provide a separate function
        for FHT only at this time."""
        dim = (124,36,4)
        outcome_f, outcome_d = run_fht_test(dim)
        self.assertTrue(outcome_d)
        self.assertTrue(outcome_f)

        dim = (124,36,2)
        outcome_f, outcome_d = run_fht_test(dim)
        self.assertTrue(outcome_d)
        self.assertTrue(outcome_f)

        dim = (124,36,32)
        outcome_f, outcome_d = run_fht_test(dim)
        self.assertTrue(outcome_d)
        self.assertTrue(outcome_f)

        dim = (3001, 5, 1024)
        outcome_f, outcome_d = run_fht_test(dim)
        self.assertTrue(outcome_d)
        self.assertTrue(outcome_f)

        dim = (250, 1, 4096)
        outcome_f, outcome_d = run_fht_test(dim)
        self.assertTrue(outcome_d)
        self.assertTrue(outcome_f)


    def test_2d_array_transform(self):
        """Tests the 2d FHT transform."""
        dim = (124, 4)
        outcomes = run_fht_2d_test(dim)
        for outcome in outcomes:
            self.assertTrue(outcome)

        dim = (124, 2)
        outcomes = run_fht_2d_test(dim)
        for outcome in outcomes:
            self.assertTrue(outcome)

        dim = (124, 32)
        outcomes = run_fht_2d_test(dim)
        for outcome in outcomes:
            self.assertTrue(outcome)

        dim = (3001, 1024)
        outcomes = run_fht_2d_test(dim)
        for outcome in outcomes:
            self.assertTrue(outcome)

        dim = (250, 4096)
        outcomes = run_fht_2d_test(dim)
        for outcome in outcomes:
            self.assertTrue(outcome)


    def test_srht(self):
        """Tests SRHT functionality. Note that this tests SRHT
        functionality by using FHT. Therefore if the FHT did not
        pass, this one will not either."""

        outcomes = run_srht_test((150,256), 128)
        for outcome in outcomes:
            self.assertTrue(outcome)

        outcomes = run_srht_test((304,512), 256)
        for outcome in outcomes:
            self.assertTrue(outcome)

        outcomes = run_srht_test((5,2048), 512)
        for outcome in outcomes:
            self.assertTrue(outcome)


def run_fht_test(dim, random_seed = 123):
    """A helper function that runs an FHT test with specified
    dimensionality on a 3d input array."""
    scipy_double, scipy_float, marr, marr_float = setup_fht_test(dim, random_seed)
    cFHT(marr, 1)
    cFHT(marr_float, 1)
    outcome_d = np.allclose(scipy_double, marr)
    outcome_f = np.allclose(scipy_float, marr_float, rtol=1e-4, atol=1e-4)
    print("Did the C extension provide the correct result for "
            f"a {dim} array for floats, doubles? {outcome_f},{outcome_d}")
    return outcome_d, outcome_f



def run_fht_2d_test(dim, random_seed = 123):
    """A helper function that runs an FHT 2d test with specified
    dimensionality."""
    scipy_double, scipy_float, marr, marr_float = setup_fht_2d_test(dim, random_seed)
    cupy_marr, cupy_marr_float = marr.copy(), marr_float.copy()
    cFHT2D(marr, 1)
    cFHT2D(marr_float, 1)
    outcome_d = np.allclose(scipy_double, marr)
    outcome_f = np.allclose(scipy_float, marr_float, rtol=1e-4, atol=1e-4)
    print("Did the C extension provide the correct result for "
            f"a {dim} 2d array for floats, doubles? {outcome_f},{outcome_d}")

    if "cupy" in sys.modules:
        cupy_marr, cupy_marr_float = cp.asarray(cupy_marr), cp.asarray(cupy_marr_float)
        cudaFHT2D(cupy_marr)
        cudaFHT2D(cupy_marr_float)
        outcome_cuda_d = np.allclose(scipy_double, cupy_marr)
        outcome_cuda_f = np.allclose(scipy_float, cupy_marr_float, rtol=1e-4, atol=1e-4)
        print("Did the cuda extension provide the correct result for "
            f"a {dim} 2d array for floats, doubles? {outcome_cuda_f},{outcome_cuda_d}")
        return outcome_d, outcome_f, outcome_cuda_d, outcome_cuda_f

    return outcome_d, outcome_f


def setup_fht_test(dim, random_seed = 123):
    """A helper function that builds a matrix to test on
    with specified dimensions and gets the 'true' result using
    Scipy's Hadamard matrix generator. This is a very slow
    way to do a transform -- we use it here only to cross-
    check the results from our routines. Generated for
    both single and double precision."""
    rng = np.random.default_rng(random_seed)
    marr = rng.uniform(low=-10.0,high=10.0, size=dim)

    hmat = hadamard(marr.shape[2])
    scipy_results = []
    for i in range(marr.shape[1]):
        scipy_transform = np.matmul(hmat, marr[:,i,:].T).T
        scipy_results.append(scipy_transform)
    scipy_double = np.stack(scipy_results, axis=1)

    marr_float = marr.copy().astype(np.float32)
    scipy_results = []
    for i in range(marr_float.shape[1]):
        scipy_transform = np.matmul(hmat, marr_float[:,i,:].T).T
        scipy_results.append(scipy_transform)
    scipy_float = np.stack(scipy_results, axis=1)
    return scipy_double, scipy_float, marr, marr_float

def setup_fht_2d_test(dim, random_seed = 123):
    """A helper function that builds a matrix to test on
    with specified dimensions and gets the 'true' result using
    Scipy's Hadamard matrix generator. This is a very slow
    way to do a transform -- we use it here only to cross-
    check the results from our routines. Generated for
    both single and double precision."""
    rng = np.random.default_rng(random_seed)
    marr = rng.uniform(low=-10.0,high=10.0, size=dim)

    hmat = hadamard(marr.shape[1])
    scipy_double = np.matmul(hmat, marr.T).T

    marr_float = marr.copy().astype(np.float32)
    scipy_float = np.matmul(hmat, marr_float.T).T
    return scipy_double, scipy_float, marr, marr_float

def run_srht_test(dim, compression_size, random_seed = 123):
    """A helper function that runs the SRHT test for
    specified input dimensions."""
    marr_gt_double, marr_test_double, marr_gt_float, marr_test_float,\
            radem, norm_constant = setup_srht_test(dim,
                                    compression_size, random_seed)

    if "cupy" in sys.modules:
        cuda_test_double = cp.asarray(marr_test_double)
        cuda_test_float = cp.asarray(marr_test_float)

    marr_gt_double = marr_gt_double * radem[None,:] * norm_constant
    cFHT2D(marr_gt_double, 2)

    marr_gt_float = marr_gt_float * radem[None,:] * norm_constant
    cFHT2D(marr_gt_float, 2)

    cSRHT(marr_test_double, radem, 2)
    cSRHT(marr_test_float, radem, 2)
    outcome_d = np.allclose(marr_gt_double, marr_test_double)
    outcome_f = np.allclose(marr_gt_float, marr_test_float, rtol=1e-3,
            atol=1e-5)
    print("Did the C extension provide the correct result for SRHT of "
            f"a {dim} 2d array for floats, doubles? {outcome_f},{outcome_d}")

    if "cupy" in sys.modules:
        radem = cp.asarray(radem)
        cudaSRHT(cuda_test_float, radem, 2)
        cudaSRHT(cuda_test_double, radem, 2)
        cuda_test_double = cp.asnumpy(cuda_test_double)
        cuda_test_float = cp.asnumpy(cuda_test_float)
        outcome_cuda_d = np.allclose(marr_gt_double, cuda_test_double)
        outcome_cuda_f = np.allclose(marr_gt_float, cuda_test_float,
                rtol=1e-3, atol=1e-5)
        print("Did the Cuda extension provide the correct result for SRHT of "
            f"a {dim} 2d array for floats, doubles? {outcome_cuda_f},{outcome_cuda_d}")
        return outcome_d, outcome_f, outcome_cuda_d, outcome_cuda_f

    return outcome_d, outcome_f



def setup_srht_test(dim, compression_size, random_seed = 123):
    """A helper function that builds the matrices required for
    the SRHT test, specified using the input dimensions."""
    radem_array = np.asarray([-1.0,1.0], dtype=np.int8)
    rng = np.random.default_rng(random_seed)

    marr_gt_double = rng.uniform(low=-10.0,high=10.0, size=dim)
    marr_test_double = marr_gt_double.copy()
    marr_gt_float = marr_gt_double.copy().astype(np.float32)
    marr_test_float = marr_gt_float.copy()

    radem = rng.choice(radem_array, size=(dim[1]), replace=True)
    norm_constant = np.log2(dim[1]) / 2
    norm_constant = 1 / (2**norm_constant)
    return marr_gt_double, marr_test_double, marr_gt_float, marr_test_float, \
            radem, norm_constant


if __name__ == "__main__":
    unittest.main()
