/*
 * heat2d-i-eigen-sparse.cpp --- 2次元熱方程式を陰解法で解く
 *                境界にある格子点における値は未知数に含めない方法
 *                係数行列が対称帯行列であることを利用して解く
 *                (GLSC, Eigen を利用)
 *
 *   コンパイルするには (/opt/X11 に X があるとして)
 *     g++ -I/opt/X11/include heat2d-i-eigen-sparse.cpp -lglscd -L/opt/X11/lib -lX11
 *       (GLSCの場所も指定する必要があるかも…)
 *   あるいは頑張って cglsc++ スクリプトを作って
 *     cglsc++ heat2d-i-eigen-sparse.cpp
 *
 *  http://nalab.mind.meiji.ac.jp/~mk/program/fdm/heat2d-i-eigen-sparse.cpp
 *  http://nalab.mind.meiji.ac.jp/~mk/program/fdm/heat2d-i.c の焼き直し
 *  http://nalab.mind.meiji.ac.jp/~mk/labo/text/heat-fdm-1.pdf
 */

#include <iostream>
#include <iomanip>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <Eigen/SparseCholesky>

using namespace std;
using namespace Eigen;

typedef Eigen::Triplet<double> T; // こうしておかないと面倒

// GLSCのヘッダーを読み込む
#ifndef G_DOUBLE
#define G_DOUBLE
#endif
extern "C" {
#include <glsc.h>
};

// 格子点につける番号 (RowMajor を採用)
#define psi(i,j) (((j)-1)*mm+(i)-1)

