# -*- coding: utf-8 -*-

import numpy as np, scipy, arviz as az, matplotlib.pyplot as plt
from cmdstanpy import CmdStanModel

data = np.loadtxt("data/mixture_data.csv")
mdl_data = {"N": len(data), "obs": data}

modelfile = "mixture.stan"
with open(modelfile, "w") as file: file.write("""
	data {
		int<lower=0> N;
		vector[N] obs;
	}

	transformed data {
		int<lower=0> n_class = 2;
	}

	parameters { // discrete parameters impossible
		simplex[n_class] class_prob;
		ordered[n_class] centers;
		vector<lower=0>[n_class] sigmas;
	}

	model {
		centers[1] ~ normal(120, 10);
		centers[2] ~ normal(190, 10);

		sigmas ~ uniform(0, 100);

		for (n in 1:N) { // marginalize out the discrete parameter
			vector[n_class] lps = log(class_prob);
			for (k in 1:n_class)
				lps[k] += normal_lpdf(obs[n] | centers[k], sigmas[k]);
			target += log_sum_exp(lps);
		}
	}
""")

sm = CmdStanModel(stan_file = modelfile)

# maximum likelihood estimation
optim = sm.optimize(data = mdl_data).optimized_params_pd
optim[optim.columns[~optim.columns.str.startswith("lp")]]

# variational inference
vb = sm.variational(data = mdl_data)
vb.variational_sample.columns = vb.variational_params_dict.keys()
vb_name = vb.variational_params_pd.columns[~vb.variational_params_pd.columns.str.startswith(("lp", "log_"))]
vb.variational_params_pd[vb_name]
vb.variational_sample[vb_name]

# Markov chain Monte Carlo
fit = sm.sample(
	data = mdl_data, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit.draws().shape # iterations, chains, parameters
fit.summary().loc[vb_name] # pandas DataFrame
print(fit.diagnose())

posterior = fit.stan_variables()

az_trace = az.from_cmdstanpy(fit)
az.summary(az_trace).loc[vb_name] # pandas DataFrame
az.plot_trace(az_trace)

# Compute the un-normalized log probabilities for each cluster
cluster_0_log_prob = scipy.stats.norm.logpdf(
	np.expand_dims(data, axis = 1),
	loc = np.expand_dims(posterior["centers"][:, 0], axis = 0),
	scale = np.expand_dims(posterior["sigmas"][:, 0], axis = 0)
) + np.log(posterior["class_prob"][:, 0])
cluster_1_log_prob = scipy.stats.norm.logpdf(
	np.expand_dims(data, axis = 1),
	loc = np.expand_dims(posterior["centers"][:, 1], axis = 0),
	scale = np.expand_dims(posterior["sigmas"][:, 1], axis = 0)
) + np.log(posterior["class_prob"][:, 1])

# Bayes rule to compute the assignment probability: P(cluster = 1 | data) ∝ P(data | cluster = 1) P(cluster = 1)
log_p_assign_1 = cluster_1_log_prob - np.logaddexp(cluster_0_log_prob, cluster_1_log_prob)

# Average across the MCMC chain
log_p_assign_1bis = scipy.special.logsumexp(log_p_assign_1, axis=-1) - np.log(log_p_assign_1.shape[-1])

p_assign_1 = np.exp(log_p_assign_1bis)
assign_trace = p_assign_1[np.argsort(data)]

plt.figure(figsize = (8, 6), tight_layout = True)
plt.scatter(data[np.argsort(data)], assign_trace, c = 1 - assign_trace, cmap = "RdBu")
plt.title("Probability of data point belonging to cluster 1")
plt.ylabel("probability")
plt.xlabel("value of data point")
