/********************************************************************************
 * Copyright (c) 2023 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include "aidge/data/Tensor.hpp"
#include "aidge/utils/Types.h"
#include "aidge/utils/ErrorHandling.hpp"

void Aidge::Tensor::copyCastFrom(const Tensor& src, std::shared_ptr<Tensor>& movedSrcPtr) {
    if (&src == this) {
        return;
    }

    // Current Tensor has necessarily a data type, but may not have backend
    if (!getImpl()) {
        // If no backend was set for the current tensor, use the same as src
        const auto deviceSrc = src.getImpl()->device();
        setBackend(deviceSrc.first, deviceSrc.second);
    }
    resize(src.dims());

    if (dataType() != src.dataType()) {
        // First move data to the target device (only if needed)
        const auto device = getImpl()->device();
        const Tensor& movedSrc = src.refFrom(movedSrcPtr, device.first, device.second);
        // Second, copy-cast data (necessary)
        getImpl()->copyCast(movedSrc.getImpl()->rawPtr(), movedSrc.size(), movedSrc.dataType());
    }
    else {
        // Directly copy, no conversion necessary
        // Avoid making a double copy if both data type and device are the same
        getImpl()->copyFrom(*(src.getImpl()), src.size());
    }
}

Aidge::Tensor& Aidge::Tensor::refCast(std::shared_ptr<Tensor>& fallback, const Aidge::DataType& dt) {
    // Scott Meyers' solution to avoid code duplication
    return const_cast<Tensor&>(static_cast<const Tensor&>(*this).refCast(fallback, dt));
}

const Aidge::Tensor& Aidge::Tensor::refCast(std::shared_ptr<Tensor>& fallback, const Aidge::DataType& dt) const {
    AIDGE_ASSERT(getImpl(), "no backend was set for tensor, cannot refCast() it");

    if (dt == dataType()) {
        return *this;
    }
    else {
        if (!fallback) {
            fallback = std::make_shared<Tensor>(dt);
        }
        else {
            fallback->setDataType(dt, false); // don't keep previous data (no copy)
        }

        const auto device = getImpl()->device();
        fallback->setBackend(device.first, device.second, false); // don't keep previous data (no copy)
        fallback->resize(dims());
        fallback->getImpl()->copyCast(getImpl()->rawPtr(), size(), dataType());
        return *fallback;
    }
}

Aidge::Tensor& Aidge::Tensor::refFrom(std::shared_ptr<Tensor>& fallback, const std::string &backend, int device) {
    // Scott Meyers' solution to avoid code duplication
    return const_cast<Tensor&>(static_cast<const Tensor&>(*this).refFrom(fallback, backend, device));
}

const Aidge::Tensor& Aidge::Tensor::refFrom(std::shared_ptr<Tensor>& fallback, const std::string &backend, int device) const {
    AIDGE_ASSERT(getImpl(), "no backend was set for tensor, cannot refFrom() it");

    if (std::make_pair(backend, device) == getImpl()->device()) {
        return *this;
    }
    else {
        if (!fallback) {
            fallback = std::make_shared<Tensor>(dataType());
        }
        else {
            fallback->setDataType(dataType(), false); // don't keep previous data (no copy)
        }

        fallback->setBackend(backend, device, false); // don't keep previous data (no copy)
        fallback->resize(dims());
        fallback->getImpl()->copyFrom(*getImpl(), size());
        return *fallback;
    }
}