#include "CImg.h"
#include <math.h>

/*---------------------------------------------------------------------------*/
/**
 * @file main.cpp
 * Este arquivo contém uma implementação do método de Horn-Schunck para determinação
 * do fluxo óptico em um par de imagens em nível de cinza
 * As imagens devem possuir a mesma dimensão.
 * O programa gera um arquivo contendo o campo vetorial gerado e tenta ler o arquivo
 * com o campo original, cujo nome pode ser passado como parâmetro.
 * @author Lenardo (lmartins@tecgraf.puc-rio.br)
 * @date 15/10/2008
 */
/*---------------------------------------------------------------------------*/


/**
  * Função que calcula a derivada Ex non ponto (i,j)
  * @param[in] img1 primeira imagem
  * @param[in] img2 segunda imagem
  * @param[in] i posição x do pixel
  * @param[in] j posição y do pixel
  * @param[in] width dimensão horizontal da imagem
  * @return o valor de Ex naquele ponto        
*/
inline double getEx(double * img1, double * img2, int i, int j, int width)
{
	return (img1[i+1+j*width] + img2[i+1+j*width] + img1[i+1+(j+1)*width] + img2[i+1+(j+1)*width] -
		img1[i+  j*width] - img2[i+  j*width] - img1[i+  (j+1)*width] - img2[i+  (j+1)*width])
		/4.0;
}

/**
  * Função que calcula a derivada Ey non ponto (i,j)
  * @param[in] img1 primeira imagem
  * @param[in] img2 segunda imagem
  * @param[in] i posição x do pixel
  * @param[in] j posição y do pixel
  * @param[in] width dimensão horizontal da imagem
  * @return o valor de Ey naquele ponto        
*/
inline double getEy(double * img1, double * img2, int i, int j, int width)
{
	return (img1[i+(j+1)*width] + img2[i+(j+1)*width] + img1[i+1+(j+1)*width] + img2[i+1+(j+1)*width] -
			img1[i+    j*width] - img2[i+    j*width] - img1[i+1+    j*width] - img2[i+1+    j*width])
			/4.0;
}

/**
  * Função que calcula a derivada Et non ponto (i,j)
  * @param[in] img1 primeira imagem
  * @param[in] img2 segunda imagem
  * @param[in] i posição x do pixel
  * @param[in] j posição y do pixel
  * @param[in] width dimensão horizontal da imagem
  * @return o valor de Et naquele ponto        
*/
inline double getEt(double * img1, double * img2, int i, int j, int width)
{
		return (img2[i+j*width] + img2[i+(j+1)*width] + img2[i+1+j*width] + img2[i+1+(j+1)*width] -
				img1[i+j*width] - img1[i+(j+1)*width] - img1[i+1+j*width] - img1[i+1+(j+1)*width])
				/4.0;
}


using namespace cimg_library;

const double PI = 3.14159265;
const float white = 255.0;

/**
  * Função que calcula o erro angular médio e o erro médio, dado nome do arquivo 
  * que contém o campo correto e o nome do arquivo do campo estimado.
  * @param[in] correctField nome do arquivo que contém o campo correto
  * @param[in] estimatedField nome do arquivo que contém o campo estimado
  * @param[out] aae Average Angular Error
  * @param[out] error erro absoluto médio      
*/
bool calcError(char* correctField, char* estimatedField, double* aae, double* error)
{
     
     FILE* f1 = fopen(correctField,"r");
     FILE* f2 = fopen(estimatedField,"r");
     
     if (!f1 || !f2) return false;
     
     char s1[256],s2[256];
     int w,h, lin, col;
     float u1,v1, u2,v2, w1, w2;
     double asoma = 0.0, esoma = 0.0, mod1, mod2;
     fscanf(f1,"%s %s",s1,s2);
     fscanf(f2,"%s %s",s1,s2);
     fscanf(f1,"%d %d",&w,&h);
     fscanf(f2,"%d %d",&w,&h);
     
     CImg<>image(w,h);//imagem 1
     CImg<>field(image.dimx(), image.dimy(),1,2,3); //campo de vetores obtido          
     
     while (!feof(f1))
     {
        fscanf(f1,"%d",&lin);          
        fscanf(f1,"%d",&col);   
        fscanf(f1,"%f",&u1);          
        fscanf(f1,"%f",&v1);
        
        field(col, lin,0) = u1;
        field(col, lin,1) = v1;
        
        fscanf(f2,"%d",&lin);
        fscanf(f2,"%d",&col);
        fscanf(f2,"%f",&u2);
        fscanf(f2,"%f",&v2);
        
        w1 = 1.0;
        w2 = 1.0;
        
        esoma += sqrt( (u1-u2)*(u1-u2) + (v1-v2)*(v1-v2));
                
        mod1 = sqrt(u1*u1 + v1*v1 + w1*w1);
        mod2 = sqrt(u2*u2 + v2*v2 + w2*w2);
        
        if (mod1 != 0.0)
        {
           u1 /= mod1;
           v1 /= mod1;
           w1 /= mod1;
           
        }
        if (mod2 != 0.0)
        {
           u2 /= mod2;
           v2 /= mod2;
           w2 /= mod2;
        }
        
        double pEscalar = u1*u2 + v1*v2 + w1*w2;       
        
        if ( pEscalar > 1.00) pEscalar = 1.00;
        
        asoma += (double)((acos( pEscalar))*(180.0/PI));
        
        
     }     
     asoma /=(w*h);
     esoma /= (w*h);     
     
     image.draw_quiver(field,&white,8,-14,0,0.7f).save("correctField.pgm");//salva o campo correto
     
     *aae = asoma;
     *error = esoma;
     return true;
}

