#include <TStopwatch.h>
#include <TMinuit.h>

#include "MParList.h"
#include "MTaskList.h"
#include "MEvtLoop.h"

#include "MReadTree.h"
#include "MHMatrix.h"
#include "MChisqEval.h"
#include "MMatrixLoop.h"
#include "MDataMember.h"
#include "MDataElement.h"
#include "MEnergyEstParam.h"

// --------------------------------------------------------------------------

void fcn(Int_t &npar, Double_t *gin, Double_t &f, Double_t *par, Int_t iflag)
{
    MEvtLoop *evtloop = (MEvtLoop*)gMinuit->GetObjectFit();

    MTaskList *tlist  = (MTaskList*)evtloop->GetParList()->FindObject("MTaskList"); //GetTaskList();

    MChisqEval      *eval = (MChisqEval*)     tlist->FindObject("MChisqEval");
    MEnergyEstParam *eest = (MEnergyEstParam*)tlist->FindObject("MEnergyEstParam");

    eest->SetCoeff(TArrayD(eest->GetNumCoeff(), par));

    evtloop->Eventloop();

    f = eval->GetChisq();
}

// --------------------------------------------------------------------------
//
//  0: fit impact parameter only
//  1: fit energy only
//  2: fit all parameters with respect to the energy resolution
//
void estfit(Int_t evalenergy=0)
{
    //
    // Fill events into a MHMatrix
    //
    MParList parlist;
    MHMatrix matrix;

    Int_t col = matrix.AddColumn(evalenergy?"MMcEvt.fEnergy":"MMcEvt.fImpact");

    MEnergyEstParam eest("Hillas");
    eest.Add("HillasSrc");
    eest.InitMapping(&matrix);

    MReadTree read("Events", "MC_ON2_short.root");
    read.DisableAutoScheme();

    if (!matrix.Fill(&parlist, &read))
        return;

    //
    // Setup the tasklist used to evaluate the needed chisq
    //
    MTaskList tasklist;
    parlist.AddToList(&tasklist);

    MMatrixLoop loop(&matrix);

    MChisqEval eval;
    eval.SetY1(new MDataElement(&matrix, col));
    eval.SetY2(new MDataMember(evalenergy ? "MEnergyEst.fEnergy" : "MEnergyEst.fImpact"));
    eval.SetOwner();

    tasklist.AddToList(&loop);
    tasklist.AddToList(&eest);
    tasklist.AddToList(&eval);

    MEvtLoop evtloop;
    evtloop.SetParList(&parlist);

    //
    // Be carefull: This is not thread safe
    //
    TMinuit minuit(12);
    minuit.SetPrintLevel(-1);
    minuit.SetFCN(fcn);

    // Ready for: minuit.mnexcm("SET ERR", arglist, 1, ierflg)
    if (minuit.SetErrorDef(1))
    {
        cout << "SetErrorDef failed." << endl;
        return;
    }

    //
    // Set initial values
    //
    TArrayD fA(5);
    TArrayD fB(7);
    fA[0] =  392957;
    fA[1] =  135;
    fA[2] =  -37449;
    fA[3] =  0.3464;
    fA[4] =  1;
    fB[0] =  1;
    fB[1] =  1;
    fB[2] =  1;
    fB[3] =  1;
    fB[4] =  1;
    fB[5] =  1;
    fB[6] =  1;
     /*
     fA[0] = 6122.97;
     fA[1] = 144.889;
     fA[2] = -601.256;
     fA[3] = 0.00171985;
     fA[4] = 116.451;
     fB[0] =  -2368.21;
     fB[1] =  1186.26;
     fB[2] =  0.147235;
     fB[3] =  144.49;
     fB[4] =  42.7681;
     fB[5] = -0.757817;
     fB[6] =  0.182482;
     */
    /*
     TArrayD fA(5);
    fA[0] = 6122.97;
    fA[1] = 144.889;
    fA[2] = -601.256;
    fA[3] = 0.00171985;
    fA[4] = 116.451;

    TArrayD fB(7);
    fB[0] =  -10445.5;
    fB[1] =  2172.05;
    fB[2] =  0.69;
    fB[3] =  491.2;
    fB[4] =  4.71444;
    fB[5] = -0.0331926;
    fB[6] = -0.014833;
    */
    // Set starting values and step sizes for parameters
    for (Int_t i=0; i<fA.GetSize(); i++)
    {
        TString name = "fA[";
        name += i;
        name += "]";
        Double_t vinit = fA[i];
        Double_t step  = fabs(fA[i]/3);

        Double_t limlo = 0; // limlo=limup=0: no limits
        Double_t limup = 0; 

        Bool_t rc = minuit.DefineParameter(i, name, vinit, step, limlo, limup);
        if (evalenergy==1)
        {
            minuit.FixParameter(i);
            cout << "Fixed Parameter #" << i << endl;
        }

        if (!rc)
            continue;

        cout << "Error in defining parameter #" << i << endl;
        return;
    }

    for (Int_t i=0; i<fB.GetSize(); i++)
    {
        TString name = "fB[";
        name += i;
        name += "]";
        Double_t vinit = fB[i];
        Double_t step  = fabs(fB[i]/3);

        Double_t limlo = 0; // limlo=limup=0: no limits
        Double_t limup = 0;

        Bool_t rc = minuit.DefineParameter(i+fA.GetSize(), name, vinit, step, limlo, limup);
        if (evalenergy==0)
        {
            minuit.FixParameter(i+fA.GetSize());
            cout << "Fixed Parameter #" << i+fA.GetSize() << endl;
        }

        if (!rc)
            continue;

        cout << "Error in defining parameter #" << i+fA.GetSize() << endl;
        return;
    }

    //
    // Setup globals used in FCN
    //
    minuit.SetObjectFit(&evtloop);

    cout << endl << "Fitting procedure running..." << endl;

    TStopwatch clock;
    clock.Start();

    // Now ready for minimization step: minuit.mnexcm("MIGRAD", arglist, 1, ierflg)
    gLog.SetNullOutput(kTRUE);
    Bool_t rc = minuit.Migrad();
    gLog.SetNullOutput(kFALSE);
    if (rc)
    {
        cout << "Migrad failed." << endl;
        return;
    }

    cout << endl;
    clock.Stop();
    clock.Print();
    cout << endl;

    for (Int_t i=(evalenergy==1?fA.GetSize():0); i<(evalenergy>0?fA.GetSize()+fB.GetSize():fA.GetSize()); i++)
    //for (Int_t i=0; i<fA.GetSize()+fB.GetSize(); i++)
    {
        Double_t val;
        Double_t er;

        if (!minuit.GetParameter(i, val, er))
            cout << "Error getting parameter #" << i << endl;

        cout << i << ":  " << val << "  +-  " << er << endl;
    }

    /*
     // Print results
     Double_t amin, edm, errdef;
     Int_t nvpar, nparx, icstat;
     minuit.mnstat(amin, edm, errdef, nvpar, nparx, icstat);
     minuit.mnprin(3, amin);
     */
}