const int debug_level = 1;

#include <iostream>
#include <TFile.h>
#include <TH1F.h>
#include <TMath.h>
#include <TTree.h>
#include <TMatrixTSym.h>
#include <TMatrixD.h>
#include "hddm_fitter.h"
#define MAXRESIDS 1024
#define MAXPOINTS 2048
#define DOUBLE_NA -1.0E10
using namespace std;
using namespace TMath;

struct fitter_event {
  Int_t run;
  Int_t event;
  Int_t nCDC;
  Int_t nFDC;
  Int_t ntracks;
  Int_t status;
  Int_t iters;
  Int_t progress;
  Double_t chisq;
  Double_t chisqParam;
  Double_t xp0True;
  Double_t z0True;
  Double_t thetaTrue;
  Double_t phiTrue;
  Double_t ptinvTrue;
  Double_t xp0Fit;
  Double_t z0Fit;
  Double_t thetaFit;
  Double_t phiFit;
  Double_t ptinvFit;
  Double_t sigmaXp0;
  Double_t sigmaZ0;
  Double_t sigmaTheta;
  Double_t sigmaPhi;
  Double_t sigmaPtinv;
};

int main()
{

  TFile *file = new TFile("fitter_tree.root", "RECREATE");
  char hddmFileName[128] = "fitter.hddm";

  int count = 0;
  fitter_HDDM_t *thisInputEvent = 0;
  fitter_iostream_t *thisInputFile = 0;
  if (! (thisInputFile = open_fitter_HDDM(hddmFileName))) {
    printf("Error - could not open input file\n");
    exit(1);
  }

  TTree *tree = new TTree("fitter_tree", "event data from fitter");
  fitter_event fe;
  string treeDef = "run/I";
  treeDef += ":event";
  treeDef += ":nCDC";
  treeDef += ":nFDC";
  treeDef += ":ntracks";
  treeDef += ":status";
  treeDef += ":iters";
  treeDef += ":progress";
  treeDef += ":chisq/D";
  treeDef += ":chisqParam/D";
  treeDef += ":xp0True/D";
  treeDef += ":z0True/D";
  treeDef += ":thetaTrue/D";
  treeDef += ":phiTrue/D";
  treeDef += ":ptinvTrue/D";
  treeDef += ":xp0Fit/D";
  treeDef += ":z0Fit/D";
  treeDef += ":thetaFit/D";
  treeDef += ":phiFit/D";
  treeDef += ":ptinvFit/D";
  treeDef += ":sigmaXp0/D";
  treeDef += ":sigmaZ0/D";
  treeDef += ":sigmaTheta/D";
  treeDef += ":sigmaPhi/D";
  treeDef += ":sigmaPtinv/D";
  tree->Branch("event", &fe.run, treeDef.c_str());
  Int_t nResids;
  Int_t iDet[MAXRESIDS];
  Double_t resids[MAXRESIDS];
  Double_t dists[MAXRESIDS];
  Double_t docas[MAXRESIDS];
  Double_t xWire[MAXRESIDS];
  Double_t yWire[MAXRESIDS];
  Double_t zWire[MAXRESIDS];
  tree->Branch("nResids", &nResids, "nResids/I");
  tree->Branch("iDet", iDet, "iDet[nResids]/I");
  tree->Branch("resids", resids, "resids[nResids]/D");
  tree->Branch("dists", dists, "dists[nResids]/D");
  tree->Branch("docas", docas, "docas[nResids]/D");
  tree->Branch("xWire", xWire, "xWire[nResids]/D");
  tree->Branch("yWire", yWire, "yWire[nResids]/D");
  tree->Branch("zWire", zWire, "zWire[nResids]/D");
  Int_t nResidsThrown;
  Double_t residsThrown[MAXRESIDS];
  Double_t distsThrown[MAXRESIDS];
  Double_t docasThrown[MAXRESIDS];
  tree->Branch("nResidsThrown", &nResidsThrown, "nResidsThrown/I");
  tree->Branch("residsThrown", residsThrown, "residsThrown[nResidsThrown]/D");
  tree->Branch("distsThrown", distsThrown, "distsThrown[nResids]/D");
  tree->Branch("docasThrown", docasThrown, "docasThrown[nResids]/D");
  Int_t nPoints;
  Double_t pointsX[MAXPOINTS];
  Double_t pointsY[MAXPOINTS];
  Double_t pointsZ[MAXPOINTS];
  Int_t rings[MAXPOINTS];
  tree->Branch("nPoints", &nPoints, "nPoints/I");
  tree->Branch("pointsX", pointsX, "pointsX[nPoints]/D");
  tree->Branch("pointsY", pointsY, "pointsY[nPoints]/D");
  tree->Branch("pointsZ", pointsZ, "pointsZ[nPoints]/D");
  tree->Branch("rings", rings, "rings[nPoints]/I");

  unsigned int nparams = 0; // number of parameters for a trajectory
  fitter_Parameter_t* paramArray = NULL;
  int label = -1;
  int progress = -1;
  // loop over events
  while (thisInputEvent = read_fitter_HDDM(thisInputFile)) {
    count++; // increment event count
    progress = 0;
    nResids = 0;
    nResidsThrown = 0;
    nPoints = 0; // assume no as-swum trajectory
    fe.nCDC = -1; // should always be overwritten
    fe.nFDC = -1; // should always be overwritten
    fe.status = -99;
    fe.xp0True = -99.9;
    fe.z0True = -99.9;
    fe.thetaTrue = -99.9;
    fe.phiTrue = -99.9;
    fe.ptinvTrue = -99.9;
    fe.xp0Fit = -99.9;
    fe.z0Fit = -99.9;
    fe.thetaFit = -99.9;
    fe.phiFit = -99.9;
    fe.ptinvFit = -99.9;
    if (debug_level > 1) printf("count = %d\n", count);
    // event number
    fe.event = thisInputEvent->events->in[0].number;
    if (debug_level > 1) printf("event number = %d\n", fe.event);
    // track multiplicity
    if (thisInputEvent->events->in[0].tracks != HDDM_NULL) {
      progress++; // there are tracks
      fe.ntracks = thisInputEvent->events->in[0].tracks->mult;
      if (debug_level > 1) cout << "ntracks = " << fe.ntracks << endl;
      // cut on one and only one track
      if (fe.ntracks == 1) {
	progress++; // there is only one track
	fitter_Track_t* trackPtr = &(thisInputEvent->events->in[0].tracks->in[0]);
	// number of CDC hits on track
	fe.nCDC = trackPtr->nCDC;
	// number of FDC hits on track
	fe.nFDC = trackPtr->nFDC;
	// number of stored trajectories
	int ntrajs = trackPtr->trajectorys->mult;
	if (debug_level > 1) cout << "ntrajs = " << ntrajs << endl;
	// get pointer to CDC hits
	fitter_Cdcdata_t* cdcdataPtr;
	if (trackPtr->nCDC > 0) {
	  cdcdataPtr = thisInputEvent->events->in[0].cdcdata;
	} else {
	  cdcdataPtr = NULL;
	}
	fe.status = trackPtr->fitStatus;
	// final chi-squared
	fe.chisq = trackPtr->chisq;
	
	TMatrixTSym<double> cov(5);
	TMatrixD deltaCol(5,1);
	TMatrixD deltaRow(1,5);
	bool caseFitFound = false;
	bool caseThrownFound = false;
	bool covFound = false;
	for (int it = 0; it < ntrajs; it++) {
	  fitter_Trajectory_t* trajPtr = &(trackPtr->trajectorys->in[it]);
	  if (debug_level > 1) cout << "index = " << trajPtr->index << endl;
	  fitter_Fit_t* fitPtr = trajPtr->fit;
	  fe.iters = fitPtr->iterations;
	  switch(trajPtr->index) {
	  case 1:
	    // get trajectory with index 1, final fit
	    caseFitFound = true;
	    progress++; // final fit parameters produced
	    // get the final fit parameters
	    nparams = trajPtr->parameters->mult;
	    if (nparams != 5) {cout << "bad number of parameters\n"; exit(1);}
	    paramArray = trajPtr->parameters->in;
	    for (unsigned int ip = 0; ip < nparams; ip++) {
	      label = paramArray[ip].label;
	      switch (label) {
	      case 0:
		fe.xp0Fit = paramArray[ip].value;
		break;
	      case 1:
		fe.z0Fit = paramArray[ip].value;
		break;
	      case 2:
		fe.thetaFit = paramArray[ip].value;
		break;
	      case 3:
		fe.phiFit = paramArray[ip].value;
		break;
	      case 4:
		fe.ptinvFit = paramArray[ip].value;
		break;
	      }
	    }
	    
	    // number of residuals
	    nResids = trajPtr->residuals->mult;
	    if (debug_level > 1) cout << "nResids = " << nResids << endl;
	    // loop over residuals
	    for (int i = 0; i < nResids; i++) {
	      fitter_Residual_t* residualPtr = trajPtr->residuals->in + i;
	      // put residuals into branch
	      resids[i] = residualPtr->value;
	      iDet[i] = residualPtr->detector_index;
	      if (debug_level > 1) cout << i << " resid = " << resids[i] << "iDet = " << iDet[i] << endl;
	      // get CDC/FDC-specific info
	      if (iDet[i] == 0) {
		if (debug_level > 1) cout << "CDC hit\n";
		fitter_ResidInfoCdc_t* residInfoCdcPtr = residualPtr->residInfoCdc;
		dists[i] = residInfoCdcPtr->dist;
		docas[i] = residInfoCdcPtr->doca;
		xWire[i] = residInfoCdcPtr->xWire;
		yWire[i] = residInfoCdcPtr->yWire;
		zWire[i] = residInfoCdcPtr->zWire;
	      } else if (iDet[i] == 1) {
		if (debug_level > 1) cout << "FDC hit\n";
		fitter_ResidInfoFdc_t* residInfoFdcPtr = residualPtr->residInfoFdc;
		docas[i] = residInfoFdcPtr->doca;
		dists[i] = DOUBLE_NA;
		xWire[i] = DOUBLE_NA;
		yWire[i] = DOUBLE_NA;
		zWire[i] = DOUBLE_NA;
	      } else {
		if (debug_level > 1) cout << "neither, cannot happen\n";
		return(1);
	      }
	    }
	    // check for good fit pointer
	    if (debug_level > 1) cout << "fit trajPtr->fit = " << trajPtr->fit << endl;
	    if (fitPtr != HDDM_NULL) {
	      if (debug_level > 1) cout << "good pointer\n";
	      fitter_CovarianceMatrix_t* covMatPtr = fitPtr->covarianceMatrix;
	      if (covMatPtr != HDDM_NULL) {
		if (debug_level > 1) cout << "good cov mat\n";
		fitter_CovarianceElements_t* covMatElsPtr = covMatPtr->covarianceElements;
		if (covMatElsPtr != HDDM_NULL) {
		  if (debug_level > 1) cout << "good elements\n";
		  int nel = covMatElsPtr->mult;
		  if (nel != 15) {
		    int error = 7;
		    throw error;
		  }
		  covFound = true;
		  int row, col;
		  for (int i = 0; i < nel; i++) {
		    fitter_CovarianceElement_t* elPtr = covMatElsPtr->in + i;
		    if (debug_level > 1) cout << elPtr->row << " " << elPtr->column << " " << elPtr->value << endl;
		    row = elPtr->row;
		    col = elPtr->column;
		    cov(row - 1, col - 1) = elPtr->value;
		    if (row != col) cov(col - 1, row - 1) = elPtr->value;
		  }
		  for (int i = 0; i < nparams; i++) {
		    for (int j = 0; j < nparams; j++) {
		      if (debug_level > 1) cout << "cov mat " << i << " " << j << " " << cov(i, j) << endl;
		    }
		  }
		}
	      }
	    }
	    break;
	  case 2:
	    // get trajectory with index 2, thrown parameters
	    caseThrownFound = true;
	    if (debug_level > 1) cout << "index = " << trajPtr->index << endl;
	    // get the thrown parameters
	    nparams = trajPtr->parameters->mult;
	    if (nparams != 5) {cout << "bad number of parameters\n"; exit(1);}
	    paramArray = trajPtr->parameters->in;
	    for (unsigned int ip = 0; ip < nparams; ip++) {
	      label = paramArray[ip].label;
	      switch (label) {
	      case 0:
		fe.xp0True = paramArray[ip].value;
		break;
	      case 1:
		fe.z0True = paramArray[ip].value;
		break;
	      case 2:
		fe.thetaTrue = paramArray[ip].value;
		break;
	      case 3:
		fe.phiTrue = paramArray[ip].value;
		break;
	      case 4:
		fe.ptinvTrue = paramArray[ip].value;
		break;
	      }
	    }
	    // number of residuals
	    nResidsThrown = trajPtr->residuals->mult;
	    if (debug_level > 1) cout << "nResidsThrown = " << nResidsThrown
				      << endl;
	    // loop over residuals
	    for (int i = 0; i < nResidsThrown; i++) {
	      fitter_Residual_t* residualPtr = trajPtr->residuals->in + i;
	      // put residuals into branch
	      residsThrown[i] = residualPtr->value;
	      if (debug_level > 1) cout << i << " resid = " << residsThrown[i]
					<< endl;
	      // check for CDC info
	      fitter_ResidInfoCdc_t* residInfoCdcPtr = residualPtr->residInfoCdc;
	      if (residInfoCdcPtr) {
		distsThrown[i] = residInfoCdcPtr->dist;
		docasThrown[i] = residInfoCdcPtr->doca;
	      } else {
		distsThrown[i] = -1.0;
		docasThrown[i] = -1.0;
	      }
	    }
	    break;
	  case 3:
	    // get trajectory with index 3, mc swum parameters
	    if (debug_level > 1) cout << "index = " << trajPtr->index << endl;
	    // number of points
	    nPoints = trajPtr->points->mult;
	    if (debug_level > 1) cout << "nPoints = " << nPoints
				      << endl;
	    // loop over points
	    for (int i = 0; i < nPoints; i++) {
	      fitter_Point_t* pointPtr = trajPtr->points->in + i;
	      // put points into branch
	      pointsX[i] = pointPtr->x;
	      pointsY[i] = pointPtr->y;
	      pointsZ[i] = pointPtr->z;
	      // not every point will have a ring
	      if (i < trackPtr->nCDC) {
		rings[i] = cdcdataPtr->cdchits->in[i].ring;
		if (debug_level > 1) cout << i << " ring = " << rings[i]
					  << endl;
	      }
	    }
	    break;
	  }
	  if (caseFitFound && caseThrownFound && covFound) {
	    fe.sigmaXp0 = sqrt(cov(0,0));
	    fe.sigmaZ0 = sqrt(cov(1,1));
	    fe.sigmaTheta = sqrt(cov(2,2));
	    fe.sigmaPhi = sqrt(cov(3,3));
	    fe.sigmaPtinv = sqrt(cov(4,4));
	    deltaCol(0,0) = fe.xp0Fit - fe.xp0True; 
	    deltaCol(1,0) = fe.z0Fit - fe.z0True; 
	    deltaCol(2,0) = fe.thetaFit - fe.thetaTrue; 
	    deltaCol(3,0) = fe.phiFit - fe.phiTrue; 
	    deltaCol(4,0) = fe.ptinvFit - fe.ptinvTrue;
	    for (int i = 0; i < 5; i++) {
	      if (debug_level > 1) cout << i << " " << deltaCol(i,0) << endl;
	      deltaRow(0,i) = deltaCol(i,0);
	    }		    
	    fe.chisqParam = (deltaRow * cov.Invert() * deltaCol)(0,0);
	    if (debug_level > 1) cout << "chisq = " << fe.chisq << " chisqParam = " << fe.chisqParam << endl;
	  }
	}
      }
    }
    fe.progress = progress;
    tree->Fill();
    if (debug_level > 1) cout << "progress = " << progress << endl;
    if (debug_level > 1) cout << "flush the event" << endl;
    flush_fitter_HDDM(thisInputEvent,0);
  }
  close_fitter_HDDM(thisInputFile);
  
  tree->Print();
  file->Write();
  
  return 0;
  
}