/*
 * NumPy arrays with associated physical units.
 *
 * Author: Malte J. Ziebarth (malte.ziebarth@tum.de)
 *
 * Copyright (C) 2026 Technische Universität München
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *
 *    1. Redistributions of source code must retain the above copyright notice, this
 *       list of conditions and the following disclaimer.
 *    2. Redistributions in binary form must reproduce the above copyright notice,
 *       this list of conditions and the following disclaimer in the documentation
 *       and/or other materials provided with the distribution.
 *    3. Neither the name of the copyright holder nor the names of its contributors
 *       may be used to endorse or promote products derived from this software
 *       without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * SPDX-FileCopyrightText: 2026 Technical University of Munich
 * SPDX-License-Identifier: BSD-3-Clause
 * SPDX-FileContributor: Malte J. Ziebarth <malte.ziebarth@tum.de>
 */


/* Can also use Funktionsklempner without Boost as a requirement;
 * this can be done by defining NO_BOOST.
 */
#ifndef FUNKTIONSKLEMPNER_NO_BOOST

#ifndef FUNKTIONSKLEMPNER_QUANTITYARRAY_HPP
#define FUNKTIONSKLEMPNER_QUANTITYARRAY_HPP

#include <ranges>
#include <funktionsklempner/ndarray.hpp>
#include <funktionsklempner/units.hpp>
#include <funktionsklempner/default_units.hpp>
#include <boost/units/quantity.hpp>
#include <boost/units/systems/si.hpp>

namespace funktionsklempner {


/* The main QuantityArray class is contained in
 * the qtyarray sub-namespace.
 * This allows us to typedef input and output
 * QuantityArrays later on. */
namespace qtyarray {


/*
 * Forward-declare some classes:
 */
template<Unit U, NumPyDType T, bool bounds_check=true>
class QuantityArray;



template<Unit U, NumPyDType T>
class QuantityView
{
public:
    using Quantity = boost::units::quantity<U, std::decay_t<T>>;

    constexpr QuantityView(T* ptr, const Quantity& unit)
       : ptr(ptr), unit(unit)
    {
        if (ptr == nullptr)
            throw std::runtime_error(
                "Found nullptr in QuantityView initialization."
            );
    }

    constexpr operator Quantity() const noexcept
    {
        return *ptr * unit;
    }

    /*
     * To be able to let this be used in place of a
     * Quantity&, implement assignment and increment
     * operators.
     * Enable them if `T` is not a constant type (otherwise,
     * we cannot assign).
     */
    template<
        typename Q,
        std::enable_if_t<
            (std::is_same_v<Q, Quantity> && !std::is_const_v<T>),
            bool
        > = true
    >
    constexpr QuantityView& operator=(const Q& q)
    {
        *ptr = q / unit;
        return *this;
    }


    template<
        typename Q,
        std::enable_if_t<
            (std::is_same_v<Q, Quantity> && !std::is_const_v<T>),
            bool
        > = true
    >
    constexpr QuantityView& operator+=(const Q& q)
    {
        *ptr += q / unit;
        return *this;
    }


    template<
        typename Q,
        std::enable_if_t<
            (std::is_same_v<Q, Quantity> && !std::is_const_v<T>),
            bool
        > = true
    >
    constexpr QuantityView& operator-=(const Q& q)
    {
        *ptr -= q / unit;
        return *this;
    }


    constexpr auto operator<=>(const Quantity& q) const noexcept
    {
        return *ptr * unit <=> q;
    }

    constexpr bool operator==(const Quantity& q) const noexcept
    {
        return *ptr * unit == q;
    }

    constexpr Quantity operator+(const Quantity& q) const noexcept
    {
        return (*ptr * unit) + q;
    }

    constexpr Quantity operator-(const Quantity& q) const noexcept
    {
        return (*ptr * unit) - q;
    }

    constexpr Quantity operator*(double d) const noexcept
    {
        return (*ptr * d) * unit;
    }

