// グレイスケール画像にFFTを使って画像透かしを入れる（取り出す）プログラム
// プログラム名：fft_watermark.cpp

#include "eps-header.hh"
#include "fft-header.hh"

const double WATERMARK_COEF = 10.0;	// 透かしを入れる係数(ε)

// 元画像に透かしを入れる関数
//
// （処理の流れ)                振幅
//  ┏━━━━┓ FFT┌────┐成分┏━━━━┓       ┌────┐逆FFT┏━━━━┓
//  ┃f(x, y) ┃ ⇒ │F(u, v) │ ＋ ┃g(x, y) ┃×ε ⇒│H(u, v) │  ⇒ ┃ h(x, y)┃
//  ┗━━━━┛    └────┘    ┗━━━━┛       └────┘     ┗━━━━┛
//    基準画像        基準画像       透かし画像       透かし混入画像   透かし混入画像
//   (空間領域)     (周波数領域)     (空間領域)        (周波数領域)      (空間領域)
//
void input_watermark(
	class image& src,		// 元画像クラス(入力)
	class image& mrk,		// 透かし画像クラス(入力)
	double eps,				// 透かし画像を入れる係数(入力)
	class image& dst)		// 処理済画像（元画像に透かしが入っている画像）クラス(出力)
{
	if (src.p == NULL) {
		printf("元画像がありません。");
		return;
	}
	if (mrk.p == NULL) {
		printf("透かし画像がありません。");
		return;
	}

	if (src.get_nr_lines() != mrk.get_nr_lines()) {
		printf("元画像と透かし画像のライン数が異なります。\n");
		return;
	}
	if (src.get_nr_pixels() != mrk.get_nr_pixels()) {
		printf("元画像と透かし画像のピクセル数が異なります。\n");
		return;
	}

	printf("\nグレイスケール画像に透かし画像を入れます。\n");

	int src_nr_lines  = src.get_nr_lines();		// 元画像のライン数   (縦の画素数)
	int src_nr_pixels = src.get_nr_pixels();	// 元画像のピクセル数 (横の画素数)

	int dst_nr_lines  = src_nr_lines;			// 処理済画像のライン数
	int dst_nr_pixels = src_nr_pixels;			// 処理済画像のピクセル数
	dst.gmake(dst_nr_lines, dst_nr_pixels);		// 処理済画像（元画像に透かしが入っている画像）の配列を確保

	class fourier src_ft(src_nr_lines, src_nr_pixels);	// 元画像用高速フーリエ変換クラスの宣言

	// 実数部を元画像の輝度値，虚数部を0.0とする複素数を２次元配列に入れる。
	for (int iy = 0; iy < src_nr_lines; iy ++) {
		for (int ix = 0; ix < src_nr_pixels; ix ++) {
			src_ft.p[iy][ix] = creal8((double)src.p[iy][ix], 0.0);
		}
	}
	src_ft.Forward_FFT();		// 元画像を２次元高速フーリエ変換

	for (int iy = 0; iy < src_nr_lines; iy ++) {
		for (int ix = 0; ix < src_nr_pixels; ix ++) {
			double re    = real(src_ft.p[iy][ix]);
			double im    = imag(src_ft.p[iy][ix]);
			double amp   = sqrt(re * re + im * im);		// 元画像の振幅
			double phase = arg(src_ft.p[iy][ix]);		// 元画像の位相

			src_ft.p[iy][ix] = (amp + eps * mrk.p[iy][ix]) * exp(creal8(0.0, phase));	// 位相を保存し，振幅に透かし画像を混入
		}
	}

	src_ft.Backward_FFT();		// ２次元高速フーリエ逆変換し処理済画像とする。
	for (int iy = 0; iy < src_nr_lines; iy ++) {
		for (int ix = 0; ix < src_nr_pixels; ix ++) {
			src_ft.p[iy][ix] /= (src_nr_lines * src_nr_pixels);
		}
	}

	src_ft.MakeImage(
		OUTPUT_AMP,
		FREQ_SPACE_SHUFFLING_OFF,
		OUTPUT_SCALE_OFF,
		dst.p);
}

// 元画像から透かし画像を取り出す関数
//
// （処理の流れ)
//  ┏━━━━┓    ┏━━━━┓
//  ┃h(x, y) ┃    ┃f(x, y) ┃
//  ┗━━━━┛    ┗━━━━┛
// 透かし混入画像     基準画像
//   (空間領域)      (空間領域)
//    FFT ↓    振幅  FFT ↓
//  ┌────┐成分┌────┐          ┏━━━━┓
// (│H(u, v) │ − │F(u, v) │) ÷ε ⇒ ┃g(x, y) ┃
//  └────┘    └────┘          ┗━━━━┛
// 透かし混入画像      基準画像            透かし画像
//   (周波数領域)    (周波数領域)          (空間領域)
//
void output_watermark(
	class image& src,		// 元画像クラス(入力)
	class image& org,		// 基準画像クラス(入力)
	double eps,				// 透かし画像を入れたときの係数(入力)
	class image& dst)		// 透かし画像クラス(出力)
{
	if (src.p == NULL) {
		printf("元画像がありません。");
		return;
	}
	if (org.p == NULL) {
		printf("基準画像がありません。");
		return;
	}

	if (src.get_nr_lines() != org.get_nr_lines()) {
		printf("元画像と基準画像のライン数が異なります。\n");
		return;
	}
	if (src.get_nr_pixels() != org.get_nr_pixels()) {
		printf("元画像と基準画像のピクセル数が異なります。\n");
		return;
	}

