/*********************************************************************/
/* File:   commutingamg.cpp                                          */
/* Author: Joachim Schoeberl                                         */
/* Date:   15. Aug. 2002                                             */
/*********************************************************************/

#include <la.hpp>
namespace ngla
{
  using namespace ngla;

  
  AMG_H1 :: AMG_H1 (ARRAY<INT<2> > & e2v,
		    ARRAY<double> & weighte,
		    int levels)
  {
    int i, j, k;

    // find number of vertices
    int ne = e2v.Size();
    int nv = 0;
    for (i = 0; i < ne; i++)
      for (j = 0; j < 2; j++)
	if (e2v[i][j] > nv)
	  nv = e2v[i][j];
    nv++;

    cout << "ne = " << ne << ", nv = " << nv << endl;

    if (nv < 20 || levels == 0)
      {
	recAMG = 0;
	prol = 0;
	return;
      }

    ARRAY<int> vcoarse(nv), connected(nv);
    ARRAY<bool> edge_collapse(ne);
    ARRAY<double> edge_collapse_weight(ne);
    edge_collapse = 0;
    ARRAY<int> v2edge(nv);
    

    /*
    // compute weight to collapse edge
    ARRAY<double> vstrength(nv);
    vstrength = 0.0;
    for (i = 0; i < weighte.Size(); i++)
      for (j = 0; j < 2; j++)
	vstrength[e2v[i][j]] += sqr (weighte[i]);

    for (i = 0; i < ne; i++)
      edge_collapse_weight[i] = 
	sqr (weighte[i]) / min2 (vstrength[e2v[i][0]], vstrength[e2v[i][1]]);
	//	sqr (weighte[i]) / sqrt (vstrength[e2v[i][0]] * vstrength[e2v[i][1]]);

	*/
    // compute weight to collapse edge
    ARRAY<double> vstrength(nv);
    vstrength = 0.0;
    for (i = 0; i < weighte.Size(); i++)
      for (j = 0; j < 2; j++)
	vstrength[e2v[i][j]] += weighte[i];

    for (i = 0; i < ne; i++)
      {
	double vstr1 = vstrength[e2v[i][0]];
	double vstr2 = vstrength[e2v[i][1]];
	edge_collapse_weight[i] = 
	  weighte[i] * (vstr1+vstr2) / (vstr1 * vstr2);
      }
    
    v2edge = -1;
    edge_collapse = 0;
    bool changed;

    // figure out best edges to collapse (iterative improvement)
    do
      {
	changed = 0;
	for (i = 0; i < ne; i++)
	  if (edge_collapse_weight[i] > 0.1)
	    {
	      if (v2edge[e2v[i][0]] == -1 && v2edge[e2v[i][1]] == -1)
		
		{
		  edge_collapse[i] = 1;
		  v2edge[e2v[i][0]] = i;
		  v2edge[e2v[i][1]] = i;
		  changed = 1;
		}	    
	    
	      else
		
		{
		  for (j = 0; j < 2; j++)
		    {
		      int pi1 = e2v[i][j];
		      int pi2 = e2v[i][1-j];
		      
		      if (v2edge[pi1] != -1 && v2edge[pi2] == -1 &&
			  edge_collapse_weight[i] > 
			  edge_collapse_weight[v2edge[pi1]])
			{
			  int remove = v2edge[pi1];
			  edge_collapse[i] = 1;
			  edge_collapse[remove] = 0;
			  v2edge[e2v[remove][0]] = -1;
			  v2edge[e2v[remove][1]] = -1;
			  v2edge[pi1] = i;
			  v2edge[pi2] = i;
			  changed = 1;
			  break;
			}
		    }
		}
	    }
      }
    while (changed);

    // compute fine vertex to coarse vertex map (vcoarse)
    for (i = 0; i < nv; i++)
      connected[i] = i;

    for (i = 0; i < ne; i++)
      if (edge_collapse[i])
	{
	  int pi1 = e2v[i][0];
	  int pi2 = e2v[i][1];
	  if (vstrength[pi1] > vstrength[pi2])
	    connected[pi2] = pi1;
	  else
	    connected[pi1] = pi2;
	}

    int ncv = 0;
    for (i = 0; i < nv; i++)
      if (connected[i] == i)
	{
	  vcoarse[i] = ncv;
	  ncv++;
	}
    for (i = 0; i < nv; i++)
      if (connected[i] != i)
	vcoarse[i] = vcoarse[connected[i]];




    // compute fine edge to coarse edge map (ecoarse)
    HashTable<INT<2>, int> ht_ecoarse(e2v.Size());
    for (i = 0; i < e2v.Size(); i++)
      {
	INT<2> ce;
	for (int j = 0; j < 2; j++)
	  ce[j] = vcoarse[e2v[i][j]];
	ce.Sort();
	if (ce[0] != ce[1])
	  ht_ecoarse.Set (ce, -1);
      }
    
    ARRAY<INT<2> > ce2v;
    for (i = 0; i < ht_ecoarse.Size(); i++)
      for (j = 0; j < ht_ecoarse.EntrySize(i); j++)
	{
	  INT<2> ce;
	  int efi;
	  ht_ecoarse.GetData (i, j, ce, efi);
	  efi = ce2v.Size();
	  ce2v.Append (ce);
	  ht_ecoarse.SetData (i, j, ce, efi);
	}

    ARRAY<int> ecoarse(ne);
    for (i = 0; i < e2v.Size(); i++)
      {
	INT<2> ce;
	for (int j = 0; j < 2; j++)
	  ce[j] = vcoarse[e2v[i][j]];
	ce.Sort();
	if (ce[0] != ce[1])
	  ecoarse[i] = ht_ecoarse.Get(ce);
	else
	  ecoarse[i] = -1;
      }

    // coarse edge weights:
    ARRAY<double> cweighte(ce2v.Size());
    cweighte = 0;
    for (i = 0; i < e2v.Size(); i++)
      if (ecoarse[i] != -1)
	cweighte[ecoarse[i]] += weighte[i];


    // compute prolongation matrix 
    // prol  ... for matrix projection

    // piecewise constant:
    ARRAY<int> nne(nv);
    nne = 1;

    prol = new SparseMatrix<double> (nne);
    for (i = 0; i < nv; i++)
      {
	prol->CreatePosition (i, vcoarse[i]);
	(*prol)(i, vcoarse[i]) = 1;
      }

    recAMG = new AMG_H1 (ce2v, cweighte, levels-1);
  }
  

