import gzip
import numpy as np
import torch
import matplotlib.pyplot as plt
from torch import nn

from feedforward_np_version import NumpyMLPClassifier
from feedforward_pytorch_version import TorchMLPClassifier

np.random.seed(2023)


class MNISTDataset:
    def __init__(self, data_prefix, train=True):
        images, labels = self._load_data(data_prefix, train)
        self.images = images.astype(np.float32) / 255.0
        self.labels = labels

    def __getitem__(self, index):
        image = self.images[index][None, :, :]
        if self.labels is None:
            return image
        return image, int(self.labels[index])

    def __len__(self):
        return len(self.images)

    def _load_data(self, data_prefix, train):
        labels = None
        if train:
            with gzip.open(data_prefix + '-labels.gz', 'rb') as f:
                labels = np.frombuffer(f.read(), np.uint8, offset=8)

        with gzip.open(data_prefix + '-images.gz', 'rb') as f:
            images = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1, 28, 28)
        return images, labels


def preview_mnist_samples():
    train_set = MNISTDataset('./data/train', train=True)
    plt.figure()
    for idx in range(6):
        plt.subplot(2, 3, idx + 1)
        plt.tight_layout()
        plt.imshow(train_set[idx][0][0], cmap='gray', interpolation='none')
        plt.title(f'Ground Truth: {train_set[idx][1]}')
        plt.xticks([])
        plt.yticks([])
    plt.show()


def load_mnist():
    train_set = MNISTDataset('./data/train', train=True)
    test_set = MNISTDataset('./data/t10k', train=False)
    return train_set, test_set


def collect_numpy_features(dataset):
    return np.stack([dataset[i][0] for i in range(len(dataset))], axis=0)


def build_one_hot_targets(dataset, num_classes=10):
    class_ids = np.array([dataset[i][1] for i in range(len(dataset))])
    one_hot = np.zeros((len(dataset), num_classes), dtype=np.float32)
    one_hot[np.arange(len(dataset)), class_ids] = 1.0
    return one_hot


def count_correct_predictions(log_prob, one_hot_labels):
    pred_class = np.argmax(log_prob, axis=1)
    true_class = np.argmax(one_hot_labels, axis=1)
    return np.sum(pred_class == true_class)


def save_predictions(predictions, save_path='data/predict.txt'):
    pred_txt = [str(label) for label in predictions]
    with open(save_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(pred_txt))


def train_numpy_mlp(train_set, test_set):
    _, width, height = train_set[0][0].shape
    model = NumpyMLPClassifier(width * height, 10)

    train_features = collect_numpy_features(train_set)
    train_targets = build_one_hot_targets(train_set, num_classes=10)

    num_epochs = 100
    for epoch in range(1, num_epochs + 1):
        log_prob = model.forward(train_features)
        model.backward(train_targets)
        model.update()

        loss = model.compute_loss(log_prob, train_targets)
        correct = count_correct_predictions(log_prob, train_targets)
        num_samples = train_targets.shape[0]

        if epoch % 10 == 0:
            print('epoch %d, loss %.4f, train acc %.3f' % (epoch, loss / num_samples, correct / num_samples))

    test_features = collect_numpy_features(test_set)
    test_log_prob = model.forward(test_features)
    predictions = np.argmax(test_log_prob, axis=1)
    save_predictions(predictions)


def train_torch_mlp(train_set):
    _, width, height = train_set[0][0].shape
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

    model = TorchMLPClassifier(width * height, 10)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()

    num_epochs = 10
    for epoch in range(1, num_epochs + 1):
        total_loss, total_correct, total_seen = 0.0, 0.0, 0
        for images, labels in train_loader:
            images = images.float()
            labels = labels.long()
            logits = model(images)
            loss = criterion(logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * labels.shape[0]
            total_correct += (logits.argmax(dim=1) == labels).float().sum().item()
            total_seen += labels.shape[0]

        if epoch % 5 == 0:
            print('epoch %d, loss %.4f, train acc %.3f' % (epoch, total_loss / total_seen, total_correct / total_seen))


if __name__ == '__main__':
    train_set, test_set = load_mnist()

    # preview_mnist_samples()

    train_numpy_mlp(train_set, test_set)
    train_torch_mlp(train_set)
