# -*- coding: utf-8 -*-
"""DD-pseudolabels.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1eYKQIIS4i7RPTeAFX5SgIZmG3S2jtnr3

## Double Descent Risk Curve for Neural Networks
The goal of this experiment is to analyze error/loss curve for Resnet18 trained using pseudo-labeling method with CIFAR-10.

- models: resnet18
- dataset: CIFAR-10
- learning algorithm: standard supervised, semi-supervised(pseudo-labeling)
- output : model complexity (number of parameters, epochs) vs. test error/loss
---
- hypothesis: in both supervised and semi-supervised, we should get Double Descent phenomenon, which is defined by 
    - having a U-shaped curve before the interpolation threshold (under-parameterized)
    - peaking at the threshold
    - decreasing again in the over-parameterized regime
"""

# Commented out IPython magic to ensure Python compatibility.
from google.colab import drive
drive.mount('/gdrive')

# %matplotlib inline
import torch
from torch.utils.data import DataLoader, Subset, random_split
from torch import nn
import torch.nn.functional as F

from torchvision import datasets, transforms
from torchvision.models.resnet import ResNet, BasicBlock
from torchvision.datasets import CIFAR10, MNIST
from torchvision.transforms import ToTensor 

import numpy as np
import pickle, time, json
import matplotlib.pyplot as plt
from PIL import Image

torch.manual_seed(0)
np.random.seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

"""Define some functions"""

def alpha_weight(epoch, T1, T2, af):
    """ calculate value of alpha used in loss based on the epoch
  
    params:
        - epoch: your current epoch
        - T1: threshold for training with only labeled data
        - T2: threshold for training with only unlabeled data
        - af: max alpha value
    """

    if epoch < T1:
        return 0.0

    elif epoch > T2:
        return af

    else:
        return ((epoch - T1) / (T2 - T1)) * af


def evaluate(model, data_loader, b):
    """ evaluate the loss and accuracy of the trained network on test data

    returns: 
        - (test_accuracy, test_loss)

    params:
        - model:
        - test_loader:
        - b:
    """
    correct = 0 
    total = 0
    running_loss = 0

    model.eval()
    with torch.no_grad():
        for data, labels in data_loader:
            data = data.to(device)
            labels = labels.to(device)

            output = model(data)
            predicted = torch.max(output, 1)[1]
            correct += (predicted == labels).sum()
            total += data.shape[0]

            loss = loss_fn(output, labels, b)
            running_loss += loss.item()
    
    test_error = 1 - correct/total
    return test_error, running_loss/len(data_loader)


def loss_fn(outputs, labels, b):
    criterion = nn.CrossEntropyLoss()

    return (criterion(outputs, labels) - b).abs() + b


