// PROGRAM tdse
// ポテンシャルに入射した波束の運動

#include "TrueBASIC.h"
int main();
void parameters(float *x0, float *k0, float *width, float *V0, float *a, 
                float *xmin, float *xmax, int *n, float *dx, float *dx2, 
                float *dt);
void initial_packet(float Re[], float Im[], float P[], float x0, float k0, 
                    float width, float xmin, int n, float dx, float dt);
void set_up_windows(float xmin, float xmax, float V0, int IC1_, int IC2_, 
                    int IC3_, int IC4_);
void draw_potential(float V0, float a, float xmin, float xmax, int IC4_);
void evolve(float Re[], float Im[], float Imold[], float *t, float V0, 
            float a, float dx, float dx2, float dt, float xmin, int n);
void plots(float Re[], float Im[], float Imold[], float P[], float xmin, 
           int n, float dx, int IC1_, int IC2_, int IC3_, 
           int IC4_);
float V(float x, float V0, float a);

float wx1, wy1, wx2, wy2;

int main()
{
    int IC1_ = 11, IC2_ = 12, IC3_ = 13, IC4_ = 14, n;
    float Im[1101], Imold[1101], Re[1101], P[1101], 
          V0, a, dt, dx, dx2, k0, t, width, x0, xmax, xmin;

    GWopen(0);
    GWmode(3, 0);

    parameters(&x0, &k0, &width, &V0, &a, &xmin, &xmax, &n, &dx, &dx2, &dt);
    initial_packet(Re, Im, P, x0, k0, width, xmin, n, dx, dt);
    set_up_windows((float)xmin, (float)xmax, (float)V0, IC1_, IC2_, IC3_, IC4_);
    draw_potential((float)V0, (float)a, (float)xmin, (float)xmax, IC4_);
    GWanchor(1);
    do
    {
        int i;
        for(i = 0; i < 100; ++i)
        evolve(Re, Im, Imold, &t, V0, a, dx, dx2, dt, xmin, n);
        plots(Re, Im, Imold, P, xmin, n, dx, IC1_, IC2_, IC3_, IC4_);
    } while(!TBkeyinput());

    GWquit();
    return 0;
}

void parameters(float *x0, float *k0, float *width, float *V0, float *a, 
                float *xmin, float *xmax, int *n, float *dx, float *dx2, 
                float *dt)
{
    (*x0) = -15;                            // 波束の初期位置
    (*width) = 1;                           // x 空間での波束の幅
    (*k0) = 2;                              // 波束の群速度
    (*xmax) = 20;                           // x の最大値
    (*xmin) = -(*xmax);                     // x の最小値
    (*V0) = 2;
    (*a) = 1;                               // 井戸型ポテンシャルの幅の半分
    (*dx) = 0.1f;                           // 刻み幅
    (*dx2) = (*dx)*(*dx);
    (*n) = (int)(((*xmax) - (*xmin))/(*dx)+0.5);
    (*dt) = 0.005f;                         // 時間間隔
}

void initial_packet(float Re[], float Im[], float P[], float x0, float k0, 
                    float width, float xmin, int n, float dx, float dt)
// 初期状態としてのガウス型の波束
{
    int i;
    float A, arg, arg2, b, delta2, e, e2, x;

    delta2 = width*width;
    A = (float)pow(2*pi*delta2, -0.25);
    b = 0.5f*k0*dt;
    for(i = 1; i <= n; ++i)
    {
        x = xmin + (i-1)*dx;
        arg = (float)(0.25f*pow(x - x0, 2.0f)/delta2);
        arg2 = (float)(0.25f*pow(x - x0 - 0.5f*k0*dt, 2.0f)/delta2);
        e = (float)exp(-arg);
        e2 = (float)exp(-arg2);
        Re[i] = A*(float)cos(k0*(x - x0))*e;       // t = 0 での値
        Im[i] = A*(float)sin(k0*(x - x0 - 0.5*b))*e2; // t = dt/2 での値
    }
    Re[0] = Re[n+1] = 0;
    Im[0] = Im[n+1] = 0;
    P[0] = P[n+1] = 0;
}

