/* Fesslix - Stochastic Analysis
 * Copyright (C) 2010-2016 Wolfgang Betz
 *
 * Fesslix 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.
 *
 * Fesslix 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 Fesslix.  If not, see <http://www.gnu.org/licenses/>. 
 */

#include "flxfem.h"

#define fesslix_flxfem_CPP


#ifdef __WINDOWS__
  #pragma comment(lib, "flxlib")
#endif



void flxFEM::assembleDOF()
{
  if (StiffnessMtx) {
    delete StiffnessMtx;
  }
  StiffnessMtx = new flxStiffnessMtxFEM(this);
}

void flxFEM::plot_ElProp(std::ostream& sout, const std::vector<FlxElementInfo::elInfo>& V, FlxMtxConstFun* ppdim, FlxSMtx* elmtx, const int prec, const int fixW)
{
  ElementList.plot(sout, V, ppdim, elmtx, prec, fixW);
}

flxFEM::flxFEM(FlxConstantBox* constBox) 
: ConstantBox(constBox),StiffnessMtx(NULL) 
{ 
  flxFace::set_FEM(this);
  LoadCaseBox.insert(0,new FlxLoadCase(0,"",false,ONE),true);
}

flxFEM::~flxFEM()
{
  if (StiffnessMtx!=NULL) {
    delete StiffnessMtx;
  }
}

flxStiffnessMtxBase::flxStiffnessMtxBase(flxFEM* FEM)
:FEM(FEM),reassembly(false),StfMtx(NULL),StfMtx_2x(new FlxMtxSpars*[4]),
ForceVec(NULL),ForceVec_fixed(NULL),ForceVec_internal(NULL),ForceVec_internal_fixed(NULL),ForceVec_2(NULL),
SolutionVec(NULL),preDisplVec(NULL),Mtx1(NULL),Mtx2x(new FlxMtx*[4])
{
  try {
    // first things frist
      for (int i = 0; i<4;++i) StfMtx_2x[i] = NULL;
      for (int i = 0; i<4;++i) Mtx2x[i] = NULL;
    // refine the mesh
      FEM->EdgeList.refineMesh();
      FEM->FaceList.refineMesh();
      FEM->BoxList.refineMesh();
    // count DOFs
      DofA = 0;
      FixA = 0;
      preDispA = 0;
      DofSum = 0;
      FEM->BoxList.countDOF(DofA,FixA,preDispA,DofSum);
      FEM->FaceList.countDOF(DofA,FixA,preDispA,DofSum);
      FEM->EdgeList.countDOF(DofA,FixA,preDispA,DofSum);
      FEM->VertexList.countDOF(DofA,FixA,preDispA,DofSum); 
    // assemble DOFs
      tnlong DofIndex = 0;
      tnlong FixPreIndex = DofA;
      tnlong FixZeroIndex = DofA+preDispA;
      FEM->VertexList.assembleDOF(DofIndex,FixPreIndex,FixZeroIndex);
      FEM->EdgeList.assembleDOF(DofIndex,FixPreIndex,FixZeroIndex);
      FEM->FaceList.assembleDOF(DofIndex,FixPreIndex,FixZeroIndex);
      FEM->BoxList.assembleDOF(DofIndex,FixPreIndex,FixZeroIndex);
    // check for problems
      if (DofIndex!=DofA || FixPreIndex != DofA+preDispA || FixZeroIndex != DofSum) {
	std::ostringstream ssV;
	  ssV << "ERROR: DofA(" << DofA << "/" << DofIndex << ") FixPre(" << DofA+preDispA << "/" << FixPreIndex << ")";
	  ssV << " FixZero(" << DofSum << "/" << FixZeroIndex << ")";
	  throw FlxException("flxNodeList::assembleDOF_1", ssV.str() );
      }
    // assemble dynamic stuff
      if (preDispA>0) {
	preDisplVec = new flxVec(preDispA);
      }
  } catch (FlxException& e) {
    FLXMSG("flxNodeList::assembleDOF_2",1);
    if (StfMtx) delete StfMtx;
    for (int i=0;i<4;++i) {
      if (StfMtx_2x[i]!=NULL) delete StfMtx_2x[i];
    }
    delete [] StfMtx_2x;
    if (ForceVec) delete ForceVec;
    if (ForceVec_fixed) delete ForceVec_fixed;
    if (ForceVec_internal) delete ForceVec_internal;
    if (ForceVec_internal_fixed) delete ForceVec_internal_fixed;
    if (ForceVec_2) delete ForceVec_2;
    if (SolutionVec) delete SolutionVec;
    if (preDisplVec) delete preDisplVec;
    if (Mtx1) delete Mtx1;
    for (int i=0;i<4;++i) {
      if (Mtx2x[i]) delete Mtx2x[i];
    }
    delete [] Mtx2x;
    throw;
  }
}

