/* Distributed under the Apache License, Version 2.0.
   See accompanying NOTICE file for details.*/

#include "VTKUtils.h"

#include <vtkAxis.h>
#include <vtkDataSet.h>
#include <vtkCellData.h>
#include <vtkCellTypes.h>
#include <vtkChartXY.h>
#include <vtkChartLegend.h>
#include <vtkContextView.h>
#include <vtkContextScene.h>
#include <vtkDataSetReader.h>
#include <vtkFieldData.h>
#include <vtkFloatArray.h>
#include <vtkGeometryFilter.h>
#include <vtkGraphicsFactory.h>
#include <vtkIntArray.h>
#include <vtkPlot.h>
#include <vtkPNGWriter.h>
#include <vtkPen.h>
#include <vtkPointData.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkResampleWithDataSet.h>
#include <vtkSmartPointer.h>
#include <vtkTable.h>
#include <vtkThreshold.h>
#include <vtkTriangleFilter.h>
#include <vtkUnstructuredGrid.h>
#include <vtkWindowToImageFilter.h>
#include <vtkXMLPolyDataWriter.h>
#include <vtkXMLReader.h>
#include <vtkXMLRectilinearGridReader.h>
#include <vtkXMLRectilinearGridWriter.h>
#include <vtksys/SystemTools.hxx>

