import numpy as np
import torch
import matplotlib.pyplot as plt
from torch import optim, nn

from pytorch.xor.multilayer_perceptron import MultilayerPerceptron
from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations

input_size = 2
output_size = len(set(LABELS))
num_hidden_layers = 0
hidden_size = 2 # isn't ever used but we still set it


seed = 24

torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)

mlp1 = MultilayerPerceptron(input_size=input_size,
                            hidden_size=hidden_size,
                            num_hidden_layers=num_hidden_layers,
                            output_size=output_size)
print(mlp1)
batch_size = 1000

x_data_static, y_truth_static = get_toy_data(batch_size)
fig, ax = plt.subplots(1, 1, figsize=(10,5))
visualize_results(mlp1, x_data_static, y_truth_static,
                  ax=ax, title='Initial Perceptron State', levels=[0.5])

plt.axis('off')
plt.savefig('images/perceptron_initial.png')
plt.show()

losses = []
batch_size = 10000
n_batches = 10
max_epochs = 10

loss_change = 1.0
last_loss = 10.0
change_threshold = 1e-3
epoch = 0
all_imagefiles = []

lr = 0.01
optimizer = optim.Adam(params=mlp1.parameters(), lr=lr)
cross_ent_loss = nn.CrossEntropyLoss()


def early_termination(loss_change, change_threshold, epoch, max_epochs):
    terminate_for_loss_change = loss_change < change_threshold
    terminate_for_epochs = epoch > max_epochs

    # return terminate_for_loss_change or
    return terminate_for_epochs


while not early_termination(loss_change, change_threshold, epoch, max_epochs):
    for _ in range(n_batches):
        # step 0: fetch the data
        x_data, y_target = get_toy_data(batch_size)

        # step 1: zero the gradients
        mlp1.zero_grad()

        # step 2: run the forward pass
        y_pred = mlp1(x_data).squeeze()

        # step 3: compute the loss
        loss = cross_ent_loss(y_pred, y_target.long())

        # step 4: compute the backward pass
        loss.backward()

        # step 5: have the optimizer take an optimization step
        optimizer.step()

        # auxillary: bookkeeping
        loss_value = loss.item()

        losses.append(loss_value)
        loss_change = abs(last_loss - loss_value)
        last_loss = loss_value

    print("epoch: {}: loss_value: {}".format(epoch, loss_value))
    fig, ax = plt.subplots(1, 1, figsize=(10, 5))
    visualize_results(mlp1, x_data_static, y_truth_static, ax=ax, epoch=epoch,
                      title=f"{loss_value:0.2f}; {loss_change:0.4f}")
    plt.axis('off')
    epoch += 1
    all_imagefiles.append(f'images/perceptron_epoch{epoch}_toylearning.png')
    plt.savefig(all_imagefiles[-1])

_, ax = plt.subplots(1,1,figsize=(10,5))
visualize_results(mlp1, x_data_static, y_truth_static, epoch=None, levels=[0.5], ax=ax)
plt.axis('off');
plt.savefig('images/perceptron_final.png')

plot_intermediate_representations(mlp1,
                                  "The Perceptron's Input and Intermediate Representation",
                                  figsize=(9, 3))
plt.savefig("images/perceptron_intermediate.png")
plt.savefig("images/figure_4_5.pdf")
[((17, 0, 17, 23), 'torch.manual_seed', 'torch.manual_seed', ({(17, 18, 17, 22): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((18, 0, 18, 32), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(18, 27, 18, 31): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((19, 0, 19, 20), 'numpy.random.seed', 'np.random.seed', ({(19, 15, 19, 19): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((21, 7, 24, 52), 'pytorch.xor.multilayer_perceptron.MultilayerPerceptron', 'MultilayerPerceptron', (), '', False, 'from pytorch.xor.multilayer_perceptron import MultilayerPerceptron\n'), ((28, 32, 28, 56), 'pytorch.xor.utils.get_toy_data', 'get_toy_data', ({(28, 45, 28, 55): 'batch_size'}, {}), '(batch_size)', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((29, 10, 29, 44), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((30, 0, 31, 72), 'pytorch.xor.utils.visualize_results', 'visualize_results', (), '', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((33, 0, 33, 15), 'matplotlib.pyplot.axis', 'plt.axis', ({(33, 9, 33, 14): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((34, 0, 34, 44), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(34, 12, 34, 43): '"""images/perceptron_initial.png"""'}, {}), "('images/perceptron_initial.png')", True, 'import matplotlib.pyplot as plt\n'), ((35, 0, 35, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((50, 17, 50, 38), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', False, 'from torch import optim, nn\n'), ((97, 8, 97, 40), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((98, 0, 98, 87), 'pytorch.xor.utils.visualize_results', 'visualize_results', (), '', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((99, 0, 99, 15), 'matplotlib.pyplot.axis', 'plt.axis', ({(99, 9, 99, 14): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((100, 0, 100, 42), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(100, 12, 100, 41): '"""images/perceptron_final.png"""'}, {}), "('images/perceptron_final.png')", True, 'import matplotlib.pyplot as plt\n'), ((102, 0, 104, 49), 'pytorch.xor.utils.plot_intermediate_representations', 'plot_intermediate_representations', (), '', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((105, 0, 105, 49), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(105, 12, 105, 48): '"""images/perceptron_intermediate.png"""'}, {}), "('images/perceptron_intermediate.png')", True, 'import matplotlib.pyplot as plt\n'), ((106, 0, 106, 36), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(106, 12, 106, 35): '"""images/figure_4_5.pdf"""'}, {}), "('images/figure_4_5.pdf')", True, 'import matplotlib.pyplot as plt\n'), ((89, 14, 89, 49), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((90, 4, 91, 69), 'pytorch.xor.utils.visualize_results', 'visualize_results', (), '', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((92, 4, 92, 19), 'matplotlib.pyplot.axis', 'plt.axis', ({(92, 13, 92, 18): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((95, 4, 95, 35), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(95, 16, 95, 34): 'all_imagefiles[-1]'}, {}), '(all_imagefiles[-1])', True, 'import matplotlib.pyplot as plt\n'), ((64, 27, 64, 51), 'pytorch.xor.utils.get_toy_data', 'get_toy_data', ({(64, 40, 64, 50): 'batch_size'}, {}), '(batch_size)', False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n')]