flxStiffnessMtxBase::~flxStiffnessMtxBase()
{
  // NOTE -> same code ABOVE !!! 
  if (StfMtx) delete StfMtx;
  for (int i=0;i<4;++i) {
    if (StfMtx_2x[i]!=NULL) delete StfMtx_2x[i];
  }
  delete [] StfMtx_2x;
  if (ForceVec) delete ForceVec;
  if (ForceVec_fixed) delete ForceVec_fixed;
  if (ForceVec_internal) delete ForceVec_internal;
  if (ForceVec_internal_fixed) delete ForceVec_internal_fixed;
  if (ForceVec_2) delete ForceVec_2;
  if (SolutionVec) delete SolutionVec;
  if (preDisplVec) delete preDisplVec;
  if (Mtx1) delete Mtx1;
  for (int i=0;i<4;++i) {
    if (Mtx2x[i]) delete Mtx2x[i];
  }
  delete [] Mtx2x;
}

void flxStiffnessMtxBase::assemble_forc(const bool errSerious)
{
  // check if stiffness matrix needs to be assembled
    if (StfMtx_2x[0]==NULL && preDispA>0) {
      FlxError(errSerious,"flxStiffnessMtxBase::assemble_forc_1","Please execute 'fem assstf' beforehand.");
    }
  // allocate memory, assign with zero
    if (ForceVec == NULL) {
      const tnlong dummy = (DofA==0)?1:DofA;
      ForceVec = new flxVec(dummy);
      ForceVec_internal = new flxVec(dummy);
      if (StfMtx_2x[0]!=NULL) {
	ForceVec_2 = new flxVec(dummy);
      }
      if (FixA>0) {
	ForceVec_fixed = new flxVec(FixA);
	ForceVec_internal_fixed = new flxVec(FixA);
      }
    } else {
      #if FLX_DEBUG
	if (ForceVec->get_N() != DofA) {
	  throw FlxException_Crude("flxStiffnessMtxBase::assemble_forc_2");
	}
      #endif
      ForceVec->set_zero();
      ForceVec_internal->set_zero();
      if (FixA>0) {
	ForceVec_fixed->set_zero();
	ForceVec_internal_fixed->set_zero();
      }
      if (StfMtx_2x[0]) { ForceVec_2->set_zero(); }
    }
  // assemble element loads (internal forces)
    FEM->EdgeList.assembleForces(*this);
    FEM->FaceList.assembleForces(*this);
    FEM->BoxList.assembleForces(*this);
    // store internal forces separately
      (*ForceVec_internal) = (*ForceVec);
      if (ForceVec_fixed) {
	(*ForceVec_internal_fixed) = (*ForceVec_fixed);
      }
  // assemble nodal forces
    FEM->VertexList.assembleForces(*this);
  // deal with predisplacements
    if (StfMtx_2x[0]!=NULL) {
      StfMtx_2x[0]->TransMultVec(*preDisplVec,*ForceVec_2);
      *ForceVec_2 -= *ForceVec;
      *ForceVec_2 *= -ONE;
    }
    ForceVec->check_TOL();
    if (ForceVec_internal) ForceVec_internal->check_TOL();
    if (ForceVec_fixed) ForceVec_fixed->check_TOL();
    if (ForceVec_internal_fixed) ForceVec_internal_fixed->check_TOL();
    if (ForceVec_2) ForceVec_2->check_TOL();
}