    constexpr Quantity operator/(double d) const noexcept
    {
        return (*ptr / d) * unit;
    }

private:
    T* ptr;
    Quantity unit;
};



/*
 * Operators with boost quantities as LHS:
 */
template<Unit U, NumPyDType T>
constexpr auto
operator+(const boost::units::quantity<U>& l, const QuantityView<U, T>& r) noexcept
{
    return l + static_cast<boost::units::quantity<U>>(r);
}

template<Unit U, NumPyDType T>
constexpr auto
operator-(const boost::units::quantity<U>& l, const QuantityView<U, T>& r) noexcept
{
    return l - static_cast<boost::units::quantity<U>>(r);
}


template<Unit Ul, NumPyDType Tl, Unit Ur, NumPyDType Tr>
constexpr auto
operator*(const boost::units::quantity<Ul, Tl>& l, const QuantityView<Ur, Tr>& r)
{
    return l * static_cast<boost::units::quantity<Ur, std::decay_t<Tr>>>(r);
}


template<Unit Ul, NumPyDType Tl, Unit Ur, NumPyDType Tr>
constexpr auto
operator*(const QuantityView<Ul, Tl>& l, const boost::units::quantity<Ur, Tr>& r)
{
    return static_cast<boost::units::quantity<Ul, std::decay_t<Tl>>>(l) * r;
}


template<Unit U, NumPyDType T>
constexpr boost::units::quantity<U, T>
operator*(double d, const QuantityView<U, T>& qv)
{
    return qv * d;
}


template<Unit Ul, NumPyDType Tl, Unit Ur, NumPyDType Tr>
constexpr auto
operator/(const boost::units::quantity<Ul, Tl>& l, const QuantityView<Ur, Tr>& r)
{
    return l / static_cast<boost::units::quantity<Ur, std::decay_t<Tr>>>(r);
}


template<Unit Ul, NumPyDType Tl, Unit Ur, NumPyDType Tr>
constexpr auto
operator/(const QuantityView<Ul, Tl>& l, const boost::units::quantity<Ur, Tr>& r)
{
    return static_cast<boost::units::quantity<Ul, std::decay_t<Tl>>>(l) / r;
}


template<Unit U, NumPyDType T>
constexpr auto
operator/(double d, const QuantityView<U, T>& qv)
{
    namespace bu = boost::units;
    return bu::quantity<bu::si::dimensionless>::from_value(d) / qv;
}




template<Unit U, NumPyDType T>
class QuantityIterator
{
    friend class QuantityArray<U, T>;
public:
    /*
     * Iterator interface:
     */
    using iterator_category = std::random_access_iterator_tag;
    using value_type = QuantityView<U, T>;
    using reference = value_type;
    using difference_type = std::ptrdiff_t;
    using pointer = value_type*;

    /* The unit: */
    using Quantity = boost::units::quantity<U>;
    constexpr static Quantity ONE = Quantity::from_value(1.0);

    constexpr QuantityIterator() noexcept
       : ptr(nullptr), unit(ONE)
    {}
    constexpr QuantityIterator(QuantityIterator&&) noexcept = default;
    constexpr QuantityIterator(const QuantityIterator&) noexcept = default;

    /*
     * Assignment:
     */
    constexpr QuantityIterator<U,T>& operator=(const QuantityIterator&) = default;
    constexpr QuantityIterator<U,T>& operator=(QuantityIterator&&) = default;


    /*
     * Increments:
     */
    constexpr QuantityIterator<U,T>& operator++() noexcept
    {
        ++ptr;
        return *this;
    }

    constexpr QuantityIterator<U,T> operator++(int) noexcept
    {
        QuantityIterator<U, T> res(*this, unit);
        ++ptr;
        return res;
    }

    constexpr QuantityIterator<U,T>& operator--() noexcept
    {
        --ptr;
        return *this;
    }

    constexpr QuantityIterator<U,T> operator--(int) noexcept
    {
        QuantityIterator<U, T> res(*this, unit);
        --ptr;
        return res;
    }

    constexpr QuantityIterator<U,T>& operator+=(std::ptrdiff_t i) noexcept
    {
        ptr += i;
        return *this;
    }

    constexpr QuantityIterator<U,T>& operator-=(std::ptrdiff_t i) noexcept
    {
        ptr -= i;
        return *this;
    }

    constexpr QuantityIterator<U,T> operator+(std::ptrdiff_t i) const noexcept
    {
        return QuantityIterator<U,T>(ptr + i, unit);
    }

    constexpr QuantityIterator<U,T> operator-(std::ptrdiff_t i) const noexcept
    {
        return QuantityIterator<U,T>(ptr - i, unit);
    }



    constexpr std::ptrdiff_t operator-(const QuantityIterator& other) const noexcept
    {
        return ptr - other.ptr;
    }


