// ===========================================================================
//
//       Filename:  utFEM4EllipticPDE.cpp
//
//        Created:  08/16/12 19:49:10
//       Compiler:  Tested with g++, icpc, and MSVC 2010
//
//         Author:  Trevor Irons (ti)
//
//   Organisation:  Colorado School of Mines (CSM)
//                  United States Geological Survey (USGS)
//
//          Email:  tirons@mines.edu, tirons@usgs.gov
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// ===========================================================================

/**
  @file
  @author   Trevor Irons
  @date     08/16/12
  @version   0.0
 **/

#include "FEM4EllipticPDE.h"
#include "vtkSphere.h"
#include "vtkRectilinearGrid.h"
#include "vtkFloatArray.h"
#include <vtkIdList.h>
#include <vtkIdTypeArray.h>
#include <vtkCell.h>
#include <vtkTriangleFilter.h>
#include <vtkDataSetSurfaceFilter.h>
#include <vtkExtractCells.h>
#include <vtkGeometryFilter.h>
#include <vtkUnstructuredGrid.h>
#include <vtkExtractEdges.h>
#include <vtkImplicitDataSet.h>

// debugging stuff
#include <vtkDataSetMapper.h>
#include <vtkSelectionNode.h>
#include <vtkExtractSelection.h>
#include <vtkSelection.h>
#include <vtkVertexGlyphFilter.h>

#define RENDERTEST
//#define CONNECTTEST

#ifdef RENDERTEST
#include "vtkRectilinearGridGeometryFilter.h"
#include "vtkRectilinearGridOutlineFilter.h"
#include "vtkPolyDataMapper.h"
#include "vtkActor.h"
#include "vtkRenderWindowInteractor.h"
#include "vtkRenderer.h"
#include "vtkRenderWindow.h"
#include "vtkCamera.h"
#include "vtkProperty.h"
#include <vtkDataSetMapper.h>
#include <vtkSelectionNode.h>
#include <vtkSelection.h>
#include <vtkExtractSelection.h>
#include <vtkExtractEdges.h>
#include <vtkVertexGlyphFilter.h>
#include <vtkTriangleFilter.h>
#include <vtkImplicitHalo.h>
#endif

#ifdef CONNECTTEST
#include <vtkVersion.h>
#include <vtkSmartPointer.h>
#include <vtkPolyData.h>
#include <vtkCellData.h>
#include <vtkDoubleArray.h>
#include <vtkDataSet.h>
#include <vtkSphereSource.h>
#include <vtkTriangleFilter.h>
#include <vtkExtractEdges.h>
#include <vtkDataSetMapper.h>
#include <vtkActor.h>
#include <vtkRenderWindow.h>
#include <vtkRenderer.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkSelectionNode.h>
#include <vtkSelection.h>
#include <vtkExtractSelection.h>
#include <vtkProperty.h>
#include <vtkVertexGlyphFilter.h>
#endif

vtkSmartPointer<vtkIdList> GetConnectedVertices(vtkSmartPointer<vtkDataSet> mesh, int id);
vtkSmartPointer<vtkIdList> GetConnectedVertices2(vtkSmartPointer<vtkPolyData> mesh, int id);

// This function is copyright (C) 2012 Joeseph Cappriotti
vtkSmartPointer<vtkIdList> GetConnectedPoints(int id0, vtkSmartPointer<vtkDataSet> grid);

using namespace Lemma;

