#include <cmath>
#include <ctime>
#include <string>
#include <fstream>
#include <iostream>
#include <iomanip>
using namespace std;

#define ColorNum 512
#define ColorDepth 3
#define max_size 1024
#define SecLength 1000000

char SecretFileName[25] = "aaa";

unsigned int Tc, Td;
unsigned int TotalEmbeddedBits;
unsigned int im_height, im_width;
unsigned int PaletteNum, indexsize;

double 		luminance[ColorNum];
double 		sortstack[ColorNum];
double		ColorDifference[ColorNum][ColorNum];
unsigned int	SecMap[5]={0,1,0,1,0};
unsigned int  	EmCode[SecLength];						//用來存放Secret data的bits
unsigned int  	IndexImage[max_size*max_size];					//Image的block指標
unsigned int 	Palette[ColorNum][ColorDepth];
unsigned char 	InImage[max_size][max_size][ColorDepth];
unsigned char 	OutImage[max_size][max_size][ColorDepth];

double RMSE();
double PSNR(double);
int main();
int Product_Palette();
int OptimalReplaceColor(int [], int, int);
void Embedding();
void SortPalette();
void PaletteIndex();
void ReadBMP(char []);							//讀出影像檔
void WriteBMP(char []);							//寫入影像檔
void ReadSigningCode();					//讀檔:Secret Data
void PaletteChange(int);
void ExtrPaletteIndex();
void PreprocessingOfPalette();
void Quicksort(const int ,const int);



//====================================================================================
void PreprocessingOfPalette()
{
	cout<<"PreprocessingOfPalette"<<endl;
	int i, j, k;
	double sum;
	
	//建立每一個Palette中的Color與其它Color的Difference
	for(i=0; i<PaletteNum; i++){
		for(j=0; j<PaletteNum; j++){
			sum = 0.0;
			for(k=0; k<ColorDepth; k++)
				sum += (Palette[i][k]-Palette[j][k])*(Palette[i][k]-Palette[j][k]);
			ColorDifference[i][j]=sqrt(sum);
		}
	}
}
//====================================================================================

//====================================================================================
int OptimalReplaceColor(int neighbor[4], int e, int nP)
{
	int i, j;
	int Opt=0;
	double sum, minDiff;
	int NSec, CountCand;
	int Candidate[ColorNum];
	
	CountCand = 0;
	
	for(i=0; i<PaletteNum; i++){
		if(i>=neighbor[3])
			NSec = SecMap[0];
		else if(neighbor[3]>i && i>=neighbor[2])
			NSec = SecMap[1];
		else if(neighbor[2]>i && i>=neighbor[1])
			NSec = SecMap[2];
		else if(neighbor[1]>i && i>=neighbor[0])
			NSec = SecMap[3];
		else
			NSec = SecMap[4];
				
		if( NSec == e){
			Candidate[CountCand]=i;							
			CountCand++;
		}
	}
	
	minDiff = 999999;
	for(i=0; i<CountCand; i++){		
		sum = ColorDifference[ nP ][ i ];		
		if(sum<minDiff){
			minDiff = sqrt(sum);
			Opt = Candidate[i];
		}
	}
	
	if(minDiff<=Td)
		return Opt;
	else
		return nP;		//not do
}
//====================================================================================

//====================================================================================
void ExtrPaletteIndex()
{
	int i,j,k,l;
	
	//針對 OutImage
	k = 0;
	for(i=0; i<im_height; i++){
		for(j=0; j<im_width; j++){			
			for(l=0; l<ColorDepth; l++){				
				OutImage[i][j][l] = Palette[IndexImage[k]][l];				
			}
			k++;
		}
	}	
}
//====================================================================================

//====================================================================================
void ReadSigningCode(char incb[25])		//讀Secret Data File
{
	char ch;		
	ifstream infile(incb,ios::binary);
	for(int i=0;i<SecLength;i++)
	{
		infile.get(ch);
		if (infile.eof())
			break;	
		else{
			if(ch=='1')	EmCode[i]=1;				
			else if(ch=='0') EmCode[i]=0;
		}
	}
	
	infile.close();
}
//====================================================================================

