# ---
# 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
# ---

# %%

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
from glob import iglob
import os

# %%
# Loads the csv files found in result as diverse dataframes
df = {}
for measure in saved_measures:
    df[measure] = pd.concat((pd.read_csv(filename) for filename in iglob(f"main/results/{measure}*.csv")))

# %%
from argparse import ArgumentParser

parser = ArgumentParser(description="Analyze the specified part of data.")
parser.add_argument("--dataset", type=str, default="CIFAR10")
parser.add_argument("--n_layers", type=int, default=6)
parser.add_argument("--loss_function", type=str, default="CE")
parser.add_argument("--network", type=str, default="FCN")
parser.add_argument("--activation", type=str, default="ReLU")
parser.add_argument("--examples", type=int, default=None)
parser.add_argument("--width", type=int, default=None)
parser.add_argument("--renormalize", type=bool, default=False)
parser.add_argument("--lr", type=float, default=None)
parser.add_argument("--correct_filter", type=float, default=0.12)
parser.add_argument("-f", type=str, default="ignore",
                    help="Ignore if we are executing in notebook")
args = parser.parse_args()

qry = " and ".join((f"{name} == {repr(value)}" for name, value in args.__dict__.items() 
                    if name not in ["f", "correct_filter"] and value is not None))
print("Query:", qry)


