import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import TensorDataset, DataLoader

from torchvision import transforms as tfs
import os
from torchvision.datasets import MNIST

data_tfs = tfs.Compose([
    tfs.ToTensor(),
    tfs.Normalize((0.5), (0.5))
])

# install for train and test
root = './'
train_dataset = MNIST(root, train=True,  transform=data_tfs, download=True)
val_dataset  = MNIST(root, train=False, transform=data_tfs, download=True)

train_dataloader =  DataLoader(train_dataset, batch_size=128)
valid_dataloader =  DataLoader(val_dataset, batch_size=128)

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # Свёртка1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding = 2)
        # Свёртка2
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        # Pooling1
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride = 2)
        # Pooling2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride = 2)
        # Линейный слой1
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # Линейный слой2
        self.fc2 = nn.Linear(120, 84)
        # Выходной слой
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Свёртка1 + активация Тангес
        x = F.tanh(self.conv1(x))
        # Pooling1
        x = self.pool1(x)
        # Свёртка1 + активация Тангес
        x = F.tanh(self.conv2(x))
        # Pooling2
        x = self.pool2(x)
        # Flatten
        size_ = int(x.nelement() / x.shape[0])
        x = x.view(-1, size_)
        # Линейный слой1 + активация Тангес
        x = F.tanh(self.fc1(x))
        # Линейный слой2 + активация Тангес
        x = F.tanh(self.fc2(x))
        # Выходной слой
        x = self.fc3(x)
        return x

device = "cuda" if torch.cuda.is_available() else "cpu"

model = LeNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
loaders = {"train": train_dataloader, "valid": valid_dataloader}

max_epochs = 10
accuracy = {"train": [], "valid": []}
for epoch in range(max_epochs):
    for k, dataloader in loaders.items():
        epoch_correct = 0
        epoch_all = 0
        for x_batch, y_batch in dataloader:
            if k == "train":
                model.train()
                optimizer.zero_grad()
                outp = model(x_batch.to(device))
            else:
                model.eval()
                with torch.no_grad():
                    outp = model(x_batch.to(device))
            preds = outp.argmax(-1)
            correct =  (preds == y_batch).sum()
            all = y_batch.shape[0]
            epoch_correct += correct.item()
            epoch_all += all
            if k == "train":
                loss = criterion(outp, y_batch.to(device))
                loss.backward()
                optimizer.step()
        if k == "train":
            print(f"Epoch: {epoch+1}")
        print(f"Loader: {k}. Accuracy: {epoch_correct/epoch_all}")
        accuracy[k].append(epoch_correct/epoch_all)

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import time
import copy

