"""
Autoencoders — comprehensive topic content with interactive graphs.
"""

import json
import numpy as np


def _autoencoder_architecture():
    """Visualize the bottleneck architecture."""
    layers = ["Input\n(784)", "Enc. 1\n(256)", "Enc. 2\n(64)", "Latent\n(16)", "Dec. 1\n(64)", "Dec. 2\n(256)", "Output\n(784)"]
    sizes = [784, 256, 64, 16, 64, 256, 784]
    normalized = [s / max(sizes) for s in sizes]
    colors = ["#6C63FF", "#a78bfa", "#a78bfa", "#FF6584", "#a78bfa", "#a78bfa", "#34d399"]

    data = [{"x": layers, "y": normalized, "type": "bar",
             "marker": {"color": colors, "line": {"width": 1.5, "color": "#fff"}},
             "text": [str(s) for s in sizes], "textposition": "outside",
             "textfont": {"color": "#e0e0e0", "size": 13}}]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Autoencoder Bottleneck Architecture", "font": {"size": 18}},
        "yaxis": {"title": "Relative Size", "gridcolor": "rgba(255,255,255,0.08)", "showticklabels": False},
        "margin": {"l": 40, "r": 30, "t": 50, "b": 80},
        "annotations": [
            {"x": 1.5, "y": 1.05, "text": "← Encoder →", "showarrow": False, "font": {"size": 14, "color": "#6C63FF"}},
            {"x": 3, "y": 0.15, "text": "Bottleneck", "showarrow": False, "font": {"size": 12, "color": "#FF6584"}},
            {"x": 4.5, "y": 1.05, "text": "← Decoder →", "showarrow": False, "font": {"size": 14, "color": "#34d399"}},
        ],
    }
    return json.dumps({"data": data, "layout": layout})