const double flxStiffnessMtxBase::solve(tdouble& eps, tnlong& iter, tuint pcn, bool startZero, const bool errSerious, const bool verboseOutput)
{
  if (StfMtx==NULL) {
    std::ostringstream ssV;
    ssV << "Please execute 'fem assstf' beforehand.";
    FlxError(errSerious, "flxStiffnessMtxBase::solve_1", "The stiffness matrix has not yet been assembled.", ssV.str() );
  }
  if (ForceVec==NULL) {
    std::ostringstream ssV;
    ssV << "Please execute 'fem assforces' beforehand.";
    FlxError(errSerious, "flxStiffnessMtxBase::solve_2", "The force vector has not yet been assembled.", ssV.str() );
  }
  allocate_SolutionVec();
  if (DofA==0) {
    GlobalVar.slog(3) << "FEM: nothing to be solved explicitely." << std::endl;
    eps=ZERO; iter=0;
    update_displVec_GeoEls();
    return ZERO;
  }
  flxVec* f = ForceVec;
  if (ForceVec_2!=NULL) {
    f = ForceVec_2;
  }
  if (StfMtx->solve_CG(*SolutionVec,*f,eps,iter,pcn,startZero) == false) {
    std::ostringstream ssV;
    ssV << "System could not be solved. (eps=" << eps << "; iter=" << iter << ")" << std::endl << "Last Solution Vector:" << std::endl;
    output_solVec(ssV,verboseOutput);
//     ssV << std::endl << "-------------" << *StfMtx << std::endl << *ForceVec << std::endl;
    FlxError("flxStiffnessMtxBase::solve_3", ssV.str() );
  }
  update_displVec_GeoEls();
  return eps/ForceVec->get_Norm2();
}

void flxStiffnessMtxBase::update_displVec_GeoEls()
{
  FEM->EdgeList.updateDisplVec(*this);
  FEM->FaceList.updateDisplVec(*this);
  FEM->BoxList.updateDisplVec(*this);
}

void flxStiffnessMtxBase::output_solVec(std::ostream& sout, bool verbose)
{
  if (SolutionVec==NULL) return;
  sout << "Solution of the system:" << std::endl;
  FEM->VertexList.print(sout,true,FEM);
  if (verbose) {
  sout << "Solution vector:" << std::endl;
  sout << *SolutionVec << std::endl;
  }
}

void flxStiffnessMtxBase::output_StfMtx(std::ostream& sout)
{
  if (StfMtx==NULL) return;
  StfMtx->output_Mtx(sout);
}

void flxStiffnessMtxBase::output_ForceVec(std::ostream& sout)
{
  if (ForceVec==NULL) return;
  sout << *ForceVec << std::endl;
}

const bool flxStiffnessMtxBase::isPositiveDefinite(tuint& r) const
{
  return StfMtx->isPositiveDefinite(r,true);
}

void flxStiffnessMtxBase::build_stf()
{
  reassembly=true;
  #if FLX_DEBUG
    if (Mtx1==NULL) {
      std::ostringstream ssV;
      ssV << "ERROR";
      throw FlxException("flxStiffnessMtxFEM::build_stf_1", ssV.str() );
    };
  #endif
  // deal with stiffness matrix
    if (StfMtx!=NULL) delete StfMtx;
    StfMtx = new FlxMtxSparsSym(*Mtx1);
    delete Mtx1;
    Mtx1=NULL;
  // deal with matrix_2x
    for (int i=0;i<4;++i) {
      if (StfMtx_2x[i]!=NULL) delete StfMtx_2x[i];
      if (Mtx2x[i]!=NULL) {
	StfMtx_2x[i] = new FlxMtxSpars(*Mtx2x[i]);
	delete Mtx2x[i];
	Mtx2x[i]=NULL;
      } else {
	StfMtx_2x[i]=NULL;
      }
    }
}

inline void flxStiffnessMtxBase::prepare_stf_for_reassembly()
{
  if (reassembly) {
    StfMtx->set_zeroMtx();
  } else {
    std::ostringstream ssV;
    ssV << "ERROR";
    throw FlxException("flxStiffnessMtxBase::prepare_stf_for_reassembly", ssV.str() );
  }
}

void flxStiffnessMtxBase::assemble_stf()
{
  reassembly=false;
  Mtx1 = new FlxMtxSym(DofA); 
  if (preDispA>0) {
    Mtx2x[0] = new FlxMtx(preDispA,DofA);
    if (FixA-preDispA>0) {
      Mtx2x[2] = new FlxMtx(preDispA,preDispA);
    }
  }
  if (FixA-preDispA>0) {
    Mtx2x[1] = new FlxMtx(FixA-preDispA,DofA);
    if (preDispA>0) {
      Mtx2x[3] = new FlxMtx(FixA-preDispA,preDispA);
    }
  }
  FEM->VertexList.assembleStf(*this);	// nodal springs and pre-displacements
  FEM->EdgeList.assembleStf(*this);
  FEM->FaceList.assembleStf(*this);
  FEM->BoxList.assembleStf(*this);
  build_stf();
}

