// $Id$
//
//    File: JEventProcessor_TrackPNG.cc
// Created: Fri Feb  1 07:09:07 EST 2019
// Creator: davidl (on Linux gluon47.jlab.org 2.6.32-642.3.1.el6.x86_64 x86_64)
//

// This plugin will create a PNG image file representing each single track.
// The image will be greyscale with each CDC and FDC wire represented 
// by a pixel. The pixel color will be 255 (not hit) or a number between
// 0-249 representing the drift time corresponding to 0-200ns. 
//
// The images will be separated into directory structures corresponding to
// the geant PID of the track. Because the hits for different mass hypothesis
// will be identical, the images will also be identical.
//
// In each PID directory will be a file "fit_parms.csv" that will contain the
// name of the file and the TBT fit parameters.


#include <iostream>
#include <vector>
#include <cstdint>
#include <string>
using namespace std;

#include <png.h>
#include <stdio.h>
#include <string.h>
#include <sstream>
#include <sys/stat.h>
#include <sys/types.h>

#include <DANA/DApplication.h>
#include <particleType.h>
#include <CDC/DCDCTrackHit.h>
#include <FDC/DFDCPseudo.h>
#include "JEventProcessor_TrackPNG.h"
using namespace jana;

// Routine used to create our JEventProcessor
#include <JANA/JApplication.h>
#include <JANA/JFactory.h>
extern "C"{
void InitPlugin(JApplication *app){
	InitJANAPlugin(app);
	app->AddProcessor(new JEventProcessor_TrackPNG());
}
} // "C"


//------------------
// JEventProcessor_TrackPNG (Constructor)
//------------------
JEventProcessor_TrackPNG::JEventProcessor_TrackPNG()
{
	// Create output directories and open files for track parameters
	mode_t mymode = S_IRUSR|S_IWUSR|S_IXUSR | S_IRGRP|S_IWGRP|S_IXGRP | S_IROTH|S_IWOTH|S_IXOTH;
	mkdir( "IMAGES", mymode);
	vector<int> pids = {2,3,8,9,11,12,14,15};
	for( int pid : pids ){
		char mydirname[256];
		sprintf(mydirname, "IMAGES/PID%02d", pid);
		mkdir( mydirname, mymode);
		
		ofstream myofs( (string(mydirname)+"/README_IMPORTANT").c_str() );
		myofs << "DO NOT DO \"ls-a\" IN THIS DIRECTORY!" << endl;
		myofs << endl;
		myofs << "There are MANY hidden files here (names start with \".\"" << endl;
		myofs << "The track_parms.csv file has a list of all of them so it is" << endl;
		myofs << "probably better to just read that." << endl;
		myofs << endl;
		myofs << "These files were generated by the TrackPNG plugin." << endl;
		myofs << "Contact davidl@jlab.org if you have questions." << endl;
		myofs.close();

		char fname[256];
		sprintf(fname, "%s/track_parms.csv", mydirname);
		ofs[ pid ] = new ofstream( fname );	
		(*ofs[pid]) << "filename,event,trackid,q_over_pt,phi,tanl,D,z,cov00,cov01,cov02,cov03,cov04,cov11,cov12,cov13,cov13,cov14,cov22,cov23,cov24,cov33,cov34,cov44,pid" << endl;

		sprintf(fname, "%s/images.raw", mydirname);
		bfs[ pid ] = new ofstream( fname );	
	}

	// Set number of wires in each CDC layer
	cdc_nwires.push_back(42);
	cdc_nwires.push_back(42);
	cdc_nwires.push_back(54);
	cdc_nwires.push_back(54);
	cdc_nwires.push_back(66);
	cdc_nwires.push_back(66);
	cdc_nwires.push_back(80);
	cdc_nwires.push_back(80);
	cdc_nwires.push_back(93);
	cdc_nwires.push_back(93);
	cdc_nwires.push_back(106);
	cdc_nwires.push_back(106);
	cdc_nwires.push_back(123);
	cdc_nwires.push_back(123);
	cdc_nwires.push_back(135);
	cdc_nwires.push_back(135);
	cdc_nwires.push_back(146);
	cdc_nwires.push_back(146);
	cdc_nwires.push_back(158);
	cdc_nwires.push_back(158);
	cdc_nwires.push_back(170);
	cdc_nwires.push_back(170);
	cdc_nwires.push_back(182);
	cdc_nwires.push_back(182);
	cdc_nwires.push_back(197);
	cdc_nwires.push_back(197);
	cdc_nwires.push_back(209);
	cdc_nwires.push_back(209);
}

