\documentclass[tikz]{standalone}

\usepackage[small]{tensorstyle}

\begin{document}

\def\colorTheta{bg-red}
\def\colorHeff{bg-blue}

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



\begin{tikzpicture}[x=\unitCm cm,y=\unitCm cm, scale=\scale, transform shape]
    \pgfmathsetmacro{\margin}{0.3}
    \pgfmathsetmacro{\scopeDistance}{1}
    \setScopeMacros{scopeA}{0}{0}{8}{2}  % Heff*Theta
    \setScopeMacros{scopeB}{0}{-5}{3}{2}  % Heffmat * Theta
    \setScopeMacros{scopeC}{8}{\scopeBy}{1}{2}  % ThetaTwiddle
    \setScopeMacros{scopeD}{2+0.5}{\scopeCy-4}{2}{0}  % SVD
    \setScopeMacros{scopeE}{-1}{\scopeDy-4}{3}{2}  % L
    \setScopeMacros{scopeF}{5}{\scopeDy-4}{5}{2}  % R
    % for positioning
    % \drawGrid{-1}{10}{-20}{1}

    % \pgfmathsetmacro{\dx}{0.3*\nodeDistance}
    % \pgfmathsetmacro{\dy}{0.3*\nodeDistance}
    % Heffmat
    \begin{scope}
        \pgfmathsetmacro{\N}{9}
        \pgfmathsetmacro{\iMax}{\N-1}
        \foreach \i in {0,...,\iMax} {
            \node[ten] (M\i) at (\i,0) {}; % $M^{[\the\numexpr\i+1]}$
            \node[mpo] (W\i) at (\i,-1) {}; % $W^{[\the\numexpr\i+1]}$
            \node[ten] (m\i) at (\i,-2) {}; % $m^{[\the\numexpr\i+1]}$
        }
        \foreach \i in {0,...,\iMax} {
            \pgfmathsetmacro{\j}{int(\i+1)}
            % horizontal, top bot, skip the legs for the theta tensors

            % \pgfmathparse{\i < \iMax && (\i>4 || \i<2 || \i==3) }  % dont connect Heff and Theta
            \pgfmathparse{\i < \iMax}
            \ifnum\pgfmathresult=1 \relax
              \draw[leg] (M\i) -- (M\j);
              \draw[leg] (m\i) -- (m\j);
            \fi
            % horizontal middle
            \pgfmathparse{\i < \iMax}
            \ifnum\pgfmathresult=1 \relax
              \draw[leg] (W\i) -- (W\j);
            \fi
            % vertical
            % \pgfmathparse{\i>4 || \i<3}  % dont connect Heff and Theta
            \pgfmathparse{1}
            \ifnum\pgfmathresult=1 \relax
                \draw[leg] (M\i) -- (W\i);
                \draw[leg] (W\i) -- (m\i);
            \fi
        }

        % \node[draw, dashed, rounded corners=5pt, fit=(M0) (M2) (m0) (m2)] (group) {};
        \begin{scope}[on background layer]
          \draw[draw=none,fill=\colorHeff,tensorcorners] 
            ($(M0)+(-\tdx, \tdx)$) --
            ($(M2)+( \tdx, \tdx)$) --
            ($(W2)+( \tdx, \tdx)$) --
            ($(W5)+(-\tdx, \tdx)$) --
            ($(M5)+(-\tdx, \tdx)$) --
            ($(M8)+( \tdx, \tdx)$) --
            ($(m8)+( \tdx,-\tdx)$) --
            ($(m5)+(-\tdx,-\tdx)$) --
            ($(W5)+(-\tdx,-\tdx)$) --
            ($(W2)+( \tdx,-\tdx)$) --
            ($(m2)+( \tdx,-\tdx)$) --
            ($(m0)+(-\tdx,-\tdx)$) --
            cycle;

          \drawAroundRect[draw=none,fill=\colorTheta,tensorcorners]{(M3)}{1}{0}
          \drawAroundRect[draw=none,fill=\colorTheta,tensorcorners]{(m3)}{1}{0}
        \end{scope}

        \node[anchor=south,above=0.1 of M1] (L) {$L^{[n]}$};
        \node[anchor=south,above=0.1 of M7] (R) {$R^{[n+1]}$};
        \node[anchor=south] (Theta)  at ($(M3)+0.5*(M4)-0.5*(M3)+(0,0.4)$) {$\Theta$};
        \node[anchor=north] (ThetaC) at ($(m3)+0.5*(m4)-0.5*(m3)-(0,0.4)$) {$\overline{\Theta}$};
        \node[anchor=west] at (\iMax+\dLabel, 0) {$\big\}M^{[i]}$};
        \node[anchor=west] at (\iMax+\dLabel,-1) {$\big\}W$};
        \node[anchor=west] at (\iMax+\dLabel,-2) {$\big\}\overline{M}^{[i]}$};
    \end{scope}

    \draw[->] (\scopeAx-1,\scopeAy-\scopeAh/2) -- ++(-1, 0) -- ++(0,-\scopeAy+\scopeBy) node[midway,right] {Reshape }-- ++(1, 0) ;

    % Heff mat * Theta
    \begin{scope}[shift={(\scopeBx, \scopeBy)}]
        % "fake" points to keep the same size
        \foreach \i in {0,...,3} {
            \node[ten,draw=none] (M\i) at (\i, 0) {}; % $M^{[\the\numexpr\i+1]}$
            \node[ten,draw=none] (W\i) at (\i,-1) {}; % $W^{[\the\numexpr\i+1]}$
            \node[ten,draw=none] (m\i) at (\i,-2) {}; % $m^{[\the\numexpr\i+1]}$
        }
        \node[ten,draw,inner sep=0,tensorcorners,fit=(M0)(m2)] {};
        \node[ten,draw,inner sep=0,tensorcorners,fit=(M3)(m3)] {};
        \begin{scope}[on background layer]
          \drawAroundPoints[fill=\colorHeff,tensorcorners]{(M0)(m2)}
          \drawAroundPoints[fill=\colorTheta,tensorcorners]{(M3)(m3)}
        \end{scope}
        \draw[leg] (W2) -- (W3);
        \draw[leg] (W0) -- ++(-\legLength,0);
        \node[anchor=south]  (Heff) at (1,\dLabel) {$H_\text{eff}$};
        \node[anchor=south] (Theta) at (3,\dLabel) {$\Theta$};
    \end{scope}

    \draw[->] (\scopeBx+\scopeBw+1,\scopeBy-\scopeBh/2) -- (\scopeCx-1,\scopeBy-\scopeBh/2) node[midway,above] {Ground} node[midway,below] {state};

    % Theta Twiddle <3
    \begin{scope}[shift={(\scopeCx, \scopeCy)}]
        \node[ten,draw=none] (M0) at (0, 0) {};
        \node[ten,draw=none] (w0) at (0,-1) {};
        \node[ten,draw=none] (m0) at (0,-2) {};

        \node[ten,draw,inner sep=0,tensorcorners,fit=(M0)(m0)] {};
        \draw[leg] (w0) -- ++(-\legLength,0);

        \begin{scope}[on background layer]
          \drawAroundPoints[fill=\colorTheta,tensorcorners]{(M0)(m0)}
        \end{scope}
        \node[anchor=south] at (0,\dLabel) {$\tilde{\Theta}$};
    \end{scope}

    \draw[->] (\scopeCx+1,\scopeCy-\scopeCh/2) -- ++(1,0) -- ++(0, +\scopeDy-\scopeCy+\scopeCh/2) node[midway,right] {SVD} -- (\scopeDx+\scopeDw+1,\scopeDy);


    % Theta Twiddle SVD
    \def\colorATensor{bg-green}
    \def\colorLTensor{bg-yellow}
    \def\colorBTensor{bg-purple}
    \begin{scope}[shift={(\scopeDx, \scopeDy)}]
        \node[ten,fill=\colorATensor] (M0) at (0,0) {};
        \node[ten,fill=\colorLTensor] (M1) at (1,0) {};
        \node[ten,fill=\colorBTensor] (M2) at (2,0) {};
        \draw[leg] (M0) -- ++(-\legLength,0);
        \draw[leg] (M0) -- (M1);
        \draw[leg] (M1) -- (M2);
        \draw[leg] (M2) -- ++(\legLength,0);

        \begin{scope}[on background layer]
            \drawAroundRect[draw=none,fill=\colorTheta,tensorcorners]{(0,0)}{2}{0}
        \end{scope}
        \node[anchor=south] at (1,\dLabel) {$\tilde{\Theta}$};
    \end{scope}

    % L n+1
    \begin{scope}[shift={(\scopeEx, \scopeEy)}]
        \pgfmathsetmacro{\N}{4}
        \pgfmathsetmacro{\iMax}{\N-1}
        \foreach \i in {0,...,\iMax} {
            \node[ten] (M\i) at (\i, 0) {};
            \node[mpo] (W\i) at (\i,-1) {};
            \node[ten] (m\i) at (\i,-2) {};
        }
        \foreach \i in {0,...,\iMax} {
            \pgfmathsetmacro{\j}{int(\i+1)}
            % horizontal
            \pgfmathparse{\j < \N}
            \ifnum\pgfmathresult=1 \relax
              \draw[leg] (M\i) -- (M\j);
              \draw[leg] (W\i) -- (W\j);
              \draw[leg] (m\i) -- (m\j);
            \fi
            % vertical
            \draw[leg] (M\i) -- (W\i);
            \draw[leg] (W\i) -- (m\i);
        }
        % sticking out legs
        \draw[leg] (M3) -- ++(\legLength,0);
        \draw[leg] (W3) -- ++(\legLength,0);
        \draw[leg] (m3) -- ++(\legLength,0);

        \node[ten,fill=\colorATensor] (M3c) at (M3) {};
        \node[ten,fill=\colorATensor] (m3c) at (m3) {};
        \begin{scope}[on background layer]
          \drawAroundRect[fill=\colorHeff,tensorcorners]{(M0)}{\iMax}{-2}
        \end{scope}
        \node[anchor=south] at (1,\dLabel) {$L^{[n+1]}$};
    \end{scope}

    % R
    \begin{scope}[shift={(\scopeFx, \scopeFy)}]
        \pgfmathsetmacro{\N}{5}
        \pgfmathsetmacro{\iMax}{\N-1}
        \foreach \i in {0,...,\iMax} {
            \node[ten] (M\i) at (\i, 0) {};
            \node[mpo] (W\i) at (\i,-1) {};
            \node[ten] (m\i) at (\i,-2) {};
        }
        \foreach \i in {0,...,\iMax} {
            \pgfmathsetmacro{\j}{int(\i+1)}
            % horizontal
            \pgfmathparse{\j < \N}
            \ifnum\pgfmathresult=1 \relax
              \draw[leg] (M\i) -- (M\j);
              \draw[leg] (W\i) -- (W\j);
              \draw[leg] (m\i) -- (m\j);
            \fi
            % vertical
            \draw[leg] (M\i) -- (W\i);
            \draw[leg] (W\i) -- (m\i);
        }
        % sticking out legs
        \draw[leg] (M0) -- ++(-\legLength,0);
        \draw[leg] (W0) -- ++(-\legLength,0);
        \draw[leg] (m0) -- ++(-\legLength,0);

        \node[ten,fill=\colorBTensor] (M0c) at (M0) {};
        \node[ten,fill=\colorBTensor] (m0c) at (m0) {};
        \begin{scope}[on background layer]
          \drawAroundRect[fill=\colorHeff,tensorcorners]{(M0)}{\iMax}{-2}
        \end{scope}
        \node[anchor=south] at (\iMax-1,\dLabel) {$R^{[n]}$};
    \end{scope}

    \draw[->] (\scopeDx+\scopeDw/2,\scopeDy-1) -- (\scopeEx+\scopeEw,\scopeEy+1);
    \draw[->] (\scopeDx+\scopeDw/2,\scopeDy-1) -- (\scopeFx,\scopeFy+1) node[midway,right] {Update};
\end{tikzpicture}

\end{document}
