/***************************************************************************

  PFitterFcnDKS.cpp

  Author: Andreas Suter
  e-mail: andreas.suter@psi.ch

***************************************************************************/

/***************************************************************************
 *   Copyright (C) 2007-2025 by Andreas Suter                              *
 *   andreas.suter@psi.ch                                                  *
 *                                                                         *
 *   This program 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 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program 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 this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <iostream>

#include "PMusr.h"
#include "PFitterFcnDKS.h"

//--------------------------------------------------------------------------
// Constructor
//--------------------------------------------------------------------------
/**
 * <p>Constructor.
 *
 * \param runList run list collection
 * \param useChi2 if true, a chisq fit will be performed, otherwise a log max-likelihood fit will be carried out.
 */
PFitterFcnDKS::PFitterFcnDKS(PRunListCollection *runList, const Bool_t useChi2, const UInt_t dksTag,
                             const std::string theo) :
  fTheoStr(theo),
  fUseChi2(useChi2),
  fRunListCollection(runList)
{
  fValid = false;

  if (fUseChi2)
    fUp = 1.0;
  else
    fUp = 0.5;

  InitDKS(dksTag);
}

//--------------------------------------------------------------------------
// Destructor
//--------------------------------------------------------------------------
/**
 * <p>Destructor
 */
PFitterFcnDKS::~PFitterFcnDKS()
{
//  FreeDKS();
}

//--------------------------------------------------------------------------
// operator()
//--------------------------------------------------------------------------
/**
 * <p>Minuit2 interface function call routine. This is the function which should be minimized.
 *
 * \param par a vector with all the parameters of the function
 */
Double_t PFitterFcnDKS::operator()(const std::vector<Double_t>& par) const
{
  Double_t value = 0.0, chisq = 0.0;

  // write parameter to GPU
  Int_t ierr = fDKS.writeParams(&par[0], par.size());

  // loop over all data sets
  PDKSParams dksp;
  Double_t norm = 1.0;
  // single histos
  for (UInt_t i=0; i<fRunListCollection->GetNoOfSingleHisto(); i++) {
    // get current values of N0, Nbkg, tau, the functions, the maps, the time resolution, the fit start time, etc.
    ierr = fRunListCollection->GetSingleHistoParams(i, par, dksp);

    // set N0, Nbkg
    ierr += fDKS.callSetConsts(dksp.fN0, dksp.fTau, dksp.fNbkg);

    // set fun values
    if (dksp.fFun.size() == 0)
      ierr += fDKS.writeFunctions(nullptr, 0);
    else
      ierr += fDKS.writeFunctions(&dksp.fFun[0], dksp.fFun.size());

    // set map values
    if (dksp.fMap.size() == 0)
      ierr += fDKS.writeMaps(nullptr, 0);
    else
      ierr += fDKS.writeMaps(&dksp.fMap[0], dksp.fMap.size());

    // calc chisq/log-mlh
    chisq = 0.0;
    ierr += fDKS.callLaunchChiSquare(FITTYPE_SINGLE_HISTO, fMemDataSingleHisto[i], fMemDataSingleHistoErr[i], dksp.fNoOfFitBins,
                                     par.size(), dksp.fFun.size(), dksp.fMap.size(),
                                     dksp.fStartTime , dksp.fPackedTimeResolution, chisq);
    value += chisq;    

    if (ierr != 0) {
      std::cerr << "PFitterFcnDKS::operator(): **ERROR** Kernel launch for single histo failed!" << std::endl;
      exit (EXIT_FAILURE);
    }
  }

  // asymmetries
  for (UInt_t i=0; i<fRunListCollection->GetNoOfAsymmetry(); i++) {
    // get current values of alpha and beta, the functions, the maps, the time resolution, the fit start time, etc.
    ierr = fRunListCollection->GetAsymmetryParams(i, par, dksp);

    // set alpha and beta
    ierr += fDKS.callSetConsts(dksp.fAlpha, dksp.fBeta);

    // set fun values
    if (dksp.fFun.size() == 0)
      ierr += fDKS.writeFunctions(nullptr, 0);
    else
      ierr += fDKS.writeFunctions(&dksp.fFun[0], dksp.fFun.size());

    // set map values
    if (dksp.fMap.size() == 0)
      ierr += fDKS.writeMaps(nullptr, 0);
    else 
      ierr += fDKS.writeMaps(&dksp.fMap[0], dksp.fMap.size());

    // calc chisq
    chisq = 0.0;
    ierr += fDKS.callLaunchChiSquare(FITTYPE_ASYMMETRY, fMemDataAsymmetry[i], fMemDataAsymmetryErr[i], dksp.fNoOfFitBins,
                                     par.size(), dksp.fFun.size(), dksp.fMap.size(),
                                     dksp.fStartTime , dksp.fPackedTimeResolution, chisq);
    value += chisq;

    if (ierr != 0) {
      std::cerr << "PFitterFcnDKS::operator(): **ERROR** Kernel launch for asymmetry failed!" << std::endl;
      exit (EXIT_FAILURE);
    }
  }

  // mu mius
  for (UInt_t i=0; i<fRunListCollection->GetNoOfMuMinus(); i++) {
    // get current values of N0, Nbkg, tau, the functions, the maps, the time resolution, the fit start time, etc.
    ierr = fRunListCollection->GetMuMinusParams(i, par, dksp);

    // set fun values
    ierr += fDKS.writeFunctions(&dksp.fFun[0], dksp.fFun.size());

    // set map values
    ierr += fDKS.writeMaps(&dksp.fMap[0], dksp.fMap.size());

    // calc chisq/log-mlh
    chisq = 0.0;
    ierr += fDKS.callLaunchChiSquare(FITTYPE_MU_MINUS, fMemDataMuMinus[i], fMemDataMuMinusErr[i], dksp.fNoOfFitBins,
                                     par.size(), dksp.fFun.size(), dksp.fMap.size(),
                                     dksp.fStartTime , dksp.fPackedTimeResolution, chisq);
    value += chisq;

    if (ierr != 0) {
      std::cerr << "PFitterFcnDKS::operator(): **ERROR** Kernel launch for mu minus failed!" << std::endl;
      exit (EXIT_FAILURE);
    }
  }

  if (dksp.fScaleN0AndBkg)
    norm = dksp.fPackedTimeResolution*1.0e3;

  return norm*value;
}