//------------------
// ~JEventProcessor_TrackPNG (Destructor)
//------------------
JEventProcessor_TrackPNG::~JEventProcessor_TrackPNG()
{
	// Close all tracking param output files
	for( auto pr : ofs ){
		lock_guard<mutex> grd( mtxs[ pr.first ] );
		pr.second->close();
		delete pr.second;
	}

	// Close all image param output files
	for( auto pr : bfs ){
		lock_guard<mutex> grd( mtxs[ pr.first ] );
		pr.second->close();
		delete pr.second;
	}
}

//------------------
// init
//------------------
jerror_t JEventProcessor_TrackPNG::init(void)
{
	// This is called once at program startup. 

	return NOERROR;
}

//------------------
// brun
//------------------
jerror_t JEventProcessor_TrackPNG::brun(JEventLoop *eventLoop, int32_t runnumber)
{
	
	return NOERROR;
}

//------------------
// evnt
//------------------
jerror_t JEventProcessor_TrackPNG::evnt(JEventLoop *loop, uint64_t eventnumber)
{
	vector<const DTrackTimeBased*> tbts;
	loop->Get(tbts);
	uint32_t trackid = 0;
	for(auto tbt : tbts) WriteTrack( tbt, trackid++, eventnumber);

	return NOERROR;
}

