#include <globals.hh>
#include <Randomize.hh>

#include <RAT/MetaMedianSeed.hh>
#include <RAT/Method.hh>
#include <RAT/OptimisedMethod.hh>
#include <RAT/SeededMethod.hh>
#include <RAT/OptimiserFactory.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DB.hh>
#include <TH1D.h>
using namespace RAT::DS;
using namespace RAT::Optimisers;

#include <string>
#include <CLHEP/Units/SystemOfUnits.h>
#include <CLHEP/Units/PhysicalConstants.h>
#include <utility>
using namespace CLHEP;
using namespace std;

MetaMedianSeed::~MetaMedianSeed()
{
  delete fOptimiser;
}

void
MetaMedianSeed::Initialise( const std::string& param )
{
  if( param.empty() == true )
    fOptimiser = OptimiserFactory::GetInstance()->GetOptimiser( "minuit" );
  else
    fOptimiser = OptimiserFactory::GetInstance()->GetOptimiser( param );
  fNSeeds = 20;
}

void
MetaMedianSeed::BeginOfRun( DS::Run& run )
{
  DB* db = DB::Get();
  fAVRadius = db->GetLink("SOLID", "acrylic_vessel_outer")->GetD("r_sphere");
  fOptimiser->BeginOfRun( run );
}

void
MetaMedianSeed::EndOfRun( DS::Run& run )
{
  fOptimiser->EndOfRun( run );
}

double
MetaMedianSeed::Minimise()
{
  fMinFactor = 1.0;
  return Optimise();
}

double
MetaMedianSeed::Maximise()
{
  fMinFactor = -1.0;
  return Optimise();
}

double
MetaMedianSeed::Optimise()
{
  // Would like to use a TUPLE: ( param | positive errors | negative errors ) but this is only in C++11!
  // For the moment, use a pair: ( param | pair (positive errors | negative errors) ) ... i.e. pair inside a pair
  vector< pair< vector<double>, pair< vector<double>, vector<double> > > > fitResults;
  vector<double> fitValues;
  FitResult seed = fMethod->GetFitResult();
  fSeededMethod->UpdateSeed( seed );

  TVector3 bestPos(0,0,0);

  int iLoopCount=0;  // Keep a count of number of "valid" loops
  vector<double> vX, vY, vZ, vT;  // Compile vectors of fitted vertices

  fOptimiser->SetComponent( fMethod );
  // Now alter the seed and optimise
  for( int iLoop = 0; iLoopCount < fNSeeds; iLoop++ )
    {
      // Now alter the seed, if not the first iteration
      if( iLoop != 0 )
        {
          FitResult newSeed = NewSeed( seed );
          fSeededMethod->UpdateSeed( newSeed );
        }
      // Now optimise
      if( fMinFactor < 0 )
        fitValues.push_back( fOptimiser->Maximise() );
      else
        fitValues.push_back( fOptimiser->Minimise() );
      pair< vector<double>, vector<double> > errorsPair( fOptimiser->GetPositiveErrors(), fOptimiser->GetNegativeErrors() );
      fitResults.push_back( pair< vector<double>, pair< vector<double>, vector<double> > >(fOptimiser->GetParams(), errorsPair) );

      // Remove invalid results
      TVector3 currentPos(fitResults[iLoopCount].first[0],fitResults[iLoopCount].first[1],fitResults[iLoopCount].first[2]);
      if( fOptimiser->GetValid() == false )
        {
          fitResults.pop_back();
          fitValues.pop_back();
	  iLoopCount--;
        }
      else if(currentPos.Mag()>fAVRadius)
	{
	  // Check result is within a sensible physical radius
	  // e.g, within the AV/PSUP
	  fitResults.pop_back();
	  fitValues.pop_back();
	  iLoopCount--;
	}
      else
	{
	  for(int m=0; m<fitValues[iLoopCount]; m++)
	  {
	    // Load results into separate vectors
	    // Results are weighted by likelihood value
	    vX.push_back(fitResults[iLoopCount].first[0]);
	    vY.push_back(fitResults[iLoopCount].first[1]);
	    vZ.push_back(fitResults[iLoopCount].first[2]);
	    vT.push_back(fitResults[iLoopCount].first[3]);
	  }
	}
      iLoopCount++;
    }
  // Set not valid and return what ever the current values are
  if( fitValues.empty() )
    {
      fParams = fOptimiser->GetParams();
      fPositiveErrors = fOptimiser->GetPositiveErrors();
      fNegativeErrors = fOptimiser->GetNegativeErrors();
      fValid = false;
      return 0.0;
    }

  // Find the best fit by taking the median values of co-ordinates
  fValid = true;

  sort(vX.begin(), vX.end());
  sort(vY.begin(), vY.end());
  sort(vZ.begin(), vZ.end());
  sort(vT.begin(), vT.end());
  sort(fitValues.begin(), fitValues.end());

  /*
  // Ian: Attempt to reject poor values by looking at the spread of
  // the co-ordinates vectors. Possibly return to later...
  double widthX= (vX[3*vX.size()/4]-vX[vX.size()/4]);
  double widthY= (vY[3*vY.size()/4]-vY[vY.size()/4]);
  double widthZ= (vZ[3*vZ.size()/4]-vZ[vZ.size()/4]);
  double limit=2000;
  if(widthX>limit || widthY>limit || widthZ>limit)
    {
      // Too much spread in results vectors to be a good result
      // Mark as invalid!
      fValid = false;
      return 0.0;
    }
  */

  // Take median values
  vector<double> results, errors;
  results.push_back(vX[vX.size()/2]);
  results.push_back(vY[vY.size()/2]);
  results.push_back(vZ[vZ.size()/2]);
  results.push_back(vT[vT.size()/2]);

  // Attempt to assign errors based on spread of results vector
  double widthX= (vX[3*vX.size()/4]-vX[vX.size()/4]);
  double widthY= (vY[3*vY.size()/4]-vY[vY.size()/4]);
  double widthZ= (vZ[3*vZ.size()/4]-vZ[vZ.size()/4]);
  double widthT= (vT[3*vT.size()/4]-vT[vT.size()/4]);
  errors.push_back(widthX);
  errors.push_back(widthY);
  errors.push_back(widthZ);
  errors.push_back(widthT);

  // Now set the method result
  fParams = results;
  fPositiveErrors = errors;
  fNegativeErrors = errors;

  return fitValues[fitValues.size()/2];
}

