# ---
# jupyter:
#   jupytext:
#     comment_magics: true
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
# %load_ext autoreload
# %autoreload 2

# %%
import math
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

sns.set_theme(style="darkgrid")
sns.set(rc={"figure.figsize": (16, 9)}, font_scale=1.5)

# %%
from utils import saved_measures, calculate_hat_sigma, calculate_expected_eigenvalue
from glob import iglob
import os

# %%
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument("--activation", type=str, default="Tanh", help="Activation to select")
parser.add_argument("-f", type=str, default="", help="Needed for jupyter notebook")
args = parser.parse_args()

# %%
# Loads the csv files found in result as diverse dataframes
df = {}
for measure in ["input_matrix", "norms"]:
    df[measure] = pd.concat((pd.read_csv(filename) for filename in iglob(f"conditioning/results/{measure}*.csv")))
    df[measure] = df[measure].query("loss_function == 'CE' and activation == @args.activation")


# %%
def save_plot(name, **kwargs):
    if not os.path.exists("images/conditioning/"):
        os.makedirs("images/conditioning/", exist_ok=True)
    nargs = {**args.__dict__, **kwargs}
    sluglist = [f"{key}_{value}" for key, value in nargs.items() if key != "f"]
    slug = "@".join(sluglist)
    plt.savefig(f"images/conditioning/{name}@{slug}.svg")


# %%
def to_integer(x):
    try:
        return int(x)
    except ValueError:
        return np.inf

def to_string(x):
    if x != np.inf:
        return str(int(x))
    else:
        return "latest"


# %%
dfn = df["norms"]
dfm = df["input_matrix"]

dfn["layer"] = dfn["layer"].apply(to_integer)
dfm["layer"] = dfm["layer"].apply(to_integer)

dfa = dfm.copy()
dfa["lowest"] = dfa["min_eigenvalue"]
dfa["highest"] = dfa["max_eigenvalue"]

# %%
dfa = dfa.query("layer > 0")
dfa.loc[:, ["layer"]] = dfa["layer"].apply(to_string)
dfm = dfm.query("layer > 0")
dfm.loc[:, ["layer"]] = dfm["layer"].apply(to_string)
dfn = dfn.query("layer > 0")
dfn.loc[:, ["layer"]] = dfn["layer"].apply(to_string)

# %% [markdown]
# # 1. Well-conditioning on wide networks

# %%
partial = dfa["lowest"].apply(lambda x: max(x, 1e-10))
dfa["conditioning"] = (dfa["highest"] / partial).apply(lambda x: math.log(x))
dfa["examples_width_ratio"] = dfa["width"] / dfa["examples"]
dfa["conditioning_factor"] = (1 / ((1 / dfa["examples"]).apply(math.sqrt) - (1 / dfa["width"]).apply(math.sqrt))).apply(math.log)

# %%
sns.relplot(x="conditioning_factor", y="conditioning", col="layer", col_wrap=2,
             data=dfa.query("renormalize == True and examples < width and layer in ['1', '11', '21', 'latest']"));
save_plot("conditioning_wrt_factor", renormalize=True)

# %%
sns.relplot(x="conditioning_factor", y="conditioning", col="layer", col_wrap=2,
             data=dfa.query("renormalize == False and examples < width and layer in ['1', '11', '21', 'latest']"));
save_plot("conditioning_wrt_factor", renormalize=False)

# %%
sns.lineplot(x="layer", y="conditioning", hue="examples_width_ratio", style="renormalize", ci=None,
            data=dfa.query("epoch == 0 and examples < width"), palette=sns.color_palette("bright", n_colors=7));
save_plot("conditioning_vs_layer")

# %%
sns.lineplot(x="layer", y="conditioning", hue="width", style="renormalize",
            data=dfa.query("epoch == 0 and examples == 500"), palette=sns.color_palette("bright", n_colors=3));
save_plot("conditioning_vs_width", examples=500)

# %%
sns.lineplot(x="layer", y="conditioning", hue="examples", style="renormalize", ci=None,
             data=dfa.query("epoch == 0 and examples < width"),
             palette=sns.color_palette("bright", n_colors=4));
save_plot("conditioning_vs_examples")

# %%
for width in [1000, 2000, 5000]:
    sns.lineplot(x="layer", y="conditioning", hue="examples", style="renormalize",
                 data=dfa.query("epoch == 0 and width == @width"),
                 palette=sns.color_palette("bright", n_colors=4));
    #save_plot("conditioning_examples", width=width)
    plt.show();

# %% [markdown]
# Thus renormalizing helps if the network has too many layers, but if it has less it does not help too much.

# %%
ax = sns.lineplot(x="layer", y="min_ondiagonal", style="width", label="Min on diagonal", data=dfm.query("renormalize == False"));
bx = sns.lineplot(x="layer", y="max_offdiagonal", style="width", label="Max off diagonal", data=dfm.query("renormalize == False"));
save_plot("on_off_diagonal_entries", renormalize=False)

# %%
sns.lineplot(x="layer", y="min_ondiagonal", style="width", data=dfm.query("renormalize == True"), label="Min on diagonal")
sns.lineplot(x="layer", y="max_offdiagonal", style="width", data=dfm.query("renormalize == True"), label="Max off diagonal");
save_plot("on_off_diagonal_entries", renormalize=True)

