So the construction of the (unrestricted) Wreath Product $A\wr B$ i'm using is such that elements are in the form $(b,f)$ with a right action $f^b(x) = f(xb^{-1})$.
So the theorem says that the embedding is given by:
$\phi_\tau (g) = (\hat{g},f_g)$ where $f_g(b) = ((b\hat{g^{-1}})\tau)^{-1}\hat{g}(b\tau)$ where $\tau$ is a function that takes a coset and produces a representative for it. Now, $\mathbb{Z}_2 \times \mathbb{Z}_2$ is an extension of $\mathbb{Z}_2$ by $\mathbb{Z}_2$ and so i opted for $\tau$ as follows:
$(\mathbb{Z}_2 \times \{0\})\tau = (0,0)$ and $(\mathbb{Z}_2 \times \{1\})\tau = (1,1).$
But when I tried using the embedding given above I got that the identity of $\mathbb{Z}_2 \times \mathbb{Z}_2$ didn't map to the identity of $\mathbb{Z}_2 \wr \mathbb{Z}_2$, so i'm doing something wrong, somewhere, but I can't work out what's going on with it. Thanks!