//------------------
// WriteTrack
//------------------
void JEventProcessor_TrackPNG::WriteTrack(const DTrackTimeBased *tbt, uint32_t trackid, uint64_t eventnumber)
{
	// Create PNG file name
	char fname[256], fullpath[256];
	auto ptype_str = ParticleType(tbt->PID());
	sprintf( fname, ".trk%06ld_%02d_%s.png", eventnumber, trackid, ptype_str);
	sprintf( fullpath, "IMAGES/PID%02d/%s", tbt->PID(), fname);
	

// 	FILE *fp = fopen(fullpath, "wb");
// 	if(!fp){
// 		cerr << "Unable to open file " << fname << " for writing!" << endl;
// 		return;
// 	}
// 	
// 	png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING
// 		, png_voidp_NULL
// 		, png_error_ptr_NULL
// 		, png_error_ptr_NULL);
// 	if( !png_ptr ){
// 		cerr << "Error creating PNG write struct! " << endl;
// 		return;
// 	}
// 	
// 	png_infop info_ptr = png_create_info_struct(png_ptr);
// 	if( !info_ptr ){
// 		png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
// 		cerr << "Unable to create PNG info struct!" << endl;
// 		return;
// 	}
// 	
// 	if( setjmp( png_jmpbuf(png_ptr) ) ){
// 		png_destroy_write_struct(&png_ptr, &info_ptr);
// 		fclose(fp);
// 		cerr << "Something went wrong and wayback was activated!" << endl;
// 		return;
// 	}
	
	// Allocate memory for image data. One byte per pixel
	png_uint_32 width  = 720;
	png_uint_32 height = 360;
	auto buff = new uint8_t[ width*height ];
	vector<png_byte*> row_pointrs;
	for(png_uint_32 irow=0; irow<height; irow++){
		row_pointrs.push_back( &buff[irow*width] );
	}

	// initialize everything as white
	memset(buff, 0xff, width*height);

// 	// Black out pixels unused by drift chambers
// 	for(int j=0; j<28; j++){
// 		for( uint32_t i=cdc_nwires[27-j]; i<width; i++) row_pointrs[j][i] = 0;
// 	}
// 	for(int j=0; j<24; j++){
// 		for( uint32_t i=96; i<width; i++) row_pointrs[j+28][i] = 0;
// 	}
	
	// Set pixels in image
	vector<const DCDCTrackHit*> cdchits;
	vector<const DFDCPseudo*> fdchits;
	tbt->Get( cdchits );
	tbt->Get( fdchits );
	
	// CDC (left half of picture)
	double x0 = width/4;
	double y0 = height/2;
	double mx = (width/2.0)/180.0; // pixels/cm in x
	double my =  height/180.0;     // pixels/cm in y
	double bias = 16.0;
	for(auto hit: cdchits){
		double x = hit->wire->origin.X();
		double y = hit->wire->origin.Y();
		uint32_t icol = floor((x*mx) + x0 + 0.5);
		uint32_t irow = floor((y*my) + y0 + 0.5);
		if( (icol>=(width/2)) || (irow>=height) ){
			_DBG_ << "pixel coordinates out of range! (irow="<<irow<<" icol="<<icol<<endl;
			_exit(-1);
		}
		double color =   bias + (255.0-bias)*hit->tdrift/200.0;
		if( color>255.0 ) color = 255.0;
		if( color<bias  ) color = bias;
		row_pointrs[irow][icol] = (uint8_t)color;
	}

	// FDC (right half of picture)
	x0 += width/2;
	for(auto hit: fdchits){
		double x = hit->xy.X();
		double y = hit->xy.Y();
		uint32_t icol = floor((x*mx) + x0 + 0.5);
		uint32_t irow = floor((y*my) + y0 + 0.5);
		if( (icol<(width/2)) || (icol>=width) || (irow>=height) ){
			_DBG_ << "pixel coordinates out of range! (irow="<<irow<<" icol="<<icol<<endl;
			_exit(-1);
		}
		double color =   bias + (255.0-bias)*hit->time/200.0;
		if( color>255.0 ) color = 255.0;
		if( color<bias  ) color = bias;
		row_pointrs[irow][icol] = (uint8_t)color;
	}
	

// 	// Write PNG file
// 	png_init_io(png_ptr, fp);
// 	png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_GRAY, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
// 	png_write_info(png_ptr, info_ptr);
// 	png_write_image(png_ptr, row_pointrs.data());
// 	png_write_end(png_ptr, info_ptr);
// 	png_destroy_write_struct(&png_ptr, &info_ptr);
// 	delete[] buff;
// 	fclose(fp);
	
	// Add tracking parameters to appropriate output file.
	double state[5];
	tbt->TrackingStateVector(state);
	auto errm = tbt->TrackingErrorMatrix();
	stringstream ss;
	ss << fname << "," ;
	ss << eventnumber << "," << trackid << ",";
	for(int i=0; i<5; i++) ss << state[i] << ",";
	for(int i=0; i<5; i++) 
		for(int j=i; j<5; j++) ss << (*errm)[i][j] << "," ;
	ss << tbt->PID();
	
	lock_guard<mutex> grd( mtxs[tbt->PID()] );
	auto of = ofs[ tbt->PID() ];
	if( of ) (*of) << ss.str() << endl;

	bfs[tbt->PID()]->write( (const char*)buff, width*height );
	
	delete[] buff;
}		


//------------------
// erun
//------------------
jerror_t JEventProcessor_TrackPNG::erun(void)
{
	// This is called whenever the run number changes, before it is
	// changed to give you a chance to clean up before processing
	// events from the next run number.
	return NOERROR;
}

//------------------
// fini
//------------------
jerror_t JEventProcessor_TrackPNG::fini(void)
{
	// Called before program exit after event processing is finished.
	return NOERROR;
}