def _latent_space_2d():
    """Show a 2D latent space with clusters."""
    np.random.seed(42)
    n = 80
    clusters = [
        (np.random.randn(n, 2) * 0.5 + [-2, 1], "0"),
        (np.random.randn(n, 2) * 0.5 + [0, -1.5], "1"),
        (np.random.randn(n, 2) * 0.5 + [2, 1], "2"),
        (np.random.randn(n, 2) * 0.6 + [-1, -2], "3"),
        (np.random.randn(n, 2) * 0.4 + [1.5, -1], "4"),
    ]
    colors_list = ["#6C63FF", "#FF6584", "#34d399", "#fbbf24", "#a78bfa"]

    data = []
    for i, (pts, label) in enumerate(clusters):
        data.append({
            "x": pts[:, 0].tolist(), "y": pts[:, 1].tolist(),
            "mode": "markers", "type": "scatter", "name": f"Digit {label}",
            "marker": {"size": 6, "color": colors_list[i], "opacity": 0.7},
        })

    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "2D Latent Space (VAE on MNIST)", "font": {"size": 18}},
        "xaxis": {"title": "Latent Dim 1", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Latent Dim 2", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def _reconstruction_loss():
    """Show reconstruction loss during training."""
    np.random.seed(42)
    epochs = list(range(1, 81))
    train_loss = [50 * np.exp(-0.06 * e) + 3 + np.random.normal(0, 0.3) for e in epochs]
    val_loss = [50 * np.exp(-0.05 * e) + 5 + np.random.normal(0, 0.4) for e in epochs]

    data = [
        {"x": epochs, "y": train_loss, "mode": "lines", "type": "scatter", "name": "Training Loss",
         "line": {"color": "#6C63FF", "width": 2.5}},
        {"x": epochs, "y": val_loss, "mode": "lines", "type": "scatter", "name": "Validation Loss",
         "line": {"color": "#FF6584", "width": 2.5}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Reconstruction Loss During Training", "font": {"size": 18}},
        "xaxis": {"title": "Epoch", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "MSE Loss", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def get_content() -> dict:
    return {
        "title": "Autoencoders",
        "icon": "🔮",
        "subtitle": "Learning compressed representations through reconstruction",
        "sections": [
            {"id": "introduction", "heading": "What are Autoencoders?", "type": "text",
             "content": ("An Autoencoder is a neural network trained to <strong>reconstruct its input</strong>. It consists of "
                         "an <strong>encoder</strong> that compresses data into a low-dimensional <strong>latent space</strong>, "
                         "and a <strong>decoder</strong> that reconstructs the original from this compressed representation.\n\n"
                         "The bottleneck forces the network to learn the most important features, making autoencoders useful for "
                         "dimensionality reduction, denoising, anomaly detection, and generative models.")},
            {"id": "architecture", "heading": "Bottleneck Architecture", "type": "graph",
             "description": "Data is compressed from 784 dimensions (28×28 image) down to just 16 in the latent space, then reconstructed. The network must learn an efficient compression.",
             "graph_json": _autoencoder_architecture()},
            {"id": "math", "heading": "The Mathematics", "type": "math",
             "content": [
                 {"label": "Encoder", "formula": "z = f_\\theta(x) = \\sigma(Wx + b)",
                  "explanation": "Maps input x to latent representation z. Can be multiple layers."},
                 {"label": "Decoder", "formula": "\\hat{x} = g_\\phi(z) = \\sigma(W'z + b')",
                  "explanation": "Reconstructs the input from the latent code. Trained to minimize ||x - x̂||²."},
                 {"label": "Reconstruction Loss", "formula": "\\mathcal{L} = \\frac{1}{n} \\sum_{i=1}^{n} \\|x_i - \\hat{x}_i\\|^2",
                  "explanation": "Mean Squared Error between input and reconstruction. The core training objective."},
                 {"label": "VAE Loss (ELBO)", "formula": "\\mathcal{L}_{VAE} = \\mathbb{E}[\\|x - \\hat{x}\\|^2] + D_{KL}(q(z|x) \\| p(z))",
                  "explanation": "VAEs add a KL divergence term that regularizes the latent space to follow a Gaussian prior. Enables generation."},
             ]},
            {"id": "latent", "heading": "Latent Space Visualization", "type": "graph",
             "description": "A VAE trained on MNIST projects digits into 2D. Similar digits cluster together. You can interpolate between clusters to generate smooth transitions between digit styles.",
             "graph_json": _latent_space_2d()},
            {"id": "loss", "heading": "Training Progress", "type": "graph",
             "description": "Reconstruction loss decreases as the autoencoder learns to compress and reconstruct data. The gap between train and val loss indicates generalization quality.",
             "graph_json": _reconstruction_loss()},
            {"id": "variants", "heading": "Autoencoder Variants", "type": "list",
             "entries": [
                 {"title": "Vanilla Autoencoder", "detail": "Basic encoder-decoder. Good for dimensionality reduction. Can't generate new data."},
                 {"title": "Variational Autoencoder (VAE)", "detail": "Learns a probability distribution in latent space. Can generate new samples by sampling from the distribution."},
                 {"title": "Denoising Autoencoder", "detail": "Trained with noisy inputs to reconstruct clean outputs. Learns robust features and removes noise."},
                 {"title": "Sparse Autoencoder", "detail": "Adds sparsity constraint (most neurons inactive). Learns disentangled, interpretable features."},
                 {"title": "Convolutional Autoencoder", "detail": "Uses conv layers. Better for images than fully-connected. Preserves spatial structure."},
             ]},
            {"id": "implementation", "heading": "Python Implementation", "type": "code", "language": "python",
             "content": ("import torch\nimport torch.nn as nn\n\nclass Autoencoder(nn.Module):\n"
                         "    def __init__(self, input_dim=784, latent_dim=16):\n"
                         "        super().__init__()\n"
                         "        self.encoder = nn.Sequential(\n"
                         "            nn.Linear(input_dim, 256),\n"
                         "            nn.ReLU(),\n"
                         "            nn.Linear(256, 64),\n"
                         "            nn.ReLU(),\n"
                         "            nn.Linear(64, latent_dim),\n"
                         "        )\n"
                         "        self.decoder = nn.Sequential(\n"
                         "            nn.Linear(latent_dim, 64),\n"
                         "            nn.ReLU(),\n"
                         "            nn.Linear(64, 256),\n"
                         "            nn.ReLU(),\n"
                         "            nn.Linear(256, input_dim),\n"
                         "            nn.Sigmoid(),  # output in [0, 1]\n"
                         "        )\n\n"
                         "    def forward(self, x):\n"
                         "        z = self.encoder(x)    # compress\n"
                         "        return self.decoder(z)  # reconstruct\n\n"
                         "model = Autoencoder()\ncriterion = nn.MSELoss()\n"
                         "# Training: loss = criterion(model(x), x)\n")},
            {"id": "applications", "heading": "Applications", "type": "list",
             "entries": [
                 {"title": "Anomaly Detection", "detail": "Train on normal data only. Anomalies have high reconstruction error. Used in fraud detection and manufacturing QC."},
                 {"title": "Denoising", "detail": "Learn to remove noise from images, audio, or sensor data."},
                 {"title": "Data Compression", "detail": "Learn task-specific compression that preserves important features."},
                 {"title": "Image Generation (VAE)", "detail": "Sample from latent space to generate new images. Smoother but blurrier than GANs."},
             ]},
        ],
    }
