import torch
import torch.optim as optim
import matplotlib.pyplot as plt

# Створимо dummy модель та optimizer для демонстрації
model = torch.nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

# StepLR scheduler
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# Симуляція 100 епох
epochs = 100
lr_history_step = []

for epoch in range(epochs):
    lr_history_step.append(optimizer.param_groups[0]['lr'])
    
    # Імітація training step (не важливо для демонстрації scheduler)
    optimizer.step()
    
    # Оновлення LR
    scheduler.step()

# Візуалізація
plt.figure(figsize=(12, 5))
plt.plot(lr_history_step, linewidth=2.5, color='#3b82f6', label='StepLR')
plt.xlabel('Epoch', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('StepLR Scheduler (step_size=30, gamma=0.1)', fontsize=14, fontweight='bold')
plt.yscale('log')  # Логарифмічна шкала для кращої візуалізації
plt.grid(alpha=0.3)
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

print("📊 StepLR:")
print(f"├─ Епохи 0-29:   LR = {lr_history_step[0]:.4f}")
print(f"├─ Епохи 30-59:  LR = {lr_history_step[30]:.4f}")
print(f"├─ Епохи 60-89:  LR = {lr_history_step[60]:.4f}")
print(f"└─ Епохи 90-99:  LR = {lr_history_step[90]:.4f}")

# MultiStepLR scheduler
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.1)

lr_history_multistep = []
for epoch in range(epochs):
    lr_history_multistep.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    scheduler.step()

# Візуалізація
plt.figure(figsize=(12, 5))
plt.plot(lr_history_multistep, linewidth=2.5, color='#8b5cf6', label='MultiStepLR')

