"""Lecture 1: observations, state, and the information missing from snapshots."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from common import make_dataset, prepare_embedding


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    X, Z, _, _ = prepare_embedding(data)
    train_sessions = set(data['session'][data['train']].tolist())
    test_sessions = set(data['session'][~data['train']].tolist())
    assert train_sessions.isdisjoint(test_sessions)
    nxt = data['next_index']
    valid = nxt >= 0
    assert np.all(data['session'][valid] == data['session'][nxt[valid]])
    assert np.all(data['time'][nxt[valid]] > data['time'][valid])
    rng = np.random.default_rng(7)
    shuffled = rng.permutation(len(X))
    # Row order changes neither the centroid nor a snapshot point cloud.
    centroid_error = np.max(np.abs(X.mean(0) - X[shuffled].mean(0)))
    fig, axes = plt.subplots(1, 3, figsize=(13, 4), constrained_layout=True)
    axes[0].scatter(data['z_true'][:, 0], data['z_true'][:, 1],
                    c=data['time'], s=5, cmap='viridis')
    axes[0].set_title('Generator state (oracle, never input)')
    axes[1].scatter(Z[:, 0], Z[:, 1], c=data['time'], s=5, cmap='viridis')
    axes[1].set_title('PCA of observed features')
    sid = sorted(test_sessions)[0]
    rows = np.flatnonzero(data['session'] == sid)
    axes[2].plot(Z[rows, 0], Z[rows, 1], 'o-', ms=4)
    axes[2].annotate('observed first step', Z[rows[0]], xytext=(-8, 8),
                     textcoords='offset points', ha='right')
    axes[2].set_title('One held-out session in time order')
    for ax in axes:
        ax.set_xlabel('coordinate 1')
        ax.set_ylabel('coordinate 2')
    fig.savefig(output_dir / 'lesson01.png', dpi=160)
    plt.close(fig)
    return {'rows': int(len(X)), 'features': int(X.shape[1]),
            'train_sessions': len(train_sessions), 'test_sessions': len(test_sessions),
            'session_overlap': len(train_sessions & test_sessions),
            'valid_transitions': int(valid.sum()),
            'permutation_centroid_error': float(centroid_error)}


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