#############################################################
### EM algorithm for univariate normal mixtures           ###
### (mixture of two normal distributions)                 ###
### R-code by Frank Miller, 2018-2022                     ###
### Translated to Python, Julia, Matlab 2025              ###
#############################################################

#############################################################
### R                                                     ###
#############################################################
emalg <- function(dat, eps=1e-6){
  n      <- length(dat)
  pi     <- rep(NA, n)   # initialize vector for prob. to belong to group 1  
  p      <- 0.5          # starting value for mixing parameter 
  sigma1 <- sd(dat)*2/3  # starting value for variances
  sigma2 <- sigma1
  mu1    <- mean(dat) - sigma1 / 2  # starting values for means
  mu2    <- mean(dat) + sigma1 / 2
  pv     <- c(p, mu1, mu2, sigma1, sigma2)  # parameter vector
  cc     <- eps + 100    # initialize conv. crit. not to stop directly 
  while (cc>eps){
    pv1  <- pv           # save previous parameter vector
    ### E step ###
    for (j in 1:n){
      pi1   <- p*dnorm(dat[j], mean=mu1, sd=sigma1)
      pi2   <- (1-p)*dnorm(dat[j], mean=mu2, sd=sigma2)
      pi[j] <- pi1/(pi1+pi2)
    }
    ### M step ###
    p      <- mean(pi)
    mu1    <- sum(pi*dat)/(p*n)
    mu2    <- sum((1-pi)*dat)/((1-p)*n)
    sigma1 <- sqrt(sum(pi*(dat-mu1)*(dat-mu1)/(p*n)))
    sigma2 <- sqrt(sum((1-pi)*(dat-mu2)*(dat-mu2)/((1-p)*n)))
    ######
    pv     <- c(p, mu1, mu2, sigma1, sigma2)
    cc     <- t(pv-pv1)%*%(pv-pv1)
  }
  pv
}

data <- c(0.1, 0.5, 0.7, 1.1, 2.5, 3.4, 3.5, 3.9, 4.0)
emalg(data)

#############################################################
### Python                                                ###
#############################################################
import numpy as np
from scipy.stats import norm

def emalg(dat, eps=1e-6):
    dat    = np.array(dat)
    n      = len(dat)
    pi     = np.empty(n)      # initialize vector for prob. to belong to group 1
    p      = 0.5              # starting value for mixing parameter
    sigma1 = np.std(dat)*2/3  # starting value for variances
    sigma2 = sigma1
    mu1    = np.mean(dat) - sigma1 / 2  # starting values for means
    mu2    = np.mean(dat) + sigma1 / 2
    pv     = np.array([p, mu1, mu2, sigma1, sigma2])  # parameter vector
    cc     = eps + 100        # initialize conv. crit. not to stop directly
    while cc > eps:
        pv1 = pv.copy()       # save previous parameter vector
        ### E step ###
        for j in range(n):
            pi1   = p * norm.pdf(dat[j], loc=mu1, scale=sigma1)
            pi2   = (1-p) * norm.pdf(dat[j], loc=mu2, scale=sigma2)
            pi[j] = pi1 / (pi1+pi2)
        ### M step ###
        p      = np.mean(pi)
        mu1    = np.sum(pi * dat) / (p * n)
        mu2    = np.sum((1-pi) * dat) / ((1-p) * n)
        sigma1 = np.sqrt(np.sum(pi * (dat-mu1)**2) / (p * n))
        sigma2 = np.sqrt(np.sum((1-pi) * (dat-mu2)**2) / ((1-p) * n))
        ######
        pv     = np.array([p, mu1, mu2, sigma1, sigma2])
        cc     = np.sum((pv-pv1)**2)
    return pv

data = [0.1, 0.5, 0.7, 1.1, 2.5, 3.4, 3.5, 3.9, 4.0]
result = emalg(data)
print(result)

#############################################################
### Julia                                                 ###
#############################################################
using Statistics
using Distributions
using LinearAlgebra

function emalg(dat; eps=1e-6)
    n      = length(dat)
    pi     = fill(NaN, n)     # initialize vector for prob. to belong to group 1
    p      = 0.5              # starting value for mixing parameter
    sigma1 = std(dat) * 2/3   # starting value for variance
    sigma2 = sigma1
    mu1    = mean(dat) - sigma1 / 2  # starting values for means
    mu2    = mean(dat) + sigma1 / 2
    pv     = [p, mu1, mu2, sigma1, sigma2]   # parameter vector
    cc     = eps + 100.0      # initialize conv. crit. not to stop directly
    while cc > eps
        pv1 = copy(pv)
        ### E-step ###
        for j in 1:n
            pi1   = p * pdf(Normal(mu1, sigma1), dat[j])
            pi2   = (1 - p) * pdf(Normal(mu2, sigma2), dat[j])
            pi[j] = pi1 / (pi1 + pi2)
        end
        ### M-step ###
        p      = mean(pi)
        mu1    = sum(pi .* dat) / (p * n)
        mu2    = sum((1 .- pi) .* dat) / ((1 - p) * n)
        sigma1 = sqrt(sum(pi .* (dat .- mu1).^2) / (p * n))
        sigma2 = sqrt(sum((1 .- pi) .* (dat .- mu2).^2) / ((1 - p) * n))
        ######
        pv     = [p, mu1, mu2, sigma1, sigma2]
        cc     = dot(pv - pv1, pv - pv1)
    end
    return pv
end

data = [0.1, 0.5, 0.7, 1.1, 2.5, 3.4, 3.5, 3.9, 4.0]
result = emalg(data)
println(result)

#############################################################
### Matlab                                                ###
#############################################################
function pv = emalg(dat, eps)
    n      = length(dat);
    pi_vec = NaN(n, 1);       % initialize vector for prob. to belong to group 1
    p      = 0.5;             % starting value for mixing parameter
    sigma1 = std(dat) * 2/3;  % starting value for variances
    sigma2 = sigma1;
    mu1    = mean(dat) - sigma1 / 2;  % starting value for means
    mu2    = mean(dat) + sigma1 / 2;
    pv     = [p; mu1; mu2; sigma1; sigma2];  % parameter vector
    cc     = eps + 100;        % initialize conv. crit. not to stop directly
    while cc > eps
        pv1 = pv;
        % E-step
        for j = 1:n
            pi1 = p * normpdf(dat(j), mu1, sigma1);
            pi2 = (1 - p) * normpdf(dat(j), mu2, sigma2);
            pi_vec(j) = pi1 / (pi1 + pi2);
        end
        % M-step
        p      = mean(pi_vec);
        mu1    = sum(pi_vec .* dat') / (p * n);
        mu2    = sum((1 - pi_vec) .* dat') / ((1 - p) * n);
        sigma1 = sqrt(sum(pi_vec .* (dat' - mu1).^2) / (p * n));
        sigma2 = sqrt(sum((1 - pi_vec) .* (dat' - mu2).^2) / ((1 - p) * n));
        %
        pv     = [p; mu1; mu2; sigma1; sigma2];
        cc     = sum((pv - pv1).^2);
    end
end

data = [0.1, 0.5, 0.7, 1.1, 2.5, 3.4, 3.5, 3.9, 4.0];
result = emalg(data, 1e-6);
disp(result)