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")
[((401, 424), 'torch.manual_seed', 'torch.manual_seed', (['seed'], {}), '(seed)\n', (418, 424), False, 'import torch\n'), ((425, 457), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', (['seed'], {}), '(seed)\n', (451, 457), False, 'import torch\n'), ((458, 478), 'numpy.random.seed', 'np.random.seed', (['seed'], {}), '(seed)\n', (472, 478), True, 'import numpy as np\n'), ((487, 621), 'pytorch.xor.multilayer_perceptron.MultilayerPerceptron', 'MultilayerPerceptron', ([], {'input_size': 'input_size', 'hidden_size': 'hidden_size', 'num_hidden_layers': 'num_hidden_layers', 'output_size': 'output_size'}), '(input_size=input_size, hidden_size=hidden_size,\n    num_hidden_layers=num_hidden_layers, output_size=output_size)\n', (507, 621), False, 'from pytorch.xor.multilayer_perceptron import MultilayerPerceptron\n'), ((765, 789), 'pytorch.xor.utils.get_toy_data', 'get_toy_data', (['batch_size'], {}), '(batch_size)\n', (777, 789), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((800, 835), 'matplotlib.pyplot.subplots', 'plt.subplots', (['(1)', '(1)'], {'figsize': '(10, 5)'}), '(1, 1, figsize=(10, 5))\n', (812, 835), True, 'import matplotlib.pyplot as plt\n'), ((835, 949), 'pytorch.xor.utils.visualize_results', 'visualize_results', (['mlp1', 'x_data_static', 'y_truth_static'], {'ax': 'ax', 'title': '"""Initial Perceptron State"""', 'levels': '[0.5]'}), "(mlp1, x_data_static, y_truth_static, ax=ax, title=\n    'Initial Perceptron State', levels=[0.5])\n", (852, 949), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((964, 979), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (972, 979), True, 'import matplotlib.pyplot as plt\n'), ((980, 1024), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""images/perceptron_initial.png"""'], {}), "('images/perceptron_initial.png')\n", (991, 1024), True, 'import matplotlib.pyplot as plt\n'), ((1025, 1035), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (1033, 1035), True, 'import matplotlib.pyplot as plt\n'), ((1273, 1294), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), '()\n', (1292, 1294), False, 'from torch import optim, nn\n'), ((2748, 2783), 'matplotlib.pyplot.subplots', 'plt.subplots', (['(1)', '(1)'], {'figsize': '(10, 5)'}), '(1, 1, figsize=(10, 5))\n', (2760, 2783), True, 'import matplotlib.pyplot as plt\n'), ((2781, 2873), 'pytorch.xor.utils.visualize_results', 'visualize_results', (['mlp1', 'x_data_static', 'y_truth_static'], {'epoch': 'None', 'levels': '[0.5]', 'ax': 'ax'}), '(mlp1, x_data_static, y_truth_static, epoch=None, levels=[\n    0.5], ax=ax)\n', (2798, 2873), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((2869, 2884), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (2877, 2884), True, 'import matplotlib.pyplot as plt\n'), ((2886, 2928), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""images/perceptron_final.png"""'], {}), "('images/perceptron_final.png')\n", (2897, 2928), True, 'import matplotlib.pyplot as plt\n'), ((2930, 3047), 'pytorch.xor.utils.plot_intermediate_representations', 'plot_intermediate_representations', (['mlp1', '"""The Perceptron\'s Input and Intermediate Representation"""'], {'figsize': '(9, 3)'}), '(mlp1,\n    "The Perceptron\'s Input and Intermediate Representation", figsize=(9, 3))\n', (2963, 3047), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((3112, 3161), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""images/perceptron_intermediate.png"""'], {}), "('images/perceptron_intermediate.png')\n", (3123, 3161), True, 'import matplotlib.pyplot as plt\n'), ((3162, 3198), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""images/figure_4_5.pdf"""'], {}), "('images/figure_4_5.pdf')\n", (3173, 3198), True, 'import matplotlib.pyplot as plt\n'), ((2406, 2441), 'matplotlib.pyplot.subplots', 'plt.subplots', (['(1)', '(1)'], {'figsize': '(10, 5)'}), '(1, 1, figsize=(10, 5))\n', (2418, 2441), True, 'import matplotlib.pyplot as plt\n'), ((2446, 2572), 'pytorch.xor.utils.visualize_results', 'visualize_results', (['mlp1', 'x_data_static', 'y_truth_static'], {'ax': 'ax', 'epoch': 'epoch', 'title': 'f"""{loss_value:0.2f}; {loss_change:0.4f}"""'}), "(mlp1, x_data_static, y_truth_static, ax=ax, epoch=epoch,\n    title=f'{loss_value:0.2f}; {loss_change:0.4f}')\n", (2463, 2572), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n'), ((2595, 2610), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (2603, 2610), True, 'import matplotlib.pyplot as plt\n'), ((2707, 2738), 'matplotlib.pyplot.savefig', 'plt.savefig', (['all_imagefiles[-1]'], {}), '(all_imagefiles[-1])\n', (2718, 2738), True, 'import matplotlib.pyplot as plt\n'), ((1726, 1750), 'pytorch.xor.utils.get_toy_data', 'get_toy_data', (['batch_size'], {}), '(batch_size)\n', (1738, 1750), False, 'from pytorch.xor.utils import LABELS, get_toy_data, visualize_results, plot_intermediate_representations\n')]