  AMG_H1 :: ~AMG_H1 ()
  {
    delete prol;
    delete recAMG;
  }
  

  void AMG_H1 :: ComputeMatrices (const BaseSparseMatrix & mat)
  {
    pmat = &mat;
    jacobi = mat.CreateJacobiPrecond ();

    if (recAMG)
      {
	coarsemat = mat.Restrict (*prol);
	recAMG -> ComputeMatrices (*coarsemat);
	inv = 0;
      }
    else 
      inv = mat.InverseMatrix();
  }

  int AMG_H1 :: NZE() const
  {
    int nze = pmat->NZE();
    if (recAMG)
      nze += recAMG->NZE();
    return nze;
  }

  
  void AMG_H1 :: Mult (const BaseVector & x, BaseVector & y) const
  {
    if (inv)
      {
	y = (*inv) * x;
	return;
      }

    BaseVector & hv = *pmat->CreateVector();
    BaseVector & wc = *coarsemat->CreateVector();
    BaseVector & dc = *coarsemat->CreateVector();

    y = 0;
    jacobi->GSSmooth (y, x);

    if (recAMG)
      {
	hv = x - (*pmat) * y;
	dc = Transpose (*prol) * hv;

	if (recAMG)
	  recAMG -> Mult (dc, wc);

	y += (*prol) * wc;
      }

    jacobi->GSSmoothBack (y, x);

    delete &hv;
    delete &wc;
    delete &dc;
  }
      



  /* ************************* HCurl AMG ********************************* */




