// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-FileCopyrightText: Copyright 2008 Sandia Corporation
// SPDX-License-Identifier: LicenseRef-BSD-3-Clause-Sandia-USGov
// .SECTION Thanks
// Thanks to Philippe Pebay and David Thompson from Sandia National Laboratories
// for implementing this test.
// Test added for Robust PCA by Tristan Coulange, Kitware SAS 2013

#include "vtkDoubleArray.h"
#include "vtkMathUtilities.h"
#include "vtkNew.h"
#include "vtkOrderStatistics.h"
#include "vtkPCAStatistics.h"
#include "vtkSmartPointer.h"
#include "vtkStatisticalModel.h"
#include "vtkStringArray.h"
#include "vtkTable.h"
#include "vtkTestUtilities.h"

#include "vtksys/SystemTools.hxx"

#include <iostream>

int TestPCA(int argc, char* argv[]);
int TestPCARobust(int argc, char* argv[]);
int TestPCAPart(int argc, char* argv[], bool RobustPCA);
int TestPCARobust2();
int TestEigen();

//=============================================================================
int TestPCAStatistics(int argc, char* argv[])
{
  int result = EXIT_SUCCESS;

  result |= TestPCA(argc, argv);
  result |= TestPCARobust(argc, argv);
  result |= TestPCARobust2();
  result |= TestEigen();

  if (result == EXIT_FAILURE)
  {
    std::cout << "FAILURE" << std::endl;
  }
  else
  {
    std::cout << "SUCCESS" << std::endl;
  }

  return result;
}

//=============================================================================
int TestPCA(int argc, char* argv[])
{
  return TestPCAPart(argc, argv, false);
}

//=============================================================================
int TestPCARobust(int argc, char* argv[])
{
  return TestPCAPart(argc, argv, true);
}

//=============================================================================
int TestPCARobust2()
{
  constexpr int nVals = 7;
  double mingledData[] = {
    0., 1.,  //
    1., 1.,  //
    2., 1.,  //
    3., 1.,  //
    4., 1.,  //
    5., 1.,  //
    10., 10. //
  };

  constexpr char m0Name[] = "M0";
  vtkNew<vtkDoubleArray> dataset1Arr;
  dataset1Arr->SetNumberOfComponents(1);
  dataset1Arr->SetName(m0Name);

  constexpr char m1Name[] = "M1";
  vtkNew<vtkDoubleArray> dataset2Arr;
  dataset2Arr->SetNumberOfComponents(1);
  dataset2Arr->SetName(m1Name);

  for (int i = 0; i < nVals; ++i)
  {
    dataset1Arr->InsertNextValue(mingledData[i * 2]);
    dataset2Arr->InsertNextValue(mingledData[i * 2 + 1]);
  }

  vtkNew<vtkTable> datasetTable;
  datasetTable->AddColumn(dataset1Arr);
  datasetTable->AddColumn(dataset2Arr);

  // Set PCA statistics algorithm and its input data port
  vtkNew<vtkPCAStatistics> pcas;

  // Prepare first test with data
  pcas->SetInputData(vtkStatisticsAlgorithm::INPUT_DATA, datasetTable);
  pcas->MedianAbsoluteDeviationOn();

  // -- Select Column Pairs of Interest ( Learn Mode ) --
  pcas->SetColumnStatus(m0Name, 1);
  pcas->SetColumnStatus(m1Name, 1);

  // Test all options but Assess
  pcas->SetLearnOption(true);
  pcas->SetDeriveOption(true);
  pcas->SetTestOption(true);
  pcas->SetAssessOption(true);
  pcas->Update();

  vtkTable* outputData = pcas->GetOutput();

  double res[] = { -3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 7.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 9.0 };

  for (vtkIdType j = 0; j < 2; j++)
  {
    for (vtkIdType i = 0; i < 7; i++)
    {
      if (outputData->GetValue(i, j + 2) != res[j * outputData->GetNumberOfRows() + i])
      {
        return EXIT_FAILURE;
      }
    }
  }

  return EXIT_SUCCESS;
}

