#include "km3net-dataformat/online/JDAQ.hh"
#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JMath/JMatrixNS.hh"
#include "Jeep/JParser.hh"

#include "TF2.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TMath.h"
#include "TMatrixDSym.h"
#include "TROOT.h"
#include "TVectorD.h"

#include <string>
#include <iostream>
#include <set>

namespace FITL1DT {

  using namespace JMATH;
  using namespace std;


  JMatrixNS removeRowColumn(const JMatrixNS& A, int C) {
    const int N = A.size() - 1;
    JMatrixNS ret(N);
    int k = 0;
    for (int i = 0; i < N+1; i++) {
      if (i == C) { 
        continue; 
      }
      int l = 0;
      for (int j = 0; j < N+1; j++) {
        if (j == C) { 
          continue; 
        }
        ret.at(k,l) = A.at(i,j);
        l++;
      }
      k++;
    }
    return ret;
  }

  // inner prod exists for root matrices: y = M * x iff the sizes of the matrix-vector couple are the same (mxn) * (n) = (m) 
  vector<double> dotprod(const JMatrixNS& A, const vector<double>& x) {
    if (A.size() != x.size()) { 
      cout << "Can not evaluate dot product: dimensions dont match" << endl; 
      return vector<double>(0, 0.0); 
    }
    vector<double> ret(A.size(), 0.0);
    for (unsigned int i = 0; i < A.size(); ++i) {
      for (unsigned int j = 0; j < A.size(); ++j) {
        ret.at(i) += A.at(i,j) * x.at(j);
      }
    }
    return ret;
  }
}


