/*
	File providing TF2 function returning interpolated PSF for given parameters.
	Requires PRF to be loaded into prf_histo
*/

#define lambda2 1.;
#define oneolambda 1.;
#define oneolambda2 1.;	

#include <iostream>
#include "TApplication.h"
#include "TROOT.h"
#include "TCanvas.h"
#include "TH2.h"
#include "TProfile2D.h"
#include "TStyle.h"
#include "TMath.h"
#include "TF2.h"
#include "TMinuit.h"
#include "TVirtualFitter.h"
#include "TFile.h"
#include "TDatime.h"
#include "TTimer.h"
#include "TSystem.h"
#include "TGraph.h"
#include "TGraph2D.h"

#include <vector>	
#include <map>
#pragma link C++ class vector<std::pair<double, double>>+;

using namespace std;

// global declarations required for analytical scale and offset calculation
TH2D *img_histo;
TF2 *mpsf;

// Global fit parameters (named)

const Int_t nnam=8;

// Maximum number of fit parameters (nnam+numzer[nord])

const Int_t npar=44;

double window;

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

const Int_t nint=10;
double intx[] = {-0.9739065, -0.865063, -0.6794096, -0.433395, -0.1488743, 
						0.1488743,	0.433395,	0.6794096,	0.865063,	0.9739065}; 
double intw[] = { 0.0666713443, 0.149451349, 0.21908636, 0.2692667,	0.29552422,
						0.29552422,	0.2692667,	0.21908636, 0.149451349, 0.0666713443};

// ----------------------------------------------------------------------------------
//
// Pixel response function
// Geometrical pixel size is (-0.5,0.5)x(-0.5,0.5)

const double xprfmax = 0.8;
const double yprfmax = 0.8;
Double_t aint_res = 20., **prf_mat;
TH2D *prf_histo;

double prf_div = aint_res/(2*xprfmax);

inline double prf(double x, double y)
{
	return prf_mat[TMath::Nint((x+xprfmax)*prf_div)][TMath::Nint((y+yprfmax)*prf_div)];
}

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


static double fact( int n ) {	// The factorial function!
	double f = 1;
	while(n>1)
		f *= n--;
	return f;
}

int factm[] = {1, 1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800, 39916800, 479002000};
int pow_m1[] = {1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1};
double ro_pow[10];
TH1D *Rmnro_hist[64];
int Rmnro_mn[64][2], Rmnro_cnt=0;
double *Rmnro_mat[16][16];

double zernike_polynomials( int m, int n, double ro, double th ) 
{
	double Rmnro;
	bool even = m>=0;
	if(!even) m = -m;
	if( (n-m)%2 ) return 0;

	// Check if Rmnro for these m,n was already calculated
	if(Rmnro_mat[n][m]!=NULL) Rmnro = Rmnro_mat[n][m][int(ro*1024)];
	else
	{	
		cout << "Calculating Rmnro matrix for m=" << m << ", n=" << n << "... " << endl;
		cout.flush();
		int nmmo2 = (n-m)/2;
		int npmo2 = (n+m)/2;
		Rmnro_mat[n][m]=new double[1024];
		double y=0.0001;
		for(int x=0; x<1024; ++x)
		{
			for(int k=0;k<=nmmo2;++k) 
			{
				Rmnro += pow(y,n-2*k)*( pow_m1[k] * factm[n-k] ) /
					( factm[k] * factm[npmo2-k] *  factm[nmmo2-k]);
			}
			Rmnro_mat[n][m][x]=Rmnro;
			Rmnro=0;
			y+=9.765625e-04;
		} 
		Rmnro = Rmnro_mat[n][m][int(ro*1024)];
	}
 

	double m_th = m*th;
	if(even) return Rmnro * cos(m_th);
	else return Rmnro * sin(m_th);

}

int n[64]={0.}; 
int m[64]={0.};
 
