//   Read the documentation to learn more about C++ code generator
//   versioning.
//	  %X% %Q% %Z% %W%

// LevMarq
#include <XSFit/FitMethod/LevMarq/LevMarq.h>
#include <XSFit/Fit/Fit.h>
#include <XSFit/Fit/StatMethod.h>
#include <XSModel/Model/Model.h>
#include <XSModel/GlobalContainer/ModelContainer.h>
#include <XSUtil/Error/Error.h>
#include <XSUtil/Signals/SignalHandler.h>
#include <XSUtil/Utils/XSutility.h>
#include <XSUtil/Utils/XSstream.h>
#include <XSstreams.h>
#include <XSContainer.h>
#include <cmath>
#include <sstream>
#include <typeinfo>
#include <algorithm>
#include <iomanip>

void invertSVD(double* matrix, int n, bool qzero, double *wvec, double *vmat);


// Class LevMarq 

LevMarq::LevMarq(const LevMarq &right)
  : FitMethod(right), 
    m_dLambda(right.m_dLambda),
    m_variableParameters(right.m_variableParameters), // non-owning
    m_fitStatistic(right.m_fitStatistic),
    m_alpha(right.m_alpha),
    m_beta(right.m_beta)
{
}

LevMarq::LevMarq (const std::string& initString)
  : FitMethod(initString),
    m_dLambda(0.001),
    m_variableParameters(), // non-owning
    m_fitStatistic(.0),
    m_alpha(),
    m_beta()
{
  firstDerivativeRequired(true);
  secondDerivativeRequired(true);  
}


