\documentclass[tikz]{standalone}

\usepackage[small]{tensorstyle}

\begin{document}

\def\colorPsi{bg-red}

\newcommand\setScopeMacros[5]{
    \directlua{
        tex.sprint("\\pgfmathsetmacro{\\#1x}{#2}")
        tex.sprint("\\pgfmathsetmacro{\\#1y}{#3}")
        tex.sprint("\\pgfmathsetmacro{\\#1w}{#4}")
        tex.sprint("\\pgfmathsetmacro{\\#1h}{#5}")
    }
}

% 1: number of tensors
% 2: shift in (x,y)
% 3: prefix of input points names
% 4: prefix of output points names
\newcommand{\drawBrickwall}[4]{
    \begin{scope}[shift={#2}]
        \pgfmathsetmacro{\N}{#1}
        \pgfmathsetmacro{\iMax}{\N/2-1}
        % invisible tensor sized nodes
        \tikzstyle{U} = [minimum size=\tensorSize cm, rounded corners=2pt]
        \foreach \i in {0,...,\iMax} {
            \pgfmathsetmacro{\j}{int(2*\i)}
            \pgfmathsetmacro{\k}{int(2*\i+1)}
            \pgfmathsetmacro{\l}{int(2*\i-1)}

            % Node points
            % First row
            \node[U] (#4\j-1) at (\j,0) {};
            \node[U] (#4\k-1) at (\k,0) {};

            % Second row, use place the outpoint for the first in last tensor in the first row 
            \pgfmathparse{\i==0}
            \ifnum\pgfmathresult=1 \relax
                \node[U] (#4\j) at (\j, 0) {};
                \node[U] (#4\k) at (\k,-1) {};
            \else
                \pgfmathparse{\i>=\iMax}
                \ifnum\pgfmathresult=1 \relax
                    \node[U] (#4\j) at (\j,-1) {};
                    \node[U] (#4\k) at (\k,0) {};
                \else
                    \node[U] (#4\j) at (\j,-1) {};
                    \node[U] (#4\k) at (\k,-1) {};
                \fi
            \fi

            % first row tensor
            % workaround for getting a two-size tensor with the same height as the 
            % single site node tensor
            \begingroup
                \node[draw, fit=(#4\j-1)(#4\k-1), inner sep=0, tensorcorners] {};
            \endgroup
            % connect to previous
            \ifstrempty{#3}{}{
                \draw[leg] (#3\j) -- (#4\j-1);
                \draw[leg] (#3\k) -- (#4\k-1);
            }

            % second row tensor
            \pgfmathparse{\i>0}
            \ifnum\pgfmathresult=1 \relax
                \node[draw, fit=(#4\l)(#4\j), inner sep=0, tensorcorners] {};
            \fi
            \pgfmathparse{\i<=\iMax}
            \ifnum\pgfmathresult=1 \relax
              \draw[leg] (#4\j-1) -- (#4\j);
              \draw[leg] (#4\k-1) -- (#4\k);
            \fi
        }
    \end{scope}
}


\begin{tikzpicture}[x=\unitCm cm,y=\unitCm cm, scale=\scale, transform shape]
    % for positioning
    % \drawGrid{-1}{10}{-20}{1}
    \pgfmathsetmacro{\N}{10}
    \pgfmathsetmacro{\iMax}{\N-1}
    % Heffmat
    \begin{scope}[shift={(0,0)}]
        % Psi
        \foreach \i in {0,...,\iMax} {
            \node[ten] (M\i) at (\i,0) {};
            \pgfmathsetmacro{\j}{int(\i-1)}
            \pgfmathparse{\j >= 0}
            \ifnum\pgfmathresult=1 \relax
                \draw[leg] (M\i) -- (M\j);
            \fi
        }
        \begin{scope}[on background layer]
            \drawAroundRect[draw=none, fill=\colorPsi,tensorcorners]{(M0)}{\iMax}{0}
        \end{scope}
        % \node[anchor=west] at (\iMax+\dLabel, 0) {$\big\}M^{[i]}$};
        \node[anchor=west] at (\iMax+\dLabel, 0) {$\psi(t=0)$};
    \end{scope}
    \drawBrickwall{\N}{(0,-1)}{M}{U1}
    \drawBrickwall{\N}{(0,-3)}{U1}{U2}
    \drawBrickwall{\N}{(0,-7)}{}{U3}
    \begin{scope}[on background layer]
        \drawAroundRect[draw=none,fill=bg2,tensorcorners]{(U10-1)}{\N-1}{-1}
        \drawAroundRect[draw=none,fill=bg2,tensorcorners]{(U20-1)}{\N-1}{-1}
        \drawAroundRect[draw=none,fill=bg2,tensorcorners]{(U30-1)}{\N-1}{-1}
        % \drawAroundPoints[U]{(#4\j-1)(#4\k-1)}
    \end{scope}
    \node[anchor=west] at (\iMax+\dLabel, -1.5) {$n = 1$};
    \node[anchor=west] at (\iMax+\dLabel, -3.5) {$n = 2$};
    \node[anchor=west] at (\iMax+\dLabel, -7.5) {$n = n_\text{steps}$};

    % connect dashed
    \foreach \i in {0,...,\iMax} {
        \draw[leg] (U2\i) -- (\i,-5);
        \draw[leg,dotted] (\i,-5) -- (\i,-6);
        \draw[leg] (\i,-6) -- (U3\i-1);
    }

    \begin{scope}[shift={(0,-9)}]
        % PsiTilde
        \foreach \i in {0,...,\iMax} {
            \node[ten] (m\i) at (\i,0) {};
            \pgfmathsetmacro{\j}{int(\i-1)}
            \pgfmathparse{\j >= 0}
            \ifnum\pgfmathresult=1 \relax
                \draw[leg] (m\i) -- (m\j);
            \fi
            % leg to Us
            \draw[leg] (U3\i) -- (m\i);
        }
        \begin{scope}[on background layer]
            \drawAroundRect[draw=none,fill=\colorPsi,tensorcorners]{(m0)}{\iMax}{0}
        \end{scope}
        % \node[anchor=west] at (\iMax+\dLabel, 0) {$\big\}\tilde{M}^{[i]}$};
        \node[anchor=west] at (\iMax+\dLabel, 0) {${\psi}(t=n_\text{steps}\delta t)$};
    \end{scope}

        % \node[draw, dashed, rounded corners=5pt, fit=(M0) (M2) (m0) (m2)] (group) {};
\end{tikzpicture}

\end{document}