def train_semisuper(epochs, model, optimizer, train_loader, unlabeled_loader, test_loader, b):
    """train model on labeled and unlabeled data using pseudo-labels/self-training

    returns: 
        - metrics:
        
    params:
        - epochs: total epochs for training
        - model:
        - optimizer: 
        - train_loader:
        - unlabeled_loader:
        - test_loader:
        - b:
    """
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    metrics = []
    start_ts = time.time()

    T1 = 20 # no contribution of unlabeled data
    T2 = 150 # less contribution of unlabeled data
    af = 3.1 # Instead of using current epoch we use a "step" variable to calculate alpha_weight. This helps the model converge faster

    # step 1: pre-train teacher model for T1 epochs only on labeled data
    print("START SUPERVISED LEARNING")
    train_super(T1, model, optimizer, train_loader, test_loader, b)
    model.train()

    print("START SEMI-SUPERVISED LEARNING")
    # step 2: train model only on unlabeled data for (T2-T1) epochs and both
    for epoch in range(epochs):
        running_loss = 0
        unlabeled_correct = 0
        unlabeled_total = 0
        labeled_correct = 0
        labeled_total = 0

        # generate pseudo labels and train model on unlabeled data
        for (x_unlabeled, y_unlabeled), (x_labeled, y_labeled) in zip(unlabeled_loader, train_loader):

            # genereate the pseudo labels (changed by every weight update) for unlabeled images
            x_unlabeled, y_unlabeled = x_unlabeled.to(device), y_unlabeled.to(device)
            x_labeled, y_labeled = x_labeled.to(device), y_labeled.to(device)

            output_unlabeled = model(x_unlabeled)
            _, pseudo_labels = torch.max(output_unlabeled, 1) # even if the absolute confidence is low, the max class becomes 1
            
            unlabeled_loss = loss_fn(output_unlabeled, pseudo_labels, b)   
            
            unlabeled_total += x_unlabeled.shape[0]
            unlabeled_correct += (pseudo_labels == y_unlabeled).sum()
            
            output_labeled = model(x_labeled)
            _, predicted = torch.max(output_labeled, 1)

            labeled_loss = loss_fn(output_labeled, y_labeled, b)   

            labeled_total += x_labeled.shape[0]
            labeled_correct += (predicted == y_labeled).sum()

            total_loss = alpha_weight(epoch, T1, T2, af)*unlabeled_loss + labeled_loss

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            running_loss += total_loss.item()
        
        print('pseudo label accuracy={}'.format(unlabeled_correct/unlabeled_total))
        
        train_loss = running_loss/len(train_loader)

        correct = unlabeled_correct + labeled_correct
        total = unlabeled_total + labeled_total
        train_error = 1 - correct/total
        
        test_error, test_loss = evaluate(model, test_loader, b)

        metrics.append([train_error.item(), test_error.item(), train_loss, test_loss])
        print('Epoch: {} : Alpha Weights : {:.3f}, Train Error : {:.3f}, Train Loss : {:.3f} | Test Error : {:.3f}, Test Loss : {:.3f} '.format(epoch+1, alpha_weight(epoch, T1, T2, af), train_error, train_loss, test_error, test_loss))
        model.train()

    print("minutes elapsed: {:.3f}".format((time.time() - start_ts)/60))

    return metrics


def train_super(epochs, model, optimizer, train_loader, test_loader, b):
    """train model in a standard supervised way
    
    returns: 
        - metrics:
        
    params:
        - epochs: total epochs for training
        - model:
        - optimizer: 
        - train_loader:
        - test_loader: 
        - b: flooding level
    """
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    metrics = []
    start_ts = time.time()
    model.train()

    # for each epoch
    for epoch in range(epochs):
        running_loss = 0
        total = 0
        correct = 0

        # for each batch
        for batch_idx, (data, labels) in enumerate(train_loader):
            data, labels = data.to(device), labels.to(device)

            output = model(data)
            _, predicted = torch.max(output, 1) 

            correct += (predicted == labels).sum()
            total += data.shape[0]
            
            loss = loss_fn(output, labels, b)
            running_loss += loss.item()
                       
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        train_loss = running_loss/len(train_loader)
        train_error = 1 - correct/total
        test_error, test_loss = evaluate(model, test_loader, b)
        
        metrics.append([train_error.item(), test_error.item(), train_loss, test_loss])
        print('Epoch: {} : Train Error : {:.3f}, Train Loss : {:.3f} | Test Error : {:.3f}, Test Loss : {:.3f} '.format(epoch+1, train_error, train_loss, test_error, test_loss))
        model.train()

    print("minutes elapsed: {:.3f}".format((time.time() - start_ts)/60))

    return metrics

"""Define your neural net model"""

class PreActBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, **kwargs):
        super(PreActBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False))

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out += shortcut
        return out


class PreActResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, init_channels=64):
        super(PreActResNet, self).__init__()
        self.in_planes = init_channels
        c = init_channels

        self.conv1 = nn.Conv2d(3, c, kernel_size=3, stride=1, padding=1, bias=False)
        self.layer1 = self._make_layer(block, c, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 2*c, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 4*c, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 8*c, num_blocks[3], stride=2)
        self.dropout = nn.Dropout(0.5)
        self.linear = nn.Linear(8 * c * block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        # eg: [2, 1, 1, ..., 1]. Only the first one downsamples.
        strides = [stride] + [1] * (num_blocks-1)
        layers = []
        
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion

        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.linear(out)
        return out


class MNISTResNet(PreActResNet):
    def __init__(self, n_classes, k):
        super(MNISTResNet, self).__init__(PreActBlock, [2, 2, 2, 2], num_classes=n_classes, init_channels=k)


class CIFARResNet(PreActResNet):
    def __init__(self, n_classes, k):
        super(CIFARResNet, self).__init__(PreActBlock, [2, 2, 2, 2], num_classes=n_classes, init_channels=k)

"""#Supervised Experiments

### Experiment 1 - Model-wise DD (Finished!)
"""

basic_setting = {
        'k': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 40, 48, 56, 64],
        'epochs': 300,
        'label_noise': 0.15,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': (10000, 40000),
        'augmentation': True
        }

super_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_model_64.json'
open_file = open(file_name, "w")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
total_samples = len(train.targets)

# assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)
for rand in rands:
    train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

n_labeled, n_unlabeled = basic_setting['n_labeled']
# split training data into labeled and unlabeled
train, val = random_split(train, [n_labeled, n_unlabeled])
print(len(train), len(val))

train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

print(basic_setting)
for k in basic_setting['k']:
    model = CIFARResNet(basic_setting['n_classes'], k) # define model with the number of parameter
    model.to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print("number of model parameters = {} when k={}".format(total_params, k))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(k)] = error_metrics
    
    # semi-supervised training using pseudo-labels
    # error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 2 - Epoch-wise DD vs. n_samples"""

basic_setting = {
        'k': 64,
        'epochs': 300,
        'label_noise': 0.2,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
        'augmentation': True
        }

super_results = {}
semisuper_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/super_epoch_n_samples.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)
n_labeled, n_unlabeled = basic_setting['n_labeled']

for ratio in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)

    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)
    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()
    
    # split training data into labeled and unlabeled
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(ratio)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 3 - Epoch-wise DD vs. flooding (Finished!)"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'label_noise': 0.2,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': [0.1, 0.15, 0.2],
        'n_labeled': (50000, 0),
        'augmentation': True
        }

super_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_epoch_flooding.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)

train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)
   
    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)

    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

    # split training data into labeled and unlabeled
    n_labeled, n_unlabeled = basic_setting['n_labeled']
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

for flood in basic_setting['b']:
    
    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, flood)
    super_results[str(flood)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 4 - Epoch-wise DD vs. label noise (Finished!)"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'noise': [0.1, 0.15, 0.2],
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.1,
        'n_labeled': (20000, 30000),
        'augmentation': True
        }

super_results = {}
semisuper_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_epoch_label_noise.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)
n_labeled, n_unlabeled = basic_setting['n_labeled']

for noise in basic_setting['noise']:
    # train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_cifar)
    # test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_test)
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)
   
    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(noise*total_samples), replace=False)

    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

    # split training data into labeled and unlabeled
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(noise)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""# Semi-Supervised Experiments

### Experiment 1 - Epoch-wise DD vs. labeled ratio
"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.1,
        'n_labeled': [(20000, 30000),(10000, 40000)],
        'augmentation': True
        }
    
# We observe all forms of double descent most strongly in settings with label noise in the train set (as is often the case when collecting train data in the real-world).
# lr = (n_unlabeled)**(-0.5) # SGD learning rate

print(basic_setting)
semisuper_results = {}
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Semi-supervised Experiments/semisuper_epoch_ratio.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
# train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_mnist)
# test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist)

