#include <fstream.h>

#include "MTask.h"
#include "MParList.h"

#include "MDataChain.h"

class MChisqEval : public MTask
{
private:
    static const TString gsDefName;
    static const TString gsDefTitle;

    Double_t fChisq; //! Evaluated chi square

    MData   *fData0; // Data Member one (monte carlo data or chisq function)
    MData   *fData1; // Data Member two (measured data)

    // --------------------------------------------------------------------------
    //
    // Implementation of SavePrimitive. Used to write the call to a constructor
    // to a macro. In the original root implementation it is used to write
    // gui elements to a macro-file.
    //
    void StreamPrimitive(ofstream &out) const
    {
        out << "   MChisqEval " << GetUniqueName() << ";";
        if (fData0)
            out << "   " << GetUniqueName() << ".SetY1(\"" << fData0->GetRule() << "\");" << endl;
        if (fData1)
            out << "   " << GetUniqueName() << ".SetY1(\"" << fData1->GetRule() << "\");" << endl;
    }

    enum { kIsOwner = BIT(14) };

public:
    MChisqEval(const char *name=NULL, const char *title=NULL) : fData0(NULL), fData1(NULL)// : fMatrix(mat), fColumn(col), fEvalE(evale)
    {
        fName  = name  ? name  : gsDefName.Data();
        fTitle = title ? title : gsDefTitle.Data();
    }

    MChisqEval(MData *y1, const char *name=NULL, const char *title=NULL) : fData0(NULL), fData1(NULL)// : fMatrix(mat), fColumn(col), fEvalE(evale)
    {
        fName  = name  ? name  : gsDefName.Data();
        fTitle = title ? title : gsDefTitle.Data();
        SetY1(y1);
    }
    MChisqEval(MData *y1, MData *y2, const char *name=NULL, const char *title=NULL) : fData0(NULL), fData1(NULL)// : fMatrix(mat), fColumn(col), fEvalE(evale)
    {
        fName  = name  ? name  : gsDefName.Data();
        fTitle = title ? title : gsDefTitle.Data();
        SetY1(y1);
        SetY2(y2);
    }
    ~MChisqEval()
    {
        if (fData0 && (fData0->TestBit(kCanDelete) || TestBit(kIsOwner)))
            delete fData0;

        if (fData1 && (fData1->TestBit(kCanDelete) || TestBit(kIsOwner)))
            delete fData1;
    }

    void SetY1(MData *data)
    {
        // Set MC value
        if (fData0 && (fData0->TestBit(kCanDelete) || TestBit(kIsOwner)))
            delete fData0;
        fData0 = data;
        fData0->SetBit(kCanDelete);
        AddToBranchList(fData0->GetDataMember());
    }
    void SetY2(MData *data)
    {
        // Set measured/estimated value
        if (fData1 && (fData1->TestBit(kCanDelete) || TestBit(kIsOwner)))
            delete fData1;
        fData1 = data;
        fData1->SetBit(kCanDelete);
        AddToBranchList(fData1->GetDataMember());
    }
    void SetY1(const TString data) { SetY1(new MDataChain(data)); }
    void SetY2(const TString data) { SetY2(new MDataChain(data)); }

    void SetOwner(Bool_t o=kTRUE) { o ? SetBit(kIsOwner) : ResetBit(kIsOwner); }

    Bool_t PreProcess(MParList *plist)
    {
        fChisq = 0;

        if (!fData0)
            return kFALSE;

        if (!fData0->PreProcess(plist))
            return kFALSE;

        if (fData1)
            if (!fData1->PreProcess(plist))
                return kFALSE;

        return kTRUE;
    }

    Bool_t Process()
    {
        const Double_t y1 = fData0->GetValue();
        const Double_t y2 = fData1 ? fData1->GetValue() : 0;

        const Double_t dy  = y2-y1;
        const Double_t err = fData1 ? y1*y1 : 1;

        fChisq += dy*dy/err;
        return kTRUE;
    }

    Bool_t PostProcess()
    {
        fChisq /= GetNumExecutions();
        return kTRUE;
    }

    Double_t GetChisq() const { return fChisq; }
 
    ClassDef(MChisqEval, 0)
};

ClassImp(MChisqEval);

const TString MChisqEval::gsDefName  = "MChisqEval";
const TString MChisqEval::gsDefTitle = "Evaluate a chisq";

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

#include <TMinuit.h>

#include "MEvtLoop.h"
#include "MTaskList.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();
}

// --------------------------------------------------------------------------
#include <TStopwatch.h>

#include "MHMatrix.h"
#include "MEnergyEst.h"

#include "MDataMember.h"
#include "MDataElement.h"

#include "MReadTree.h"
#include "MMatrixLoop.h"

void estfit(Bool_t evalenergy=kFALSE)
{
    //
    // Fill events into a MHMatrix
    //
    MParList parlist;
    MHMatrix matrix;

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

    MEnergyEstParam eest;
    eest.Add("MHillasSrc");
    eest.InitMapping(&matrix);

    MReadTree read("Events", "gammas.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.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);
    fA[0] =  -2539; // [cm]
    fA[1] =    900; // [cm]
    fA[2] =   17.5; // [cm]
    fA[3] =  1;//    4;
    fA[4] =   58.3;

    TArrayD fB(7);
    fB[0] =    -8.64; // [GeV]
    fB[1] =   -0.069; // [GeV]
    fB[2] =  0.00066; // [GeV]
    fB[3] =    0.033; // [GeV]
    fB[4] = 0.000226; // [GeV]
    fB[5] =  4.14e-8; // [GeV]
    fB[6] =    -0.06;

    // 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)
            minuit.FixParameter(i);

        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)
            minuit.FixParameter(i+fA.GetSize());

        if (!rc)
            continue;

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

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

    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?fA.GetSize():0); i<(evalenergy?fA.GetSize()+fB.GetSize():fA.GetSize()); i++)
    {
        Double_t val;
        Double_t er;

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

        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);
     */
}