#include <RAT/SmearResult.hh>
#include <RAT/Processor.hh>
#include <RAT/Log.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/FitVertex.hh>
#include <CLHEP/Random/RandFlat.h>

using namespace RAT;
using namespace RAT::Methods;

SmearResult::SmearResult() {}

void SmearResult::Initialise(const std::string&)
{
  // map enum types to strings
  fStringToTypeMap.insert(std::make_pair("energy", energy));
  fStringToTypeMap.insert(std::make_pair("position", position));
  fStringToTypeMap.insert(std::make_pair("posx", posx));
  fStringToTypeMap.insert(std::make_pair("posy", posy));
  fStringToTypeMap.insert(std::make_pair("posz", posz));
  fStringToTypeMap.insert(std::make_pair("posr", posr));
  fStringToTypeMap.insert(std::make_pair("postheta", postheta));
  fStringToTypeMap.insert(std::make_pair("posphi", posphi));
  fStringToTypeMap.insert(std::make_pair("direction", direction));

  for(std::map< std::string, SmearType >::iterator it=fStringToTypeMap.begin();
      it != fStringToTypeMap.end(); ++it)
    fTypeToStringMap.insert(std::make_pair(it->second, it->first));

  for(std::map< std::string, SmearType >::iterator it = fStringToTypeMap.begin();
      it != fStringToTypeMap.end(); ++it)
    fMapOfSmears.insert(std::make_pair(it->second, false));

  // Set all smears OFF by default
  fAppliedPosition = false;
  fAppliedDirection = false;
  fAppliedEnergy = false;
}

void SmearResult::SetActivePosition()
{
  fCurrentSmear = position;
}

void SmearResult::SetActiveDirection()
{
  fCurrentSmear = direction;
}

void SmearResult::SetActiveEnergy()
{
  fCurrentSmear = energy;
}

void SmearResult::SetD(const std::string& param, const double value)
{
  if(param == "low")
  {
    double xhi = fMapOfFunctions[fCurrentSmear].GetXmax();
    fMapOfFunctions[fCurrentSmear].SetRange(value, xhi);
  }
  else if(param == "high")
  {
    double xlo = fMapOfFunctions[fCurrentSmear].GetXmin();
    fMapOfFunctions[fCurrentSmear].SetRange(xlo, value);
  }
  else
    throw Processor::ParamUnknown(param);
}

void SmearResult::SetS(const std::string& param_string, const std::string& value)
{
  if( fStringToTypeMap.count(param_string) )
  {
    SmearType param = fStringToTypeMap[param_string];
    if(fMapOfSmears.count(param))
    {
      if(value == "on")
      {
        fMapOfSmears[param] = 1;
        fCurrentSmear = param;
        SetApplied(param_string);
      }
      else if(value == "off")
        fMapOfSmears[param] = 0;
      else
        throw Processor::ParamInvalid(param_string, "must be on or off");
    }
  }
  else if(param_string == "function")
    fMapOfFunctions[fCurrentSmear] =
      TF1( fTypeToStringMap[fCurrentSmear].c_str(), value.c_str(), -1, 1);
  else
  {
    throw Processor::ParamUnknown(param_string);
  }
}

void SmearResult::SetApplied(const std::string& param)
{
  // Split params into position, direction, and energy
  if( param == "energy" )
    fAppliedEnergy = true;
  else if( param == "direction" )
    fAppliedDirection = true;
  else if( param == "position" ||
           param == "posx"     ||
           param == "posy"     ||
           param == "posz"     ||
           param == "posr"     ||
           param == "postheta" ||
           param == "posphi" )
    fAppliedPosition = true;
}

void SmearResult::BeginOfRun(DS::Run&)
{
  DefaultSeed();
}

void SmearResult::EndOfRun(DS::Run&)
{
}

DS::FitResult SmearResult::GetBestFit()
{
  // Set the fFitResult to the given seed
  fFitResult.Reset();
  fFitResult = GetSeed();
  if( !fMapOfSmears[fCurrentSmear] )
    return fFitResult;

  DS::FitVertex vertex = fFitResult.GetVertex(0);
  TVector3 seedPosition = vertex.GetPosition();
  TVector3 seedDirection = vertex.GetDirection();
  double seedEnergy = vertex.GetEnergy();

  if( fCurrentSmear == energy )
  {
    fMapOfFunctions[fCurrentSmear].SetParameter(0, seedEnergy);
    double newEnergy = fMapOfFunctions[fCurrentSmear].GetRandom();
    vertex.SetEnergy(newEnergy);
  }
  if( fCurrentSmear == position )
  {
    if( fMapOfSmears[position] )
    {
      //double x,y,z;
      double sTheta = CLHEP::RandFlat::shoot(-TMath::Pi(), TMath::Pi());
      double sPhi = TMath::ACos( CLHEP::RandFlat::shoot(-1.0, 1.0) );
      double sR = fMapOfFunctions[fCurrentSmear].GetRandom();
      TVector3 smearVec;
      smearVec.SetMagThetaPhi(sR, sTheta, sPhi);
      //TVector3 smearVec(x,y,z);
      vertex.SetPosition(smearVec + seedPosition);
    }
    if( fMapOfSmears[posx] )
    {
      TVector3 smearVec(fMapOfFunctions[fCurrentSmear].GetRandom(), 0, 0);
      vertex.SetPosition(smearVec + seedPosition);
    }
    if( fMapOfSmears[posy] )
    {
      TVector3 smearVec(0, fMapOfFunctions[fCurrentSmear].GetRandom(), 0);
      vertex.SetPosition(smearVec + seedPosition);
    }
    if( fMapOfSmears[posz] )
    {
      TVector3 smearVec(0, 0, fMapOfFunctions[fCurrentSmear].GetRandom());
      vertex.SetPosition(smearVec + seedPosition);
    }
    if( fMapOfSmears[posr] )
    {
      seedPosition.SetMag( seedPosition.Mag() + fMapOfFunctions[fCurrentSmear].GetRandom() );
      vertex.SetPosition(seedPosition);
    }
    if( fMapOfSmears[postheta] )
    {
      seedPosition.SetTheta( seedPosition.Theta() + fMapOfFunctions[fCurrentSmear].GetRandom() );
      vertex.SetPosition(seedPosition);
    }
    if( fMapOfSmears[posphi] )
    {
      seedPosition.SetPhi( seedPosition.Phi() + fMapOfFunctions[fCurrentSmear].GetRandom() );
      vertex.SetPosition(seedPosition);
    }
  }
  if( fCurrentSmear == direction )
  {
    TRotation rMatrix;
    rMatrix.Rotate( CLHEP::RandFlat::shoot(-TMath::Pi(), TMath::Pi()), seedDirection );
    double randomAngle = fMapOfFunctions[fCurrentSmear].GetRandom();
    seedDirection.SetTheta( randomAngle + seedDirection.Theta() );
    seedDirection = rMatrix * seedDirection;
    vertex.SetDirection(seedDirection);
  }

  fFitResult.SetVertex(0, vertex);
  return fFitResult;
}

void SmearResult::DefaultSeed()
{
  DS::FitVertex seedVertex;
  seedVertex.SetPosition(TVector3(0, 0, 0), false, true);
  seedVertex.SetEnergy(2.0, false, true);
  fSeedResult.SetVertex(0, seedVertex);
}