# %%
def save_plot(name, **kwargs):
    if not os.path.exists("images/general/"):
        os.makedirs("images/general/", 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/general/{name}@{slug}.svg")


# %%
print("Total records:", df["general"].size)

# Filter the loaded datasets depending on the setted arguments
for measure in saved_measures:
    df[measure] = df[measure].query(qry)
print("Filtered records:", df["general"].size)

# %%
df["general"].query("epoch == epochs and correct < @args.correct_filter")

# %%
# Filter only the examples for which the final correct value is >= filter (i.e. they did not completely diverge).
correct_filter = args.correct_filter
"These are the columns that determine the run we are into"
case_columns = ["network", "lr", "seed", "dataset", "n_layers", "width", "examples", "batch_size", "batch",
                "epochs", "loss_function", "activation", "kernel_size", "renormalize"]
to_retain = df["general"].query("epoch == epochs and correct >= @correct_filter")[case_columns]
print(f"Retained examples (correct >= {correct_filter}): {to_retain.size}/{df['general'].query('epoch == epochs').size}")
for lr in df["general"]["lr"].unique():
    print(f"Retained examples with {lr=}: {to_retain.query('lr == @lr').size}/{df['general'].query('epoch == epochs and lr == @lr').size}")

for measure in saved_measures:
    df[measure] = pd.merge(df[measure], to_retain, how="inner", on=tuple(case_columns),
                           suffixes=("", ""), validate="many_to_one")

# %%
dfg = df["general"]
dfi = df["iconditioning"]
dfd = df["distance"]
dfn = df["norms"]
dfm = df["input_matrix"]
dfa = df["input_matrix"].copy()
dfa["lowest"] = dfa["min_eigenvalue"]
dfa["highest"] = dfa["max_eigenvalue"]

# %% [markdown]
# Since we first measure the train loss, then update weights, then measure test loss, the train loss refers to pre-updated weights.
# We thus have to refer `train_loss` to `epoch - 1`.

# %%
case_columns = ["epoch", "batch", "dataset", "n_layers", "width", "examples", 
                "batch_size", "epochs", "lr", "seed", "network", "loss_function", 
                "activation", "kernel_size", "renormalize"]
dfgg = dfg.loc[:, case_columns + ["train_loss"]]
dfgg["epoch"] -= 1

dfg = pd.merge(dfg.drop(columns=["train_loss"]), dfgg, on=tuple(case_columns), how="inner", suffixes=("", ""))

# %% [markdown]
# ## General introductory plots

# %%
for examples in dfg["examples"].unique():
    print(f"Examples: {examples}")
    sns.lineplot(x="epoch", y="correct", hue="lr", data=dfg.query("examples == @examples"))#, palette=sns.color_palette("bright", n_colors=4));
    save_plot("correct_epoch", examples=examples)
    plt.show();

# %%
sns.lineplot(x="epoch", y="train_loss", label="train_loss", style="lr", data=dfg);
sns.lineplot(x="epoch", y="test_loss", label="test_loss", style="lr", ci=None, data=dfg);

# %% [markdown]
# # 1. Testing well-conditioning theory
# We measure conditioning at initialization and compare its improvement with depth with the amount predicted by the theory.

# %%
from utils import calculate_expected_eigenvalue, calculate_parameters

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

# %% [markdown]
# We can see that there is a correlation between good conditioning of the latest layer and higher correctness percentage.

# %%
sns.lineplot(x="layer", y="conditioning", hue="examples", style="width", data=dfa.query("epoch == 0 and examples < width"));
# TODO: What would be the predicted amount for the given width and number of examples?
# TODO: Plot also the expected decrease with infinite width


# %%
sns.lineplot(x="layer", y="lowest", hue="examples", style="width", data=dfa.query("epoch == 0"));

# %% [markdown]
# We do see that conditioning decreases, as expected from the [AAK20] theory, but in the following plot we see that the off diagonal values are not really decreasing to zero and that the norm of the ondiagonal terms is not being preserved, possibly because of small width...

# %%
sns.lineplot(x="layer", y="min_ondiagonal", style="width", data=dfm.query("epoch == 0"), label="min on diagonal")
sns.lineplot(x="layer", y="max_offdiagonal", style="width", data=dfm.query("epoch == 0"), label="max off diagonal");

# %%
dfm["minondiagonal_maxsumoffdiagonal_ratio"] = dfm["max_sumoffdiagonal"] / dfm["min_ondiagonal"]
sns.lineplot(x="epoch", y="minondiagonal_maxsumoffdiagonal_ratio", style="lr", hue="layer", data=dfm, ci=None);

# %%
# But what happens to such things over the course of training?
sns.lineplot(x="epoch", y="min_ondiagonal", hue="lr", data=dfm.query("examples < width"));

# %%
sns.lineplot(x="epoch", y="max_offdiagonal", hue="lr", data=dfm);

# %%
dfm["minondiag_maxoffdiag_ratio"] = dfm["min_ondiagonal"] / dfm["max_offdiagonal"]
sns.lineplot(x="epoch", y="minondiag_maxoffdiag_ratio", hue="lr", data=dfm);

# %% [markdown]
# ## Evolution of conditioning during training
# How does conditioning of the layers evolve during the network training?
# Given our theory we should expect it to get worse, depending on amount of distance from initialization traveled.

# %%
sns.lineplot(x="epoch", y="conditioning", style="lr", hue="layer", data=dfa.query("layer != '0' and examples < width and lr < 0.01"), ci=None);
save_plot("conditioning_vs_epochs")

# %%
sns.lineplot(x="epoch", y="conditioning", style="lr", units="seed", estimator=None, hue="examples", data=dfa.query("layer == '1' and examples < width"));

# %% [markdown]
# We can see how the learning rate has a devastating difference in the outcome of the conditioning.
# Is it because we fall in regions where matrices are near-singular for too high learning rates?

# %%
sns.lineplot(x="epoch", y="conditioning", style="lr", hue="layer", data=dfa.query("layer != '0' and examples == width"), ci=None);

# %%
sns.lineplot(x="epoch", y="highest", style="lr", hue="layer", data=dfa.query("layer != '0' and examples < width"), ci=None);

# %%
sns.lineplot(x="epoch", y="lowest", style="lr", hue="layer", data=dfa.query("layer != '0'"), ci=None);

# %% [markdown]
# We can notice how conditioning does get worse in the beginning, but then the lowest eigenvalues begin to rise, and also the highest ones begin to decrease.

# %%
sns.lineplot(x="epoch", y="wdistance", style="lr", data=dfd, ci=None);

# %%
sns.lineplot(x="epoch", y="wdistanceinf", style="lr", hue="layer", data=dfd, ci=None);

# %%
sns.lineplot(x="epoch", y="wdistanceinf", style="lr", hue="layer", data=dfd.query("lr < 0.001"), ci=None);

# %% [markdown]
# There must be something else at work in the second phase that makes conditioning get better.

# %%
for measurement in ["vector_norm_2", "vector_norm_inf", "matrix_norm_2", "matrix_norm_inf"]:
    sns.lineplot(x="epoch", hue="layer", style="lr", y=measurement, data=dfn);
    plt.show();

# %% [markdown]
# We can see that the norms do not really change during training, and can effectively be supposed constants.

# %% [markdown]
# For what concerns our relation between conditioning and distance from initialization we recall that obtained
# $$\zeta(D_\alpha f_M(\alpha, x)) \le \frac1{\lambda_\alpha - G_\phi \|\alpha M\|_2 \|x - x_0\|}$$
# where $\lambda_\alpha$ is the minimum eigenvalue at initialization, $G_\phi$ is the Lipschitz constant of the activation function.

# %%
sns.lineplot(x="epoch", y="imaxsingledistance", hue="layer", style="lr", data=dfd.query("layer == '4'"))

# %%
sns.lineplot(x="epoch", y="idistance", hue="examples", style="lr", data=dfd.query("layer == '4'"))

# %%
dfcur = dfa.copy()
# Add original lowest eigenvalue to the current dataframe
case_columns = ["network", "lr", "seed", "layer", "dataset", "width", "examples",
                "batch_size", "epochs", "loss_function", "activation", "kernel_size", "renormalize"]
dfcur = pd.merge(dfcur, dfcur[dfcur["epoch"] == 0][case_columns + ["lowest", "highest"]],
             how="right", on=tuple(case_columns), sort=False,
             suffixes=("", "_original"), validate="many_to_one")
# Need to add the distance from beginning and the current norm of the weights
dfcur = pd.merge(dfcur, dfd[case_columns + ["imaxsingledistance", "idistance", "wdistance", "epoch"]], how="inner",
                 on=tuple(case_columns + ["epoch"]), sort=False,
                 suffixes=("", ""), validate="one_to_one");
# Also we add the norm of the weight matrix
dfcur = pd.merge(dfcur, dfn[case_columns + ["vector_norm_2", "epoch"]], how="inner",
                 on=tuple(case_columns + ["epoch"]), sort=False,
                 suffixes=("", ""), validate="one_to_one")

# %%
actmean, actvar, actlip, actmeansmoothness = calculate_parameters(args.activation)
activation_lipschitzianity = actlip / math.sqrt(actvar)
dfcur["predicted_lowest"] = (dfcur["lowest_original"] - activation_lipschitzianity * dfcur["idistance"] / dfcur["width"].apply(math.sqrt)).apply(lambda x: max(x, 0))

data = dfcur.query("epoch <= 80 and layer == '2' and examples < width")
sns.lineplot(x="epoch", y="lowest", style="examples", data=data)
sns.lineplot(x="epoch", y="predicted_lowest", style="examples", data=data, legend=False,
            )#style=True, dashes=[(12, 2)]);
save_plot("lowest_predicted_bounds")

# %%
data["lowest_ratio"] = data["predicted_lowest"] / data["lowest"]
sns.lineplot(x="epoch", y="lowest_ratio", style="examples", data=data)

# %%
dfcur["highest_original_p"] = dfcur["highest_original"] / (dfcur["examples"] * dfcur["width"]).apply(math.sqrt)
dfcur["highest_p"] = dfcur["highest"] / (dfcur["examples"] * dfcur["width"]).apply(math.sqrt)

dfcur["predicted_highest"] = dfcur["highest_original_p"] + activation_lipschitzianity * dfcur["idistance"] / (dfcur["examples"] * dfcur["width"]).apply(math.sqrt) 

data = dfcur.query("epoch <= 80 and layer == '1'")
sns.lineplot(x="epoch", y="highest_p", style="examples", data=data, ci=None)
sns.lineplot(x="epoch", y="predicted_highest", style="examples", data=data, legend=False,
            )#style=True, dashes=[(12, 2)]);

# %%
data["hratio"] = data["highest_p"] / data["predicted_highest"]
sns.lineplot(x="epoch", y="hratio", data=data)

# %%
# Here again we can see that after the first epochs we begin to see an increase in the lowest eigenvalue which we
# would not expect from the distance traveled by the layers and 

# %% [markdown]
# All of the above plots show that we have a non-zero lowest eigenvalue in all layers, and thus that during the whole training time the network model is PL with respect to the last layer at least...
#
# We now plot conditioning with respect to the previous input layer, and then we continue with studying the PL theory of optimization of networks.

# %%
sns.lineplot(x="epoch", y="lowest", hue="layer", data=dfi);

# %%
sns.lineplot(x="epoch", y="highest", hue="layer", data=dfi);

# %% [markdown]
# # 2. Estimating decrease in Loss via PL coefficients
# From the given measurements of the largest and smallest eigenvalues of each layer, we compute estimated PL and smoothness coefficient, and we observe the prediction for the decrease in training loss with the actual decrease wrt the previous time-step.
#
# When we have $\mu$-PL function $f$ and we use $\eta$ as learning rate, i.e. $x^{k+1} = x^k - \eta \nabla f(x^k)$ we get an expected improvement of
# $$(f(x^{k+1}) - f^*) \le (1 - \eta \mu (2 - L \eta)) (f(x^k) - f^*)$$
# where $L$ is the smoothness constant of $f$.
#

# %%
if args.loss_function.endswith("MSE"):
    # What we do in code is first apply softmax, then use MSE wrt one-hot encoded target
    # Moreover MSE in pytorch is different and has: Lipschitz <= 4/N MSE(x), Smoothness <= 4/N, PLNess = 2/N
    
    smoothness = 4.0 / 10.0 # Number of classes
    plness = 2.0 / 10.0
    # To get the lipschitz constant, we can consider directly the value of the MSE loss value
    # lipschitz = 4.0 / 10.0 * mse
    
    # We need to merge things to calculate rates for the last layer only
    case_columns = ["network", "lr", "seed", "dataset", "width", "examples",
                "batch_size", "epochs", "loss_function", "activation", "kernel_size", "renormalize"]
    # Calculate resulting rates for the last layer only
    dfcur = pd.merge(dfn.query("layer == 'latest'"), dfa.query("layer == 'latest'"),
                     how="inner", on=tuple(case_columns + ["layer", "epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", ""), validate="one_to_one")
    dfcur = pd.merge(dfcur, dfg, how="inner",
                     on=tuple(case_columns + ["epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", ""), validate="one_to_one")
    dfcur["latest_lipschitzianity"] = activation_lipschitzianity * dfcur["matrix_norm_2"]
    dfcur["latest_total_smoothness"] = smoothness * dfcur["latest_lipschitzianity"]**2
    dfcur["latest_total_plness"] = plness * dfcur["lowest"]
    expected_decrease = 1.0 - dfcur["lr"] * dfcur["latest_total_plness"] * (2.0 - dfcur["latest_total_smoothness"] * dfcur["lr"])
    dfcur["expected_decrease"] = expected_decrease
    dfcur["expected_next_loss"] = expected_decrease * dfcur["train_loss"]

    # Now join prediction in the previous epoch with value in the next epoch to compare them
    dfprev = dfcur.copy()
    dfprev["epoch"] += 1

    dfcur = pd.merge(dfcur, dfprev, how="inner",
                     on=tuple(case_columns + ["epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", "_prev"), validate="one_to_one")
    dfcur["loss_prediction_ratio"] = dfcur["expected_next_loss_prev"] / dfcur["train_loss"]

# %% [markdown]
# Formula for the decrease of weakly PL functions:
# $$\xi(x_{k+1}) \le \left(1 - \frac{\mu \xi(x_k)}{2 L \|x_k - x_*\|^2}\right) \xi(x_k)$$

# %%
if args.loss_function == "CE":
    # We use the formula for weakly PL functions and use |x_k - x_*| <= R
    mu = 1.0
    smoothness = 2.0
    
    # We need to merge things to calculate rates for the last layer only
    case_columns = ["network", "lr", "seed", "dataset", "width", "examples",
                "batch_size", "epochs", "loss_function", "activation", "kernel_size", "renormalize"]
    # Calculate resulting rates for the last layer only
    dfcur = pd.merge(dfn.query("layer == 'latest'"), dfa.query("layer == 'latest'"),
                     how="inner", on=tuple(case_columns + ["layer", "epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", ""), validate="one_to_one")
    dfcur = pd.merge(dfcur, dfg, how="inner",
                     on=tuple(case_columns + ["epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", ""), validate="one_to_one")
    dfcur["latest_lipschitzianity"] = activation_lipschitzianity * dfcur["matrix_norm_2"]
    dfcur["latest_total_smoothness"] = smoothness * dfcur["latest_lipschitzianity"]**2
    lipschitz = dfcur["latest_lipschitzianity"]
    R = 2.0 * dfcur["train_loss"] / lipschitz
    expected_decrease = 1.0 - dfcur["lr"] * mu * dfcur["train_loss"] / (2.0 * R**2)
    dfcur["expected_decrease"] = expected_decrease
    dfcur["expected_next_loss"] = expected_decrease * dfcur["train_loss"]

    # Now join prediction in the previous epoch with value in the next epoch to compare them
    dfprev = dfcur.copy()
    dfprev["epoch"] += 1

    dfcur = pd.merge(dfcur, dfprev, how="inner",
                     on=tuple(case_columns + ["epoch", "n_layers", "batch"]), sort=False,
                     suffixes=("", "_prev"), validate="one_to_one")
    dfcur["loss_prediction_ratio"] = dfcur["expected_next_loss_prev"] / dfcur["train_loss"]

# %%
sns.lineplot(x="epoch", y="latest_total_smoothness", style="lr", data=dfcur)

# %%
sns.lineplot(x="epoch", y="expected_decrease", style="lr", data=dfcur.query("lr < 0.005"));

# %%
sns.lineplot(x="epoch", y="loss_prediction_ratio", style="lr", data=dfcur.query("lr < 0.005"));
save_plot("loss_prediction_ratio")

# %%
cols = ["epochs", "lr", "seed", "batch", "n_layers", "width", "examples", "batch_size", "loss_function"]
ndfcur = dfcur.copy()
ndfcur["cumlosspred"] = dfcur.groupby(cols)["loss_prediction_ratio"].transform(pd.Series.cumprod)
sns.lineplot(x="epoch", y="cumlosspred", style="lr", data=ndfcur.query("lr < 0.005"))
save_plot("cumulative_loss_prediction_ratio")

# %% [markdown]
# # 3. Estimating Generalization from PL coefficients

# %%
sns.lineplot(x="epoch", y="lowest", hue="layer", data=dfi)
save_plot("lowest_iconditioning")

# %%
sns.lineplot(x="epoch", y="highest", hue="layer", data=dfi)

# %%
from utils import calculate_parameters

actmean, actvar, actlip, _ = calculate_parameters(args.activation)
actlipschitz = actlip / math.sqrt(actvar)
print(f"{args.activation} rescaled lipschitzianity: {actlipschitz}")

# %%
# Compute lipschitz constant for the whole network
cols = ["epochs", "lr", "seed", "batch", "n_layers", "width", "examples", "batch_size", "loss_function"]

case_columns = ["network", "lr", "seed", "dataset", "width", "examples",
                "batch_size", "epochs", "loss_function", "activation", "kernel_size", "renormalize"]
# Perform merge of all per-layer datasets
dfcur = pd.merge(dfn, dfa,
                 how="inner", on=tuple(case_columns + ["layer", "epoch", "n_layers", "batch"]), sort=False,
                 suffixes=("", ""), validate="one_to_one")
dfcur = pd.merge(dfcur, dfg, how="inner",
                 on=tuple(case_columns + ["epoch", "n_layers", "batch"]), sort=False,
                 suffixes=("", ""), validate="many_to_one")
dfcur = pd.merge(dfcur, dfi, how="inner",
                 on=tuple(case_columns + ["layer", "epoch", "n_layers", "batch"]), sort=False,
                 suffixes=("", "_i"), validate="one_to_one")
ndfcur = dfcur.copy()
# Perform the product of lipschitz constant of all layers
ndfcur["net_lipschitz"] = dfcur.groupby(cols + ["epoch"])["matrix_norm_2"].transform(pd.Series.cumprod)
ndfcur["net_lipschitz"] *= actlipschitz**args.n_layers
# ndfcur["net_plness"] = dfcur.groupby(cols + ["epoch"])[""].transform(pd.Series.cumprod)
ndfcur.query("layer == 'latest'")

# %%
dfg["loss_diff"] = (dfg["test_loss"] - dfg["train_loss"]).apply(abs)
sns.lineplot(x="epoch", y="loss_diff", hue="lr", data=dfg)

# %%
# Calculate lipschitzianity only for the last layer, assuming the previous ones just act
# as fixed transformations on the data.
ddf = dfcur.query("layer == 'latest'")
ddf["lipschitz"] = ddf["matrix_norm_2"] * 1.0 # Max lipschitz constant of MSE in our case
ddf["pl"] = ddf["lowest_i"] * plness
ddf["optimization"] = ddf["train_loss"] # We know the minimum is at zero since we are OP
g2mu = 2 * ddf["lipschitz"]**2 / ddf["pl"]
ddf["ptw_stability"] = 2 * (g2mu * ddf["optimization"]).apply(math.sqrt) + g2mu / (ddf["examples"] - 1)
M = 1.0
delta = 0.05 # Probability 95%
ddf["generalization"] = ((M**2 + 12 * M * ddf["examples"] * ddf["ptw_stability"]) / (2 * ddf["examples"] * delta)).apply(math.sqrt)

# %%
sns.lineplot(x="epoch", y="ptw_stability", hue="examples", data=ddf)

# %%
sns.lineplot(x="epoch", y="generalization", hue="examples", data=ddf)
save_plot("generalization")

# %% [markdown]
# Estrapolated generalization bounds are essentially vacuous when applied to MSE and considering only the effect of the last layer...
# We possibly need to use more examples, since these can give vastly better bounds.
# However our networks already do generalize quite a lot, we have to persue another try.

# %%

# %%
