#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Example: Using Augmented Lagrangian for Per-Dimension KL Floor Enforcement

This script demonstrates how to use the augmented Lagrangian (AL) approach
to enforce per-dimension KL divergence floors in a VAE, preventing posterior collapse.

Key features:
- AL adds NO penalty when KL >= τ (constraint satisfied)
- AL adds linear + quadratic penalty when KL < τ (constraint violated)
- Each latent dimension is treated separately
- Dual variables (λ) adapt automatically

Usage:
    python augmented_lagrangian_example.py
"""

import numpy as np
import torch
from deeplatent import Corpus, DeepLatent

# =============================================================================
# 1. PREPARE DATA
# =============================================================================
# Load your data using the Corpus class
# train_corpus = Corpus(...)
# test_corpus = Corpus(...)

# For this example, assume you have already loaded your data
print("Note: This example assumes you have train_corpus and test_corpus loaded.")
print("Replace the initialization below with your actual data loading code.\n")

# =============================================================================
# 2. CREATE MODEL WITH AUGMENTED LAGRANGIAN ENABLED AT INITIALIZATION
# =============================================================================
# The cleanest way: enable AL directly in the constructor

model = DeepLatent(
    train_data=train_corpus,
    test_data=test_corpus,
    n_factors=20,  # number of latent dimensions
    ae_type="vae",  # REQUIRED: must be VAE for AL
    vi_type="mean_field",  # or "full_rank" or "iaf"
    latent_factor_prior="gaussian",  # or "logistic_normal" or "dirichlet"
    
    # Encoder configuration
    encoder_args={
        "default_bow": {
            "hidden_dims": [512, 256],
            "activation": "relu",
            "dropout": 0.1
        }
    },
    
    # Decoder configuration
    decoder_args={
        "default_bow": {
            "hidden_dims": [],
            "activation": None,
        }
    },
    
    # Optimizer configuration
    optim_args={
        "main": {"lr": 1e-3, "weight_decay": 0.0},
        "prior": {"lr": 1e-4, "weight_decay": 0.01}
    },
    
    # ELBO configuration
    w_prior=1.0,  # β=1 for standard ELBO
    kl_annealing_start=-1,  # no KL annealing
    kl_annealing_end=-1,
    free_bits=0.0,  # disable free_bits when using AL (they serve similar purposes)
    
    # AUGMENTED LAGRANGIAN CONFIGURATION
    use_augmented_lagrangian=True,  # Enable AL
    tau_floor=0.5,                  # Per-dimension KL floor (nats). Try 0.1-0.5
    rho_al=2.0,                     # Penalty parameter. Start with 0.5-5.0
    ema_momentum_al=0.95,           # Smoothing for batch KL estimates (0.9-0.99)
    
    # Training configuration
    num_steps=10000,
    batch_size=64,
    print_every_n_steps=500,
    patience=2000,
    device="cuda" if torch.cuda.is_available() else "cpu"
)

# Note: Training starts automatically when the model is initialized.
# The AL will be active throughout training.
# Get detailed diagnostics about AL state

print("\n" + "="*80)
print("FINAL AUGMENTED LAGRANGIAN DIAGNOSTICS")
print("="*80)

diag = model.get_augmented_lagrangian_diagnostics()

if diag['enabled']:
    print(f"\nConfiguration:")
    print(f"  τ (per-dim floor): {diag['tau_floor']:.4f} nats")
    print(f"  ρ (penalty):       {diag['rho_al']:.4f}")
    print(f"  Total min KL:      {diag['tau_floor'] * model.n_factors:.2f} nats")
    
    print(f"\nKL Statistics:")
    print(f"  Total KL:  {diag['total_kl']:.4f} nats")
    print(f"  Mean KL:   {diag['mean_kl']:.4f} nats (per dimension)")
    print(f"  Min KL:    {diag['min_kl']:.4f} nats")
    print(f"  Max KL:    {diag['max_kl']:.4f} nats")
    
    print(f"\nConstraint Status:")
    print(f"  Violations:    {diag['num_violated']}/{model.n_factors} dimensions")
    print(f"  Avg violation: {diag['avg_violation']:.4f} nats")
    
    print(f"\nDual Variables:")
    print(f"  Mean λ: {diag['mean_lambda']:.4f}")
    print(f"  Max λ:  {diag['max_lambda']:.4f}")
    
    # Visualize per-dimension KL
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))
    
    # Plot 1: Per-dimension KL with floor
    dims = np.arange(model.n_factors)
    axes[0].bar(dims, diag['ema_kl_per_dim'], alpha=0.7, label='KL per dimension')
    axes[0].axhline(y=diag['tau_floor'], color='r', linestyle='--', 
                    label=f'Floor (τ={diag["tau_floor"]:.2f})')
    axes[0].set_xlabel('Latent Dimension')
    axes[0].set_ylabel('KL Divergence (nats)')
    axes[0].set_title('Per-Dimension KL Divergence')
    axes[0].legend()
    axes[0].grid(alpha=0.3)
    
    # Plot 2: Dual variables
    axes[1].bar(dims, diag['lambda_dual'], alpha=0.7, color='orange', label='Dual variable (λ)')
    axes[1].set_xlabel('Latent Dimension')
    axes[1].set_ylabel('λ Value')
    axes[1].set_title('Dual Variables (Higher = More Enforcement Needed)')
    axes[1].legend()
    axes[1].grid(alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('augmented_lagrangian_diagnostics.png', dpi=150)
    print("\nDiagnostic plots saved to: augmented_lagrangian_diagnostics.png")
    plt.show()

# =============================================================================
# 6. TIPS FOR HYPERPARAMETER TUNING
# =============================================================================
print("\n" + "="*80)
print("HYPERPARAMETER TUNING TIPS")
print("="*80)
print("""
1. τ (tau_floor): Per-dimension KL floor
   - Start small: 0.1-0.5 nats per dimension
   - If you still see collapse: gradually increase to 1.0-2.0
   - Total KL floor = τ × n_factors
   - Too high: model can't satisfy constraint (λ grows indefinitely)