FitResult
MetaMedianSeed::NewSeed( const FitResult& startSeed )
{
  FitResult newSeed = startSeed;
  FitVertex vertex = newSeed.GetVertex(0);
  if( vertex.ContainsPosition() )
	{
	  const double cosTheta = -1.0 + 2.0 * G4UniformRand();
	  const double phi = twopi * G4UniformRand();
	  const double r = pow( G4UniformRand(), 1.0 / 3.0 ) * fAVRadius;
	  vertex.SetPosition( TVector3( r * sin( acos( cosTheta ) ) * cos( phi ),
					r * sin( acos( cosTheta ) ) * sin( phi ),
				        r * cosTheta ) );
	  vertex.SetPositionErrors( ROOT::Math::XYZVectorF(fAVRadius, fAVRadius, fAVRadius) );
	}
  if( vertex.ContainsTime() )
	{
	  vertex.SetTime( 240.0 + G4UniformRand() * 20.0 );
	  vertex.SetTimeErrors( 20.0 );
	}
  if( vertex.ContainsDirection() )
	{
	  const double cosTheta = -1.0 + 2.0 * G4UniformRand();
	  const double phi = twopi * G4UniformRand();
	  vertex.SetDirection( TVector3( sin( acos( cosTheta ) ) * cos( phi ),
									 sin( acos( cosTheta ) ) * sin( phi ),
									 cosTheta ) );
	  vertex.SetDirectionErrors( TVector3( 1.0, 1.0, 1.0 ) );
	}
  if( vertex.ContainsEnergy() )
	{
	  // Random numbers
	  vertex.SetEnergy( 6.0 * G4UniformRand() );
	  vertex.SetEnergyErrors( 12.0 );
	}
  newSeed.SetVertex( 0, vertex );
  return newSeed;
}