#include <TStopwatch.h>
using namespace ROOT;

#include <RAT/FitterProcessor.hh>
#include <RAT/MethodFactory.hh>
#include <RAT/PDFFactory.hh>
#include <RAT/OptimiserFactory.hh>
#include <RAT/PMTSelectorFactory.hh>
#include <RAT/Method.hh>
#include <RAT/OptimisedMethod.hh>
#include <RAT/SelectorMethod.hh>
#include <RAT/PDFMethod.hh>
#include <RAT/SeededMethod.hh>
#include <RAT/Optimiser.hh>
#include <RAT/SeedModifier.hh>
#include <RAT/SeedModifierFactory.hh>
#include <RAT/FitterPMT.hh>
#include <RAT/Log.hh>
#include <RAT/MetaInformation.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/DataNotFound.hh>
#include <RAT/ListHelp.hh>
using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::PDFs;
using namespace RAT::Optimisers;
using namespace RAT::PMTSelectors;
using namespace RAT::SeedModifiers;
using namespace RAT::DS;

#include <sstream>
using namespace std;

FitterProcessor::FitterProcessor()
  : Processor("fitter"), fMethod(NULL), fPDF(NULL), fOptimiser(NULL), fSeedModifier(NULL)
{
  fPMTCalSelector = PMTSelectorFactory::Get()->GetPMTSelector( "PMTCalSelector" );
}

FitterProcessor::~FitterProcessor()
{
  delete fMethod;
  delete fPDF;
  delete fOptimiser;
  for (std::vector< PMTSelector* >::iterator iSelector = fSelectors.begin() ; iSelector != fSelectors.end(); ++iSelector)
    delete (*iSelector);
  fSelectors.clear();
  delete fSeedModifier;
  delete fPMTCalSelector;
}

void
FitterProcessor::SetI( const std::string& param,
                       const int value )
{
  vector<string> parts = split( param, "." );
  if( parts.size() != 2 )
    throw Processor::ParamUnknown( param );

  // For options that may be set multiple times (e.g. selector) this
  // sets the options for the most recent component.
  if( parts[0] == string( "method" ) && fMethod !=NULL )
    fMethod->SetI( parts[1], value );
  else if( parts[0] == string( "pdf" ) && fPDF !=NULL )
    fPDF->SetI( parts[1], value );
  else if( parts[0] == string( "optimiser" ) && fOptimiser !=NULL )
    fOptimiser->SetI( parts[1], value );
  else if( parts[0] == string( "selector" ) && fSelectors.size() != 0 )
    fSelectors.back()->SetI( parts[1], value );
  else if( parts[0] == string( "seed" ) && parts[1] == string( "pass" )
           && !fSeedPass.empty() )
    fSeedPass.back() = value;
  else if( parts[0] == string( "seedModifier" ) && fSeedModifier != NULL )
    fSeedModifier->SetI( parts[1], value );
  else
    throw Processor::ParamUnknown( param );
}

void
FitterProcessor::SetD( const std::string& param,
                       const double value )
{
  vector<string> parts = split( param, "." );
  if( parts.size() != 2 )
    throw Processor::ParamUnknown( param );

  if( parts[0] == string( "method" ) && fMethod !=NULL )
    fMethod->SetD( parts[1], value );
  else if( parts[0] == string( "pdf" ) && fPDF !=NULL )
    fPDF->SetD( parts[1], value );
  else if( parts[0] == string( "optimiser" ) && fOptimiser !=NULL )
    fOptimiser->SetD( parts[1], value );
  else if( parts[0] == string( "selector" ) && fSelectors.size() != 0 )
    fSelectors.back()->SetD( parts[1], value );
  else if( parts[0] == string( "seedModifier" ) && fSeedModifier != NULL )
    fSeedModifier->SetD( parts[1], value );
  else
    throw Processor::ParamUnknown( param );
}

void
FitterProcessor::SetS( const std::string& param,
                       const std::string& value )
{
  if( param == string( "method" ) )
    {
      try { fMethod = MethodFactory::Get()->GetMethod( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "FitterProcessor::SetS: Error method " + e.id + " is unknown." ); }
    }
  else if( param == string( "pdf" ) )
    {
      try { fPDF = PDFFactory::Get()->GetPDF( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "FitterProcessor::SetS: Error pdf " + e.id + " is unknown." ); }
      dynamic_cast< PDFMethod* >( fMethod )->SetPDF( fPDF );
    }
  else if( param == string( "optimiser" ) )
    {
      try { fOptimiser = OptimiserFactory::Get()->GetOptimiser( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "FitterProcessor::SetS: Error optimiser " + e.id + " is unknown." ); }
      dynamic_cast< OptimisedMethod* >( fMethod )->SetOptimiser( fOptimiser );
    }
  else if( param == string( "selector" ) )
    {
      PMTSelector* selector = NULL;
      try { selector = PMTSelectorFactory::Get()->GetPMTSelector( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "FitterProcessor::SetS: Error selector " + e.id + " is unknown." ); }
      dynamic_cast< SelectorMethod* >( fMethod )->AddPMTSelector( selector );
      fSelectorName.push_back( selector->GetName() );
      fSelectors.push_back( selector );
    }
  else if( param == string( "seed" ) )
    {
      fSeeds.push_back( value );
      fSeedPass.push_back( -1 ); // default to the current pass
    }
  else if( param == string( "name" ) )
    fNickName = value;
  else if( param == string( "seedModifier" ) )
    {
      try { fSeedModifier = SeedModifierFactory::Get()->GetSeedModifier( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "FitterProcessor::SetS: Error seedModifier " + e.id + " is unknown."); }
    }
  else
    {
      vector<string> parts = split( param, "." );
      if( parts.size() != 2 )
        throw Processor::ParamUnknown( param );

      if( parts[0] == string( "method" ) && fMethod !=NULL )
        fMethod->SetS( parts[1], value );
      else if( parts[0] == string( "pdf" ) && fPDF !=NULL )
        fPDF->SetS( parts[1], value );
      else if( parts[0] == string( "optimiser" ) && fOptimiser !=NULL )
        fOptimiser->SetS( parts[1], value );
      else if( parts[0] == string( "selector" ) && fSelectors.size() != 0 )
        fSelectors.back()->SetS( parts[1], value );
      else if( parts[0] == string( "seedModifier" ) && fSeedModifier != NULL )
        fSeedModifier->SetS( parts[1], value );
      else
        throw Processor::ParamUnknown( param );
    }
}