//=============================================================================
int TestPCAPart(int argc, char* argv[], bool robustPCA)
{
  char* normScheme = vtkTestUtilities::GetArgOrEnvOrDefault(
    "-normalize-covariance", argc, argv, "VTK_NORMALIZE_COVARIANCE", "None");
  int testStatus = 0;

  /* */
  double mingledData[] = {
    46, 45, //
    47, 49, //
    46, 47, //
    46, 46, //
    47, 46, //
    47, 49, //
    49, 49, //
    47, 45, //
    50, 50, //
    46, 46, //
    51, 50, //
    48, 48, //
    52, 54, //
    48, 47, //
    52, 52, //
    49, 49, //
    53, 54, //
    50, 50, //
    53, 54, //
    50, 52, //
    53, 53, //
    50, 51, //
    54, 54, //
    49, 49, //
    52, 52, //
    50, 51, //
    52, 52, //
    49, 47, //
    48, 48, //
    48, 50, //
    46, 48, //
    47, 47  //
  };
  int nVals = 32;

  constexpr char m0Name[] = "M0";
  vtkDoubleArray* dataset1Arr = vtkDoubleArray::New();
  dataset1Arr->SetNumberOfComponents(1);
  dataset1Arr->SetName(m0Name);

  constexpr char m1Name[] = "M1";
  vtkDoubleArray* dataset2Arr = vtkDoubleArray::New();
  dataset2Arr->SetNumberOfComponents(1);
  dataset2Arr->SetName(m1Name);

  constexpr char m2Name[] = "M2";
  vtkDoubleArray* dataset3Arr = vtkDoubleArray::New();
  dataset3Arr->SetNumberOfComponents(1);
  dataset3Arr->SetName(m2Name);

  for (int i = 0; i < nVals; ++i)
  {
    int ti = i << 1;
    dataset1Arr->InsertNextValue(mingledData[ti]);
    dataset2Arr->InsertNextValue(mingledData[ti + 1]);
    dataset3Arr->InsertNextValue(i != 12 ? -1. : -1.001);
  }

  vtkTable* datasetTable = vtkTable::New();
  datasetTable->AddColumn(dataset1Arr);
  dataset1Arr->Delete();
  datasetTable->AddColumn(dataset2Arr);
  dataset2Arr->Delete();
  datasetTable->AddColumn(dataset3Arr);
  dataset3Arr->Delete();

  // Set PCA statistics algorithm and its input data port
  vtkPCAStatistics* pcas = vtkPCAStatistics::New();
  pcas->SetMedianAbsoluteDeviation(robustPCA);

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

  // Prepare first test with data
  pcas->SetInputData(vtkStatisticsAlgorithm::INPUT_DATA, datasetTable);
  pcas->SetNormalizationSchemeByName(normScheme);
  pcas->SetBasisSchemeByName("FixedBasisEnergy");
  pcas->SetFixedBasisEnergy(1. - 1e-8);

  datasetTable->Delete();

  // -- Select Column Pairs of Interest ( Learn Mode ) --
  pcas->SetColumnStatus(m0Name, 1);
  pcas->SetColumnStatus(m1Name, 1);
  pcas->RequestSelectedColumns();
  pcas->ResetAllColumnStates();
  pcas->SetColumnStatus(m0Name, 1);
  pcas->SetColumnStatus(m1Name, 1);
  pcas->SetColumnStatus(m2Name, 1);
  pcas->SetColumnStatus(m2Name, 0);
  pcas->SetColumnStatus(m2Name, 1);
  pcas->RequestSelectedColumns();
  pcas->RequestSelectedColumns(); // Try a duplicate entry. This should have no effect.
  pcas->SetColumnStatus(m0Name, 0);
  pcas->SetColumnStatus(m2Name, 0);
  pcas->SetColumnStatus("Metric 3",
    1); // An invalid name. This should result in a request for metric 1's self-correlation.
  // pcas->RequestSelectedColumns(); will get called in RequestData()

  // Test all options but Assess
  pcas->SetLearnOption(true);
  pcas->SetDeriveOption(true);
  pcas->SetTestOption(true);
  pcas->SetAssessOption(false);
  pcas->Update();

  auto* outputMetaDS = pcas->GetOutputModel();
  vtkTable* outputTest = pcas->GetOutput(vtkStatisticsAlgorithm::OUTPUT_TEST);

  std::cout << "## Calculated the following statistics for data set:\n";
  vtkTable* primaryTab = outputMetaDS->GetTable(vtkStatisticalModel::Learned, 0);
  std::cout << "Primary Statistics\n";
  primaryTab->Dump();
  for (int b = 0; b < outputMetaDS->GetNumberOfTables(vtkStatisticalModel::Derived); ++b)
  {
    vtkTable* derivedTab = outputMetaDS->GetTable(vtkStatisticalModel::Derived, b);
    std::cout << "Derived Statistics " << b << "\n";
    derivedTab->Dump();
  }

  // Check some results of the Test option
  std::cout << "\n## Calculated the following Jarque-Bera-Srivastava statistics for pseudo-random "
               "variables (n="
            << nVals;

#ifdef USE_GNU_R
  int nNonGaussian = 1;
  int nRejected = 0;
  double alpha = .01;

  std::cout << ", null hypothesis: binormality, significance level=" << alpha;
#endif // USE_GNU_R

  std::cout << "):\n";

  // Loop over Test table
  for (vtkIdType r = 0; r < outputTest->GetNumberOfRows(); ++r)
  {
    std::cout << "   ";
    for (int i = 0; i < outputTest->GetNumberOfColumns(); ++i)
    {
      std::cout << outputTest->GetColumnName(i) << "=" << outputTest->GetValue(r, i).ToString()
                << "  ";
    }

#ifdef USE_GNU_R
    // Check if null hypothesis is rejected at specified significance level
    double p = outputTest->GetValueByName(r, "P").ToDouble();
    // Must verify that p value is valid (it is set to -1 if R has failed)
    if (p > -1 && p < alpha)
    {
      std::cout << "N.H. rejected";

      ++nRejected;
    }
#endif // USE_GNU_R

    std::cout << "\n";
  }

#ifdef USE_GNU_R
  if (nRejected < nNonGaussian)
  {
    vtkGenericWarningMacro("Rejected only "
      << nRejected << " null hypotheses of binormality whereas " << nNonGaussian
      << " variable pairs are not Gaussian");
    testStatus = 1;
  }
#endif // USE_GNU_R

  // Test Assess option
  auto* paramsTables = vtkStatisticalModel::New();
  paramsTables->ShallowCopy(outputMetaDS);

  pcas->SetInputData(vtkStatisticsAlgorithm::INPUT_MODEL, paramsTables);
  paramsTables->Delete();

  // Test Assess only (Do not recalculate nor rederive nor retest a model)
  pcas->SetLearnOption(false);
  pcas->SetDeriveOption(false);
  pcas->SetTestOption(false);
  pcas->SetAssessOption(true);
  pcas->Update();

  std::cout << "\n## Assessment results:\n";
  vtkTable* outputData = pcas->GetOutput();
  outputData->Dump();

  pcas->Delete();
  delete[] normScheme;

  return testStatus;
}

