clear
% lock exchange flowの計算

% 
% --- 条件
% 
sg = [1.0068 1.0151 1.023];  %比重[１％　２％　３％]
dimension = 1;  %塩分濃度と密度を近似する次数
salinity = 2;  %塩分濃度[％]
temp = 17.1;  %水温[℃]
depth = 10;  %水深[cm]
width = 70;  %幅[cm]
dep_int = 100;  %水深方向格子数
wid_int = 700;  %幅方向格子数
sal_dep = 2.5;  %塩水の水深
sal_wid = 20;  %塩水の幅
T = 1;  %計算終了時間[s]
dt = 0.001;  %時間間隔
er = 0.1;  %補正時の微小値
% 以下は変更不要な条件
dx = width/wid_int;
dz = depth/dep_int;
D = 1.24*10^(-5);  %水中でのNaClの拡散係数[cm^2/s]
V = [1.306 1.139 1.003 0.893];
[poly_v,polyfit_profile] = polyfit(10:5:25,V,2);
v = polyval(poly_v,temp)*10^(-2);  %動粘性係数[cm^2/s]
g = -9.8;  %重力加速度

% 
% --- 計算
% 
Drx = D*dt/dx^2;
Drz = D*dt/dz^2;
vrx = v*dt/dx^2;
vrz = v*dt/dz^2;
if Drx >= 0.5 || Drz >= 0.5 || vrx >= 0.5 || vrz >= 0.5
    disp('拡散数が0.5を超えています')
    error('dep_int，wid_int，dtを変更してください')
end
% % 第一段階　塩分濃度変化の密度式決定　％で定義
Cexpansion = (temp-5)/67*10^(-3);
r(1) = 1 - Cexpansion*(temp-3.98);
r(2:4) = r(1)*sg;
[rho,polyfit_profile] = polyfit(0:3,r,dimension);
d0 = polyval(rho,0);  %淡水の密度

% % 第二段階　初期条件の決定
% ベクトルとスカラーの座標
[uX,uY] = meshgrid(dz/2:dz:depth-dz/2,0:dx:width);
[wX,wY] = meshgrid(0:dz:depth,dx/2:dx:width-dx/2);
[scalarX,scalarY] = meshgrid(dz/2:dz:depth,dx/2:dx:width-dx/2);
% 変数の初期条件
u = zeros(wid_int+1,dep_int);
temp_u = u;
uk = u;
w = zeros(wid_int,dep_int+1);
temp_w = w;
wk = w;
s = zeros(wid_int,dep_int);
s(1:sal_wid/dx,1:sal_dep/dz) = salinity;
temp_s = s;
d = polyval(rho,s);
p = zeros(wid_int,dep_int);
for i = 1:wid_int
    for j = dep_int:-1:1
        if j == dep_int
            p(i,j) = dz/2*d(i,j)/2;
        else
            p(i,j) = p(i,j+1) + dz*mean(d(i,j:j+1));
        end
    end
end
pk = p;
Dk = zeros(wid_int,dep_int);