namespace LBM
{
namespace VTK
{

vtkSmartPointer<vtkPolyData> ExtractSurface(vtkSmartPointer<vtkRectilinearGrid> vtr, int label)
{
  return ExtractSurface(vtr, label, label);
}
vtkSmartPointer<vtkPolyData> ExtractSurface(vtkSmartPointer<vtkRectilinearGrid> vtr, int min_label, int max_label)
{
  auto thresh = vtkSmartPointer<vtkThreshold>::New();
  thresh->SetInputData(vtr);
  thresh->SetInputArrayToProcess(0, 0, 0, vtkDataObject::FIELD_ASSOCIATION_CELLS, BOUNDARY_LABELS);
  // TODO How do you know a label exists?
  thresh->ThresholdBetween(min_label, max_label);
  thresh->Update();

  auto geom = vtkSmartPointer<vtkGeometryFilter>::New();
  geom->SetInputData(thresh->GetOutput());
  geom->Update();

  auto tri = vtkSmartPointer<vtkTriangleFilter>::New();
  tri->SetInputData(geom->GetOutput());
  tri->Update();

  auto res = vtkSmartPointer<vtkResampleWithDataSet>::New();
  res->SetInputData(tri->GetOutput());
  res->SetSourceData(tri->GetOutput());
  res->Update();

  return dynamic_cast<vtkPolyData*>(res->GetOutput());
}

std::unique_ptr<LBM::VTK::FlowMetrics> ComputeMetrics(LBM::Run const& lbm)
{
  std::unique_ptr<LBM::VTK::FlowMetrics> metrics(new LBM::VTK::FlowMetrics());

  auto vtr = ToVTR(lbm);
  int min_label, max_label;

  min_label = 100;
  max_label = 199;
  lbm.Info("Extracting surface label "+std::to_string(min_label)+" to "+std::to_string(max_label)+"...");
  metrics->inlet_surface = ExtractSurface(vtr, min_label, max_label);
  min_label = 200;
  max_label = 299;
  lbm.Info("Extracting surface label "+std::to_string(min_label)+" to "+std::to_string(max_label)+"...");
  metrics->outlet_surface = ExtractSurface(vtr, min_label, max_label);
  min_label = 1;
  max_label = 99;
  lbm.Info("Extracting surface label "+std::to_string(min_label)+" to "+std::to_string(max_label)+"...");
  metrics->wall_surface = ExtractSurface(vtr, min_label, max_label);

  float vx, vy, vz;
  float voxel_area = lbm.in.grid_spacing * lbm.in.grid_spacing;

  vtkDataArray* inlet_pressures = metrics->inlet_surface->GetPointData()->GetArray(PRESSURE);
  vtkDataArray* inlet_x_flows = metrics->inlet_surface->GetPointData()->GetArray(X_FLOW);
  vtkDataArray* inlet_y_flows = metrics->inlet_surface->GetPointData()->GetArray(Y_FLOW);
  vtkDataArray* inlet_z_flows = metrics->inlet_surface->GetPointData()->GetArray(Z_FLOW);
  int num_inlet_voxels = inlet_pressures->GetSize();

  vtkDataArray* outlet_pressures = metrics->outlet_surface->GetPointData()->GetArray(PRESSURE);
  vtkDataArray* outlet_x_flows = metrics->outlet_surface->GetPointData()->GetArray(X_FLOW);
  vtkDataArray* outlet_y_flows = metrics->outlet_surface->GetPointData()->GetArray(Y_FLOW);
  vtkDataArray* outlet_z_flows = metrics->outlet_surface->GetPointData()->GetArray(Z_FLOW);
  int num_outlet_voxels = outlet_pressures->GetSize();

  // Compute the flow on the inlet and outlet surfaces
  // |v| = sqrt(vx*vx + vy*vy + vz*vz)
  // velocity = sum(|v|) / num_voxels
  // flow = velocity * area
  // area = num_voxels * voxel_area
  // Compute the pressure on the inlet and outlet surfaces
  // |p| = sum(p) / num_voxels
  // Compute the resistance
  // resistance = deltaP / flow

  // Inlet
  float inlet_velocity = 0;
  float inlet_pressure = 0;
  for (int idx = 0; idx < num_inlet_voxels; idx++)
  {
    inlet_pressure += inlet_pressures->GetTuple1(idx);
    vx = inlet_x_flows->GetTuple1(idx);
    vy = inlet_y_flows->GetTuple1(idx);
    vz = inlet_z_flows->GetTuple1(idx);
    inlet_velocity += std::sqrt((vx * vx) + (vy * vy) + (vz * vz));
  }
  inlet_pressure /= num_inlet_voxels;
  float inlet_flow = inlet_velocity * voxel_area;

  // Outlet
  float outlet_velocity = 0;
  float outlet_pressure = 0;
  for (int idx = 0; idx < num_outlet_voxels; idx++)
  {
    outlet_pressure += outlet_pressures->GetTuple1(idx);
    vx = outlet_x_flows->GetTuple1(idx);
    vy = outlet_y_flows->GetTuple1(idx);
    vz = outlet_z_flows->GetTuple1(idx);
    outlet_velocity += std::sqrt((vx * vx) + (vy * vy) + (vz * vz));
  }
  outlet_pressure /= num_outlet_voxels;
  float outlet_flow = outlet_velocity * voxel_area;
  if (inlet_flow != outlet_flow)
  {
    lbm.Warning("Outlet flow does not match inlet flow");
  }

  float deltaP = std::abs(inlet_pressure - outlet_pressure);

  metrics->flow_rate = inlet_flow;
  metrics->resistance = deltaP / inlet_flow;

  return std::move(metrics);
}

bool PlotCSV(std::string const& base_name, std::vector<CSVPlotSource>const& srcs, Logger& log)
{
  return PlotCSV(base_name, srcs, false, log);
}
bool PlotCSV(std::string const& base_name, std::vector<CSVPlotSource> const& srcs, bool const_y_axis, Logger& log)
{
  if (srcs.empty())
    return false;

  size_t num_columns = srcs[0].csv.size();
  size_t num_values = srcs[0].csv[0].second.size();
  float y_min = std::numeric_limits<float>::max();
  float y_max = -std::numeric_limits<float>::max();

  // Make sure all the CSV's have the same number of columns, column names, and value lengths

  for (size_t i=0; i<srcs.size(); i++)
  {
    auto& src = srcs[i];
    if (src.csv.size() != num_columns)
    {
      log.Error("All CSV's need to have the same number of columns");
      return false;
    }
    // Make sure all the columns are the same
    for (size_t c = 1; c < num_columns; c++)
    {
      auto& csv = src.csv[c];
      if (csv.first != srcs[0].csv[c].first)
      {
        log.Error("CSVPlotSource column names do not match");
        return false;
      }
      if (csv.second.size() != srcs[0].csv[c].second.size())
      {
        log.Error("CSVPlotSource columns size does not match");
        return false;
      }

      if (const_y_axis)
      {
        // Make sure all our plots have the same y scale
        for (size_t v = 0; v < csv.second.size(); v++)
        {
          float y_val = csv.second[v];
          if (y_val > y_max)
            y_max = y_val;
          if (y_val < y_min)
            y_min = y_val;
        }
      }
    }
  }
  if (const_y_axis)
  {
    // Add 2% on each side
    float pad = (std::fabs(y_min) + std::fabs(y_max)) * 0.02;
    y_min -= pad;
    y_max += pad;
  }


  for (size_t c = 1; c < num_columns; c++)
  {
    std::string plot_name = base_name + srcs[0].csv[c].first;

    auto table = vtkSmartPointer<vtkTable>::New();
    auto time_data = vtkSmartPointer<vtkFloatArray>::New();
    time_data->SetName("X-Axis");
    table->AddColumn(time_data);
    // Create Columns for each line
    for (size_t s = 0; s < srcs.size(); s++)
    {
      auto y_values = vtkSmartPointer<vtkFloatArray>::New();
      y_values->SetName(srcs[s].name.c_str());
      table->AddColumn(y_values);
    }
    // Add data to each column (must be done after all columns are created)
    table->SetNumberOfRows(num_values);
    for (size_t s = 0; s < srcs.size(); s++)
    {
      const LBM::CSV& csv = srcs[s].csv;
      for (size_t y = 0; y < num_values; y++)
      {
        if (s == 0)// Add X values from the first axis
          table->SetValue(y, 0, csv[0].second[y]);
        table->SetValue(y, s+1, csv[c].second[y]);
      }
    }

    //Create our chart
    auto chart = vtkSmartPointer<vtkChartXY>::New();
    //chart->SetTitle(srcs[0].csv[0].first + " vs. " + srcs[0].csv[c].first);
    chart->GetAxis(vtkAxis::LEFT)->SetTitle(srcs[0].csv[c].first);
    if (const_y_axis)
    {
      chart->GetAxis(vtkAxis::LEFT)->SetMinimum(y_min);
      chart->GetAxis(vtkAxis::LEFT)->SetMaximum(y_max);
      chart->GetAxis(vtkAxis::LEFT)->SetBehavior(vtkAxis::FIXED);
    }
    chart->GetAxis(vtkAxis::BOTTOM)->SetTitle(srcs[0].csv[0].first);
    // Add some lines to the chart
    for (size_t s = 0; s < srcs.size(); s++)
    {
      auto& src = srcs[s];
      vtkPlot* c_line = chart->AddPlot(vtkChart::LINE);
      c_line->SetInputData(table, 0, s+1);// Always using Column 0 as our x coordinate
      c_line->SetColor(src.r, src.g, src.b);
      c_line->SetWidth(src.w);
      c_line->GetPen()->SetLineType(src.p);
    }
    // Add the Legend
    // Set up the legend to be off to the top right of the viewport.
    chart->SetShowLegend(true);
    chart->GetLegend()->SetInline(false);
    chart->GetLegend()->SetHorizontalAlignment(vtkChartLegend::CENTER);
    chart->GetLegend()->SetVerticalAlignment(vtkChartLegend::TOP);

    // Set up the view
    auto view = vtkSmartPointer<vtkContextView>::New();
    view->GetRenderer()->SetBackground(1.0, 1.0, 1.0);
    view->GetRenderer()->SetBackgroundAlpha(1.0);
    view->GetRenderWindow()->SetSize(640, 480);
    view->GetRenderWindow()->SetOffScreenRendering(1);
    view->GetScene()->AddItem(chart);
    view->GetRenderWindow()->Render();
    // Screenshot
    auto windowToImageFilter = vtkSmartPointer<vtkWindowToImageFilter>::New();
    windowToImageFilter->SetInput(view->GetRenderWindow());
    windowToImageFilter->SetInputBufferTypeToRGBA(); //also record the alpha (transparency) channel
    windowToImageFilter->ReadFrontBufferOff(); // read from the back buffer
    windowToImageFilter->Update();
    // Write to disk
    log.Info("Writing " + plot_name);
    auto writer = vtkSmartPointer<vtkPNGWriter>::New();
    writer->SetFileName(std::string(plot_name +".png").c_str());
    writer->SetInputConnection(windowToImageFilter->GetOutputPort());
    writer->Write();
  }
  return true;
}

bool WriteMetrics(std::string const& base_name, LBM::VTK::FlowMetrics const& m)
{
  ofstream results_file;
  results_file.open(base_name+"_metrics.json", std::ofstream::out | std::ofstream::trunc);
  results_file << "{" << "\n";
  results_file << "\"Obstructed Volumetric Airflow (L/min)\": " << "\"" << m.flow_rate << "\"," << "\n";
  results_file << "\"Obstructed Resistance Ratio\": " << "\"" << m.resistance << "\"," << "\n";
  results_file << "\"Obstructed Heat Flux (W/m^2)\": " << "\"" << m.heat_flux << "\"," << "\n";
  results_file << "\"Obstructed Wall Shear\": " << "\"" << m.wall_shear << "\"" << "\n";
  results_file << "}" << "\n";
  results_file.close();

  if (!WriteVTP(base_name+"_inlet_surface.vtp", m.inlet_surface))
    return false;
  if (!WriteVTP(base_name + "_outlet_surface.vtp", m.outlet_surface))
    return false;
  if (!WriteVTP(base_name + "_wall_surface.vtp", m.wall_surface))
    return false;
  return true;
}

bool WriteVTP(std::string const& filename, vtkSmartPointer<vtkPolyData> pd)
{
  auto w = vtkSmartPointer<vtkXMLPolyDataWriter>::New();
  w->SetInputData(pd);
  w->SetFileName(filename.c_str());
  return w->Write() == 1;
}

vtkSmartPointer<vtkRectilinearGrid> ToVTR(LBM::Run const& lbm)
{
  // Create the point lattice
  vtkSmartPointer<vtkRectilinearGrid> grid =
    vtkSmartPointer<vtkRectilinearGrid>::New();

  grid->SetDimensions((int)lbm.in.dimensions[0]+1,
                      (int)lbm.in.dimensions[1]+1,
                      (int)lbm.in.dimensions[2]+1);
  vtkNew<vtkFloatArray> xCoords;
  xCoords->SetName("X");
  for (size_t i = 0; i < lbm.in.dimensions[0] + 1; i++)
  {
    xCoords->InsertNextValue(i * lbm.in.grid_spacing);
  }
  vtkNew<vtkFloatArray> yCoords;
  yCoords->SetName("Y");
  for (size_t i = 0; i < lbm.in.dimensions[1] + 1; i++)
  {
    yCoords->InsertNextValue(i * lbm.in.grid_spacing);
  }
  vtkNew<vtkFloatArray> zCoords;
  zCoords->SetName("Z");
  for (size_t i = 0; i < lbm.in.dimensions[2] + 1; i++)
  {
    zCoords->InsertNextValue(i * lbm.in.grid_spacing);
  }
  grid->SetXCoordinates(xCoords);
  grid->SetYCoordinates(yCoords);
  grid->SetZCoordinates(zCoords);

  if (!lbm.in.source_labels.empty())
  {
    vtkNew<vtkIntArray> sbl;
    sbl->SetName(SOURCE_LABELS);
    for (int i : lbm.in.source_labels)
    {
      sbl->InsertNextValue(i);
    }
    grid->GetCellData()->AddArray(sbl);
  }
  else if (!lbm.in.labels.empty())
  {
    vtkNew<vtkIntArray> sbl;
    sbl->SetName(SOURCE_LABELS);
    for (LBM::BoundaryTypes t : lbm.in.labels)
    {
      sbl->InsertNextValue((int)t);
    }
    grid->GetCellData()->AddArray(sbl);
  }

  if (lbm.out.labels != nullptr)
  {
    vtkNew<vtkIntArray> bl;
    bl->SetName(BOUNDARY_LABELS);
    for (size_t i = 0; i < lbm.out.num_cells; i++)
    {
      bl->InsertNextValue(lbm.out.labels[i]);
    }
    grid->GetCellData()->AddArray(bl);
  }

  if (lbm.out.pressure_flows_and_stress != nullptr)
  {
    vtkNew<vtkFloatArray>p;
    p->SetName(PRESSURE);
    grid->GetCellData()->AddArray(p);
    vtkNew<vtkFloatArray>vx;
    vx->SetName(X_FLOW);
    grid->GetCellData()->AddArray(vx);
    vtkNew<vtkFloatArray>vy;
    vy->SetName(Y_FLOW);
    grid->GetCellData()->AddArray(vy);
    vtkNew<vtkFloatArray>vz;
    vz->SetName(Z_FLOW);
    grid->GetCellData()->AddArray(vz);


    vtkNew<vtkFloatArray>Sxy;
    vtkNew<vtkFloatArray>Sxz;
    vtkNew<vtkFloatArray>Syz;
    vtkNew<vtkFloatArray>Sxx;
    vtkNew<vtkFloatArray>Syy;
    vtkNew<vtkFloatArray>Szz;
    if (lbm.cfg.expanded_results)
    {
      Sxy->SetName(XY_STRESS);
      grid->GetCellData()->AddArray(Sxy);
      Sxz->SetName(XZ_STRESS);
      grid->GetCellData()->AddArray(Sxz);
      Syz->SetName(YZ_STRESS);
      grid->GetCellData()->AddArray(Syz);
      Sxx->SetName(XX_STRESS);
      grid->GetCellData()->AddArray(Sxx);
      Syy->SetName(YY_STRESS);
      grid->GetCellData()->AddArray(Syy);
      Szz->SetName(ZZ_STRESS);
      grid->GetCellData()->AddArray(Szz);
    }

    for (size_t i = 0; i < lbm.out.num_cells; i++)
    {
      p->InsertNextValue(lbm.out.pressure_flows_and_stress[i]);
      vx->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells+i]);
      vy->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells*2+i]);
      vz->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells*3+i]);
      if (lbm.cfg.expanded_results)
      {
        Sxy->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 4 + i]);
        Sxz->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 5 + i]);
        Syz->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 6 + i]);
        Sxx->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 7 + i]);
        Syy->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 8 + i]);
        Szz->InsertNextValue(lbm.out.pressure_flows_and_stress[lbm.out.num_cells * 9 + i]);
      }
    }
  }

    if (lbm.out.wall_shear_stress != nullptr)
    {
        vtkNew<vtkFloatArray> stress;
        stress->SetName(STRESS);
        grid->GetCellData()->AddArray(stress);

        vtkNew<vtkFloatArray> legacy_stress;
        if (lbm.cfg.expanded_results)
        {
          legacy_stress->SetName(LEGACY_STRESS);
          grid->GetCellData()->AddArray(legacy_stress);
        }

        for (size_t i = 0; i < lbm.out.num_cells; i++)
        {
          if(lbm.cfg.expanded_results)
            legacy_stress->InsertNextValue(lbm.out.wall_shear_stress[i]);
          stress->InsertNextValue(lbm.out.wall_shear_stress[lbm.out.num_cells+i]);
        }
    }

  if (lbm.out.temperature != nullptr)
  {
    vtkNew<vtkFloatArray> t;
    t->SetName(TEMPERATURE);
    for (size_t i = 0; i < lbm.out.num_cells; i++)
    {
      t->InsertNextValue(lbm.out.temperature[i]);
    }
    grid->GetCellData()->AddArray(t);
  }

  return grid;
}
bool WriteVTR(std::string const& filename, LBM::Run const& lbm)
{
  vtkSmartPointer<vtkRectilinearGrid> grid = ToVTR(lbm);
  return WriteVTR(filename, grid);
}
bool WriteVTR(std::string const& filename, vtkSmartPointer<vtkRectilinearGrid> vtr)
{
  // Write file
  vtkSmartPointer<vtkXMLRectilinearGridWriter> writer =
  vtkSmartPointer<vtkXMLRectilinearGridWriter>::New();
  writer->SetFileName(filename.c_str());
  writer->SetInputData(vtr);
  return writer->Write() == 1;
}