void set_up_windows(float xmin, float xmax, float V0, int IC1_, int IC2_, 
                    int IC3_, int IC4_)
{
    GWvport(0.0f, 0.75f, GWaspect(1), 1.0f);    // 確率
    GWindow(xmin, -0.1f, xmax, 0.5f);
    GWsetpen(BLACK, -1, -1, 4);
    GWline(xmin, 0.0f, xmax, 0.0f);
    GWsetpen(BLUE, -1, -1, -1);
    GWsavevp(IC1_);
    GWvport(0.0f, 0.5f, GWaspect(1), 0.75f);    // 実部
    GWsetpen(BLACK, -1, -1, 4);
    GWindow(xmin, -1.0f, xmax, 1.0f);
    GWsetpen(BLACK, -1, -1, 4);
    GWline(xmin, 0.0f, xmax, 0.0f);
    GWsetpen(BLUE, -1, -1, -1);
    GWsavevp(IC2_);
    GWvport(0.0f, 0.25f, GWaspect(1), 0.5f);    // 虚部
    GWindow(xmin, -1.0f, xmax, 1.0f);
    GWsetpen(BLACK, -1, -1, 4);
    GWline(xmin, 0.0f, xmax, 0.0f);
    GWsetpen(BLUE, -1, -1, -1);
    GWsavevp(IC3_);
    GWvport(0.0f, 0.0f, GWaspect(1), 0.22f);    // ポテンシャル
    GWindow(xmin, -V0, xmax, V0);
    GWsavevp(IC4_);
}

void draw_potential(float V0, float a, float xmin, float xmax, int IC4_)
{
    GWselvp(IC4_);
    GWsetpen(RED, -1, -1, -1);
    GWmove2(xmin, 0.0f);
    GWline2(a, 0.0f);
    GWline2(a, V0);
    GWline2(xmax, V0);
    GWindow(0.0f, 20.0f, 80.0f, 0.0f);
    GWsavevp(IC4_);
}

void evolve(float Re[], float Im[], float Imold[], float *t, float V0, 
            float a, float dx, float dx2, float dt, float xmin, int n)
{
    int i;
    float HIm, HRe, x;

    for(i = 1; i <= n; ++i)
    {
        x = xmin + (i-1)*dx;
        HIm = V(x, V0, a)*Im[i] -0.5f*(Im[i + 1] -2*Im[i] +Im[i - 1])/dx2;
        // dt の整数倍の時刻で定義された実部
        Re[i] += HIm*dt;
    }
    for(i = 1; i <= n; ++i)
    {
        x = xmin + (i-1)*dx;
        // 実部より dt/2 だけ前の値
        Imold[i] = Im[i];
        HRe = V(x, V0, a)*Re[i] -0.5f*(Re[i + 1] -2*Re[i] +Re[i - 1])/dx2;
        // 実部より dt/2 だけ後の値
        Im[i] -= HRe*dt;
    }
    (*t) += dt;                             // 実部の時刻
}

void plots(float Re[], float Im[], float Imold[], float P[], float xmin, 
           int n, float dx, int IC1_, int IC2_, int IC3_, 
           int IC4_)
{
    int i;
    float Psum = 0, x;
    char Stmp1_[_LBUFF_];

    GWselvp(IC2_);
    GWsetogn(0, -IC2_);
    GWplot1(-1, n, xmin, xmin+n*dx, 0.0f, 0.0f, 1.0f, 0.0f, Re+1);
    GWflush(-IC2_);
    GWselvp(IC3_);
    GWsetogn(0, -IC3_);
    GWplot1(-1, n, xmin, xmin+n*dx, dx/2, 0.0f, 1.0f, 0.0f, Im+1);
    GWflush(-IC3_);
    Psum = 0;
    for(i = 1; i <= n; ++i)
    {
        x = xmin + (i - 1)*dx;
        P[i] = Re[i]*Re[i] + Im[i]*Imold[i];
        Psum += P[i]*dx;
    }
    GWselvp(IC1_);
    GWsetogn(0, -IC1_);
    GWplot1(-1, n, xmin, xmin+n*dx, dx/2, 0.0f, 1.0f, 0.0f, P+1);
    GWflush(-IC1_);
    GWselvp(IC4_);
    GWsetogn(0, -IC4_);
    sprintf(Stmp1_, "total probability = %f", Psum);
                                        // 全確率は 1 になるはず
    GWputtxt(1.0f, 1.0f, Stmp1_);
    GWflush(-IC4_);
    GWsetogn(0, 1);
}

float V(float x, float V0, float a)
// 階段状ポテンシャル
{
    float V_;

    if(x > a)
        V_ = V0;
    else
        V_ = 0;
    return V_;
}