% % 第三段階　計算
for t = 1:T/dt
    
    % 仮のuを求める
    for i = 1:wid_int+1
        for j = 1:dep_int
            if i == 1
                ux = 0;
                wx = 0;
                px = 0;
                dd = d(i,j);
            elseif i == wid_int+1
                ux = 0;
                wx = 0;
                px = 0;
                dd = d(i-1,j);
            else
                ux = u(i+1,j)-2*u(i,j)+u(i-1,j);
                wx = sum(sum(w(i-1:i,j:j+1)))/4;
                px = p(i,j)-p(i-1,j);
                dd = (d(i,j) + d(i-1,j))/2;
            end
            if j == 1 || j == dep_int
                uz = 0;
            else
                uz = u(i,j+1)-2*u(i,j)+u(i,j-1);
            end
            temp_u(i,j) = u(i,j)...
                        - u(i,j)*dt/2/dx*ux...
                        - wx*dt/2/dz*uz...
                        - dt/dd/2*dx*px...
                        + v*dt/dx^2*ux...
                        + v*dt/dz^2*uz;
        end
    end
    % 仮のwを求める
    for i = 1:wid_int
        for j = 1:dep_int+1
            if i == 1 || i == wid_int
                wx = 0;
            else
                wx = w(i+1,j)-2*w(i,j)+w(i-1,j);
            end
            if j == 1
                uz = 0;
                wz = 0;
                pz = 0;
                dd = d(i,j);
            elseif j == dep_int+1
                uz = 0;
                wz = 0;
                pz = 0;
                dd = d(i,j-1);
            else
                uz = sum(sum(u(i:i+1,j-1:j)))/4;
                wz = w(i,j+1)-2*w(i,j)+w(i,j-1);
                pz = p(i,j)-p(i,j-1);
                dd = (d(i,j) + d(i,j-1))/2;
            end
            temp_w(i,j) = w(i,j)...
                        - uz*dt/2/dx*wx...
                        - w(i,j)*dt/2/dz*wz...
                        - dt/dd/2*dx*pz...
                        - dt*g*(dd-d0)/d0...
                        + v*dt/dx^2*wx...
                        + v*dt/dz^2*wz;
        end
    end
    % u，w，pの補正
    for k = 1:100
        for i = 2:wid_int-1
            for j = 2:dep_int-1
                Dk(i,j) = (temp_u(i+1,j)-2*temp_u(i,j)+temp_u(i-1,j))/dx...
                        + (temp_w(i,j+1)-2*temp_w(i,j)+temp_w(i,j-1))/dz;
            end
        end
        if sum(sum(abs(Dk) >= er)) == 0
            break
        end
        dp = -Dk/(2*(1/dx^2+1/dz^2));
        for i = 2:wid_int
            for j = 1:dep_int
                uk(i,j) = temp_u(i,j) + (dp(i,j)+dp(i-1,j))/2/dx;
            end
        end
        temp_u = uk;
        for i = 1:wid_int
            for j = 2:dep_int
                wk(i,j) =temp_w(i,j) + (dp(i,j)+dp(i,j-1))/2/dz;
            end
        end
        temp_w = wk;
        pk = p + dp;
        p = pk;
    end
    u = temp_u;
    w = temp_w;
    % 塩分濃度の分布を求める
    for i = 1:wid_int
        for j = 1:dep_int
            if i == 1
                sx = s(i+1,j)-2*s(i,j)+s(i,j);
            elseif i == wid_int
                sx = s(i,j)-2*s(i,j)+s(i-1,j);
            else
                sx = s(i+1,j)-2*s(i,j)+s(i-1,j);
            end
            if j == 1
                sz = s(i,j+1)-2*s(i,j)+s(i,j);
            elseif j == dep_int
                sz = s(i,j)-2*s(i,j)+s(i,j-1);
            else
                sz = s(i,j+1)-2*s(i,j)+s(i,j-1);
            end
            temp_s(i,j) = s(i,j)...
                        - (u(i+1,j)+u(i,j))/2*dt/2/dx*sx...
                        - (w(i,j+1)+w(i,j))/2*dt/2/dz*sz...
                        + D*dt/dx^2*sx...
                        + D*dt/dz^2*sz;
        end
    end
    s = temp_s;
    % 密度分布を求める
    d = polyval(rho,s);
    
    t
end
% 
% --- 図を描く
% 
% 座標に注意　ベクトルとスカラーは座標が違う
figure,imagesc(0:dx:width,dz/2:dz:depth-dz/2,rot90(u))
title('u')
figure,imagesc(dx/2:dx:width-dx/2,0:dz:depth,rot90(w))
title('w')
figure,imagesc(dx/2:dx:width-dx/2,dz/2:dz:depth-dz/2,rot90(d))
title('Density')
figure,imagesc(dx/2:dx:width-dx/2,dz/2:dz:depth-dz/2,rot90(p))
title('Pressure')
figure,imagesc(dx/2:dx:width-dx/2,dz/2:dz:depth-dz/2,rot90(s))
title('Salinity')
xlabel('width [cm]')
ylabel('depth [cm]')