/*----------------------------MegaWave2 module----------------------------*/
/* mwcommand
name = {tv_besov};
version = {"1.0"};
author = {"Triet Le"};
function = {"Image decomposition using TV(u) + lambda*||G*(f-u)||_1. Example of commands to run the program:"};
usage = {
'w':[w=2.0]->w                "lambda, weight on fidelity term (default = 2.0)",
's':[teta=1.0]->teta		"Choices on the sime group: teta=1 is for heat kernel, teta=0.5 is for Poisson kernel (default = 1.0)",
'b':[beta=1e-10]->beta    "beta in the tv term (default = 1e-10)",
't':[k=-240]->k			"t = 1.05^k is the scale for the kernel (heat or Poisson, ...) (default = -240)",
n->n                  		"perform a fixed number of iterations",
 f->f                           	"input fimage",
 u<-u                           	"output fimage u (TV part)",
 v<-v                           	"outut fimage v (texture part)"
        };
*/
/*
1) command = tv_besov 500 barbara barbaara_u barbara_v
Here w, teta, beta and k all have default values, n=500, f=barbara, u=barbara_u and v=barbara_v
2) command = tv_besov -w 1.5 -t -45 500 barbara barbara_u barbara_v 
Here w=1.5, teta has default value, beta has defualt value, k=-45, n=500, f=barbara, u=barbara_u and v=barbara_v
*/
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "mw.h"

extern void fft2d();

const double PI = 3.1415926535;
const double TAU = 1.05;

double sign(x)
	double x;
{
	if(x < 0)
		return -1;
	else if(x > 0)
		return 1;
	else
		return 0;
}

double lp_norm(f, p)
Fimage f;
double p;
{
	int i, nx, ny;
	double norm;
	nx = f->ncol; ny = f->nrow;
	norm = 0;
	for(i=0; i<nx*ny; i++) {
		norm += pow(fabs(f->gray[i]), p);
	}
	norm = pow(norm/(nx*ny), 1.0/p);
	return norm;
}


void reflect(u)
/*
Reflecting the boundaryr of an image
*/
Fimage u;
{
	int i, j, nx, ny;
	nx = u->ncol; ny = u->nrow;
	for(i=0; i<nx; i++) {
		u->gray[i] = u->gray[nx + i];
		u->gray[(ny-1)*nx + i] = u->gray[(ny-2)*nx + i];
	}
	for(j=0; j<ny; j++) {
		u->gray[j*nx] = u->gray[j*nx + 1];
		u->gray[j*nx + nx-1] = u->gray[j*nx + nx-2];
	}
}

void discretize_tv(u, i, j, iter, beta, c1, c2, c3, c4)
Fimage u;
int i, j, iter;
double beta;
double *c1, *c2, *c3, *c4;
{
	int type = iter%4;
	int nx, ny;
	double ux, uy;
	nx = u->ncol; ny = u->nrow;
	if(beta == 0)
		beta = 1e-10;
	if(type == 0) {
		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[(j+1)*nx + i-1] - u->gray[j*nx+i-1];
		*c2 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c3 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[(j-1)*nx+i+1] - u->gray[(j-1)*nx+i];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c4 = 1.0/sqrt(ux*ux + uy*uy + beta);
	}
	else if(type == 1) {
		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[j*nx+i-1] - u->gray[(j-1)*nx+i-1];
		*c2 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[(j+1)*nx+i+1] - u->gray[(j+1)*nx+i];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c3 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c4 = 1.0/sqrt(ux*ux + uy*uy + beta);
	}
	else if(type == 2) {
		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[(j+1)*nx+i+1] - u->gray[j*nx+i+1];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c2 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c3 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[(j-1)*nx+i] - u->gray[(j-1)*nx + i-1];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c4 = 1.0/sqrt(ux*ux + uy*uy + beta);
	}
	else {
		ux = u->gray[j*nx+i+1] - u->gray[j*nx+i];
		uy = u->gray[j*nx+i+1] - u->gray[(j-1)*nx+i+1];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c2 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[(j+1)*nx+i] - u->gray[(j+1)*nx+i-1];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c3 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+i-1];
		uy = u->gray[j*nx+i] - u->gray[(j-1)*nx+i];
		*c4 = 1.0/sqrt(ux*ux + uy*uy + beta);
	}
}

void fftshift(f)
		Fimage f;
{
	int i, j, nx, ny;
	double a1, a2, a3, a4;
	nx = f->ncol; ny = f->nrow;
	for(i=0; i<nx/2; i++) {
		for(j=0; j<ny/2; j++) {
			a1 = f->gray[(j+ny/2)*nx + (i+nx/2)];
			a2 = f->gray[j*nx+i];
			a3 = f->gray[j*nx + (i+nx/2)];
			a4 = f->gray[(j+ny/2)*nx + i];
			f->gray[j*nx+i] = a1;
			f->gray[(j+ny/2)*nx + (i+nx/2)] = a2;
			f->gray[(j+ny/2)*nx + i] = a3;
			f->gray[j*nx + (i+nx/2)] = a4;
		}
	}
}


void make_periodic(im_in, im_out)
/*
im_out is the periodic extension and twice the size of im_in.
*/
Fimage im_in, im_out;
{
	int i, j, k, l, nx, ny;
	nx = im_in->ncol; ny = im_in->nrow;
	/*
	if(mw_change_fimage(im_out, 2*ny, 2*nx) == NULL)
	mwerror(FATAL, 1,
	"make_periodic --> Not enough memory to allocate the periodic image\n");
	*/
	for(i=0; i<2*nx; i++) {
		k = i%nx;
		if(i >= nx)
			k = nx-1 - k;
		for(j=0; j<2*ny; j++) {
			l = j%ny;
			if(j >= ny)
				l = ny-1 - l;
			im_out->gray[j*(2*nx)+i] = im_in->gray[l*nx+k];
		}
	}
}

