#pragma once

#include "common.h"
#include "Tensor.h"

namespace pytorch_compat {
    inline void TORCH_CHECK(bool cond, const std::string &msg = "") {
        assert (cond);
    }

    template<typename T>
    inline void C10_CUDA_CHECK(T ret) {
        return checkCUDA(ret);
    }

    namespace at {
        using ::Tensor;

        constexpr auto kFloat32 = Tensor::FP32;
        constexpr auto kFloat = Tensor::FP32;
        constexpr auto kFloat16 = Tensor::FP16;
        constexpr auto kBFloat16 = Tensor::BF16;
        constexpr auto kInt32 = Tensor::INT32;
        constexpr auto kInt64 = Tensor::INT64;

        struct Generator {
            Generator() { throw std::runtime_error("Not implemented"); }
            std::mutex mutex_;
        };

        namespace cuda {
            using ::getCurrentDeviceProperties;

            struct StreamWrapper {
                cudaStream_t st;
                cudaStream_t stream() const { return st; }
            };
            inline StreamWrapper getCurrentCUDAStream() { 
                return StreamWrapper(::getCurrentCUDAStream());
            }

            struct CUDAGuard {
                int dev;
            };

            namespace detail {
                inline Generator getDefaultCUDAGenerator() {
                    return Generator();
                }
            }
        }

        using CUDAGeneratorImpl = Generator;

        template<typename T>
        std::unique_ptr<Generator> get_generator_or_default(std::optional<Generator> gen, T gen2) {
            throw std::runtime_error("Not implemented");
        }
    }

    namespace torch {
        using at::kFloat32;
        using at::kFloat;
        using at::kFloat16;
        using at::kBFloat16;
        using at::kInt32;
        using at::kInt64;
        constexpr Device kCUDA = Device::cuda();

        using IntArrayRef = std::vector<int>;
        using TensorOptions = Tensor::TensorOptions;

        inline Tensor empty_like(const Tensor &tensor) {
            return Tensor::empty_like(tensor);
        }
        inline Tensor empty(TensorShape shape, Tensor::TensorOptions options) {
            return Tensor::empty(shape, options.dtype(), options.device());
        }
        inline Tensor zeros(TensorShape shape, Tensor::TensorOptions options) {
            return Tensor::empty(shape, options.dtype(), options.device()).zero_();
        }

        namespace nn {
            namespace functional {
                using PadFuncOptions = std::vector<int>;
                inline Tensor pad(Tensor x, PadFuncOptions options) {
                    throw std::runtime_error("Not implemented");
                }
            }
        }

        namespace indexing {
            constexpr int None = 0;
            struct Slice {
                int a;
                int b;
            };
        }
    }

    namespace c10 {
        using std::optional;
    }

}