//====================================================================================
void Embedding()
{
	int i, j, k, l, p;
	int DistinctColor, appear, NowSec, CountCand;
	int ReAp[4];
	int neighbor[4];
	int Candidate[ColorNum];
	double MaxD, sum, CandMaxD, MinDiff;
	
	TotalEmbeddedBits = 0;
	
	cout<<"Tc:(1~4)";
	do{
		cin>>Tc;
	}while(Tc>4 || Tc<1);
	cout<<"Td:";
	do{
		cin>>Td;
	}while(Td>40 || Td<10);
	
	for(i=0; i<im_height; i++){
		for(j=0; j<im_width; j++){
			
			//判斷(i,j)是否符合藏入條件
			DistinctColor = 0;
			MaxD = 0.0;
					
			//計算Maximum color difference between (i,j) and (i-1, j-1), (i-1, j), (i-1, j+1), (i, j-1)
			if((i-1)>0 && (j-1)>0){			//For (i-1, j-1)				
				MaxD = ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[(i-1)*im_width+(j-1)] ];
				sortstack[0] = IndexImage[(i-1)*im_width + (j-1)];			
			}
			if((i-1)>0){					//For (i-1, j)
				if(ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[(i-1)*im_width+j] ]>MaxD)
					MaxD = ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[(i-1)*im_width+j] ];
				sortstack[1] = IndexImage[(i-1)*im_width + j];				
			}
			if((i-1)>0 && (j+1)<im_width){	//For (i-1, j+1)
				if(ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[(i-1)*im_width+(j+1)] ]>MaxD)
					MaxD = ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[(i-1)*im_width+(j+1)] ];
				sortstack[2] = IndexImage[(i-1)*im_width+(j+1)];				
			}
			if((j-1)>0){					//For (i, j-1)
				if(ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[i*im_width+(j-1)] ]>MaxD)
					MaxD = ColorDifference[ IndexImage[i*im_width+j] ][ IndexImage[i*im_width+(j-1)] ];
				sortstack[3] = IndexImage[i*im_width+(j-1)];
			}
			Quicksort(0,3);
			
			DistinctColor = 1;
			for(k=1; k<4; k++){
				if(sortstack[k]!=sortstack[k-1])
					DistinctColor++;
			}
			
			MinDiff = 999999;
			CountCand = 0;
			
			if(IndexImage[i*im_width+j]>=sortstack[3])
				NowSec = SecMap[0];
			else if(sortstack[3]>IndexImage[i*im_width+j] && IndexImage[i*im_width+j]>=sortstack[2])
				NowSec = SecMap[1];
			else if(sortstack[2]>IndexImage[i*im_width+j] && IndexImage[i*im_width+j]>=sortstack[1])
				NowSec = SecMap[2];
			else if(sortstack[1]>IndexImage[i*im_width+j] && IndexImage[i*im_width+j]>=sortstack[0])
				NowSec = SecMap[3];
			else
				NowSec = SecMap[4];
					
			for(k=0; k<PaletteNum; k++){															
				if(k>=sortstack[3])
					l = SecMap[0];
				else if(sortstack[3]>k && k>=sortstack[2])
					l = SecMap[1];
				else if(sortstack[2]>k && k>=sortstack[1])
					l = SecMap[2];
				else if(sortstack[1]>k && k>=sortstack[0])
					l = SecMap[3];
				else
					l = SecMap[4];
								
				if(l != NowSec){
					//計算Maximum color difference between (i,j) and (i-1, j-1), (i-1, j), (i-1, j+1), (i, j-1)
					if((i-1)>0 && (j-1)>0)		//For (i-1, j-1)				
						CandMaxD = ColorDifference[ k ][ IndexImage[(i-1)*im_width+(j-1)] ];						
					if((i-1)>0){					//For (i-1, j)
						if(ColorDifference[ k ][ IndexImage[(i-1)*im_width+j] ]>CandMaxD)
							CandMaxD = ColorDifference[ k ][ IndexImage[(i-1)*im_width+j] ];									
					}
					if((i-1)>0 && (j+1)<im_width){	//For (i-1, j+1)
						if(ColorDifference[ k ][ IndexImage[(i-1)*im_width+(j+1)] ]>CandMaxD)
							CandMaxD = ColorDifference[ k ][ IndexImage[(i-1)*im_width+(j+1)] ];												
					}
					if((j-1)>0){					//For (i, j-1)
						if(ColorDifference[ k ][ IndexImage[i*im_width+(j-1)] ]>CandMaxD)
						CandMaxD = ColorDifference[ k ][ IndexImage[i*im_width+(j-1)] ];								
					}
				}
				if(CandMaxD<MinDiff){
					MinDiff = CandMaxD;
					p = k;
				}					
			}
			
			if(DistinctColor>=Tc && MaxD<=Td && MinDiff<=Td){
				
				if( NowSec != EmCode[TotalEmbeddedBits]){
					//cout<<"Embedding 1 bits! changed"<<TotalEmbeddedBits<<endl;
					IndexImage[i*im_width+j] = p;
					TotalEmbeddedBits++;								
				}
				else{
					//cout<<"Embedding 1 bits!  nonchanged"<<TotalEmbeddedBits<<endl;
					TotalEmbeddedBits++;
				}					
			}
		}
	}
	cout<<"TotalEmbeddedBits:"<<TotalEmbeddedBits<<endl;
}
//====================================================================================

