#include "GmshConfig.h"
#include "GModel.h"
#include "GVertex.h"
#include "GEdge.h"
#include "frameSolver.h"
#include "linearSystemCSR.h"
#include "linearSystemPETSc.h"
#include "linearSystemGMM.h"
#include "linearSystemFull.h"
#if defined(HAVE_POST)
#include "PView.h"
#include "PViewData.h"
#endif

frameSolver2d::frameSolver2d (GModel *gm)
  : pAssembler(0),_myModel(gm)
{
}


void frameSolver2d::addFixations (const std::vector<int> & dirs, const std::vector<int> &modelVertices, double value)
{
  for (unsigned int j=0;j<modelVertices.size();j++){
    GVertex *gv = _myModel->getVertexByTag(modelVertices[j]);
    if (gv){
      for (unsigned int i=0;i<dirs.size();i++){
	_fixations.push_back(gmshFixation(gv,dirs[i],value));
      }
    }
  }
}

void frameSolver2d::addNodalForces (const std::vector<int> &modelVertices, const std::vector<double> & force)
{
  for (unsigned int j=0;j<modelVertices.size();j++){
    GVertex *gv = _myModel->getVertexByTag(modelVertices[j]);
    if (gv){
      _nodalForces.push_back(std::make_pair(gv,force));
    }
  }
}

void frameSolver2d::addBeamsOrBars (const std::vector<int> &modelEdges,
				    double E, double I, double A, int r[2]){
  int r_middle[2] = {1,1} , r_left[2] = {r[0],1} , r_right[2] = {0,r[1]};
  //  printf("adding %d beams\n",modelEdges.size());
  for (unsigned int i=0;i<modelEdges.size();i++){
    GEdge *ge = _myModel->getEdgeByTag(modelEdges[i]);
    if (ge){
      //      printf("model edge %d found\n",ge->tag());
      for (unsigned int j=0; j < ge->lines.size(); ++j){
	MLine *l = ge->lines[j];
	if (j == 0 && j == ge->lines.size()-1) _beams.push_back(gmshBeam2d ( l , E , I , A , r));
	else if (j == 0) _beams.push_back(gmshBeam2d ( l , E , I , A , r_left));
	else if (j == ge->lines.size()-1) _beams.push_back(gmshBeam2d ( l , E , I , A , r_right));
	else _beams.push_back(gmshBeam2d ( l , E , I , A , r_middle));
      }
    }
  }
}

void frameSolver2d::addBeams (const std::vector<int> &modelEdges,
			      double E, double I, double A)
{
  int r[2] = {1,1};
  addBeamsOrBars(modelEdges,E,I,A,r);
}

void frameSolver2d::addBars (const std::vector<int> &modelEdges,
			     double E, double I, double A)
{
  int r[2] = {0,0};
  addBeamsOrBars(modelEdges,E,I,A,r);
}



// solve any time a parameter is modified
/*

  A vertex is connected to beams. The question
  is how many bars are rotuled

  We define 2 dofs per node

 */
void frameSolver2d::createDofs ()
{
  //  printf("coucou %d fixations\n",_fixations.size());
  for (unsigned int i = 0; i<_fixations.size(); ++i){
    const  gmshFixation &f = _fixations[i];
    //    printf("f._vertex(%d) = %p %d %g\n",i,f._vertex,f._direction,f._value);
    MVertex *v = f._vertex->mesh_vertices[0];
    Dof DOF (v->getNum(),f._direction);
    pAssembler->fixDof(DOF, f._value);
  }

  //  printf("coucou2\n");
  computeRotationTags ();
  //  printf("coucou3\n");
  for (unsigned int i=0;i<_beams.size(); i++){
    //    printf("beam[%d] rot %d %d\n",i,_beams[i]._rotationTags[0],_beams[i]._rotationTags[1]);
    for (unsigned int j=0;j<2;j++){
      MVertex *v = _beams[i]._element->getVertex(j);
      Dof theta (v->getNum(),
		 Dof::createTypeWithTwoInts (2,_beams[i]._rotationTags[j]));
      pAssembler->numberDof(theta);
      Dof U (v->getNum(),0);
      pAssembler->numberDof(U);
      Dof V (v->getNum(),1);
      pAssembler->numberDof(V);
    }
  }
  //  printf("%d dofs\n",pAssembler->sizeOfR());
}

