/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2010 University of Bonn. All rights reserved.
 * Please see the LICENSE file or "Copyright notice" in builder.cpp for details.
 */

/*
 * SubspaceFactorizer.cpp
 *
 *  Created on: Nov 23, 2010
 *      Author: heber
 */

// include config.h
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <cmath>

#include <gsl/gsl_eigen.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <boost/foreach.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/timer.hpp>

#include "CodePatterns/Assert.hpp"
#include "Helpers/defs.hpp"
#include "CodePatterns/Log.hpp"
#include "CodePatterns/toString.hpp"
#include "CodePatterns/Verbose.hpp"
#include "LinearAlgebra/Eigenspace.hpp"
#include "LinearAlgebra/MatrixContent.hpp"
#include "LinearAlgebra/Subspace.hpp"
#include "LinearAlgebra/VectorContent.hpp"

typedef std::set<std::set<size_t> > SetofIndexSets;
typedef std::set<size_t> IndexSet;
typedef std::multimap< size_t, boost::shared_ptr<Subspace> > SubspaceMap;
typedef std::multimap< size_t, boost::shared_ptr<IndexSet> > IndexMap;
typedef std::list< boost::shared_ptr<VectorContent> > VectorList;
typedef std::list< std::pair<boost::shared_ptr<VectorContent>, double> > VectorValueList;
typedef std::vector< boost::shared_ptr<VectorContent> > VectorArray;
typedef std::vector< double > ValueArray;


/** Iterative function to generate all power sets of indices of size \a maxelements.
 *
 * @param SetofSets Container for all sets
 * @param CurrentSet pointer to current set in this container
 * @param Indices Source set of indices
 * @param maxelements number of elements of each set in final SetofSets
 * @return true - generation continued, false - current set already had
 *         \a maxelements elements
 */
bool generatePowerSet(
    SetofIndexSets &SetofSets,
    SetofIndexSets::iterator &CurrentSet,
    IndexSet &Indices,
    const size_t maxelements)
{
  if (CurrentSet->size() < maxelements) {
    // allocate the needed sets
    const size_t size = Indices.size() - CurrentSet->size();
    std::vector<std::set<size_t> > SetExpanded;
    SetExpanded.reserve(size);

    // copy the current set into each
    for (size_t i=0;i<size;++i)
      SetExpanded.push_back(*CurrentSet);

    // expand each set by one index
    size_t localindex=0;
    BOOST_FOREACH(size_t iter, Indices) {
      if (CurrentSet->count(iter) == 0) {
        SetExpanded[localindex].insert(iter);
        ++localindex;
      }
    }

    // insert set at position of CurrentSet
    for (size_t i=0;i<size;++i) {
      //DoLog(1) && (Log() << Verbose(1) << "Inserting set #" << i << ": " << SetExpanded[i] << std::endl);
      SetofSets.insert(CurrentSet, SetExpanded[i]);
    }
    SetExpanded.clear();

    // and remove the current set
    //SetofSets.erase(CurrentSet);
    //CurrentSet = SetofSets.begin();

    // set iterator to a valid position again
    ++CurrentSet;
    return true;
  } else {
    return false;
  }
}



/** Prints the scalar product of each possible pair that is not orthonormal.
 * We use class logger for printing.
 * @param AllIndices set of all possible indices of the eigenvectors
 * @param CurrentEigenvectors array of eigenvectors
 * @return true - all are orthonormal to each other,
 *        false - some are not orthogonal or not normalized.
 */
bool checkOrthogonality(const IndexSet &AllIndices, const VectorArray &CurrentEigenvectors)
{
  size_t nonnormalized = 0;
  size_t nonorthogonal = 0;
  // check orthogonality
  BOOST_FOREACH( size_t firstindex, AllIndices) {
    BOOST_FOREACH( size_t secondindex, AllIndices) {
      const double scp = (*CurrentEigenvectors[firstindex])*(*CurrentEigenvectors[secondindex]);
      if (firstindex == secondindex) {
        if (fabs(scp - 1.) > MYEPSILON) {
          nonnormalized++;
          Log() << Verbose(2) << "Vector " << firstindex << " is not normalized, off by "
              << fabs(1.-(*CurrentEigenvectors[firstindex])*(*CurrentEigenvectors[secondindex])) << std::endl;
        }
      } else {
        if (fabs(scp) > MYEPSILON) {
          nonorthogonal++;
          Log() << Verbose(2) << "Scalar product between " << firstindex << " and " << secondindex
              << " is " << (*CurrentEigenvectors[firstindex])*(*CurrentEigenvectors[secondindex]) << std::endl;
        }
      }
    }
  }

  if ((nonnormalized == 0) && (nonorthogonal == 0)) {
    DoLog(1) && (DoLog(1) && (Log() << Verbose(1) << "All vectors are orthonormal to each other." << std::endl));
    return true;
  }
  if ((nonnormalized == 0) && (nonorthogonal != 0))
    DoLog(1) && (DoLog(1) && (Log() << Verbose(1) << "All vectors are normalized." << std::endl));
  if ((nonnormalized != 0) && (nonorthogonal == 0))
    DoLog(1) && (DoLog(1) && (Log() << Verbose(1) << "All vectors are orthogonal to each other." << std::endl));
  return false;
}

