export  GaussianMixtureModel, PoissonMixtureModel, fit!, log_likelihood, sample, E_Step, M_Step!, MixtureModel

"""
    sample(model::MixtureModel, n::Int)

Draw `n` samples from the given mixture model.

# Returns
- A matrix or vector of samples (with shape 'n' by 'data dimension'), depending on mixture model type. 

# Examples
```julia
# Gaussian Mixture Model example
gmm = GaussianMixtureModel(3, 2)  # Create a Gaussian Mixture Model with 3 clusters and 2-dimensional data
gmm_samples = sample(gmm, 100)  # Draw 100 samples from the Gaussian Mixture Model
```
"""
function sample(mixture_model::MixtureModel, n::Int) end



"""
    fit!(model::MixtureModel, data::AbstractMatrix; <keyword arguments>)

Fit the given mixture model to the data using the Expectation-Maximization (EM) algorithm. 

# Arguments
- `model::MixtureModel`: The mixture model to fit.
- `data::AbstractMatrix`: The data matrix where rows correspond to data points.
- `maxiter::Int=50`: The maximum number of iterations for the EM algorithm (default is 50).
- `tol::Float64=1e-3`: The convergence tolerance for the log-likelihood (default is 1e-3).
- `initialize_kmeans::Bool=true`: Whether to initialize the model parameters using KMeans (default is true).

# Returns
- The class probabilities (responsibilities) for each data point.

# Examples
```julia
gmm = GaussianMixtureModel(3, 2)  # Create a Gaussian Mixture Model with 3 clusters and 2-dimensional data
fit!(gmm, data)  # Fit the model to the data
```
"""
function fit!(gmm::MixtureModel, data::AbstractMatrix; maxiter::Int=50, tol::Float64=1e-3, initialize_kmeans::Bool=true) end




"""
    GaussianMixtureModel

A Gaussian Mixture Model for clustering and density estimation.

# Fields
- `k::Int`: Number of clusters.
- `μₖ::Matrix{<:Real}`: Means of each cluster (dimensions: data_dim x k).
- `Σₖ::Array{Matrix{<:Real}, 1}`: Covariance matrices of each cluster.
- `πₖ::Vector{Float64}`: Mixing coefficients for each cluster.


# Examples
```julia
gmm = GaussianMixtureModel(3, 2) # Create a Gaussian Mixture Model with 3 clusters and 2-dimensional data
fit!(gmm, data)
```
"""
mutable struct GaussianMixtureModel <: MixtureModel
    k::Int # Number of clusters
    μₖ::Matrix{<:Real} # Means of each cluster
    Σₖ::Array{Matrix{<:Real}, 1} # Covariance matrices of each cluster
    πₖ::Vector{Float64} # Mixing coefficients
end


"""
    GaussianMixtureModel(k::Int, data_dim::Int)
    
Constructor for GaussianMixtureModel. Initializes Σₖ's covariance matrices to the 
identity, πₖ to a uniform distribution, and μₖ's means to zeros.

"""
function GaussianMixtureModel(k::Int, data_dim::Int)
    Σ = [Matrix{Float64}(I, data_dim, data_dim) for _ = 1:k]
    πs = ones(k) ./ k
    μ = zeros(Float64, k, data_dim)  # Mean of each cluster initialized to zero matrix
    return GaussianMixtureModel(k, μ, Σ, πs)
end

"""
Draw 'n' samples from gmm. Returns a Matrix{<:Real}, where each row is a data point.
"""
function sample(gmm::GaussianMixtureModel, n::Int)
    # Determine the number of samples from each component
    component_samples = rand(Multinomial(n, gmm.πₖ), 1)
    
    # Initialize a container for all samples
    samples = Matrix{Float64}(undef, n, size(gmm.μₖ, 2))
    start_idx = 1
    
    for i in 1:gmm.k
        num_samples = component_samples[i]
        if num_samples > 0
            # Sample all at once from the i-th Gaussian component
            dist = MvNormal(gmm.μₖ[i, :], gmm.Σₖ[i])
            samples[start_idx:(start_idx + num_samples - 1), :] = transpose(rand(dist, num_samples))
            start_idx += num_samples
        end
    end
    return samples
