#include <comp.hpp>
#include <solve.hpp>

using namespace std;
ostream * testst;

#include "../../../pebbles_old/libsrc/matrix/mydefs.hh"
#include "../../../pebbles_old/pebbles/extern.hh"
#include "../../../pebbles_old/pebbles/pebbles.hh"

/*
#include "../../../pebbles_old/libsrc/matrix/mydefs.hh"
#include "../../../pebbles_old/libsrc/prepro/outback.hh"
*/

namespace ngcomp
{
  using namespace ngcomp;



  AMGPreconditioner :: AMGPreconditioner (PDE * pde, Flags & flags)
  {
    bfa = pde->GetBilinearForm (flags.GetStringFlag ("bilinearform", NULL));
    bfl = pde->GetBilinearForm (flags.GetStringFlag ("bilinearformlaplace", NULL));
  }
  
  AMGPreconditioner :: ~AMGPreconditioner()
  {
    ;
  }
  
  void AMGPreconditioner :: Update ()
  {
    cout << "Update AMG Preconditioner" << endl;
    if (dynamic_cast<const NedelecFESpace*> (&bfa->GetFESpace()))
      pre = new PebblesEdgeMatrix (*bfa, *bfl);
    else
      pre = new PebblesMatrix (*bfa);
  }
  


  PebblesMatrix :: PebblesMatrix (const BilinearForm & bfa)
  {
    using namespace RS_AMG;

    const FESpace & fespace = bfa.GetFESpace();
    const MeshAccess & ma = fespace.GetMeshAccess();

    ARRAY<int> dnums;
    ElementTransformation eltrans;
	  
    int ndof = fespace.GetNDof();
    BitArray useddof(ndof);
    useddof.Clear();
	  
    int ne = ma.GetNE();
	  
    cout << "Assemble pebbles matrix " << flush;
	  
    LocalHeap locheap (2000000);
	  


    outback = new SparseOutBack(ndof, ne, 0, 0, 1, 0, 0);
    outback -> Init();
    outback -> Create (1);

    cout << ndof << " " << ne << endl;
    cout << "prepare matrix graph " << endl;

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
		
	fespace.GetDofNrs (i, dnums);
	
	for (int j=0; j<dnums.Size(); j++)
	  dnums[j]++;
	
	outback -> GetConnection(&dnums[0], i+1, dnums.Size(), 1, 1); 
      }

    outback -> CreateMatrix();
    //outback -> Print ();

    cout << "assemble matrix" << endl;

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
	locheap.CleanUp();
		
	const FiniteElement & fel = fespace.GetFE (i);
	ma.GetElementTransformation (i, eltrans, locheap);
	fespace.GetDofNrs (i, dnums);
	
	for (int k=0; k<dnums.Size(); k++)
	  dnums[k]++;
	    