void LevMarq::doPerform (Fit* fit)
{
   // this is the driver program for levenberg-marquadt.
   // it should do exactly what lvmrun.f does for the fortran version.

  using namespace std;
  using namespace XSContainer;

  double oldFtstat = fit->statMethod()->statistic();
  m_fitStatistic = oldFtstat;

  m_dLambda = 0.001;  

  int iError = 0;
  int iErrorPar = 0;

  int savConVerbose = tpout.conVerbose();
  int savLogVerbose = tpout.logVerbose();  
  if (!fit->errorCalc())
  {
        XSstream::verbose(tpout,15,15);
        tcout << "Number of trials and critical delta: " << numberOfTrials() 
                << "  " << deltaCrit() << std::endl;
        XSstream::verbose(tpout, savConVerbose, savLogVerbose);
  }

  IntegerArray tempFrozenParams;

  bool writeHeader (true);

  if (!fit->errorCalc())  fit->reportParameters(tcout,writeHeader);

  int trialsRemaining = int(numberOfTrials());
  Real curDel = 2.*deltaCrit();

  const std::map<int,ModParam*>::const_iterator vpBegin = fit->variableParameters().begin();
  std::map<int,ModParam*>::const_iterator vp      = vpBegin;
  const std::map<int,ModParam*>::const_iterator vpEnd   = fit->variableParameters().end();

  // initialize parameter pegging flags.
  while ( vp != vpEnd )
  {
          vp->second->unpeg();
          ++vp;
  }

  ModelMapConstIter mi;
  ModelMapConstIter miEnd (models->modelSet().end());


  bool done (false);
  bool unpeggedParams(false);

  // If calling code has registered the SIGINT_Handler, this
  // will allow user to break with ctrl-C.  Otherwise, 
  // intHandler = 0 and has no affect.

  const SIGINT_Handler *intHandler = dynamic_cast<const SIGINT_Handler*>
                (SignalHandler::instance()->getHandler(SIGINT));
  bool ctrl_c_exit = false;
  do  {    
     --trialsRemaining;          
     tempFrozenParams.clear();
     m_variableParameters.clear();

     mi = models->modelSet().begin();
     while (mi != miEnd)
     {
        Model* m (mi->second);
        if ( m->isActive()) m->checkNorms(tempFrozenParams);     
        ++mi;       
     }

     size_t M (0);
     if ( (M = tempFrozenParams.size()) != 0)
     {
        tcout << " Due to zero model norms, the following fit parameters are"
              << " temporarily frozen:";
        for (size_t k = 0; k < M; ++k)
        {
           tcout << "\n  "  
                 << XSContainer::ModelContainer::parameterLabel(tempFrozenParams[k]);
        }
        tcout << std::endl;
     }          


     // create temporary param array consisting of variable parameters
     // that are not temporarily frozen.
     // all parameters found in the array of temporarily  frozen
     // parameters are NOT loaded into the algorithm engine.
     vp =  vpBegin;
     while (vp != vpEnd)
     {
        if (!std::binary_search(tempFrozenParams.begin(),tempFrozenParams.end(),
                                                        static_cast<int>(vp->first))
            && !vp->second->isPegged() )
        {
           m_variableParameters[vp->first] = vp->second;
        }
        ++vp;
     }

     // escape if there are no free parameters after the freezing exercise.
     if (m_variableParameters.size() == 0)  throw Fit::NotVariable();

     fit->statMethod()->resetDerivativeArrays();

     curFitXspec(&iError, &iErrorPar);

     // iError == 0,-1,+1 

     if (iError == 0)
     {
        // error checking and parameter pegging: not done if any parameters are  
        // pegged.
        if ( !unpeggedParams ) fit->checkPegged();

        // report parameter values
        fit->reportParameters(tcout);
        // set current and reset previous value.
        curDel = oldFtstat - m_fitStatistic;
        oldFtstat = m_fitStatistic;
        // check for convergence.
        done = ( curDel >= 0 && curDel <= deltaCrit() );

        // keep parameters that have been pegged
        // in that state until  convergence criterion
        // has been met.
        if (!done ) unpeggedParams = false;

     }
     else if ( iError == -1 )
     {
        // peg parameter and do at least one more iteration.
        ModParam* ppar = fit->variableParameters(iErrorPar);
        ppar->peg();
        tcout << " Parameter " 
           << XSContainer::ModelContainer::parameterLabel(iErrorPar) 
           << " is pegged due to zero (or negative) pivot element." << std::endl;
        done = false;    
     }
     else if (iError == +1)
     {
        // get out of the loop, convergence failed, unless there
        // are pegged parameters, in which case we unpeg them and 
        // try again 
        done = true;
        vp =  vpBegin;
        while ( vp != vpEnd )
        {
           if ( vp->second->isPegged() )
           {
              done = false;
              break;
           } 
           ++vp;
        }
     }

     if ( iError >= 0 )
     {
        size_t numberUnpegged (0);
        vp =  vpBegin;
        while ( vp != vpEnd )
        {
           if ( vp->second->isPegged() ) 
           {
              vp->second->unpeg();
              ++numberUnpegged;
           }
           ++vp;
        }       
        if ( numberUnpegged == 0 || !unpeggedParams  )
        {
           unpeggedParams = numberUnpegged > 0;
           if ( unpeggedParams )
           {
              if ( tpout.maxChatter() > 20 )
              {
                 tcout << " Unpegged  " << numberUnpegged 
                                 << " parameters" << std::endl;  
              }
              done = false;
              m_dLambda = 0.001;     
           }   

        } 
     }


     if ( trialsRemaining == 0 && !done)
     {
         // check for number of iterations exceeded.
         // prompt user for continuation -
         if ( fit->queryMode() == Fit::ON)
         {
             if (XSutility::yesToQuestion("Number of trials exceeded: continue fitting? ",1,
                         tcin))
             {
                trialsRemaining = int(numberOfTrials());
             }
             else 
             {
                done = true;
             }         
         }
         else
         {
             if  ( fit->queryMode() == Fit::YES )
             {
                trialsRemaining = int(numberOfTrials());
                done = false;
             }    
             else
             {
                done = true;
             }   
         }

     }

     tcout << std::flush;
     if (intHandler)
     {
        ctrl_c_exit = static_cast<bool>(intHandler->interrupted());
     }
  }   while (!done && !ctrl_c_exit); 
  if (ctrl_c_exit)
  {
     throw Fit::FitInterrupt();
  }     

  // If not called from error calculations, compute covariance matrix.
  if (!fit->errorCalc())
  {
     invertCorrelationMatrix();
     reportEigenvectors();
  }
}

LevMarq* LevMarq::clone () const
{
  return new LevMarq(*this);
}

void LevMarq::copy (const FitMethod& right)
{
  if ( typeid(*this) != typeid(right) ) 
          throw RedAlert("*** Programmer Error: copying to wrong type ***");  
  LevMarq that(static_cast<const LevMarq&>(right));
  swap(that);   
}

void LevMarq::swap (FitMethod& right)
{
  FitMethod::swap(right);
  LevMarq&  that = static_cast<LevMarq&>(right);
  std::swap(m_dLambda,that.m_dLambda);
  std::swap(m_variableParameters,that.m_variableParameters);
  std::swap(m_fitStatistic,that.m_fitStatistic);
  std::swap(m_alpha,that.m_alpha);
  std::swap(m_beta,that.m_beta);
}