Double_t myfun2d(Double_t *x, Double_t *par) 
{
// par[0] can be negative to skip integration over pixel surface
	bool bpar0 = (par[0]>0);
	int nzer=bpar0?par[0]+0.5:0.5-par[0];

	double x0=x[0]-par[1];
	double y0=x[1]-par[2];
	double dphi=par[3];
	double pwr=par[4];
	double norm=par[5];
	double bg=par[6];	
	double lambda=par[7];
	
	double val=0.;

	int nx=bpar0?nint:1;
	int ny=nx;

	double intw_ix, intx_ix;
	double rpwr;

	for(int ix=0; ix<nx; ++ix)
	{
		intw_ix=intw[ix];
		intx_ix=intx[ix];
		double dx = x0 + xprfmax*intx_ix*bpar0;
		double dx2 = dx*dx;
		for(int iy=0; iy<ny; ++iy)
		{
			double dy = y0 + yprfmax*intx[iy]*bpar0;

			double r2=rpwr=dx2+dy*dy;
			double r=TMath::Sqrt(r2);

			double r2l2=r2*oneolambda2;
			double r2l=r*oneolambda;

			double u=1.-TMath::Exp(-r2l);
			double phi=TMath::ATan2(dy,dx)-dphi;

			double zer=0.;

			for(int j = 0; j<nzer; ++j)
			{
				if(!TMath::AreEqualAbs(par[nnam+j], 0, 0.0000001) && par[nnam+j]!=0.)
				{
//					if(nnam+j==8)
					zer+= par[nnam+j] * zernike_polynomials(m[j], n[j], u, phi );
				}
			}


			if(par[0]>0)
			{
				double wint = intw_ix*intw[iy]*prf(xprfmax*intx_ix,yprfmax*intx[iy]);
				double epow = -0.5*rpwr*zer;
				// so small value means that fit here is wrong, so the value does not matter
				if(epow>10) epow=10;
				val+=TMath::Exp(epow)*wint;
			}
			else
			{
				double epow = -0.5*rpwr*zer;
				// so small value means that fit here is wrong, so the value does not matter
				if(epow>10) epow=10;
				val+=TMath::Exp(epow);

			}
		}
	}
	return	norm*val + bg;
}

vector< vector< vector<Double_t> > > myfun2d_mat(0, vector< vector<Double_t> > (41, vector<Double_t>(41,0)) );
vector<Double_t> myfun2d_posx;
vector<Double_t> myfun2d_posy;
int nf=0;
Int_t nbx, nby, nbxo2, nbyo2;
bool disp=false;

int scale_fit_range=7, old_scale_fit_range=0;
bool refit=false;

void calc_scale_offset(Double_t &mscale, Double_t &moffset, Double_t *psf_p)
{
	Double_t sf2=0, sf=0, s1=0, sfz=0;
	vector<Double_t> vy(0,0);
	vector< vector<Double_t> > vx(0, vector<Double_t>(0,0));
	Double_t _1oerror2;
	// First calculate needed sums
	for(int i=-nbxo2; i<=nbxo2; ++i)
	{
		for(int j=-nbyo2; j<=nbyo2; ++j)
		{
			vy.push_back(1);
			if(i>=-scale_fit_range && i<=scale_fit_range && j>=-scale_fit_range && j<=scale_fit_range)
			{
				Int_t ci = i;
				Int_t cj = j;
				// Skip proper pixels for different pix numbers
				if((ci==-2 && cj==-2) || (ci==2 && cj==2) || (ci==-2 && cj==2) || (ci==2 && cj==-2)) continue;
				Double_t x[]={(Double_t)i,(Double_t)j};
				Double_t f_val = myfun2d(x, psf_p);
				Double_t h_val = img_histo->GetBinContent(i+nbxo2+1,j+nbyo2+1);

				vy.pop_back();
				vy.push_back(f_val);
				// f - fitted function value, z - data pixel value

				if(h_val>0) _1oerror2 = 1./h_val;
				else _1oerror2=1;
				// This is assuming operation on dark subtracted frames and dark value is about 1600
				_1oerror2=1./(h_val+1500);
				sf+=f_val*_1oerror2; sf2+=(f_val*f_val)*_1oerror2; s1+=_1oerror2;
				
				sfz+=f_val*h_val*_1oerror2;
			}
		}
		vx.push_back(vy);
		vy.resize(0);
	}
	myfun2d_mat.push_back(vx);

	// only if mscale was fixed to 1. otherwise assume a function with with other scale was requested
	// however, need to empty matrices to get proper result
	if(TMath::AreEqualAbs(psf_p[5], 1, 0.00000001) || psf_p[5]==1)
		mscale = (sfz-moffset*sf)/sf2;
	else
		mscale=1;

	// Recalculate values matrix to contain rescaled values
	for(int i=-nbxo2; i<=nbxo2; ++i)
		for(int j=-nbyo2; j<=nbyo2; ++j)
		{
			Double_t f_val = myfun2d_mat[nf][i+nbxo2][j+nbyo2];
			// only if function was calculated in this point
			if(f_val!=1 || !TMath::AreEqualAbs(f_val, 1, 0.00000001)) myfun2d_mat[nf][i+nbxo2][j+nbyo2]=mscale*f_val+moffset;
			// otherwise keep 1 as its value
			else myfun2d_mat[nf][i+nbxo2][j+nbyo2]=1;
		}
	// Needed for caching of function values
	myfun2d_posx.push_back(psf_p[1]);
	myfun2d_posy.push_back(psf_p[2]);
	nf++;
	
}