end

# E-Step (Confirmed in Python)
"""
    Returns 
        - `class_probabilities::Matrix{<:Real}`: Probability of each class for each data point.

"""
function E_Step(gmm::GaussianMixtureModel, data::Matrix{<:Real})
    N, _ = size(data)
    K = gmm.k
    γ = zeros(N, K)
    log_γ = zeros(N, K)
    class_probabilities = zeros(N, K)

    for n in 1:N
        for k in 1:K
            distribution = MvNormal(gmm.μₖ[k, :], gmm.Σₖ[k])
            log_γ[n, k] = log(gmm.πₖ[k]) + logpdf(distribution, data[n, :])
        end
        logsum = logsumexp(log_γ[n, :])
        γ[n, :] = exp.(log_γ[n, :] .- logsum)
    end
    # Return probabilities of each class for each point
    return γ
end

function M_Step!(gmm::GaussianMixtureModel, data::Matrix{<:Real}, class_probabilities::Matrix{<:Real})
    N, D = size(data)
    K = gmm.k
    γ = class_probabilities  

    N_k = zeros(K)
    μₖ = zeros(K, D)
    Σₖ = zeros(D, D, K)

    for k in 1:K
        N_k[k] = sum(γ[:, k])
        μₖ[k, :] = (γ[:, k]' * data) ./ N_k[k]
    end

    for k in 1:K
        x_n = data .- μₖ[k, :]'
        Σₖ[:,:,k] = ((γ[:, k] .* x_n)' * x_n ./ (N_k[k] + I*1e-6)) + (I * 1e-6)
        if !ishermitian(Σₖ[:, :, k])
            Σₖ[:,:,k] = 0.5 * (Σₖ[:,:,k] + Σₖ[:,:,k]')
        end
        gmm.πₖ[k] = N_k[k] / N
    end

    gmm.μₖ = μₖ
    gmm.Σₖ = [Σₖ[:,:,k] for k in 1:K]
end


"""
    log_likelihood(gmm::GaussianMixtureModel, data::Matrix{<:Real})

Compute the log-likelihood of the data given the Gaussian Mixture Model (GMM). The data matrix should be of shape (# observations, # features).

# Returns
- `Float64`: The log-likelihood of the data given the model.
"""
function log_likelihood(gmm::GaussianMixtureModel, data::Matrix{<:Real})
    N, K = size(data, 1), gmm.k
    ll = 0.0
    for n in 1:N
        log_probabilities = [log(gmm.πₖ[k]) + logpdf(MvNormal(gmm.μₖ[k, :], gmm.Σₖ[k]), data[n, :]) for k in 1:K]
        max_log_prob = maximum(log_probabilities)
        ll_n = max_log_prob + log(sum(exp(log_prob - max_log_prob) for log_prob in log_probabilities))
        ll += ll_n
    end
    return ll
end

"""
    fit!(gmm::GaussianMixtureModel, data::Matrix{<:Real}; <keyword arguments>)
Fits a Gaussian Mixture Model (GMM) to the given data using the Expectation-Maximization (EM) algorithm.

# Arguments
- `gmm::GaussianMixtureModel`: The Gaussian Mixture Model to be fitted.
- `data::Matrix{<:Real}`: The dataset on which the model will be fitted, where each row represents a data point.
- `maxiter::Int=50`: The maximum number of iterations for the EM algorithm (default: 50).
- `tol::Float64=1e-3`: The tolerance for convergence. The algorithm stops if the change in log-likelihood between iterations is less than this value (default: 1e-3).
- `initialize_kmeans::Bool=false`: If true, initializes the means of the GMM using K-means++ initialization (default: false).

# Returns
- `class_probabilities`: A matrix where each entry (i, k) represents the probability of the i-th data point belonging to the k-th component of the mixture model.

# Example
```julia
data = rand(2, 100)  # Generate some random data
gmm = GaussianMixtureModel(k=3, d=2)  # Initialize a GMM with 3 components and 2-dimensional data
class_probabilities = fit!(gmm, data, maxiter=100, tol=1e-4, initialize_kmeans=true)
```
"""
function fit!(gmm::GaussianMixtureModel, data::Matrix{<:Real}; maxiter::Int=50, tol::Float64=1e-3, initialize_kmeans::Bool=false)
    prev_ll = -Inf  # Initialize to negative infinity

    if initialize_kmeans
        gmm.μₖ = permutedims(kmeanspp_initialization(data, gmm.k))
    end

    for i = 1:maxiter
        # E-Step
        class_probabilities = E_Step(gmm, data)
        # M-Step
        M_Step!(gmm, data, class_probabilities)
        # Calculate current log-likelihood
        curr_ll = log_likelihood(gmm, data)
        # Debug: Output log-likelihood
        println("Iteration: $i, Log-likelihood: $curr_ll")
        # Check for convergence
        if abs(curr_ll - prev_ll) < tol
            println("Convergence reached at iteration $i")
            break
        end
        prev_ll = curr_ll
    end

    class_probabilities = E_Step(gmm, data)
    return class_probabilities
end

# Handle vector data by reshaping it into a 2D matrix with a single column
function E_Step(gmm::GaussianMixtureModel, data::Vector{Float64})
    E_Step(gmm, reshape(data, :, 1))
end

function M_Step!(gmm::GaussianMixtureModel, data::Vector{Float64}, class_probabilities::Matrix{<:Real})
    M_Step!(gmm, reshape(data, :, 1), class_probabilities::Matrix{<:Real})
end

function log_likelihood(gmm::GaussianMixtureModel, data::Vector{Float64})
    log_likelihood(gmm, reshape(data, :, 1))
end

function fit!(gmm::GaussianMixtureModel, data::Vector{Float64}; maxiter::Int=50, tol::Float64=1e-3, initialize_kmeans::Bool=true)
    fit!(gmm, reshape(data, :, 1); maxiter=maxiter, tol=tol, initialize_kmeans=initialize_kmeans)
end

"""
    PoissonMixtureModel

A Poisson Mixture Model for clustering and density estimation.

## Fields
- `k::Int`: Number of poisson-distributed clusters.
- `λₖ::Vector{Float64}`: Means of each cluster.
- `πₖ::Vector{Float64}`: Mixing coefficients for each cluster.


## Examples
```julia
pmm = PoissonMixtureModel(3) # 3 clusters, 2-dimensional data
fit!(pmm, data)```
"""
mutable struct PoissonMixtureModel <: MixtureModel
    k::Int # Number of clusters
    λₖ::Vector{Float64} # Means of each cluster
    πₖ::Vector{Float64} # Mixing coefficients
end


"""
    PoissonMixtureModel(k::Int)
    
Constructor for PoissonMixtureModel. Initializes λₖ's means to 
ones and πₖ to a uniform distribution.

"""
function PoissonMixtureModel(k::Int)
    λs = ones(k)
    πs = ones(k) ./ k
    return PoissonMixtureModel(k, λs, πs)
end

"""E-Step for PMM"""
function E_Step(pmm::PoissonMixtureModel, data::Matrix{Int})
    N, _ = size(data)
    γ = zeros(N, pmm.k)
    
    for n in 1:N
        for k in 1:pmm.k
            λk = pmm.λₖ[k]
            log_γnk = log(pmm.πₖ[k]) + logpdf(Poisson(λk), data[n, 1])
            γ[n, k] = exp(log_γnk)  # Direct computation of responsibilities
        end
        γ[n, :] /= sum(γ[n, :])  # Normalize responsibilities
    end
    
    return γ  # Return the responsibility matrix
end

"""M-Step for PMM"""
function M_Step!(pmm::PoissonMixtureModel, data::Matrix{Int}, γ::Matrix{<:Real})
    N, _ = size(data)
    
    for k in 1:pmm.k
        Nk = sum(γ[:, k])
        pmm.λₖ[k] = sum(γ[:, k] .* data) / Nk  # Update λk
        pmm.πₖ[k] = Nk / N  # Update mixing coefficient
    end
end

"""
    fit!(pmm::PoissonMixtureModel, data::Matrix{Int}; <keyword arguments>)

Fits a Poisson Mixture Model (PMM) to the given data using the Expectation-Maximization (EM) algorithm.

# Arguments
- `pmm::PoissonMixtureModel`: The Poisson Mixture Model to be fitted.
- `data::Matrix{Int}`: The dataset on which the model will be fitted, where each row represents a data point.
- `maxiter::Int=50`: The maximum number of iterations for the EM algorithm (default: 50).
- `tol::Float64=1e-3`: The tolerance for convergence. The algorithm stops if the change in log-likelihood between iterations is less than this value (default: 1e-3).
- `initialize_kmeans::Bool=false`: If true, initializes the means of the PMM using K-means++ initialization (default: false).

# Returns
- `class_probabilities`: A matrix where each entry (i, k) represents the probability of the i-th data point belonging to the k-th component of the mixture model.

# Example
```julia
data = rand(1:10, 100, 1)  # Generate some random integer data
pmm = PoissonMixtureModel(k=3)  # Initialize a PMM with 3 components
class_probabilities = fit!(pmm, data, maxiter=100, tol=1e-4, initialize_kmeans=true)
```
"""
function fit!(pmm::PoissonMixtureModel, data::Matrix{Int}; maxiter::Int=50, tol::Float64=1e-3, initialize_kmeans::Bool=false)
    prev_ll = -Inf  # Initialize previous log likelihood to negative infinity

    if initialize_kmeans
        λₖ_matrix = permutedims(kmeanspp_initialization(Float64.(data), pmm.k))
        pmm.λₖ = vec(λₖ_matrix)
    end

    for iter in 1:maxiter
        γ = E_Step(pmm, data)  # E-Step
        M_Step!(pmm, data, γ)  # M-Step
        curr_ll = log_likelihood(pmm, data)  # Current log likelihood

        println("Iteration: $iter, Log-likelihood: $curr_ll")

        if abs(curr_ll - prev_ll) < tol  # Check for convergence
            println("Convergence reached at iteration $iter")
            break
        end
        prev_ll = curr_ll  # Update previous log likelihood
    end
end

"""
    log_likelihood(pmm::PoissonMixtureModel, data::Matrix{Int})

Compute the log-likelihood of the data given the Poisson Mixture Model (PMM). The data matrix should be of shape (# observations, # features).

# Returns
- `Float64`: The log-likelihood of the data given the model.
"""
function log_likelihood(pmm::PoissonMixtureModel, data::Matrix{Int})
    ll = 0.0
    for n in 1:size(data, 1)
        ll_n = log(sum([pmm.πₖ[k] * pdf(Poisson(pmm.λₖ[k]), data[n, 1]) for k in 1:pmm.k]))
        ll += ll_n
    end
    return ll
end

"""
    sample(pmm::PoissonMixtureModel, n)

Draw 'n' samples from pmm. Returns a Vector{Int} of lenth n.

"""
function sample(pmm::PoissonMixtureModel, n::Int)
    # Determine the number of samples from each component
    component_samples = rand(Multinomial(n, pmm.πₖ), 1)
    
    # Initialize a container for all samples
    samples = Vector{Int}(undef, n)
    start_idx = 1
    
    for i in 1:pmm.k
        num_samples = component_samples[i]
        if num_samples > 0
            # Sample all at once from the i-th Poisson component
            λ = pmm.λₖ[i] # λ for the i-th component
            samples[start_idx:(start_idx + num_samples - 1)] = rand(Poisson(λ), num_samples)
            start_idx += num_samples
        end
    end
    return samples
end


# Handle vector data by reshaping it into a 2D matrix with a single column
function E_Step(pmm::PoissonMixtureModel, data::Vector{Int})
    E_Step(pmm, reshape(data, :, 1))
end

function M_Step!(pmm::PoissonMixtureModel, data::Vector{Int}, class_probabilities::Matrix{<:Real})
    M_Step!(pmm, reshape(data, :, 1), class_probabilities)
end

function log_likelihood(pmm::PoissonMixtureModel, data::Vector{Int})
    log_likelihood(pmm, reshape(data, :, 1))
end

function fit!(pmm::PoissonMixtureModel, data::Vector{Int}; maxiter::Int=50, tol::Float64=1e-3, initialize_kmeans::Bool=true)
    fit!(pmm, reshape(data, :, 1); maxiter=maxiter, tol=tol, initialize_kmeans=initialize_kmeans)
end