for n_labeled, n_unlabeled in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=int(len(train)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    unlabeled_loader = DataLoader(val, batch_size=int(len(val)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    # total_params = sum(p.numel() for p in model.parameters())
    # print("number of model parameters = {} when k={}".format(total_params, basic_setting['k']))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
 
    # semi-supervised training using pseudo-labels
    error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])

semisuper_results = {}
semisuper_results[str(n_labeled)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(semisuper_results, f)

"""### Experiment 2 - Model-wise DD"""

basic_setting = {
        'k': 64,
        'epochs': 300,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': [0.6, 0.4, 0.2],
        'augmentation': True
        }
# We observe all forms of double descent most strongly in settings with label noise in the train set (as is often the case when collecting train data in the real-world).
# lr = (n_unlabeled)**(-0.5) # SGD learning rate

print(basic_setting)
semisuper_results = {}
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/semisuper_epoch.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
# train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_mnist)
# test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist)

for ratio in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    
    # split training data into labeled and unlabeled
    n_labeled = int(total_samples*ratio)
    n_unlabeled = int(total_samples*(1-ratio))
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=int(len(train)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    unlabeled_loader = DataLoader(val, batch_size=int(len(val)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    # total_params = sum(p.numel() for p in model.parameters())
    # print("number of model parameters = {} when k={}".format(total_params, basic_setting['k']))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)
 
    # semi-supervised training using pseudo-labels
    error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])
    semi_super_results[str(ratio)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(semi_super_results, f)

"""# Plotting"""

# matplotlib.rcParams.update({'font.size': 25})

def plot_modelwise(results, fname):
    titles = ['error', 'loss']
    colors = ['blue', 'lime']
    labels = ['train', 'test']

    k = [] # fixed x-axis
    metrics =. []
    exp_type = fname.split('/')[-2]
    exp_name = fname.split('/')[-1].split('.')[0]

    # extract model size and metrics into separate lists
    for model_size, metric in results.items():    
        k.append(int(model_size))
        metrics.append(metric[-1])

    metrics = np.array(metrics)

    # test/train error and loss
    for i, title in enumerate(titles):
        fig, axes = plt.subplots(figsize=(8, 6), dpi=300)
        axes.grid()

        axes.plot(k, metrics[:, 2*i], label=labels[0], color=colors[0]) # train
        axes.plot(k, metrics[:, 2*i+1], label=labels[1], color=colors[1]) # test           

        axes.set_xlabel('resnet width=k')
        axes.set_ylabel(title)
        axes.set_title("model width vs. {}".format(title))
        axes.legend(loc='upper right', prop={'size': 15})
        fig.savefig('/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/{}/{}_{}.png'.format(exp_type, exp_name, title), dpi=300)


def plot_epochwise(results, fname):
    titles = ['error', 'loss']
    colors = ['red', 'lime', 'blue']
    labels = ['train', 'test']

    exp_type = fname.split('/')[-2]
    exp_name = fname.split('/')[-1].split('.')[0]

    for key in results.keys():
        epochs = np.arange(1, len(results[key])+1) # fixed x-axis
    
    for i, title in enumerate(titles):
        fig, axes = plt.subplots(figsize=(8, 6))

        for j, (param, metrics) in enumerate(results.items()):
            metrics = np.array(results[param])
            axes.plot(epochs, metrics[:, ], label=param, color=colors[0])
            axes.plot(epochs, metrics[:, ], label=param, color=colors[1])

        axes.set_xlabel('Epochs')
        axes.set_ylabel(title)
        axes.set_title('Epochs vs. {}'.format(title))
        axes.legend(loc='upper right', prop={'size': 15})
        fig.savefig('/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/{}/{}_{}.png'.format(exp_type, exp_name, title), dpi=300)

file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Semi-supervised Experiments/semisuper_epoch_ratio.json'

with open(file_name) as json_file:
    data = json.load(json_file)
    print(data)
    #plot_modelwise(data, file_name)
    plot_epochwise(data, file_name)