/*
 * heat2d-i-band3.c --- 2次元熱方程式を陰解法で解く (帯行列版)
 * http://nalab.mind.meiji.ac.jp/~mk/program/fdm/heat2d-i-band3.c
 * http://nalab.mind.meiji.ac.jp/~mk/program/fdm/smallmatrix.h
 * http://nalab.mind.meiji.ac.jp/~mk/program/linear/luband3.c
 * http://nalab.mind.meiji.ac.jp/~mk/program/linear/luband3.h

 * To compile
 *     cc -O3 -c luband3.c
 *     cglsc heat2d-i-band3.c luband3.o
 *   あるいは
 *     ccmg heat2d-i-band3.c luband3.o
 */

/*
 * 係数行列 Aは、半バンド幅mの帯行列であるので、それを利用して
 * 連立方程式を解くのを高速化+省メモリ化する。
 * つまり、係数行列の帯の部分を取り出して計算していく。
 * 元の係数行列を(a_{i,j})とすると、取り出した行列Aとの関係は,
 *   A[i][j] = a_{i,i+j-n}  (  m ≦ i+j < m+n )
 *           = 0            (  i+j < m , m+n ≦ i+j )
 * または、
 *   a_{i,j} = A[i][j-i+n] (  0 ≦ j-i+m < 2m+1 )
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#ifdef OLD
#include <matrix.h>
#else
#include "smallmatrix.h"
#endif
#include "luband3.h"
#ifndef G_DOUBLE
#define G_DOUBLE
#endif
#include <glsc.h>

#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, k, kmax;
    matrix Uk, A;
    double *B, *vector_U, cond;
    int skip;
    double h_x, h_y, lambda_x, lambda_y, lambda, lambda_limit, tau,
     Tmax, dt;
    double u0(double, double), alpha(double, double);
    double theta, beta_0, beta_1, beta_2, beta_00, beta_10, beta_20;
    double x, y, t;

    a = 0.0;
    b = 1.0;
    c = 0.0;
    d = 1.0;

    printf("Nx, Ny: ");
    scanf("%d %d", &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;

    if ((Uk = new_matrix(N_x + 1, N_y + 1)) == NULL) {
	fprintf(stderr, "数列 U^k を記憶する領域の確保に失敗\n");
	exit(1);
    }
    if ((A = new_matrix(N, (2 * m + 1))) == NULL) {
	fprintf(stderr, "係数行列 A を記憶する領域の確保に失敗\n");
	exit(1);
    }
    if ((B = malloc(sizeof(double) * N)) == NULL) {
	fprintf(stderr, "B を記憶する領域の確保に失敗\n");
	exit(1);
    }
    if ((vector_U = malloc(sizeof(double) * N)) == NULL) {
	fprintf(stderr, "vector_U を記憶する領域の確保に失敗\n");
	exit(1);
    }
    printf("θ (0≦θ≦1): ");
    scanf("%lf", &theta);

    if (theta == 1.0) {
	printf("τ: ");
	scanf("%lf", &tau);
    } else {
	printf("τ(≦%g): ",
	       0.5 / (1 - theta) / (1 / (h_x * h_x) + 1 / (h_y * h_y)));
	scanf("%lf", &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)
	printf("注意: λ=%g>1/2(1-θ) となっています。\n", lambda);

    for (i = 0; i <= N_x; i++)
	for (j = 0; j <= N_y; j++)
	    Uk[i][j] = u0(a + i * h_x, c + j * h_y);

    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;

    for (p = 0; p < N; p++) {
	for (q = 0; q < 2 * m + 1; q++)		/* N */
	    A[p][q] = 0.0;	/*A[p][q] = 0.0; */
	A[p][m] = 1.0;		/*A[p][p] = 1.0; */
    }
    for (i = 1; i < N_x; i++)
	for (j = 1; j < N_y; j++) {
	    L = phi(i, j);
	    A[L][L - m - L + m] = beta_2;	/*A[L][L - m] = beta_2; */
	    A[L][L - 1 - L + m] = beta_1;	/*A[L][L - 1] = beta_1; */
	    A[L][L - L + m] = beta_0;	/*A[L][L] = beta_0;    */
	    A[L][L + 1 - L + m] = beta_1;	/*A[L][L + 1] = beta_1; */
	    A[L][L + m - L + m] = beta_2;	/*A[L][L + m] = beta_2; */
	}

#ifdef  PRINT
    printf("素朴に作った連立1次方程式の行列\n");
    for (p = 0; p < N; p++) {
	for (q = 0; q < 2 * m + 1; q++)
	    printf(" %+7.5f", A[p][q]);
	printf("\n");
    }
