#include <iostream>
#include <iomanip>
#include <sstream>
#include <fstream>
#include <string>
#include <cassert>

#include "TFile.h"
#include "TH1F.h"
#include "TMath.h"

#include "Math/Minimizer.h"
#include "Math/Factory.h"
#include "Math/Functor.h"

// Number of experimental spectra to fit and number of components
int global_num_histograms_to_fit;
int global_num_components;

// Global variables for experiment and simulation
TH1F* exp_histogram;
std::vector<TH1F*> sim_histograms;

// Global chi2 value and number of bins
double global_chi2;
double global_nBins;

double calculate_chi2(TH1F* data, std::vector<TH1F*> components, const double* parameters, double& bin_count){
  
  double chi2 = 0;
  bin_count = 0;

  // [1, NbinsX] inclusive to avoid underflow/overflow bins
  for (int i=1; i<=data->GetNbinsX(); i++){
    double obs = data->GetBinContent(i);
    double exp = 0;
    if (obs > 0){
      for (int j=0; j<global_num_components; j++){
	exp += parameters[j] * (components[j]->GetBinContent(i));
      }
      chi2 += (obs - exp) * (obs - exp) / obs;
      ++ bin_count;
    }
  }
  return chi2;
}

// The function that will be minimized
double fitting_function(const double* parameters){

  std::vector<double> bins(global_num_histograms_to_fit);
  std::vector<double> chi2_values(global_num_histograms_to_fit);

  // Find the local chi2 for each experimental spectrum being fit
  for (int i=0; i<global_num_histograms_to_fit; i++){
    chi2_values[i] = calculate_chi2(exp_histogram, sim_histograms, parameters, bins[i]);
  }

  // Sum the contibutions from each function (calculate the global chi2)
  global_chi2  = 0;
  global_nBins = 0;
  for (int i=0; i<global_num_histograms_to_fit; i++){
    global_chi2 += chi2_values[i];
    global_nBins += bins[i];
  }
  
  return global_chi2;

}

int fit_spectra(int input_num_histograms_to_fit,
		int input_num_components,
		std::string minimizer_name = "Minuit",
		std::string algorithm_name = "Migrad",
		int num_max_iterations = 1e9,
		int num_max_function_calls = 1e9,
		double tolerance = 0.01,
		double error_definition = 1.0,
		int print_level = 0){
  
  // Set global variables
  global_num_histograms_to_fit = input_num_histograms_to_fit;
  global_num_components = input_num_components;

  // Input ROOT files
  TFile* fInData = TFile::Open("experiment_histograms.root","READ");
  exp_histogram = (TH1F*) fInData->Get("experimental_histogram");

  TFile* fInSim = TFile::Open("simulation_histograms.root","READ");
  // The simulation component histograms
  for (int i=0; i<global_num_components; i++){
    sim_histograms.push_back((TH1F*) fInSim->Get(Form("component_%03d",i)));
  }
  
  // Initialize the minimizer and algorithm
  ROOT::Math::Minimizer* minimizer = ROOT::Math::Factory::CreateMinimizer(minimizer_name, algorithm_name);
  ROOT::Math::Functor minimization_function(&fitting_function, global_num_components);

  // Set minimizer properties
  minimizer->SetPrintLevel(print_level);
  minimizer->SetMaxIterations(num_max_iterations);
  minimizer->SetMaxFunctionCalls(num_max_function_calls);
  minimizer->SetTolerance(tolerance);
  minimizer->SetErrorDef(error_definition);
  minimizer->SetFunction(minimization_function);

  // Parameters for the minimizer
  std::vector<std::string> parameter_names;
  std::vector<double> start_values;
  std::vector<double> step_sizes;
  std::vector<double> lower_bounds;
  std::vector<double> upper_bounds;

  // Set parameter values
  for (int i=0; i<global_num_components; i++){
    parameter_names.push_back(Form("parameter_%03d",i));
    start_values.push_back(1.0/global_num_components); // Guess at true value
    step_sizes.push_back(0.001); // This is just a starting value as Minuit will adjust this if necessary
    lower_bounds.push_back(0.00000);
    upper_bounds.push_back(10.0);
  }

  // Set parameters
  for (int i=0; i<global_num_components; i++){
    minimizer->SetLimitedVariable(i, parameter_names[i].c_str(), start_values[i], step_sizes[i], lower_bounds[i], upper_bounds[i]);
  }

  // Minimize
  minimizer->Minimize();
 
  // Get results and save them
  const double* minimum_parameters = minimizer->X(); // Pointer to the parameters at the minimum
  const double* minimum_errors = minimizer->Errors(); // Pointer to the errors at the minimum
  
  // Normalize parameters
  double minimum_parameters_sum = 0.0;
  for (int i=0; i<global_num_components; i++){
    minimum_parameters_sum += minimum_parameters[i];
  }
  std::vector<double> normalized_minimum_parameters;
  for (int i=0; i<global_num_components; i++){
    normalized_minimum_parameters.push_back(minimum_parameters[i] / minimum_parameters_sum);
  }
  
  // Unnormalized parameters
  for (int i=0; i<global_num_components; i++){
    std::cout << "Par " << std::right << std::setfill('0') << std::setw(3) << i << "    " << std::left << std::setfill(' ') << std::setw(15) << minimum_parameters[i] << std::endl;
  }
  // Normalized parameters
  for (int i=0; i<global_num_components; i++){
    std::cout << "normalized Par " << std::right << std::setfill('0') << std::setw(3) << i << "    " << std::left << std::setfill(' ') << std::setw(15) << normalized_minimum_parameters[i] << std::endl;
  }

  // Finished
  minimizer->Clear();
  
  fInData->Close();
  delete fInData;

  fInSim->Close();
  delete fInSim;

  return 0;

}