// Copyright (c) 2021, Justin Kim

// We are making my contributions/submissions to this project solely in our
// personal capacity and are not conveying any rights to any intellectual
// property of any third parties.

#include <Sapphire/compute/Initialize.hpp>
#include <Sapphire/compute/dense/cuda/Initialize.cuh>
#include <Sapphire/compute/dense/naive/NaiveInitialize.hpp>
#include <chrono>
#include <cmath>

namespace Sapphire::Compute::Initialize
{
void Normal(TensorUtil::TensorData& data, float mean, float sd)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Normal(data.CudaMutableRawPtr(), mean, sd,
                            data.DenseTotalLengthCuda,
                            static_cast<int>(std::clock()));
    }
    else
    {
        Dense::Naive::Normal(data.HostMutableRawPtr(), mean, sd,
                             data.GetShape());
    }
}

void Uniform(TensorUtil::TensorData& data, float min, float max)
{
    if (const auto device = data.GetCudaDevice();
        data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Uniform(data.CudaMutableRawPtr(), min, max,
                             data.DenseTotalLengthCuda,
                             static_cast<int>(std::clock()));
    }
    else
    {
        Dense::Naive::Uniform(data.HostMutableRawPtr(), min, max,
                              data.GetShape());
    }
}

void Ones(TensorUtil::TensorData& data)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Scalar(data.CudaMutableRawPtr(), 1.0f,
                            data.DenseTotalLengthCuda);
    }
    else
    {
        Dense::Naive::Scalar(data.HostMutableRawPtr(), 1.0f, data.GetShape());
    }
}

void Zeros(TensorUtil::TensorData& data)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Scalar(data.CudaMutableRawPtr(), 0.0f,
                            data.DenseTotalLengthCuda);
    }
    else
    {
        Dense::Naive::Scalar(data.HostMutableRawPtr(), 0.0f, data.GetShape());
    }
}

void Scalar(TensorUtil::TensorData& data, float value)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Scalar(data.CudaMutableRawPtr(), value,
                            data.DenseTotalLengthCuda);
    }
    else
    {
        Dense::Naive::Scalar(data.HostMutableRawPtr(), value, data.GetShape());
    }
}

void HeNormal(TensorUtil::TensorData& data, int fanIn)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Normal(
            data.CudaMutableRawPtr(), 0.0,
            2.0f / std::sqrt(static_cast<float>(fanIn)),
            data.DenseTotalLengthCuda, static_cast<int>(std::clock()));
    }
    else
    {
        Dense::Naive::Normal(data.HostMutableRawPtr(), 0.0,
                             2.0f / std::sqrt(static_cast<float>(fanIn)),
                             data.GetShape());
    }
}

void Xavier(TensorUtil::TensorData& data, int fanIn, int fanOut)
{
    const auto device = data.GetCudaDevice();
    if (data.Mode() == ComputeMode::Cuda)
    {
        Dense::Cuda::Normal(
            data.CudaMutableRawPtr(), 0.0,
            1.0f / std::sqrt(static_cast<float>(fanIn + fanOut)),
            data.DenseTotalLengthCuda, static_cast<int>(std::clock()));
    }
    else
    {
        Dense::Naive::Normal(data.HostMutableRawPtr(), 0.0,
                             1.0f / std::sqrt(
                                 static_cast<float>(fanIn + fanOut)),
                             data.GetShape());
    }
}
} // namespace Sapphire::Compute::Initialize