void LevMarq::reportProgress (std::ostream& s, Fit* fit) const
{
  using namespace std;
  // Calling function will deal with restoring default format/precision.
  int logLevel = int(std::log10(m_dLambda) + 0.5);
  s.precision(6);
  s << setw(14) << left<< fit->statMethod()->statistic() << setw(6) << logLevel;
}

string LevMarq::settingString () const
{
    std::ostringstream out;
    out << numberOfTrials() << ' ' << deltaCrit();
    return out.str();
}

string LevMarq::fullName () const
{
  return string("Levenberg-Marquardt");
}

void LevMarq::curFitXspec (int* iError, int* iErrorPar)
{
   // This was adapted from the Fortran crfitx.f:  A version of
   // the Marquardt proceedure for minimization, following the lead 
   // of the Bevington program CURFIT.  The calling sequence,
   // intermediate value storage, and logic has been optimized for 
   // use with XSPEC.

   const size_t MAXITER = 10;
   RealArray parVals(0.0,m_variableParameters.size());
   std::map<int,ModParam*>::const_iterator itPar = m_variableParameters.begin();
   std::map<int,ModParam*>::const_iterator itParEnd = m_variableParameters.end();
   size_t i=0;
   while (itPar != itParEnd)
   {
      parVals[i] = itPar->second->value('a');
      ++itPar, ++i;
   }

   Real oldFitStat = m_fitStatistic;

   //  Save the original parameter values, in case the initial matrix
   //  inversion does not decrease the fit statistic and it must
   //  be redone with an increased value of dlamda.
   const RealArray savedParVals(parVals);

   try
   {
      // calcStatAndDerivatives can throw during model calculations.
      // Table models are the most likely to do this.

      // calculate the derivatives of the quantity T == (Obs-Mod)/sigma
      // (folded through the response) wrt the free parameters then 
      // calculate the alpha matrix and beta vectors.
      calcStatAndDerivatives(parVals, true, iError, iErrorPar);

      // iError = -1 is the only possible error return from 
      // calcStatAndDerivatives and occurs if a pivot element is zero.
      if (*iError == -1)
         return;

      //  loop round changing the parameters and recalculating the
      //  fit parameter until a better fit is achieved. After each
      //  failure the lambda parameter is multiplied by ten.

      bool isDone = false;
      size_t nIter = 0;
      m_dLambda *= 0.1;
      while (!isDone)
      {
         ++nIter;
         if (m_dLambda > 1.0e300)
            *iError = 1;
         else
         {
            m_dLambda *= 10.0;
            calcNewPars(parVals, savedParVals, iError, iErrorPar);

            // Evaluate the the new model, fold through the response, 
            // and compare with the data.
            calcStatAndDerivatives(parVals, false, iError, iErrorPar);

            // If the new fit is still worse than the original then 
            // ignore any parameter bounds errors. 
            // If the number of iterations has been exceeded then 
            // return a non-convergence error.
            if (m_fitStatistic > oldFitStat)
            {
               if (*iError == -1)
                  *iError = 0;
               if (nIter > MAXITER)
               {
                  *iErrorPar = static_cast<int>((m_fitStatistic-oldFitStat)*1.0e4 + 0.5);
                  *iError = 1;
               }
            }
         }
         isDone = (m_fitStatistic <= oldFitStat) || (*iError != 0);
      } // end while loop

      // Reset the lambda parameter
      if (m_dLambda > 1.0e-30)
         m_dLambda *= 0.1;

      // If error is 1 then reset to the situation on 
      // entry to the routine. (-1 for zero alpha matrix element
      // is OK at this point.)
      if (*iError == 1)
      {
         int savConVerbose = tpout.conVerbose();
         int savLogVerbose = tpout.logVerbose();
         XSstream::verbose(tpout, 15, 15);
         tcout << " Iteration failed - backing up " << *iError << " "
              << *iErrorPar <<std::endl;
         XSstream::verbose(tpout, savConVerbose, savLogVerbose);
         parVals = savedParVals;
         calcStatAndDerivatives(parVals, false, iError, iErrorPar);
      }
   }
   catch (YellowAlert&)
   {
      // If we're in here, presumably some model calculation went astray
      // in calcStatAndDerivatives.  Reset pars and calculations to the
      // previous attempt.  
      parVals = savedParVals;
      try
      {
         calcStatAndDerivatives(parVals, false, iError, iErrorPar);
      }
      catch (...)
      {
         // Hard to see how this can happen, but if it throws again 
         // then something is seriously wrong.
         throw RedAlert("Exit from LevMarq::curFitXspec");
      }
      throw;
   }
}

