#include "MapsTrackManager.hh"
#include "MapsTrack.hh"
#include "MapsSensor.hh"
#include "TFile.h"
#include "TTree.h"
#include "Rtypes.h"
#include "TBranch.h"
#include <map>
#include <vector>
#include <string>
#include "ToString.h"
#include <cassert>

MapsTrackManager::MapsTrackManager(std::vector<MapsSensor*> sensors) {
	mySensors = sensors;
	for (std::vector<MapsSensor*>::const_iterator it = sensors.begin(); it
			!= sensors.end(); it++)
		mySensorIds.push_back((*it)->id());
	myRequiredLeft = 0;
	myRequiredRight = 0;
}
MapsTrackManager::MapsTrackManager() {
	myRequiredLeft = 0;
	myRequiredRight = 0;

}
MapsTrackManager::~MapsTrackManager() {

}
void MapsTrackManager::addTrack(MapsTrack* mt) {
	myTracks.push_back(mt);
}
std::vector<MapsTrack*>& MapsTrackManager::getTracks() {
	return myTracks;
}
std::vector<MapsSensor*>& MapsTrackManager::getSensors() {
	return mySensors;
}

void MapsTrackManager::addSensor(MapsSensor* s) {
	mySensors.push_back(s);
	mySensorIds.push_back(s->id());
}

void MapsTrackManager::setSensors(std::vector<MapsSensor*> sensors) {
	mySensors = sensors;
}

void MapsTrackManager::setRequiredLeftSensor(MapsSensor* left) {
	myRequiredLeft = left;
}

void MapsTrackManager::setRequiredRightSensor(MapsSensor* right) {
	myRequiredRight = right;
}

void MapsTrackManager::recreateFromRootFile(char* rootFileName)
		throw (MapsException) {
	TFile* file = new TFile(rootFileName);
	if (file == 0) {
		std::string desc(__PRETTY_FUNCTION__);
		desc.append(": Couldn't open root file!\n");
		MapsException me(desc.c_str());
		throw me;
	}

	TTree* tree = (TTree*) file->Get("mapsTracks");

	//Extract values from the tree
	UInt_t sensorIds[4];
	Double_t sensorZs[4];
	Double_t sensorPhis[4];
	Double_t xAlign[4];
	Double_t yAlign[4];

	TBranch* sensorBr = tree->GetBranch("sensorIds");
	sensorBr->SetAddress(sensorIds);
	TBranch* sensorZBr = tree->GetBranch("sensorZs");
	sensorZBr->SetAddress(sensorZs);
	TBranch* sensorPhiBr = tree->GetBranch("sensorPhis");
	sensorPhiBr->SetAddress(sensorPhis);
	TBranch* sensorXAl = tree->GetBranch("xAlign");
	sensorXAl->SetAddress(xAlign);
	TBranch* sensorYAl = tree->GetBranch("yAlign");
	sensorYAl->SetAddress(yAlign);

	sensorBr->GetEntry(0);
	sensorZBr->GetEntry(0);
	sensorPhiBr->GetEntry(0);
	sensorXAl->GetEntry(0);
	sensorYAl->GetEntry(0);

	//create sensors
	std::vector<MapsSensor*> sensors;
	std::map<unsigned, MapsSensor*> sensorsAndKeys;
	std::vector<unsigned> ids;

	for (unsigned k(0); k < 4; k++) {
		MapsSensor::physXY align(xAlign[k], yAlign[k]);
		MapsSensor* s = new MapsSensor(sensorIds[k], sensorZs[k], align, sensorPhis[k]);
		//MapsSensor* s = new MapsSensor(sensorIds[k], sensorZs[k], align);
		sensors.push_back(s);
		ids.push_back(sensorIds[k]);
		sensorsAndKeys[sensorIds[k]] = s;
		std::cout << "Recreated sensor:\t" << *s << "\n";
	}
	mySensors = sensors;
	mySensorIds = ids;

	std::cout << "Recreated sensors successfully!\n";

	//Now, let's make some tracks :-)
	UInt_t bx;
	UInt_t fourthSensor;
	UInt_t fourthSensorThresh;
	UInt_t nHits;
	UInt_t sensorHit[4];
	UInt_t xHit[4];
	UInt_t yHit[4];

	TBranch* bxBr = tree->GetBranch("bx");
	bxBr->SetAddress(&bx);

	TBranch* fourthSensorBr = tree->GetBranch("fourthSensor");
	fourthSensorBr->SetAddress(&fourthSensor);

	TBranch* fourthSensorThreshBr = tree->GetBranch("fourthSensorThresh");
	fourthSensorThreshBr->SetAddress(&fourthSensorThresh);

	TBranch* nHitsBr = tree->GetBranch("nHits");
	nHitsBr->SetAddress(&nHits);

	TBranch* sensorHitBr = tree->GetBranch("sensorHit");
	sensorHitBr->SetAddress(sensorHit);

	TBranch* xHitBr = tree->GetBranch("xHit");
	xHitBr->SetAddress(xHit);

	TBranch* yHitBr = tree->GetBranch("yHit");
	yHitBr->SetAddress(yHit);
	std::cout << "Created branches.\n";

	double chiXSum(0);
	std::vector<MapsTrack*> tracks;
	std::cout << "Tree has " << tree->GetEntries() << " entries\n";
	for (unsigned entries(0); entries < tree->GetEntries(); entries++) {
		tree->GetEntry(entries);
		std::map<MapsSensor*, MapsTrack::coord> hits;

		for (unsigned k(0); k < nHits; k++) {
			MapsSensor* relevantS = sensorsAndKeys[sensorHit[k]];

			MapsTrack::coord* relevantHit = new MapsTrack::coord(xHit[k], yHit[k]);
			hits[relevantS] = *relevantHit;
			//std::cout << "In: \t" << *relevantS << ": " << *relevantHit << "\n";
		}
		MapsTrack* t = new MapsTrack(bx, hits, sensorsAndKeys[fourthSensor], fourthSensorThresh);
		//diagnose(std::cout, *t);

		tracks.push_back(t);
		chiXSum += t->chiSq(0);
	}
	std::cout << "ChiXSum coming in: " << chiXSum << "\n";
	myTracks = tracks;
	std::cout << "Initialised myTracks with " << myTracks.size()
			<< " entries.\n";
}

