Let $D$ be the diagonal matrix and $A^2=D$
So far I've taken 2 approaches, the first being a direct algebraic approach. For a $2 \times 2$ matrix this reduces to $$ \left( \begin{matrix} a_{11} & a_{12} \\ a_{21} & a_{22} \\ \end{matrix} \right )^2 = \left( \begin{matrix} a_{11}^2 +a_{21}a_{12} & a_{12}(a_{11} + a_{22}) \\ a_{21}(a_{11} + a_{22}) & a_{22}^2 +a_{21}a_{12} \\ \end{matrix} \right ) = \left( \begin{matrix} d_1 & 0 \\ 0 & d_2 \\ \end{matrix} \right )$$ where $d_1 \ne d_2$.
From the top right and bottom left squares it's easy to see that $a_{12}=a_{21}=0$ or $a_{11} +a_{22}=0$
If $a_{12}=a_{21}=0$ then, $$d_1=a_{11}^2 +a_{21}a_{12}=a_{22}^2 +a_{21}a_{12}=d_2$$ which means the eigenvalues aren't distinct. Thus $A$ is diagonal.
This proves it when $A$ is $2 \times 2 $.
For $n \times n $ $A$, with $n$ linearly independent eigenvectors $\mathbf v_1 , ..., \mathbf v_n$ with corresponding eigenvalues $ \lambda_1 , ... , \lambda_n$ such that for $j \ne k$ $ \lambda^2_j \ne \lambda^2_k$ Now $$ D \mathbf v_k =A^2 \mathbf v_k= \lambda^2_k \mathbf v_k$$
thus $\lambda_k^2$ is and eigenvalue of $D$ and $\mathbf v_k$ is the corresponding basis vector, hence the eigenvectors of $A$ are just the standard basis vectors and hence $A$ is diagonal
Is my conjecture that a square root of a diagonal matrix with distinct eigenvalues is diagonal wrong or is there a general proof?