# %% [markdown]
# # Testing python-som with audio dataset

# %% [markdown]
# # Imports

# %%
import matplotlib.pyplot as plt
# import librosa as lr
# import librosa.display as lrdisp
import numpy as np
import pandas as pd
import pickle
import seaborn as sns
import sklearn.preprocessing

from python_som import SOM

FILE_PREFIX = 'som64_u_grupo1'

# %% [markdown]
# # Loading dataset

# %%
df = pd.read_csv('features_means.csv', index_col=0, verbose=True)
df.index = pd.to_datetime(df.index)
df['rac'] = False
df.loc['2020-09-22':, 'rac'] = True  # type: ignore
df.sort_index(inplace=True)

# %% [markdown]
# ## Checking for and dropping duplicates

# %%
# Resetting index for duplicate analysis
df.reset_index(inplace=True)
print("Duplicates by filename:",
      df.duplicated(subset=['file_name']).value_counts(),
      sep='\n')
df.drop_duplicates(subset=['file_name'], inplace=True)
print("Duplicates by (datetime, ala, grupo):",
      df.duplicated(subset=['datetime', 'ala', 'grupo']).value_counts(),
      sep='\n')
df.drop_duplicates(subset=['datetime', 'ala', 'grupo'], inplace=True)
# Rebuilding dataframe index
df.set_index('datetime', inplace=True)

# %%
# Filtering dataset by 'group'
df = df[df['grupo'] == 1]

# %%
# Dropping tail of dataset for class balancing
# tail_size = abs(
#     len(df[df['rac'].astype(int) == 1]) - len(df[df['rac'].astype(int) == 0]))
# df.drop(df.tail(tail_size).index, inplace=True)

# %% [markdown]
# ## Visualizing distribution of sample dates

# %%
df_tmp = pd.DataFrame(df['file_name'].resample('1D').count())
df_tmp['count'] = df_tmp['file_name']
del df_tmp['file_name']
df_tmp['rac'] = False
df_tmp.loc['2020-09-22':, 'rac'] = True  # type: ignore

plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.barplot(y=df_tmp.index, x=df_tmp['count'], hue=df_tmp['rac'])
plt.draw()

df_tmp = pd.DataFrame(df['file_name'].resample('1H').count())
df_tmp['count'] = df_tmp['file_name']
del df_tmp['file_name']
df_tmp['rac'] = False
df_tmp.loc['2020-09-22':, 'rac'] = True  # type: ignore
df_tmp = df_tmp.reset_index()
df_tmp['hour'] = df_tmp['datetime'].dt.hour

plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
sns.barplot(y=df_tmp['hour'], x=df_tmp['count'], hue=df_tmp['rac'], orient='h')
plt.draw()

# %%
df_melt = pd.melt(df, value_vars=['rac'], value_name='ractopamine')
plt.figure(figsize=(10, 10))
sns.set(style="whitegrid",
        palette=sns.color_palette("muted", n_colors=6, desat=1.0))
ax = sns.countplot(data=df_melt, x='ractopamine', hue='ractopamine')

for p in ax.patches:
    ax.annotate(f'\n{p.get_height()}', (p.get_x() + 0.2, p.get_height()),
                ha='center',
                va='top',
                color='white',
                size=18)

plt.draw()

# %%
# using sklearn's MinMaxScaler
scaler = sklearn.preprocessing.MinMaxScaler(feature_range=(0, 1))

df_train = df.iloc[:, 3:-1].copy()
df_train = scaler.fit_transform(df_train)

# %%
# Defining first element of SOM shape
# Second element will be assigned based on the ratio between the
# first two principal components of the train dataset
som_x: int = 64
try:
    with open(f'./{FILE_PREFIX}.obj', 'rb') as f:
        som = pickle.load(f)
except FileNotFoundError:
    som = SOM(x=som_x,
              y=None,
              input_len=df_train.shape[1],
              learning_rate=0.5,
              neighborhood_radius=1.0,
              neighborhood_function='gaussian',
              cyclic_x=True,
              cyclic_y=True,
              data=df_train)
    # Training SOM
    som.weight_initialization(mode='linear', data=df_train)
    som.train(data=df_train, mode='random', verbose=True)
    with open(f'./{FILE_PREFIX}.obj', 'wb') as f:
        pickle.dump(som, f)

# %%
som_x, som_y = som.get_shape()
print('SOM shape:', (som_x, som_y))

# %%
# Visualizing distance matrix and activation matrix
umatrix = som.distance_matrix()
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 9))
sns.heatmap(umatrix.T, cmap='bone_r', ax=ax1, robust=True)
sns.heatmap(som.activation_matrix(data=df_train).T,
            cmap='mako',
            ax=ax2,
            robust=True)
ax1.invert_yaxis()
ax2.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_umatrix_activation.png',
            bbox_inches='tight',
            transparent=True)
plt.draw()

# %%
# Visualizing distance matrix anc activation matrix separately
fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(umatrix.T, cmap='bone_r', robust=True)
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_umatrix.png',
            bbox_inches='tight',
            transparent=True)

fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(som.activation_matrix(data=df_train).T,
                 cmap='mako',
                 robust=True)
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_activation_matrix.png',
            bbox_inches='tight',
            transparent=True)

# %% [markdown]
# ## Visualizing distribution of features

# %%
for column in df.iloc[:, 3:-1].columns:
    hmap = som.get_weights()[:, :, df.iloc[:, 3:-1].columns.get_loc(column)].T
    fig = plt.figure(figsize=(16, 9))
    ax = sns.heatmap(hmap, robust=True, cmap='BrBG')
    ax.invert_yaxis()
    fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_{column}.png',
                bbox_inches='tight',
                transparent=True)
    plt.close(fig=fig)

# %% [markdown]
# ## Visualizing distribution of audios by metadata (day, hour, ...)
# Each node is colorized according to its most frequent label

# %%
df['days'] = df.index.date
df['days'] = (df['days'] - df['days'][0])
df['days'] = df['days'].apply(lambda x: x.days)
df['hour'] = df.index.hour

# %%
# Visualizing 'rac' distribution
class_assignments = som.label_map(np.array(df_train), np.array(df['rac']))
hmap = np.zeros((som_x, som_y))
for i, j in sorted(class_assignments.keys()):
    try:
        hmap[i][j] = class_assignments[(i, j)].most_common()[0][0] + 1
    except Exception:
        continue
hmap = hmap.T
fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(hmap,
                 cmap=sns.color_palette(palette=["#000000", "blue", "orange"],
                                        n_colors=3),
                 cbar_kws={'ticks': [0, 1, 2]})
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_rac.png',
            bbox_inches='tight',
            transparent=True)
plt.show()

# %%
# Visualizing by 'grupo'
print(df.groupby('grupo')['rac'].count())
column = 'grupo'
class_assignments = som.label_map(np.array(df_train), np.array(df[column]))
hmap = np.zeros((som_x, som_y))
for i, j in sorted(class_assignments.keys()):
    try:
        hmap[i][j] = class_assignments[(i, j)].most_common()[0][0]
    except Exception:
        hmap[i][j] = 0
hmap = hmap.T
fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(hmap,
                 cmap=sns.color_palette(palette=["#000000", "blue", "orange"],
                                        n_colors=3),
                 cbar_kws={'ticks': [0, 1, 2]})
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_{column}.png',
            bbox_inches='tight',
            transparent=True)
plt.show()

# %%
# Visualizing by 'days'
print(df.groupby('days')['rac'].count())
column = 'days'
class_assignments = som.label_map(np.array(df_train), np.array(df[column]))
hmap = np.zeros((som_x, som_y))
for i, j in sorted(class_assignments.keys()):
    try:
        hmap[i][j] = class_assignments[(i, j)].most_common()[0][0]
    except Exception:
        hmap[i][j] = -1
hmap = hmap.T
fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(hmap, cmap='viridis')
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_{column}.png',
            bbox_inches='tight',
            transparent=True)
plt.show()

# %%
# Visualizing by 'hour'
print(df.groupby('hour')['rac'].count())
column = 'hour'
class_assignments = som.label_map(np.array(df_train), np.array(df[column]))
hmap = np.zeros((som_x, som_y))
for i, j in sorted(class_assignments.keys()):
    try:
        hmap[i][j] = class_assignments[(i, j)].most_common()[0][0]
    except Exception:
        hmap[i][j] = -1
hmap = hmap.T
fig = plt.figure(figsize=(16, 9))
ax = sns.heatmap(hmap,
                 cmap=sns.diverging_palette(150,
                                            250,
                                            s=100,
                                            l=20,
                                            sep=1,
                                            n=26,
                                            center='light'),
                 center=12)
ax.invert_yaxis()
fig.savefig(f'./output_{FILE_PREFIX}/{FILE_PREFIX}_{column}.png',
            bbox_inches='tight',
            transparent=True)
plt.show()

# %%