  void SortFace (INT<4> & face)
  {
    if (face[3] == -1)
      {
	if (face[0] > face[1]) swap (face[0], face[1]);
	if (face[1] > face[2]) swap (face[1], face[2]);
	if (face[0] > face[1]) swap (face[0], face[1]);
      }
    else
      {
	while (face[1] < face[0] || face[2] < face[0] || face[3] < face[0])
	  {
	    int hi = face[0];
	    face[0] = face[1];
	    face[1] = face[2];
	    face[2] = face[3];
	    face[3] = hi;
	  }
      }
  }


  
  AMG_HCurl :: AMG_HCurl (ARRAY<INT<2> > & e2v,
			  ARRAY<INT<4> > & f2v,
			  ARRAY<double> & weighte,
			  ARRAY<double> & weightf,
			  int levels)
  {
    int i, j, k;

    // find number of vertices
    int ne = e2v.Size();
    int nf = f2v.Size();
    int nv = 0;
    for (i = 0; i < ne; i++)
      for (j = 0; j < 2; j++)
	if (e2v[i][j] > nv)
	  nv = e2v[i][j];
    nv++;

    cout << "nfa = " << nf << ", ned = " << ne << ", nv = " << nv << endl;

    ARRAY<int> vcoarse(nv), connected(nv);
    ARRAY<bool> edge_collapse(ne);
    ARRAY<double> edge_collapse_weight(ne);
    edge_collapse = 0;
    ARRAY<int> v2edge(nv);



    // compute face 2 edge table
    HashTable<INT<2>, int> ht_edge(e2v.Size());
    ARRAY<INT<4> > f2e(nf);

    for (i = 0; i < ne; i++)
      if (e2v[i][0] != -1)
	{
	  INT<2> ce = e2v[i];
	  ce.Sort();
	  ht_edge.Set (ce, i);
	}

    for (i = 0; i < f2v.Size(); i++)
      if (f2v[i][0] != -1)
	{
	  int nfv = (f2v[i][3] == -1) ? 3 : 4;
	  f2e[i][3] = -1;
	  for (j = 0; j < nfv; j++)
	    {
	      INT<2> ce;
	      ce[0] = f2v[i][j];
	      ce[1] = f2v[i][(j+1)%nfv];
	      ce.Sort();
	      if (!ht_edge.Used (ce))
		{
		  /*
		    cout << "Err: unused edge, " 
		    << "face = " << f2v[i] << endl;
		*/
		  (*testout) << "Err: unused edge, " 
			     << "face = " << f2v[i] << endl;
		  f2e[i][j] = -1;
		}
	      else
		f2e[i][j] = ht_edge.Get(ce);
	    }
	}
      else
	for (j = 0; j < 4; j++)
	  f2e[i][j] = -1;
	  
    
    //    (*testout) << "weightf = " << weightf << endl;
    
    ARRAY<double> sume(ne);
    sume = 0;
    for (i = 0; i < nf; i++)
      for (j = 0; j < 4; j++)
	if (f2e[i][j] != -1)
	  sume[f2e[i][j]] += sqr (weightf[i]);

    //    (*testout) << "sume = " << sume << endl;
    
    ARRAY<double> face_collapse_weight(nf);
    for (i = 0; i < nf; i++)
      {
	double mine = 1e99;
	for (j = 0; j < 4; j++)
	  if (f2e[i][j] != -1)
	    mine = min2 (mine, sume[f2e[i][j]]);
	face_collapse_weight[i] = sqr (weightf[i]) / mine;
	/*
	double maxe = 0;
	for (j = 0; j < 4; j++)
	  if (f2e[i][j] != -1)
	    maxe = max2 (maxe, sume[f2e[i][j]]);
	face_collapse_weight[i] = sqr (weightf[i]) / maxe;
	*/
      }

    edge_collapse_weight = 1e99;
    for (i = 0; i < nf; i++)
      for (j = 0; j < 4; j++)
	if (f2e[i][j] != -1)
	  edge_collapse_weight[f2e[i][j]] = 
	    min2 (edge_collapse_weight[f2e[i][j]],
		  face_collapse_weight[i]);

    // figure out best edges to collapse (iterative improvement)
    v2edge = -1;
    edge_collapse = 0;
    bool changed;
    do
      {
	changed = 0;
	for (i = 0; i < ne; i++)
	  if (edge_collapse_weight[i] > 0.05 && 
	      e2v[i][0] != -1)
	    {
	      if (v2edge[e2v[i][0]] == -1 && v2edge[e2v[i][1]] == -1)
		
		{
		  edge_collapse[i] = 1;
		  v2edge[e2v[i][0]] = i;
		  v2edge[e2v[i][1]] = i;
		  changed = 1;
		}	    
	    
	      else
		
		{
		  for (j = 0; j < 2; j++)
		    {
		      int pi1 = e2v[i][j];
		      int pi2 = e2v[i][1-j];
		      
		      if (v2edge[pi1] != -1 && v2edge[pi2] == -1 &&
			  edge_collapse_weight[i] > 
			  edge_collapse_weight[v2edge[pi1]])
			{
			  int remove = v2edge[pi1];
			  edge_collapse[i] = 1;
			  edge_collapse[remove] = 0;
			  v2edge[e2v[remove][0]] = -1;
			  v2edge[e2v[remove][1]] = -1;
			  v2edge[pi1] = i;
			  v2edge[pi2] = i;
			  changed = 1;
			  break;
			}
		    }
		}
	    }
      }
    while (changed);

    // compute fine vertex to coarse vertex map (vcoarse)
    for (i = 0; i < nv; i++)
      connected[i] = i;

    for (i = 0; i < ne; i++)
      if (edge_collapse[i])
	{
	  int pi1 = e2v[i][0];
	  int pi2 = e2v[i][1];
	  int mini = min2(pi1, pi2);
	  connected[pi1] = mini;
	  connected[pi2] = mini;
	}

    int ncv = 0;
    for (i = 0; i < nv; i++)
      if (connected[i] == i)
	{
	  vcoarse[i] = ncv;
	  ncv++;
	}
      else
	vcoarse[i] = vcoarse[connected[i]];


    // compute fine edge to coarse edge map (ecoarse)
    HashTable<INT<2>, int> ht_ecoarse(ne);
    for (i = 0; i < ne; i++)
      if (e2v[i][0] != -1)
	{
	  INT<2> ce;
	  for (int j = 0; j < 2; j++)
	    ce[j] = vcoarse[e2v[i][j]];
	  if (ce[0] != ce[1])
	    {
	      if (ce[0] > ce[1]) Swap (ce[0], ce[1]);
	      ht_ecoarse.Set (ce, -1);
	    }
	}
    
    ARRAY<INT<2> > ce2v;
    for (i = 0; i < ht_ecoarse.Size(); i++)
      for (j = 0; j < ht_ecoarse.EntrySize(i); j++)
	{
	  INT<2> ce;
	  int efi;
	  ht_ecoarse.GetData (i, j, ce, efi);
	  efi = ce2v.Size();
	  ce2v.Append (ce);
	  ht_ecoarse.SetData (i, j, ce, efi);
	}

    ARRAY<int> ecoarse(ne);
    ARRAY<short> ecoarsesign(ne);
    ecoarsesign = 1;
    for (i = 0; i < ne; i++)
      {
	if (e2v[i][0] == -1)
	  {
	    ecoarsesign[i] = 0;
	    ecoarse[i] = -1;
	    continue;
	  }

	INT<2> ce;
	for (int j = 0; j < 2; j++)
	  ce[j] = vcoarse[e2v[i][j]];
	
	if (ce[0] == ce[1])
	  {
	    ecoarsesign[i] = 0;
	    ecoarse[i] = -1;
	  }
	else 
	  {
	    if (ce[0] > ce[1]) 
	      {
		Swap (ce[0], ce[1]);
		ecoarsesign[i] = -1;
	      }
	    ecoarse[i] = ht_ecoarse.Get(ce);
	  }
      }

    // coarse edge weights:
    ARRAY<double> cweighte(ce2v.Size());
    cweighte = 0;
    for (i = 0; i < e2v.Size(); i++)
      if (ecoarse[i] != -1)
	cweighte[ecoarse[i]] += weighte[i];






    // compute fine face to coarse face map (fcoarse)
    HashTable<INT<4>, int> ht_fcoarse(nf);
    for (i = 0; i < nf; i++)
      {
	bool valid = 1;
	for (j = 0; j < 3; j++)
	  if (f2e[i][j] == -1)
	    valid = 0;
	if (!valid) continue;


	INT<4> cf;
	for (j = 0; j < 4; j++)
	  if (f2v[i][j] != -1)
	    cf[j] = vcoarse[f2v[i][j]];
	  else
	    cf[j] = -1;
	SortFace(cf);
	bool degenerated = 0;
	for (j = 0; j < 3; j++)
	  for (k = j+1; k < 4; k++)
	    if (cf[j] == cf[k])
	      degenerated = 1;
	//	if (cf[0] != cf[1] && cf[1] != cf[2] && cf[2] != cf[3] && cf[3] != cf[0])
	if (!degenerated)
	  ht_fcoarse.Set (cf, -1);
      }
    
    ARRAY<INT<4> > cf2v;
    for (i = 0; i < ht_fcoarse.Size(); i++)
      for (j = 0; j < ht_fcoarse.EntrySize(i); j++)
	{
	  INT<4> cf;
	  int cfi;
	  ht_fcoarse.GetData (i, j, cf, cfi);
	  cfi = cf2v.Size();
	  cf2v.Append (cf);
	  ht_fcoarse.SetData (i, j, cf, cfi);
	}

    ARRAY<int> fcoarse(nf);
    fcoarse = -1;
    for (i = 0; i < nf; i++)
      {
	bool valid = 1;
	for (j = 0; j < 3; j++)
	  if (f2e[i][j] == -1)
	    valid = 0;
	if (!valid) continue;


	INT<4> cf;
	for (j = 0; j < 4; j++)
	  if (f2v[i][j] != -1)
	    cf[j] = vcoarse[f2v[i][j]];
	  else
	    cf[j] = -1;
	SortFace(cf);
	bool degenerated = 0;
	for (j = 0; j < 3; j++)
	  for (k = j+1; k < 4; k++)
	    if (cf[j] == cf[k])
	      degenerated = 1;
	// if (cf[0] != cf[1] && cf[1] != cf[2] && cf[2] != cf[3] && cf[3] != cf[0])
	if (!degenerated)
	  fcoarse[i] = ht_fcoarse.Get(cf);
      }

    // coarse face weights:
    ARRAY<double> cweightf(cf2v.Size());
    cweightf = 0;
    for (i = 0; i < nf; i++)
      if (fcoarse[i] != -1)
	cweightf[fcoarse[i]] += weightf[i];





    // compute prolongation matrix
    ARRAY<int> nne(ne);
    nne = 0;
    for (i = 0; i < ne; i++)
      if (ecoarse[i] != -1)
	nne[i] = 1;

    prol = new SparseMatrix<double> (nne);
    for (i = 0; i < ne; i++)
      if (ecoarse[i] != -1)
	{
	  prol->CreatePosition (i, ecoarse[i]);
	  (*prol)(i, ecoarse[i]) = ecoarsesign[i];
	}

    // compute gradient matrix:
    for (i = 0; i < ne; i++)
      if (e2v[i][0] == -1)
	nne[i] = 0;
      else
	nne[i] = 2;

    grad = new SparseMatrix<double> (nne);
    for (i = 0; i < ne; i++)
      if (e2v[i][0] != -1)
      {
	grad->CreatePosition (i, e2v[i][0]);
	grad->CreatePosition (i, e2v[i][1]);
	(*grad)(i, e2v[i][0]) = 1;
	(*grad)(i, e2v[i][1]) = -1;
      }

    if (nv > 50 && ncv < nv && levels != 0)
      {
	recAMG = new AMG_HCurl (ce2v, cf2v, cweighte, cweightf, levels-1);
	h1AMG = new AMG_H1 (e2v, weighte, levels);
      }
    else
      {
	recAMG = 0;
	h1AMG = 0;
      }
  }