void frameSolver2d::computeStiffnessMatrix (int iBeam,
					    fullMatrix<double> &K)
{
  const gmshBeam2d &b = _beams[iBeam];
  const double BS = b._E * b._I / (b._L*b._L*b._L);
  const double TS = b._E * b._A / b._L;
  const MVertex *v1 = b._element->getVertex(0);
  const MVertex *v2 = b._element->getVertex(1);
  const double alpha = atan2 (v2->y() - v1->y(),
			      v2->x() - v1->x());
  const double C = cos(alpha);
  const double S = sin(alpha);

  printf("beam %d %g %g %g\n",iBeam,alpha,C,S);

  fullMatrix<double> R (6,6);
  R(0,0) = R(1,1) = R (3,3) = R(4,4) = C;
  R(0,1) = R (3,4) = S;
  R(1,0) = R (4,3) = -S;
  R(2,2) = R(5,5) = 1.0;

  fullMatrix<double> k (6,6);

  // tensile stiffness
  k(0,0) = k (3,3) =  TS;
  k(0,3) = k (3,0) = -TS;
  // bending stiffness
  k(1,1) = k(4,4) = 12 * BS;
  k(2,2) = k(5,5) = 4. * BS * b._L * b._L;
  k(1,2) = k(2,1) = k(1,5) = k(5,1) =  6 * BS * b._L;
  k(4,2) = k(2,4) = k(4,5) = k(5,4) = -6 * BS * b._L;
  k(4,1) = k(1,4) = -12 * BS ;
  k(5,2) = k(2,5) = -2 * BS * b._L * b._L;

  fullMatrix<double> Rt (R), temp(6,6);
  Rt.transposeInPlace();
  Rt.mult(k,temp);temp.mult(R,K);
}

void frameSolver2d::solve ()
{
#if defined(HAVE_TAUCS)
  linearSystemCSRTaucs<double> *lsys = new linearSystemCSRTaucs<double>;
#elif defined(HAVE_PETSC)
  linearSystemPETSc<double> *lsys = new linearSystemPETSc<double>;
#else
  linearSystemGmm<double> *lsys = new linearSystemGmm<double>;
  lsys->setNoisy(2);
#endif

  if (pAssembler)delete pAssembler;
  pAssembler = new dofManager<double>(lsys);

  // fix dofs and create free ones
  createDofs();

  // force vector
  std::vector<std::pair<GVertex*,std::vector<double> > >::iterator it =
    _nodalForces.begin();
  for (; it != _nodalForces.end(); ++it){
    MVertex *v = it->first->mesh_vertices[0];
    const std::vector<double> & F = it->second;
    Dof DOFX (v->getNum(),0);
    Dof DOFY (v->getNum(),1);
    pAssembler->assemble(DOFX, F[0]);
    pAssembler->assemble(DOFY, F[1]);
  }

  // stifness matrix
  for (unsigned int i=0;i<_beams.size(); i++){
    fullMatrix<double> K(6,6);
    computeStiffnessMatrix (i,K);
    _beams[i]._stiffness = K;
    MVertex *v0 = _beams[i]._element->getVertex(0);
    MVertex *v1 = _beams[i]._element->getVertex(1);
    Dof theta0 (v0->getNum(),
		Dof::createTypeWithTwoInts (2,_beams[i]._rotationTags[0]));
    Dof theta1 (v1->getNum(),
		Dof::createTypeWithTwoInts (2,_beams[i]._rotationTags[1]));
    Dof U0 (v0->getNum(),0);
    Dof U1 (v1->getNum(),0);
    Dof V0 (v0->getNum(),1);
    Dof V1 (v1->getNum(),1);
    Dof DOFS[6]  = {U0,V0,theta0,U1,V1,theta1};
    for (int j=0;j<6;j++){
      for (int k=0;k<6;k++){
	pAssembler->assemble(DOFS[j],DOFS[k],K(j,k));
      }
    }
  }
  lsys->systemSolve();

  // save the solution
  for (unsigned int i=0;i<_beams.size(); i++){
    MVertex *v0 = _beams[i]._element->getVertex(0);
    MVertex *v1 = _beams[i]._element->getVertex(1);
    Dof theta0 (v0->getNum(),
		Dof::createTypeWithTwoInts (2,_beams[i]._rotationTags[0]));
    Dof theta1 (v1->getNum(),
		Dof::createTypeWithTwoInts (2,_beams[i]._rotationTags[1]));
    Dof U0 (v0->getNum(),0);
    Dof U1 (v1->getNum(),0);
    Dof V0 (v0->getNum(),1);
    Dof V1 (v1->getNum(),1);
    Dof DOFS[6]  = {U0,V0,theta0,U1,V1,theta1};
    for (int j=0;j<6;j++){
      pAssembler->getDofValue(DOFS[j],  _beams[i]._displacement[j]);
    }
  }
  delete lsys;
  delete pAssembler;
}