void MapsTrackManager::exportToRootFile(char* rootFileName) {
	TFile* file = new TFile(rootFileName, "recreate");
	TTree* tree = new TTree("mapsTracks", "Maps Tracks");
	std::cout << "Created file and tree...\n";
	
	UInt_t sensorIds[4];
	Double_t sensorZs[4];
	Double_t sensorPhis[4];
	Double_t xAlign[4];
	Double_t yAlign[4];
	UInt_t bx;

	UInt_t fourthSensor;
	UInt_t fourthSensorThresh;

	UInt_t nHits;
	UInt_t sensorHit[4];
	UInt_t xHit[4];
	UInt_t yHit[4];
	Double_t physX[4];
	Double_t physY[4];

	//these guys have to be variable length arrays since there isn't always a fourth hit defined
	//so we can't save it in that case.
	Double_t residFourthX[1];
	Double_t residFourthY[1];
	Double_t chiSqProbX3Hit[1];
	Double_t chiSqProbY3Hit[1];

	UInt_t fourthHitPresent(0);

	Double_t chiXProb, chiYProb;
	Double_t p0, p1, q0, q1;

	Double_t chiX, chiY, theta;
	Double_t meanX, meanY;

	std::cout << "Defining branches...\n";
	tree->Branch("sensorIds", sensorIds, "sensorIds[4]/I");
	tree->Branch("sensorZs", sensorZs, "sensorZs[4]/D");
	tree->Branch("sensorPhis", sensorPhis, "sensorPhis[4]/D");
	tree->Branch("xAlign", xAlign, "xAlign[4]/D");
	tree->Branch("yAlign", yAlign, "yAlign[4]/D");

	tree->Branch("bx", &bx, "bx/I");
	tree->Branch("fourthSensor", &fourthSensor, "fourthSensor/I");
	tree->Branch("fourthSensorThresh", &fourthSensorThresh,
			"fourthSensorThresh/I");

	tree->Branch("nHits", &nHits, "nHits/I");
	tree->Branch("fourthHitPresent", &fourthHitPresent, "fourthHitPresent/I");

	tree->Branch("chiX", &chiX, "chiX/D");
	tree->Branch("chiY", &chiY, "chiY/D");
	tree->Branch("chiXProb", &chiXProb, "chiXProb/D");
	tree->Branch("chiYProb", &chiYProb, "chiYProb/D");
	tree->Branch("theta", &theta, "theta/D");

	tree->Branch("meanX", &meanX, "meanX/D");
	tree->Branch("meanY", &meanY, "meanY/D");
	tree->Branch("p0", &p0, "p0/D");
	tree->Branch("p1", &p1, "p1/D");
	tree->Branch("q0", &q0, "q0/D");
	tree->Branch("q1", &q1, "q1/D");

	//variable length arrays
	tree->Branch("sensorHit", sensorHit, "sensorHit[nHits]/I");
	tree->Branch("xHit", xHit, "xHit[nHits]/I");
	tree->Branch("yHit", yHit, "yHit[nHits]/I");
	tree->Branch("physX", physX, "physX[nHits]/D");
	tree->Branch("physY", physY, "physY[nHits]/D");
	tree->Branch("residFourthX", residFourthX, "fourthHitX[fourthHitPresent]/D");
	tree->Branch("residFourthY", residFourthY, "fourthHitY[fourthHitPresent]/D");
	tree->Branch("chiSqProbX3Hit", chiSqProbX3Hit,
			"chiSqProbX3Hit[fourthHitPresent]/D");
	tree->Branch("chiSqProbY3Hit", chiSqProbY3Hit,
			"chiSqProbY3Hit[fourthHitPresent]/D");

	std::cout << "Populating sensors...\n";
	//populate sensors here
	unsigned count(0);
	for (std::vector<MapsSensor*>::const_iterator it = mySensors.begin(); it
			!= mySensors.end(); it++) {
		MapsSensor* s = *it;
		sensorIds[count] = s->id();
		sensorZs[count] = s->zPosition();
		sensorPhis[count] = s->phi();
		MapsSensor::physXY align = s->getAlignment();
		xAlign[count] = align.first;
		yAlign[count] = align.second;
		count++;
		std::cout << "\tExporting: " << *s <<"\n";
	}
	//make sure the first entry contains all the sensors!
	//tree->Fill();

	std::cout << "Writing tracks: " << myTracks.size() << "\n";
	//loop over each track
	double chiXSum(0);
	unsigned fourHitTracks(0);
	unsigned ok(0);
	unsigned anomaly(0);
	for (std::vector<MapsTrack*>::const_iterator it = myTracks.begin(); it
			!= myTracks.end(); it++) {
		MapsTrack* t = *it;
		if (myRequiredLeft == 0 || myRequiredRight == 0)
			t->tellSensorsOfResiduals();
		else
			t->tellSensorsOfResiduals(myRequiredLeft, myRequiredRight);
		std::map<MapsSensor*, MapsTrack::coord> hits = t->getHits();
		std::map<MapsSensor*, MapsTrack::physXY> physicals = t->getGlobalHits();

		bx = t->timeStamp();

		if (t->fourthSensor() == 0) {

			fourthSensor = 0;
			fourthSensorThresh = 0;
		} else {

			fourthSensor = t->fourthSensor()->id();
			fourthSensorThresh = t->fourthSensorThresh();
		}

		nHits = hits.size();

		if (nHits == 4) {
			fourthHitPresent = 1;
			std::pair<double, double> fourthHitResid;
			MapsTrack threeHitVariant;
			t->make3HitTrack(threeHitVariant);
			//diagnose(std::cout, threeHitVariant);
			unsigned status = t->getFourthHitResidual(threeHitVariant,
					fourthHitResid);
			assert(status == 0);
			residFourthX[0] = fourthHitResid.first;
			residFourthY[0] = fourthHitResid.second;
			chiSqProbX3Hit[0] = threeHitVariant.chiSqProb(0);
			chiSqProbY3Hit[0] = threeHitVariant.chiSqProb(1);
			fourHitTracks++;
		} else {
			fourthHitPresent = 0;
		}

		count = 0;
		for (std::map<MapsSensor*, MapsTrack::coord>::iterator hitIt =
				hits.begin(); hitIt != hits.end(); hitIt++) {
			sensorHit[count] = (*hitIt).first->id();
			xHit[count] = (*hitIt).second.first;
			yHit[count] = (*hitIt).second.second;
			physX[count] = physicals.at((*hitIt).first).first;
			physY[count] = physicals.at((*hitIt).first).second;
			
//			//something bonkers is going on...
//			if((*hitIt).first->id() == 7) {
//				if(xHit[count] < 84 && physX[count] < 0) {
//					std::cout << "Anomaly! " << xHit[count] << ",\t --> " << physX[count] << "\n";
//					diagnose(std::cout, *t);
//					++anomaly;
//				} else if(xHit[count] > 84 && physX[count] > 0) {
//					std::cout << "Anomaly! " << xHit[count] << ",\t --> " << physX[count] << "\n"; 
//					diagnose(std::cout, *t);
//					++anomaly;
//				} else {
//					++ok;
//				}
//			}
			count++;
		}
		//loop over global hits

		chiXSum += t->chiSq(0);
		chiX = t->chiSq(0);
		chiY = t->chiSq(1);
		chiXProb = t->chiSqProb(0);
		chiYProb = t->chiSqProb(1);
		meanX = t->meanX();
		meanY = t->meanY();
		theta = t->theta();
		std::pair<double, double> p = t->fitParameters(0);
		p0 = p.first;
		p1 = p.second;

		std::pair<double, double> q = t->fitParameters(1);
		q0 = q.first;
		q1 = q.second;

		//if (theta < 0.1)
		tree->Fill();
	}
	std::cout << "\tFound " << fourHitTracks << " 4 hit tracks.\n";
	std::cout << "ok: \t" << ok << ",\t anomolies: \t" << anomaly << "\n";
 	for (std::vector<MapsSensor*>::const_iterator it = mySensors.begin(); it
			!= mySensors.end(); it++) {
		TH2F resid;
		(*it)->getResidualXYPlot(resid);
		resid.Write();
	}

	file->Write();
	file->Close();

}