void flxStiffnessMtxBase::Stf_set_value(tnlong i, tnlong j, const tdouble& v)
{
  if (j>i) std::swap(i,j);
  if (!reassembly) {
    switch (StfCoeff_is_stored(i,j)) {
      case 1:
	(*Mtx1)(i,j)+=v;
	break;
      case 2:
	if (Mtx2x[0]==NULL) return;
	(*Mtx2x[0])(i-DofA,j)+=v;
	break;
      case 3:
	if (Mtx2x[1]==NULL) return;
	(*Mtx2x[1])(i-DofA-preDispA,j)+=v;
	break;
      case 4:
	if (Mtx2x[2]==NULL) return;
	(*Mtx2x[2])(i-DofA,j-DofA)+=v;
	break;
      case 5:
	if (Mtx2x[3]==NULL) return;
	(*Mtx2x[3])(i-DofA-preDispA,j-DofA)+=v;
	break;
    }
  } else {
    switch (StfCoeff_is_stored(i,j)) {
      case 1:
	StfMtx->add_value(i,j,v);
	break;
      case 2:
	StfMtx_2x[0]->add_value(i-DofA,j,v);
	break;
      case 3:
	StfMtx_2x[1]->add_value(i-DofA-preDispA,j,v);
	break;
      case 4:
	StfMtx_2x[2]->add_value(i-DofA,j-DofA,v);
	break;
      case 5:
	StfMtx_2x[3]->add_value(i-DofA-preDispA,j-DofA,v);
	break;
    }
  };
}

void flxStiffnessMtxBase::Stf_set_value(tnlong i, tnlong j, const tdouble* v, const tuint N)
{
  if (N==0) return;
  if (j>i) {
    std::swap(i,j);
    if (!reassembly) {
      switch (StfCoeff_is_stored(i,j)) {
	case 1:
	  for (tuint k=0;k<N;++k) {
	    (*Mtx1)(i+k,j)+=v[k];
	  }
	  break;
	case 2:
	  if (Mtx2x[0]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[0])(i-DofA+k,j)+=v[k];
	  }
	  break;
	case 3:
	  if (Mtx2x[1]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[1])(i-DofA-preDispA+k,j)+=v[k];
	  }
	  break;
	case 4:
	  if (Mtx2x[2]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[2])(i-DofA+k,j-DofA)+=v[k];
	  }
	  break;
	case 5:
	  if (Mtx2x[3]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[3])(i-DofA-preDispA+k,j-DofA)+=v[k];
	  }
	  break;
      }
    } else {
      throw FlxException_NotImplemented("flxStiffnessMtxBase::Stf_set_value_1");
    }
  } else {
    #if FLX_DEBUG
      if (j+N-1>i) throw FlxException_Crude("flxStiffnessMtxBase::Stf_set_value_1");
    #endif
    if (!reassembly) {
      switch (StfCoeff_is_stored(i,j)) {
	case 1:
	  for (tuint k=0;k<N;++k) {
	    (*Mtx1)(i,j+k)+=v[k];
	  }
	  break;
	case 2:
	  if (Mtx2x[0]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[0])(i-DofA,j+k)+=v[k];
	  }
	  break;
	case 3:
	  if (Mtx2x[1]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[1])(i-DofA-preDispA,j+k)+=v[k];
	  }
	  break;
	case 4:
	  if (Mtx2x[2]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[2])(i-DofA,j-DofA+k)+=v[k];
	  }
	  break;
	case 5:
	  if (Mtx2x[3]==NULL) return;
	  for (tuint k=0;k<N;++k) {
	    (*Mtx2x[3])(i-DofA-preDispA,j-DofA+k)+=v[k];
	  }
	  break;
      }
    } else {
      throw FlxException_NotImplemented("flxStiffnessMtxBase::Stf_set_value_2");
    }
  }
}

void flxStiffnessMtxBase::Stf_set_predispl_value(tnlong i, const tdouble& v)
{
  if (v==ZERO) return;
  #if FLX_DEBUG
    if (preDisplVec==NULL) {
      throw FlxException_Crude("flxStiffnessMtxBase::Stf_set_predispl_value_1");
    }
    
  #endif
  #if FLX_DEBUG
    if (i<DofA) {
      throw FlxException_Crude("flxStiffnessMtxBase::Stf_set_predispl_value_2");
    }
  #endif
  i -= DofA;
  #if FLX_DEBUG
    if (i>preDispA) {
      throw FlxException_Crude("flxStiffnessMtxBase::Stf_set_predispl_value_3");
    }
  #endif
  preDisplVec->operator[](i) =  v;
}