/**
  * Função main    
*/
int main(int argc, char *argv[])
{
    
    if (argc < 6)
    {
       printf("Uso: %s <img1> <img2> alpha espacamento nIteracoes <correctField (opcional)>\n",argv[0]);
       system("pause");
       exit(0);
    }
        
    const CImg<>image1(argv[1]);//imagem 1
    
    const CImg<>image2(argv[2]);//imagem 2
    CImg<>image3(image1.dimx(), image1.dimy()); //imagem de saida    
            
    CImg<>field(image1.dimx(), image1.dimy(),1,2,3); //campo de vetores obtido
            
    double* Ex, * Ey, * Et, * U, * V, * im1, * im2, * U1, * U2, * V1, * V2;
    int width = image1.dimx(), height = image1.dimy();//dimensões da imagem
    int iterations = atoi(argv[5]);//numero de iteracoes
    int place = 0;//posição no vetor da imagem
    double udash=0.0, vdash=0.0, multiplier=0.0,temp=0.0;//variáveis auxiliares para o algoritmo
	double Uerr = 0.0, Verr = 0.0;//erro medio entre as iterações

	double alpha = (double)atof(argv[3]);//parâmetro alpha
	int espacamento = atoi(argv[4]);//especifica o espaçamento entre os vetores na visualização do campo.
                                    //quanto menor este valor, mais denso é o campo visualizado   
            
    double* rgb1 = im1, *rgb2 = im2;//vetores auxiliares

    int size = width*height;//tamanho da imagem em pixels
    
    //alocação de memoria
    Ex = new double[size];   
    Ey = new double[size];
    Et = new double[size];
    U = new double[size];
    V = new double[size];
    U1 = new double[size];
    V1 = new double[size];
    U2 = new double[size];
    V2 = new double[size];
    
    im1 = new double[size];
    im2 = new double[size];
    
    //inicialização    
    unsigned long int cont = 0;    
    for (int lin = 0; lin < height; lin++)
    {
        for (int col = 0; col < width; col++)
        {
            im1[cont] = image1(col, lin);
            im2[cont] = image2(col, lin);
            cont++;            
        }
    }    
    
    for (int i = 0; i < size; i++)
    {
        Ex[i] = Ey[i] = Et[i] = U[i] = V[i] = U1[i] = V1[i] = U2[i] = V2[i] = 0;
    }
    
    //calculo de derivadas parciais
    for (int j = 2; j < height-2; j++)
    {
		for (int i = 2; i < width-2; i++)
        {
            Ex[j*width+i] = getEx(im1, im2, i,j,width);
			Ey[j*width+i] = getEy(im1, im2, i,j,width);
			Et[j*width+i] = getEt(im1, im2, i,j,width);
		}
	}
	
    //preenche arestas das derivadas parciais
	for(int n=2; n<width-2; n++)
    {
		Ex[n]=Ex[width+n]=Ex[2*width+n];
		Ey[n]=Ey[width+n]=Ey[2*width+n];
		Et[n]=Et[width+n]=Et[2*width+n];

		Ex[(height-1)*width+n]=Ex[(height-2)*width+n]=Ex[(height-3)*width+n];
		Ey[(height-1)*width+n]=Ey[(height-2)*width+n]=Ey[(height-3)*width+n];
		Et[(height-1)*width+n]=Et[(height-2)*width+n]=Et[(height-3)*width+n];
	}
	
    for(int m=0; m<height; m++)
    {
		Ex[m*width]=Ex[m*width+1]=Ex[m*width+2];
		Ey[m*width]=Ey[m*width+1]=Ey[m*width+2];
		Et[m*width]=Et[m*width+1]=Et[m*width+2];

		Ex[m*width+width-1]=Ex[m*width+width-2]=Ex[m*width+width-3];
		Ey[m*width+width-1]=Ey[m*width+width-2]=Ey[m*width+width-3];
		Et[m*width+width-1]=Et[m*width+width-2]=Et[m*width+width-3];
	}
	
	//Loop principal 
 	 for(int iter=0; iter < iterations; iter++)
     {
	   for (int j = 1; j < height-1; j++)
       {
	   	  for (int i = 1; i < width-1; i++)
          {
             place= i+j*width;
		   	 udash = (U[i-1+j*width]+U[i+1+j*width]+U[i+(j-1)*width]+U[i+(j+1)*width])/4.0f;
		     vdash = (V[i-1+j*width]+V[i+1+j*width]+V[i+(j-1)*width]+V[i+(j+1)*width])/4.0f;
		   	 multiplier = (Ex[place]*udash + Ey[place]*vdash + Et[place]) / ((alpha)*(alpha) + Ex[place]*Ex[place] + Ey[place]*Ey[place]);

             temp = U2[place];
			 U2[place] = udash - (multiplier * Ex[place]);
			 Uerr += (U2[place]-temp)*(U2[place]-temp);

			 temp = V2[place];
			 V2[place] = vdash - (multiplier * Ey[place]);
			 Verr += (V2[place]-temp)*(V2[place]-temp);
		  }
		}
		
		for (int j = 0; j < height; j++)
        {
		   for (int i = 0; i < width; i++)
           {
				U[i+j*width] = U2[i+j*width];
				V[i+j*width] = V2[i+j*width];
			}
		}
		//Preenche os vetores U e V (bordas)
		for(int n=1; n<width-1; n++)
        {
			U[n]=U[width+n];
			V[n]=V[width+n];

			U[(height-1)*width+n]=U[(height-2)*width+n];
			V[(height-1)*width+n]=V[(height-2)*width+n];
		}
		
        for(int m=0; m<height; m++)
        {
			U[m*width]=U[m*width+1];
			V[m*width]=V[m*width+1];

			U[m*width+width-1]=U[m*width+width-2];
			V[m*width+width-1]=V[m*width+width-2];
		}
		//exibe os erros médios
		Uerr /= (double)((height-2)*(width-2));
		Verr /= (double)((height-2)*(width-2));
		fprintf(stderr,"errors u:%e v:%e iter:%d\n",Uerr,Verr,iter);
		Uerr=0.0;
		Verr=0.0;
	}//fim loop principal
	
	
	char estimatedField[540];
	//grava o campo gerado em arquivo
	sprintf(estimatedField,"%s.estimated.field",argv[1]);	
	FILE * opticField = fopen(estimatedField,"w");	
	fprintf(opticField,"%s %s\n",argv[1],argv[2]);
	fprintf(opticField,"%d %d\n",image1.dimx(),image1.dimy());
	//preenche a imagem do campo
	cont = 0;
	for (int lin = 0; lin < height; lin++)
    {
        for (int col = 0; col < width; col++)
        {
            field(col, lin,0) = U[cont];
            field(col, lin,1) = V[cont];
            fprintf(opticField,"%d %d %.10f %.10f\n",lin, col,U[cont],V[cont]);            
            cont++;
        }
    }
    //fim da gravação
    fclose(opticField);
    
    //exibe a imagem com os vetores
    image3.draw_quiver(field,&white,espacamento,-14,0,0.7f).display(0);
    //salva a imagem com os vetores
    char outstr[512];    
    sprintf(outstr,"%s.%d.field.pgm",argv[1],(int)(alpha));    
    image3.draw_quiver(field,&white,espacamento,-14,0,0.7f).save(outstr);
    
    // se foi especificado um arquivo com os vetores corretos, calcula os erros    
    if (argc >=7)
    {
       double aae, error;           
       if (calcError(argv[6], estimatedField, &aae, &error) == true)
       {    
         printf("o erro angular medio eh: %f\n",aae);
         printf("o erro absoluto medio eh: %f\n",error);
        }
       else printf("Nao foi possivel calcular os erros\n"); 
    }
    
    //liberação de memória   
    delete [] Ex;
    delete [] Ey;
    delete [] Et;
    delete [] U;
    delete [] V;
    delete [] im1;
    delete [] im2;
    delete [] U1;
    delete [] V1;
    delete [] U2;
    delete [] V2;
           
    system("PAUSE");
    return EXIT_SUCCESS;
}