    /*
     * Dereference:
     */
    QuantityView<U, T> operator*() const
    {
        return QuantityView<U, T>(ptr, unit);
    }

    QuantityView<U,T> operator[](std::ptrdiff_t i) const
    {
        return QuantityView<U,T>(ptr + i, unit);
    }


    /*
     * Comparison:
     */
    constexpr bool operator==(const QuantityIterator& other) const noexcept
    {
        return ptr == other.ptr;
    }

    constexpr bool operator==(const std::decay_t<T>* sentinel) const noexcept
    {
        return ptr == sentinel;
    }

    constexpr auto operator<=>(const QuantityIterator& other) const noexcept
    {
        return ptr <=> other.ptr;
    }

    constexpr auto operator<=>(const std::decay_t<T>* sentinel) const noexcept
    {
        return ptr <=> sentinel;
    }


private:
    T* ptr;
    Quantity unit;



    constexpr QuantityIterator(T* ptr, const Quantity& unit) noexcept
       : ptr(ptr), unit(unit)
    {}
};



template<Unit U, NumPyDType T>
QuantityIterator<U,T> operator+(std::ptrdiff_t i, const QuantityIterator<U,T>& it)
{
    return it + i;
}






/*
 * QuantityArray
 * =============
 * An array of quantities with common unit.
 */
template<Unit U, NumPyDType T, bool bounds_check>
class QuantityArray : private ndarray::NDArray<T>
{
public:
    constexpr QuantityArray(
        T* data, size_t N, const char* dtype, const char* unit
    ) : ndarray::NDArray<T>(data, N, dtype),
        unit(parse_si_unit<U>(unit).value())
    {}

    constexpr QuantityArray(
        T* data, size_t N, const char* dtype
    ) : ndarray::NDArray<T>(data, N, dtype),
        unit(boost::units::quantity<U>::from_value(1.0))
    {}


    constexpr QuantityIterator<U, T> begin() const
    {
        return QuantityIterator<U, T>(
            ndarray::NDArray<T>::begin(), unit
        );
    }

    constexpr const std::decay_t<T>* end() const noexcept
    {
        return ndarray::NDArray<T>::end();
    }


    constexpr QuantityView<U, T> operator[](size_t i) const
    {
        if constexpr (bounds_check){
            if (i >= size()){
                std::string msg(
                    "Tried to access element outside bounds ("
                    "index "
                );
                msg += std::to_string(i);
                msg += " for QuantityArray with size ";
                msg += std::to_string(size());
                msg += ").";
                throw std::runtime_error(msg);
            }
        }
        return QuantityView<U, T>(
            ndarray::NDArray<T>::begin() + i, unit
        );
    }


    constexpr size_t size() const noexcept
    {
        return ndarray::NDArray<T>::size();
    }

    template<
        typename U2,
        std::enable_if_t<
            (!std::is_const_v<T> && std::is_same_v<U,U2>),
            bool
        > = true
    >
    constexpr void set_unit(const boost::units::quantity<U2>& u)
    {
        /* Ensure that the unit is valid: */
        constexpr boost::units::quantity<U> ZERO
            = boost::units::quantity<U>::from_value(0.0);
        if (u <= ZERO)
            throw std::runtime_error(
                "Given unit value in base units is smaller than or "
                "equal to zero."
            );

        /* First transform the values:: */
        const double scale = unit / u;
        std::transform(
            ndarray::NDArray<T>::begin(),
            ndarray::NDArray<T>::end(),
            ndarray::NDArray<T>::begin(),
            [scale](T& val) -> T
            {
                return val * scale;
            }
        );

        /* Set the unit: */
        unit = u;
    }


private:
    boost::units::quantity<U> unit;

};


} // namespace qtyarray


/*
 * Specialize the QuantityArray for input and output.
 * This allows specifying input and output arguments
 * via type names.
 *
 * InputQuantityArray:
 *    for input to the C++ world from the Python world
 *    (i.e. just read access in C++).
 * OutputNDArray:
 *    for output from the C++ world to the Python world
 *    (i.e. read and write access in C++).
 */
template<Unit U, NonConstNumPyDType T, bool bounds_check=true>
using InputQuantityArray = qtyarray::QuantityArray<U, const T, bounds_check>;

template<Unit U, NonConstNumPyDType T, bool bounds_check=true>
using OutputQuantityArray = qtyarray::QuantityArray<U, T, bounds_check>;



}

#endif
#endif