#pragma once
#include "helper_math.hpp"

namespace CUDA {
    struct ArrayValues {
        int size;
        float fullIntegral;
        float length;
    };

    struct ArrayValuesSmall {
        int size;
        float fullIntegral;
    };
} // namespace CUDA

/*
 * x = position
 * y = value
 * z = size
 * w = left integral
 */

__host__ __device__ inline float getIntegralArray(const float pos, const float4* array, const int size,
                                                  const unsigned int offset, const float tMin) {
    if (size == 0) {
        return 0.f;
    }
    auto length = size;
    int first = 0;
    while (length > 0) {
        auto rem = length % 2;
        length /= 2;
        if (array[offset + first + length].x < pos) {
            first += length + rem;
        }
    }

    if (first >= size) {
        // behind last value
        float4 element = array[offset + size - 1];
        return element.w + element.y * element.z + element.y * 0.5f * (pos - element.x);
    }
    // element
    float4 element = array[offset + first];

    if (element.x - element.z > pos) {
        if (first == 0) {
            // infront of first element
            return (pos - tMin) / (element.x - tMin) * element.w;
        }

        // inbetween elements
        float4 dataLeft = array[offset + first - 1];
        float xLeft = dataLeft.x;
        float yLeft = dataLeft.w + dataLeft.y * dataLeft.z;

        float4 dataRight = element;
        float xRight = dataRight.x - dataRight.z;
        float yRight = dataRight.w;

        return yLeft + (yRight - yLeft) / (xRight - xLeft) * (pos - xLeft);
    }

    // inside element
    if (element.z == 0.f) {
        return element.w;
    }
    float a = element.w * (element.x - pos);
    float b = (element.w + element.z * element.y) * (pos - (element.x - element.z));
    return (a + b) / element.z;
}

__host__ __device__ inline float4 internal_insert(float4* newValues, unsigned int offset, float4 candidate, int& k,
                                                  float& len, float& integral, float4& curr, float tMin) {
    // no merge on first insertion
    if (k == 0) {
        // fill empty space between border and first element with average value
        integral += (candidate.x - candidate.z - tMin) * candidate.y * 0.5f;

        candidate.w = integral;
        k++;
        curr = candidate;

        len = candidate.z;
        integral += candidate.z * candidate.y;
        return make_float4(-1);
    }

    if (curr.y == candidate.y) {
        if (curr.x - curr.z < candidate.x && curr.x > candidate.x) {
            // new value redundant
            return make_float4(-1);
        }

        if (candidate.x - candidate.z - curr.x < 0.5) {
            // merge elements
            float oldSize = curr.z;
            float newSize = candidate.x - (curr.x - oldSize);
            integral -= oldSize * curr.y;
            curr = make_float4(candidate.x, curr.y, newSize, integral);

            len += newSize - oldSize;
            integral += newSize * curr.y;
            return make_float4(-1);
        }
    } else {
        if (curr.x - curr.z < candidate.x && curr.x > candidate.x) {
            // split element
            float pos = curr.x;
            float imp = curr.y;
            len -= curr.z;
            integral -= curr.z * imp;
            float newPos = curr.x - curr.z;
            newValues[offset + k - 1] = make_float4(newPos, imp, 0, integral);

            integral += (candidate.x - candidate.z - newPos) * (imp + candidate.y) * 0.5f;
            candidate.w = integral;
            newValues[offset + k++] = candidate;
            curr = candidate;

            integral += candidate.z * candidate.y;
            // after split need to keep end of split range
            return make_float4(pos, imp, 0, 0);
        }
    }
    // insert new element
    // fill empty space between 2 elements with average value
    integral += (candidate.x - candidate.z - curr.x) * (curr.y + candidate.y) * 0.5f;

    candidate.w = integral;
    newValues[offset + k - 1] = curr;
    k++;
    curr = candidate;

    len += candidate.z;
    integral += candidate.z * candidate.y;
    return make_float4(-1);
}

__host__ __device__ inline CUDA::ArrayValues insertArray(int size, int offsetOld, int offsetNew, int offsetTemp,
                                                         float4* oldValues, float4* newValues, float2* temp,
                                                         int newSamples, float tMin, float tMax, int threads) {
    bool coalesced = false;
    if (threads > -1) {
        coalesced = true;
    }
    int i = 0;
    int j = 0;
    int k = 0;
    float len = 0;
    float integral = 0;
    float4 curr = make_float4(-1);

    // Traverse both arrays
    float4 val1 = make_float4(-1);
    while (i < size && j < newSamples) {
        bool newValue = false;
        if (val1.y < 0) {
            val1 = oldValues[offsetOld + i];
            newValue = true;
        }
        int index = j;
        if (coalesced) {
            index = j * threads;
        }
        auto val2 = temp[offsetTemp + index];
        float4 candidate;
        if (val1.x - val1.z < val2.x) {
            candidate = val1;
            if (newValue) {
                ++i;
            }
        } else {
            candidate = make_float4(val2);
            ++j;
        }
        val1 = internal_insert(newValues, offsetNew, candidate, k, len, integral, curr, tMin);
    }

    // Store remaining elements of old array
    while (i < size) {
        if (val1.y > -1) {
            val1 = internal_insert(newValues, offsetNew, val1, k, len, integral, curr, tMin);
        }
        internal_insert(newValues, offsetNew, oldValues[offsetOld + i++], k, len, integral, curr, tMin);
    }

    // Store remaining elements of new array
    while (j < newSamples) {
        if (val1.y > -1) {
            int index = j;
            if (coalesced) {
                index = j * threads;
            }
            float4 candidate = make_float4(temp[offsetTemp + index]);
            if (candidate.x < val1.x) {
                internal_insert(newValues, offsetNew, candidate, k, len, integral, curr, tMin);
                j++;
            } else {
                val1 = internal_insert(newValues, offsetNew, val1, k, len, integral, curr, tMin);
            }
        } else {
            int index = j;
            if (coalesced) {
                index = j * threads;
            }
            val1 = internal_insert(newValues, offsetNew, make_float4(temp[offsetTemp + index]), k, len, integral, curr,
                                   tMin);
            j++;
        }
    }
    if (val1.y > -1) {
        // Last element splitted
        newValues[offsetNew + k - 1] = curr;
        integral += curr.y * curr.z + 0.5f * (curr.y + val1.y) * (val1.x - curr.x);
        val1.w = integral;
        ++k;
        curr = val1;
    }
    newValues[offsetNew + k - 1] = curr;
    integral += curr.y * 0.5f * (tMax - curr.x);

    return {k, integral, len};
}

__host__ __device__ inline void insert_light(float4* newValues, int offset, float2 candidate, int& size, float& len,
                                             float& integral, float4& curr) {
    if (curr.y == candidate.y && candidate.x - curr.x < 0.5) {
        // merge elements
        float oldSize = curr.z;
        float newSize = candidate.x - (curr.x - oldSize);
        integral -= oldSize * curr.y;

        curr.x = candidate.x;
        curr.z = newSize;
        curr.w = integral;

        len += newSize - oldSize;
        integral += newSize * curr.y;
        return;
    }
    // insert new element
    newValues[offset + size - 1] = curr;
    size++;

    // fill empty space between 2 elements with average value
    integral += (candidate.x - curr.x) * (curr.y + candidate.y) * 0.5f;
    curr = make_float4(candidate.x, candidate.y, 0, integral);
}
