Derivation of the M-step in EM algorithm for a three-dimensional panel mixture model I have a question regarding the estimation of a latent-class gaussian mixture model, where the model is for three dimensional panel data set with individuals $i$, in country $j$ in time $t$.
I want the classes to vary over the individuals, but within the classes there to be a specific effect for the country dimension $j$. However, I have some trouble deriving all the steps for the M-step with the additional dimensions and the parameters. Especially, I don't know how to do the maximization step in the last part.
My specification is as follows:
\begin{equation}
\label{eq: mixture spec}
    y_{ijt} = x_{{s_i}jt}' \beta_{s_i} + \alpha_{s_i} + \gamma_{{s_i}j}  +  \epsilon_{ijt}
\end{equation}
\noindent
where $s_i$ is a latent, unobserved variable that we treat as a stochastic variable with $P[S_i = s] = p_s$ for $s = 1, ..., K$ and $\sum_{s=1}^K p_s = 1$. Mixture models are often estimated with the EM algorithm. This is an iterative algorithm that in the E-step calculates the posterior probabilities $\tilde{p}_s = P[s_i = s | y, \hat{\theta}]$ given the current set of parameter estimates $\theta$ and in the M-step maximises the expected log-likelihood function with respect to the set of parameters $\theta$.
Given this specification we have the following likelihood function:
\begin{multline}
    L(\theta) = \prod_{i=1}^N \sum_{s=1}^K p_s \left(\prod_{j=1}^M \prod_{t=1}^T \phi(y_{ijt}, x_{ijt}; \beta_{s_i}, \alpha_{s_i}, \gamma_{{s_i}j}) \right)  
    \\ = \prod_{i=1}^N \sum_{s=1}^K p_s \left(\prod_{j=1}^M \prod_{t=1}^T 
    \frac{1}{\sigma_{\varepsilon} \sqrt{2 \pi}} \exp \left(-\frac{1}{2 \sigma_{\varepsilon}^{2}}\left( y_{ijt} - x_{{s_i}jt}' \beta_{s_i} - \alpha_{s_i} - \gamma_{{s_i}j}\right)^{2}\right) \right)
\end{multline}
For the EM algorithm we consider the complete data likelihood function:
\begin{equation}
    L_j(\theta) = \prod_{i=1}^N  \prod_{s} \left(p_{s} \prod_{j=1}^M \prod_{t=1}^T \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\epsilon}) \right) ^{I(s_i = s)}
\end{equation}
and the subsequent log-complete likelihood:
\begin{equation}
\label{eq: log complete}
    \ell_{j}(\theta)=\log L_{j}(\theta)= \sum_{i=1}^N \sum_{s=1}^K I(s_i = s) (\log \: p_s  + \sum_{j=1}^M \sum_{t=1}^T \left(\log \: \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\epsilon}) \right) )
\end{equation}
In the E-step we calculate:
\begin{equation}
    \pi_{i s} \equiv \mathrm{E}\left[I\left(s_{i}=s\right) \mid y_{ijt}, x_{ijt} \right]=   \frac{  \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\epsilon}) p_{s}}   {\sum_{k=1}^K \phi(y_{ijt}, x_{ijt}; \beta_{k}, \alpha_{k}, \gamma_{{k}j}) p_k}
\end{equation}
Substituting this in the log-complete likelihood gives:
\begin{equation}
    \ell_{j}(\theta)=\log L_{j}(\theta)= \sum_{i=1}^N \sum_{s=1}^K \pi_{i s} (\log \: p_s  + \sum_{j=1}^M \sum_{t=1}^T \left(\log \: \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\epsilon}) \right) )
\end{equation}
And the M-step results in maximising:
\begin{equation}
        \max _{p, \alpha, \beta, \gamma, \sigma} \left(  \left(\sum_{s=1}^{K} \sum_{i=1}^{N} \pi_{i s} \log p_{s}\right)+\left(\sum_{s=1}^{K} \sum_{i=1}^{N} \pi_{is}  \sum_{j=1}^M \sum_{t=1}^T \log \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\epsilon})\right) \right)
\end{equation}
Solving
\begin{equation}
    max_{p}\left(\sum_{s=1}^{K} \sum_{i=1}^{N} \pi_{i s} \log p_{s}\right) \text { subject to } \sum_{s=1}^{K} p_{s}=1
\end{equation}
yields
\begin{equation}
p_{s}=\frac{1}{i} \sum_{i=1}^{i}\pi_{is}
\end{equation}
For the second part:
\begin{multline}
    \max _{p, \alpha, \beta, \gamma, \sigma}  \left(\sum_{s=1}^{K} \sum_{i=1}^{N} \pi_{is}  \sum_{j=1}^M \sum_{t=1}^T \log \phi(y_{ijt}, x_{ijt}; \beta_{s}, \alpha_{s}, \gamma_{{s}j}, \sigma_{\varepsilon})\right) 
    \\ =  max _{p, \alpha, \beta, \gamma, \sigma}  \left(  \sum_{s=1}^{K} \sum_{i=1}^{N}  \pi_{is} \sum_{j=1}^M \sum_{t=1}^T -\frac{1}{2} \log \left( 2 \pi \sigma_{\varepsilon}^{2}\right) - \frac{1}{2 \sigma_{\varepsilon}^{2}} \left( y_{ijt} - x_{sjt}' \beta_{s} - \alpha_{s} - \gamma_{{s}j}\right)^{2} \right)
\end{multline}

One can rewrite the target as
$$
\sum_{i=1}^{N}  \pi_{is} \sum_{j=1}^M \sum_{t=1}^T \left( y_{ijt} - x_{sjt}' \beta_{s} - \alpha_{s} - \gamma_{{s}j}\right)^{2}=\sum_{i=1}^{N}  \pi_{is} 
\vert\vert y_{i\cdot\cdot}-X^\sf{T}\theta_s\vert\vert^2
$$
where is the $M\cdot T$ vector with components $y_{ijt}$ and $X$ is the $(M\cdot T,M+\ell+1)$ matrix with rows $(x_{sjt}^\prime,1,\mathbb I_1(j),\ldots,\mathbb I_M(j))$, $\ell$ being the dimension of $x_{sjt}$