Double_t old_par[100]={-1000, -1000, -1000, -1000, -1000, -1000, -1000, -1000}, scale, offset;

TH2D *par_histo[44];

Double_t old_mscale=0, old_moffset=1;

bool debug=false;

// Function returning value of myfun2D for parameters calculated from paraboles parameters
double multifun(double *x, double *par)
{
	// Reject points chosen for rejection
	if(x[0]>scale_fit_range || x[0]<-scale_fit_range || x[1]>scale_fit_range || x[1]<-scale_fit_range)
		{ TF2::RejectPoint(); return 0;}

	int roundx, roundy;
	if(x[0]>=0 && x[0]-TMath::Nint(x[0])>=0.5) roundx=TMath::Nint(x[0])+1;
	else if(x[0]>=0 && x[0]-TMath::Nint(x[0])<0.5) roundx=TMath::Nint(x[0]);
	else if(x[0]<0 && x[0]-TMath::Nint(x[0])<-0.5) roundx=TMath::Nint(x[0])-1;
	else if(x[0]<0 && x[0]-TMath::Nint(x[0])>=-0.5) roundx=TMath::Nint(x[0]);
	if(x[1]>0 && x[1]-TMath::Nint(x[1])>=0.5) roundy=TMath::Nint(x[1])+1;
	else if(x[1]>=0 && x[1]-TMath::Nint(x[1])<0.5) roundy=TMath::Nint(x[1]);
	else if(x[1]<0 && x[1]-TMath::Nint(x[1])<-0.5) roundy=TMath::Nint(x[1])-1;
	else if(x[1]<0 && x[1]-TMath::Nint(x[1])>=-0.5) roundy=TMath::Nint(x[1]);

	if((roundx==-2 && roundy==-2) || (roundx==2 && roundy==2) || (roundx==-2 && roundy==2) || (roundx==2 && roundy==-2)) 
	{ TF2::RejectPoint();  return 0; }

	// Check if function was not already calculated
	for(int i=0; i<nf; i++) if((TMath::AreEqualAbs(myfun2d_posx[i], par[1], 0.00000001) || myfun2d_posx[i]==par[1]) && (TMath::AreEqualAbs(myfun2d_posy[i], par[2], 0.00000001) ||myfun2d_posy[i]==par[2]))
	{
		return myfun2d_mat[i][TMath::Nint(x[0]+nbx/2)][TMath::Nint(x[1]+nby/2)];
	}

	// parameters matrix for called myfun2D
	double psf_p[45];

	psf_p[1]=par[1]; psf_p[2]=par[2]; 
	psf_p[5]=par[4]; psf_p[6]=par[5];

	// nzer, axis, pow, lambda
	psf_p[0]=par[8]*36; psf_p[3]=par[3]; psf_p[4]=par[6]; psf_p[7]=par[7];
	
	// List of defined parameters
	Int_t parlist[] = {8, 12, 13, 14, 15, 20, 21, 22, 24, 25, 32, 33, 34, 35, 37, 38, 39}, parcnt = 17;
	for(int i=8; i<45; i++)
	{
		psf_p[i]=0;
		for(int j=0; j<parcnt; j++) 
		{
			if(i==parlist[j]) 
			{
				if(par[3]<0) par[3]+=TMath::TwoPi();
				psf_p[i]=par_histo[i]->GetBinContent(par_histo[i]->FindBin(par[3],par[0]));
				break;
			}				
		}
	}

	Double_t nx[]={x[0]-par[1], x[1]-par[2]};

	// Settings for analytical scale and offset calculations
	Double_t mscale=old_mscale, moffset=psf_p[6];
	psf_p[6]=0;
	
	// Calc only if centre position changed
	if((!TMath::AreEqualAbs(old_par[1], par[1], 0.00000001) && old_par[1]!=par[1]) || (!TMath::AreEqualAbs(old_par[2], par[2], 0.00000001) && old_par[2]!=par[2]) || nf==0)
	{ 
		calc_scale_offset(mscale, moffset, psf_p);
		old_par[1]=par[1];
		old_par[2]=par[2];
		old_mscale=mscale; old_moffset=moffset;
	}

	psf_p[5]=mscale; psf_p[6]=moffset;

	return myfun2d_mat[nf-1][TMath::Nint(x[0]+nbx/2)][TMath::Nint(x[1]+nby/2)];
}
	
