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

// Fit
#include <XSFit/Fit/Fit.h>
// Background
#include <XSModel/Data/BackCorr/Background.h>
// SpectralData
#include <XSModel/Data/SpectralData.h>
// Weight
#include <XSModel/GlobalContainer/Weight.h>
// Model
#include <XSModel/Model/Model.h>
// ChiSquare
#include <XSFit/StatMethod/ChiSquare/ChiSquare.h>

#include <XSFit/Fit/FitMethod.h>
#include <XSModel/GlobalContainer/ModelContainer.h>
#include <XSModel/GlobalContainer/DataSetTypes.h>
#include <XSModel/Parameter/ModParam.h>
#include <XSModel/Data/Detector/Response.h>
#include <XSUtil/Utils/XSstream.h>
#include <XSUtil/Utils/XSutility.h>
#include <XSUtil/Numerics/Gamma.h>
#include <XSUtil/Error/Error.h>
#include <XSstreams.h>
#include <XSContainer.h>
#include <XSsymbol.h>
#include <iostream>
#include <iomanip>
#include <typeinfo>
#include <set>
#include <numeric>
#include <sstream>
#include <cmath>
using namespace XSContainer;


// Class ChiSquare 

ChiSquare::ChiSquare()
  : StatMethod("chi"), 
    m_difference(),
    m_variance(),
    m_fullName("Chi-Squared"),
    m_scriptName("\\gx\\u2\\d"),
    m_dSigma()
{
   isDerivAnalytic(true);
}

ChiSquare::ChiSquare(const ChiSquare &right)
  : StatMethod(right), 
    m_difference(right.m_difference), 
    m_variance(right.m_variance),
    m_fullName(right.m_fullName),
    m_scriptName(right.m_scriptName),
    m_dSigma(right.m_dSigma)
{
}


void ChiSquare::doPerform (Fit* fit)
{

  // alright, now we perform the sum of the statistic over the models.
  Real chiSquare(0.);

  std::set<string>::iterator ss (fit->activeModels().begin());
  std::set<string>::iterator ssEnd (fit->activeModels().end());
  while ( ss != ssEnd)
  {
        models->calculate(*ss);
        models->fold(*ss);
	++ss;
  }   

  doReset(fit);


  ModelMapConstIter ml (models->modelSet().begin());      
  ModelMapConstIter  mlEnd (models->modelSet().end());      
  while ( ml != mlEnd )
  {
        // D - F for each model. In the case where there are multiple
        // sources difference will be iteratively defined.
        Model* m (ml->second);
        if (m->isActive())
        {
                m->difference(m_difference);                  
                weight()->dataVariance(*m,m_variance);
        }
        ++ml;
  }    

  SpectralDataMapConstIt   sm (fit->spectraToFit().begin());
  SpectralDataMapConstIt   smEnd (fit->spectraToFit().end());

  int savConVerbose = tpout.conVerbose();
  int savLogVerbose = tpout.logVerbose();
  XSstream::verbose(tpout,35,35);         
  while ( sm != smEnd )
  {
        size_t index = sm->first;
        const RealArray& diff = m_difference[index];
        RealArray& var  = m_variance[index];
        applyMinVariance(sm->second, var);
        RealArray quot  (diff*diff/var);
        if ( tpout.maxChatter() >= 35 )
        {
                tcout << "ChiSquare::doPerform output" <<std::endl;
                tcout << " i data  diff  var \n" ;
                size_t N (diff.size());
	        std::vector<Real> PS(N);
	        std::partial_sum(&quot[0],&quot[0]+N,PS.begin());

                for (size_t i = 0; i < N; ++i)
                {
                        tcout << i << "  " << sm->second->spectrum()[i] 
                              << " " << diff[i]
                              << " " << var[i] << " "  << quot[i]
                              << " " << PS[i]  << '\n';
                }
                tcout << std::endl;


        }
        chiSquare += quot.sum();             
        ++sm;       
  } 
  XSstream::verbose(tpout, savConVerbose, savLogVerbose);
  // output section, to be encapsulated.

  Real bayesianPrior = bayes().lPrior(fit);           
  statistic(chiSquare - 2.0*bayesianPrior);
}

