Two different versions for exam, solution file needs to have a figure. Float(s) lost error with current code

Question

I'm writing an exam for my class. The problems file needs to have an empty box after each question, while the solutions file should have the solution (preferably in a different colour). I got the code to do this from here.

Unfortunately, this only works for those solutions that have only text/math in them. I also have figures, and that's throwing a "float(s) lost" error. Here is the MWE.

\documentclass{article}
\usepackage{etex}
\usepackage{pdfsync}
\usepackage{amsfonts, amsmath, amssymb, amsthm, thmtools}
\usepackage{hyperref} 
\usepackage[dvipsnames,svgnames,x11names,hyperref]{xcolor} 
\usepackage{subcaption}
\usepackage{pgfplots}

%%%%%%%%%%%%%%%%%%%%%% 
% STUFF FOR SOLUTION ENVIRONMENT
\usepackage{ifthen}

\newenvironment{solution}[1]{%
    \ifthenelse{\isundefined{\teacher}}{%
        \fbox{\begin{minipage}{\linewidth}\hfill\vspace*{#1}\end{minipage}}%
    }{}%
    \newbox\tempbox%
    \begin{lrbox}{\tempbox}\begin{minipage}{\linewidth}%
}{%
    \end{minipage}\end{lrbox}%
    \ifthenelse{\not\isundefined{\teacher}}{%
        \medskip%
        \fbox{\usebox{\tempbox}}%
        \medskip%
    }{}%
}

\def\teacher{true} %comment this out for empty box (in student's version)
%%%%%%%%%%%%%%%%%%%%%%%%

\begin{document}
\begin{enumerate}
    \item This is the first problem on the test. \\
        \begin{solution}{5cm}
            This is the first solution. 
        \end{solution}

    \item This is the second problem. \\
        \begin{solution}{5cm}
            This is the second solutiuon.
            %%%%%%%%%%%%%%%%%%%% 
            % COMMENT OUT FIGURE BELOW TO SEE IT WORK
                \begin{figure}[h]
                \centering
                \begin{tikzpicture}
                \begin{axis} [
                height=5cm, width=7cm, xlabel={$t$}, ylabel={$x(t)$},
                xmin=-4, xmax=4, ymin=-1.5, ymax=1.5, no markers,
                xtick={-2.5, -.5, .5, 2.5}, ytick={1, -1},
                axis lines=middle
                ]
                \addplot[thick, black] coordinates {(-1, 0) (0, 1) (1, 0)};
                \end{axis}
                \end{tikzpicture}
                \caption{Output signal for Problem 3a}\label{fig:samplefig}
                \end{figure}
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        \end{solution}
\end{enumerate}

\end{document}

Answer 1

Unfortunately, it's not obvious how to use figure in a minipage environment, see How to use figure inside a minipage?.

If you can avoid the figure environment, you can simply do this:

\documentclass{article}
\usepackage{etex}
\usepackage{pdfsync}
\usepackage{amsfonts}
\usepackage{amsmath}
\usepackage{hyperref} 
\usepackage[dvipsnames,svgnames,x11names,hyperref]{xcolor} 
\usepackage{amssymb}
\usepackage{amsthm, thmtools}
\usepackage{subcaption}
\usepackage{pgfplots}
a%%%%%%%%%%%%%%%%%%%%%% 
% STUFF FOR SOLUTION ENVIRONMENT
\usepackage{ifthen}

\newenvironment{solution}[1]{%
    \ifthenelse{\isundefined{\teacher}}{%
        \fbox{\begin{minipage}{\linewidth}\hfill\vspace*{#1}\end{minipage}}%
    }{}%
    \newbox\tempbox%
    \begin{lrbox}{\tempbox}\begin{minipage}{\linewidth}%
}{%
    \end{minipage}\end{lrbox}%
    \ifthenelse{\not\isundefined{\teacher}}{%
        \medskip%
        \fbox{\usebox{\tempbox}}%
        \medskip%
    }{}%
}

\def\teacher{true} %comment this out for empty box
%%%%%%%%%%%%%%%%%%%%%%%%

\begin{document}
\begin{enumerate}
    \item This is the first problem on the test. \\
        \begin{solution}{5cm}
            This is the first solution. 
        \end{solution}

    \item This is the second problem. \\
        \begin{solution}{5cm}
            This is the second solutiuon.
            %%%%%%%%%%%%%%%%%%%% 
            % COMMENT OUT FIGURE BELOW TO SEE IT WORK
                \begin{center}
                    \begin{tikzpicture}
                    \begin{axis} [
                    height=5cm, width=7cm, xlabel={$t$}, ylabel={$x(t)$},
                    xmin=-4, xmax=4, ymin=-1.5, ymax=1.5, no markers,
                    xtick={-2.5, -.5, .5, 2.5}, ytick={1, -1},
                    axis lines=middle
                    ]
                    \addplot[thick, black] coordinates {(-1, 0) (0, 1) (1, 0)};
                    \end{axis}
                    \end{tikzpicture}
                    {Output signal for Problem 3a}\label{fig:samplefig}
                \end{center}
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        \end{solution}
\end{enumerate}

\end{document}