[((25, 5, 25, 65), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((26, 11, 26, 35), 'pandas.to_datetime', 'pd.to_datetime', ({(26, 26, 26, 34): 'df.index'}, {}), '(df.index)', True, 'import pandas as pd\n'), ((68, 0, 68, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((71, 0, 71, 65), 'seaborn.barplot', 'sns.barplot', (), '', True, 'import seaborn as sns\n'), ((72, 0, 72, 10), 'matplotlib.pyplot.draw', 'plt.draw', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((82, 0, 82, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((85, 0, 85, 79), 'seaborn.barplot', 'sns.barplot', (), '', True, 'import seaborn as sns\n'), ((86, 0, 86, 10), 'matplotlib.pyplot.draw', 'plt.draw', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((89, 10, 89, 67), 'pandas.melt', 'pd.melt', (), '', True, 'import pandas as pd\n'), ((90, 0, 90, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((93, 5, 93, 68), 'seaborn.countplot', 'sns.countplot', (), '', True, 'import seaborn as sns\n'), ((102, 0, 102, 10), 'matplotlib.pyplot.draw', 'plt.draw', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((142, 18, 142, 53), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((143, 0, 143, 58), 'seaborn.heatmap', 'sns.heatmap', (), '', True, 'import seaborn as sns\n'), ((153, 0, 153, 10), 'matplotlib.pyplot.draw', 'plt.draw', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((157, 6, 157, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((158, 5, 158, 55), 'seaborn.heatmap', 'sns.heatmap', (), '', True, 'import seaborn as sns\n'), ((164, 6, 164, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((200, 7, 200, 31), 'numpy.zeros', 'np.zeros', ({(200, 16, 200, 30): '(som_x, som_y)'}, {}), '((som_x, som_y))', True, 'import numpy as np\n'), ((207, 6, 207, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((216, 0, 216, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((223, 7, 223, 31), 'numpy.zeros', 'np.zeros', ({(223, 16, 223, 30): '(som_x, som_y)'}, {}), '((som_x, som_y))', True, 'import numpy as np\n'), ((230, 6, 230, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((239, 0, 239, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((246, 7, 246, 31), 'numpy.zeros', 'np.zeros', ({(246, 16, 246, 30): '(som_x, som_y)'}, {}), '((som_x, som_y))', True, 'import numpy as np\n'), ((253, 6, 253, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((254, 5, 254, 38), 'seaborn.heatmap', 'sns.heatmap', (), '', True, 'import seaborn as sns\n'), ((259, 0, 259, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((266, 7, 266, 31), 'numpy.zeros', 'np.zeros', ({(266, 16, 266, 30): '(som_x, som_y)'}, {}), '((som_x, som_y))', True, 'import numpy as np\n'), ((273, 6, 273, 33), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((287, 0, 287, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((179, 10, 179, 37), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((180, 9, 180, 52), 'seaborn.heatmap', 'sns.heatmap', (), '', True, 'import seaborn as sns\n'), ((185, 4, 185, 22), 'matplotlib.pyplot.close', 'plt.close', (), '', True, 'import matplotlib.pyplot as plt\n'), ((199, 34, 199, 52), 'numpy.array', 'np.array', ({(199, 43, 199, 51): 'df_train'}, {}), '(df_train)', True, 'import numpy as np\n'), ((199, 54, 199, 73), 'numpy.array', 'np.array', ({(199, 63, 199, 72): "df['rac']"}, {}), "(df['rac'])", True, 'import numpy as np\n'), ((222, 34, 222, 52), 'numpy.array', 'np.array', ({(222, 43, 222, 51): 'df_train'}, {}), '(df_train)', True, 'import numpy as np\n'), ((222, 54, 222, 74), 'numpy.array', 'np.array', ({(222, 63, 222, 73): 'df[column]'}, {}), '(df[column])', True, 'import numpy as np\n'), ((245, 34, 245, 52), 'numpy.array', 'np.array', ({(245, 43, 245, 51): 'df_train'}, {}), '(df_train)', True, 'import numpy as np\n'), ((245, 54, 245, 74), 'numpy.array', 'np.array', ({(245, 63, 245, 73): 'df[column]'}, {}), '(df[column])', True, 'import numpy as np\n'), ((265, 34, 265, 52), 'numpy.array', 'np.array', ({(265, 43, 265, 51): 'df_train'}, {}), '(df_train)', True, 'import numpy as np\n'), ((265, 54, 265, 74), 'numpy.array', 'np.array', ({(265, 63, 265, 73): 'df[column]'}, {}), '(df[column])', True, 'import numpy as np\n'), ((70, 16, 70, 65), 'seaborn.color_palette', 'sns.color_palette', (), '', True, 'import seaborn as sns\n'), ((84, 16, 84, 65), 'seaborn.color_palette', 'sns.color_palette', (), '', True, 'import seaborn as sns\n'), ((92, 16, 92, 65), 'seaborn.color_palette', 'sns.color_palette', (), '', True, 'import seaborn as sns\n'), ((118, 14, 118, 28), 'pickle.load', 'pickle.load', ({(118, 26, 118, 27): 'f'}, {}), '(f)', False, 'import pickle\n'), ((120, 10, 128, 28), 'python_som.SOM', 'SOM', (), '', False, 'from python_som import SOM\n'), ((209, 22, 210, 51), 'seaborn.color_palette', 'sns.color_palette', (), '', True, 'import seaborn as sns\n'), ((232, 22, 233, 51), 'seaborn.color_palette', 'sns.color_palette', (), '', True, 'import seaborn as sns\n'), ((275, 22, 281, 59), 'seaborn.diverging_palette', 'sns.diverging_palette', (), '', True, 'import seaborn as sns\n'), ((133, 8, 133, 27), 'pickle.dump', 'pickle.dump', ({(133, 20, 133, 23): 'som', (133, 25, 133, 26): 'f'}, {}), '(som, f)', False, 'import pickle\n')]