%matplotlib notebook
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import pandas as pd

from matplotlib.ticker import MaxNLocator

"""
Here we give a file, with each line being a dictionary of the parameters

NUM_BASIS_NOUN, NUM_BASIS_VERB, BASIS_NOUN_DIST_CUTOFF, BASIS_VERB_DIST_CUTOFF, VERB_NOUN_DIST_CUTOFF

where each parameter above is a key with associated value of the variable run with.

Additionally, the result values 
sentences, patterns

hold the number of observed values for both quantities from the run.
"""
dict_file = "path_to_file.out"

df = pd.DataFrame({})
with open(dict_file) as file:
    for l in file:
        dict_line = eval(l)
        dict_line = {k:[v] for k,v in dict_line.items()}
        df_tmp = pd.DataFrame(dict_line)
        df = df.append(df_tmp)

df = df.reset_index(drop=True)
df

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "sentences"]
p3d = ax.scatter3D(
    df[plot_order[0]],
    df[plot_order[1]],
    df[plot_order[2]],
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "BASIS_NOUN_DIST_CUTOFF", "sentences"]
p3d = ax.scatter3D(
    df[plot_order[0]],
    df[plot_order[1]],
    df[plot_order[2]],
    s=df["BASIS_VERB_DIST_CUTOFF"]*10,
    c=df[plot_order[3]],
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])
cb = plt.colorbar(p3d)
cb.set_label(plot_order[-1], rotation=90)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "BASIS_NOUN_DIST_CUTOFF", "VERB_NOUN_DIST_CUTOFF", "patterns"]
p3d = ax.scatter3D(
    df[plot_order[0]],
    df[plot_order[1]],
    df[plot_order[2]],
    s=df[plot_order[3]]*10,
    c=df[plot_order[4]],
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])
cb = plt.colorbar(p3d)
cb.set_label(plot_order[-1], rotation=90)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "patterns"]
p3d = ax.scatter3D(
    df[plot_order[0]],
    df[plot_order[1]],
    df[plot_order[2]],
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])

max_patterns = 30
df_sub = df.where(df["patterns"] < max_patterns).dropna()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "BASIS_NOUN_DIST_CUTOFF", "patterns", "BASIS_VERB_DIST_CUTOFF"]
p3d = ax.scatter3D(
    df_sub[plot_order[0]],
    df_sub[plot_order[1]],
    df_sub[plot_order[2]],
    c = df_sub[plot_order[3]],    
    s = df_sub[plot_order[4]]*10,
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])
cb = plt.colorbar(p3d)
cb.set_label(plot_order[3], rotation=90)
ax.plot(
    df_sub[plot_order[0]], 
    df_sub[plot_order[1]], 
    df_sub[plot_order[2]],
    'k--',
    alpha=0.8,
    linewidth=0.5
)

max_patterns = 30
min_patterns = 25

#Filter within range, and ensure patterns is less than sentences
df_sub = df.where(df["patterns"] < max_patterns)
df_sub = df_sub.where(df["patterns"] > min_patterns)
df_sub = df_sub.where(df["patterns"] < df["sentences"])
df_sub = df_sub.dropna()

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

#Use only integer values
ax.xaxis.set_major_locator(MaxNLocator(integer=True))
ax.yaxis.set_major_locator(MaxNLocator(integer=True))
ax.zaxis.set_major_locator(MaxNLocator(integer=True))


ax.set_proj_type('ortho')
plot_order = ["NUM_BASIS_NOUN", "NUM_BASIS_VERB", "BASIS_NOUN_DIST_CUTOFF", "patterns", "BASIS_VERB_DIST_CUTOFF"]
p3d = ax.scatter3D(
    df_sub[plot_order[0]],
    df_sub[plot_order[1]],
    df_sub[plot_order[2]],
    c = df_sub[plot_order[3]],    
    s = df_sub[plot_order[4]]*10,
    cmap="RdBu_r",
    edgecolors='k'
)
ax.set_xlabel(plot_order[0])
ax.set_ylabel(plot_order[1])
ax.set_zlabel(plot_order[2])

cb = plt.colorbar(p3d)
cb.set_label(plot_order[3], rotation=90)
cb.set_ticks(range(min_patterns, max_patterns+1))

ax.plot(
    df_sub[plot_order[0]], 
    df_sub[plot_order[1]], 
    df_sub[plot_order[2]],
    'k--',
    alpha=0.8,
    linewidth=0.5
)

df_sub

df_sub.where(df_sub["NUM_BASIS_NOUN"] == 10).dropna()

df.iloc[2655]

scr="""
#!/bin/bash
#SBATCH -J qubits_full
#SBATCH -N 1

#SBATCH -p GpuQ
#SBATCH -t 18:00:00
#SBATCH -A "ichec001"
#SBATCH --mail-user=lee.oriordan@ichec.ie
#SBATCH --mail-type=ALL

#no extra settings
cd /ichec/work/ichec001/loriordan_scratch/intel-qnlp-iqs2/build
module load intel/2019u5 gcc cmake3
source ../load_env.sh

for nbn in $(seq 2 2 10);do
for nbv in $(seq 2 2 10);do
for bndc in $(seq 1 1 5);do
for bvdc in $(seq 1 1 5);do
for vndc in $(seq 1 1 5);do

echo "NUM_BASIS_NOUN=${nbn} NUM_BASIS_VERB=${nbv} BASIS_NOUN_DIST_CUTOFF=${bndc} BASIS_VERB_DIST_CUTOFF=${bvdc} VERB_NOUN_DIST_CUTOFF=${vndc}"
NUM_BASIS_NOUN=${nbn} NUM_BASIS_VERB=${nbv} BASIS_NOUN_DIST_CUTOFF=${bndc} BASIS_VERB_DIST_CUTOFF=${bvdc} VERB_NOUN_DIST_CUTOFF=${vndc} srun --ntasks 32 -c 1 --cpu-bind=cores -m plane=32 python ../modules/py/scripts/QNLP_EndToEnd_MPI.py
echo ""

done
done
done
done
done
"""

note="""
The ../modules/py/scripts/QNLP_EndToEnd_MPI.py file was modified to exit upon calculating the 
vec_to_encode (i.e. the patterns) list, and prints the following:

    d ={"sentences" : len(sentences),
        "patterns" : len(vec_to_encode),
        "NUM_BASIS_NOUN" : NUM_BASIS_NOUN,
        "NUM_BASIS_VERB" : NUM_BASIS_VERB,
        "BASIS_NOUN_DIST_CUTOFF" : BASIS_NOUN_DIST_CUTOFF,
        "BASIS_VERB_DIST_CUTOFF" : BASIS_VERB_DIST_CUTOFF,
        "VERB_NOUN_DIST_CUTOFF" : VERB_NOUN_DIST_CUTOFF
    }
    print(d)

"""