/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * 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.
 *  * Neither the name of NVIDIA CORPORATION 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 ``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 OWNER 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.
 */

#include "Mandelbrot_gold.h"

#define ABS(n) ((n) < 0 ? -(n) : (n))

/* dfloat class declaration */
class dfloat
{
private:
    float val[2];

public:
    dfloat() { val[0] = val[1] = 0; }
    dfloat(float a, float b)
    {
        val[0] = a;
        val[1] = b;
    }
    dfloat(double b);
    inline float operator[](unsigned idx) const { return val[idx]; }
};
inline dfloat operator+(const dfloat &dsa, const dfloat &dsb);
inline dfloat operator-(const dfloat &dsa, const dfloat &dsb);
inline dfloat operator*(const dfloat &dsa, const dfloat &dsb);
inline int    operator<(const dfloat &a, float b) { return a[0] < b; }

// The core Mandelbrot calculation function template
template <class T>
inline int
CalcMandelbrot(const T xPos, const T yPos, const T xJParam, const T yJParam, const int crunch, const bool isJulia)
{
    T   x, y, xx, yy, xC, yC;
    int i = crunch;

    if (isJulia) {
        xC = xJParam;
        yC = yJParam;
        y  = yPos;
        x  = xPos;
        yy = y * y;
        xx = x * x;
    }
    else {
        xC = xPos;
        yC = yPos;
        x = y = 0;
        xx = yy = 0;
    }

    while (--i && (xx + yy < 4.0f)) {
        y  = x * y + x * y + yC;
        x  = xx - yy + xC;
        yy = y * y;
        xx = x * x;
    }

    return i;
} // CalcMandelbrot

inline void updatePixel(uchar4 &dst, const uchar4 &color, int frame)
{
    int frame1 = frame + 1;
    int frame2 = frame1 / 2;
    dst.x      = (dst.x * frame + color.x + frame2) / frame1;
    dst.y      = (dst.y * frame + color.y + frame2) / frame1;
    dst.z      = (dst.z * frame + color.z + frame2) / frame1;
}

inline void setColor(uchar4 &dst, const uchar4 &colors, int &m, const int animationFrame)
{
    if (m == 0) {
        dst.x = 0;
        dst.y = 0;
        dst.z = 0;
        return;
    }

    m += animationFrame;
    dst.x = m * colors.x;
    dst.y = m * colors.y;
    dst.z = m * colors.z;
}

template <class T, class T_>
void runMandelbrotGold0(uchar4      *dst,
                        const int    imageW,
                        const int    imageH,
                        const int    crunch,
                        const T      xOff,
                        const T      yOff,
                        const T      xJParam,
                        const T      yJParam,
                        const T      scale,
                        const uchar4 colors,
                        const int    frame,
                        const int    animationFrame,
                        const bool   isJulia)
{
    for (int iy = 0; iy < imageH; iy++)
        for (int ix = 0; ix < imageW; ix++) {
            // Calculate the location
            const T_ xPos = (T)ix * scale + xOff;
            const T_ yPos = (T)iy * scale + yOff;

            // Calculate the Mandelbrot index for the current location
            int m = CalcMandelbrot<T_>(xPos, yPos, xJParam, yJParam, crunch, isJulia);
            m     = m > 0 ? crunch - m : 0;

            // Convert the Mandelbrot index into a color
            uchar4 color;

            setColor(color, colors, m, animationFrame);

            // Output the pixel
            int pixel = imageW * iy + ix;

            if (frame == 0) {
                color.w    = 0;
                dst[pixel] = color;
            }
            else
                updatePixel(dst[pixel], color, frame);
        }

} // runMandelbrotGold0_

// Determine if two pixel colors are within tolerance
inline int CheckColors(const uchar4 &color0, const uchar4 &color1)
{
    int x = color1.x - color0.x;
    int y = color1.y - color0.y;
    int z = color1.z - color0.z;
    return (ABS(x) > 10) || (ABS(y) > 10) || (ABS(z) > 10);
} // CheckColors

template <class T, class T_>
void runMandelbrotGold1(uchar4      *dst,
                        const int    imageW,
                        const int    imageH,
                        const int    crunch,
                        const T      xOff,
                        const T      yOff,
                        const T      xJParam,
                        const T      yJParam,
                        const T      scale,
                        const uchar4 colors,
                        const int    frame,
                        const int    animationFrame,
                        const bool   isJulia)
{
    for (int iy = 0; iy < imageH; iy++)
        for (int ix = 0; ix < imageW; ix++) {
            // Get the current pixel color
            int    pixel      = imageW * iy + ix;
            uchar4 pixelColor = dst[pixel];
            int    count      = 0;

            // Search for pixels out of tolerance surrounding the current pixel
            if (ix > 0)
                count += CheckColors(pixelColor, dst[pixel - 1]);

            if (ix + 1 < imageW)
                count += CheckColors(pixelColor, dst[pixel + 1]);

            if (iy > 0)
                count += CheckColors(pixelColor, dst[pixel - imageW]);

            if (iy + 1 < imageH)
                count += CheckColors(pixelColor, dst[pixel + imageW]);

            if (count) {
                // Calculate the location
                const T_ xPos = (T)ix * scale + xOff;
                const T_ yPos = (T)iy * scale + yOff;

                // Calculate the Mandelbrot index for the current location
                int m = CalcMandelbrot<T_>(xPos, yPos, xJParam, yJParam, crunch, isJulia);
                m     = m > 0 ? crunch - m : 0;

                // Convert the Mandelbrot index into a color
                uchar4 color;

                setColor(color, colors, m, animationFrame);

                // Output the pixel
                updatePixel(dst[pixel], color, frame);
            }
        }
} // RunMandelbrotGold1_