2. ρ (rho_al): Penalty parameter
   - Start moderate: 1.0-5.0
   - Higher ρ: tighter enforcement, faster convergence to constraint
   - Lower ρ: gentler enforcement, may need higher λ
   - If violations persist: increase ρ

3. EMA momentum: Smoothing for batch noise
   - Default: 0.95
   - Higher (0.99): more smoothing, slower adaptation
   - Lower (0.90): less smoothing, faster adaptation

4. When to enable AL:
   - Option A: Enable from start (after model initialization)
   - Option B: Train normally first, then enable AL if collapse detected
   - Option C: Gradually increase τ during training (warm-up)

5. Monitoring:
   - Check diag['num_violated']: should decrease over time
   - Check diag['max_lambda']: should stabilize (not grow indefinitely)
   - If λ keeps growing: τ might be too high or ρ too low

6. Combining with other techniques:
   - Use AL instead of free_bits (they serve similar purposes)
   - Can combine with KL annealing (but usually not needed)
   - Works well with lagging inference (separate encoder updates)
   - Compatible with any vi_type: mean_field, full_rank, or iaf
""")

# =============================================================================
# 7. COMPARISON: Before vs After AL
# =============================================================================
print("\n" + "="*80)
print("INTERPRETING RESULTS")
print("="*80)
print("""
Success indicators:
✓ KL per dimension stays above τ (or close to it)
✓ Few violations (< 20% of dimensions)
✓ Dual variables λ stabilize (don't grow indefinitely)
✓ Latent codes show meaningful variation
✓ Reconstruction quality maintained or improved

Failure indicators:
✗ Many dimensions with KL << τ
✗ Dual variables λ keep growing
✗ Reconstruction quality severely degraded
✗ Model can't learn (loss not decreasing)

Troubleshooting:
- If reconstruction degrades: decrease τ or ρ
- If collapse persists: increase τ or ρ
- If training unstable: decrease ρ, increase EMA momentum
- If AL has no effect: check that ae_type='vae' and β=1
""")

print("\n" + "="*80)
print("EXAMPLE COMPLETE!")
print("="*80)