int TestEigen()
{
  constexpr char m0Name[] = "M0";
  vtkSmartPointer<vtkDoubleArray> dataset1Arr = vtkSmartPointer<vtkDoubleArray>::New();
  dataset1Arr->SetNumberOfComponents(1);
  dataset1Arr->SetName(m0Name);
  dataset1Arr->InsertNextValue(0);
  dataset1Arr->InsertNextValue(1);
  dataset1Arr->InsertNextValue(0);

  constexpr char m1Name[] = "M1";
  vtkSmartPointer<vtkDoubleArray> dataset2Arr = vtkSmartPointer<vtkDoubleArray>::New();
  dataset2Arr->SetNumberOfComponents(1);
  dataset2Arr->SetName(m1Name);
  dataset2Arr->InsertNextValue(0);
  dataset2Arr->InsertNextValue(0);
  dataset2Arr->InsertNextValue(1);

  constexpr char m2Name[] = "M2";
  vtkSmartPointer<vtkDoubleArray> dataset3Arr = vtkSmartPointer<vtkDoubleArray>::New();
  dataset3Arr->SetNumberOfComponents(1);
  dataset3Arr->SetName(m2Name);
  dataset3Arr->InsertNextValue(0);
  dataset3Arr->InsertNextValue(0);
  dataset3Arr->InsertNextValue(0);

  vtkSmartPointer<vtkTable> datasetTable = vtkSmartPointer<vtkTable>::New();
  datasetTable->AddColumn(dataset1Arr);
  datasetTable->AddColumn(dataset2Arr);
  datasetTable->AddColumn(dataset3Arr);

  vtkSmartPointer<vtkPCAStatistics> pcaStatistics = vtkSmartPointer<vtkPCAStatistics>::New();
  pcaStatistics->SetInputData(vtkStatisticsAlgorithm::INPUT_DATA, datasetTable);

  pcaStatistics->SetColumnStatus("M0", 1);
  pcaStatistics->SetColumnStatus("M1", 1);
  pcaStatistics->SetColumnStatus("M2", 1);
  pcaStatistics->RequestSelectedColumns();

  pcaStatistics->SetDeriveOption(true);

  pcaStatistics->Update();

  vtkSmartPointer<vtkStatisticalModel> outputMetaDS = vtkStatisticalModel::SafeDownCast(
    pcaStatistics->GetOutputDataObject(vtkStatisticsAlgorithm::OUTPUT_MODEL));

  vtkSmartPointer<vtkTable> outputMeta = outputMetaDS->GetTable(vtkStatisticalModel::Derived, 0);
  outputMeta->Dump();

  ///////// Eigenvalues ////////////
  vtkSmartPointer<vtkDoubleArray> eigenvalues = vtkSmartPointer<vtkDoubleArray>::New();
  pcaStatistics->GetEigenvalues(eigenvalues);
  double eigenvaluesGroundTruth[3] = { .5, .166667, 0 };
  vtkIdType eigenvaluesCount = eigenvalues->GetNumberOfTuples();
  if (eigenvaluesCount > 3)
  {
    return EXIT_FAILURE;
  }
  for (vtkIdType i = 0; i < eigenvaluesCount; i++)
  {
    std::cout << "Eigenvalue " << i << " = " << eigenvalues->GetValue(i) << std::endl;
    if (!vtkMathUtilities::FuzzyCompare<double>(
          eigenvalues->GetValue(i), eigenvaluesGroundTruth[i], 0.0001))
    {
      std::cerr << "Eigenvalues (GetEigenvalues) are not correct! (" << eigenvalues->GetValue(i)
                << " vs " << eigenvaluesGroundTruth[i] << ")" << std::endl;
      return EXIT_FAILURE;
    }

    if (!vtkMathUtilities::FuzzyCompare<double>(
          pcaStatistics->GetEigenvalue(i), eigenvaluesGroundTruth[i], 0.0001))
    {
      std::cerr << "Eigenvalues (GetEigenvalue) are not correct! ("
                << pcaStatistics->GetEigenvalue(i) << " vs " << eigenvaluesGroundTruth[i] << ")"
                << std::endl;
      return EXIT_FAILURE;
    }
  }

  std::vector<std::vector<double>> eigenvectorsGroundTruth;
  std::vector<double> e0(3);
  e0[0] = -.707107;
  e0[1] = .707107;
  e0[2] = 0;
  std::vector<double> e1(3);
  e1[0] = .707107;
  e1[1] = .707107;
  e1[2] = 0;
  std::vector<double> e2(3);
  e2[0] = 0;
  e2[1] = 0;
  e2[2] = 1;
  eigenvectorsGroundTruth.push_back(e0);
  eigenvectorsGroundTruth.push_back(e1);
  eigenvectorsGroundTruth.push_back(e2);

  vtkSmartPointer<vtkDoubleArray> eigenvectors = vtkSmartPointer<vtkDoubleArray>::New();

  pcaStatistics->GetEigenvectors(eigenvectors);
  for (vtkIdType i = 0; i < eigenvectors->GetNumberOfTuples(); i++)
  {
    std::cout << "Eigenvector " << i << " : ";
    double* evec = new double[eigenvectors->GetNumberOfComponents()];
    eigenvectors->GetTuple(i, evec);
    int iamax = 0;
    double vamax = fabs(eigenvectorsGroundTruth[i][0]);
    for (vtkIdType j = 1; j < eigenvectors->GetNumberOfComponents(); j++)
    {
      double tmp = fabs(eigenvectorsGroundTruth[i][j]);
      if (tmp > vamax)
      {
        iamax = j;
        vamax = tmp;
      }
    }
    double factor = (vamax == eigenvectorsGroundTruth[i][iamax]) ? +1 : -1;
    for (vtkIdType j = 0; j < eigenvectors->GetNumberOfComponents(); j++)
    {
      std::cout << evec[j] << " ";
      vtkSmartPointer<vtkDoubleArray> eigenvectorSingle = vtkSmartPointer<vtkDoubleArray>::New();
      pcaStatistics->GetEigenvector(i, eigenvectorSingle);
      if (!vtkMathUtilities::FuzzyCompare<double>(
            factor * eigenvectorsGroundTruth[i][j], evec[j], 0.0001) ||
        !vtkMathUtilities::FuzzyCompare<double>(
          factor * eigenvectorsGroundTruth[i][j], eigenvectorSingle->GetValue(j), 0.0001))
      {
        std::cerr << "Eigenvectors do not match!" << std::endl;
        return EXIT_FAILURE;
      }
    }
    delete[] evec;
    std::cout << std::endl;
  }

  return EXIT_SUCCESS;
}