# Проверка GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ---------------------------------------------------------
# TODO: Реализуйте архитектуру сети согласно заданию
# ---------------------------------------------------------
class CustomCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(CustomCNN, self).__init__()

        # Свёрточный блок 1
        self.conv1 = nn.Conv2d(____, ____, kernel_size=3, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Свёрточный блок 2
        self.conv2 = nn.Conv2d(____, ____, kernel_size=3, padding=0)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Полносвязные слои
        self.fc1 = nn.Linear(____, 256)  # TODO: вычислите входную размерность
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.out = nn.Linear(64, num_classes)

        self.relu = nn.ReLU()

    def forward(self, x):
        # TODO: Реализуйте forward pass
        # Conv1 -> ReLU -> Pool1
        # Conv2 -> ReLU -> Pool2
        # Flatten
        # FC1 -> ReLU -> FC2 -> ReLU -> FC3 -> ReLU -> Output
        return x

# Проверка архитектуры
model = CustomCNN(num_classes=10)
print(model)

# Тестовый прогон
test_input = torch.randn(1, 3, 32, 32)
test_output = model(test_input)
print(f"Output shape: {test_output.shape}")  # Должно быть: torch.Size([1, 10])и

# ---------------------------------------------------------
# Подготовка данных
# ---------------------------------------------------------
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

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

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

# ---------------------------------------------------------
# Функция обучения
# ---------------------------------------------------------
def train_model(model, train_loader, test_loader, num_epochs=15, lr=0.001):
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

    history = {'train_loss': [], 'test_acc': [], 'train_time': []}

    for epoch in range(num_epochs):
        start_time = time.time()
        model.train()
        running_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        scheduler.step()

        # Оценка на тесте
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        epoch_time = time.time() - start_time
        test_acc = 100 * correct / total

        history['train_loss'].append(running_loss / len(train_loader))
        history['test_acc'].append(test_acc)
        history['train_time'].append(epoch_time)

        print(f'Epoch [{epoch+1}/{num_epochs}] | Loss: {running_loss/len(train_loader):.4f} | '
              f'Acc: {test_acc:.2f}% | Time: {epoch_time:.2f}s')

    return model, history

# ---------------------------------------------------------
# TODO: Обучите модель с нуля
# ---------------------------------------------------------
print("=" * 50)
print("ЭКСПЕРИМЕНТ 1: Обучение с нуля на CIFAR-10")
print("=" * 50)

model_scratch = CustomCNN(num_classes=10)
model_scratch, history_scratch = train_model(model_scratch, train_loader, test_loader,
                                              num_epochs=15, lr=0.001)

# Сохранение весов
torch.save(model_scratch.state_dict(), 'cifar10_scratch.pth')
print("Веса модели сохранены: cifar10_scratch.pth")

# ---------------------------------------------------------
# TODO: Постройте графики обучения
# ---------------------------------------------------------
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# График потерь
axes[0].plot(history_scratch['train_loss'], 'b-', label='Train Loss')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')
axes[0].set_title('Training Loss (Scratch)')
axes[0].legend()
axes[0].grid(True)

# График точности
axes[1].plot(history_scratch['test_acc'], 'g-', label='Test Accuracy')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Accuracy (%)')
axes[1].set_title('Test Accuracy (Scratch)')
axes[1].legend()
axes[1].grid(True)

plt.tight_layout()
plt.show()

train_dataset_100 = torchvision.datasets.CIFAR100(root='./data', train=True,
                                                  download=True, transform=transform_train)
test_dataset_100 = torchvision.datasets.CIFAR100(root='./data', train=False,
                                                 download=True, transform=transform_test)

train_loader_100 = DataLoader(train_dataset_100, batch_size=64, shuffle=True, num_workers=2)
test_loader_100 = DataLoader(test_dataset_100, batch_size=64, shuffle=False, num_workers=2)

train_dataset_100 = torchvision.datasets.CIFAR100(root='./data', train=True,
                                                  download=True, transform=transform_train)
test_dataset_100 = torchvision.datasets.CIFAR100(root='./data', train=False,
                                                 download=True, transform=transform_test)

train_loader_100 = DataLoader(train_dataset_100, batch_size=64, shuffle=True, num_workers=2)
test_loader_100 = DataLoader(test_dataset_100, batch_size=64, shuffle=False, num_workers=2)

# ---------------------------------------------------------
# TODO: Обучите модель на CIFAR-100 (Source Task)
# ---------------------------------------------------------
print("=" * 50)
print("ЭКСПЕРИМЕНТ 2A: Предобучение на CIFAR-100")
print("=" * 50)

model_pretrained = CustomCNN(num_classes=100)  # 100 классов для CIFAR-100
model_pretrained, history_pretrained = train_model(model_pretrained,
                                                    train_loader_100,
                                                    test_loader_100,
                                                    num_epochs=15, lr=0.001)

# Сохранение предобученных весов
torch.save(model_pretrained.state_dict(), 'cifar100_pretrained.pth')
print("Предобученные веса сохранены: cifar100_pretrained.pth")

# ---------------------------------------------------------
# Функция загрузки предобученных весов
# ---------------------------------------------------------
def load_pretrained_for_transfer(model, pretrained_path, freeze_conv=True):
    """
    Загружает веса предобученной модели, исключая последний слой,
    и опционально замораживает сверточные слои.
    """
    pretrained_dict = torch.load(pretrained_path)
    model_dict = model.state_dict()

    # Исключаем веса выходного слоя (разная размерность: 100 vs 10)
    pretrained_dict = {k: v for k, v in pretrained_dict.items()
                       if k not in ['out.weight', 'out.bias']}

    # Обновляем словарь весов
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)

    # Замораживаем сверточные слои если нужно
    if freeze_conv:
        for param in model.conv1.parameters():
            param.requires_grad = False
        for param in model.conv2.parameters():
            param.requires_grad = False
        print("✅ Сверточные слои заморожены (Feature Extraction)")
    else:
        print("✅ Все слои разморожены (Fine-Tuning)")

    return model