# Вертикальні лінії на milestones
for milestone in [30, 60, 90]:
    plt.axvline(x=milestone, color='red', linestyle='--', alpha=0.5, linewidth=1.5)
    plt.text(milestone, 0.15, f'Milestone\n{milestone}', ha='center', fontsize=9, 
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.xlabel('Epoch', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('MultiStepLR Scheduler (milestones=[30, 60, 90], gamma=0.1)', fontsize=14, fontweight='bold')
plt.yscale('log')
plt.grid(alpha=0.3)
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

# ExponentialLR scheduler
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)

lr_history_exp = []
for epoch in range(epochs):
    lr_history_exp.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    scheduler.step()

# Візуалізація
plt.figure(figsize=(12, 5))
plt.plot(lr_history_exp, linewidth=2.5, color='#f59e0b', label='ExponentialLR (gamma=0.95)')
plt.xlabel('Epoch', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('ExponentialLR Scheduler — плавне зменшення', fontsize=14, fontweight='bold')
plt.yscale('log')
plt.grid(alpha=0.3)
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

print(f"📊 ExponentialLR (gamma=0.95):")
print(f"├─ Епоха 0:   LR = {lr_history_exp[0]:.6f}")
print(f"├─ Епоха 25:  LR = {lr_history_exp[25]:.6f}")
print(f"├─ Епоха 50:  LR = {lr_history_exp[50]:.6f}")
print(f"├─ Епоха 75:  LR = {lr_history_exp[75]:.6f}")
print(f"└─ Епоха 99:  LR = {lr_history_exp[99]:.6f}")

# CosineAnnealingLR scheduler
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

lr_history_cosine = []
for epoch in range(epochs):
    lr_history_cosine.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    scheduler.step()

# Візуалізація всіх schedulers разом для порівняння
fig, axes = plt.subplots(1, 2, figsize=(16, 5))

# Графік 1: CosineAnnealingLR
ax = axes[0]
ax.plot(lr_history_cosine, linewidth=2.5, color='#22c55e', label='CosineAnnealingLR')
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('CosineAnnealingLR — плавна косинусоїда', fontsize=14, fontweight='bold')
ax.grid(alpha=0.3)
ax.legend(fontsize=11)

# Графік 2: Порівняння всіх schedulers
ax = axes[1]
ax.plot(lr_history_step, linewidth=2, label='StepLR', alpha=0.8)
ax.plot(lr_history_multistep, linewidth=2, label='MultiStepLR', alpha=0.8)
ax.plot(lr_history_exp, linewidth=2, label='ExponentialLR', alpha=0.8)
ax.plot(lr_history_cosine, linewidth=2, label='CosineAnnealingLR', alpha=0.8)

ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('Порівняння всіх schedulers', fontsize=14, fontweight='bold')
ax.set_yscale('log')
ax.grid(alpha=0.3)
ax.legend(fontsize=10)

plt.tight_layout()
plt.show()

# ReduceLROnPlateau scheduler
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',      # Моніторимо мінімізацію loss (для accuracy — mode='max')
    factor=0.5,      # Множити LR на 0.5 (зменшення удвічі)
    patience=10,     # Чекати 10 епох без покращення перед зменшенням
    verbose=False,   # Виводити повідомлення при зміні LR
    min_lr=1e-6      # Мінімальне значення LR
)

# Симуляція val_loss (спочатку зменшується, потім plateau, потім знову зменшується)
np.random.seed(42)
simulated_val_loss = []
lr_history_plateau = []

# Епохи 0-30: швидке зменшення loss
simulated_val_loss.extend(np.linspace(2.3, 0.8, 30) + np.random.randn(30) * 0.05)

# Епохи 30-50: plateau (loss майже не змінюється)
simulated_val_loss.extend(np.ones(20) * 0.8 + np.random.randn(20) * 0.03)

# Епохи 50-70: знову зменшення після зменшення LR
simulated_val_loss.extend(np.linspace(0.8, 0.4, 20) + np.random.randn(20) * 0.02)

# Епохи 70-100: ще один plateau
simulated_val_loss.extend(np.ones(30) * 0.4 + np.random.randn(30) * 0.02)

for epoch in range(100):
    lr_history_plateau.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    
    # Передаємо поточну val_loss у scheduler
    scheduler.step(simulated_val_loss[epoch])

# Візуалізація
fig, axes = plt.subplots(2, 1, figsize=(14, 10), sharex=True)

# Графік 1: Val Loss
ax = axes[0]
ax.plot(simulated_val_loss, linewidth=2, color='#ef4444', label='Validation Loss')
ax.set_ylabel('Validation Loss', fontsize=12)
ax.set_title('ReduceLROnPlateau — адаптивне реагування на plateau', fontsize=14, fontweight='bold')
ax.grid(alpha=0.3)
ax.legend(fontsize=11)

# Графік 2: Learning Rate
ax = axes[1]
ax.plot(lr_history_plateau, linewidth=2.5, color='#8b5cf6', label='Learning Rate')
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_yscale('log')
ax.grid(alpha=0.3)
ax.legend(fontsize=11)

# Позначимо моменти зміни LR
lr_changes = []
for i in range(1, len(lr_history_plateau)):
    if lr_history_plateau[i] != lr_history_plateau[i-1]:
        lr_changes.append(i)
        ax.axvline(x=i, color='red', linestyle='--', alpha=0.5, linewidth=1.5)

plt.tight_layout()
plt.show()

print(f"📊 ReduceLROnPlateau:")
print(f"├─ Початковий LR: {lr_history_plateau[0]:.6f}")
for change_epoch in lr_changes:
    print(f"├─ Епоха {change_epoch}: LR зменшено до {lr_history_plateau[change_epoch]:.6f}")
print(f"└─ Фінальний LR: {lr_history_plateau[-1]:.6f}")

# Реалізація Warm-up + CosineAnnealing вручну
class WarmupCosineScheduler:
    def __init__(self, optimizer, warmup_epochs, total_epochs, eta_min=1e-6, warmup_start_lr=1e-5):
        self.optimizer = optimizer
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.eta_min = eta_min
        self.warmup_start_lr = warmup_start_lr
        self.base_lr = optimizer.param_groups[0]['lr']
        self.current_epoch = 0
    
    def step(self):
        if self.current_epoch < self.warmup_epochs:
            # Лінійний warm-up
            lr = self.warmup_start_lr + (self.base_lr - self.warmup_start_lr) * \
                 (self.current_epoch / self.warmup_epochs)
        else:
            # Cosine annealing після warm-up
            progress = (self.current_epoch - self.warmup_epochs) / \
                      (self.total_epochs - self.warmup_epochs)
            lr = self.eta_min + (self.base_lr - self.eta_min) * \
                 0.5 * (1 + np.cos(np.pi * progress))
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        
        self.current_epoch += 1
        return lr

# Симуляція
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = WarmupCosineScheduler(optimizer, warmup_epochs=5, total_epochs=100)

lr_history_warmup = []
for epoch in range(100):
    lr = scheduler.step()
    lr_history_warmup.append(lr)

# Візуалізація
fig, axes = plt.subplots(1, 2, figsize=(16, 5))

# Графік 1: Повний діапазон
ax = axes[0]
ax.plot(lr_history_warmup, linewidth=2.5, color='#f59e0b', label='Warm-up + Cosine')
ax.axvline(x=5, color='red', linestyle='--', alpha=0.5, linewidth=2, label='Кінець warm-up')
ax.fill_between([0, 5], 0, 0.15, alpha=0.2, color='orange', label='Warm-up фаза')
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('Warm-up (5 епох) + Cosine Annealing', fontsize=14, fontweight='bold')
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

# Графік 2: Детальний перегляд warm-up фази
ax = axes[1]
ax.plot(range(10), lr_history_warmup[:10], 'o-', linewidth=2.5, markersize=8, 
        color='#f59e0b', label='Warm-up детально')
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('Деталі Warm-up фази (епохи 0-9)', fontsize=14, fontweight='bold')
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()

print("📊 Warm-up + Cosine Scheduler:")
print(f"├─ Епоха 0 (старт warm-up): LR = {lr_history_warmup[0]:.8f}")
print(f"├─ Епоха 5 (кінець warm-up): LR = {lr_history_warmup[5]:.8f}")
print(f"├─ Епоха 50 (середина cosine): LR = {lr_history_warmup[50]:.8f}")
print(f"└─ Епоха 99 (фініш): LR = {lr_history_warmup[99]:.8f}")

# PyTorch CosineAnnealingWarmRestarts
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0=10,      # Перший restart через 10 епох
    T_mult=2,    # Кожен наступний період подвоюється
    eta_min=1e-6
)

