%% em3d-MT
%% Three-dimensional forward modeling of magnetotelluric / radio-magnetotelluric data considering diffusion and wave fields in anisotropic media
%
% November 13, 2024 Version: 1.0.0
% Submit to GEOPHYSICS
%
%% 1 Input parameters
% Input: mesh, material parameters, frequency, number of threads

t = tic;
% Input;

%
%% 2 FE 预处理 FE pre-processing
% get_mesh_Connect ：获取网格连接信息 Get mesh connection information

[mesh] = get_mesh_Connect(mesh);

%  get_boundary：获取边界的棱边信息 Get boundary information

[Bnd] = get_boundary(mesh.node2coord,mesh.edgeCen,mesh.elemCen,mesh.edgeUnit,mesh.elem2edge);

% get_abcd_volume：获得基函数待定系数 Get the undetermined coefficients and volume

[mesh.abcd,mesh.vol] = get_abcd_volume(mesh.node2coord,mesh.elem2node) ;

%
%% 3 插值矩阵 Field Interpolation matrix
% GetInterpMatrix：计算测点处电磁场的插值矩阵 Get the interpolation matrix of the electromagnetic field

clear rec
load rec.dat
nrec = size(rec,1);

Q = GetInterpMatrix(rec,mesh);

%
%% 4 有限元系数矩阵组装 FE assembly
%%  单元矩阵M组装
% K_{ij}^e ={\int ~}_{\Omega e} \left(\nabla \times N_i^e \right)\cdot {\mu_r }^{-1} \cdot \left(\nabla \times N_j^e \right)dV;
%  get_K：刚度矩阵k组装 Assembled stiffness matrix

K = get_K_v2(mesh.abcd,mesh.vol,mesh.elem2edge,mesh.edgeLength,mesh.elemEdgeLength,mesh.sign,P.mur);

%  get_M：质量矩阵组装 Assemble mass matrix
% {M_{\sigma } }_{ij}^e ={\int ~}_{\Omega ~e} N_i^e \cdot \sigma^e {\cdot N}_j^e dV
%  get_M:  Assemble the diffusion field mass matrix

M1  = get_M_v2(mesh.abcd,mesh.vol,mesh.elem2edge,mesh.edgeLength,mesh.elemEdgeLength, mesh.sign,P.sig);

%  {M_{\varepsilon } }_{ij}^e ={\int ~}_{\Omega e} N_i^e \cdot \varepsilon^e {\cdot N}_j^e dV;
%  get_M: Assemble the wave field mass matrix

M2  = get_M_v2(mesh.abcd,mesh.vol,mesh.elem2edge,mesh.edgeLength,mesh.elemEdgeLength, mesh.sign,P.eps);

%
%% 5 边界条件 boundary condition
% get_Boundary1DModel：The 1D background model was extracted

[boundModel,P.sigdiff]  = get_Boundary1DModel(P.sig,P.epsr,P.mur,mesh.elemCen,Bnd.elemID,mesh.elem2node,mesh.node2coord,0,0);
rhs = zeros(mesh.nedge,1);
col = jet(nfreq);
for ifreq = 1:nfreq

%  Assemble the coefficient matrix \mathbf{K}-({{\mu_0 \omega }^2 \mathbf{M}}_{\varepsilon } +i\omega \mu_0 {\mathbf{M}}_{\sigma } )

    A = K - 1i*omega(ifreq)*P.mu0*M1  - omega(ifreq)^2*P.mu0*M2;

% get_BoundaryFieldAniEpsMu: Calculation of boundary electric fields

    beField(ifreq) = get_BoundaryFieldAniEpsMu(freq(ifreq),boundModel,Bnd.edgeCen,Bnd.unit,0,col(ifreq,:));
    [FE(ifreq).A,FE(ifreq).rhs] = set_Boundary(beField(ifreq),Bnd.edgeID,A,rhs);
end

clear A K M1 M2  beField Bnd rhs

%
%% 6 Solve Ax = b
% solver: 'Paradiso'，'umfpack'，'iluQMR'，'iluBicgstab'

% solver = 'Paradiso';
x1  = zeros(mesh.nedge,nfreq);
x2  = x1;

try; pool = parpool('local',2); end
parfor ifreq = 1:nfreq
    t2(ifreq) = tic;
    if  strcmp(solver,'Paradiso')
        xtemp = pardiso_MT(FE(ifreq).A,[FE(ifreq).rhs.TE,FE(ifreq).rhs.TM],OMP);
        x1(:,ifreq) =  xtemp(:,1);
        x2(:,ifreq) =  xtemp(:,2);
    elseif strcmp(solver,'umfpack')
        xtemp = umfpack_MT(FE(ifreq).A,[FE(ifreq).rhs.TE,FE(ifreq).rhs.TM,Qt]);
        x1(:,ifreq) =  xtemp(:,1);
        x2(:,ifreq) =  xtemp(:,2);
        Aiqt{ifreq} =  xtemp(:,3:end);
    elseif strcmp(solver,'iluQMR')
        [L,U]=ilu(FE(ifreq).A);
        [x1(:,ifreq),flag1(ifreq),relres1(ifreq),iter1(ifreq),resvec1{ifreq}] = qmr(FE(ifreq).A,FE(ifreq).rhs.TE,1e-7,3000,L,U);
        [x2(:,ifreq),flag2(ifreq),relres2(ifreq),iter2(ifreq),resvec2{ifreq}] = qmr(FE(ifreq).A,FE(ifreq).rhs.TM,1e-7,3000,L,U);
    elseif strcmp(solver,'iluBicgstab')
        [L,U]=ilu(FE(ifreq).A);
        [x1(:,ifreq),flag1(ifreq),relres1(ifreq),iter1(ifreq),resvec1{ifreq}] = bicgstab(FE(ifreq).A,FE(ifreq).rhs.TE,1e-7,3000,L,U);
        [x2(:,ifreq),flag2(ifreq),relres2(ifreq),iter2(ifreq),resvec2{ifreq}] = bicgstab(FE(ifreq).A,FE(ifreq).rhs.TM,1e-7,3000,L,U);
    end
    erro(ifreq).TE = norm( FE(ifreq).A*x1(:,ifreq) -FE(ifreq).rhs.TE) / norm(FE(ifreq).rhs.TE) ;
    erro(ifreq).TM = norm( FE(ifreq).A*x2(:,ifreq) - FE(ifreq).rhs.TM) /norm(FE(ifreq).rhs.TM) ;

    tFE(ifreq)=toc(t2(ifreq));
end

if ismember(solver,['iluQMR','iluBicgstab'])
    plot_iterative;
end
clear xtemp
time.SolveFE=  tFE;
time.MT3d = toc(t);

%
%% 7  后处理 Post-processing
% get_data: Data function, Get MT response

RX   = (1:size(rec,1))';
data = get_data_MT(x1, x2, Q, freq, RX);

%
%%