// 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 from Sandia National Laboratories
// for implementing this test.

#include "vtkContingencyStatistics.h"
#include "vtkDataAssembly.h"
#include "vtkStatisticalModel.h"
#include "vtkStringArray.h"
#include "vtkTable.h"
#include "vtkVariantArray.h"

#include <cmath>

#include <iostream>

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

  vtkVariant mingledData[] = {
    123, 456, 80, "HTTP",   //
    123, 789, 80, "HTTP",   //
    123, 789, 80, "HTTP",   //
    123, 456, 80, "HTTP",   //
    456, 123, 80, "HTTP",   //
    456, 123, 80, "HTTP",   //
    456, 123, 8080, "HTTP", //
    789, 123, 1122, "HTTP", //
    456, 789, 80, "HTTP",   //
    456, 789, 25, "SMTP",   //
    456, 789, 25, "SMTP",   //
    456, 789, 25, "SMTP",   //
    456, 789, 25, "SMTP",   //
    123, 789, 25, "SMTP",   //
    789, 123, 80, "SMTP",   //
    123, 456, 20, "FTP",    //
    789, 456, 20, "FTP",    //
    789, 123, 20, "FTP",    //
    789, 123, 122, "FTP",   //
    789, 456, 20, "FTP",    //
    789, 456, 20, "FTP"     //
  };
  int nVals = 21;

  vtkVariantArray* dataset0Arr = vtkVariantArray::New();
  dataset0Arr->SetNumberOfComponents(1);
  dataset0Arr->SetName("Source");

  vtkVariantArray* dataset1Arr = vtkVariantArray::New();
  dataset1Arr->SetNumberOfComponents(1);
  dataset1Arr->SetName("Destination");

  vtkVariantArray* dataset2Arr = vtkVariantArray::New();
  dataset2Arr->SetNumberOfComponents(1);
  dataset2Arr->SetName("Port");

  vtkVariantArray* dataset3Arr = vtkVariantArray::New();
  dataset3Arr->SetNumberOfComponents(1);
  dataset3Arr->SetName("Protocol");

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

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

  int nMetricPairs = 2;

  // Entropies in the summary table should normally be retrieved as follows:
  //   column 2: H(X,Y)
  //   column 3: H(Y|X)
  //   column 4: H(X|Y)
  int iEntropies[] = { 2, 3, 4 };
  int nEntropies = 3; // correct number of entropies reported in the summary table
  double* H = new double[nEntropies];

  // Set contingency statistics algorithm and its input data port
  vtkContingencyStatistics* cs = vtkContingencyStatistics::New();

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

  // Prepare first test with data
  cs->SetInputData(vtkStatisticsAlgorithm::INPUT_DATA, datasetTable);
  vtkTable* outputData = cs->GetOutput(vtkStatisticsAlgorithm::OUTPUT_DATA);

  datasetTable->Delete();

  // Select Column Pair of Interest ( Learn Option )
  // 1.1: a valid pair
  cs->AddColumnPair("Port", "Protocol");
  // 1.2: the same valid pair, just reversed -- should thus be ignored
  cs->AddColumnPair("Protocol", "Port");
  // 2: another valid pair
  cs->AddColumnPair("Source", "Port");
  // 3: an invalid pair
  cs->AddColumnPair("Source", "Dummy");

  // Test Learn, Derive, Assess, and Test options
  cs->SetLearnOption(true);
  cs->SetDeriveOption(true);
  cs->SetAssessOption(true);
  cs->SetTestOption(true);
  cs->Update();

  auto* outputModelDS = vtkStatisticalModel::SafeDownCast(
    cs->GetOutputDataObject(vtkStatisticsAlgorithm::OUTPUT_MODEL));
  vtkTable* outputSummary = outputModelDS->FindTableByName(vtkStatisticalModel::Learned, "Summary");
  vtkTable* outputContingency =
    outputModelDS->FindTableByName(vtkStatisticalModel::Learned, "Contingency Table");
  vtkTable* outputTest =
    vtkTable::SafeDownCast(cs->GetOutputDataObject(vtkStatisticsAlgorithm::OUTPUT_TEST));

  int testIntValue = 0;
  double testDoubleValue = 0;

  vtkIdType n = outputContingency->GetValueByName(0, "Cardinality").ToInt();
  std::cout << "## Calculated the following information entropies (grand total: " << n << "):\n";

  testIntValue = outputSummary->GetNumberOfColumns();

  if (testIntValue != nEntropies + 2)
  {
    vtkGenericWarningMacro("Reported an incorrect number of columns in the summary table: "
      << testIntValue << " != " << nEntropies + 2 << ".");
    testStatus = 1;
  }
  else
  {
    // For each row in the summary table, fetch variable names and information entropies
    for (vtkIdType r = 0; r < outputSummary->GetNumberOfRows(); ++r)
    {
      // Variable names
      std::cout << "   (X,Y) = (" << outputSummary->GetValue(r, 0).ToString() << ", "
                << outputSummary->GetValue(r, 1).ToString() << ")";

      // Information entropies
      for (vtkIdType c = 0; c < nEntropies; ++c)
      {
        H[c] = outputSummary->GetValue(r, iEntropies[c]).ToDouble();

        std::cout << ", " << outputSummary->GetColumnName(iEntropies[c]) << "=" << H[c];
      }
      std::cout << "\n";

      // Make sure that H(X,Y) > H(Y|X)+ H(X|Y)
      testDoubleValue = H[1] + H[2]; // H(Y|X)+ H(X|Y)

      if (testDoubleValue > H[0])
      {
        vtkGenericWarningMacro("Reported inconsistent information entropies: H(X,Y) = "
          << H[0] << " < " << testDoubleValue << " = H(Y|X)+ H(X|Y).");
        testStatus = 1;
      }
    }
  }
  std::cout
    << "   where H(X,Y) = - Sum_{x,y} p(x,y) log p(x,y) and H(X|Y) = - Sum_{x,y} p(x,y) log "
       "p(x|y).\n";

  std::cout << "\n";

  std::cout
    << "## Calculated the following joint and conditional probabilities and mutual information:\n";
  testIntValue = 0;

  // Skip first row which contains data set cardinality
  vtkIdType key;
  for (vtkIdType r = 1; r < outputContingency->GetNumberOfRows(); ++r)
  {
    key = outputContingency->GetValue(r, 0).ToInt();
    std::cout << "   (" << outputSummary->GetValue(key, 0).ToString() << ","
              << outputSummary->GetValue(key, 1).ToString() << ") = ("
              << outputContingency->GetValue(r, 1).ToString() << ","
              << outputContingency->GetValue(r, 2).ToString() << ")";

    for (vtkIdType c = 3; c < outputContingency->GetNumberOfColumns(); ++c)
    {
      std::cout << ", " << outputContingency->GetColumnName(c) << "="
                << outputContingency->GetValue(r, c).ToDouble();
    }

    std::cout << "\n";

    // Update total cardinality
    testIntValue += outputContingency->GetValueByName(r, "Cardinality").ToInt();
  }

  if (testIntValue != nVals * nMetricPairs)
  {
    vtkGenericWarningMacro("Reported an incorrect total cardinality: "
      << testIntValue << " != " << nVals * nMetricPairs << ".");
    testStatus = 1;
  }
  std::cout << "\n";

  std::cout << "## Calculated the following marginal probabilities:\n";
  testIntValue = 0;

  for (int b = 0; b < outputModelDS->GetNumberOfTables(vtkStatisticalModel::Derived); ++b)
  {
    outputContingency = outputModelDS->GetTable(vtkStatisticalModel::Derived, b);

    for (vtkIdType r = 0; r < outputContingency->GetNumberOfRows(); ++r)
    {
      std::cout << "   " << outputContingency->GetColumnName(0) << " = "
                << outputContingency->GetValue(r, 0).ToString() << ", "
                << outputContingency->GetColumnName(1) << "="
                << outputContingency->GetValue(r, 1).ToDouble() << ", "
                << outputContingency->GetColumnName(2) << "="
                << outputContingency->GetValue(r, 2).ToDouble() << "\n";
    }
    std::cout << "\n";

    // Update total cardinality
    testIntValue += 0; // outputContingency->GetValueByName( r, "Cardinality" ).ToInt();
  }

  // Now inspect results of the Assess option by looking for outliers
  key = 0;
  std::string varX = outputSummary->GetValue(key, 0).ToString();
  std::string varY = outputSummary->GetValue(key, 1).ToString();

  // List of columns used for outlier detection
  std::string outlierColumn[] = { "P", "Px|y", "PMI" };
  // Corresponding threshold (low) values
  double threshold[] = { .2, .2, .0 };

  // Corresponding known number of outliers
  int nOutliers[] = { 4, 4, 1 };

  int nOutlierTypes = 3;
  for (int i = 0; i < nOutlierTypes; ++i)
  {
    std::string colName = outlierColumn[i] + "(" + varX + "," + varY + ")";

    std::cout << "## Found the following outliers such that " << colName << " < " << threshold[i]
              << ":\n";

    double val;
    testIntValue = 0;
    for (vtkIdType r = 0; r < outputData->GetNumberOfRows(); ++r)
    {
      val = outputData->GetValueByName(r, colName.c_str()).ToDouble();
      if (val >= threshold[i])
      {
        continue;
      }

      ++testIntValue;

      std::cout << "   " << outlierColumn[i] << "("
                << outputData->GetValueByName(r, varX.c_str()).ToString() << ","
                << outputData->GetValueByName(r, varY.c_str()).ToString() << ") = " << val << "\n";
    }

    if (testIntValue != nOutliers[i])
    {
      vtkGenericWarningMacro("Reported an incorrect number of outliers: " << testIntValue << " != "
                                                                          << nOutliers[i] << ".");
      testStatus = 1;
    }
    std::cout << "\n";
  }

  // Last, check some results of the Test option
  std::cout << "## Chi square statistics:\n";

  // Reference values
  double testValues[] = {
    // (Port,Protocol)
    10.,    // number of degrees of freedom
    36.896, // Chi square statistic
    22.35,  // Chi square statistic with Yates correction
#ifdef USE_GNU_R
    .00005899, // p-valued of Chi square statistic
    .01341754, // p-value of Chi square statistic with Yates correction
#endif         // USE_GNU_R
    // (Port,Source)
    10.,    // number of degrees of freedom
    17.353, // Chi square statistic
    7.279,  // Chi square statistic with Yates correction
#ifdef USE_GNU_R
    .06690889, // p-valued of Chi square statistic
    .69886917  // p-value of Chi square statistic with Yates correction
#endif         // USE_GNU_R
  };

#ifdef USE_GNU_R
  double alpha = .05;
  vtkIdType nv = 5;
#else  // USE_GNU_R
  vtkIdType nv = 3;
#endif // USE_GNU_R

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

    for (vtkIdType c = 0; c < nv; ++c)
    {
      double x = outputTest->GetValue(r, c).ToDouble();
      std::cout << ", " << outputTest->GetColumnName(c) << "=" << x;

      // Verify calculated results
      if (fabs(x - testValues[r * nv + c]) > 1.e-4 * x)
      {
        vtkGenericWarningMacro("Incorrect " << outputTest->GetColumnName(c) << ": " << x
                                            << " != " << testValues[r * nv + c]);
        testStatus = 1;
      }
    }

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

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

  // Clean up
  delete[] H;
  cs->Delete();

  return testStatus;
}