int main(void)
{
  double a, b, c, d;
  int N_x, N_y, mm, NN, i, j, p, q, L, n, nMax;
  int skip;
  double h_x, h_y, lambda_x, lambda_y, lambda, lambda_limit, tau, Tmax, dt;
  double f(double, double), alpha(double, double);
  double theta, beta_0, beta_1, beta_2, beta_00, beta_10, beta_20;
  double x, y, t;

  /* 問題を考える区間 [a,b]×[c,d] */
  a = 0.0; b = 1.0; c = 0.0; d = 1.0;

  /* 区間の分割数 */
  cout << "Nx, Ny: "; cin >> N_x >> N_y;

  mm = N_x - 1;
  NN = (N_x - 1) * (N_y - 1);
  /* 空間の刻み幅 */
  h_x = (b - a) / N_x; h_y = (d - c) / N_y;

  /* θ法の重みの決定 */
  printf("θ (0≦θ≦1): "); scanf("%lf", &theta);

  if (theta == 1.0) {
    cout << "τ: "; cin >> tau;
  }
  else {
    cout << "τ(≦" << 0.5 / (1 - theta) / (1 / (h_x * h_x) + 1 / (h_y * h_y))
	 << "≡最大値ノルムに関する安定性条件を満たすτの上限): ";
    cin >> tau;
  }

  lambda_x = tau / (h_x * h_x); lambda_y = tau / (h_y * h_y);
  lambda = lambda_x + lambda_y;

  /* 最大値ノルムに関する安定性を満たすλの上限 */
  lambda_limit = 1.0 / (2.0 * (1.0 - theta));

  if (lambda > lambda_limit)
    cout << "注意: λ=" << lambda
	 << ">1/2(1-θ) となっています。" << endl;
  else
    cout << "λ=" << lambda << endl;

  /* 2次元格子上の数値データ、ベクトルを記憶するオブジェクトの用意 */
  Matrix <double, Dynamic, Dynamic, RowMajor> Uk(N_x + 1, N_y + 1);
  /* 初期値の設定 */
  for (i = 0; i <= N_x; i++) {
    x = a + i * h_x;
    for (j = 0; j <= N_y; j++)
      Uk(i, j) = f(x, c + j * h_y);
  }

  /* 連立1次方程式の係数行列に現れる非零成分 */
  beta_0 = 1.0 + 2.0 * theta * lambda;
  beta_1 = - theta * lambda_x;
  beta_2 = - theta * lambda_y;

  beta_00 = 1.0 - 2.0 * (1.0 - theta) * lambda;
  beta_10 = (1.0 - theta) * lambda_x;
  beta_20 = (1.0 - theta) * lambda_y;

  /********************* 係数行列の作成とCholesky分解 *********************/
  // 行列を記憶するオブジェクトと作成準備
  vector<T> tripletList;
  SparseMatrix<double> A(NN,NN);
  tripletList.reserve(5 * NN); // ちょっと多め
  // 連立１次方程式の右辺、解を記憶するオブジェクトの準備
  VectorXd B(NN), X(NN);
  // 係数行列の非零成分を収めたリスト作成
  for (i = 1; i < N_x; i++)
    for (j = 1; j < N_y; j++) {
      L = psi(i, j);
      tripletList.push_back(T(L,L, beta_0));
      if (i != 1)
	tripletList.push_back(T(L,L-1, beta_1));
      if (i != N_x - 1)
	tripletList.push_back(T(L,L+1, beta_1));
      if (j != 1)
	tripletList.push_back(T(L, L-mm, beta_2));
      if (j != N_y - 1)
	tripletList.push_back(T(L, L+mm, beta_2));
    }
  // リストから行列オブジェクトを作る
  A.setFromTriplets(tripletList.begin(), tripletList.end());
  // 係数行列の Cholesky 分解
  SimplicialLDLT<SparseMatrix<double> > solver;
  solver.compute(A);
  if (solver.info() != Success) {
    cerr << "分解に失敗" << endl;
    exit(1);
  }
  /* ********************* 係数行列の作成と Cholesky分解 終了 **********************/

  /* 連立1次方程式の係数行列を表示する (小さい時だけ) */
  if (NN < 20) {
    cout << "連立1次方程式の行列" << endl;
    for (p = 0; p < NN; p++) {
      for (q = 0; q < NN; q++)
	cout << fixed << setw(5) << setprecision(2)
	     << A.coeffRef(p,q);
      cout << endl;
    }
  }
  cout << "備考: 1+2θλ=" << fixed << setw(5) << setprecision(2)
       << beta_0 << ", -θλx=" << fixed << setw(5) << setprecision(2)
       << beta_1 << ", -θλy=" << fixed << setw(5) << setprecision(2)
       << beta_2 << endl;

  cout << "Tmax: "; cin >> Tmax;
  cout << "Δt: "; cin >> dt;
  skip = rint(dt / tau);
  if (skip == 0) {
    skip = 1;
  }
  dt = skip * tau;

  nMax = rint(Tmax / tau);

  /* グラフィックス・ライブラリィ GLSC の呼び出し */
  g_init((char *) "Meta", 250.0, 160.0);
  g_device(G_BOTH);
  g_def_scale(0, 0.0, 1.0, 0.0, 1.0, 30.0, 70.0, 100.0, 72.0);
  g_def_scale(4, -1.0, 1.0, -1.0, 1.0, 30.0, 30.0, 100.0, 100.0);
  g_def_line(0, G_BLACK, 0, G_LINE_SOLID);
  g_sel_scale(0);

  g_cls();
  g_hidden(1.0, 1.0, 0.4, -1.0, 1.0, 5.0, 25.0, 20.0, 20.0, 20.0,
	   150.0, 100.0, Uk.data(), N_x + 1, N_y + 1,
	   1, G_SIDE_NONE, 2, 1);


  /* 時間に関するループ */
  for (n = 1; n <= nMax; n++) {

    /* まず、素朴な連立1次方程式の右辺を用意する */
    /* 第 n ステップにおける値の寄与 */
    for (i = 1; i < N_x; i++)
      for (j = 1; j < N_y; j++) {
	L = psi(i, j);
	B(L) = beta_00 * Uk(i, j)
	  + beta_10 * (Uk(i + 1, j) + Uk(i - 1, j))
	  + beta_20 * (Uk(i, j + 1) + Uk(i, j - 1));
      }
    /* 第 n+1 ステップの境界上の格子点における値 (既知) の寄与 */
    for (j = 1; j < N_y; j++) {
      y = c + j * h_y;
      /* (1,j) のとき */
      L = psi(1, j);
      B(L) -= beta_1 * alpha(a, y);
      /* (N_x-1,j) のとき */
      L = psi(N_x - 1, j);
      B(L) -= beta_1 * alpha(b, y);
    }
    for (i = 1; i < N_x; i++) {
      x = a + i * h_x;
      /* (i,1) のとき */
      L = psi(i, 1);
      B(L) -= beta_2 * alpha(x, c);
      /* (i,N_y-1) のとき */
      L = psi(i, N_y - 1);
      B(L) -= beta_2 * alpha(x, d);
    }

    /* A vector_U = B を解く */
    X = solver.solve(B);
    /* 領域内部にある格子点での値 */
    for (i = 1; i < N_x; i++)
      for (j = 1; j < N_y; j++)
	Uk(i, j) = X(psi(i, j));

    /* 境界その1: 下の辺、上の辺にある格子点 (角の点も含める) */
    for (i = 0; i <= N_x; i++) {
      x = a + i * h_x;
      Uk(i, 0) = alpha(x, c);
      Uk(i, N_y) = alpha(x, d);
    }
    /* 境界その2: 左の辺、右の辺にある格子点 (角の点は含めない) */
    for (j = 1; j < N_y; j++) {
      y = c + j * h_y;
      Uk(0, j) = alpha(a, y);
      Uk(N_x, j) = alpha(b, y);
    }

    if (n % skip == 0) { // 時間間隔 dt で結果の表示
#ifdef  PRINT
      /* データを数値で表示 */
      for (i = 0; i <= N_x; i++) {
	for (j = 0; j <= N_y; j++)
	  cout << " " << fixed << setw(5) << serprecision(2)
	       << Uk(i, j);
	cout << endl;
      }
#endif
      /* 鳥瞰図を描く */
      g_cls();
      g_hidden(1.0, 1.0, 0.4, -1.0, 1.0, 5.0, 25.0, 20.0, 20.0, 20.0,
	       150.0, 100.0, Uk.data(), N_x + 1, N_y + 1,
	       1, G_SIDE_NONE, 2, 1);
    }
  }
  /* マウスでクリックされるのを待つ */
  g_sleep(-1.0);
  /* ウィンドウを消す */
  g_term();
  cout << endl;
  return 0;
}

/* 初期値 */
double f(double x, double y)
{
  /* ピラミッド型の関数 */
  if (y > 0.5)
    y = 1 - y;
  if (x > 0.5)
    x = 1 - x;
  if (y < x)
    return 5 * y;
  else
    return 5 * x;
}

/* 境界値 */
double alpha(double x, double y)
{
  /* 同次 Dirichlet 境界条件 */
  return 0.0;
}