#endif

    /* 対称化するための作業 1 */
    for (j = 1; j < N_y; j++) {
	/* (1,j) */
	L = phi(1, j);
	A[L][L - 1 - L + m] = 0.0;	/*A[L][L - 1] = 0.0; */
	/* (N_x-1,j) */
	L = phi(N_x - 1, j);
	A[L][L + 1 - L + m] = 0.0;	/*A[L][L + 1] = 0.0; */
    }

    /* 対称化するための作業 2 */
    for (i = 1; i < N_x; i++) {
	/* (i,1) */
	L = phi(i, 1);
	A[L][L - m - L + m] = 0.0;	/*A[L][L - m] = 0.0; */
	/* (i,N_y-1) */
	L = phi(i, N_y - 1);
	A[L][L + m - L + m] = 0.0;	/*A[L][L + m] = 0.0; */
    }

#ifdef  PRINT
    printf("対称化した連立1次方程式の行列\n");
    for (p = 0; p < N; p++) {
	for (q = 0; q < 2 * m + 1; q++)
	    printf(" %+7.5f", A[p][q]);
	printf("\n");
    }
#endif

    printf("備考: 1+2θλ=%4.1f, -θλx=%5.1f, -θλy=%5.1f\n",
	   beta_0, beta_1, beta_2);

    printf("Tmax: ");
    scanf("%lf", &Tmax);
    printf("Δt: ");
    scanf("%lf", &dt);
    skip = dt / tau;

    kmax = (Tmax + 0.1 * tau) / tau;

    g_init("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();
#ifdef OLD
    g_hidden2(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, N_x + 1, N_y + 1,
	      1, G_SIDE_NONE, 2, 1);
#else
    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[0][0], N_x + 1, N_y + 1,
	     1, G_SIDE_NONE, 2, 1);
#endif
    
    /* LU分解をする */
    bandlu(N, m, A);		/* decomp(N, A, &cond, iwork, B); */

    for (k = 1; k <= kmax; k++) {

	/* まず、素朴な連立一次方程式の右辺を用意する */
	/* 内部の格子点 */
	for (i = 1; i < N_x; i++)
	    for (j = 1; j < N_y; j++) {
		L = phi(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]);
	    }
	/* 下の辺、上の辺にある格子点 (角の点も含める) */
	for (i = 0; i <= N_x; i++) {
	    x = a + i * h_x;
	    /* (i, 0) */
	    L = phi(i, 0);
	    B[L] = alpha(x, c);
	    /* (i, N_y) */
	    L = phi(i, N_y);
	    B[L] = alpha(x, d);
	}
	/* 左の辺、右の辺にある格子点 (角の点は含めない) */
	for (j = 1; j < N_y; j++) {
	    y = c + j * h_y;
	    /* (0, j) */
	    L = phi(0, j);
	    B[L] = alpha(a, y);
	    /* (N_x, j) */
	    L = phi(N_x, j);
	    B[L] = alpha(b, y);
	}

	/* 対称化する */
	/* 対称化するための作業 1 */
	for (j = 1; j < N_y; j++) {
	    y = c + j * h_y;
	    /* (1,j) */
	    L = phi(1, j);
	    B[L] -= beta_1 * alpha(a, y);
	    /* (N_x-1,j) */
	    L = phi(N_x - 1, j);
	    B[L] -= beta_1 * alpha(b, y);
	}
	/* 対称化するための作業 2 */
	for (i = 1; i < N_x; i++) {
	    x = a + i * h_x;
	    /* (i,1) */
	    L = phi(i, 1);
	    B[L] -= beta_2 * alpha(x, c);
	    /* (i,N_y-1) */
	    L = phi(i, N_y - 1);
	    B[L] -= beta_2 * alpha(x, d);
	}

	/* A vector_U = B を解く */
	bandsol(N, m, A, B);	/* solve(N, A, B, iwork); */
	/* */
	for (i = 0; i <= N_x; i++)
	    for (j = 0; j <= N_y; j++)
		Uk[i][j] = B[phi(i, j)];
	/* */
	if (k % skip == 0) {
#ifdef  PRINT
	    printf("%15.8f\n", k * dt);
	    for (i = 0; i <= N_x; i++) {
		for (j = 0; j <= N_y; j++)
		    printf(" %+7.5f", Uk[i][j]);
		printf("\n");
	    }
#endif

	    g_cls();
#ifdef OLD
	    g_hidden2(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, N_x + 1, N_y + 1,
		      1, G_SIDE_NONE, 2, 1);
#else
	    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[0][0], N_x + 1, N_y + 1,
		     1, G_SIDE_NONE, 2, 1);
#endif
	}
    }
    /* マウスでクリックされるのを待つ */
    g_sleep(-1.0);
    /* ウィンドウを消す */
    g_term();

}

double u0(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;
}