/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkOpenGLVertexBufferObjectGroup.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkOpenGLVertexBufferObjectGroup.h"

#include <cassert>

#include "vtkObjectFactory.h"
#include "vtkDataArray.h"
#include "vtkOpenGLRenderWindow.h"
#include "vtkOpenGLVertexArrayObject.h"
#include "vtkOpenGLVertexBufferObject.h"
#include "vtkOpenGLVertexBufferObjectCache.h"
#include "vtkViewport.h"
#include "vtkShaderProgram.h"

// STL headers
#include <algorithm>

// ----------------------------------------------------------------------------
vtkStandardNewMacro(vtkOpenGLVertexBufferObjectGroup);

typedef std::map<std::string, vtkOpenGLVertexBufferObject *>::iterator vboIter;

// typedef std::map<const std::string, vtkDataArray *>::iterator arrayIter;
typedef std::map<std::string, std::vector<vtkDataArray *> >::iterator arrayIter;

// ----------------------------------------------------------------------------
vtkOpenGLVertexBufferObjectGroup::vtkOpenGLVertexBufferObjectGroup()
{
}

// ----------------------------------------------------------------------------
vtkOpenGLVertexBufferObjectGroup::~vtkOpenGLVertexBufferObjectGroup()
{
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    i->second->Delete();
  }
  this->UsedVBOs.clear();
}

int vtkOpenGLVertexBufferObjectGroup::GetNumberOfComponents(
  const char *attribute)
{
  vboIter i = this->UsedVBOs.find(attribute);
  if (i != this->UsedVBOs.end())
  {
    return i->second->GetNumberOfComponents();
  }
  return 0;
}

int vtkOpenGLVertexBufferObjectGroup::GetNumberOfTuples(
  const char *attribute)
{
  vboIter i = this->UsedVBOs.find(attribute);
  if (i != this->UsedVBOs.end())
  {
    return i->second->GetNumberOfTuples();
  }
  return 0;
}

vtkOpenGLVertexBufferObject *vtkOpenGLVertexBufferObjectGroup::GetVBO(
  const char *attribute)
{
  vboIter i = this->UsedVBOs.find(attribute);
  if (i != this->UsedVBOs.end())
  {
    return i->second;
  }
  return nullptr;
}

void vtkOpenGLVertexBufferObjectGroup::RemoveAttribute(
  const char *attribute)
{
  // empty array, means delete any existing entries
  arrayIter diter = this->UsedDataArrays.find(attribute);
  if (diter != this->UsedDataArrays.end())
  {
    vboIter viter = this->UsedVBOs.find(attribute);
    if (viter != this->UsedVBOs.end())
    {
      viter->second->UnRegister(this);
      this->UsedVBOs.erase(viter);
      this->Modified();
    }
    std::vector<vtkDataArray *> &vec = diter->second;
    for (size_t j = 0; j < vec.size(); j++)
    {
      if (vec[j])
      {
        vec[j]->Delete();
      }
    }
    diter->second.clear();
    this->UsedDataArrays.erase(diter);
  }
  return;
}

void vtkOpenGLVertexBufferObjectGroup::CacheDataArray(
  const char *attribute,
  vtkDataArray *da,
  vtkViewport *vp,
  int destType)
{
  vtkOpenGLRenderWindow *renWin =
    vtkOpenGLRenderWindow::SafeDownCast(vp->GetVTKWindow());
  vtkOpenGLVertexBufferObjectCache *cache = renWin->GetVBOCache();
  this->CacheDataArray(attribute, da, cache, destType);
}

void vtkOpenGLVertexBufferObjectGroup::CacheDataArray(
  const char *attribute,
  vtkDataArray *da,
  vtkOpenGLVertexBufferObjectCache *cache,
  int destType)
{
  // empty array, means delete any existing entries
  if (!da || da->GetNumberOfTuples() == 0)
  {
    this->RemoveAttribute(attribute);
    return;
  }

  // OK we have a data array
  arrayIter diter = this->UsedDataArrays.find(attribute);

  // if a different array is already setup for this attribute
  // clear it out. Remember that UsedArrays is cleared
  // after upload. So if there is another array here
  // it means the same attrribute has been set twice.
  // so we delete the prior setting, last one wins.
  if (diter != this->UsedDataArrays.end() &&
      (diter->second.size() != 1 || diter->second[0] != da))
  {
    for (size_t j = 0; j < diter->second.size(); j++)
    {
      if (diter->second[j])
      {
        diter->second[j]->Delete();
      }
    }
    diter->second.clear();
  }

  // make sure we add this DA to our list of arrays
  da->Register(this);
  std::vector<vtkDataArray *> &vec = this->UsedDataArrays[attribute];
  vec.push_back(da);

  // get the VBO for this DA
  vtkOpenGLVertexBufferObject* vbo = cache->GetVBO(da, destType);
  vboIter viter = this->UsedVBOs.find(attribute);

  // if this VBO is the same as previous for this attribute
  // then just return
  if (viter != this->UsedVBOs.end() && viter->second == vbo)
  {
    vbo->UnRegister(this); // GetVBO increments the ref count
    return;
  }

  this->Modified();

  // if this VBO is different from the prior VBO for this
  // attribute then free the prior VBO
  if (viter != this->UsedVBOs.end())
  {
    viter->second->UnRegister(this);
    this->UsedVBOs.erase(viter);
  }

  // store the VBO and upload
  this->UsedVBOs[attribute] = vbo;
}

