from sklearn import datasets
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import pandas as pd
from simulation.network import AttractorNetwork, Langevin, relax, inverse_Langevin
from simulation.utils import fetch_digits_data, preprocess_digits_data, continous_inference_and_learning, get_deterministic_attractors
from joblib import Parallel, delayed
from copy import deepcopy
from tqdm import tqdm

digits = fetch_digits_data()
train_data, _ = preprocess_digits_data(digits)

data = train_data[1:4].copy()
data *= 20

fig, axes = plt.subplots(nrows=10, ncols=10, figsize=(10, 10))
axes = axes.flatten()

for i, ax in enumerate(axes):
    if i < data.shape[0]:
        image = data[i].reshape(8, 8) + np.random.normal(0, 0.1, (8,8)) 
        for j in range(image.shape[0]):
            for k in range(image.shape[1]):
                image[j, k] = Langevin(image[j, k])

        ax.imshow(image, cmap="gray_r", interpolation="nearest", vmin=-np.max(image), vmax=np.max(image))
        ax.set_axis_off()
    else:
        ax.set_visible(False)

plt.show()
sns.histplot(data.flatten())

size = data.shape[1]
num_variables = data.shape[0]
J = np.zeros((size, size))
nw = AttractorNetwork(J, biases = np.zeros(J.shape[0]))

num_iterations = 2000

for i in tqdm(range(num_iterations)):
    di = i % data.shape[0] #np.random.randint(0, data.shape[0])
    activations, weight_change, accuracy, complexity, vfe = continous_inference_and_learning(
        nw=nw, 
        data=data[di], 
        inverse_temperature=1, 
        learning_rate=0.001, 
        num_steps=1)
    
sns.heatmap(nw.get_J(), cmap="coolwarm", center=0)
plt.show()

#for i, pattern_idx in enumerate(pattern):
#    plt.text(i*num_runs, max(error[i*num_runs:i*num_runs+num_runs]), str(pattern_idx), ha='center', va='bottom')

# Symmetrize the weight matrix J
J = nw.get_J()
J_symmetrized = 0.5 * (J + J.T)  # J_ij† = (1/2) * (J_ij + J_ji)

# Display the original and symmetrized weight matrices
fig, axes = plt.subplots(1, 3, figsize=(10, 3))
sns.heatmap(J, cmap="coolwarm", center=0, ax=axes[0])
axes[0].set_title("Original Weight Matrix J")

sns.heatmap(J_symmetrized, cmap="coolwarm", center=0, ax=axes[1])
axes[1].set_title("Symmetric Component J†")


# Calculate the asymmetric component of the weight matrix
J_asymmetric = 0.5 * (J - J.T)  # J_ij‡ = (1/2) * (J_ij - J_ji)

# Display the asymmetric component
sns.heatmap(J_asymmetric, cmap="coolwarm", center=0, ax=axes[2])
axes[2].set_title("Antisymmetric Component")
plt.tight_layout()
plt.show()



# Calculate and print the asymmetry measure
asymmetry = np.linalg.norm(J - J.T) / np.linalg.norm(J)
print(f"Asymmetry measure: {asymmetry:.6f}")

acts, weight_change, accuracy, complexity, vfe = continous_inference_and_learning(
    nw, np.zeros(data.shape[1]), inverse_temperature=1, 
    learning_rate=0, num_steps=100)

fig, axes = plt.subplots(5, 10, figsize=(20, 10))
axes = axes.flatten()
for i, ax in enumerate(axes):
    sns.heatmap(np.array(acts[i+10]).reshape(8,8), cmap="gray_r", ax=ax, cbar=False)
    ax.set_xticks([])
    ax.set_yticks([])
plt.tight_layout()
plt.show()

nw_sym = AttractorNetwork(J_symmetrized, biases = np.zeros(nw.get_J().shape[0]))
get_deterministic_attractors(nw_sym, data, noise_levels=[0], inverse_temperature=1 )