/*
 * Copyright 2008 Sandia Corporation.
 * Under the terms of Contract DE-AC04-94AL85000, there is a non-exclusive
 * license for use of this work by or on behalf of the
 * U.S. Government. Redistribution and use in source and binary forms, with
 * or without modification, are permitted provided that this Notice and any
 * statement of authorship are reproduced on all copies.
 */
// .SECTION Thanks
// Thanks to Janine Bennett, Philippe Pebay, and David Thompson from Sandia National Laboratories
// for implementing this test.

#include "vtkDataSetAttributes.h"
#include "vtkDoubleArray.h"
#include "vtkIdTypeArray.h"
#include "vtkKMeansStatistics.h"
#include "vtkMath.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkStringArray.h"
#include "vtkTable.h"
#include "vtkTimerLog.h"
#include "vtkUnsignedCharArray.h"

#include <limits>
#include <sstream>

//=============================================================================
int TestKMeansStatistics(int, char*[])
{
  int testStatus = 0;

  const int nDim = 4;
  int nVals = 50;
  int numberOfGhosts = 3;

  // Seed random number generator
  vtkMath::RandomSeed(static_cast<int>(vtkTimerLog::GetUniversalTime()));

  // Generate an input table that contains samples of mutually independent random variables over [0,
  // 1]
  vtkTable* inputData = vtkTable::New();
  vtkDoubleArray* doubleArray;

  vtkNew<vtkUnsignedCharArray> ghosts;
  ghosts->SetNumberOfValues(nVals);
  ghosts->SetName(vtkDataSetAttributes::GhostArrayName());
  ghosts->Fill(0);
  ghosts->SetValue(nVals / 2, 1);
  ghosts->SetValue(nVals / 3, 1);
  ghosts->SetValue(nVals / 4, 1);
  inputData->AddColumn(ghosts);

  int numComponents = 1;
  for (int c = 0; c < nDim; ++c)
  {
    std::ostringstream colName;
    colName << "coord " << c;
    doubleArray = vtkDoubleArray::New();
    doubleArray->SetNumberOfComponents(numComponents);
    doubleArray->SetName(colName.str().c_str());
    doubleArray->SetNumberOfTuples(nVals);

    double x;
    for (int r = 0; r < nVals; ++r)
    {
      // x = vtkMath::Gaussian();
      x = vtkMath::Random();
      doubleArray->SetValue(r, x);
    }

    doubleArray->SetValue(nVals / 2, std::numeric_limits<double>::quiet_NaN());
    doubleArray->SetValue(nVals / 3, std::numeric_limits<double>::quiet_NaN());
    doubleArray->SetValue(nVals / 4, std::numeric_limits<double>::quiet_NaN());

    inputData->AddColumn(doubleArray);
    doubleArray->Delete();
  }

  vtkTable* paramData = vtkTable::New();
  vtkIdTypeArray* paramCluster;
  vtkDoubleArray* paramArray;
  const int numRuns = 5;
  const int numClustersInRun[] = { 5, 2, 3, 4, 5 };
  paramCluster = vtkIdTypeArray::New();
  paramCluster->SetName("K");

  for (int curRun = 0; curRun < numRuns; curRun++)
  {
    for (int nInRun = 0; nInRun < numClustersInRun[curRun]; nInRun++)
    {
      paramCluster->InsertNextValue(numClustersInRun[curRun]);
    }
  }
  paramData->AddColumn(paramCluster);
  paramCluster->Delete();

  for (int c = 0; c < 5; ++c)
  {
    std::ostringstream colName;
    colName << "coord " << c;
    paramArray = vtkDoubleArray::New();
    paramArray->SetNumberOfComponents(numComponents);
    paramArray->SetName(colName.str().c_str());

    double x;
    for (int curRun = 0; curRun < numRuns; curRun++)
    {
      for (int nInRun = 0; nInRun < numClustersInRun[curRun]; nInRun++)
      {
        // x = vtkMath::Gaussian();
        x = vtkMath::Random();
        paramArray->InsertNextValue(x);
      }
    }
    paramData->AddColumn(paramArray);
    paramArray->Delete();
  }

  // Set k-means statistics algorithm and its input data port
  vtkKMeansStatistics* haruspex = vtkKMeansStatistics::New();

  // First verify that absence of input does not cause trouble
  cout << "## Verifying that absence of input does not cause trouble... ";
  haruspex->Update();
  cout << "done.\n";

  // Prepare first test with data
  haruspex->SetInputData(vtkStatisticsAlgorithm::INPUT_DATA, inputData);
  haruspex->SetColumnStatus(inputData->GetColumnName(0), 1);
  haruspex->SetColumnStatus(inputData->GetColumnName(2), 1);
  haruspex->SetColumnStatus("Testing", 1);
  haruspex->RequestSelectedColumns();
  haruspex->SetDefaultNumberOfClusters(3);

  // Test Learn and Derive options
  haruspex->SetLearnOption(true);
  haruspex->SetDeriveOption(true);
  haruspex->SetTestOption(false);
  haruspex->SetAssessOption(false);
  haruspex->Update();

  // Get output data and meta tables
  vtkMultiBlockDataSet* outputMetaDS = vtkMultiBlockDataSet::SafeDownCast(
    haruspex->GetOutputDataObject(vtkStatisticsAlgorithm::OUTPUT_MODEL));
  for (unsigned int b = 0; b < outputMetaDS->GetNumberOfBlocks(); ++b)
  {
    vtkTable* outputMeta = vtkTable::SafeDownCast(outputMetaDS->GetBlock(b));
    if (!b)
    {

      vtkIdType testIntValue = 0;
      for (vtkIdType r = 0; r < outputMeta->GetNumberOfRows(); r++)
      {
        testIntValue += outputMeta->GetValueByName(r, "Cardinality").ToInt();
      }

      cout << "## Computed clusters (cardinality: " << testIntValue << " / run):\n";

      if (testIntValue != nVals - numberOfGhosts)
      {
        vtkGenericWarningMacro("Sum of cluster cardinalities is incorrect: "
          << testIntValue << " != " << nVals - numberOfGhosts << ".");
        testStatus = 1;
      }
    }
    else
    {
      cout << "## Ranked cluster: "
           << "\n";
    }

    outputMeta->Dump();
    cout << "\n";
  }

  haruspex->SetInputData(vtkStatisticsAlgorithm::LEARN_PARAMETERS, paramData);
  cout << "## Testing with input table:"
       << "\n";

  paramData->Dump();
  cout << "\n";

  // Test Assess option only
  haruspex->SetLearnOption(true);
  haruspex->SetDeriveOption(true);
  haruspex->SetTestOption(false);
  haruspex->SetAssessOption(false);

  haruspex->Update();
  outputMetaDS = vtkMultiBlockDataSet::SafeDownCast(
    haruspex->GetOutputDataObject(vtkStatisticsAlgorithm::OUTPUT_MODEL));
  for (unsigned int b = 0; b < outputMetaDS->GetNumberOfBlocks(); ++b)
  {
    vtkTable* outputMeta = vtkTable::SafeDownCast(outputMetaDS->GetBlock(b));
    if (b == 0)
    {
      vtkIdType r = 0;
      vtkIdType testIntValue = 0;
      for (int curRun = 0; curRun < numRuns; curRun++)
      {
        testIntValue = 0;
        for (int nInRun = 0; nInRun < numClustersInRun[curRun]; nInRun++)
        {
          testIntValue += outputMeta->GetValueByName(r, "Cardinality").ToInt();
          r++;
        }
      }

      if (r != outputMeta->GetNumberOfRows())
      {
        vtkGenericWarningMacro("Inconsistency in number of rows: "
          << r << " != " << outputMeta->GetNumberOfRows() << ".");
        testStatus = 1;
      }

      cout << "## Computed clusters (cardinality: " << testIntValue << " / run):\n";

      if (testIntValue != nVals - numberOfGhosts)
      {
        vtkGenericWarningMacro("Sum of cluster cardinalities is incorrect: "
          << testIntValue << " != " << nVals - numberOfGhosts << ".");
        testStatus = 1;
      }
    }
    else
    {
      cout << "## Ranked cluster: "
           << "\n";
    }

    outputMeta->Dump();
    cout << "\n";
  }

  cout << "=================== ASSESS ==================== " << endl;
  vtkMultiBlockDataSet* paramsTables = vtkMultiBlockDataSet::New();
  paramsTables->ShallowCopy(outputMetaDS);

  haruspex->SetInputData(vtkStatisticsAlgorithm::INPUT_MODEL, paramsTables);

  // Test Assess option only (do not recalculate nor rederive a model)
  haruspex->SetLearnOption(false);
  haruspex->SetDeriveOption(false);
  haruspex->SetTestOption(false);
  haruspex->SetAssessOption(true);
  haruspex->Update();
  vtkTable* outputData = haruspex->GetOutput();
  outputData->Dump();
  paramsTables->Delete();
  paramData->Delete();
  inputData->Delete();
  haruspex->Delete();

  return testStatus;
}