//--------------------------------------------------------------------------
// CalcExpectedChiSquare (public)
//--------------------------------------------------------------------------
/**
 * <p>Calculates the expected chisq, expected chisq per run, and chisq per run, if applicable.
 *
 * \param par
 * \param totalExpectedChisq expected chisq for all run blocks
 * \param expectedChisqPerRun expected chisq vector for all the run blocks
 */
void PFitterFcnDKS::CalcExpectedChiSquare(const std::vector<Double_t> &par, Double_t &totalExpectedChisq, std::vector<Double_t> &expectedChisqPerRun)
{
  // init expected chisq related variables
  totalExpectedChisq = 0.0;
  expectedChisqPerRun.clear();

  // only do something for chisq
  Double_t value = 0.0;
  if (fUseChi2) {
    // single histo
    for (UInt_t i=0; i<fRunListCollection->GetNoOfSingleHisto(); i++) {
      value = fRunListCollection->GetSingleRunChisqExpected(par, i); // calculate the expected chisq for single histo run block 'i'
      expectedChisqPerRun.push_back(value);
      totalExpectedChisq += value;
    }
    // mu minus
    for (UInt_t i=0; i<fRunListCollection->GetNoOfMuMinus(); i++) {
      value = fRunListCollection->GetSingleRunChisqExpected(par, i); // calculate the expected chisq for mu minus run block 'i'
      expectedChisqPerRun.push_back(value);
      totalExpectedChisq += value;
    }
  } else {
    // single histo
    for (UInt_t i=0; i<fRunListCollection->GetNoOfSingleHisto(); i++) {
      value = fRunListCollection->GetSingleRunMaximumLikelihoodExpected(par, i); // calculate the expected maxLH for single histo run block 'i'
      expectedChisqPerRun.push_back(value);
      totalExpectedChisq += value;
    }
    // mu minus
    for (UInt_t i=0; i<fRunListCollection->GetNoOfMuMinus(); i++) {
      value = fRunListCollection->GetSingleRunMaximumLikelihoodExpected(par, i); // calculate the expected maxLH for mu minus run block 'i'
      expectedChisqPerRun.push_back(value);
      totalExpectedChisq += value;
    }
  }
}