void ChiSquare::setWeight (const std::string& weightName)
{
  // weightName string should already have been validated as
  // corresponding to an actual weight.
  Weight* input (weightScheme(weightName));
  if ( input )
  {
        weight(input);

        tcout << " Chi-Squared statistic: variance weighted using " << input->name() 
                        << " weighting" << std::endl;
  }
  else
  {
     throw RedAlert("Invalid weight scheme in ChiSquare::setWeight");
  } 
}

void ChiSquare::doInitialize (Fit* fit)
{

    // Important: initialize has two definitions, one where allocations 
    // take place (called for each fit) and one per iteration where
    // arrays for variance, difference are reset to their values for
    // a new iteration.
    // SO NEED to ADD new function reset() to do the 

    // deal with the common operations: initialize arrays needed for holding
    // derivatives (if required by fit's FitMethod).

    StatMethod::doInitialize(fit);   

    // now, initialize arrays that are specific to ChiSquare.
    // each is an array of arrays of size # of noticed channels,
    // one per spectrum loaded. 

    // for every weighting technique except model weighting, the variance
    // is constant (independent of the model) apart from a possible systematic
    // term we can add later. If the weighting is Model, then the data variance
    // is given by the sum of the model weight and the background variance. 
    // so here we need to store the background variance in the variance array.
    // If model systematic errors or model errors are called for these will be
    // added during iteration by the Weight::dataVariance() function.
    ArrayContainer::const_iterator pw (protoWorkspace().begin());
    ArrayContainer::const_iterator pwEnd  (protoWorkspace().end());




    while ( pw != pwEnd )
    {
          m_variance.insert(ArrayContainer::value_type(pw->first,pw->second));
          m_difference.insert(ArrayContainer::value_type(pw->first,pw->second));
          ++pw;
    }

    doReset(fit);
}

void ChiSquare::doDeallocate ()
{
  StatMethod::doDeallocate();      

  m_variance.clear();
  m_difference.clear();
  m_dSigma.clear();
}

void ChiSquare::doReset (Fit* fit)
{

  // StatMethod::doReset(fit);        

  // reset quantities specific to ChiSquare object for new iteration.

  SpectralDataMapConstIt   sm (fit->spectraToFit().begin());
  SpectralDataMapConstIt   smEnd (fit->spectraToFit().end());
  while ( sm != smEnd )
  {
        size_t index = sm->first;
        SpectralData* spectrum = sm->second;
        const std::valarray<size_t>& IN = spectrum->indirectNotice();
        // Subtle note: Because we are retrieving a subset from
        // a non-const valarray, if we don't do the static cast
        // it will remain as type indirect_array<Real>.  For some
        // reason on Solaris, this then causes a memory leak when 
        // assigned to m_difference[index].  The cast seems to get
        // rid of this.  
        m_difference[index] = RealArray(spectrum->spectrum()[IN]);
        if (weight()->name() == string("model"))
           m_variance[index] = RealArray(.0, IN.size());
        else
           m_variance[index] = RealArray(spectrum->variance()[IN]);
        if (spectrum->background())
        {
           m_difference[index] -= 
                           spectrum->background()->spectrum()[IN];
           if (spectrum->background()->variance().size() > 0)
           {
                   m_variance[index] += 
                           spectrum->background()->variance()[IN];
           }

        }
        if (spectrum->correction() && spectrum->correctionScale() != 0) 
        {
                m_difference[index] -=  spectrum->correctionScale()
                        *spectrum->correction()->spectrum()[IN];         
        }     
        ++sm;
  }

  ArrayContainer::iterator dsig = m_dSigma.begin();
  ArrayContainer::iterator dsigEnd = m_dSigma.end();
  while (dsig != dsigEnd)
  {
     dsig->second = 0.;
     ++dsig;
  }      
}

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