void
FitterProcessor::BeginOfRun( DS::Run& run )
{
  if( fMethod )
    fMethod->BeginOfRun( run );
  if( fPDF )
    fPDF->BeginOfRun( run );
  if( fOptimiser )
    fOptimiser->BeginOfRun( run );
  for (std::vector< PMTSelector* >::iterator iSelector = fSelectors.begin() ; iSelector != fSelectors.end(); ++iSelector)
    (*iSelector)->BeginOfRun( run );
  if( fSeedModifier )
    fSeedModifier->BeginOfRun( run );

  // Now construct the result name
  stringstream resultName;
  if( fNickName.empty() )
    {
      resultName << fMethod->GetName();
      if( fOptimiser != NULL )
        resultName << ":" << fOptimiser->GetName();
      if( fPDF != NULL )
        resultName << ":" << fPDF->GetName();
      if( fSelectorName.empty() == false )
	for( vector< string >::const_iterator iSelector = fSelectorName.begin(); iSelector != fSelectorName.end(); ++iSelector)
          resultName << ":" << *iSelector;
      if( fSeeds.empty() == false )
        for( vector< string >::const_iterator iSeed = fSeeds.begin(); iSeed != fSeeds.end(); ++iSeed )
          resultName << ":" << *iSeed;
      fNickName = resultName.str();
    }
}

Processor::Result
FitterProcessor::DSEvent( DS::Run& run,
                          DS::Entry& ds )
{
  const size_t pass = MetaInformation::Get()->GetCurrentPass();
  for( size_t iEV = 0; iEV < ds.GetEVCount(); iEV++ )
    {
      DS::EV& ev = ds.GetEV( iEV );
      ev.AddFitterPass( pass );
      vector<FitterPMT> pmtData;
      for( size_t iPMTCal = 0; iPMTCal < ev.GetCalPMTs().GetCount(); iPMTCal++ )
        pmtData.push_back( FitterPMT( ev.GetCalPMTs().GetPMT( iPMTCal ) ) );
      /// Additionally select pmtData so it *only* contains PMTs which pass the PMTCal groups recommended selector cuts.
      DS::FitVertex dummyVertex;
      pmtData = fPMTCalSelector->GetSelectedPMTs( pmtData, dummyVertex );

      fMethod->SetEventData( pmtData, &ev, &run, &ds );

      SeededMethod* seedCast = dynamic_cast< SeededMethod* >( fMethod );
      if( seedCast != NULL )
        {
          seedCast->DefaultSeed();
          if( fSeeds.empty() == false )
            {
              for( unsigned int iSeed = 0; iSeed < fSeeds.size(); iSeed++ )
                {
                  try
                    {
                      if( fSeedPass[ iSeed ] == -1 )
                        seedCast->SetSeed( ev.GetFitResult( fSeeds[ iSeed ] ) );
                      else
                        seedCast->SetSeed( ev.GetFitResults( fSeedPass[ iSeed ] ).GetResult( fSeeds[ iSeed ] ) );
                    }
                  catch( DataNotFound& ) { warn << "FitterProcessor::Event : Cannot seed from " << fSeeds[ iSeed ] << newline; }
                }
            }
          // if requested, adjust the seed
          // do this after all the seeds have been set (merged into a single seed in the method)
          // to ensure that the seed used is adjusted correctly.
          if( fSeedModifier != NULL )
            {
              seedCast->SetSeed( fSeedModifier->ModifySeed( ev, seedCast->GetSeed() ) );
            }

        }

      // Try to fit
      try
        {
          TStopwatch timer;
          timer.Start( true );
          FitResult result = fMethod->GetBestFit();
          timer.Stop();
          result.SetExecutionTime( timer.RealTime() );
          ev.SetFitResult( pass, fNickName, result );
        }
      catch( Method::MethodFailError& error )
        {
          debug << "FitterProcessor::Event " << error.what() << newline;
        }
      catch( Optimiser::OptimiserFailError& error )
        {
          debug << "FitterProcessor::Event " << error.what() << newline;
        }
    }
  return Processor::OK;
}

void
FitterProcessor::EndOfRun( DS::Run& run )
{
  if( fMethod )
    fMethod->EndOfRun( run );
  if( fPDF )
    fPDF->EndOfRun( run );
  if( fOptimiser )
    fOptimiser->EndOfRun( run );
  for (std::vector< PMTSelector* >::iterator iSelector = fSelectors.begin() ; iSelector != fSelectors.end(); ++iSelector)
    (*iSelector)->EndOfRun( run );
  if( fSeedModifier )
    fSeedModifier->EndOfRun( run );
}