/*----------------------------MegaWave2 module----------------------------*/
/* mwcommand
name = {tv_sobolev};
version = {"1.0"};
author = {"Triet Le"};
function = {"Image decomposition using J(u) = |u|_{BV} + lambda |I_s*(f-u)|_{L1}"};
usage = {
'w':[w=1.0]->w  		"weight on fidelity term (default 1.0)",
's':[s=-0.25]->s			"the parameter s<0 for the Riesz potential kernel I_s",
'b':[beta=1e-10]->beta   "beta in the tv term",
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)"
        };
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "mw.h"

extern void fft2d();

const double PI = 3.1415926535;

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

void reflect(u)
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, im1, jm1;
	double ux, uy;
	nx = u->ncol; ny = u->nrow;
	if(beta == 0)
		beta = 1e-10;
	im1 = i-1;
	jm1 = j-1;
	if(im1<0)
		im1 = nx-1;
	if(jm1<0)
		jm1 = ny-1;
	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+im1];
		uy = u->gray[(j+1)*nx + im1] - u->gray[j*nx+im1];
		*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[(jm1)*nx+i+1] - u->gray[(jm1)*nx+i];
		uy = u->gray[j*nx+i] - u->gray[(jm1)*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[(jm1)*nx+i];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[j*nx+i] - u->gray[j*nx+im1];
		uy = u->gray[j*nx+im1] - u->gray[(jm1)*nx+im1];
		*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[(jm1)*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+im1];
		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+im1];
		uy = u->gray[(j+1)*nx+i] - u->gray[j*nx+i];
		*c3 = 1.0/sqrt(ux*ux + uy*uy + beta);

		ux = u->gray[(jm1)*nx+i] - u->gray[(jm1)*nx + im1];
		uy = u->gray[j*nx+i] - u->gray[(jm1)*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[(jm1)*nx+i+1];
		*c1 = 1.0/sqrt(ux*ux + uy*uy + beta);

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

		ux = u->gray[(j+1)*nx+i] - u->gray[(j+1)*nx+im1];
		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+im1];
		uy = u->gray[j*nx+i] - u->gray[(jm1)*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)
		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_sobolev(f,u,v,n,w, s,beta)
Fimage f,u,v;
int n;
double *w;
double *s, *beta;
{
	const double PI = 3.14159265;
	int i, j, nx, ny, index, iter, NX, NY, im1, jm1;
	double coef, lambda, t, val;
	double c1, c2, c3, c4, x, y, norm;
	double dt = 1;
	Fimage per_v, per_re1, per_im1, per_re2, per_im2, per_K_v;
	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);
	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)
		mwerror(FATAL, 1, "Not enough memory.\n");

	
	lambda = *w;
	reflect(f);
	/*---Initialize u---*/
	mw_copy_fimage(f,u);
	

	
	/*--- the main loop ---*/
	for(iter=0; iter<n; iter++) {
		/*---Computing I_s*(sign(I_s*v))---*/
		for(j=0; j<ny; j++) {
			for(i=0; i<nx; i++) {
				v->gray[j*nx+i] = f->gray[j*nx+i] - u->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);
				if(x == 0 && y == 0) {
					per_re2->gray[j*NX+i] = 0;
					per_im2->gray[j*NX+i] = 0;
				}
				else {
					per_re2->gray[j*NX+i] = per_re1->gray[j*NX+i]*pow(2*PI*sqrt(x*x + y*y), *s);
					per_im2->gray[j*NX+i] = per_im1->gray[j*NX+i]*pow(2*PI*sqrt(x*x + y*y), *s);
				}
			}
		}
		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);
				if(x == 0 && y == 0) {
					per_re2->gray[j*NX+i] = 0;
					per_im2->gray[j*NX+i] = 0;
				}
				else {
					per_re2->gray[j*NX+i] = per_re1->gray[j*NX+i]*pow(2*PI*sqrt(x*x + y*y), *s);
					per_im2->gray[j*NX+i] = per_im1->gray[j*NX+i]*pow(2*PI*sqrt(x*x + y*y), *s);
				}
			}
		}
		fftshift(per_re2);
		fftshift(per_im2);
		fft2d(per_re2, per_im2, per_K_v, NULL, 1);
		/*---------------------*/
		/*---Computing the next u---*/
		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(j=0; j<ny; j++) {
		for(i=0; i<nx; i++) {
			v->gray[j*nx+i] = f->gray[j*nx+i] - u->gray[j*nx+i] + 100;
		}
	}
	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);
}