# ---------------------------------------------------------
# TODO: Реализуйте Transfer Learning
# ---------------------------------------------------------
print("=" * 50)
print("ЭКСПЕРИМЕНТ 2B: Transfer Learning на CIFAR-10")
print("=" * 50)

# Создаём модель для CIFAR-10
model_tl = CustomCNN(num_classes=10)

# Загружаем предобученные веса
model_tl = load_pretrained_for_transfer(model_tl, 'cifar100_pretrained.pth', freeze_conv=True)

# Оптимизатор только для обучаемых параметров
optimizer_tl = optim.Adam(filter(lambda p: p.requires_grad, model_tl.parameters()), lr=0.001)

# Модифицированная функция обучения для TL
def train_model_tl(model, train_loader, test_loader, optimizer, num_epochs=10):
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()

    history = {'train_loss': [], 'test_acc': []}

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        # Оценка
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        test_acc = 100 * correct / total
        history['train_loss'].append(running_loss / len(train_loader))
        history['test_acc'].append(test_acc)

        print(f'Epoch [{epoch+1}/{num_epochs}] | Loss: {running_loss/len(train_loader):.4f} | Acc: {test_acc:.2f}%')

    return model, history

# Обучение с Transfer Learning
model_tl, history_tl = train_model_tl(model_tl, train_loader, test_loader,
                                       optimizer_tl, num_epochs=10)

torch.save(model_tl.state_dict(), 'cifar10_transfer.pth')

# Стратегия 1: Feature Extraction (свертки заморожены)
model_fe = CustomCNN(num_classes=10)
model_fe = load_pretrained_for_transfer(model_fe, 'cifar100_pretrained.pth', freeze_conv=True)
optimizer_fe = optim.Adam(filter(lambda p: p.requires_grad, model_fe.parameters()), lr=0.001)
model_fe, history_fe = train_model_tl(model_fe, train_loader, test_loader, optimizer_fe, num_epochs=10)

# Стратегия 2: Fine-Tuning (все слои обучаются)
model_ft = CustomCNN(num_classes=10)
model_ft = load_pretrained_for_transfer(model_ft, 'cifar100_pretrained.pth', freeze_conv=False)
optimizer_ft = optim.Adam(model_ft.parameters(), lr=0.0001)  # Меньший LR для Fine-Tuning
model_ft, history_ft = train_model_tl(model_ft, train_loader, test_loader, optimizer_ft, num_epochs=10)

# ---------------------------------------------------------
# TODO: Постройте сравнительные графики
# ---------------------------------------------------------
fig, axes = plt.subplots(1, 2, figsize=(15, 5))

# Сравнение точности
axes[0].plot(history_scratch['test_acc'], 'r-', label='Scratch', linewidth=2)
axes[0].plot(history_tl['test_acc'], 'b-', label='Transfer (Frozen)', linewidth=2)
axes[0].plot(history_fe['test_acc'], 'g--', label='Feature Extraction', linewidth=2)
axes[0].plot(history_ft['test_acc'], 'm--', label='Fine-Tuning', linewidth=2)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Test Accuracy (%)')
axes[0].set_title('Сравнение стратегий обучения')
axes[0].legend()
axes[0].grid(True)

# Сравнение потерь
axes[1].plot(history_scratch['train_loss'], 'r-', label='Scratch', linewidth=2)
axes[1].plot(history_tl['train_loss'], 'b-', label='Transfer (Frozen)', linewidth=2)
axes[1].plot(history_fe['train_loss'], 'g--', label='Feature Extraction', linewidth=2)
axes[1].plot(history_ft['train_loss'], 'm--', label='Fine-Tuning', linewidth=2)
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Train Loss')
axes[1].set_title('Сравнение функций потерь')
axes[1].legend()
axes[1].grid(True)

plt.tight_layout()
plt.show()