/* This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
 *
 * MAUS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MAUS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <algorithm>
#include <cmath>
#include <utility>
#include <functional>

// ROOT headers
#include "TLinearFitter.h"
#include "TMatrixTBase.h"
#include "TMath.h"
#include "Math/Functor.h"
#include "Fit/Fitter.h"
#include "Minuit2/Minuit2Minimizer.h"

// MAUS headers
#include "src/common_cpp/Recon/SciFi/RootFitter.hh"

#define PI 3.14159265

namespace RootFitter {

bool FitLineLinear(const std::vector<double>& x, const std::vector<double>& y,
                   const std::vector<double>& yerr, MAUS::SimpleLine& line, TMatrixD& cov_matrix) {

  TLinearFitter lf(1);
  lf.SetFormula("pol1");
  // Have to copy data into new elements as AssignData wants non-const arguments for some reason
  std::vector<double> x1 = x;
  std::vector<double> y1 = y;
  std::vector<double> yerr1 = yerr;
  lf.AssignData(x.size(), 1, &x1[0], &y1[0], &yerr1[0]);
  lf.Eval();

  TVectorD params;
  TVectorD errors;
  lf.GetParameters(params);
  lf.GetErrors(errors);

  line.set_c(params[0]);
  line.set_m(params[1]);
  line.set_c_err(errors[0]);
  line.set_m_err(errors[1]);
  line.set_chisq(lf.GetChisquare());
  lf.GetCovarianceMatrix(cov_matrix);

  return true;
}

bool FitCircleMinuit(const std::vector<double>& x, const std::vector<double>& y,
                     const std::vector<double>& err,
                     MAUS::SimpleCircle& circ, TMatrixD& cov_matrix) {

  auto Chi2Function = [&x, &y, &err](const double *par) {
    // Minimisation function computing the sum of squares of residuals
    // looping over the points
    double f = 0.0;
    for (size_t i = 0; i < x.size(); i++) {
        double u = (x[i] - par[0]);
        double v = y[i] - par[1];
        double dr = par[2] - std::sqrt(u*u+v*v);
        f += (dr*dr) / (err[i]*err[i]);
    }
    return f;
  };

  // Wrap chi2 function in a function object for the fit
  // 3 is the number of fit parameters (size of array par)
  ROOT::Math::Functor fcn(Chi2Function, 3);
  ROOT::Fit::Fitter fitter;
  double pStart[3] = {0, 0, 1};
  fitter.SetFCN(fcn, pStart);
  fitter.Config().ParSettings(0).SetName("x0");
  fitter.Config().ParSettings(1).SetName("y0");
  fitter.Config().ParSettings(2).SetName("R");

  // do the fit
  bool ok = fitter.FitFCN();
  if (!ok) {
    return ok;
  }
  const ROOT::Fit::FitResult & result = fitter.Result();
  result.Print(std::cout);

  // Create a circle object with the results
  circ = MAUS::SimpleCircle(result.Parameter(0), result.Parameter(1),
                            result.Parameter(2));
  circ.set_x0_err(result.Error(0));
  circ.set_y0_err(result.Error(1));
  circ.set_R_err(result.Error(2));
  // circ.set_chisq(result.Chi2());
  circ.set_chisq(result.MinFcnValue());
  circ.set_pvalue(result.Prob());
  result.GetCovarianceMatrix(cov_matrix);

  return true;
}

bool FitHelixMinuit(const std::vector<double>& x, const std::vector<double>& y,
                    const std::vector<double>& z, const std::vector<double>& err,
                    const double* pStart, const MinuitSZParameters& cfg,
                    MAUS::SimpleHelix& helix) {

  // Pull out x and y at the tracker reference station
  double x0 = x[0];
  double y0 = y[0];

  // This is a C++11 lambda function, which defines the quantity to be minimised by MINUIT,
  // using a sum over squared residuals looping over the seed spacepoints
  std::function<double(const double*)> //NOLINT(*)
      Chi2Function = [&x, &y, &z, &err, &x0, &y0](const double *par) {
    double xc = par[0];
    double yc = par[1];
    double rad = par[2];
//    double dsdz = par[3];
    double dphidz = par[3];

    double chisq = 0.0;
    for (size_t i = 0; i < x.size(); i++) {
        double theta = (z[i]*dphidz);
        double f = xc + (x0 - xc)*cos(theta) - (y0 - yc)*sin(theta); // x model prediction
        double g = yc + (y0 - yc)*cos(theta) + (x0 - xc)*sin(theta); // y model prediction
//        double theta = z[i]*dphidz + phi0;
//        double f = xc + rad*cos(theta); // x model prediction
//        double g = yc + rad*sin(theta); // y model prediction
        double dx = f - x[i];
        double dy = g - y[i];
        double dchisq = ((dx*dx) + (dy*dy)) / (err[i]*err[i]);
        chisq += dchisq;
    }
    return chisq;
  };

  // Set up the minimiser - ROOT::Math::Minimizer interface to MINUIT2, Migrad algorithm
  ROOT::Minuit2::Minuit2Minimizer min(ROOT::Minuit2::kMigrad);
  min.SetPrintLevel(0); // Keep Quiet!
  min.SetMaxFunctionCalls(cfg.maxFunctionCalls);
  min.SetMaxIterations(cfg.maxIterations);
  min.SetTolerance(cfg.tolerance);
  ROOT::Math::Functor fcn(Chi2Function, 4);
  min.SetFunction(fcn);

  // Determine the approximate dsdz value which corresponds to the global minimum of the chisq func
  double dphidz_seed = LocateGlobalChiSqMinimum(pStart, cfg, Chi2Function);

  // Set the variables to be minimized
  min.SetFixedVariable(0, "xc", pStart[0]);
  min.SetFixedVariable(1, "yc", pStart[1]);
  min.SetFixedVariable(2, "R",  pStart[2]);
  double step[4] = {cfg.stepSize, cfg.stepSize, cfg.stepSize, cfg.stepSize};
  min.SetLimitedVariable(3, "dphidz", dphidz_seed, 0.00000001, -0.018, 0.018);
//  min.SetLimitedVariable(4, "phi0", dsdz_seed, 0.000001, -3.14159265, 3.14159265);

  // Do the fit
  bool ok = min.Minimize();
  if (!ok) {
    // std::cerr << "Full helix fit did not converge\n";
    return ok;
  }
  const double *xs = min.X();
  const double *errors = min.Errors();

  // Calculate s0
  double phi = atan2((y0 - xs[1]), (x0 - xs[0])); // Use atan2 not atan, or lose half circle range
  if (phi < 0.0) {
    phi = phi + 2*PI;
  }
  double s0 = xs[2] * phi;

  // Pull out the covariance matrix
  TMatrixD cov(4, 4);
  for (size_t i = 0; i < 4; ++i) {
    for (size_t j = 0; j < 4; ++j) {
      cov(i, j) = min.CovMatrix(i, j);
    }
  }

  // Populate the helix
  std::vector<double> params = {xs[0], xs[1], xs[2], xs[2]*xs[3], s0};
  std::vector<double> errs = {errors[0], errors[1], xs[2]*errors[2], errors[3], 0.0};
  int zsize = static_cast<int>(z.size());
  std::vector<int> ndfs = {(2*zsize - 3), (2*zsize - 1), (2*zsize - 5)};
  std::vector<double> chisqs = {-1.0, min.MinValue(), -1.0};
  helix = MAUS::SimpleHelix(params, errs, ndfs, chisqs, -1.0, cov);
  return true;
}

double LocateGlobalChiSqMinimum(const double* pStart, const MinuitSZParameters& cfg,
                                std::function<double(const double*)> Chi2Function) { //NOLINT(*)
  double nsteps = 50.0;

  double lower_limit = -0.01125; // 80MeV/c
  double upper_limit = 0.01125;
  if (cfg.handedness == 1) {
    lower_limit = 0.003; // 300MeV/c
  } else if (cfg.handedness == -1) {
    upper_limit = -0.003;
  }

  double threshold = 100.0;
//  double step_size = (upper_limit - lower_limit) / nsteps;
  double step_size = 2.0e-4;
  double min_chisq = 1000.0;
  double dphidz_seed = pStart[3];

  std::vector<double> candidate_seeds;

  double params[4] = {pStart[0], pStart[1], pStart[2], lower_limit};
  double x_0 = Chi2Function(params);
  params[3] =  lower_limit+step_size;
  double x_1 = Chi2Function(params);
  params[3] =  lower_limit+step_size+step_size;
  double x_2 = Chi2Function(params);
  params[3] =  lower_limit+step_size+step_size+step_size;
  double x_3 = Chi2Function(params);

  std::vector<std::pair<double, double>> found_seeds;

  TMatrixD X(4, 3);
  X(0, 0) = 1.0;
  X(1, 0) = 1.0;
  X(2, 0) = 1.0;
  X(3, 0) = 1.0;

  TMatrixD Y(4, 1);

  double seed_0 = lower_limit;
  double seed_1 = lower_limit + step_size;
  double seed_2 = lower_limit + step_size + step_size;

  min_chisq = threshold;

  for (double seed = lower_limit+step_size*4;
                       seed < upper_limit + step_size + step_size + step_size; seed += step_size) {

    double seed_avg = 0.5*(x_1 + x_2);

    X(0, 1) = seed_0;
    X(1, 1) = seed_1;
    X(2, 1) = seed_2;
    X(3, 1) = seed;

    X(0, 2) = seed_0*seed_0;
    X(1, 2) = seed_1*seed_1;
    X(2, 2) = seed_2*seed_2;
    X(3, 2) = seed*seed;

    TMatrixD XT(TMatrixD::kTransposed, X);

    Y(0, 0) = x_0;
    Y(1, 0) = x_1;
    Y(2, 0) = x_2;
    Y(3, 0) = x_3;

    TMatrixD temp = (XT * X).Invert();
    TMatrixD beta = temp * XT * Y;

    double alpha = beta(2, 0);
    double minimum = - 0.5 * ( beta(1, 0) / beta(2, 0) );
    double value = beta(0, 0) + beta(1, 0)*minimum + beta(2, 0)*minimum*minimum;

    if ((value < threshold) && (minimum > seed_1) && (minimum <= seed_2)) {
      found_seeds.push_back(std::make_pair(minimum, value));
      if (value < min_chisq) {
        min_chisq = value;
        dphidz_seed = minimum;
      }
    }

    params[3] = seed;

    x_0 = x_1;
    x_1 = x_2;
    x_2 = x_3;
    x_3 = Chi2Function(params);

    seed_0 = seed_1;
    seed_1 = seed_2;
    seed_2 = seed;
  }

  return dphidz_seed;
}
}