void flxStiffnessMtxBase::ForceVec_set_value(const tnlong i, const tdouble& v)
{
  if (i<DofA) {
    ForceVec->operator[](i) += v;
  } else {
    ForceVec_fixed->operator[](i-DofA) += v;
  }
}

void flxStiffnessMtxBase::ForceVec_set_value(const tnlong i, const tdouble* v, const tuint N)
{
  if (N==0) return;
  tdouble *fv;
  if (i<DofA) {
    fv = &(ForceVec->operator[](i));
  } else {
    fv = &(ForceVec_fixed->operator[](i-DofA));
  }
  for (tuint k=0;k<N;++k) {
    fv[k] += v[k];
  }
}

const int flxStiffnessMtxBase::StfCoeff_is_stored(const tnlong i, const tnlong j) const
{
  if (i<DofA && j<DofA) return 1;				// symmetric matrix
  else if (i<DofA+preDispA && j<DofA) return 2;			// nxm matrix
  else if (i<DofSum && j<DofA) return 3;			// nxm matrix
  else if (i<DofA+preDispA && j<DofA+preDispA) return 4;	// symmetric matrix
  else if (i<DofSum && j<DofA+preDispA) return 5;		// nxm matrix
  else return 0;
}

void flxStiffnessMtxFEM::allocate_SolutionVec()
{
  if (SolutionVec==NULL) {
    const tnlong dummy = (DofA==0)?1:DofA;
    SolutionVec=new flxVec(dummy);
  }
}

const tdouble flxStiffnessMtxFEM::get_U(const tnlong dofindex) const
{
  if (SolutionVec==NULL) {
    std::ostringstream ssV;
    ssV << "Solution vector is empty.";
    throw FlxException("flxStiffnessMtxFEM::get_U", ssV.str() );
  }
  if (dofindex < DofA) {
    return (*SolutionVec)[dofindex];
  } else if (dofindex < DofA+preDispA) {
    return preDisplVec->operator[](dofindex-DofA);
  } else {
    return ZERO;
  }
}

void flxStiffnessMtxFEM::get_U(const tnlong dofindex, tdouble* v, const tuint N) const
{
  if (SolutionVec==NULL) {
    std::ostringstream ssV;
    ssV << "Solution vector is empty.";
    throw FlxException("flxStiffnessMtxFEM::get_U", ssV.str() );
  }
  const tdouble* sv;
  if (dofindex < DofA) {
    sv = &((*SolutionVec)[dofindex]);
  } else if (dofindex < DofA+preDispA) {
    sv = &(preDisplVec->operator[](dofindex-DofA));
  } else {
    for (tuint i=0;i<N;++i) {
      v[i] = ZERO;
    }
    return;
  }
  for (tuint i=0;i<N;++i) {
    v[i] = sv[i];
  }
}

const tdouble flxStiffnessMtxFEM::get_F(const flxDOFn* Dof)
{
  // some preliminary checks
    if (ForceVec==NULL) {
      std::ostringstream ssV;
      ssV << "Force vector is empty.";
      throw FlxException("flxStiffnessMtxFEM::get_F_1", ssV.str() );
    }
    if (SolutionVec==NULL) {
      std::ostringstream ssV;
      ssV << "Solution vector is empty.";
      throw FlxException("flxStiffnessMtxFEM::get_F_2", ssV.str() );
    }
  const tnlong DOFINDEX = Dof->get_DOFindex();
  if (DOFINDEX < DofA) {
    if (Dof->get_state() == FlxDof::dof_spring) {
      return -ONE*(*SolutionVec)[DOFINDEX]*(Dof->get_state_value());
    }
    return (*ForceVec)[DOFINDEX] - ForceVec_internal->operator[](DOFINDEX);
  } else {
    if (DOFINDEX < DofA+preDispA) {
      return StfMtx_2x[0]->MultRowV(*SolutionVec,DOFINDEX-DofA) + StfMtx_2x[2]->MultRowV(*preDisplVec,DOFINDEX-DofA) - ForceVec_fixed->operator[](DOFINDEX-DofA);
    } else {
      tdouble f = StfMtx_2x[1]->MultRowV(*SolutionVec,DOFINDEX-DofA-preDispA) - ForceVec_fixed->operator[](DOFINDEX-DofA);
      if (preDispA>0) {
	f += StfMtx_2x[3]->MultRowV(*preDisplVec,DOFINDEX-DofA-preDispA);
      }
      return  f;
    }
  }
}




