"""Lecture 2: centered PCA, SVD identity, and unseen-session reconstruction."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from common import make_dataset


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    train = data['train']
    scaler = StandardScaler().fit(data['X'][train])
    X = scaler.transform(data['X'])
    pca = PCA(n_components=2, svd_solver='full').fit(X[train])
    Z = pca.transform(X)
    reconstruction = pca.inverse_transform(Z)
    centered = X[train] - pca.mean_
    _, singular_values, vt = np.linalg.svd(centered, full_matrices=False)
    W = vt[:2].T
    svd_projection = centered @ W @ W.T
    sklearn_projection = reconstruction[train] - pca.mean_
    observed_sse = np.sum((centered - svd_projection) ** 2)
    discarded_sse = np.sum(singular_values[2:] ** 2)
    assert np.allclose(observed_sse, discarded_sse)
    errors = []
    dimensions = list(range(1, X.shape[1] + 1))
    for d in dimensions:
        model = PCA(n_components=d, svd_solver='full').fit(X[train])
        estimate = model.inverse_transform(model.transform(X[~train]))
        errors.append(float(np.mean((X[~train] - estimate) ** 2)))
    fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)
    axes[0].scatter(Z[~train, 0], Z[~train, 1], c=data['time'][~train], s=9)
    axes[0].set(xlabel='PC1', ylabel='PC2', title='Held-out sessions')
    axes[1].plot(dimensions, errors, 'o-')
    axes[1].set(xlabel='Retained dimensions', ylabel='Standardized test MSE',
                title='Reconstruction is one evaluation target')
    fig.savefig(output_dir / 'lesson02.png', dpi=160)
    plt.close(fig)
    return {'test_mse_2d': float(np.mean((X[~train] - reconstruction[~train]) ** 2)),
            'train_mse_2d': float(np.mean((X[train] - reconstruction[train]) ** 2)),
            'variance_ratio_2d': float(pca.explained_variance_ratio_.sum()),
            'svd_projection_max_error': float(np.max(np.abs(svd_projection - sklearn_projection))),
            'discarded_energy_relative_error': float(abs(observed_sse-discarded_sse) / max(observed_sse, 1e-12)),
            'test_mse_by_dimension': errors}


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))