//====================================================================================
void PaletteIndex()							//將影像轉成Palette Index
{
	int i, j, k;		
	int CountOfIndex;	
	
	CountOfIndex=0;
	
	for(i=0; i<im_height; i++){
		for(j=0; j<im_width; j++){
			for(k=0;k<PaletteNum;k++){
				if((InImage[i][j][0]==Palette[k][0]) && (InImage[i][j][1]==Palette[k][1]) && (InImage[i][j][2]==Palette[k][2])){					
					IndexImage[CountOfIndex]=k;
					CountOfIndex++;
				}
			}
		}
	}
	//cout<<CountOfIndex<<"::"<<indexsize<<endl;
	//if(CountOfIndex != indexsize){
	//	cout<<"Error!!Size isn't same!! - VQIndex\n";
	//	exit(1);
	//}	
}
//====================================================================================

//====================================================================================
void PaletteChange(int p)
{
	int i;
	int ColorTemp[ColorDepth];
	
	for(i; i<ColorDepth; i++)
		ColorTemp[i]=Palette[p][i];
	for(i; i<ColorDepth; i++)
		Palette[p][i]=Palette[p][p+1];
	for(i; i<ColorDepth; i++)
		Palette[p+1][i]=ColorTemp[i];
}
//====================================================================================

//====================================================================================
void Quicksort(const int left,const int right)
{
	int i, j;
	double pivot,temp;
	
	if(left<right)
	{
		i=left;
		j=right+1;
		pivot=sortstack[left];
		do{
    		do{
				i++;			
			}while(sortstack[i]<pivot);
			do{
				j--;
			}while(sortstack[j]>pivot);
			if(i<j)
			{
				temp=sortstack[i];
				sortstack[i]=sortstack[j];
				sortstack[j]=temp;
			}
		}while(i<j);
    	temp=sortstack[left];		
		sortstack[left]=sortstack[j];
		sortstack[j]=temp;		
	    
		Quicksort(left,j-1);
    	Quicksort(j+1,right);
	}
}
//====================================================================================

//====================================================================================
double RMSE()							//計算Mean Square Error
{
	double _mse = 0;	
	for(int i=0;i<im_height;i++){
		for(int j=0;j<im_width;j++){
			for(int k=0;k<ColorDepth;k++)
				_mse += (InImage[i][j][k]-OutImage[i][j][k])*(InImage[i][j][k]-OutImage[i][j][k]);
		}
	}	
	_mse = sqrt(_mse/(im_height*im_width));		
	return _mse;
}
//====================================================================================

//====================================================================================
void SortPalette()
{
	int i, j, k, l;
	int appear = 0;
	int Reco[ColorNum];
	int Temp[ColorNum][ColorDepth];
	
	for(i=0; i<PaletteNum; i++){
		for(j=0; j<ColorDepth; j++){
			Temp[i][0] = Palette[i][0];
		    Temp[i][1] = Palette[i][1];
		    Temp[i][2] = Palette[i][2];
		}
	}
	
	for(i=0; i<PaletteNum; i++){
		luminance[i] = 0.3*Palette[i][0]+0.59*Palette[i][1]+0.11*Palette[i][2];
		sortstack[i] = 0.3*Palette[i][0]+0.59*Palette[i][1]+0.11*Palette[i][2];		
	}
	
	Quicksort(0, (PaletteNum-1));
	
	l = 0;
	for(i=0; i<=PaletteNum; i++){	
		for(j=0; j<PaletteNum; j++){			
			if(sortstack[(PaletteNum-1)-i]==luminance[j]){
				appear = 0;
				for(k=0; k<l; k++){
					if(Reco[k]==j)
						appear = 1;
				}
				if(appear==0){
					for(k=0; k<ColorDepth; k++)
						Palette[i][k] = Temp[j][k];
					Reco[l] = j;
					l++;
					break;
				}
			}
		}
		//cout<<endl;
	}
	
	for(i=0; i<PaletteNum; i++)
		luminance[i] = sortstack[(PaletteNum-1)-i];
	
	for(i=0; i<(PaletteNum-1); i++){		
		if(luminance[i]==luminance[i+1]){
			if(Palette[i][0]<Palette[i+1][0]){
				PaletteChange(i);
			}
			else if(Palette[i][0]==Palette[i+1][0]){
				if(Palette[i][1]<Palette[i+1][1]){
					PaletteChange(i);
				}
				else if(Palette[i][1]==Palette[i+1][1]){
					if(Palette[i][2]<Palette[i+1][2])
						PaletteChange(i);
				}
			}
		}
	}	
	//cout<<"Palette:"<<PaletteNum<<endl;
	//for(i=0;i<PaletteNum;i++){
	//	for(j=0; j<ColorDepth; j++)
	//		cout<<(int)Palette[i][j]<<" ";
	//	cout<<endl;
	//}
}
//====================================================================================

