#include "radiosity.h"

#include <omp.h>

inline unsigned long mwc()
{
  static unsigned long x[8] = { 123456789, 123456789, 123456789, 123456789, 123456789, 123456789, 123456789, 123456789 };
  static unsigned long y[8] = { 362436069, 362436069, 362436069, 362436069, 362436069, 362436069, 362436069, 362436069 }; 
  static unsigned long z[8] = { 77465321, 77465321, 77465321, 77465321, 77465321, 77465321, 77465321, 77465321 }; 
  static unsigned long c[8] = { 13579, 13579, 13579, 13579, 13579, 13579, 13579, 13579 };
  
  int tid = omp_get_thread_num();
  unsigned long long t = 916905990LL * x[tid] + c[tid];
  
  x[tid] = y[tid];
  y[tid] = z[tid];
  c[tid] = (t >> 32);
  z[tid] = (t & 0xffffffff);

  return z[tid];
}

inline float rnd_uniform(float s, float e)
{
  //float t = (float)rand() / 32767.0f;
  const float invMul = 1.0f / 16777214.0f;
  float t = invMul*(float)(mwc() & 0x00FFFFFF);
  return s + t*(e - s);
}

template<typename T>
float3 _dirX(T data)
{
  switch (data.axis)
  {
  case PlaneType::PLUS_X:
  case PlaneType::MINUS_X:
    return float3(0, 0, 1);

  case PlaneType::PLUS_Y:
  case PlaneType::MINUS_Y:
    return float3(1, 0, 0);

  case PlaneType::PLUS_Z:
  case PlaneType::MINUS_Z:
    return float3(1, 0, 0);
  };
}

template<typename T>
float3 _dirY(T data)
{
  switch (data.axis)
  {
  case PlaneType::PLUS_X:
  case PlaneType::MINUS_X:
    return float3(0, 1, 0);

  case PlaneType::PLUS_Y:
  case PlaneType::MINUS_Y:
    return float3(0, 0, 1);

  case PlaneType::PLUS_Z:
  case PlaneType::MINUS_Z:
    return float3(0, 1, 0);
  };
}

float3 QuadAA::boxMin() const { return centerPos - size.x*dirX() - size.y*dirY(); }
float3 QuadAA::boxMax() const { return centerPos + size.x*dirX() + size.y*dirY(); }

float3 QuadAA::dirX() const { return _dirX<QuadAA>(*this); }
float3 QuadAA::dirY() const { return _dirY<QuadAA>(*this); }

float3 QuadAA::normal() const
{
  switch (axis)
  {
  case PlaneType::PLUS_X:   return float3(+1, 0, 0);
  case PlaneType::MINUS_X:  return float3(-1, 0, 0);
  
  case PlaneType::PLUS_Y:   return float3(0, +1, 0);
  case PlaneType::MINUS_Y:  return float3(0, -1, 0);

  case PlaneType::PLUS_Z:   return float3(0, 0, +1);
  case PlaneType::MINUS_Z:  return float3(0, 0, -1);
  };
}

float3 PatchGeom::dirX() const { return _dirX<PatchGeom>(*this); }
float3 PatchGeom::dirY() const { return _dirY<PatchGeom>(*this); }

float PatchGeom::sizeX() const
{
  switch (axis)
  {
  case PlaneType::PLUS_X:
  case PlaneType::MINUS_X:
    return vmax.z - vmin.z; // (0, 0, 1);

  case PlaneType::PLUS_Y:
  case PlaneType::MINUS_Y:
    return vmax.x - vmin.x; // (1, 0, 0);

  case PlaneType::PLUS_Z:
  case PlaneType::MINUS_Z:
    return vmax.x - vmin.x; //(1, 0, 0);
  };
}

