/* 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.
 */

#ifndef __BODYSYSTEM_H__
#define __BODYSYSTEM_H__

#include <algorithm>

enum NBodyConfig {
  NBODY_CONFIG_RANDOM,
  NBODY_CONFIG_SHELL,
  NBODY_CONFIG_EXPAND,
  NBODY_NUM_CONFIGS
};

enum BodyArray {
  BODYSYSTEM_POSITION,
  BODYSYSTEM_VELOCITY,
};

template <typename T>
struct vec3 {
  typedef float Type;
};  // dummy
template <>
struct vec3<float> {
  typedef float3 Type;
};
template <>
struct vec3<double> {
  typedef double3 Type;
};

template <typename T>
struct vec4 {
  typedef float Type;
};  // dummy
template <>
struct vec4<float> {
  typedef float4 Type;
};
template <>
struct vec4<double> {
  typedef double4 Type;
};

class string;

// BodySystem abstract base class
template <typename T>
class BodySystem {
 public:  // methods
  BodySystem(int numBodies) {}
  virtual ~BodySystem() {}

  virtual void loadTipsyFile(const std::string &filename) = 0;

  virtual void update(T deltaTime) = 0;

  virtual void setSoftening(T softening) = 0;
  virtual void setDamping(T damping) = 0;

  virtual T *getArray(BodyArray array) = 0;
  virtual void setArray(BodyArray array, const T *data) = 0;

  virtual unsigned int getCurrentReadBuffer() const = 0;

  virtual unsigned int getNumBodies() const = 0;

  virtual void synchronizeThreads() const {};

 protected:        // methods
  BodySystem() {}  // default constructor

  virtual void _initialize(int numBodies) = 0;
  virtual void _finalize() = 0;
};

inline float3 scalevec(float3 &vector, float scalar) {
  float3 rt = vector;
  rt.x *= scalar;
  rt.y *= scalar;
  rt.z *= scalar;
  return rt;
}

inline float normalize(float3 &vector) {
  float dist =
      sqrtf(vector.x * vector.x + vector.y * vector.y + vector.z * vector.z);

  if (dist > 1e-6) {
    vector.x /= dist;
    vector.y /= dist;
    vector.z /= dist;
  }

  return dist;
}

inline float dot(float3 v0, float3 v1) {
  return v0.x * v1.x + v0.y * v1.y + v0.z * v1.z;
}

inline float3 cross(float3 v0, float3 v1) {
  float3 rt;
  rt.x = v0.y * v1.z - v0.z * v1.y;
  rt.y = v0.z * v1.x - v0.x * v1.z;
  rt.z = v0.x * v1.y - v0.y * v1.x;
  return rt;
}