//====================================================================================
double PSNR(double _pm)					//計算PSNR
{	
	double _psnr;		
	_psnr = 10*log10(255*255/_pm);	
	return _psnr;
}
//====================================================================================

//====================================================================================
int Product_Palette()
{
	int i, j, k, pn, reseen;
	pn = 0;
	for(i=0;i<im_height;i++){
		for(j=0;j<im_width;j++){
			reseen = 0;
			for(k=0; k<pn; k++){							
				if((Palette[k][0]==InImage[i][j][0]) && (Palette[k][1]==InImage[i][j][1]) && (Palette[k][2]==InImage[i][j][2]))
					reseen = 1;
		    }
		    if(reseen == 0){
		    	Palette[pn][0] = InImage[i][j][0];
		    	Palette[pn][1] = InImage[i][j][1];
		    	Palette[pn][2] = InImage[i][j][2];
		    	pn++;
		    }		    
		}
	}		
	//cout<<"Palette:"<<pn<<endl;
	//for(i=0;i<pn;i++){
	//	for(j=0; j<ColorDepth; j++)
	//		cout<<(int)Palette[i][j]<<" ";
	//	cout<<endl;
	//}
	return pn;
}
//====================================================================================

//====================================================================================
void WriteBMP(char outputname[40])	//寫影像檔(.Raw)
{	
	int i, j, k;
	ofstream outfile(outputname,ios::binary);
    for(i=0;i<im_height;i++){
		for(j=0;j<im_width;j++){		
			for(k=0; k<ColorDepth; k++)
		    	outfile<<(unsigned char)OutImage[i][j][k];		
		}
	}
	outfile.close();		
}
//====================================================================================

//====================================================================================
void ReadBMP(char inputname[25])	//讀影像檔(.Raw)
{
	int i, j, k;
	char ch;		
	ifstream infile(inputname,ios::binary);
	for(i=0;i<im_height;i++){
		for(j=0;j<im_width;j++){
			for(k=0; k<ColorDepth; k++){
				infile.get(ch);				
				if (infile.eof())
					break;	
				else
					InImage[i][j][k] = ch;	
			}
		}
	}
	infile.close();	
	
	//for(k=0; k<ColorDepth; k++)			//test
	//	cout<<(int)InImage[0][0][k]<<" ";	
	//cout<<endl;
}
//====================================================================================

//====================================================================================
int main()
{
	int i, j, k;
	double ms, ps;
	char InImageName[25];
	char OutImageName[40]="COM_Cover_";
	
	cout<<"image file name (.raw):";
	cin>>InImageName;
	cout<<"height:";
	do{
		cin>>im_height;
	}while(im_height>max_size || im_height<0);
	cout<<"width:";
	do{
		cin>>im_width;
	}while(im_width>max_size || im_width<0);
	
	indexsize = im_height*im_width;
		
	ReadBMP(InImageName);
	ReadSigningCode(SecretFileName);
	PaletteNum = Product_Palette();
	cout<<"PaletteNum"<<PaletteNum<<endl;
	for(i=0;i<im_height;i++){
		for(j=0;j<im_width;j++){		
			for(k=0; k<ColorDepth; k++)
		    	OutImage[i][j][k]=InImage[i][j][k];
		}
	}
	
	SortPalette();
	
	PaletteIndex();
	
	PreprocessingOfPalette();
	
	Embedding();
	
	ExtrPaletteIndex();
	
	ms = RMSE();	//==計算MSE==	
	cout<<"MSE = "<<ms;	
	if( ms != 0){
		ps = PSNR(ms);  //==計算PSNR==
		cout<<"	PSNR = "<<ps;
	}
	
	strcat(OutImageName, InImageName);
	WriteBMP(OutImageName);
	
	return 0;
}
//====================================================================================