void vtkOpenGLVertexBufferObjectGroup::AppendDataArray(
  const char *attribute,
  vtkDataArray *da,
  int destType)
{
  if (!da)
  {
    return;
  }

  std::vector<vtkDataArray *> &arrays = this->UsedDataArrays[attribute];
  if (std::find(arrays.begin(), arrays.end(), da) == arrays.end())
  {
    da->Register(this);
    arrays.push_back(da);

    // make sure we have a VBO for this array
    // we do not use the cache when appending
    if (this->UsedVBOs.find(attribute) == this->UsedVBOs.end())
    {
      vtkOpenGLVertexBufferObject* vbo = vtkOpenGLVertexBufferObject::New();
      vbo->SetDataType(destType);
      this->UsedVBOs[attribute] = vbo;
    }
  }
}

void vtkOpenGLVertexBufferObjectGroup::ReleaseGraphicsResources(vtkWindow *)
{
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    i->second->Delete();
  }
  this->UsedVBOs.clear();
}

void vtkOpenGLVertexBufferObjectGroup::AddAllAttributesToVAO(
  vtkShaderProgram *program,
  vtkOpenGLVertexArrayObject *vao)
{
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    std::string dataShaderName = i->first;
    if (program->IsAttributeUsed(dataShaderName.c_str()))
    {
      vtkOpenGLVertexBufferObject* vbo = i->second;
      if (!vao->AddAttributeArray(
        program, vbo, dataShaderName,
        0, // offset see assert later in this file
        (vbo->GetDataType() == VTK_UNSIGNED_CHAR))) //TODO: fix tweak. true for colors.
      {
        vtkErrorMacro(<< "Error setting '" << dataShaderName << "' in shader VAO.");
      }
    }
  }
}

void vtkOpenGLVertexBufferObjectGroup::ClearAllDataArrays()
{
  for (arrayIter i = this->UsedDataArrays.begin();
       i != this->UsedDataArrays.end(); ++i)
  {
    std::vector<vtkDataArray *> &vec = i->second;
    for (size_t j = 0; j < vec.size(); j++)
    {
      if (vec[j])
      {
        vec[j]->Delete();
      }
    }
    i->second.clear();
  }
  this->UsedDataArrays.clear();
}

void vtkOpenGLVertexBufferObjectGroup::ClearAllVBOs()
{
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    i->second->Delete();
  }
  this->UsedVBOs.clear();
}

void vtkOpenGLVertexBufferObjectGroup::BuildAllVBOs(vtkViewport *vp)
{
  vtkOpenGLRenderWindow *renWin =
    vtkOpenGLRenderWindow::SafeDownCast(vp->GetVTKWindow());
  vtkOpenGLVertexBufferObjectCache *cache = renWin->GetVBOCache();
  this->BuildAllVBOs(cache);
}

// ----------------------------------------------------------------------------
void vtkOpenGLVertexBufferObjectGroup::BuildAllVBOs(
  vtkOpenGLVertexBufferObjectCache *)
{
  // free any VBOs for unused attributes
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end();)
  {
    if (this->UsedDataArrays.find(i->first) == this->UsedDataArrays.end())
    {
      i->second->UnRegister(this);
      vboIter toErase = i;
      ++i;
      this->UsedVBOs.erase(toErase);
    }
    else
    {
      ++i;
    }
  }

  // we always upload appended data :-(
  for (arrayIter i = this->UsedDataArrays.begin();
       i != this->UsedDataArrays.end(); ++i)
  {
    std::string attribute = i->first;
    std::vector<vtkDataArray *> &vec = i->second;
    vtkOpenGLVertexBufferObject *vbo = this->UsedVBOs[attribute];

    if (vec.size() > 1)
    {
      for (size_t j = 0; j < vec.size(); j++)
      {
        // only append if needed?
        vbo->AppendDataArray(vec[j]);
      }
      vbo->UploadVBO();
    }
  }

  // for everything else we upload based on mtimes
  for (arrayIter i = this->UsedDataArrays.begin();
       i != this->UsedDataArrays.end(); ++i)
  {
    std::string attribute = i->first;
    std::vector<vtkDataArray *> &vec = i->second;
    vtkOpenGLVertexBufferObject *vbo = this->UsedVBOs[attribute];

    if (vec.size() == 1 &&
        vec[0]->GetMTime() > vbo->GetUploadTime())
    {
      vbo->UploadDataArray(vec[0]);
    }
  }

  // Upload updated VBOs
  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    vtkOpenGLVertexBufferObject* vbo = i->second;
    if(vbo->GetMTime() > vbo->GetUploadTime())
    {
      vbo->UploadVBO();
    }
  }

  this->ClearAllDataArrays();
}

//----------------------------------------------------------------------------
vtkMTimeType vtkOpenGLVertexBufferObjectGroup::GetMTime()
{
  vtkMTimeType mTime = this->Superclass::GetMTime();

  for (vboIter i = this->UsedVBOs.begin(); i != this->UsedVBOs.end(); ++i)
  {
    vtkMTimeType time = i->second->GetMTime();
    if (time > mTime)
    {
      mTime = time;
    }
  }

  return mTime;
}

// ----------------------------------------------------------------------------
void vtkOpenGLVertexBufferObjectGroup::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os,indent);
}