void ChiSquare::doReport (Fit* fit) const
{
  using namespace std;
  int degrees (0);
  size_t bins(0);
  degreesOfFreedom(fit,degrees,bins);
  ios_base::fmtflags save(tcout.flags());
  streamsize savePrec(tcout.precision());
  if (!fit->spectraToFit().empty())
  {
     Numerics::GammaQ GQ;
     Real nullHyp = GQ( degrees/2. , statistic()/2.);
     Real truncatedStat = setOutputFormat(fit);
     tcout << "\n Chi-Squared = " << setw(14) << truncatedStat << " using " << bins << " PHA bins.";
     truncatedStat = setOutputFormat(fit, degrees);
     tcout << "\n Reduced chi-squared = " << setw(14) << truncatedStat << " for " 
           << setw(6) << degrees << " degrees of freedom " << setprecision(6) 
           << "\n Null hypothesis probability = " << setw(14) << scientific << nullHyp << endl;  
  }
  else
  {
     tcout.setf(ios_base::fixed);
     tcout << "\n Chi-Squared = " << setw(14) << "0.0" << " using " << 0 << " PHA bins." << endl;     
  }
  tcout.precision(savePrec);
  tcout.flags(save);
}

Real ChiSquare::plotChi (Real obs, Real mod, Real error, Real areaTime, const Real& NODATA) const
{
  Real point (obs - mod > 0 ? (obs - mod)*(obs - mod) : -(obs - mod)*(obs - mod) );
  Real msyst (models->modelSystematicError());
  Real var   (msyst > 0 ? error*error + msyst*msyst*mod*mod : error*error);
  point /= var;

  return point;       
}

Real ChiSquare::plotDeltaChi (Real obs, Real mod, Real error, const Real& NODATA) const
{
  Real point ( obs - mod );
  Real msyst (models->modelSystematicError());
  Real var ( msyst > 0 ?  sqrt(error*error + msyst*msyst*mod*mod) : error ); 
  point /= var;
  return point;
}

ArrayContainer ChiSquare::getDifferences (const Fit* fit) const
{
  // seems a waste perhaps to return just a copy of previously calculated array,
  // but not all statistics will have calculated this.
  // the object of providing this virtual function is to support
  // peak residual calculation - may be there are other uses also.
  return m_difference;
}

