#include <TStopwatch.h>
using namespace ROOT;

#include <RAT/SOCFitterProc.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/SOCPositionTimeChiSquared.hh>
#include <RAT/Optimiser.hh>
#include <RAT/FitterPMT.hh>
#include <RAT/Log.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/SOC.hh>
#include <RAT/DS/Run.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.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::DS;

#include <sstream>
using namespace std;

SOCFitterProc::SOCFitterProc()
  : Processor("fitter"), fMethod(NULL), fPDF(NULL), fOptimiser(NULL), fSelector(NULL)
{

}

SOCFitterProc::~SOCFitterProc()
{
  delete fMethod;
  delete fPDF;
  delete fOptimiser;
  delete fSelector;
}

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

  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" ) && fSelector !=NULL )
    fSelector->SetI( parts[1], value );
  else
    throw Processor::ParamUnknown( param );
}

void
SOCFitterProc::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" ) && fSelector !=NULL )
    fSelector->SetD( parts[1], value );
  else
    throw Processor::ParamUnknown( param );
}

void
SOCFitterProc::SetS( const std::string& param,
                     const std::string& value )
{
  if( param == string( "method" ) )
    {
      try { fMethod = MethodFactory::Get()->GetMethod( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "SOCFitterProc::SetS: Error method " + e.id + " is unknown." ); }
    }
  else if( param == string( "pdf" ) )
    {
      try { fPDF = PDFFactory::Get()->GetPDF( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "SOCFitterProc::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( "SOCFitterProc::SetS: Error optimiser " + e.id + " is unknown." ); }
      dynamic_cast< OptimisedMethod* >( fMethod )->SetOptimiser( fOptimiser );
    }
  else if( param == string( "selector" ) )
    {
      try { fSelector = PMTSelectorFactory::Get()->GetPMTSelector( value ); }
      catch( FactoryUnknownID& e ) { Log::Die( "SOCFitterProc::SetS: Error selector " + e.id + " is unknown." ); }
      dynamic_cast< SelectorMethod* >( fMethod )->AddPMTSelector( fSelector );
    }
  else if( param == string( "seed" ) )
    fSeeds.push_back( value );
  else if( param == string( "name" ) )
    fNickName = value;
  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" ) && fSelector !=NULL )
        fSelector->SetS( parts[1], value );
      else
        throw Processor::ParamUnknown( param );
    }
}

void
SOCFitterProc::BeginOfRun( DS::Run& run )
{
  if( fMethod )
    fMethod->BeginOfRun( run );
  if( fPDF )
    fPDF->BeginOfRun( run );
  if( fOptimiser )
    fOptimiser->BeginOfRun( run );
  if( fSelector )
    fSelector->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( fSelector != NULL )
        resultName << ":" << fSelector->GetName();
      if( fSeeds.empty() == false )
        for( vector< string >::const_iterator iSeed = fSeeds.begin(); iSeed != fSeeds.end(); ++iSeed )
          resultName << ":" << *iSeed;
      fNickName = resultName.str();
    }
}

void
SOCFitterProc::EndOfRun( DS::Run& run )
{
  // Loop over all the sources and fit each separately
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
  vector<string> socIDs = run.GetSOCSourceIDs();
  for( size_t isoc = 0; isoc < socIDs.size(); isoc++ )
    {
      DS::SOC& soc = run.GetSOCData( socIDs[isoc] );
      vector<FitterPMT> pmtData;
      std::vector<UInt_t> pmtIDs = soc.GetSOCPMTIDs();
      for( size_t ipmt = 0; ipmt < pmtIDs.size(); ipmt++ )
        {
          const DS::SOCPMT& pmt = soc.GetSOCPMT( pmtIDs[ipmt] );
          // Use only normal PMTs with good occupancy and a calculated peak
          if( pmt.GetPeakFindOK() == 0 && pmt.GetPromptOccupancy() > 5.0 && pmtInfo.GetType( pmt.GetID() ) == DU::PMTInfo::NORMAL )
            pmtData.push_back( FitterPMT( pmt ) );
        }
      // Now can call the method
      ExecuteFitter( pmtData, soc, run );
    }

  if( fMethod )
    fMethod->EndOfRun( run );
  if( fPDF )
    fPDF->EndOfRun( run );
  if( fOptimiser )
    fOptimiser->EndOfRun( run );
  if( fSelector )
    fSelector->EndOfRun( run );
}

Processor::Result
SOCFitterProc::ExecuteFitter( std::vector<FitterPMT>& pmtData,
                              DS::SOC& soc,
                              DS::Run& run )
{
  fMethod->SetEventData( pmtData, NULL, &run );

  SeededMethod* seedCast = dynamic_cast< SeededMethod* >( fMethod );
  if( seedCast != NULL )
    {

      // If the 'socPositionTimeChiSquared' method is being used
      // we need to seed using the SOC object...
      if ( fMethod->GetName() == "socPositionTimeChiSquared" ){
        SOCPositionTimeChiSquared* posTChi = dynamic_cast< SOCPositionTimeChiSquared* >( fMethod );
        posTChi->SeedBySOC( soc );
      }
      // ...otherwise use a default seed.
      else{
        seedCast->DefaultSeed();
      }

      if( fSeeds.empty() == false )
        {
          for( vector< string >::const_iterator iSeed = fSeeds.begin(); iSeed != fSeeds.end(); ++iSeed )
            {
              try { seedCast->SetSeed( soc.GetFitResult( *iSeed ) ); }
              catch( FitResult::NoFitResultError& e ) { warn << "SOCFitterProc::Event: Cannot seed from " << *iSeed << newline; }
            }
        }
    }

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