# %%
sns.lineplot(x="layer", y="max_sumoffdiagonal", style="examples", hue="renormalize", data=dfm);
save_plot("maxsumoffdiagonal_examples")

# %% [markdown]
# ## Comparison with Theoretical well-conditioning
#

# %%
initial_eigs = df["input_matrix"].rename(columns={
    "max_eigenvalue": "highest",
    "min_eigenvalue": "lowest",
}).query("layer == 0")[
    ["examples", "highest", "lowest", "max_offdiagonal", "max_sumoffdiagonal"]
].groupby("examples").mean().reset_index()
initial_eigs


# %%
def calculated_min_layers(layers, start):
    delta = start
    yield delta
    for _ in range(layers - 1):
        delta = calculate_expected_eigenvalue(args.activation, delta=delta)
        yield delta
    
def calculated_max_layers(layers, start):
    amount = start
    yield amount
    for _ in range(layers - 1):
        amount = calculate_hat_sigma(args.activation, amount)
        yield amount
    
dataframe = []
in_features = 3072
layers = dfa["layer"].unique()
for _, element in initial_eigs.iterrows():
    newdf = pd.DataFrame.from_dict({
        **element,
        "layer": layers,
        "predicted_lowest": list(calculated_min_layers(len(layers), start=element["lowest"] / math.sqrt(in_features))),
        "predicted_max_offdiagonal": list(calculated_max_layers(len(layers), start=element["max_offdiagonal"])),
    })
    dataframe.append(newdf)

# %%
newdf = pd.concat(dataframe)[["examples", "predicted_lowest", "predicted_max_offdiagonal", "layer"]]
newdf = newdf.rename(columns={"predicted_lowest": "lowest", "predicted_max_offdiagonal": "max_offdiagonal"})
newdf["width"] = np.inf
newdf["renormalize"] = True
newdf["highest"] = (1.0 + (newdf["examples"] - 1) * newdf["max_offdiagonal"]).apply(math.sqrt)
ndfa = dfa.copy()
newdf["conditioning"] = (newdf["highest"] / newdf["lowest"]).apply(math.log)
ndfa["lowest"] /= ndfa["width"].apply(math.sqrt) - ndfa["examples"].apply(math.sqrt)
ndfa["highest"] /= ndfa["width"].apply(math.sqrt)
newdf = pd.concat([
    ndfa[["width", "examples", "lowest", "highest", "conditioning", "layer", "renormalize"]],
    newdf.drop(columns=["max_offdiagonal"]),
])
newdf

# %%
sns.lineplot(x="layer", y="lowest", hue="width",
             data=newdf.query("examples < width and renormalize == True"), palette=sns.color_palette("bright", n_colors=4));
save_plot("lowest_eigenvalue", renormalize=True)

# %%
sns.lineplot(x="layer", y="lowest", hue="width",
             data=newdf.query("examples < width and renormalize == False"), palette=sns.color_palette("bright", n_colors=3));

# %% [markdown]
# From the obtained data we can observe that
# $$\lambda_\text{max} \sim \sqrt{\text{examples} \cdot \text{width}}$$
# $$\lambda_\text{min} \sim \sqrt{\text{width}} - \sqrt{\text{examples}}$$
#
# Cite: Rudelson, Mark and Vershynin, Roman - "Smallest singular value of a random rectangular matrix" for the formula for smallest eigenvalue.

# %%

# %%
sns.lineplot(x="layer", y="highest", hue="width", data=newdf.query("renormalize == True"), ci=None,
            palette=sns.color_palette("bright", n_colors=4));
save_plot("highest_eigenvalue", renormalize=True)

# %%
sns.lineplot(x="layer", y="highest", style="width", data=newdf.query("renormalize == False"));
save_plot("highest_eigenvalue", renormalize=False)

# %%
dfa["examples_times_width_sqrt"] = (dfa["examples"] * dfa["width"]).apply(math.sqrt)

# %%
sns.relplot(data=dfa.query("layer in ['1', '11', '21', 'latest']"), x="examples_times_width_sqrt",
            y="highest", col="layer")
save_plot("highest_wrt_width_examples_difference")

# %%
# This is an estimate from above from the collected data
dfa["estimated_maxeig_above"] = (dfa["width"] * (1.0 + dfa["max_sumoffdiagonal"])).apply(math.sqrt)
sns.lineplot(x="layer", y="estimated_maxeig_above", style="width", data=dfa.query("examples == 500 and renormalize == False"), label="estimated")
sns.lineplot(x="layer", y="highest", style="width", data=dfa.query("examples == 500 and renormalize == False"), label="real")
save_plot("estimated_maxeig_above", examples=500)

# %%
dfa["estimated_maxeig_below"] = (dfa["width"] * (dfa["mean_sumoffdiagonal"] + dfa["min_ondiagonal"])).apply(math.sqrt)
sns.lineplot(x="layer", y="estimated_maxeig_below", style="width", data=dfa.query("examples == 500 and renormalize == False"), label="estimated")
sns.lineplot(x="layer", y="highest", style="width", data=dfa.query("examples == 500 and renormalize == False"), label="real")
save_plot("estimated_maxeig_below", examples=500)

# %%
sns.lineplot(x="layer", y="conditioning", style="width", hue="examples",
             data=newdf.query("renormalize == True and examples < width"),
             palette=sns.color_palette("bright", n_colors=4));
save_plot("conditioning_layer", renormalize=True)