void frameSolver2d::exportFrameData(const char * DISPL, const char * M) {
#if defined(HAVE_POST)
  {
    std::map<int, std::vector<double> > data;
    for (unsigned int i=0;i<_beams.size(); i++){
      std::vector<double> tmp;
      //      tmp.push_back(_beams[i]._E);
      //      tmp.push_back(_beams[i]._I);
      //      tmp.push_back(_beams[i]._A);
      for (int j=0;j<6;j++){
	tmp.push_back(_beams[i]._displacement[j]);
      }
      data[_beams[i]._element->getNum()] = tmp;
    }
    PView *pv = new PView ("displacements", "Beam", _myModel, data, 0.0, 6);
    pv->getData()->writeMSH(DISPL);
    delete pv;
  }
  {
    std::map<int, std::vector<double> > data;
    for (unsigned int i=0;i<_beams.size(); i++){
      std::vector<double> tmp;
      fullVector<double> d (_beams[i]._displacement,6), F(6);
      _beams[i]._stiffness.mult(d,F);
      tmp.push_back(-F(2));
      tmp.push_back(F(5));
      data[_beams[i]._element->getNum()] = tmp;
    }
    PView *pv = new PView ("Momentum", "ElementNodeData", _myModel, data, 0.0, 1);
    pv->getData()->writeMSH(M);
    delete pv;
  }
#endif
}

void frameSolver2d::computeRotationTags ()
{
  std::multimap<MVertex*,gmshBeam2d*> v2b;
  for (unsigned int i=0;i<_beams.size(); i++){
    v2b.insert(std::make_pair(_beams[i]._element->getVertex(0),&_beams[i]));
    v2b.insert(std::make_pair(_beams[i]._element->getVertex(1),&_beams[i]));
  }

  std::multimap<MVertex*,gmshBeam2d*>::iterator s_it;
  for (std::multimap<MVertex*,gmshBeam2d*>::iterator it = v2b.begin();
       it != v2b.end();  it = s_it){
    MVertex *theKey = it->first;

    std::pair<std::multimap<MVertex*,gmshBeam2d*>::iterator,
	      std::multimap<MVertex*,gmshBeam2d*>::iterator>
      keyRange = v2b.equal_range(theKey);
    int countRotules = 0;
    for (s_it = keyRange.first;  s_it != keyRange.second;  ++s_it){
      gmshBeam2d *b = s_it->second;
      if (!b->isRigid(theKey)){
	b->setRotationTag(theKey,++countRotules);
      }
    }
  }
}