Real ChiSquare::sumSecondDerivs (int pi, int pj, bool isPiResp, bool isPjResp)
{
  Real alpha_ij=.0;
  ArrayContainer::const_iterator varEnd = m_variance.end();
  ArrayContainer::const_iterator sigmaEnd = m_dSigma.end();
  if (!isPiResp && !isPjResp)
  {
     // Calculate (2/var)*(dF/dpi)(dF/dpj) summed over all fit spectra.
     const ArrayContainer& piBins = dModel().find(pi)->second;
     const ArrayContainer& pjBins = dModel().find(pj)->second;
     ArrayContainer::const_iterator piSpec = piBins.begin();
     ArrayContainer::const_iterator piEnd = piBins.end();
     ArrayContainer::const_iterator pjEnd = pjBins.end();
     while (piSpec != piEnd)
     {
        // ASSUMES piBins and pjBins containers have same number of
        // RealArrays, which are indexed by the same spectra numbers.
        // Does NOT ASSUME that iterating through both in parallel will
        // grab RealArrays corresponding to the same spectrum.
        size_t specNum = piSpec->first;
        ArrayContainer::const_iterator pjSpec = pjBins.find(specNum);
        ArrayContainer::const_iterator var = m_variance.find(specNum);
        if (pjSpec == pjEnd || var == varEnd)
        {
           throw RedAlert("Chi-Square calc: Second derivatives array container size mismatch.");
        }
        RealArray secondDeriv(piSpec->second);
        secondDeriv *= pjSpec->second;
        secondDeriv /= var->second;
        ArrayContainer::const_iterator sigmaSpec = m_dSigma.find(specNum);
        if (sigmaSpec != sigmaEnd)
        {
           RealArray sigmaCorrection;
           calcD2Sigma(specNum, sigmaCorrection);
           secondDeriv *= sigmaCorrection;
        } 
        alpha_ij += secondDeriv.sum();
        ++piSpec;
     }
     alpha_ij *= 2.0;
  }
  else
  {
     // 1 or both pars are response pars.   In all 3 possible
     // cases, only 1 spectrum (at most) applies.
     size_t onlySpec = 0;
     std::map<size_t,ArrayContainer>::const_iterator piContainer = 
                dModel().find(pi);
     std::map<size_t,ArrayContainer>::const_iterator pjContainer = 
                dModel().find(pj);
     if (isPiResp)
     {
        piContainer = dRM().find(pi);
        if (piContainer == dRM().end())
        {
           throw RedAlert("Chi-Square 2nd deriv: missing response parameter");
        }
        onlySpec = piContainer->second.begin()->first;
     }
     if (isPjResp)
     {
        pjContainer = dRM().find(pj);
        if (pjContainer == dRM().end())
        {
           throw RedAlert("Chi-Square 2nd deriv: missing response parameter");
        }
        size_t tmp = pjContainer->second.begin()->first;
        if (onlySpec && (tmp != onlySpec))
        {
           // 2 response pars belonging to different spectra.
           // No terms contribute to 2nd derivative.
           return 0.0;
        }
        onlySpec = tmp;
     }
     const RealArray& piSpec = piContainer->second.find(onlySpec)->second;
     const RealArray& pjSpec = pjContainer->second.find(onlySpec)->second;
     ArrayContainer::const_iterator var = m_variance.find(onlySpec);
     if (var == varEnd)
     {
        throw RedAlert("Chi-Square calc: Second derivatives array container size mismatch.");
     }
     RealArray secondDeriv(piSpec);
     secondDeriv *= pjSpec;
     secondDeriv /= var->second;
     ArrayContainer::const_iterator sigmaSpec = m_dSigma.find(onlySpec);
     if (sigmaSpec != sigmaEnd)
     {
        RealArray sigmaCorrection;
        calcD2Sigma(onlySpec, sigmaCorrection);
        secondDeriv *= sigmaCorrection;
     }
     alpha_ij = 2.0*secondDeriv.sum(); 
  }
  return alpha_ij;
}