float PatchGeom::sizeY() const
{
  switch (axis)
  {
  case PlaneType::PLUS_X:
  case PlaneType::MINUS_X:
    return vmax.y - vmin.y; // (0, 1, 0);

  case PlaneType::PLUS_Y:
  case PlaneType::MINUS_Y:
    return vmax.z - vmin.z; // (0, 0, 1);

  case PlaneType::PLUS_Z:
  case PlaneType::MINUS_Z:
    return vmax.y - vmin.y; // (0, 1, 0);
  };
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

std::vector<PatchGeom> subdivideQuads(const std::vector<QuadAA>& a_quads)
{
  if (a_quads.size() == 0)
    return std::vector<PatchGeom>();

  int a_approxRes = a_quads[0].subdivs;

  std::vector<PatchGeom> patches(0);
  patches.reserve(a_approxRes*a_approxRes*a_quads.size());

  for (auto q : a_quads)
  {
    float patchSizeX = (q.size.x*2.0f) / float(q.subdivs);
    float patchSizeY = (q.size.y*2.0f) / float(q.subdivs);

    float3 bMin = q.boxMin();
    float3 bMax = q.boxMax();

    float3 bOx = q.dirX();
    float3 bOy = q.dirY();

    float3 norm = q.normal();

    for (int y = 0; y < q.subdivs; y++)
    {
      for (int x = 0; x < q.subdivs; x++)
      {
        PatchGeom p;

        p.vmin = bMin + float(y+0)*patchSizeY*bOy + float(x+0)*patchSizeX*bOx;
        p.vmax = bMin + float(y+1)*patchSizeY*bOy + float(x+1)*patchSizeX*bOx;

        p.normal       = norm;
        p.surfaceArea  = patchSizeX*patchSizeY;
        p.axis         = q.axis;
        p.color        = q.diffuseColor;
        p.diffuseColor = q.diffuseColor;

        patches.push_back(p);
      }
    }
  }

  return patches;
}

float V(float3 x0, float3 x1, const std::vector<QuadAA>& a_quads) // visibility func, perform ray tracing
{
  float3 ray_dir = normalize(x1 - x0);
  float3 ray_pos = x0 + ray_dir*1e-5f;

  float3 dirInv(1.0f / fmax(ray_dir.x, 1e-6f), 
                1.0f / fmax(ray_dir.y, 1e-6f), 
                1.0f / fmax(ray_dir.z, 1e-6f));
  
  float tmin = 100000000.0f;
  float tmax = length(x1 - x0);

  for (auto q : a_quads)
  {
    float3 boxMin   = q.boxMin();
    float3 boxMax   = q.boxMax();
    float2 fullSize = q.size*2.0f;
 
    float  t, hitX, hitY;
    float3 hit;

    switch (q.axis)
    {
      case PlaneType::PLUS_X:
      case PlaneType::MINUS_X:
      {
        t    = dirInv.x*(boxMin.x - ray_pos.x);
        hit  = ray_pos + t*ray_dir;

        hitX = hit.z - boxMin.z;
        hitY = hit.y - boxMin.y;
      }
      break;
      
      case PlaneType::PLUS_Y:
      case PlaneType::MINUS_Y:
      {
        t    = dirInv.y*(boxMin.y - ray_pos.y);
        hit  = ray_pos + t*ray_dir;

        hitX = hit.x - boxMin.x;
        hitY = hit.z - boxMin.z;
      }
      break;
      
      case PlaneType::PLUS_Z:
      case PlaneType::MINUS_Z:
      {
        t    = dirInv.z*(boxMin.z - ray_pos.z);
        hit  = ray_pos + t*ray_dir;

        hitX = hit.x - boxMin.x;
        hitY = hit.y - boxMin.y;
      }
      break;

    };

    if (0.0f <= hitX && hitX <= fullSize.x && 0.0f <= hitY && hitY <= fullSize.y && t > 1e-5f && t + 1e-5f < tmax)
      tmin = fmin(tmin, t);
  }

  if (tmin > 1e-4f && tmin + 1e-4f < tmax)
    return 0.0f;
  else
    return 1.0f;
}

float integrateFormFactor(PatchGeom a, PatchGeom b, const std::vector<QuadAA>& a_quads)
{
  if (length(a.normal - b.normal) < 1e-5f)
    return 0.0f;

  float3 aCenter = 0.5f*(a.vmin + a.vmax);
  float3 bCenter = 0.5f*(a.vmin + a.vmax);

  float3 ab = normalize(bCenter - aCenter);
  float3 ba = ab*(-1.0f);

  if (dot(ab, a.normal) < 0.0f || dot(ba, b.normal) < 0.0f)
    return 0.0f;

  float aSizeX = a.sizeX();
  float aSizeY = a.sizeY();

  float bSizeX = b.sizeX();
  float bSizeY = b.sizeY();

  float summ = 0.0f;

  const int iters   = 4;
  const int subdivs = 4;
  const float scale = 1.0f / float(subdivs);

  for (int x0 = 0; x0 < subdivs; x0++)           // stratified sampling
  {
    for (int y0 = 0; y0 < subdivs; y0++)
    {
      for (int x1 = 0; x1 < subdivs; x1++)
      {
        for (int y1 = 0; y1 < subdivs; y1++)
        {
          float offx0 = float(x0)*scale;
          float offy0 = float(y0)*scale;

          float offx1 = float(x1)*scale;
          float offy1 = float(y1)*scale;

          for (int i = 0; i < iters; i++)
          {
            float t0x = offx0 + rnd_uniform(0.0f, 1.0f)*scale;
            float t0y = offy0 + rnd_uniform(0.0f, 1.0f)*scale;

            float t1x = offx1 + rnd_uniform(0.0f, 1.0f)*scale;
            float t1y = offy1 + rnd_uniform(0.0f, 1.0f)*scale;

            float3 aPos = a.vmin + t0x*a.dirX()*aSizeX + t0y*a.dirY()*aSizeY;
            float3 bPos = b.vmin + t1x*b.dirX()*bSizeX + t1y*b.dirY()*bSizeY;

            float3 ab = normalize(bPos - aPos);
            float  r = length(bPos - aPos);

            float cosThetaI = fmax(+dot(ab, a.normal), 0.0f);
            float cosThetaJ = fmax(-dot(ab, b.normal), 0.0f);

            float fval = V(aPos, bPos, a_quads)*cosThetaI*cosThetaJ / fmax(3.141592654f*r*r, 1e-5f);

            if (isfinite(fval))
              summ += fval;
          }
        }
      }
    }
  }

  return (summ / float(iters*subdivs*subdivs*subdivs*subdivs))*a.surfaceArea*b.surfaceArea;
  
}


std::vector<float> computeFormFactors(const std::vector<PatchGeom>& a_patches, const std::vector<QuadAA>& a_quads)
{
  std::vector<float> formFactors(a_patches.size()*a_patches.size());

  auto w = a_patches.size();

  int precisionOld = std::cout.precision(2);
  std::cout << std::endl;

  for (int i = 0; i < a_patches.size(); i++)
  {
    #pragma omp parallel for num_threads(8)
    for (int j = i; j < a_patches.size(); j++)
    {
      formFactors[i*w + j] = integrateFormFactor(a_patches[i], a_patches[j], a_quads);
      formFactors[j*w + i] = formFactors[i*w + j];
    }

    if (i % 10 == 0)
      std::cout << "computing form factors : " << 100.0f*float(i) / float(a_patches.size()) << "% \r";
  }

  std::cout << std::endl;
  std::cout.precision(precisionOld);

  for (auto i = 0; i < a_patches.size(); i++)
    for (auto j = 0; j < a_patches.size(); j++)
      formFactors[i*w + j] /= a_patches[i].surfaceArea;

  return formFactors;
}

std::vector<float3> radiosityBruteForce(const std::vector<PatchGeom>& a_patchesG, const std::vector<float>& a_formFactors, const std::vector<EmissivePatch>& a_lightSources, int a_numPasses)
{
  std::vector<PatchRadiosity> patchesR(a_patchesG.size());

  // init phase
  //
  for (auto i = 0; i < patchesR.size(); i++)
  {
    PatchRadiosity& p = patchesR[i];

    p.emmision    = float3(0, 0, 0); 
    p.deposit     = float3(0, 0, 0);
    p.excident    = float3(0, 0, 0);
    p.incident    = float3(0, 0, 0);
    p.reflectance = a_patchesG[i].diffuseColor;
  }

  for (auto light : a_lightSources)
    patchesR[light.patchId].emmision = light.color; // if current patch is light then emmision is not black

  for (PatchRadiosity& p : patchesR)
  {
    p.excident = p.emmision;
    p.emmision = float3(0, 0, 0);
    p.deposit += p.excident;
  }

  // do radiosity
  //
  for (int pass = 0; pass < a_numPasses; pass++)
  {
    #pragma omp parallel for
    for (int i = 0; i < patchesR.size(); i++)  // each patch collects light from the scene
    {
      patchesR[i].incident = float3(0, 0, 0);
      auto offset = patchesR.size()*i;

      for (auto j = 0; j < patchesR.size(); j++)
        patchesR[i].incident += patchesR[j].excident*a_formFactors[offset + j];
    }

    for (PatchRadiosity& patch : patchesR)
    {
      patch.excident = (patch.incident*patch.reflectance);
      patch.deposit += patch.excident;
    }
  }

  std::vector<float3> finalColors(patchesR.size());

  for (auto i = 0; i < patchesR.size(); i++)
    finalColors[i] = patchesR[i].deposit;

  return finalColors; // colors
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

void checkFormFactors(std::vector<PatchGeom>& a_patches, std::vector<float>& a_formFactors, int a_w)
{
  for (int i = 0; i < a_w; i++)
  {
    float summ = 0.0f;
    int offset = a_w*i;

    for (int j = 0; j < a_w; j++)
      summ += a_formFactors[offset + j];
    
    if (summ >= 1.1f || summ < 0.9f)
    {
      std::cout << "! summ of form factors [" << i << "] = " << summ << std::endl;
      a_patches[i].color = float3(1, 0, 0);
    }
    else
      a_patches[i].color = float3(0, 1, 0);
  }
}