int main() {

    // Define Sigma term
    // Defaults to 1, (Poisson equation)

    // Define source (G) term
    //vtkSphere *Sphere = vtkSphere::New();
    vtkImplicitHalo *Sphere = vtkImplicitHalo::New();
        Sphere->SetCenter (0, 0, 0);
        Sphere->SetRadius (60);
        Sphere->SetFadeOut(.95);

    ////////////////////////////////////////////
    // Define solution mesh

    // Create a rectilinear grid by defining three arrays specifying the
    // coordinates in the x-y-z directions.

    int    nx = 160;
    double dx = 2.5;
    double ox = -200;
    vtkFloatArray *xCoords = vtkFloatArray::New();
    for (int i=0; i<nx+1; i++) xCoords->InsertNextValue(ox + i*dx);

    int    ny = 160;
    double dy = 2.5;
    double oy = -200;
    vtkFloatArray *yCoords = vtkFloatArray::New();
    for (int i=0; i<ny+1; i++) yCoords->InsertNextValue(oy + i*dy);

    int    nz = 160;
    double dz = 2.5;
    double oz = -200;
    vtkFloatArray *zCoords = vtkFloatArray::New();
    for (int i=0; i<nz+1; i++) zCoords->InsertNextValue(oz + i*dz);

    vtkRectilinearGrid *rGrid = vtkRectilinearGrid::New();
        rGrid->SetDimensions(nx+1, ny+1, nz+1);
        rGrid->SetExtent(0, nx, 0, ny, 0,  nz);
        rGrid->SetXCoordinates(xCoords);
        rGrid->SetYCoordinates(yCoords);
        rGrid->SetZCoordinates(zCoords);

    vtkRectilinearGrid *sigGrid = vtkRectilinearGrid::New();
        sigGrid->SetDimensions(nx+1, ny+1, nz+1);
        sigGrid->SetExtent(0, nx, 0, ny, 0, nz);
        sigGrid->SetXCoordinates(xCoords);
        sigGrid->SetYCoordinates(yCoords);
        sigGrid->SetZCoordinates(zCoords);
        vtkDoubleArray *sigArray = vtkDoubleArray::New();
        sigArray->SetNumberOfComponents(1);
        for (int ii=0; ii<(nx+1)*(ny+1)*(nz+1); ++ii) {
            sigArray->InsertTuple1(ii, 1.);
        }
        sigArray->SetName("sigma");
        sigGrid->GetPointData()->AddArray(sigArray);
        sigGrid->GetPointData()->SetScalars(sigArray);

    vtkImplicitDataSet* implSigma = vtkImplicitDataSet::New();
        implSigma->SetDataSet(sigGrid);
        implSigma->SetOutValue(0.);
        implSigma->SetOutGradient(0., 0., 0.);

    FEM4EllipticPDE *Solver = FEM4EllipticPDE::New();
        Solver->SetGFunction(Sphere);
        //Solver->SetSigmaFunction(implSigma);
        Solver->SetBoundaryStep(.5);
        Solver->SetGrid(rGrid);
        Solver->Solve("FEM_results.vtr");

    vtkXMLDataSetWriter *Writer = vtkXMLDataSetWriter::New();
        Writer->SetInputData(sigGrid);
        std::string fname("sigma.vtr");
        //Writer->Update();
        //fname.append( Writer->GetDefaultFileExtension() );
        Writer->SetFileName(fname.c_str());
        Writer->Write();
        Writer->Delete();

    // Clean up
    Sphere->Delete();
    Solver->Delete();
    rGrid->Delete();
    xCoords->Delete();
    yCoords->Delete();
    zCoords->Delete();

}

// This function is copyright (C) 2012 Joeseph Cappriotti
vtkSmartPointer<vtkIdList> GetConnectedPoints(int id0, vtkSmartPointer<vtkDataSet> grid){
    vtkSmartPointer<vtkIdList> pointIds = vtkSmartPointer<vtkIdList>::New();
    vtkSmartPointer<vtkIdList> cellList = vtkSmartPointer<vtkIdList>::New();
    grid->GetPointCells(id0, cellList);
    for(int i=0;i<cellList->GetNumberOfIds();++i){
        vtkCell* cell = grid->GetCell(cellList->GetId(i));
        if(cell->GetNumberOfEdges() > 0){
            for(int j=0; j<cell->GetNumberOfEdges(); ++j){
                vtkCell* edge = cell->GetEdge(j);
                vtkIdList* edgePoints=edge->GetPointIds();
                if(edgePoints->GetId(0)==id0){
                    pointIds->InsertUniqueId(edgePoints->GetId(1));
                }else if(edgePoints->GetId(1)==id0){
                    pointIds->InsertUniqueId(edgePoints->GetId(0));
                }
            }
        }
    }
    return pointIds;
}