// utility function
template <typename T>
void randomizeBodies(NBodyConfig config, T *pos, T *vel, float *color,
                     float clusterScale, float velocityScale, int numBodies,
                     bool vec4vel) {
  switch (config) {
    default:
    case NBODY_CONFIG_RANDOM: {
      float scale = clusterScale * std::max<float>(1.0f, numBodies / (1024.0f));
      float vscale = velocityScale * scale;

      int p = 0, v = 0;
      int i = 0;

      while (i < numBodies) {
        float3 point;
        // const int scale = 16;
        point.x = rand() / (float)RAND_MAX * 2 - 1;
        point.y = rand() / (float)RAND_MAX * 2 - 1;
        point.z = rand() / (float)RAND_MAX * 2 - 1;
        float lenSqr = dot(point, point);

        if (lenSqr > 1) continue;

        float3 velocity;
        velocity.x = rand() / (float)RAND_MAX * 2 - 1;
        velocity.y = rand() / (float)RAND_MAX * 2 - 1;
        velocity.z = rand() / (float)RAND_MAX * 2 - 1;
        lenSqr = dot(velocity, velocity);

        if (lenSqr > 1) continue;

        pos[p++] = point.x * scale;  // pos.x
        pos[p++] = point.y * scale;  // pos.y
        pos[p++] = point.z * scale;  // pos.z
        pos[p++] = 1.0f;             // mass

        vel[v++] = velocity.x * vscale;  // pos.x
        vel[v++] = velocity.y * vscale;  // pos.x
        vel[v++] = velocity.z * vscale;  // pos.x

        if (vec4vel) vel[v++] = 1.0f;  // inverse mass

        i++;
      }
    } break;

    case NBODY_CONFIG_SHELL: {
      float scale = clusterScale;
      float vscale = scale * velocityScale;
      float inner = 2.5f * scale;
      float outer = 4.0f * scale;

      int p = 0, v = 0;
      int i = 0;

      while (i < numBodies)  // for(int i=0; i < numBodies; i++)
      {
        float x, y, z;
        x = rand() / (float)RAND_MAX * 2 - 1;
        y = rand() / (float)RAND_MAX * 2 - 1;
        z = rand() / (float)RAND_MAX * 2 - 1;

        float3 point = {x, y, z};
        float len = normalize(point);

        if (len > 1) continue;

        pos[p++] =
            point.x * (inner + (outer - inner) * rand() / (float)RAND_MAX);
        pos[p++] =
            point.y * (inner + (outer - inner) * rand() / (float)RAND_MAX);
        pos[p++] =
            point.z * (inner + (outer - inner) * rand() / (float)RAND_MAX);
        pos[p++] = 1.0f;

        x = 0.0f;  // * (rand() / (float) RAND_MAX * 2 - 1);
        y = 0.0f;  // * (rand() / (float) RAND_MAX * 2 - 1);
        z = 1.0f;  // * (rand() / (float) RAND_MAX * 2 - 1);
        float3 axis = {x, y, z};
        normalize(axis);

        if (1 - dot(point, axis) < 1e-6) {
          axis.x = point.y;
          axis.y = point.x;
          normalize(axis);
        }

        // if (point.y < 0) axis = scalevec(axis, -1);
        float3 vv = {(float)pos[4 * i], (float)pos[4 * i + 1],
                     (float)pos[4 * i + 2]};
        vv = cross(vv, axis);
        vel[v++] = vv.x * vscale;
        vel[v++] = vv.y * vscale;
        vel[v++] = vv.z * vscale;

        if (vec4vel) vel[v++] = 1.0f;

        i++;
      }
    } break;

    case NBODY_CONFIG_EXPAND: {
      float scale = clusterScale * numBodies / (1024.f);

      if (scale < 1.0f) scale = clusterScale;

      float vscale = scale * velocityScale;

      int p = 0, v = 0;

      for (int i = 0; i < numBodies;) {
        float3 point;

        point.x = rand() / (float)RAND_MAX * 2 - 1;
        point.y = rand() / (float)RAND_MAX * 2 - 1;
        point.z = rand() / (float)RAND_MAX * 2 - 1;

        float lenSqr = dot(point, point);

        if (lenSqr > 1) continue;

        pos[p++] = point.x * scale;   // pos.x
        pos[p++] = point.y * scale;   // pos.y
        pos[p++] = point.z * scale;   // pos.z
        pos[p++] = 1.0f;              // mass
        vel[v++] = point.x * vscale;  // pos.x
        vel[v++] = point.y * vscale;  // pos.x
        vel[v++] = point.z * vscale;  // pos.x

        if (vec4vel) vel[v++] = 1.0f;  // inverse mass

        i++;
      }
    } break;
  }

  if (color) {
    int v = 0;

    for (int i = 0; i < numBodies; i++) {
      // const int scale = 16;
      color[v++] = rand() / (float)RAND_MAX;
      color[v++] = rand() / (float)RAND_MAX;
      color[v++] = rand() / (float)RAND_MAX;
      color[v++] = 1.0f;
    }
  }
}

#endif  // __BODYSYSTEM_H__