Polynomial equations for matrices over integers modulo a prime power and the cokernel of a random matrix

Given a prime $p$ and a positive integer $k$, let $\mathrm{M}_{n}(\mathbb{Z}/p^{k}\mathbb{Z})$ be the ring of $n \times n$ matrices over $\mathbb{Z}/p^{k}\mathbb{Z}$. We consider the number of solutions $X \in \mathrm{M}_{n}(\mathbb{Z}/p^{k}\mathbb{Z})$ to the polynomial equation $P(X) = 0$, where $P(t)$ is a monic polynomial in $(\mathbb{Z}/p^{k}\mathbb{Z})[t]$ whose reduction modulo $p$ is square-free over the finite field $\mathbb{F}_{p}$ of $p$ elements. Noting that $P(X) = 0$ if and only if $\mathrm{cok}(P(X)) \simeq (\mathbb{Z}/p^{k}\mathbb{Z})^{n}$, we give a conjectural generalization of counting solutions to $P(X) = 0$ as the distribution of the cokernel $\mathrm{cok}(P(X))$ of $P(X)$ up to isomorphisms, where $X$ is a uniform random matrix in $\mathrm{M}_{n}(\mathbb{Z}/p^{k}\mathbb{Z})$. This distribution involves an explicit formula when we fix the residue class of $X$ modulo $p$. We prove this conjecture for the special case when the image of $P(t)$ in $\mathbb{F}_{p}[t]$ modulo $p$ is irreducible. We explain how the distribution we obtain is closely related to the Cohen-Lenstra distribution. Our proof involves algebraic and combinatorial arguments in linear algebra over $\mathbb{Z}/p^{k}\mathbb{Z}$ and builds upon a previous work of Cheong and Kaplan.