// ----------------------------------------------------------------------------------

TF2 *init_mpsf(double axis=0., double dist=0., double window=3., double scale=1., double offset=0., double power=0., double lambda=0., double bgx=0., double bgy=0.)
{	
	TDatime *tm = new TDatime();
	tm->Print();
	
	// Open file with parameter values
	Int_t pnum, ccdx, ccdy;
	
	int histo_nums[]={8, 12, 13, 14, 15, 20, 21, 22, 24, 25, 32, 33, 34, 35, 37, 38, 39};
	const int histo_cnt=17;
	TFile *f = new TFile("par_histos.root", "READ");

	// Read parameter histograms into memory
	for(int i=0; i<histo_cnt; i++)
	{
		par_histo[histo_nums[i]] = (TH2D*)(f->Get("par"+TString::Format("%d", histo_nums[i])));
		par_histo[histo_nums[i]]->SetDirectory(0);
	}
	f->Close();
	
	// NULL the Rmnro matrix
	cout << "nulling matrix" << endl;
	for(int a=0; a<16; a++) for(int b=0; b<16; b++) Rmnro_mat[a][b]=NULL;

	mpsf = new TF2("mpsf", multifun, -window, window, -window, window, 9);
	mpsf->SetParNames("Distance", "X0", "Y0", "Axis", "Scale", "Offset", "Power", "Lambda", "Convolution");
	mpsf->SetParameters(dist, 0, 0, axis, scale, offset, power, lambda, 1);


	// Generate PRF matrix
	int i, j;
	prf_mat = new Double_t *[TMath::Nint(aint_res)*TMath::Nint(aint_res)];
	for(i=0; i<TMath::Nint(aint_res)*TMath::Nint(aint_res); i++) prf_mat[i] = new Double_t[TMath::Nint(aint_res)*TMath::Nint(aint_res)];
	Double_t x1, y1;

	cout << "Generating PRF matrix" << endl;

	for(x1=-xprfmax, i=0; x1<=xprfmax; x1+=2*xprfmax/aint_res, i++) for(y1=-yprfmax, j=0; y1<=yprfmax; y1+=2*yprfmax/aint_res, j++)
	{
		// Take values from the middle, not the edge of the range
		prf_mat[i][j]=prf_histo->Interpolate(x1+xprfmax/aint_res, y1+yprfmax/aint_res);
	}

	// Calculate zernike polynomials numbers
	for(int j = 0; j<36; ++j)
	{
		n[j] = TMath::Sqrt(9.+8.*j)/2. - 0.500001;
		m[j] = 2*j - n[j]*(n[j]+2);
	}

	return mpsf;		
}