  AMG_HCurl :: ~AMG_HCurl ()
  {
    delete prol;
    delete recAMG;
  }
  



  void AMG_HCurl :: ComputeMatrices (const BaseSparseMatrix & mat)
  {
    cout << "compute HCurl matrices" << endl;

    pmat = &mat;
    coarsemat = mat.Restrict (*prol);
    jacobi = mat.CreateJacobiPrecond ();

    h1mat = mat.Restrict (*grad);
    //    dynamic_cast<SparseMatrixSymmetric<Mat<1,1> >&> (*h1mat) (0,0) += 1;
    dynamic_cast<SparseMatrixSymmetric<double>&> (*h1mat)(0,0) += 1;

    if (recAMG)
      {
	recAMG -> ComputeMatrices (*coarsemat);
	h1AMG -> ComputeMatrices (*h1mat);
	inv = 0;
	/*
	EigenSystem eigen (mat, *this);
	eigen.Calc();
	eigen.PrintEigenValues (cout);
	*/
      }
    else 
      {
	cout << "cal inverse, size = " << mat.Height() << endl;
	inv = mat.InverseMatrix();
      }
  }



  int AMG_HCurl :: NZE() const
  {
    int nze = pmat->NZE() + h1mat->NZE();
    if (recAMG) 
      nze += recAMG->NZE() + h1AMG->NZE();

    return nze;
  }


  
  void AMG_HCurl :: Mult (const BaseVector & x, BaseVector & y) const
  {
    if (inv)
      {
	y = (*inv) * x;
	return;
      }

    BaseVector & hv = *pmat->CreateVector();
    BaseVector & wc = *coarsemat->CreateVector();
    BaseVector & dc = *coarsemat->CreateVector();

    BaseVector & wh1 = *h1mat->CreateVector();
    BaseVector & dh1 = *h1mat->CreateVector();

    y = 0;
    jacobi->GSSmooth (y, x);

    // do potential smoothing:

    hv = x - (*pmat) * y;
    dh1 = Transpose (*grad) * hv;
    wh1 = (*h1AMG) * dh1;
    y += (*grad) * wh1;

    if (recAMG)
      {
	hv = x - (*pmat) * y;
	dc = Transpose (*prol) * hv;

	if (recAMG)
	  recAMG -> Mult (dc, wc);

	y += (*prol) * wc;
      }

    hv = x - (*pmat) * y;
    dh1 = Transpose (*grad) * hv;
    wh1 = (*h1AMG) * dh1;
    y += (*grad) * wh1;

    jacobi->GSSmoothBack (y, x);

    delete &hv;
    delete &wc;
    delete &dc;
    delete &wh1;
    delete &dh1;
  }


};