lr_history_sgdr = []
for epoch in range(150):  # Довше навчання для демонстрації
    lr_history_sgdr.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    scheduler.step()

# Візуалізація
plt.figure(figsize=(14, 6))
plt.plot(lr_history_sgdr, linewidth=2.5, color='#8b5cf6', label='SGDR (T_0=10, T_mult=2)')

# Позначимо restart точки
restart_epochs = [10, 30, 70]  # 10, 10+20, 10+20+40
for restart_epoch in restart_epochs:
    if restart_epoch < len(lr_history_sgdr):
        plt.axvline(x=restart_epoch, color='red', linestyle='--', alpha=0.5, linewidth=2)
        plt.text(restart_epoch, 0.12, f'Restart', ha='center', fontsize=10, rotation=90,
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.xlabel('Epoch', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('SGDR (Warm Restarts) — періодичні стрибки LR', fontsize=14, fontweight='bold')
plt.grid(alpha=0.3)
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

print("📊 SGDR Warm Restarts:")
print(f"├─ Перший період (епохи 0-9): cosine decay")
print(f"├─ Restart на епосі 10: LR → 0.1")
print(f"├─ Другий період (епохи 10-29, довжина 20): cosine decay")
print(f"├─ Restart на епосі 30: LR → 0.1")
print(f"└─ Третій період (епохи 30-69, довжина 40): cosine decay")

# OneCycleLR
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

# Потрібно знати загальну кількість кроків (не епох!)
total_steps = 100 * 391  # 100 епох × 391 батч/епоха (CIFAR-10 з batch_size=128)

scheduler = optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=0.1,           # Максимальний LR у піку циклу
    total_steps=total_steps,
    pct_start=0.3,        # 30% часу — збільшення, 70% — зменшення
    anneal_strategy='cos', # Косинусоїдальна форма
    div_factor=25,        # Початковий LR = max_lr / 25 = 0.004
    final_div_factor=10000 # Фінальний LR = max_lr / 10000 = 0.00001
)

# Симуляція (кожен "epoch" тут — це насправді батч)
lr_history_onecycle = []
for step in range(1000):  # Симулюємо 1000 кроків
    lr_history_onecycle.append(optimizer.param_groups[0]['lr'])
    optimizer.step()
    scheduler.step()

# Візуалізація
plt.figure(figsize=(14, 6))
plt.plot(lr_history_onecycle, linewidth=2.5, color='#22c55e', label='OneCycleLR')
plt.axvline(x=300, color='orange', linestyle='--', alpha=0.5, linewidth=2, label='Пік (30%)')
plt.xlabel('Training Step (батч)', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('OneCycleLR — один великий цикл (super-convergence)', fontsize=14, fontweight='bold')
plt.grid(alpha=0.3)
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

print("📊 OneCycleLR:")
print(f"├─ Старт (крок 0):      LR = {lr_history_onecycle[0]:.6f}")
print(f"├─ Пік (крок 300):      LR = {lr_history_onecycle[300]:.6f}")
print(f"├─ Середина (крок 500): LR = {lr_history_onecycle[500]:.6f}")
print(f"└─ Фініш (крок 999):    LR = {lr_history_onecycle[999]:.6f}")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import time

# Підготовка даних
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"🖥️ Пристрій: {device}\n")

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

# Імпортуємо ResNet-18 (припускаємо, що ми реалізували його у попередній статті)
from torchvision import models
resnet18 = models.resnet18(num_classes=10)
# Адаптація для CIFAR-10: зменшимо початковий conv
resnet18.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
resnet18.maxpool = nn.Identity()  # Вимкнемо maxpool для малих зображень

def train_with_scheduler(model, scheduler_name, scheduler, num_epochs=100):
    """Тренування моделі з заданим scheduler"""
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    
    history = {'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': [], 'lr': []}
    
    print(f"\n🚀 Початок навчання: {scheduler_name}")
    print("=" * 80)
    start_time = time.time()
    
    best_acc = 0
    
    for epoch in range(1, num_epochs + 1):
        # Тренування
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * images.size(0)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        train_loss = running_loss / total
        train_acc = 100. * correct / total
        
        # Валідація
        model.eval()
        running_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                
                running_loss += loss.item() * images.size(0)
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        test_loss = running_loss / total
        test_acc = 100. * correct / total
        
        # Збереження історії
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['test_loss'].append(test_loss)
        history['test_acc'].append(test_acc)
        history['lr'].append(optimizer.param_groups[0]['lr'])
        
        # Оновлення scheduler
        if scheduler_name == "ReduceLROnPlateau":
            scheduler.step(test_loss)
        elif scheduler is not None:
            scheduler.step()
        
        if test_acc > best_acc:
            best_acc = test_acc
        
        # Виведення прогресу
        if epoch % 20 == 0 or epoch == 1:
            print(f"Epoch {epoch:3d} | Train: {train_loss:.4f} / {train_acc:.2f}% | "
                  f"Test: {test_loss:.4f} / {test_acc:.2f}% | LR: {history['lr'][-1]:.6f}")
    
    elapsed = time.time() - start_time
    print("=" * 80)
    print(f"✅ Завершено за {elapsed/60:.1f} хв | Найкраща accuracy: {best_acc:.2f}%\n")
    
    return history

# Експеримент 1: Фіксований LR (baseline)
model_fixed = models.resnet18(num_classes=10)
model_fixed.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model_fixed.maxpool = nn.Identity()

history_fixed = train_with_scheduler(model_fixed, "Fixed LR", scheduler=None, num_epochs=100)

# Експеримент 2: StepLR
model_step = models.resnet18(num_classes=10)
model_step.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model_step.maxpool = nn.Identity()

optimizer_step = optim.SGD(model_step.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler_step = optim.lr_scheduler.StepLR(optimizer_step, step_size=30, gamma=0.1)

history_step = train_with_scheduler(model_step, "StepLR", scheduler_step, num_epochs=100)

# Експеримент 3: CosineAnnealingLR
model_cosine = models.resnet18(num_classes=10)
model_cosine.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model_cosine.maxpool = nn.Identity()

optimizer_cosine = optim.SGD(model_cosine.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer_cosine, T_max=100, eta_min=1e-6)

history_cosine = train_with_scheduler(model_cosine, "CosineAnnealingLR", scheduler_cosine, num_epochs=100)

# Експеримент 4: ReduceLROnPlateau
model_plateau = models.resnet18(num_classes=10)
model_plateau.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model_plateau.maxpool = nn.Identity()

optimizer_plateau = optim.SGD(model_plateau.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler_plateau = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer_plateau, mode='min', factor=0.1, patience=10, verbose=True
)

history_plateau = train_with_scheduler(model_plateau, "ReduceLROnPlateau", scheduler_plateau, num_epochs=100)

# Експеримент 5: Warm-up + CosineAnnealing
model_warmup = models.resnet18(num_classes=10)
model_warmup.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model_warmup.maxpool = nn.Identity()

optimizer_warmup = optim.SGD(model_warmup.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler_warmup = WarmupCosineScheduler(optimizer_warmup, warmup_epochs=5, total_epochs=100)

history_warmup = train_with_scheduler(model_warmup, "Warm-up + Cosine", scheduler_warmup, num_epochs=100)

# Порівняльна візуалізація
fig, axes = plt.subplots(2, 2, figsize=(16, 10))

results = {
    'Fixed LR': history_fixed,
    'StepLR': history_step,
    'CosineAnnealing': history_cosine,
    'ReduceLROnPlateau': history_plateau,
    'Warm-up + Cosine': history_warmup
}

colors = {
    'Fixed LR': '#94a3b8',
    'StepLR': '#3b82f6',
    'CosineAnnealing': '#22c55e',
    'ReduceLROnPlateau': '#f59e0b',
    'Warm-up + Cosine': '#8b5cf6'
}

# Графік 1: Test Accuracy (найважливіший!)
ax = axes[0, 0]
for name, hist in results.items():
    ax.plot(hist['test_acc'], label=name, color=colors[name], linewidth=2, alpha=0.8)
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Test Accuracy (%)', fontsize=12)
ax.set_title('Test Accuracy Comparison', fontsize=14, fontweight='bold')
ax.legend(fontsize=10)
ax.grid(alpha=0.3)

# Графік 2: Test Loss
ax = axes[0, 1]
for name, hist in results.items():
    ax.plot(hist['test_loss'], label=name, color=colors[name], linewidth=2, alpha=0.8)
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Test Loss', fontsize=12)
ax.set_title('Test Loss Comparison', fontsize=14, fontweight='bold')
ax.legend(fontsize=10)
ax.grid(alpha=0.3)

# Графік 3: Learning Rate
ax = axes[1, 0]
for name, hist in results.items():
    ax.plot(hist['lr'], label=name, color=colors[name], linewidth=2, alpha=0.8)
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('Learning Rate Schedule', fontsize=14, fontweight='bold')
ax.set_yscale('log')
ax.legend(fontsize=10)
ax.grid(alpha=0.3)

# Графік 4: Фінальна таблиця результатів
ax = axes[1, 1]
ax.axis('off')

final_results = []
for name, hist in results.items():
    final_acc = hist['test_acc'][-1]
    best_acc = max(hist['test_acc'])
    best_epoch = hist['test_acc'].index(best_acc) + 1
    final_results.append([name, f"{final_acc:.2f}%", f"{best_acc:.2f}%", best_epoch])

table_data = [['Scheduler', 'Final Acc', 'Best Acc', 'Best Epoch']] + final_results
table = ax.table(cellText=table_data, cellLoc='center', loc='center',
                colWidths=[0.35, 0.2, 0.2, 0.2])
table.auto_set_font_size(False)
table.set_fontsize(10)
table.scale(1, 2)

# Колір заголовка
for i in range(4):
    table[(0, i)].set_facecolor('#e2e8f0')
    table[(0, i)].set_text_props(weight='bold')

# Виділимо найкращий результат
best_scheduler = max(results.items(), key=lambda x: max(x[1]['test_acc']))[0]
for i, (name, _, _, _) in enumerate(final_results, start=1):
    if name == best_scheduler:
        for j in range(4):
            table[(i, j)].set_facecolor('#dcfce7')

plt.tight_layout()
plt.savefig('scheduler_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# Детальна таблиця
print("\n" + "="*90)
print("📊 ФІНАЛЬНІ РЕЗУЛЬТАТИ ЕКСПЕРИМЕНТУ")
print("="*90)
print(f"{'Scheduler':<25} {'Final Acc':<12} {'Best Acc':<12} {'Best Epoch':<12} {'Покращення':<15}")
print("-"*90)

baseline_acc = history_fixed['test_acc'][-1]

for name, hist in results.items():
    final_acc = hist['test_acc'][-1]
    best_acc = max(hist['test_acc'])
    best_epoch = hist['test_acc'].index(best_acc) + 1
    improvement = final_acc - baseline_acc
    
    marker = "🏆" if name == best_scheduler else ""
    print(f"{name:<25} {final_acc:<12.2f}% {best_acc:<12.2f}% {best_epoch:<12} +{improvement:.2f}%  {marker}")

print("="*90)