#include <RAT/PartialWaterFitter.hh>
#include <RAT/MethodFactory.hh>
#include <RAT/PDFFactory.hh>
#include <RAT/OptimiserFactory.hh>
#include <RAT/ClassifierFactory.hh>
#include <RAT/PMTSelectorFactory.hh>
#include <RAT/Method.hh>
#include <RAT/Classifier.hh>
#include <RAT/OptimisedMethod.hh>
#include <RAT/PDFMethod.hh>
#include <RAT/SeededMethod.hh>
#include <RAT/SelectorMethod.hh>
#include <RAT/Optimiser.hh>
#include <RAT/SeededClassifier.hh>
#include <RAT/FitterPMT.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/Processor.hh>
#include <RAT/ListHelp.hh>
#include <RAT/MetaInformation.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/Entry.hh>
using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::PDFs;
using namespace RAT::Optimisers;
using namespace RAT::PMTSelectors;
using namespace RAT::DS;
using namespace RAT::Classifiers;

#include <sstream>
using namespace std;

#include <TStopwatch.h>
using namespace ROOT;

PartialWaterFitter::PartialWaterFitter()
  : Processor("partialWaterFitter")
{

  // Initialse the quad fitter
  fQuadSeed = MethodFactory::Get()->GetMethod( "quad" );

  // Initialise the simple direction fitter
  fDirectionSeed = MethodFactory::Get()->GetMethod( "simpleDirection" );

  // Initialise the position-time fitter and options
  fPositionTime = MethodFactory::Get()->GetMethod( "positionTimeLikelihood" );
  fMetaDrivePowell = OptimiserFactory::Get()->GetOptimiser( "metaDriveCorrectSeed-powell" );
  fMetaDrivePowell->SetD("drive", 400.0);
  fMetaDrivePowell->SetD("time", 2.0);
  // zCut defaults to -8m to 0m; for FTB the water level is expected to be -2.4 m
  fZCut = PMTSelectorFactory::Get()->GetPMTSelector( "zCut" );
  fZCut->SetD("highLimit", -2500.0);
  fGV1D = PDFFactory::Get()->GetPDF( "gv1d" );

  // Initialise the direction fitter and options
  fDirection = MethodFactory::Get()->GetMethod( "directionLikelihood" );
  fMetaDirection = OptimiserFactory::Get()->GetOptimiser( "metaDirectionSeed-minuit" );
  fModeCut = PMTSelectorFactory::Get()->GetPMTSelector( "modeCut" );
  fDirectionPDF = PDFFactory::Get()->GetPDF( "directionPDF" );

  fEnergyLookup = MethodFactory::Get()->GetMethod( "energyLookup" );

  fBeta14 = ClassifierFactory::Get()->GetClassifier( "beta14" );
  fITR = ClassifierFactory::Get()->GetClassifier( "ITR" );
  fQPDT = ClassifierFactory::Get()->GetClassifier( "QPDT" );

  /// Now combine the components where appropriate
  dynamic_cast< OptimisedMethod* >( fPositionTime )->SetOptimiser( fMetaDrivePowell );
  dynamic_cast< SelectorMethod* >( fPositionTime )->AddPMTSelector( fZCut );
  dynamic_cast< PDFMethod* >( fPositionTime )->SetPDF( fGV1D );

  dynamic_cast< OptimisedMethod* >( fDirection )->SetOptimiser( fMetaDirection );
  dynamic_cast< SelectorMethod* >( fDirection )->AddPMTSelector( fModeCut );
  dynamic_cast< PDFMethod* >( fDirection )->SetPDF( fDirectionPDF );

  fCutOff = 3.0; // Maximum nhit for which quad cannot run

}

PartialWaterFitter::~PartialWaterFitter()
{
  delete fQuadSeed;
  delete fPositionTime;
  delete fMetaDrivePowell;
  delete fZCut;
  delete fGV1D;
  delete fDirectionSeed;
  delete fDirection;
  delete fMetaDirection;
  delete fModeCut;
  delete fDirectionPDF;
  delete fEnergyLookup;
  delete fBeta14;
  delete fITR;
  delete fQPDT;
}

void
PartialWaterFitter::BeginOfRun( DS::Run& run )
{
  /// First call the begin of run functions
  fQuadSeed->BeginOfRun( run );
  fDirectionSeed->BeginOfRun( run );
  fPositionTime->BeginOfRun( run );
  fMetaDrivePowell->BeginOfRun( run );
  fZCut->BeginOfRun( run );
  fGV1D->BeginOfRun( run );
  fDirection->BeginOfRun( run );
  fMetaDirection->BeginOfRun( run );
  fModeCut->BeginOfRun( run );
  fDirectionPDF->BeginOfRun( run );
  fEnergyLookup->BeginOfRun( run );
  fBeta14->BeginOfRun( run );
  fITR->BeginOfRun( run );
  fQPDT->BeginOfRun( run );
}