/*----------------------- MAIN MODULE ---------------*/

void tv_besov(f,u,v,n,w, k, teta,beta)
Fimage f,u,v;
int n;
int *k;
double *w;
double *teta, *beta;
{
	const double PI = 3.14159265;
	int i, j, nx, ny, index, iter, NX, NY, im1, jm1;
	double coef, lambda, val, t;
	double c1, c2, c3, c4, x, y, norm;
	const double dt = 1.0;
	const double p = 1.0;
	Fimage per_v, per_re1, per_im1, per_re2, per_im2, per_K_v, per_K;
	nx = f->ncol; ny = f->nrow;
	NX = 2*nx; NY = 2*ny;
	
	per_v = mw_change_fimage(NULL, NY, NX);
	per_re1 = mw_change_fimage(NULL, NY, NX);
	per_im1 = mw_change_fimage(NULL, NY, NX);
	per_re2 = mw_change_fimage(NULL, NY, NX);
	per_im2 = mw_change_fimage(NULL, NY, NX);
	per_K_v = mw_change_fimage(NULL, NY, NX);
	per_K = mw_change_fimage(NULL, NY, NX);
	u = mw_change_fimage(u, ny, nx);
	v = mw_change_fimage(v,ny,nx);
	if(!v || !per_v || !per_re1 || !per_im1 || !per_re2 || !per_im2 || !u || !per_K_v ||!per_K) {
		if(v) mw_delete_fimage(v);
		if(per_v) mw_delete_fimage(per_v);
		if(per_re1) mw_delete_fimage(per_re1);
		if(per_im1) mw_delete_fimage(per_im1);
		if(per_re2) mw_delete_fimage(per_re2);
		if(per_im2) mw_delete_fimage(per_im2);
		if(u) mw_delete_fimage(u);
		if(per_K_v) mw_delete_fimage(per_K_v);
		if(per_K) mw_delete_fimage(per_K);
		mwerror(FATAL, 1, "Not enough memory.\n");
	}
	
	/*Computing the kernel*/
	t = pow(TAU,*k);
	for(j=0; j<NY; j++) {
		y = (j - (NY)/2);
		for(i=0; i<NX; i++) {
			x = (i - (NX)/2);
			per_K->gray[j*NX+i] = exp(-2*PI*t*pow(x*x + y*y, *teta));
		}
	}
	/*--- the main loop ---*/
	reflect(f);
	mw_copy_fimage(f,u);
	lambda = *w;
	for(iter=0; iter<n; iter++) {
		/*---Computing per_v---*/
		for(i=0; i<nx*ny; i++) {
			v->gray[i] = f->gray[i] - u->gray[i];
		}
		reflect(v);
		make_periodic(v,per_v);
		fft2d(per_v, NULL, per_re1, per_im1, 0);
		fftshift(per_re1); fftshift(per_im1);
		for(j=0; j<NY; j++) {
			y = (j - (NY)/2);
			for(i=0; i<NX; i++) {
				x = (i - (NX)/2);
				per_re2->gray[j*NX+i] = per_re1->gray[j*NX+i]*per_K->gray[j*NX+i];
				per_im2->gray[j*NX+i] = per_im1->gray[j*NX+i]*per_K->gray[j*NX+i];
			}
		}
		fftshift(per_re2); fftshift(per_im2);
		fft2d(per_re2, per_im2, per_v, NULL, 1);
		
		for(j=0; j<ny; j++) {
			for(i=0; i<nx; i++) {
				v->gray[j*nx+i] = sign(per_v->gray[j*NX+i]); 
			}
		}
		make_periodic(v,per_v);
		fft2d(per_v, NULL, per_re1, per_im1, 0);
		fftshift(per_re1); fftshift(per_im1);
		for(j=0; j<NY; j++) {
			y = (j - (NY)/2);
			for(i=0; i<NX; i++) {
				x = (i - (NX)/2);
				per_re2->gray[j*NX+i] = per_re1->gray[j*NX+i]*per_K->gray[j*NX+i];
				per_im2->gray[j*NX+i] = per_im1->gray[j*NX+i]*per_K->gray[j*NX+i];
			}
		}
		fftshift(per_re2); fftshift(per_im2);
		fft2d(per_re2, per_im2, per_K_v, NULL, 1);
		/*---------------------*/
		
		if(iter%10 == 0) {
			printf("iter = %d\n", iter);
		}
		
		for(j=1; j<ny-1; j++) {
			for(i=1; i<nx-1; i++) {
				discretize_tv(u, i, j, iter, *beta, &c1, &c2, &c3, &c4);
				coef = 1.0/(1.0 + dt*(c1+c2+c3+c4));
				u->gray[j*nx+i] = coef*(u->gray[j*nx+i] + dt*(c1*(u->gray[j*nx+i+1]) +
						c2*(u->gray[j*nx+i-1]) + c3*(u->gray[(j+1)*nx+i]) +
						c4*(u->gray[(j-1)*nx+i]) + lambda*(per_K_v->gray[j*NX+i])));
			}
		}
		reflect(u);
	}
	for(i=0; i<nx*ny; i++) {
		v->gray[i] = f->gray[i] - u->gray[i];
	}
	reflect(v);
	/*----Deleting allocated memory--------*/
	mw_delete_fimage(per_v);
	mw_delete_fimage(per_re1);
	mw_delete_fimage(per_im1);
	mw_delete_fimage(per_re2);
	mw_delete_fimage(per_im2);
	mw_delete_fimage(per_K_v);
	mw_delete_fimage(per_K);
}