/*
 * heat2n-i-eigen-sparse.cpp --- 2次元熱方程式(Neumann境界条件)を陰解法で解く
 *              (GLSC, Eigenを利用)
 * http://nalab.mind.meiji.ac.jp/~mk/program/fdm/heat2n-i-eigen-sparse.cpp
 *
 * To compile (/opt/X11 に X があるとして)
 *     g++ -I/opt/X11/include heat2n-i-eigen-sparse.cpp -lglscd -L/opt/X11/lib -lX11
 *       (GLSCの場所も指定する必要があるかも…)
 *   あるいは頑張って cglsc++ スクリプトを作って
 *     cglsc++ heat2n-i-eigen-sparce.cpp
 *
 *  http://nalab.mind.meiji.ac.jp/~mk/program/fdm/heat2n-i.c の焼き直し
 *  http://nalab.mind.meiji.ac.jp/~mk/program/ から入手可能
 *  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 phi(i,j) (j)*m+(i)

int main(void)
{
  double a, b, c, d;
  int N_x, N_y, m, N, 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), Phi(double, double, double);
  double theta, gamma, alpha, beta, gamma_p, alpha_p, beta_p;
  double x, y, t, t_p;

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

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

  m = N_x + 1;
  N = (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次方程式の係数行列に現れる非零成分
     γU_{ij}^{n+1}
         +α(U_{i+1,j}^{n+1}+U_{i-1,j}^{n+1})
         +β(U_{i,j+1}^{n+1}+U_{i,j-1}^{n+1})
      =γ'U_{ij}^n
         +α'(U_{i+1,j}^n+U_{i-1,j}^n)
         +β'(U_{i,j+1}^n+U_{i,j-1}^n)
     と書いたときのα,β,γ,α',β',γ' */
  gamma = 1.0 + 2.0 * theta * lambda;
  alpha = - theta * lambda_x;
  beta  = - theta * lambda_y;

  gamma_p = 1.0 - 2.0 * (1.0 - theta) * lambda;
  alpha_p = (1.0 - theta) * lambda_x;
  beta_p  = (1.0 - theta) * lambda_y;

  /********************* 係数行列の作成とCholesky分解 *********************/
  // 行列を記憶するオブジェクトと作成準備
  vector<T> tripletList;
  SparseMatrix<double> A(N,N);
  tripletList.reserve(5 * N); // ちょっと多め
  // 連立１次方程式の右辺、解を記憶するオブジェクトの準備
  VectorXd B(N), X(N);
  // 係数行列の非零成分を収めたリスト作成
  for (i = 1; i < N_x; i++)
    for (j = 1; j < N_y; j++) {
      L = phi(i, j);
      tripletList.push_back(T(L,L-m,beta));
      tripletList.push_back(T(L,L-1,alpha));
      tripletList.push_back(T(L,L,gamma));
      tripletList.push_back(T(L,L+1,alpha));
      tripletList.push_back(T(L,L+m,beta));
    }

  /* 左の辺 (角点除く) */
  for (j = 1; j < N_y; j++) {
    /* (0,j) */
    L = phi(0,j);
    tripletList.push_back(T(L,L-m,beta / 2.0));
    tripletList.push_back(T(L,L,gamma / 2.0));
    tripletList.push_back(T(L,L+1,alpha));
    tripletList.push_back(T(L,L+m,beta / 2.0));
  }
  /* 右の辺 (角点除く) */
  for (j = 1; j < N_y; j++) {
    /* (N_x,j) */
    L = phi(N_x,j);
    tripletList.push_back(T(L,L-m, beta / 2.0));
    tripletList.push_back(T(L,L-1,alpha));
    tripletList.push_back(T(L,L,gamma / 2.0));
    tripletList.push_back(T(L,L+m,beta / 2.0));
  }
  /* 下の辺 (角点除く) */
  for (i = 1; i < N_x; i++) {
    /* (i,0) */
    L = phi(i,0);
    tripletList.push_back(T(L,L-1, alpha / 2.0));
    tripletList.push_back(T(L,L, gamma / 2.0));
    tripletList.push_back(T(L,L+1, alpha / 2.0));
    tripletList.push_back(T(L,L+m,beta));
  }
  /* 上の辺 (角点除く) */
  for (i = 1; i < N_x; i++) {
    /* (i,N_y) */
    L = phi(i,N_y);
    tripletList.push_back(T(L,L-m, beta));
    tripletList.push_back(T(L,L-1, alpha / 2.0));
    tripletList.push_back(T(L,L, gamma / 2.0));
    tripletList.push_back(T(L,L+1, alpha / 2.0));
  }
  /* 角点 */
  /* 左下 */
  L = phi(0,0);
  tripletList.push_back(T(L,L, gamma / 4.0));
  tripletList.push_back(T(L,L+1, alpha / 2.0));
  tripletList.push_back(T(L,L+m, beta / 2.0));
  /* 右下 */
  L = phi(N_x,0);
  tripletList.push_back(T(L,L, gamma / 4.0));
  tripletList.push_back(T(L,L-1, alpha / 2.0));
  tripletList.push_back(T(L,L+m, beta / 2.0));
  /* 左上 */
  L = phi(0,N_y);
  tripletList.push_back(T(L,L, gamma / 4.0));
  tripletList.push_back(T(L,L+1, alpha / 2.0));
  tripletList.push_back(T(L,L-m, beta / 2.0));
  /* 右上 */
  L = phi(N_x,N_y);
  tripletList.push_back(T(L,L, gamma / 4.0));
  tripletList.push_back(T(L,L-1, alpha / 2.0));
  tripletList.push_back(T(L,L-m, beta / 2.0));

  // リストから行列オブジェクトを作る
  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 (N < 20) {
    cout << "連立1次方程式の行列" << endl;
    for (p = 0; p < N; p++) {
      for (q = 0; q < N; q++)
        cout << fixed << setw(5) << setprecision(2)
             << A.coeffRef(p,q);
      cout << endl;
    }
  }

  cout << "備考: 1+2θλ=" << fixed << setw(5) << setprecision(2) << gamma
       << ", -θλx=" << fixed << setw(5) << setprecision(2) << alpha
       << ", -θλy=" << fixed << setw(5) << setprecision(2) << beta << 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++) {
    t = n * tau; t_p = (n - 1) * tau;
    /* まず、素朴な連立1次方程式の右辺を用意する */
    /* 内部の格子点 */
    for (i = 1; i < N_x; i++)
      for (j = 1; j < N_y; j++) {
	L = phi(i,j);
	B(L) = gamma_p * Uk(i,j)
	  + alpha_p * (Uk(i + 1,j) + Uk(i - 1,j))
	  + beta_p * (Uk(i,j + 1) + Uk(i,j - 1));
      }
    /* 以下、境界にある格子点での方程式を立てる。
     * 仮想格子点での値は境界条件を中心差分近似した方程式を
     * 用いて消去する
     *   ……右辺にも移項する量があるので、右辺について後で処理する */
    /* 下の辺、上の辺にある格子点 (角の点は含めない) */
    for (i = 1; i < N_x; i++) {
      x = a + i * h_x;
      /* (i, 0) */
      L = phi(i,0);
      B(L) = gamma_p * Uk(i,0)
	+ alpha_p * (Uk(i + 1,0) + Uk(i - 1,0))
	+ 2 * beta_p * Uk(i,1)
	+ 2 * h_y  * (beta * Phi(x, c, t) - beta_p * Phi(x, c, t_p));
      B(L) /= 2;
      /* (i, N_y) */
      L = phi(i,N_y);
      B(L) = gamma_p * Uk(i,N_y)
	+ alpha_p * (Uk(i + 1,N_y) + Uk(i - 1,N_y))
	+ 2 * beta_p * Uk(i,N_y - 1)
	+ 2 * h_y * (beta_p * Phi(x, d, t_p) - beta * Phi(x, d, t));
      B(L) /= 2;
    }
    /* 左の辺、右の辺にある格子点 (角の点は含めない) */
    for (j = 1; j < N_y; j++) {
      y = c + j * h_y;
      /* (0, j) */
      L = phi(0,j);
      B(L) = gamma_p * Uk(0,j)
	+ 2 * alpha_p * Uk(1,j)
	+ beta_p * (Uk(0,j + 1) + Uk(0,j - 1))
	+ 2 * h_x * (alpha * Phi(a, y, t) - alpha_p * Phi(a, y, t_p));
      B(L) /= 2;
      /* (N_x, j) */
      L = phi(N_x,j);
      B(L) = gamma_p * Uk(N_x,j)
	+ 2 * alpha_p * Uk(N_x - 1,j)
	+ beta_p * (Uk(N_x,j + 1) + Uk(N_x,j - 1))
	+ 2 * h_x * (alpha_p * Phi(b, y, t_p) - alpha * Phi(b, y, t));
      B(L) /= 2;
    }
    /* 左下 */
    L = phi(0,0);
    B(L) = gamma_p * Uk(0,0)
      + 2 * alpha_p * Uk(1,0) + 2 * beta_p * Uk(0,1)
      + 2 * h_x * (alpha * Phi(a,c,t) - alpha_p * Phi(a,c,t_p))
      + 2 * h_y * (beta * Phi(a,c,t) - beta_p * Phi(a,c,t_p));
    B(L) /= 4;
    /* 左上 */
    L = phi(0,N_y);
    B(L) = gamma_p * Uk(0,N_y)
      + 2 * alpha_p * Uk(1,N_y) + 2 * beta_p * Uk(0,N_y - 1)
      + 2 * h_x * (alpha * Phi(a,d,t) - alpha_p * Phi(a,d,t_p))
      + 2 * h_y * (beta_p * Phi(a,d,t_p) - beta * Phi(a,d,t));
    B(L) /= 4;
    /* 右下 */
    L = phi(N_x,0);
    B(L) = gamma_p * Uk(N_x,0)
      + 2 * alpha_p * Uk(N_x - 1,0) + 2 * beta_p * Uk(N_x,1)
      + 2 * h_x * (alpha_p * Phi(b,c,t_p) - alpha * Phi(b,c,t))
      + 2 * h_y * (beta * Phi(b,c,t) - beta_p * Phi(b,c,t_p));
    B(L) /= 4;
    /* 右上 */
    L = phi(N_x,N_y);
    B(L) = gamma_p * Uk(N_x,N_y)
      + 2 * alpha_p * Uk(N_x - 1,N_y) + 2 * beta_p * Uk(N_x,N_y - 1)
      + 2 * h_x * (alpha_p * Phi(b,d,t_p) - alpha * Phi(b,d,t))
      + 2 * h_y * (beta_p * Phi(b,d,t_p) - beta * Phi(b,d,t));
    B(L) /= 4;

    /* A vector_U = B を解く */
    X = solver.solve(B);
    /* */
    for (i = 0; i <= N_x; i++)
      for (j = 0; j <= N_y; j++)
	Uk(i,j) = X(phi(i,j));

    /* データを数値で表示 */
    if (n % skip == 0) {
#ifdef  PRINT
      for (i = 0; i <= N_x; i++) {
	for (j = 0; j <= N_y; j++)
	  cout << fixed << setw(5) << setprecision(1) << 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();

  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;
}

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