	for (int j = 0; j < bfa.NumIntegrators(); j++)
	  {
	    const BilinearFormIntegrator & bfi = *bfa.GetIntegrator(j);
	    if (bfi.BoundaryForm()) continue;
		   
	    FlatMatrix<double> elmat;
	    bfi.AssembleElementMatrix (fel, eltrans, elmat, locheap);

	    fespace.TransformMat (i, false, elmat, TRANSFORM_MAT_LEFT_RIGHT);

	    outback -> GetElement (&elmat(0,0), &dnums[0], i+1, dnums.Size(),
				   0, 0, 1, 1);
	  }
      }


    cout << " boundary terms ";
    
    int nse = ma.GetNSE();
    for (int i = 0; i < nse; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
	
	locheap.CleanUp();
	
	
	const FiniteElement & fel = fespace.GetSFE (i);
	ma.GetSurfaceElementTransformation (i, eltrans, locheap);
	fespace.GetSDofNrs (i, dnums);

	for (int k=0; k<dnums.Size(); k++)
	  dnums[k]++;
		
	for (int j = 0; j < bfa.NumIntegrators(); j++)
	  {
	    const BilinearFormIntegrator & bfi = *bfa.GetIntegrator(j);
	    if (!bfi.BoundaryForm()) continue;
		    
	    FlatMatrix<double> elmat;
	    bfi.AssembleElementMatrix (fel, eltrans, elmat, locheap);
		    
	    fespace.TransformMat (i, true, elmat, TRANSFORM_MAT_LEFT_RIGHT);

	    
	    outback -> GetElement (&elmat(0,0), &dnums[0], -17, dnums.Size(),
				   0, 0, 1, 1);
	  }
      }
	  
    cout << endl;
	  
    cout << "matrix done" << endl;
    
    //outback -> Print();
    outback -> SetCount();
    outback -> ConstructMatrixHierarchy();

  }

  void PebblesMatrix :: 
  Mult (const BaseVector & x, BaseVector & y) const
  {
    outback -> ApplyPrecond(const_cast<double*> (static_cast<const double*> (x.Data())),
			    static_cast<double*> (y.Data()),
			    1);
  }














  PebblesEdgeMatrix :: PebblesEdgeMatrix (const BilinearForm & bfa,
					  const BilinearForm & bfl)
  {
    using namespace RS_AMG;

    const NedelecFESpace & fespace = 
      dynamic_cast<const NedelecFESpace&> (bfa.GetFESpace());
    const FESpace & fespacelap = bfl.GetFESpace();
    const MeshAccess & ma = fespace.GetMeshAccess();

    ARRAY<int> vnums;
    ElementTransformation eltrans;
	  
    int ndof = fespace.GetNDof();
    int ne = ma.GetNE();
    int nv = ma.GetNV();
	 
    int level = ma.GetNLevels();
    int cnt = 0;
    for (int i = 0; i < ndof; i++)
      if (fespace.FineLevelOfEdge(i) == level)
	cnt++;

    cout << "Assemble pebbles matrix " << flush;
	  
    LocalHeap locheap (2000000);
	  


    outback = new WhitneyOutBack(nv, ne, ne, 0, cnt, 1);
    outback -> Init();
    outback -> Create ();

    cout << "ndof = " << ndof << " cnt = " << cnt << " ne = " << ne << endl;
    cout << "prepare matrix graph " << endl;

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
		
	ma.GetElPNums (i, vnums);
	
	for (int j=0; j< vnums.Size(); j++)
	  vnums[j]++;
	
	outback -> GetConnection(&vnums[0], i+1, 4, 2, 1); 
      }

    outback -> CreateMatrix(2);

    cout << "assemble help matrix" << endl;

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
	locheap.CleanUp();
		
	const FiniteElement & fel = fespacelap.GetFE (i);
	ma.GetElementTransformation (i, eltrans, locheap);
	fespacelap.GetDofNrs (i, vnums);

	for (int k=0; k<vnums.Size(); k++)
	  vnums[k]++;

	for (int j = 0; j < bfl.NumIntegrators(); j++)
	  {
	    const BilinearFormIntegrator & bfi = *bfl.GetIntegrator(j);
	    if (bfi.BoundaryForm()) continue;
		   
	    FlatMatrix<double> elmat;
	    bfi.AssembleElementMatrix (fel, eltrans, elmat, locheap);

	    fespacelap.TransformMat (i, false, elmat, TRANSFORM_MAT_LEFT_RIGHT);

	    outback -> GetElement (&elmat(0,0), &vnums[0], i+1, 4, 
				   0, 0, 2, 1);
	  }
      }

    //    outback -> UpdateAuxMat();
    outback -> CreateNode2Edge();


    const EDGE * edges = ElementTopology::GetEdges (ET_TET);

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
		
	ma.GetElPNums (i, vnums);
	
	ARRAY<int> edgenr(6);
	for (int l = 0; l < 6; l++)
	  {
	    int pi1 = vnums[edges[l][0]];
	    int pi2 = vnums[edges[l][1]];
	    if (pi1 > pi2) swap (pi1, pi2);
	    edgenr[l] = outback->Node2Edge (pi1+1, pi2+1);
	  }
	outback -> GetConnection(&edgenr[0], i+1, 6, 1, 1); 
      }

    outback -> CreateMatrix(3);

    for (int i = 0; i < ne; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
	locheap.CleanUp();
		
	const FiniteElement & fel = fespace.GetFE (i);
	ma.GetElementTransformation (i, eltrans, locheap);
		
	ma.GetElPNums (i, vnums);
	
	int edgenr[6];
	for (int l = 0; l < 6; l++)
	  {
	    int pi1 = vnums[edges[l][0]];
	    int pi2 = vnums[edges[l][1]];
	    if (pi1 > pi2) swap (pi1, pi2);
	    edgenr[l] = outback->Node2Edge (pi1+1, pi2+1);
	  }

	for (int j = 0; j < bfa.NumIntegrators(); j++)
	  {
	    const BilinearFormIntegrator & bfi = *bfa.GetIntegrator(j);
	    if (bfi.BoundaryForm()) continue;
		   
	    FlatMatrix<double> elmat;
	    bfi.AssembleElementMatrix (fel, eltrans, elmat, locheap);

	    fespace.TransformMat (i, false, elmat, TRANSFORM_MAT_LEFT_RIGHT);
	    
	    outback -> GetElement (&elmat(0,0), edgenr, i+1, 6,
				   0, 0, 1, 1);
	  }
      }

    /*
    cout << " boundary terms ";
    
    int nse = ma.GetNSE();
    for (int i = 0; i < nse; i++)
      {
	if (i % 100 == 0)
	  cout << "." << flush;
	
	locheap.CleanUp();
	
	
	const FiniteElement & fel = fespace.GetSFE (i);
	ma.GetSurfaceElementTransformation (i, eltrans, locheap);
	fespace.GetSDofNrs (i, dnums);
		
	for (int k=0; k<dnums.Size(); k++)
	dnums[k]++;
	    
	for (int j = 0; j < bfa.NumIntegrators(); j++)
	  {
	    const BilinearFormIntegrator & bfi = *bfa.GetIntegrator(j);
	    if (!bfi.BoundaryForm()) continue;
		    
	    FlatMatrix<double> elmat;
	    bfi.AssembleElementMatrix (fel, eltrans, elmat, locheap);
		    
	    fespace.TransformMat (i, true, elmat, TRANSFORM_MAT_LEFT_RIGHT);

	    outback -> GetElement (&elmat(0,0), &dnums[0], -17, dnums.Size(),
				   0, 0, 1, 1);
	  }
      }
	  
    cout << endl;
    */


    reorder.SetSize (fespace.GetNDof());
    for (int i = 0; i < reorder.Size(); i++)
      {
	int pi1 = fespace.EdgePoint1(i);
	int pi2 = fespace.EdgePoint2(i);
	if (pi1 > pi2) swap (pi1, pi2);
	if (fespace.FineLevelOfEdge(i) == level)
	  reorder[i] = outback->Node2Edge (pi1+1, pi2+1)-1;
	else
	  reorder[i] = -1;
      }

    (*testout) << "reorder = " << reorder << endl;

    cout << "matrix done" << endl;
    
    outback -> SetCount();
    outback -> ConstructMatrixHierarchy();

  }

  void PebblesEdgeMatrix :: 
  Mult (const BaseVector & x, BaseVector & y) const
  {
    FlatVector<double> fx = dynamic_cast<const VVector<double>&> (x).FV();
    FlatVector<double> fy = dynamic_cast<const VVector<double>&> (y).FV();
    
    Vector<double> hx(fx.Size());
    Vector<double> hy(fy.Size());

    hx = 0.0;

    for (int i = 0; i < fx.Size(); i++)
      if (reorder[i] != -1)
	hx(reorder[i]) = fx(i);

    hy = 0.0;

    outback -> ApplyPrecond(&hx(0), &hy(0), 1);

    fy = 0.0;
    for (int i = 0; i < fx.Size(); i++)
      if (reorder[i] != -1)
	fy(i) = hy(reorder[i]);
  }






}