vtkSmartPointer<vtkIdList> GetConnectedVertices(vtkSmartPointer<vtkDataSet> mesh, int id) {

    std::cout << "number of points " << mesh->GetNumberOfPoints() << std::endl;
    std::cout << "number of cells " << mesh->GetNumberOfCells() << std::endl;

    vtkSmartPointer<vtkIdList> connectedVertices = vtkSmartPointer<vtkIdList>::New();

    //get all cells that vertex 'id' is a part of
    vtkSmartPointer<vtkIdList> cellIdList = vtkSmartPointer<vtkIdList>::New();
    mesh->GetPointCells(id, cellIdList);   // cells using point
    //mesh->GetCellPoints(id, cellIdList); // points defining cell

    cout << "Vertex " << id << " is used in cell(s) ";
    for(vtkIdType i = 0; i < cellIdList->GetNumberOfIds(); i++) {
        cout << cellIdList->GetId(i) << ", ";
    }
    cout << endl;

    // TODO this is where things don't work
    // for each cell, figure out what vertices are connected
    for(vtkIdType i = 0; i < cellIdList->GetNumberOfIds(); i++) {

        cout << "\tcell id " << i << " : " << cellIdList->GetId(i) << endl;
        std::cout << "\tcell has " << mesh->GetCell(i)->GetNumberOfPoints() << " points\n";

        vtkExtractCells *Cell = vtkExtractCells::New();
            Cell->SetInputData(mesh);
            Cell->AddCellRange(i, i);
            Cell->Update();

        //std::cout << *Cell->GetOutput() << std::endl;

        // extract surface
        vtkDataSetSurfaceFilter *surf = vtkDataSetSurfaceFilter::New();
            surf->UseStripsOn();
            surf->SetInputData(Cell->GetOutput());
            surf->Update();

        //std::cout << *surf->GetOutput() << std::endl;

        vtkSmartPointer<vtkTriangleFilter> triangleFilter =
             vtkSmartPointer<vtkTriangleFilter>::New();
        triangleFilter->SetInputData( surf->GetOutput() );
            triangleFilter->Update();

        vtkSmartPointer<vtkExtractEdges> extractEdges =
            vtkSmartPointer<vtkExtractEdges>::New();
        extractEdges->SetInputConnection(triangleFilter->GetOutputPort());
            extractEdges->Update();

        vtkSmartPointer<vtkPolyData> pmesh = extractEdges->GetOutput();
        //vtkSmartPointer<vtkPolyData> pmesh = triangleFilter->GetOutput();
        //vtkSmartPointer<vtkIdList> pconnectedVertices = GetConnectedVertices2(surf->GetOutput(), id);
//
//         //vtkSmartPointer<vtkIdTypeArray> ids =
//         //vtkSmartPointer<vtkIdTypeArray>::New();
//         //ids->SetNumberOfComponents(1);
//
        vtkSmartPointer<vtkIdList> pointIdList = vtkSmartPointer<vtkIdList>::New();
        //mesh->GetPointCells(cellIdList->GetId(i), pointIdList); // returns cells using point
        mesh->GetCellPoints(cellIdList->GetId(i), pointIdList);   // points defining cell

        //vtkSmartPointer<vtkIdList> pconnectedVertices = GetConnectedVertices2(pmesh, id);
        vtkSmartPointer<vtkIdList> pconnectedVertices = GetConnectedVertices2(surf->GetOutput(), id);
        //vtkSmartPointer<vtkIdList> pconnectedVertices = GetConnectedVertices2(triangleFilter->GetOutput(), 1);
        std::cout << "\tPoly Connected vertices: ";
        for(vtkIdType ip = 0; ip < pconnectedVertices->GetNumberOfIds(); ip++) {
            std::cout << pconnectedVertices->GetId(ip) << " ";
            //ids->InsertNextValue(connectedVertices->GetId(i));
            connectedVertices->InsertNextId( pointIdList->GetId( pconnectedVertices->GetId(ip) ) );
        }
        std::cout << std::endl;


        cout << "\tcell " << i << " contains these points ";
        for(vtkIdType ii = 0; ii < pointIdList->GetNumberOfIds(); ii++) {
            cout << pointIdList->GetId(ii) << ", ";
        }
        cout << endl;

        /*
        cout << "\tEnd points are " << pointIdList->GetId(0) << " and " << pointIdList->GetId(1) << endl;

        if(pointIdList->GetId(0) != id) {
            //cout << "Connected to " << pointIdList->GetId(0) << endl;
            connectedVertices->InsertNextId(pointIdList->GetId(0));
        } else {
            //cout << "Connected to " << pointIdList->GetId(1) << endl;
            connectedVertices->InsertNextId(pointIdList->GetId(1));
        }
        */


        #ifdef RENDERTEST
        vtkSmartPointer<vtkDataSetMapper> sphereMapper =
        vtkSmartPointer<vtkDataSetMapper>::New();
            sphereMapper->SetInputConnection(extractEdges->GetOutputPort());
        vtkSmartPointer<vtkActor> sphereActor =
            vtkSmartPointer<vtkActor>::New();
            sphereActor->SetMapper(sphereMapper);


        vtkSmartPointer<vtkIdTypeArray> ids =
            vtkSmartPointer<vtkIdTypeArray>::New();
        ids->SetNumberOfComponents(1);

        std::cout << "RENDERTEST Connected vertices: ";
        for(vtkIdType i = 0; i < connectedVertices->GetNumberOfIds(); i++) {
            std::cout << connectedVertices->GetId(i) << " ";
            ids->InsertNextValue(connectedVertices->GetId(i));
        }
        std::cout << std::endl;

        vtkSmartPointer<vtkDataSetMapper> connectedVertexMapper =
        vtkSmartPointer<vtkDataSetMapper>::New();

        {
            vtkSmartPointer<vtkSelectionNode> selectionNode =
                vtkSmartPointer<vtkSelectionNode>::New();
            selectionNode->SetFieldType(vtkSelectionNode::POINT);
            selectionNode->SetContentType(vtkSelectionNode::INDICES);
            selectionNode->SetSelectionList(ids);

            vtkSmartPointer<vtkSelection> selection =
                vtkSmartPointer<vtkSelection>::New();
            selection->AddNode(selectionNode);

            vtkSmartPointer<vtkExtractSelection> extractSelection =
                vtkSmartPointer<vtkExtractSelection>::New();

            extractSelection->SetInputConnection(0, extractEdges->GetOutputPort());
            extractSelection->SetInputData(1, selection);
            extractSelection->Update();

            vtkSmartPointer<vtkVertexGlyphFilter> glyphFilter =
                vtkSmartPointer<vtkVertexGlyphFilter>::New();
            glyphFilter->SetInputConnection(extractSelection->GetOutputPort());
            glyphFilter->Update();

            connectedVertexMapper->SetInputConnection(glyphFilter->GetOutputPort());
        }

        vtkSmartPointer<vtkActor> connectedVertexActor =
            vtkSmartPointer<vtkActor>::New();
        connectedVertexActor->SetMapper(connectedVertexMapper);
        connectedVertexActor->GetProperty()->SetColor(1,0,0);
        connectedVertexActor->GetProperty()->SetPointSize(5);

        vtkSmartPointer<vtkDataSetMapper> queryVertexMapper =
            vtkSmartPointer<vtkDataSetMapper>::New();

        {
            vtkSmartPointer<vtkIdTypeArray> ids =
                vtkSmartPointer<vtkIdTypeArray>::New();
            ids->SetNumberOfComponents(1);
            ids->InsertNextValue(id);

            vtkSmartPointer<vtkSelectionNode> selectionNode =
                vtkSmartPointer<vtkSelectionNode>::New();
            selectionNode->SetFieldType(vtkSelectionNode::POINT);
            selectionNode->SetContentType(vtkSelectionNode::INDICES);
            selectionNode->SetSelectionList(ids);

            vtkSmartPointer<vtkSelection> selection =
                vtkSmartPointer<vtkSelection>::New();
            selection->AddNode(selectionNode);

            vtkSmartPointer<vtkExtractSelection> extractSelection =
                vtkSmartPointer<vtkExtractSelection>::New();

            extractSelection->SetInputConnection(0, extractEdges->GetOutputPort());
            extractSelection->SetInputData(1, selection);
            extractSelection->Update();

            vtkSmartPointer<vtkVertexGlyphFilter> glyphFilter =
                vtkSmartPointer<vtkVertexGlyphFilter>::New();
            glyphFilter->SetInputConnection(extractSelection->GetOutputPort());
            glyphFilter->Update();

            queryVertexMapper->SetInputConnection(glyphFilter->GetOutputPort());
        }

        vtkSmartPointer<vtkActor> queryVertexActor =
            vtkSmartPointer<vtkActor>::New();
        queryVertexActor->SetMapper(queryVertexMapper);
        queryVertexActor->GetProperty()->SetColor(0,1,0);
        queryVertexActor->GetProperty()->SetPointSize(5);

        // Create the usual rendering stuff.
        vtkRenderer *renderer = vtkRenderer::New();
        vtkRenderWindow *renWin = vtkRenderWindow::New();
        renWin->AddRenderer(renderer);
        vtkRenderWindowInteractor *iren = vtkRenderWindowInteractor::New();
        iren->SetRenderWindow(renWin);

        renderer->AddActor(queryVertexActor);
        renderer->AddActor(connectedVertexActor);
        renderer->AddActor(sphereActor);

        //renderer->AddActor(wireActor);
        //renderer->AddActor(owireActor);
        renderer->SetBackground(.3,.2,.1);
        renderer->ResetCamera();
        renderer->GetActiveCamera()->Elevation(-150.0);
        renderer->GetActiveCamera()->Azimuth(0.0);
        renderer->GetActiveCamera()->Roll(0.0);
        //renderer->GetActiveCamera()->Pitch(-45.0);
        renderer->GetActiveCamera()->Zoom(1.0);

        renWin->SetSize(300,300);

        // interact with data
        renWin->Render();
        iren->Start();

        // clean up rendering stuff
        //plane->Delete();
        //rgridMapper->Delete();
        //wireActor->Delete();
        //owireActor->Delete();
        //outline->Delete();
        //outlineMapper->Delete();
        renderer->Delete();
        renWin->Delete();
        iren->Delete();
        #endif // RENDERTEST

        surf->Delete();
        Cell->Delete();
    }

    return connectedVertices;
}

