\begin{tikzpicture}[font=\small]

% Input box 
\node[rectangle,draw,fill=white] (input) at (0,0) {$X_i$};
\node[below=0.2cm of input, text width=2cm, align=center] {\large Calibration input};

% Neural network box with pi(y|x)
\node[rectangle,draw,fill=black,text=white] (pi) at ($(input)+(2.5,2)$) {\large Generative model $\hat{\pi}(y|x)$};

% Neural network diagram
\node[rectangle,draw,minimum width=2cm,minimum height=2.5cm, fill=white] (nn) at ($(input)+(2.5,0)$) {\input{gps-overview-figure/cal/nn}};

% Output Y values
\matrix [right=0.4cm of nn] (m) {
   \node[draw,fill=red!20] (y1) {$Y_1$}; \\[0.2cm]
   \node[draw,fill=red!20] (y2) {$Y_2$}; \\[0.2cm]
   \node[draw,fill=red!20] (y3) {$Y_3$}; \\[0.2cm]
   \node[draw,fill=green!20] (y4) {$Y_4$}; \\
};

% Add right brace (positioned to the left of Y values)
\draw[decorate,decoration={brace,mirror,amplitude=10pt}] 
  ($(y1.north west)+(-0.2,0)$) -- ($(y4.south west)+(-0.2,0)$);

% Add downward arrow
\draw[->,thick] ($(y1.north)+(0.6,0)$) -- ($(y4.south)+(0.6,0)$);

% K_i label aligned with Y_4 vertically
\node[right=0.5cm of y4, draw, fill=green!20] {$K_i=4$};

% Legend (loaded from file)
\node[below=0.5cm of nn] {\input{gps-overview-figure/cal/admissible-legend}};

% Add arrow
\draw[->] (input) -- (nn);

\end{tikzpicture} 