void
PartialWaterFitter::SetD( const std::string& param, double value )
{
  if( param == "zCutUpper" )
    {
      fZCut->SetD("highLimit", value);
    }
  else
    throw Processor::ParamUnknown( param );
}

Processor::Result
PartialWaterFitter::DSEvent( DS::Run& run,
                      DS::Entry& ds )
{
  for( size_t iEV = 0; iEV < ds.GetEVCount(); iEV++ )
    Event( run, ds.GetEV( iEV ) );
  return OK;
}

void
PartialWaterFitter::EndOfRun( DS::Run& run )
{
  // First call the begin of run functions
  fQuadSeed->EndOfRun( run );
  fDirectionSeed->EndOfRun( run );
  fPositionTime->EndOfRun( run );
  fMetaDrivePowell->EndOfRun( run );
  fZCut->EndOfRun( run );
  fGV1D->EndOfRun( run );
  fDirection->EndOfRun( run );
  fMetaDirection->EndOfRun( run );
  fModeCut->EndOfRun( run );
  fDirectionPDF->EndOfRun( run );
  fEnergyLookup->EndOfRun( run );
  fBeta14->EndOfRun( run );
  fITR->EndOfRun( run );
  fQPDT->EndOfRun( run );
}

Processor::Result
PartialWaterFitter::Event( DS::Run& run,
                    DS::EV& ev )
{
  /// PartialWaterFitter Logic:
  ///  This is a simple overview of the logic that is expressed within the lines of this function
  ///
  ///  1. Fit the position & time as the seedResult from quad,
  ///      if nhits < fNhitCutoff then abort the fit and return no FitResult. \n
  ///  2. Fit the direction as dirSeed using simpleDirection with the quad as seed. \n
  ///  3. Fit the position & time as waterResult using the positionTimeLikelihood,
  ///      metaDriveCorrectSeed-powell, zCut, gv1d-lightwater-sno with the seedResult and dirSeed as the seeds
  ///     If this fails then abort the fit and return no FitResult. \n
  ///  4. Fit the direction as dirResult using directionLikelihood, metaDirectionSeed-minuit,
  ///      modeCut, directionPDF with waterResult and dirSeed as seeds
  ///     If this fails then abort the fit and return no FitResult. \n
  ///  5. Fit the energy using energyLookup and the waterResult as seed
  ///     If this fails then abort the fit and return no FitResult. \n
  ///  6. Combine results (pos, time, direction) into waterResult. \n
  ///  7. Run beta14 with waterResult as seed. \n
  ///  8. Run itr with waterResult as seed. \n
  ///  9. Run qpdt with waterResult as seed. \n

  TStopwatch timer;
  timer.Start( true );

  const size_t currentPass = MetaInformation::Get()->GetCurrentPass();
  ev.AddFitterPass( currentPass ); // Ensure the EV knows fitters were run for this pass
  ev.AddClassifierPass( currentPass ); // Ensure the EV knows classifiers were run for this pass

  vector<FitterPMT> pmtData;
  for( size_t iPMTCal =0; iPMTCal < ev.GetCalPMTs().GetCount(); iPMTCal++ )
    pmtData.push_back( FitterPMT( ev.GetCalPMTs().GetPMT( iPMTCal ) ) );

  // Initialize the seeds
  fQuadSeed->SetEventData( pmtData, &ev, &run );
  // Run the seed method
  FitResult seedResult;
  try
    {
      if( ev.GetCalPMTs().GetCount() > fCutOff )
        seedResult = fQuadSeed->GetBestFit();
      else
        {
          warn << "PartialWaterFitter::Event: insufficient points for quad to run, exiting" << newline;
          return Processor::FAIL;
        }
    }
  catch( Method::MethodFailError& error ) { warn << "PartialWaterFitter::Event: Seed failed " << error.what() << ", continuing." << newline; }

  /// Now initialise the direction method seed
  fDirectionSeed->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededMethod* >( fDirectionSeed )->DefaultSeed();
  dynamic_cast< SeededMethod* >( fDirectionSeed )->SetSeed( seedResult );
  /// Run the direction seed
  FitResult directionSeedResult;
  try { directionSeedResult = fDirectionSeed->GetBestFit(); }
  catch( Method::MethodFailError& error ) { warn << error.what() << newline; }

  /// Now initialise the position time method
  fPositionTime->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededMethod* >( fPositionTime )->DefaultSeed();
  dynamic_cast< SeededMethod* >( fPositionTime )->SetSeed( seedResult );
  dynamic_cast< SeededMethod* >( fPositionTime )->SetSeed( directionSeedResult );
  /// Run the position time method
  FitResult waterResult;
  try { waterResult = fPositionTime->GetBestFit(); }
  catch( Method::MethodFailError& error )
    {
      warn << "PartialWaterFitter::Event: Main method failed " << error.what() << ", exiting" << newline;
      return Processor::FAIL;
    }

  /// Now initialise the position time direction method
  fDirection->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededMethod* >( fDirection )->DefaultSeed();
  dynamic_cast< SeededMethod* >( fDirection )->SetSeed( waterResult );
  dynamic_cast< SeededMethod* >( fDirection )->SetSeed( directionSeedResult );
  /// Run the position time direction method

  try
    {
      FitVertex vertex = waterResult.GetVertex(0);
      FitVertex directionVertex = fDirection->GetBestFit().GetVertex(0);
      vertex.SetDirection( directionVertex.GetDirection(), directionVertex.ValidDirection() );
      vertex.SetPositiveDirectionError( directionVertex.GetPositiveDirectionError(), directionVertex.ValidDirection() );
      vertex.SetNegativeDirectionError( directionVertex.GetNegativeDirectionError(), directionVertex.ValidDirection() );
      waterResult.SetVertex( 0, vertex );
    }
  catch( Method::MethodFailError& error )
    {
      warn << "PartialWaterFitter::Event: Direction method failed " << error.what() << ", exiting" << newline;
      return Processor::FAIL;
    }
  catch( Optimiser::OptimiserFailError& error )
    {
      warn << "PartialWaterFitter::Event: Direction method failed " << error.what() << ", exiting" << newline;
      return Processor::FAIL;
    }

  // Now initialise the energy method
  fEnergyLookup->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededMethod* >( fEnergyLookup )->DefaultSeed();
  if( waterResult.GetVertexCount() && waterResult.GetVertex(0).ValidPosition() )
    dynamic_cast< SeededMethod* >( fEnergyLookup )->SetSeed( waterResult );
  else if( seedResult.GetVertexCount() && seedResult.GetVertex(0).ValidPosition() )
    dynamic_cast< SeededMethod* >( fEnergyLookup )->SetSeed( seedResult );
  else
    warn << "PartialWaterFitter::Event: No seed for the energy lookup method." << newline;
  // Run the energy lookup method
  try
    {
      FitVertex vertex = waterResult.GetVertex(0);
      FitVertex energyVertex = fEnergyLookup->GetBestFit().GetVertex(0);
      vertex.SetEnergy( energyVertex.GetEnergy(), energyVertex.ValidEnergy() );
      vertex.SetPositiveEnergyError( energyVertex.GetPositiveEnergyError(), energyVertex.ValidEnergy() );
      vertex.SetNegativeEnergyError( energyVertex.GetNegativeEnergyError(), energyVertex.ValidEnergy() );
      waterResult.SetVertex( 0, vertex );
    }
  catch( Method::MethodFailError& error )
    {
      warn << "PartialWaterFitter::Event: Energy method failed " << error.what() << ", exiting" << newline;
      return Processor::FAIL;
    }

  // Now initialise and run the Beta14 classifier
  fBeta14->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededClassifier* >( fBeta14 )->DefaultSeed();
  dynamic_cast< SeededClassifier* >( fBeta14 )->SetSeed( waterResult );
  try { ev.SetClassifierResult( currentPass, fBeta14->GetName() + ":partialWaterFitter", fBeta14->GetClassification() ); }
  catch( Classifier::ClassifierFailError& error ) { warn << error.what() << newline; }
  // Now initialise and run the ITR classifier
  fITR->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededClassifier* >( fITR )->DefaultSeed();
  dynamic_cast< SeededClassifier* >( fITR )->SetSeed( waterResult );
  try { ev.SetClassifierResult( currentPass, fITR->GetName() + ":partialWaterFitter", fITR->GetClassification() ); }
  catch( Classifier::ClassifierFailError& error ) { warn << error.what() << newline; }
  // Now initialise and run the QPDT classifier
  fQPDT->SetEventData( pmtData, &ev, &run );
  dynamic_cast< SeededClassifier* >( fQPDT )->DefaultSeed();
  dynamic_cast< SeededClassifier* >( fQPDT )->SetSeed( waterResult );
  try { ev.SetClassifierResult( currentPass, fQPDT->GetName() + ":partialWaterFitter", fQPDT->GetClassification() ); }
  catch( Classifier::ClassifierFailError& error ) { warn << error.what() << newline; }

  timer.Stop();
  waterResult.SetExecutionTime( timer.RealTime() );

  return SetResult(waterResult, ev, currentPass);
}

Processor::Result
PartialWaterFitter::SetResult( DS::FitResult& waterResult,
                        DS::EV& ev,
                        size_t currentPass)
{
  ev.SetFitResult( currentPass, "partialWaterFitter", waterResult );
  // Now set the 0 vertex as the default fit vertex, if the fit is valid
  ev.SetDefaultFitVertex( "partialWaterFitter", waterResult.GetVertex(0) );

  return Processor::OK;
}