/** Calculate the sum of the scalar product of each possible pair.
 * @param AllIndices set of all possible indices of the eigenvectors
 * @param CurrentEigenvectors array of eigenvectors
 * @return sum of scalar products between all possible pairs
 */
double calculateOrthogonalityThreshold(const IndexSet &AllIndices, const VectorArray &CurrentEigenvectors)
{
  double threshold = 0.;
  // check orthogonality
  BOOST_FOREACH( size_t firstindex, AllIndices) {
    BOOST_FOREACH( size_t secondindex, AllIndices) {
      const double scp = (*CurrentEigenvectors[firstindex])*(*CurrentEigenvectors[secondindex]);
      if (firstindex == secondindex) {
        threshold += fabs(scp - 1.);
      } else {
        threshold += fabs(scp);
      }
    }
  }
  return threshold;
}


/** Operator for output to std::ostream operator of an IndexSet.
 * @param ost output stream
 * @param indexset index set to output
 * @return ost output stream
 */
std::ostream & operator<<(std::ostream &ost, const IndexSet &indexset)
{
  ost << "{ ";
  for (IndexSet::const_iterator iter = indexset.begin();
      iter != indexset.end();
      ++iter)
    ost << *iter << " ";
  ost << "}";
  return ost;
}


int main(int argc, char **argv)
{
  size_t matrixdimension = 8;
  size_t subspacelimit = 4;

  if (argc < 2) {
    std::cerr << "Usage: " << argv[0] << " <matrixdim> <subspacelimit>" << std::endl;
    return 255;
  } else {
    {
      std::stringstream s(toString(argv[1]));;
      s >> matrixdimension;
    }
    {
      std::stringstream s(toString(argv[2]));;
      s >> subspacelimit;
    }
  }

  MatrixContent *matrix = new MatrixContent(matrixdimension,matrixdimension);
  matrix->setZero();
  for (size_t i=0; i<matrixdimension ; i++) {
    for (size_t j=0; j<= i; ++j) {
      //const double value = 10. * rand() / (double)RAND_MAX;
      //const double value = i==j ? 2. : 1.;
      if (i==j)
        matrix->set(i,i, 2.);
      else if (j+1 == i) {
        matrix->set(i,j, 1.);
        matrix->set(j,i, 1.);
      } else {
        matrix->set(i,j, 0.);
        matrix->set(j,i, 0.);
      }
    }
  }

  Eigenspace::eigenvectorset CurrentEigenvectors;
  Eigenspace::eigenvalueset CurrentEigenvalues;

  setVerbosity(3);

  boost::timer Time_generatingfullspace;
  DoLog(0) && (Log() << Verbose(0) << std::endl << std::endl);
  // create the total index set
  IndexSet AllIndices;
  for (size_t i=0;i<matrixdimension;++i)
    AllIndices.insert(i);
  Eigenspace FullSpace(AllIndices, *matrix);
  DoLog(1) && (Log() << Verbose(1) << "Generated full space: " << FullSpace << std::endl);
  DoLog(0) && (Log() << Verbose(0) << "Full space generation took " << Time_generatingfullspace.elapsed() << " seconds." << std::endl);

  // generate first set of eigenvectors
  // set to first guess, i.e. the unit vectors of R^matrixdimension
  BOOST_FOREACH( size_t index, AllIndices) {
    boost::shared_ptr<VectorContent> EV(new VectorContent(matrixdimension));
    EV->setZero();
    EV->at(index) = 1.;
    CurrentEigenvectors.push_back(EV);
    CurrentEigenvalues.push_back(0.);
  }

  boost::timer Time_generatingsubsets;
  DoLog(0) && (Log() << Verbose(0) << "Generating sub sets ..." << std::endl);
  SetofIndexSets SetofSets;
  // note that starting off empty set is unstable
  IndexSet singleset;
  BOOST_FOREACH(size_t iter, AllIndices) {
    singleset.insert(iter);
    SetofSets.insert(singleset);
    singleset.clear();
  }
  SetofIndexSets::iterator CurrentSet = SetofSets.begin();
  while (CurrentSet != SetofSets.end()) {
    //DoLog(2) && (Log() << Verbose(2) << "Current set is " << *CurrentSet << std::endl);
    if (!generatePowerSet(SetofSets, CurrentSet, AllIndices, subspacelimit)) {
      // go to next set
      ++CurrentSet;
    }
  }
  DoLog(0) && (Log() << Verbose(0) << "Sub set generation took " << Time_generatingsubsets.elapsed() << " seconds." << std::endl);

  // create a subspace to each set and and to respective level
  boost::timer Time_generatingsubspaces;
  DoLog(0) && (Log() << Verbose(0) << "Generating sub spaces ..." << std::endl);
  SubspaceMap Dimension_to_Indexset;
  BOOST_FOREACH(std::set<size_t> iter, SetofSets) {
    boost::shared_ptr<Subspace> subspace(new Subspace(iter, FullSpace));
    DoLog(1) && (Log() << Verbose(1) << "Current subspace is " << *subspace << std::endl);
    Dimension_to_Indexset.insert( make_pair(iter.size(), boost::shared_ptr<Subspace>(subspace)) );
  }

  for (size_t dim = 1; dim<=subspacelimit;++dim) {
    BOOST_FOREACH( SubspaceMap::value_type subspace, Dimension_to_Indexset.equal_range(dim)) {
      if (dim != 0) { // from level 1 and onward
        BOOST_FOREACH( SubspaceMap::value_type entry, Dimension_to_Indexset.equal_range(dim-1)) {
          if (subspace.second->contains(*entry.second)) {
            // if contained then add ...
            subspace.second->addSubset(entry.second);
            // ... and also its containees as they are all automatically contained as well
            BOOST_FOREACH(boost::shared_ptr<Subspace> iter, entry.second->getSubIndices()) {
              subspace.second->addSubset(iter);
            }
          }
        }
      }
    }
  }
  DoLog(0) && (Log() << Verbose(0) << "Sub space generation took " << Time_generatingsubspaces.elapsed() << " seconds." << std::endl);

  // create a file handle for the eigenvalues
  std::ofstream outputvalues("eigenvalues.dat", std::ios_base::trunc);
  ASSERT(outputvalues.good(),
      "SubspaceFactorizerUnittest::EigenvectorTest() - failed to open eigenvalue file!");
  outputvalues << "# iteration ";
  BOOST_FOREACH(size_t iter, AllIndices) {
    outputvalues << "\teigenvalue" << iter;
  }
  outputvalues << std::endl;

  DoLog(0) && (Log() << Verbose(0) << "Solving ..." << std::endl);
  boost::timer Time_solving;
  size_t run=1; // counting iterations
  double threshold = 1.;  // containing threshold value
  while ((threshold > MYEPSILON) && (run < 20)) {
    // for every dimension
    for (size_t dim = 1; dim <= subspacelimit;++dim) {
      // for every index set of this dimension
      DoLog(1) && (Log() << Verbose(1) << std::endl << std::endl);
      DoLog(1) && (Log() << Verbose(1) << "Current dimension is " << dim << std::endl);
      std::pair<SubspaceMap::const_iterator,SubspaceMap::const_iterator> Bounds = Dimension_to_Indexset.equal_range(dim);
      for (SubspaceMap::const_iterator IndexsetIter = Bounds.first;
          IndexsetIter != Bounds.second;
          ++IndexsetIter) {
        Subspace& subspace = *(IndexsetIter->second);
        // show the index set
        DoLog(2) && (Log() << Verbose(2) << std::endl);
        DoLog(2) && (Log() << Verbose(2) << "Current subspace is " << subspace << std::endl);

        // solve
        subspace.calculateEigenSubspace();

        // note that assignment to global eigenvectors all remains within subspace
      }
    }

    // print list of similar eigenvectors
    DoLog(2) && (Log() << Verbose(2) << std::endl);
    BOOST_FOREACH( size_t index, AllIndices) {
      DoLog(2) && (Log() << Verbose(2) << "Similar to " << index << "th current eigenvector " << *(CurrentEigenvectors[index]) << " are:" << std::endl);
      BOOST_FOREACH( SubspaceMap::value_type iter, Dimension_to_Indexset) {
        const VectorContent & CurrentEV = (iter.second)->getEigenvectorParallelToFullOne(index);
        if (!CurrentEV.IsZero())
          Log() << Verbose(2)
          << "dim" << iter.first
          << ", subspace{" << (iter.second)->getIndices()
          << "}: "<< CurrentEV << std::endl;
      }
      DoLog(2) && (Log() << Verbose(2) << std::endl);
    }

    // create new CurrentEigenvectors from averaging parallel ones.
    BOOST_FOREACH(size_t index, AllIndices) {
      CurrentEigenvectors[index]->setZero();
      CurrentEigenvalues[index] = 0.;
      size_t count = 0;
      BOOST_FOREACH( SubspaceMap::value_type iter, Dimension_to_Indexset) {
        const VectorContent CurrentEV = (iter.second)->getEigenvectorParallelToFullOne(index);
        *CurrentEigenvectors[index] += CurrentEV; // * (iter.second)->getEigenvalueOfEigenvectorParallelToFullOne(index);
        CurrentEigenvalues[index] += (iter.second)->getEigenvalueOfEigenvectorParallelToFullOne(index);
        if (!CurrentEV.IsZero())
          count++;
      }
      *CurrentEigenvectors[index] *= 1./CurrentEigenvalues[index];
      //CurrentEigenvalues[index] /= (double)count;
    }

    // check orthonormality
    threshold = calculateOrthogonalityThreshold(AllIndices, CurrentEigenvectors);
    bool dontOrthonormalization = checkOrthogonality(AllIndices, CurrentEigenvectors);

    // orthonormalize
    if (!dontOrthonormalization) {
      DoLog(1) && (Log() << Verbose(1) << "Orthonormalizing ... " << std::endl);
      for (IndexSet::const_iterator firstindex = AllIndices.begin();
          firstindex != AllIndices.end();
          ++firstindex) {
        for (IndexSet::const_iterator secondindex = firstindex;
            secondindex != AllIndices.end();
            ++secondindex) {
          if (*firstindex == *secondindex) {
            (*CurrentEigenvectors[*secondindex]) *= 1./(*CurrentEigenvectors[*secondindex]).Norm();
          } else {
            (*CurrentEigenvectors[*secondindex]) -=
                ((*CurrentEigenvectors[*firstindex])*(*CurrentEigenvectors[*secondindex]))
                *(*CurrentEigenvectors[*firstindex]);
          }
        }
      }
    }

//    // check orthonormality again
//    checkOrthogonality(AllIndices, CurrentEigenvectors);

    // put obtained eigenvectors into full space
    FullSpace.setEigenpairs(CurrentEigenvectors, CurrentEigenvalues);

    // show new ones
    DoLog(1) && (Log() << Verbose(1) << "Resulting new eigenvectors and -values, run " << run << " are:" << std::endl);
    outputvalues << run;
    BOOST_FOREACH( size_t index, AllIndices) {
      DoLog(1) && (Log() << Verbose(1) << *CurrentEigenvectors[index] << " with " << CurrentEigenvalues[index] << std::endl);
      outputvalues << "\t" << CurrentEigenvalues[index];
    }
    outputvalues << std::endl;

    // and next iteration
    DoLog(0) && (Log() << Verbose(0) << "\titeration #" << run << std::endl);
    run++;
  }
  DoLog(0) && (Log() << Verbose(0) << "Solving took " << Time_solving.elapsed() << " seconds." << std::endl);
  // show final ones
  DoLog(0) && (Log() << Verbose(0) << "Resulting new eigenvectors and -values, run " << run << " are:" << std::endl);
  outputvalues << run;
  BOOST_FOREACH( size_t index, AllIndices) {
    DoLog(0) && (Log() << Verbose(0) << *CurrentEigenvectors[index] << " with " << CurrentEigenvalues[index] << std::endl);
    outputvalues << "\t" << CurrentEigenvalues[index];
  }
  outputvalues << std::endl;
  outputvalues.close();

  setVerbosity(2);

  DoLog(0) && (Log() << Verbose(0) << "Solving full space ..." << std::endl);
  boost::timer Time_comparison;
  MatrixContent tempFullspaceMatrix = FullSpace.getEigenspaceMatrix();
  gsl_vector *eigenvalues = tempFullspaceMatrix.transformToEigenbasis();
  tempFullspaceMatrix.sortEigenbasis(eigenvalues);
  DoLog(0) && (Log() << Verbose(0) << "full space solution took " << Time_comparison.elapsed() << " seconds." << std::endl);

  delete matrix;

  return 0;
}