#include <RAT/PositionANN.hh>
#include <RAT/DB.hh>
#include <RAT/Processor.hh>
#include <RAT/json.hh>

#include <RAT/DS/FitResult.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/DS/Entry.hh>

#include <TVector3.h>

using namespace std;
using namespace RAT;
using namespace RAT::Methods;


void PositionANN::BeginOfRun(DS::Run&) {

  DB* db = DB::Get();

  DBLinkPtr dbLink = db->GetLink("FIT_POSITION_ANN", "labppo_scintillator");

  fNDots = dbLink->GetD("nDots");
  fNTimes = dbLink->GetD("nTimes");
  fDotLow = dbLink->GetD("dotLow");
  fDotHigh = dbLink->GetD("dotHigh");
  fDotWidth = (fDotHigh - fDotLow) / fNDots;
  fTimeLow = dbLink->GetD("timeLow");
  fTimeHigh = dbLink->GetD("timeHigh");
  fTimeWidth = (fTimeHigh - fTimeLow) / fNTimes;
  fRadiusScale = dbLink->GetD("radiusScale");

  // The following DB values are variable size; get JSON values initially and then pass these to vectors
  json::Value weights = dbLink->GetJSON("weights");
  json::Value biases = dbLink->GetJSON("biases");

  // Reset the vector sizes
  fWeights.clear();
  fBiases.clear();

  fNLayers = weights.getArraySize() + 1;
  fWeights.resize(fNLayers-1);
  fBiases.resize(fNLayers-1);

  for(size_t iLayer = 0; iLayer < fNLayers - 1; iLayer++) {
    size_t layerSize = weights[iLayer].getArraySize();
    fBiases[iLayer] = biases[iLayer].toVector<double>();
    fWeights[iLayer].resize(layerSize);

    for(size_t iInput = 0; iInput < layerSize; iInput++) {
      fWeights[iLayer][iInput] = weights[iLayer][iInput].toVector<double>();
    }
  }

}


DS::FitResult PositionANN::GetBestFit() {

  fFitResult.Reset();
  if (fPMTData.empty())
  {
    return fFitResult;
  }

  // First estimate the event orientation (this could be separated into a fitter
  // and this classifier would then be seeded)
  TVector3 eventAxis(0.0, 0.0, 0.0);
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();

  vector <double> pmtTimes(fPMTData.size(), 0);
  for(vector<FitterPMT>::const_iterator iPMT = fPMTData.begin(); iPMT != fPMTData.end(); ++iPMT) {
    size_t index = iPMT - fPMTData.begin();
    pmtTimes[index] = iPMT->GetTime();
  }

  int pmtStart = static_cast<int>(0.1 * fPMTData.size());
  nth_element(pmtTimes.begin(), pmtTimes.begin() + pmtStart, pmtTimes.end());
  double windowStart = pmtTimes[pmtStart] - 10; // ns
  double windowEnd = pmtTimes[pmtStart] + 10; // ns
  double windowMid = pmtTimes[pmtStart];

  for(vector<FitterPMT>::const_iterator iPMT = fPMTData.begin(); iPMT != fPMTData.end(); ++iPMT) {
    double pmtTime = iPMT->GetTime();
    if(pmtTime >= windowStart && pmtTime <= windowEnd) {
      TVector3 pmtPos = pmtInfo.GetPosition(iPMT->GetID()).Unit();
      eventAxis += pmtPos;
    }
  }

  eventAxis = eventAxis.Unit();

  // Now populate an array (c.f. 2d pixel image) of PMTpos.Eventpos vs PMT time
  // Each pixel will be one input into the ANN
  vector<double> inputs = vector<double>(fNDots * fNTimes, 0.0);
  int nEntries = 0; // Need to normalise at end
  for(vector<FitterPMT>::const_iterator iPMT = fPMTData.begin(); iPMT != fPMTData.end(); ++iPMT) {
    const TVector3 pmtVector = pmtInfo.GetPosition(iPMT->GetID());
    double dotProduct = eventAxis.Dot(pmtVector.Unit());
    double deltaTime = iPMT->GetTime() - windowMid; // time since the 10% hit
    // Get the array bin
    int iDot = static_cast<int>((dotProduct - fDotLow) / fDotWidth);
    int iTime = static_cast<int>((deltaTime - fTimeLow)) / fTimeWidth;
    if(iDot >= 0 && iDot < fNDots && iTime >=0 && iTime < fNTimes) {
      inputs[iDot * fNTimes + iTime] += 1.0;
      nEntries++;
    }
  }

  // Normalise the inputs
  for(int i = 0; i < (fNTimes * fNDots); i++) {
    inputs[i] /= nEntries;
  }

  vector<double> outputs;
  // Loop over each level of the ANN, feed forward outputs
  for(size_t iLayer = 0; iLayer < fNLayers-1; iLayer++) {
    outputs = FeedForward(inputs, fWeights[iLayer], fBiases[iLayer]);
    // Update the inputs for the next iteration
    inputs = outputs;
  }

  // Output of final layer should be a one entry vector in range {0..1}
  // Scale this to the estimated radius
  double radiusEstimate = outputs[0] * fRadiusScale;

  DS::FitVertex fitVertex;
  eventAxis.SetMag(radiusEstimate);
  fitVertex.SetPosition(eventAxis);
  fFitResult.SetVertex(0, fitVertex);
  return fFitResult;

}


vector<double> PositionANN::FeedForward(const vector<double>& inputs, const vector< vector<double> >& weights, const vector<double>& biases) {
  // inputs: one per input
  // weights: list per input of one per output
  // biases: one per output
  // Note that we're using ReLU activation functions
  vector<double> outputs(biases.size(), 0);
  for(size_t iOutput = 0; iOutput < biases.size(); iOutput++) {
    double z = 0;
    for(size_t iInput = 0; iInput < weights.size(); iInput++) {
      z += (inputs[iInput] * weights[iInput][iOutput]);
    }
    z += biases[iOutput];
    // When training the output layer was = z, but relu ensures the radius is never negative
    outputs[iOutput] = max(0.0, z);
  }
  return outputs;
}