void LevMarq::calcNewPars (RealArray& newVals, const RealArray& savedVals, int* iError, int* iErrorPar)
{
   // Adapted from the Fortran clcnst.f.
   const int MAXTRY = 100;
   // ASSUMES m_alpha and m_beta are correctly sized and properly
   // filled in, and that newVals and savedVals are both the
   // same size as m_beta. No checking performed here.
   const size_t nPar = m_beta.size();
   // Normalize by the diagonal elements to improve the numerical 
   // accuracy of the inversion (basically this corrects for the 
   // large differences in magnitude between parameter values).
   RealArray diags(.0, nPar);
   RealArray invArray(.0, nPar*nPar);
   for (size_t i=0; i<nPar; ++i)
      diags[i] = m_alpha[i*nPar+i];
   for (size_t i=0; i<nPar; ++i)
      for (size_t j=0; j<nPar; ++j)
         invArray[i*nPar+j] = m_alpha[i*nPar+j]/sqrt(diags[i]*diags[j]);

   Real da = 1.0 + m_dLambda;
   for (size_t i=0; i<nPar; ++i)
      invArray[i*nPar+i] = da;

   // These 2 arrays are only needed for the invertSVD call.
   // This function won't be using their return values.
   XSutility::auto_array_ptr<double> pWvec(new double[nPar]);
   XSutility::auto_array_ptr<double> pVmat(new double[nPar*nPar]);
   // invert the alpha array. Now uses svd and trapping of any singularity.
   invertSVD(&invArray[0], nPar, true, pWvec.get(), pVmat.get());

   // Calculate the new model parameters.
   std::map<int,ModParam*>::const_iterator itPar = m_variableParameters.begin();
   for (size_t i=0; i<nPar; ++i, ++itPar)
   {
      da = 0.0;
      for (size_t j=0; j<nPar; ++j)
         da += m_beta[j]*invArray[i*nPar+j]/sqrt(diags[i]*diags[j]);

      // Make the correction to the old parameter value - tries to
      // handle the case when the new parameter value lies outside 
      // the limits.

      int iTry = 1;
      bool continueLoop = true;
      const ModParam* modPar = itPar->second;
      while (continueLoop)
      {
         newVals[i] = savedVals[i] + da;
         if (!modPar->isPeriodic() && (newVals[i] < modPar->value('l') ||
                newVals[i] > modPar->value('h')))
         {
            da /= 2.0;
            ++iTry;
            // If we can't find a valid new parameter value then return
            // an error code -1 so that the parameter gets pegged.
            if (iTry > MAXTRY)
            {
               continueLoop = false;
               *iError = -1;
               *iErrorPar = i;
            }
         }
         else
            continueLoop = false;
      }
   } // end loop over pars
}

void LevMarq::invertCorrelationMatrix ()
{
   // Adapted from ivermt.f
   // set up the correlation matrix from the alpha matrix - unlike in calcNewPars
   // we do not normalize the matrix so in this case we really do end up with
   // the correlation matrix.
   RealArray& covar = covariance();
   RealArray& eValue = evalue();
   RealArray& eVector = evector();

   covar.resize(m_alpha.size());
   covar = m_alpha;
   const size_t nPar = m_beta.size();
   eValue.resize(nPar);
   eVector.resize(m_alpha.size());

   invertSVD(&covar[0], nPar, false, &eValue[0], &eVector[0]);

   // If a parameter is not also included in the subset m_variableParameters, it will be left
   // with a sigma of -1.  This indicates something went wrong, ie. presumably it was left 
   // pegged or belongs to a group with a zero norm.
   std::map<int,ModParam*>::const_iterator vp   = XSContainer::fit->variableParameters().begin();
   std::map<int,ModParam*>::const_iterator vpEnd   = XSContainer::fit->variableParameters().end();
   while (vp != vpEnd)
   {
      vp->second->setValue(-1.0,'s');
      ++vp;
   }

   std::map<int,ModParam*>::iterator itPar = m_variableParameters.begin();
   std::map<int,ModParam*>::iterator itParEnd = m_variableParameters.end();
   size_t i=0;
   while (itPar != itParEnd)
   {
      Real var = covar[i*nPar+i];
      if (var < 0.0)
      {
         tcout << "***Warning: LevMarq::invertCorrelationMatrix: "
             << "\n            Negative diagonal element for parameter "
             << itPar->second->index() << std::endl;
      }
      else
      {
         itPar->second->setValue(sqrt(var),'s');
      }
      ++i, ++itPar;
   }
}