bool ReadVTR(std::string const& filename, LBM::Run& lbm)
{
  vtkSmartPointer<vtkRectilinearGrid> vtr = ReadVTR(filename);
  return FromVTR(vtr, lbm);
}
vtkSmartPointer<vtkRectilinearGrid> ReadVTR(std::string const& filename)
{
  vtkSmartPointer<vtkXMLRectilinearGridReader> reader =
    vtkSmartPointer<vtkXMLRectilinearGridReader>::New();
  reader->SetFileName(filename.c_str());
  reader->Update();
  reader->GetOutput()->Register(reader);
  vtkSmartPointer<vtkRectilinearGrid> vtr = reader->GetOutput();
  return vtr;
}

bool FromVTR(vtkSmartPointer<vtkRectilinearGrid> vtr, LBM::Run& lbm)
{
  lbm.in.labels.clear(); // Not stored in the VTR
  lbm.in.dimensions[0] = vtr->GetDimensions()[0]-1;
  lbm.in.dimensions[1] = vtr->GetDimensions()[1]-1;
  lbm.in.dimensions[2] = vtr->GetDimensions()[2]-1;
  lbm.in.grid_spacing = 0; // Cannot assume same spacing in all dimensions from VTR

  lbm.out.Allocate();
  size_t num_cells = lbm.out.num_cells;

  // Now check for cell data
  vtkCellData* cd = vtr->GetCellData();
  if (cd)
  {
    for (int a=0; a<cd->GetNumberOfArrays(); a++)
    {
      std::string name = cd->GetArrayName(a);
      vtkDataArray* array = cd->GetArray(name.c_str());
      if (name.compare(BOUNDARY_LABELS) == 0 || name.compare("iBC") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.labels[i] = (int)(*array->GetTuple(i));
      }
      if (name.compare(PRESSURE) == 0 || name.compare("p") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(X_FLOW) == 0 || name.compare("u") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(Y_FLOW) == 0 || name.compare("v") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells*2+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(Z_FLOW) == 0 || name.compare("w") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells*3+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(TEMPERATURE) == 0 || name.compare("t") == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.temperature[i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(STRESS) == 0 )
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.wall_shear_stress[num_cells+i] = (float)(*array->GetTuple(i));
      }

      // Expanded Results (may not be available, but that is ok!)
      if (name.compare(LEGACY_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.wall_shear_stress[i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(XY_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells*4+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(XZ_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells*5+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(YZ_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells*6+i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(XX_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells * 4 + i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(YY_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells * 5 + i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(ZZ_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells * 6 + i] = (float)(*array->GetTuple(i));
      }
      if (name.compare(ZZ_STRESS) == 0)
      {
        for (int i = 0; i < array->GetSize(); i++)
          lbm.out.pressure_flows_and_stress[num_cells * 6 + i] = (float)(*array->GetTuple(i));
      }
    }
  }
  return lbm.out.CountCellTypes();
}

bool CheckVTR(LBM::Run const& lbm)
{
  auto vtr = ToVTR(lbm);
  LBM::Run copy;
  FromVTR(vtr, copy);
  // Compare what we read in to what we wrote
  return LBM::CompareDataSet(lbm, copy, true);
}

bool TestCSVPlot()
{
  Logger log("Test.log");
  std::vector<float> time;
  std::vector<float> analytic_values;
  std::vector<float> baseline_values;
  std::vector<float> computed_values;

  for (size_t t=0; t<360; t++)
  {
    time.push_back(t);
    analytic_values.push_back(sin(t*3.1415/180));
    baseline_values.push_back(cos(t*3.1415/180));
    computed_values.push_back(2+cos(t*3.1415/180));
  }

  LBM::CSV analytic_csv;
  analytic_csv.push_back(std::pair <std::string, std::vector<float>>("X", time));
  analytic_csv.push_back(std::pair <std::string, std::vector<float>>("Y", analytic_values));

  LBM::CSV baseline_csv;
  baseline_csv.push_back(std::pair <std::string, std::vector<float>>("X", time));
  baseline_csv.push_back(std::pair <std::string, std::vector<float>>("Y", baseline_values));

  LBM::CSV computed_csv;
  computed_csv.push_back(std::pair <std::string, std::vector<float>>("X", time));
  computed_csv.push_back(std::pair <std::string, std::vector<float>>("Y", computed_values));

  std::vector<VTK::CSVPlotSource> plot_sources;
  plot_sources.push_back(VTK::CSVPlotSource("Analytic", analytic_csv, 0.0, 1.0, 0.0, 2.f, vtkPen::SOLID_LINE));
  plot_sources.push_back(VTK::CSVPlotSource("Baseline", baseline_csv, 1.0, 0.0, 0.0, 5.f, vtkPen::DASH_LINE));
  plot_sources.push_back(VTK::CSVPlotSource("Computed", computed_csv, 0.0, 0.0, 1.0, 5.f, vtkPen::DASH_DOT_LINE));
  return PlotCSV("./test_plot", plot_sources, true, log);
}

}
}