	printf("\nグレイスケール画像から透かし画像を取り出します。\n");

	int src_nr_lines  = src.get_nr_lines();		// 元画像のライン数   (縦の画素数)
	int src_nr_pixels = src.get_nr_pixels();	// 元画像のピクセル数 (横の画素数)

	int dst_nr_lines  = src_nr_lines;			// 透かし画像のライン数
	int dst_nr_pixels = src_nr_pixels;			// 透かし画像のピクセル数
	dst.gmake(dst_nr_lines, dst_nr_pixels);		// 透かし画像の配列を確保

	class fourier src_ft(src_nr_lines, src_nr_pixels);	// 元画像用高速フーリエ変換クラスの宣言
	class fourier org_ft(src_nr_lines, src_nr_pixels);	// 基準画像用高速フーリエ変換クラスの宣言

	// 実数部を元画像１，２の輝度値，虚数部を0.0とする複素数を２次元配列に入れる。
	for (int iy = 0; iy < src_nr_lines; iy ++) {
		for (int ix = 0; ix < src_nr_pixels; ix ++) {
			src_ft.p[iy][ix] = creal8((double)src.p[iy][ix], 0.0);
			org_ft.p[iy][ix] = creal8((double)org.p[iy][ix], 0.0);
		}
	}
	src_ft.Forward_FFT();		// 元画像を２次元高速フーリエ変換
	org_ft.Forward_FFT();		// 基準画像を２次元高速フーリエ変換

	for (int iy = 0; iy < src_nr_lines; iy ++) {
		for (int ix = 0; ix < src_nr_pixels; ix ++) {
			double re1    = real(src_ft.p[iy][ix]);
			double im1    = imag(src_ft.p[iy][ix]);
			double amp1   = sqrt(re1 * re1 + im1 * im1);		// 元画像の振幅

			double re2    = real(org_ft.p[iy][ix]);
			double im2    = imag(org_ft.p[iy][ix]);
			double amp2   = sqrt(re2 * re2 + im2 * im2);		// 基準画像の振幅

			double a = (amp1 - amp2) / eps;						// 透かし画像を取り出す。
			if (a < MIN_GRAY_LEVEL) a = MIN_GRAY_LEVEL;
			if (a > MAX_GRAY_LEVEL) a = MAX_GRAY_LEVEL;
			dst.p[iy][ix] = (unsigned char)a;
		}
	}
}

// プログラムの開始位置
int main()
{
	int watermark_flag;

	printf("1 : 画像透かしを入れる。\n");
	printf("2 : 画像透かしを取り出す。\n");
	printf("数値を入力してください。 ");
	char linebuf[MAX_FILENAME_LEN];
	fgets(linebuf, MAX_FILENAME_LEN - 1, esp_in);
	sscanf(linebuf, "%d", &watermark_flag);

	if (watermark_flag == 1) {
		class image src;	// 元画像クラスを宣言
		class image mrk;	// 透かし画像クラスを宣言
		class image dst;	// 処理済画像（元画像に透かしが入っている画像）クラスを宣言

		printf("\n元画像に画像透かしを入れます。\n");

		printf("\n元画像を読み込みます。\n");
		src.gload();		// 元画像をBMPファイルから読み込む
							// BMPファイルを指定したい場合は，
							//   src.gload("パス付きBMPファイル名");
							// とする。
		printf("\n透かし画像を読み込みます。\n");
		mrk.gload();		// 透かし画像をBMPファイルから読み込む
							// BMPファイルを指定したい場合は，
							//   mrk.gload("パス付きBMPファイル名");
							// とする。

		// 元画像に画像透かしを入れる
		input_watermark(
			src,			// 元画像クラス(入力)
			mrk,			// 透かし画像クラス(入力)
			WATERMARK_COEF,	// 透かしを入れる係数(入力)
			dst);			// 処理済画像（元画像に透かしが入っている画像）クラス(出力)

		printf("\n元画像に透かしが入っている画像を保存します。\n");
		dst.gsave();		// 処理済画像（元画像に透かしが入っている画像）をBMPファイルに保存する
							// BMPファイルを指定したい場合は，
							//   dst.gsave("パス付きBMPファイル名");
							// とする。

	} else if (watermark_flag == 2) {
		class image src;	// 透かし画像が入っている元画像クラスを宣言
		class image org;	// 基準画像クラスを宣言
		class image dst;	// 取り出した透かし画像クラスを宣言

		printf("\n元画像から画像透かしを取り出します。\n");

		printf("\n透かし入りの元画像を読み込みます。\n");
		src.gload();		// 透かし入りの元画像をBMPファイルから読み込む
							// BMPファイルを指定したい場合は，
							//   src.gload("パス付きBMPファイル名");
							// とする。
		printf("\n基準画像を読み込みます。\n");
		org.gload();		// 基準画像をBMPファイルから読み込む
							// BMPファイルを指定したい場合は，
							//   org.gload("パス付きBMPファイル名");
							// とする。

		// 元画像から透かし画像を取り出す
		output_watermark(
			src,			// 元画像クラス(入力)
			org,			// 基準画像クラス(入力)
			WATERMARK_COEF,	// 透かしを入れたときの係数(入力)
			dst);			// 取り出した透かし画像クラス(出力)


		printf("\n取り出した透かし画像を保存します。\n");
		dst.gsave();		// 取り出した透かし画像をBMPファイルに保存する
							// BMPファイルを指定したい場合は，
							//   dst.gsave("パス付きBMPファイル名");
							// とする。

	} else {
		printf("認識できない数値が入力されました。\n");
	}
}