void LevMarq::calcStatAndDerivatives (const RealArray& parValues, bool doDerivatives, int* iError, int* iErrorPar)
{
   using namespace XSContainer;
   *iError = 0;
   *iErrorPar = 0;

   StatMethod* stat = fit->statMethod();
   // 'z' set adjusted value without setting recompute flags.
   // 'a' set adjusted value and force recomputation.
   char key('z');
   if (!doDerivatives)
      key = 'a';
   const size_t nPar = parValues.size();
   IntegerArray diffParams(nPar);
   size_t i=0;
   std::map<int,ModParam*>::iterator itPar = m_variableParameters.begin();
   std::map<int,ModParam*>::iterator itParEnd = m_variableParameters.end();
   while (itPar != itParEnd)
   {
      itPar->second->setValue(parValues[i],key);
      diffParams[i] = itPar->first;
      ++itPar, ++i;
   }

   // differentiate statistic if required  
   if (doDerivatives)
   {
      stat->differentiate(fit,diffParams);

      // Can't assume stat's beta and alpha arrays are of size 
      // nPar and nPar*nPar respectively, since nPar may be less
      // than fit->variableParameters.size() due to pegged pars
      // and/or zero norms.  However CAN assume that the first
      // nPar and nPar*nPar blocks in beta and alpha correspond to
      // the actual varying parameters in parValues array. 
      // (See how these get filled in StatMethod::analyticDifferentiate.) 
      const RealArray& b = stat->beta();
      const RealArray& a = stat->alpha();
      m_beta.resize(nPar);
      m_alpha.resize(nPar*nPar);
      for (size_t i=0; i<nPar; ++i)
      {
         m_beta[i] = b[i];
         size_t ni = nPar*i;
         for (size_t j=i; j<nPar; ++j)
         {
            size_t nj = nPar*j;
            m_alpha[ni+j] = a[ni+j];
            if (i != j)  m_alpha[nj+i] = a[nj+i];
         }
      }

      std::map<int,ModParam*>::const_iterator itPar = m_variableParameters.begin();
      for (size_t i=0; i<nPar; ++i, ++itPar)
      {
         Real diagElem = a[(nPar+1)*i];
         if (diagElem < .0)
         {
            // This can occur when model-dependent syst error is used.
            *iError = -1;
            *iErrorPar = itPar->first;
            tcout << "***Warning: Negative alpha-matrix diagonal element "
              << "for parameter " << XSContainer::ModelContainer::parameterLabel(*iErrorPar) 
              << std::endl;
         }
         else if ( diagElem < SMALL )
         {
            *iError = -1;
            *iErrorPar = itPar->first;
            tcout << "***Warning: Zero alpha-matrix diagonal element for parameter " 
                    << XSContainer::ModelContainer::parameterLabel(*iErrorPar) <<std::endl; 
         }
      }

   }
   stat->perform(fit);       

   m_fitStatistic = stat->statistic();        
}

void LevMarq::reportEigenvectors () const
{
   using namespace std;
   const RealArray& eValue = evalue();
   const RealArray& eVector = evector();
   const size_t nPar = eValue.size();
   ios_base::fmtflags saveFormat(tcout.flags());
   const int savePrecision = tcout.precision();

   const size_t linelen(10+min(10*nPar,static_cast<size_t>(80)));
   const string topBar(linelen,'=');
   const string bottomBar(linelen,'-');
   const size_t MAXSHOW = 40;
   const size_t nShowPar = min(nPar, MAXSHOW);

   tcout << topBar << "\n Variances and Principal Axes" << endl;
   tcout << "          ";
   size_t i=0;
   map<int,ModParam*>::const_iterator itPar = m_variableParameters.begin();
   while (i < nShowPar)
   {
      tcout << right << setw(10) 
                << XSContainer::ModelContainer::parameterLabel(itPar->first);
      ++i, ++itPar;
   }
   string ellsp = (nShowPar == MAXSHOW) ? " ... " : "  ";
   tcout << ellsp << endl;

   i=0;
   while (i < nShowPar)
   {
      tcout << ' ';
      tcout << uppercase << scientific << setprecision(2) << left
          << setw(9) << eValue[i] << "|" << fixed;
      for (size_t j=0; j<nShowPar; ++j)
      {
         tcout << setw(8) << right << showpoint 
            << eVector[j*nPar + i] << "  ";               
      } 
      if (nShowPar < nPar)
         tcout << " ... ";
      tcout << endl;
      ++i;
   }

   tcout << bottomBar << endl;

   tcout.flags(saveFormat);
   tcout.precision(savePrecision);
}

// Additional Declarations