/**
 * \file
 *
 * Auxiliary program to find the detector time offsets from FitL1dtSlices output
 * 
 * \author kmelis, lnauta
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  using namespace FITL1DT;

  string inputFile;
  string outputFile;
  string detectorFile;
  bool   overwriteDetector;
  bool   noInterDU;
  int    maxFloors;
  string option;
  int    idom_ref;
  int    debug;

  try { 

    JParser<> zap("Program to calculate time offsets from FitL1dtSlices output");
    
    zap['f'] = make_field(inputFile,         "input file")                   = "fitslices.root";
    zap['o'] = make_field(outputFile,        "output file")                  = "fit.root";
    zap['a'] = make_field(detectorFile,      "detector file");
    zap['A'] = make_field(overwriteDetector, "overwrite the input detector file");
    zap['D'] = make_field(noInterDU,         "do not use inter DU information in the calculation");
    zap['r'] = make_field(idom_ref,          "reference DOM set to t=0")     = 9;
    zap['F'] = make_field(maxFloors,         "number of floors used in fit") = 9999;
    zap['d'] = make_field(debug) = 0;

    if (zap.read(argc, argv) != 0) {
      return 1;
    }
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }


  JDetector detector;

  try {
    load(detectorFile, detector);
  }
  catch(const JException& error) {
    FATAL(error);
  }

  set<int> DU_IDs;
  for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {
    DU_IDs.insert(moduleA->getString());
    if (!noInterDU) {
      break;
    }
  }

  TFile* in = TFile::Open(inputFile.c_str(), "exist");

  if (in == NULL || !in->IsOpen()) {
    FATAL("File: " << inputFile << " not opened." << endl);
  }

  // Get fitting histograms
  TH2D* h2A = (TH2D*)in->Get("Aij");
  TH2D* h2B = (TH2D*)in->Get("Bij");

  if (h2A == NULL || h2B == NULL) {
    FATAL("File does not contain histogram with matrix. Please use FitL1dtSlices script first." << endl);
  }

  TFile out(outputFile.c_str(), "recreate");

  h2A->Write();
  h2B->Write();

  vector<double> dt0(detector.size(), 0.0);
  vector<double> vart0(detector.size(), 0.0);
  vector<int> status(detector.size(), 0);

  for (set<int>::const_iterator DU_ID = DU_IDs.begin(); DU_ID != DU_IDs.end(); ++DU_ID) {
    JMatrixNS fitA(detector.size());
    vector<double> fitB(detector.size());
    vector<double> diag(detector.size());

    // Fill the matrix and vector
    for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {
      const int idom = distance(detector.begin(), moduleA);
      const int ibin = h2A->GetXaxis()->FindBin(TString(Form("%i", moduleA->getID())));

      if (noInterDU && (moduleA->getString() != *DU_ID)) {
        continue;
      }

      for (JDetector::iterator moduleB = detector.begin(); moduleB != detector.end(); ++moduleB) {
        const int jdom = distance(detector.begin(), moduleB);
        const int jbin = h2A->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID())));

        if (idom == jdom) {
          continue;
        }

        if (noInterDU && (moduleB->getString() != *DU_ID)) {
          continue;
        }

        if (fabs(moduleA->getFloor() - moduleB->getFloor()) > maxFloors) {
          continue;
        }

        double Aij = h2A->GetBinContent(ibin, jbin);
        double Bij = h2B->GetBinContent(ibin, jbin);

        // If the parabola does not have a maximum: Aij > 0, this means the fit in FitL1dtSlices was unsuccesful.
        if (Aij >= 0.0) {
          continue;
        }

        NOTICE("  " << idom << " " << jdom << "  " << moduleA->getID() << " " << moduleB->getID() << " Aij: " << Aij << " Bij: " << Bij << " bf: " << -0.5*Bij/Aij << endl);

        fitA.at(idom, jdom)  = -4.0*Aij;
        fitA.at(idom, idom) +=  4.0*Aij;
        fitB[idom]          +=  2.0*Bij; // Here Bij = -2*Aij*Bij in the mathematical derivation due to the form of the TF1 in FitL1Slices: [0]x^2+[1]x+[2] instead of [0](x-[1])^2+[2]
        diag[idom]          +=  4.0*Aij; // Wilks theorem: d^2L/dp^2 = 1/var(p), from this  the diagonal contains the second derivative of logL.
      }
    }

    // remove empty columns and rows (no data and thus no time info)
    unsigned int icolumn = 0;
    int idom = 0;
    vector<bool> changet0(detector.size(), true);
    while (icolumn < fitA.size()) {
      if (fitA.at(icolumn, icolumn) == 0) {
        fitA = removeRowColumn(fitA, icolumn);
        fitB.erase(fitB.begin() + icolumn);

        changet0[idom] = false;
      } else {
        icolumn++;
      }
      idom++;
    }

    // no data for this string
    if (fitA.size() <= 1) {
      continue;
    }

    // Use ROOT matrices from here to have more options 
    // runtime is not important for this program and the matrix functionality defined before main() is not trivial to port to ROOT.

    TMatrixDSym fit_A(fitA.size());
    fit_A.SetTol(1e-26); // Det is near 1E-40, so set the tolerance 10 orders lower for safety.
    TVectorD fit_B(fitB.size());

    for (unsigned int i = 0; i < fitA.size(); i++) {
      for (unsigned int j = 0; j < fitA.size(); j++) {
        fit_A(i,j) = fitA.at(i,j);
      }
      fit_B(i) = fitB[i];
    }

    // fix reference DOM time offset (the first DOM with data, i.e. with off-diagonal elements > 0.0)
    for (int i = 0; i < fit_A.GetNrows(); ++i) {
      fit_A(idom_ref-1, i) = 0.0;
      fit_A(i, idom_ref-1) = 0.0;
    }
    fit_A(idom_ref-1, idom_ref-1) = 1.0;
    fit_B(idom_ref-1) = 0.0;

    if (debug >= 2) fit_A.Print();

    // Fit best time offsets via matrix inversion
    Double_t fit_A_det = -999; // dont initialize to 0. det == 0 means the matrix is non-invertable.
    fit_A.Invert(&fit_A_det);
    if (fit_A_det == 0) {
      NOTICE("Matrix not invertible." << endl);
    }

    if (debug >= 2) fit_A.Print();

    // dt = A^-1*B
    TVectorD dt_i = fit_A * fit_B;

    // This shape for writing out the dt0 is chose because row+columns are removed when they have a diagonal element of 0, i.e. make the matrix singular.
    int k = 0;
    for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {
      const int idom = distance(detector.begin(), moduleA);
      if (changet0[idom]) {
        dt0[idom] = dt_i(k);
        vart0[idom] = TMath::Sqrt(-1 / diag[idom]);
        status[idom] = 2;   // fitted
        if (k == 0) {
          status[idom] = 1;
        }  // fixed t0
        k++;
      }
    }
  }

  TH1D h_dt0("dt0", "dt0", detector.size(), -0.5, detector.size()-0.5);

  for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {
    const int idom = distance(detector.begin(), moduleA);
    NOTICE("Changing t0 of DOM " << moduleA->getID() << " (S"<<setw(2) << moduleA->getString() << "F" << setw(2) << moduleA->getFloor() << ") by " << setw(13) << dt0[idom] << " [ns] ");

    if (status[idom] == 0) {
      NOTICE(" ** unable to fit **" << endl);
    }
    else if (moduleA->getFloor() == idom_ref) {
      NOTICE(" ** fixed **" << endl);
    }
    else {
      NOTICE(endl);
    }

    if (overwriteDetector) {
      moduleA->add(dt0[idom]); // minus sign difference with JHobbit.cc
    }
    h_dt0.GetXaxis()->SetBinLabel(idom+1, Form("%i", moduleA->getID()));
    h_dt0.SetBinContent(idom+1, dt0[idom]);
    h_dt0.SetBinError(idom+1, vart0[idom]);
  }

  h_dt0.Write();

  out.Close();

  if (overwriteDetector) {
    NOTICE("Store calibration data on file " << detectorFile << endl);

    store(detectorFile, detector);
  }
}