/* Implementation of exported functions */
void RunMandelbrotGold1(uchar4      *dst,
                        const int    imageW,
                        const int    imageH,
                        const int    crunch,
                        const float  xOff,
                        const float  yOff,
                        const float  xJParam,
                        const float  yJParam,
                        const float  scale,
                        const uchar4 colors,
                        const int    frame,
                        const int    animationFrame,
                        const bool   isJulia)
{
    runMandelbrotGold1<float, float>(
        dst, imageW, imageH, crunch, xOff, yOff, xJParam, yJParam, scale, colors, frame, animationFrame, isJulia);

} // RunMandelbrotGold1

void RunMandelbrotDSGold1(uchar4      *dst,
                          const int    imageW,
                          const int    imageH,
                          const int    crunch,
                          const double xOff,
                          const double yOff,
                          const double xJParam,
                          const double yJParam,
                          const double scale,
                          const uchar4 colors,
                          const int    frame,
                          const int    animationFrame,
                          const bool   isJulia)
{
    runMandelbrotGold1<double, dfloat>(
        dst, imageW, imageH, crunch, xOff, yOff, xJParam, yJParam, scale, colors, frame, animationFrame, isJulia);

} // RunMandelbrotDSGold1

void RunMandelbrotGold0(uchar4      *dst,
                        const int    imageW,
                        const int    imageH,
                        const int    crunch,
                        const float  xOff,
                        const float  yOff,
                        const float  xJParam,
                        const float  yJParam,
                        const float  scale,
                        const uchar4 colors,
                        const int    frame,
                        const int    animationFrame,
                        const bool   isJulia)
{
    runMandelbrotGold0<float, float>(
        dst, imageW, imageH, crunch, xOff, yOff, xJParam, yJParam, scale, colors, frame, animationFrame, isJulia);
} // RunMandelbrotGold0

void RunMandelbrotDSGold0(uchar4      *dst,
                          const int    imageW,
                          const int    imageH,
                          const int    crunch,
                          const double xOff,
                          const double yOff,
                          const double xJParam,
                          const double yJParam,
                          const double scale,
                          const uchar4 colors,
                          const int    frame,
                          const int    animationFrame,
                          const bool   isJulia)
{
    runMandelbrotGold0<double, dfloat>(
        dst, imageW, imageH, crunch, xOff, yOff, xJParam, yJParam, scale, colors, frame, animationFrame, isJulia);
} // RunMandelbrotDSGold0

/*dfloat operations implementation */

/* Construct a DS number equal to a double precision floating point number b*/
dfloat::dfloat(double b)
{
    val[0] = (float)b;
    val[1] = (float)(b - val[0]);
}

inline dfloat operator+(const dfloat &dsa, const dfloat &dsb)
{
    // Compute dsa + dsb using Knuth's trick.
    float t1 = dsa[0] + dsb[0];
    float e  = t1 - dsa[0];
    float t2 = ((dsb[0] - e) + (dsa[0] - (t1 - e))) + dsa[1] + dsb[1];

    // The result is t1 + t2, after normalization.
    e = t1 + t2;
    return dfloat(e, t2 - (e - t1));
}

inline dfloat operator-(const dfloat &dsa, const dfloat &dsb)
{
    // Compute dsa - dsb using Knuth's trick.
    float t1 = dsa[0] - dsb[0];
    float e  = t1 - dsa[0];
    float t2 = ((-dsb[0] - e) + (dsa[0] - (t1 - e))) + dsa[1] - dsb[1];

    // The result is t1 + t2, after normalization.
    e = t1 + t2;
    return dfloat(e, t2 - (e - t1));
}

inline dfloat operator*(const dfloat &dsa, const dfloat &dsb)
{
    // This splits dsa(1) and dsb(1) into high-order and low-order words.
    float c11 = dsa[0] * dsb[0];
    float c21 = dsa[0] * dsb[0] - c11;

    // Compute dsa[0] * dsb[1] + dsa[1] * dsb[0] (only high-order word is needed).
    float c2 = dsa[0] * dsb[1] + dsa[1] * dsb[0];

    // Compute (c11, c21) + c2 using Knuth's trick, also adding low-order product.
    float t1 = c11 + c2;
    float e  = t1 - c11;
    float t2 = ((c2 - e) + (c11 - (t1 - e))) + c21 + dsa[1] * dsb[1];

    // The result is t1 + t2, after normalization.
    e = t1 + t2;
    return dfloat(e, t2 - (e - t1));
}
