"""Lesson 8: within-session Markov counts and Gaussian HMM inference."""
from pathlib import Path
import argparse
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.special import logsumexp
from sklearn.cluster import KMeans
from common import make_dataset, prepare_embedding


def estimate_transitions(labels, next_index, train, n_states, smoothing=0.5):
    counts = np.full((n_states, n_states), smoothing, dtype=float)
    rows = np.flatnonzero(train & (next_index >= 0))
    rows = rows[train[next_index[rows]]]
    np.add.at(counts, (labels[rows], labels[next_index[rows]]), 1)
    return counts / counts.sum(axis=1, keepdims=True)


def gaussian_log_emission(observations, means, variances):
    return -0.5 * np.sum(np.log(2 * np.pi * variances)[None, :, :] +
                         (observations[:, None, :] - means[None, :, :]) ** 2 / variances[None, :, :], axis=2)


def infer_hmm(log_emission, transition, initial):
    """Log-domain forward/backward and Viterbi; fixed model, no EM fitting."""
    T, K = log_emission.shape
    log_A, log_pi = np.log(transition), np.log(initial)
    forward = np.empty((T, K))
    forward[0] = log_pi + log_emission[0]
    score = forward.copy()
    backpointer = np.zeros((T, K), dtype=int)
    for t in range(1, T):
        forward[t] = log_emission[t] + logsumexp(forward[t - 1, :, None] + log_A, axis=0)
        candidates = score[t - 1, :, None] + log_A
        backpointer[t] = np.argmax(candidates, axis=0)
        score[t] = log_emission[t] + np.max(candidates, axis=0)
    log_likelihood = logsumexp(forward[-1])
    backward = np.zeros((T, K))
    for t in range(T - 2, -1, -1):
        backward[t] = logsumexp(log_A + log_emission[t + 1][None, :] + backward[t + 1][None, :], axis=1)
    posterior = np.exp(forward + backward - log_likelihood)
    path = np.empty(T, dtype=int)
    path[-1] = np.argmax(score[-1])
    for t in range(T - 2, -1, -1):
        path[t] = backpointer[t + 1, path[t + 1]]
    return float(log_likelihood), posterior, path


def run(output_dir: Path) -> dict:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    _, Z, _, _ = prepare_embedding(data, n_components=6)
    train = data["train"]
    K = 6
    km = KMeans(n_clusters=K, n_init=10, random_state=7).fit(Z[train])
    labels = km.predict(Z)
    A = estimate_transitions(labels, data["next_index"], train, K)
    variances = np.array([Z[train & (labels == k)].var(axis=0) + 0.05 for k in range(K)])
    initial_counts = np.full(K, 0.5)
    for session in np.unique(data["session"][train]):
        rows = np.flatnonzero(data["session"] == session)
        first = rows[np.argmin(data["time"][rows])]
        initial_counts[labels[first]] += 1
    initial = initial_counts / initial_counts.sum()
    next_index = data["next_index"]
    test_pairs = np.flatnonzero(~train & (next_index >= 0))
    next_labels = labels[next_index[test_pairs]]
    markov_nll = -np.log(A[labels[test_pairs], next_labels]).mean()
    train_pairs = np.flatnonzero(train & (next_index >= 0))
    marginal = np.bincount(labels[next_index[train_pairs]], minlength=K) + 0.5
    marginal = marginal / marginal.sum()
    marginal_nll = -np.log(marginal[next_labels]).mean()
    total_loglik, total_rows, max_posterior_error = 0.0, 0, 0.0
    first_result = None
    for session in np.unique(data["session"][~train]):
        rows = np.flatnonzero(data["session"] == session)
        rows = rows[np.argsort(data["time"][rows])]
        emission = gaussian_log_emission(Z[rows], km.cluster_centers_, variances)
        likelihood, posterior, path = infer_hmm(emission, A, initial)
        total_loglik += likelihood
        total_rows += len(rows)
        max_posterior_error = max(max_posterior_error, float(np.max(np.abs(posterior.sum(axis=1) - 1))))
        if first_result is None:
            first_result = (rows, posterior, path)
    rows, posterior, path = first_result
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    im = axes[0].imshow(A, vmin=0, vmax=1, cmap="Blues")
    axes[0].set(xlabel="Next cluster", ylabel="Current cluster", title="Training transition matrix")
    fig.colorbar(im, ax=axes[0])
    im = axes[1].imshow(posterior.T, origin="lower", aspect="auto", vmin=0, vmax=1, cmap="Blues")
    axes[1].plot(np.arange(len(path)), path, "r.-", label="Viterbi path")
    axes[1].set(xlabel="Within-session step", ylabel="Hidden state", title="Held-out session: smoothed posterior")
    axes[1].legend()
    fig.colorbar(im, ax=axes[1])
    fig.tight_layout()
    fig.savefig(output_dir / "lesson08.png", dpi=150)
    plt.close(fig)
    return {"heldout_transition_nll": float(markov_nll), "heldout_marginal_nll": float(marginal_nll),
            "heldout_hmm_log_density_per_observation": float(total_loglik / total_rows),
            "posterior_row_sum_max_error": max_posterior_error,
            "training_transition_pairs": int(len(train_pairs)), "heldout_transition_pairs": int(len(test_pairs)),
            "evaluation": "HMM parameters initialized from training k-means labels; no Baum-Welch optimization"}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    print(json.dumps(run(parser.parse_args().output), indent=2, allow_nan=False))