#include <vtkDoubleArray.h>
#include <vtkGeometryFilter.h>
#include <vtkInformation.h>
#include <vtkMultiBlockDataSet.h>
#include <vtkPointData.h>
#include <vtkSmartPointer.h>
#include <vtkStreamingDemandDrivenPipeline.h>
#include <vtkTimeSourceExample.h>
#include <vtkTrailingFrame.h>

#include <iostream>
#include <iomanip>


const double epsilon = 1e-9;


double value_fonction(double x)
{
    /// function used in the vtkTimeSourceExample to generate the array "Point Value"
    return sin(2 * vtkMath::Pi() * x);
}


double get_points_values(vtkPolyData* p)
{
    // return the value associated to the first points contained in the array "Point Value"
    if (!p)
        return -1.0;
    else
    {
        auto array =  vtkDoubleArray::SafeDownCast(p->GetPointData()->GetArray("Point Value"));
        return array->GetValue(0);
    }
}

bool check_multi_block_frames(vtkMultiBlockDataSet* multiblock, std::vector<double> const& times)
{
    // this functions compares the values of the pointdata of each polydata contained in the multiblock
    // and check that it corresponds to the value calculated with the previous times steps.

    if (multiblock->GetNumberOfBlocks() != times.size())
    {
        std::cerr << "Trailing frame test failed: \n";
        std::cerr << "Expected " << times.size() << " trailing frames, found " << multiblock->GetNumberOfBlocks() << "\n";
        return false;
    }

    for (int i = 0; i < multiblock->GetNumberOfBlocks(); ++i)
    {
        auto poly_data_ptr = multiblock->GetBlock(i);
        if (poly_data_ptr == nullptr)
        {
            if (times[i] == -1.0)
                continue;
            else
            {
                std::cerr << "Trailing frame test failed: \n";
                std::cerr << "Invalid frame at index " << i << "\n";
                return false;
            }
        }
        double multi_block_value = get_points_values(vtkPolyData::SafeDownCast(multiblock->GetBlock(i)));
        double expected_value = value_fonction(times[i]);
        if (std::abs(multi_block_value - expected_value) > epsilon)
        {
            std::cerr << "Trailing frame test failed: \n";
            std::cerr << "Expected value " << expected_value << " found " << multi_block_value << "\n";
            return false;
        }
    }
    return true;
}


bool check_trailing_frames(vtkTrailingFrame* tf, vtkInformation* info, double *time_steps, int time_index)
{
    int nb_trailing_frames = tf->GetNumberOfTrailingFrames();
    // go to specified time index
    info->Set(vtkStreamingDemandDrivenPipeline::UPDATE_TIME_STEP(), time_steps[time_index]);
    tf->Update();
    auto tf_out=  vtkMultiBlockDataSet::SafeDownCast(tf->GetOutputDataObject(0));
    // extract the previous frames time steps
    std::vector<double> sub_time_steps(nb_trailing_frames + 1, -1.0);
    for (int i = 0; i < sub_time_steps.size() && time_index - i >= 0; ++i)
    {
        sub_time_steps[i] = time_steps[time_index - i];
    }
    // check the constructed multiblock
    return check_multi_block_frames(tf_out, sub_time_steps);
}


int main(int argc, char* argv[])
{
    // vtkTimeSourceExample generates 10 time steps of points with a data array called "Point Value".
    // At time t, each point has the value: sin(2*PI*t). This test checks that the points of the different
    // polydata contained in the multiblock generated by the trailingframe filter have the correct values
    // at different time steps.

    auto source = vtkSmartPointer<vtkTimeSourceExample>::New();
    auto grid_to_poly = vtkSmartPointer<vtkGeometryFilter>::New();      // converts vtkUnstructuredGrid to polydata
    auto tf = vtkSmartPointer<vtkTrailingFrame>::New();

    grid_to_poly->SetInputConnection(source->GetOutputPort());
    tf->SetInputConnection(grid_to_poly->GetOutputPort());

    int const N1 = 2;
    int const N2 = 4;
    tf->SetNumberOfTrailingFrames(N1);

    // get available time steps
    tf->UpdateInformation();
    vtkInformation* outInfo = tf->GetOutputInformation(0);
    int len = outInfo->Length(vtkStreamingDemandDrivenPipeline::TIME_STEPS());
    double *time_steps  = outInfo->Get(vtkStreamingDemandDrivenPipeline::TIME_STEPS());

    bool res = true;

    // check at timestep 0
    res = res && check_trailing_frames(tf, outInfo, time_steps, 0);
    // go to 7 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 7);
    // go back to 2 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 2);
    // go to 5 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 5);

    // change number of trailing frames
    tf->SetNumberOfTrailingFrames(N2);
    tf->Update();
    tf->UpdateInformation();
    outInfo = tf->GetOutputInformation(0);
    time_steps = outInfo->Get(vtkStreamingDemandDrivenPipeline::TIME_STEPS());

    // check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 5);
    // go back to 3 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 3);
    // go back to 0 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 0);
    // go to 7 and check
    res = res && check_trailing_frames(tf, outInfo, time_steps, 7);

    return res ? 0 : -1;
}