Real ChiSquare::sumDerivs (int parameterIndex, bool isRespar)
{
    // The m_difference array will contain the quantity
    // d = (D-F), the m_variance array will contain the variance var, 
    // corrected for any model systematic terms. Calculation of m_difference
    // is done in Model::difference(ArrayContainer&) which is invoked by
    // ChiSquare::doPerform

    // the derivative of X^2 wrt each parameter j for each folded model F is then 
    // -2 Sigma_i ( d_i/var_i +  d_i/var_i [dsigma_i/dF] dF_i/dpj )
    // the second term is the 'weight derivative' and is provided by the
    // Weight::varianceDeriv function for each array in the model.

    // ChiSquare can work with the sum of the derivative components from
    // each of multiple sources.
    ArrayContainer& dm = dModel()[parameterIndex];

    Real dSdpj (0.);
    int savConVerbose = tpout.conVerbose();
    int savLogVerbose = tpout.logVerbose();
    if (!isRespar)
    {
       ArrayContainer::const_iterator d (dm.begin());
       ArrayContainer::const_iterator dEnd (dm.end());
       XSstream::verbose(tpout,33,33);         
       while (d != dEnd)
       {
	   size_t n = d->first;
	   RealArray dSdpA (m_difference[n]);
           dSdpA *= d->second;

	   dSdpA /= m_variance[n];
           ArrayContainer::const_iterator itSig = m_dSigma.find(n);
	   if ( itSig != m_dSigma.end() )
	   {
               const RealArray& dSig_n = itSig->second;
	       dSdpA *= (1. + m_difference[n]/m_variance[n]
		         *dSig_n);
	   }
	   dSdpj += -2.0*dSdpA.sum();
           if ( tpout.maxChatter() >= 33 )
	   {
	       size_t N = m_difference[n].size();
	       std::vector<Real> PS(N);
	       std::partial_sum(&dSdpA[0],&dSdpA[0]+N,PS.begin());
	       tcout << " Spectrum " << d->first << " Parameter# " << parameterIndex  
                               << " D:  " << dSdpj << '\n';
	       for (size_t j = 0; j < N; ++ j)
	       {
		   tcout << j << " " << std::setw(12) << m_difference[n][j]  << " "
		         << std::setw(12) << m_variance[n][j] << " "
		         << std::setw(12) << (d->second)[j]  << " "
		         << std::setw(12) << dSdpA[j] << " "
		         << std::setw(12) << PS[j]  << '\n'; 
	       }
	       tcout << " d " << n << "  " << dSdpj << '\n';  
	   }          
	   ++d;
       } // end spectra loop.
    } // end if not a Response parameter.
    else
    {
       std::map<size_t,ArrayContainer>::const_iterator dFdPContainer = dRM().find(parameterIndex);
       // For response parameters, dFdPContainer should only have
       // 1 entry, since only 1 spectrum is associated with reponse.
       int specNum = dFdPContainer->second.begin()->first;
       const RealArray& dFdP = dFdPContainer->second.begin()->second;
       RealArray dSdpA(m_difference[specNum]);
       dSdpA /= m_variance[specNum];
       dSdpA *= dFdP;
       ArrayContainer::const_iterator itSig = m_dSigma.find(specNum);
       if ( itSig != m_dSigma.end() )
       {
           const RealArray& dSig_n = itSig->second;
	   dSdpA *= (1. + m_difference[specNum]/m_variance[specNum]
		     *dSig_n);
       }
       dSdpj = -2.0*dSdpA.sum();
    }
    XSstream::verbose(tpout, savConVerbose, savLogVerbose);

    return dSdpj;
}

void ChiSquare::calcD2Sigma (size_t specNum, RealArray& correction) const
{
  // If variance has a model dependency, calculate contribution
  // to second derivatives.  This assumes dSigma has already
  // been determined to have array corresponding to specNum.
  // Let g = ((D-F)*sigma*d(sigma)/dF)/sigma^2,  then correction
  // = (1 + 4g + 3g^2).
    RealArray g(m_dSigma.find(specNum)->second);
    g *= m_difference.find(specNum)->second;
    g /= m_variance.find(specNum)->second;
    correction.resize(g.size());
    correction = g;
    correction *= g;
    correction *= 3.0;
    correction += 4.0*g;
    correction += 1.0;
}

void ChiSquare::getWeightDerivs (Fit* fit)
{
    // the loop around models computes a dF_i array for each parameter that
    // the model depends on. We must also produce a varianceDeriv ArrayContainer
    // object that returns -0.5*sigma_i*dsigma_i/dF (the extra factor of sigma is
    // there so we can compute just d_i and var_i). 
    ModelMapConstIter mm (models->modelSet().begin());
    ModelMapConstIter mmEnd (models->modelSet().end());
    while ( mm != mmEnd )
    {
        Model& m (*mm->second);
        if ( m.isActive()) weight()->varianceDeriv(m,m_dSigma);
	++mm;       
    }
}