// from http://www.vtk.org/Wiki/VTK/Examples/Cxx/PolyData/VertexConnectivity
vtkSmartPointer<vtkIdList> GetConnectedVertices2(vtkSmartPointer<vtkPolyData> mesh, int id) {

    std::cout << "mesh points " << mesh->GetNumberOfPoints() << std::endl;

    vtkSmartPointer<vtkIdList> connectedVertices =
        vtkSmartPointer<vtkIdList>::New();

    //get all cells that vertex 'id' is a part of
    vtkSmartPointer<vtkIdList> cellIdList =
        vtkSmartPointer<vtkIdList>::New();
    mesh->GetPointCells(id, cellIdList);

    /*
    for(vtkIdType i = 0; i < cellIdList->GetNumberOfIds(); i++) {
        cout << cellIdList->GetId(i) << ", ";
    }
    cout << endl;
    */

    for(vtkIdType i = 0; i < cellIdList->GetNumberOfIds(); i++) {
        //cout << "id " << i << " : " << cellIdList->GetId(i) << endl;

        vtkSmartPointer<vtkIdList> pointIdList =
            vtkSmartPointer<vtkIdList>::New();
        mesh->GetCellPoints(cellIdList->GetId(i), pointIdList);
        //cout << "End points are " << pointIdList->GetId(0) << " and " << pointIdList->GetId(1) << endl;
        if(pointIdList->GetId(0) != id) {
            //cout << "Connected to " << pointIdList->GetId(0) << endl;
            connectedVertices->InsertNextId(pointIdList->GetId(0));
        } else {
            //cout << "Connected to " << pointIdList->GetId(1) << endl;
            connectedVertices->InsertNextId(pointIdList->GetId(1));
        }
    }
    return connectedVertices;
}