//--------------------------------------------------------------------------
// GetDeviceName (public)
//--------------------------------------------------------------------------
/**
 * <p> get from DKS the device name used.
 *
 * \param devName deivce name, if status == 0
 *
 * \return 0 if OK, 1 otherwise
 */
int PFitterFcnDKS::GetDeviceName(std::string &devName)
{
  int status = 1;
  
  devName = "??";
  if (!fValid)
    return status;
  
  status = fDKS.getDeviceName(devName);
  
  return status;
}

//--------------------------------------------------------------------------
// InitDKS (private)
//--------------------------------------------------------------------------
/**
 * <p>initializes the DKS interface
 *
 */
void PFitterFcnDKS::InitDKS(const UInt_t dksTag)
{
  Int_t ierr = 0;
  Int_t fitType = FITTYPE_UNDEFINED;
  Int_t startTimeBin = -1, endTimeBin = -1, length = -1;

  // if any device was allocated before, free the device resources
  FreeDKS();

  // select framework
  if (dksTag == DKS_GPU_CUDA)
    fDKS.setAPI("Cuda");
  else
    fDKS.setAPI("OpenCL");

  // select device
  if (dksTag == DKS_CPU_OPENCL)
    fDKS.setDevice("-cpu");
  else
    fDKS.setDevice("-gpu");

  // init device
  fDKS.initDevice();

  // init chisq buffer on the GPU

  // 1) calculate the maximum size for the data needed.
  Int_t parSize = -1, mapSize = -1, funSize = -1;
  Int_t maxSize = 0, size = -1;
  for (UInt_t i=0; i<fRunListCollection->GetNoOfSingleHisto(); i++) {
    size = fRunListCollection->GetNoOfBinsFitted(i);
    if (maxSize < size)
      maxSize = size;
  }
  for (UInt_t i=0; i<fRunListCollection->GetNoOfAsymmetry(); i++) {
    size = fRunListCollection->GetNoOfBinsFitted(i);
    if (maxSize < size)
      maxSize = size;
  }
  for (UInt_t i=0; i<fRunListCollection->GetNoOfMuMinus(); i++) {
    size = fRunListCollection->GetNoOfBinsFitted(i);
    if (maxSize < size)
      maxSize = size;
  }
  if (maxSize == 0) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get data size to be fitted." << std::endl;
    fValid = false;
    return;
  }

  // 2) get number of parameters / functions / maps
  parSize = fRunListCollection->GetNoOfParameters();
  if (parSize == -1) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get number of fit parameters." << std::endl;
    fValid = false;
    return;
  }
  funSize = fRunListCollection->GetNoOfFunctions();
  if (funSize == -1) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get number of functions." << std::endl;
    fValid = false;
    return;
  }
  mapSize = fRunListCollection->GetNoOfMaps();
  if (mapSize == -1) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get number of maps." << std::endl;
    fValid = false;
    return;
  }

  // now ready to init the chisq buffer on the GPU
  ierr = fDKS.initChiSquare(maxSize, parSize, funSize, mapSize);
  if (ierr != 0) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to allocate the necessary chisq buffer on the GPU." << std::endl;
    fValid = false;
    return;
  }

  // allocate memory for the data on the GPU/CPU and transfer the data sets
  PRunData *runData=0;

  // single histos
  fMemDataSingleHisto.resize(fRunListCollection->GetNoOfSingleHisto());
  fMemDataSingleHistoErr.resize(fRunListCollection->GetNoOfSingleHisto());
  for (UInt_t i=0; i<fRunListCollection->GetNoOfSingleHisto(); i++) {
    runData = fRunListCollection->GetSingleHisto(i);
    if (runData == 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get single histo data set (i=" << i << ") from fRunListCollection." << std::endl;
      fValid = false;
      return;
    }
    startTimeBin = fRunListCollection->GetStartTimeBin(MSR_FITTYPE_SINGLE_HISTO, i);
    endTimeBin   = fRunListCollection->GetEndTimeBin(MSR_FITTYPE_SINGLE_HISTO, i);
    length = endTimeBin-startTimeBin;
    if (startTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** startTimeBin undefind (single histo fit)." << std::endl;
      fValid = false;
      return;
    }
    if (endTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** endTimeBin undefind (single histo fit)." << std::endl;
      fValid = false;
      return;
    }
    if (length < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed for single histo data set (i=" << i << ") from fRunListCollection." << std::endl;
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** (endTimeBin=" << endTimeBin << ") < (startTimeBin=" << startTimeBin << ")." << std::endl;
      fValid = false;
      return;
    }
    fMemDataSingleHisto[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    fMemDataSingleHistoErr[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    if (ierr != 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to allocated single histo data set memory (i=" << i << ") on the GPU" << std::endl;
      fValid = false;
      return;
    }
    fDKS.writeData<double>(fMemDataSingleHisto[i], &runData->GetValue()->at(startTimeBin), length);
    fDKS.writeData<double>(fMemDataSingleHistoErr[i], &runData->GetError()->at(startTimeBin), length);
  }
  if (fRunListCollection->GetNoOfSingleHisto() > 0)
    fitType = FITTYPE_SINGLE_HISTO;

  // asymmetry
  fMemDataAsymmetry.resize(fRunListCollection->GetNoOfAsymmetry());
  fMemDataAsymmetryErr.resize(fRunListCollection->GetNoOfAsymmetry());
  for (UInt_t i=0; i<fRunListCollection->GetNoOfAsymmetry(); i++) {
    runData = fRunListCollection->GetAsymmetry(i);
    if (runData == 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get asymmetry data set (i=" << i << ") from fRunListCollection." << std::endl;
      fValid = false;
      return;
    }
    startTimeBin = fRunListCollection->GetStartTimeBin(MSR_FITTYPE_ASYM, i);
    endTimeBin   = fRunListCollection->GetEndTimeBin(MSR_FITTYPE_ASYM, i);
    length = endTimeBin-startTimeBin;
    if (startTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** startTimeBin undefind (asymmetry fit)." << std::endl;
      fValid = false;
      return;
    }
    if (endTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** endTimeBin undefind (asymmetry fit)." << std::endl;
      fValid = false;
      return;
    }
    if (length < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed for asymmetry data set (i=" << i << ") from fRunListCollection." << std::endl;
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** (endTimeBin=" << endTimeBin << ") < (startTimeBin=" << startTimeBin << ")." << std::endl;
      fValid = false;
      return;
    }
    fMemDataAsymmetry[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    fMemDataAsymmetryErr[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    if (ierr != 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to allocated asymmetry data set memory (i=" << i << ") on the GPU" << std::endl;
      fValid = false;
      return;
    }
    fDKS.writeData<double>(fMemDataAsymmetry[i], &runData->GetValue()->at(startTimeBin), length);
    fDKS.writeData<double>(fMemDataAsymmetryErr[i], &runData->GetError()->at(startTimeBin), length);
  }
  if (fRunListCollection->GetNoOfAsymmetry() > 0) {
    if (fitType == FITTYPE_UNDEFINED) {
      fitType = FITTYPE_ASYMMETRY;
    } else {
      std::cerr << ">>PFitterFcnDKS::InitDKS: **ERROR** mixed fit types found. This is currently not supported!" << std::endl;
      fValid = false;
      return;
    }
  }

  // mu minus
  fMemDataMuMinus.resize(fRunListCollection->GetNoOfMuMinus());
  fMemDataMuMinusErr.resize(fRunListCollection->GetNoOfMuMinus());
  for (UInt_t i=0; i<fRunListCollection->GetNoOfMuMinus(); i++) {
    runData = fRunListCollection->GetMuMinus(i);
    if (runData == 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to get mu minus data set (i=" << i << ") from fRunListCollection." << std::endl;
      fValid = false;
      return;
    }
    startTimeBin = fRunListCollection->GetStartTimeBin(MSR_FITTYPE_MU_MINUS, i);
    endTimeBin   = fRunListCollection->GetEndTimeBin(MSR_FITTYPE_MU_MINUS, i);
    length = endTimeBin-startTimeBin;
    if (startTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** startTimeBin undefind (mu minus fit)." << std::endl;
      fValid = false;
      return;
    }
    if (endTimeBin < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** endTimeBin undefind (mu minus fit)." << std::endl;
      fValid = false;
      return;
    }
    if (length < 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed for mu minus data set (i=" << i << ") from fRunListCollection." << std::endl;
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** (endTimeBin=" << endTimeBin << ") < (startTimeBin=" << startTimeBin << ")." << std::endl;
      fValid = false;
      return;
    }
    fMemDataMuMinus[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    fMemDataMuMinusErr[i] = fDKS.allocateMemory<Double_t>(length, ierr);
    if (ierr != 0) {
      std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to allocated mu minus data set memory (i=" << i << ") on the GPU" << std::endl;
      fValid = false;
      return;
    }
    fDKS.writeData<double>(fMemDataMuMinus[i], &runData->GetValue()->at(startTimeBin), length);
    fDKS.writeData<double>(fMemDataMuMinusErr[i], &runData->GetError()->at(startTimeBin), length);
  }
  if (fRunListCollection->GetNoOfMuMinus() > 0) {
    if (fitType == FITTYPE_UNDEFINED) {
      fitType = FITTYPE_MU_MINUS;
    } else {
      std::cerr << ">>PFitterFcnDKS::InitDKS: **ERROR** mixed fit types found. This is currently not supported!" << std::endl;
      fValid = false;
      return;
    }
  }

  // set the function string and compile the program
  ierr = fDKS.callCompileProgram(fTheoStr, !fUseChi2);
  if (ierr != 0) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** failed to compile theory!" << std::endl;
    fValid = false;
    return;
  }

  // checks device properties if openCL
  ierr = fDKS.checkMuSRKernels(fitType);
  if (ierr != 0) {
    std::cerr << ">> PFitterFcnDKS::InitDKS: **ERROR** muSR kernel checks failed!" << std::endl;
    fValid = false;
    return;
  }

  fValid = true;
}

//--------------------------------------------------------------------------
// FreeDKS (private)
//--------------------------------------------------------------------------
/**
 * <p>cleanup DKS/GPU memory
 *
 */
void PFitterFcnDKS::FreeDKS()
{
  PRunData *runData=0;
  // single histo
  for (UInt_t i=0; i<fMemDataSingleHisto.size(); i++) {
    runData = fRunListCollection->GetSingleHisto(i);
    fDKS.freeMemory<Double_t>(fMemDataSingleHisto[i], runData->GetValue()->size());
    fDKS.freeMemory<Double_t>(fMemDataSingleHistoErr[i], runData->GetValue()->size());
  }
  fMemDataSingleHisto.clear();
  fMemDataSingleHistoErr.clear();

  // asymmetry
  for (UInt_t i=0; i<fMemDataAsymmetry.size(); i++) {
    runData = fRunListCollection->GetAsymmetry(i);
    fDKS.freeMemory<Double_t>(fMemDataAsymmetry[i], runData->GetValue()->size());
    fDKS.freeMemory<Double_t>(fMemDataAsymmetryErr[i], runData->GetValue()->size());
  }
  fMemDataAsymmetry.clear();
  fMemDataAsymmetryErr.clear();

  // mu minus
  for (UInt_t i=0; i<fMemDataMuMinus.size(); i++) {
    runData = fRunListCollection->GetMuMinus(i);
    fDKS.freeMemory<Double_t>(fMemDataMuMinus[i], runData->GetValue()->size());
    fDKS.freeMemory<Double_t>(fMemDataMuMinusErr[i], runData->GetValue()->size());
  }
  fMemDataMuMinus.clear();
  fMemDataMuMinusErr.clear();
}