std::pair<Real,Real> ChiSquare::adjustForFrozen (const SpectralData* spec, const std::vector<Model*>& models, const std::map<Model*,ArrayContainer >& savedValsForFrozen)
{
    // NOTE:  This ASSUMES m_difference and m_variance have been reset to their
    // initial values in doReset, prior to their usage here.  It is perhaps less
    // dangerous to simply call doReset from here, but that would result in many
    // redundant calls (see how adjustForFrozen is called from StatMethod::renormalize).
   const std::valarray<size_t>& IN = spec->indirectNotice();
   const size_t N = IN.size(); 
   size_t specNum = spec->spectrumNumber();
   std::pair<Real,Real> fSums(0.,0.);
   const RealArray& pha = m_difference[specNum];
   RealArray totalModel(0.,N);
   RealArray totModFroz(.0,N);
   RealArray var(m_variance[specNum]);
   applyMinVariance(spec, var);
   for (size_t i=0; i<models.size(); ++i)
   {
      Model* mod = models[i];
      std::map<Model*,ArrayContainer >::const_iterator itSaved =
                savedValsForFrozen.find(mod);
      if (itSaved != savedValsForFrozen.end())
      {
         totalModel += itSaved->second.find(specNum)->second[IN];
         totModFroz += mod->foldedModel(specNum)[IN];
      }
      else
      {
         totalModel += mod->foldedModel(specNum)[IN];
      }
   }
   RealArray totModNonFroz(totalModel - totModFroz);
   RealArray obsNonFroz((pha - totModFroz)*totModNonFroz);
   fSums.first = (obsNonFroz/var).sum();
   fSums.second = (totModNonFroz*totModNonFroz/var).sum();

   //print frozen and total arrays to log file for debugging purposes.
   int savConVerbose = tpout.conVerbose();
   int savLogVerbose = tpout.logVerbose();
   XSstream::verbose(tpout,33,33);         
   if ( tpout.maxChatter() >= 33 )
   {
      using std::setw;
      tcout.precision(6);
      for ( size_t j = 0; j < N ; ++j)
      {
         tcout   <<  " i " << setw(5) << j 
                 << " obs " << setw(12) << pha[j]
                 << " mod " << setw(12) << totalModel[j] ;
            tcout << " fro " << setw(12) << totModFroz[j];
         tcout << '\n';
      }
      tcout << " ofSum " << fSums.first << " mfSum " << fSums.second 
                << std::endl;
    }        
   XSstream::verbose(tpout, savConVerbose, savLogVerbose);

   return fSums;
}

void ChiSquare::applyMinVariance (const SpectralData* sd, RealArray& variance)
{
     const std::valarray<size_t>& IN = sd->indirectNotice();
     RealArray areaTime = sd->exposureTime()*sd->areaScale()[IN];
     const size_t sz = variance.size();
     RealArray backAreaTime(-1.0, areaTime.size());
     RealArray bscaleRatio(1.0, areaTime.size());
     if (sd->background())
     {
        const SpectralData* bckSpec = sd->background()->data();
        backAreaTime = bckSpec->exposureTime()*bckSpec->areaScale()[IN];
        bscaleRatio = sd->backgroundScale()[IN]/bckSpec->backgroundScale()[IN];
     }
     for (size_t i=0; i<sz; ++i)
     {
        if (std::fabs(variance[i]) <= SMALL)
        {
           variance[i] = ChiSquare::calcMinVariance(areaTime[i], backAreaTime[i],
                                bscaleRatio[i]);
        }
     }
}

Real ChiSquare::calcMinVariance (Real areaTime, Real backAreaTime, Real bscaleRatio)
{
   Real minVar = .0;
   // Use negative backAreaTime as a flag indicating no background file.
   if (backAreaTime < .0)
   {
      minVar = 1.0/(areaTime*areaTime);
   }
   else
   {
      Real factor = 1.0/(areaTime*areaTime);
      Real backFactor = (bscaleRatio*bscaleRatio)/(backAreaTime*backAreaTime);
      minVar = std::min(factor, backFactor);
   }
   return minVar;
}

// Additional Declarations