# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2021 Google LLC.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
"""


# Download the codebase
!git clone https://github.com/google-research/google-research.git --depth=1

import os
os.chdir("google-research")

# Install Python packages
!pip install flax optax

# See https://github.com/google/jax/blob/master/cloud_tpu_colabs/JAX_demo.ipynb
import requests
import os

if 'COLAB_TPU_ADDR' not in os.environ:
  raise RuntimeError("Please connect to a TPU runtime first!")

if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'
  resp = requests.post(url)
  TPU_DRIVER_MODE = 1

# Use TPU Driver as JAX's backend.
from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)

import jax
print(jax.devices())

"""
Explanation: Paper experiments for "Learning generalized Gumbel-max causal mechanisms"
This notebook shows how to reproduce the experiments for Section 7.1 and the first part of Section 7.2 (the first two columns of Table 1).
Setting up the environment
End of explanation
"""


import functools
import time
from typing import *

import numpy as np
import jax
import jax.numpy as jnp
import optax
import flax
import flax.linen as nn

%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib as mpl

plt.ion()
np.set_printoptions(linewidth=150)

from gumbel_max_causal_gadgets import coupling_util
from gumbel_max_causal_gadgets import gadget_1
from gumbel_max_causal_gadgets import gadget_2
from gumbel_max_causal_gadgets import experiment_util

"""
Explanation: Imports and configuration
End of explanation
"""


def perturbed_range_logit_pair_distribution_fn(rng, dim, base_scale=.1, noise_scale=.1):
  p_rng, q_rng = jax.random.split(rng, 2)
  p_base = jnp.arange(dim) - (dim - 1.0) / 2
  q_base = -p_base
  p_logits = base_scale * p_base + noise_scale * jax.random.normal(p_rng, (dim,))
  q_logits = base_scale * q_base + noise_scale * jax.random.normal(q_rng, (dim,))
  return p_logits, q_logits

def maximal_coupling_loss_matrix_fn(logits1, logits2):
  return 1.0 - jnp.eye(logits1.shape[0])

experiments = []
S_dim = 10
Z_dim = 100
for lr in [1e-5]: # 3e-5,
  for noise_scale in [0, .1, .4, 1.6, 6.4, 25.6, 102.4]:
    for train_seed in [1, 2, 3, 4, 5]:
      ex = experiment_util.CouplingExperimentConfig(
        name=f"noise_scale={noise_scale} Z_dim={Z_dim} lr={lr} train_seed={train_seed}",
        model=(
            gadget_2.GadgetTwoMLPPredictor(
                S_dim=S_dim,
                Z_dim=Z_dim,
                hidden_features=[1024, 1024],
                relaxation_temperature=1.0,
                learn_prior=False)
        ),
        logit_pair_distribution_fn=functools.partial(
            perturbed_range_logit_pair_distribution_fn,
            dim=S_dim,
            base_scale=.1,
            noise_scale=noise_scale),
        coupling_loss_matrix_fn=maximal_coupling_loss_matrix_fn,
        inner_num_samples=16,
        batch_size=64,
        use_transpose=False,
        tx=optax.adam(lr),
        num_steps=10_000,
        print_every=1000,
        metadata={"train_seed": train_seed},
      )
      experiments.append(ex)


results = []
test_results = []
for ex in experiments:
  print("=" * 20)
  print(f"Training: {ex.name}")
  res = ex.train(jax.random.PRNGKey(ex.metadata["train_seed"]))
  results.append(res)
  time.sleep(0.5)
  (logits_1, logits_2), couplings = experiment_util.get_coupling_estimates(
      [ex], [res], 1000 + ex.metadata["train_seed"],
      num_joint_samples=100_000,
      logit_kwargs={"noise_scale":0.0})
  test_loss = experiment_util.compute_coupling_losses([ex], logits_1, logits_2, couplings)
  print(test_loss)
  test_results.append((test_loss, logits_1, logits_2, couplings))

import collections
import re

table = collections.defaultdict(lambda: collections.defaultdict(list))
for task_results in test_results:
  for name, loss in task_results[0].items():
    if name.startswith('noise_scale'):# and "lr=1e-05" in name:
      m = re.search('noise_scale=(\S+) Z_dim=(\S+)', name)
      if m is not None:
        noise_scale, Z_dim = m.groups()
        table[Z_dim][float(noise_scale)].append(float(loss))
table

noise_scales, losses = zip(*table['100'].items())
noise_scales, losses

losses = np.array(list(losses))
loss_avgs = np.mean(losses, axis=-1)
loss_stds = np.std(losses, axis=-1)

_, logits_1, logits_2, _ = test_results[0]
score_matrix = maximal_coupling_loss_matrix_fn(logits_1, logits_2)

gumbel_max_coupling = coupling_util.joint_from_samples(
    coupling_util.gumbel_max_sampler,
    logits_1=logits_1,
    logits_2=logits_2,
    rng=jax.random.PRNGKey(0),
    num_samples=100_000)
gumbel_max_score = jnp.sum(score_matrix * gumbel_max_coupling)

maximal_coupling = coupling_util.joint_from_samples(
    coupling_util.maximal_coupling_sampler,
    logits_1=logits_1,
    logits_2=logits_2,
    rng=jax.random.PRNGKey(0),
    num_samples=100_000)
maximal_score = jnp.sum(score_matrix * maximal_coupling)

gumbel_max_score, maximal_score

fig, ax = plt.subplots(figsize=(3, 3))
ax.figure.set_dpi(300)

plt.errorbar(noise_scales, loss_avgs, yerr=loss_stds, capsize=3, label="Ours")
ax.set_xscale('symlog', linthreshx=.01)

plt.axhline(gumbel_max_score, c='r', ls='--', label="Gumbel-max")
plt.axhline(maximal_score, c='k', ls='-.', label="Optimal")

plt.legend(loc='upper left')
plt.ylim([.24, .305])

plt.xlabel('Noise scale during training')
plt.ylabel('Coupling loss (test)')

i = 0
ex = experiments[i]
experiment_util.visualize_coupling_experiments(*test_results[i])

"""
Explanation: Section 7.1: Optimizing for Maximality
End of explanation
"""


def softmax_uniform_logits(rng, dim=10):
  """Sample logits in a way that is sparse."""
  return jax.nn.log_softmax(8 * jax.random.uniform(rng, shape=[dim]))

def softmax_uniform_logit_pair_distribution_fn(key, dim, p_q_mode, fixed_p_q_seed=123456):
  if p_q_mode == "fixed":
    key = jax.random.PRNGKey(fixed_p_q_seed)
    key_p, key_q = jax.random.split(key, 2)
    p_logits = softmax_uniform_logits(key_p, dim)
    q_logits = softmax_uniform_logits(key_q, dim)
  elif p_q_mode == "independent":
    key_p, key_q = jax.random.split(key, 2)
    p_logits = softmax_uniform_logits(key_p, dim)
    q_logits = softmax_uniform_logits(key_q, dim)
  elif p_q_mode == "reverse":
    key_p = key
    p_logits = softmax_uniform_logits(key_p, dim)
    q_logits = p_logits[::-1]
  else:
    raise NotImplementedError(p_q_mode)

  return p_logits, q_logits

def squared_loss_matrix_fn(logits1, logits2):
  seq = jnp.arange(logits1.shape[0]).astype(jnp.float32)
  return jnp.square(seq[None, :] - seq[:, None])

"""
Explanation: Section 7.2 part 1: Minimizing variance with random p and q
End of explanation
"""


S_dim = 10
experiments = []
for lr in [1e-5, 1e-4, 3e-4, 1e-3, 3e-3, 1e-2]:
  for p_q_mode in ["reverse", "independent", "fixed"]:
    for is_gadget_2 in [True, False]:
      ex = experiment_util.CouplingExperimentConfig(
        name=f"G{2 if is_gadget_2 else 1} relaxed {p_q_mode} lr={lr}",
        model=(
            gadget_2.GadgetTwoMLPPredictor(
                S_dim=10, Z_dim=20, hidden_features=[1024, 1024],
                relaxation_temperature=1.0, learn_prior=False)
            if is_gadget_2 else
            gadget_1.GadgetOneMLPPredictor(
                S_dim=10, hidden_features=[1024, 1024],
                relaxation_temperature=1.0)
        ),
        logit_pair_distribution_fn=functools.partial(
            softmax_uniform_logit_pair_distribution_fn,
            dim=S_dim,
            p_q_mode=p_q_mode),
        coupling_loss_matrix_fn=squared_loss_matrix_fn,
        inner_num_samples=16,
        batch_size=64,
        use_transpose=(not is_gadget_2),
        tx=optax.adam(lr),
        num_steps=5_000,
        print_every=2_500,
        metadata=dict(lr=lr, p_q_mode=p_q_mode, is_gadget_2=is_gadget_2),
      )
      experiments.append(ex)

results = []
for ex in experiments:
  print("=" * 20)
  print(f"Training: {ex.name}")
  res = ex.train(jax.random.PRNGKey(0))
  time.sleep(0.1)
  client = jax.lib.xla_bridge.get_backend()
  client.defragment()
  results.append(res)

plt.figure(figsize=(15,15))
WINDOW_SIZE = 25
for ex, res in zip(experiments, results):
  plt.plot(np.arange(0, ex.num_steps, WINDOW_SIZE), np.reshape([m["loss"] for m in res.all_metrics], [-1, WINDOW_SIZE]).mean(-1), label=ex.name)
plt.legend()

"""
Explanation: Phase 1: Tuning learning rate
End of explanation
"""


indices = [i for i, ex in enumerate(experiments) if (ex.metadata["p_q_mode"] == "independent")]
plt.figure(figsize=(5,5))
WINDOW_SIZE = 25
for i in indices:
  plt.plot(np.arange(0, experiments[i].num_steps, WINDOW_SIZE), np.reshape([m["loss"] for m in results[i].all_metrics], [-1, WINDOW_SIZE]).mean(-1), label=experiments[i].name)
plt.legend()

eval_results_independent = experiment_util.evaluate_all([experiments[i] for i in indices], [results[i] for i in indices],
                                    seed=765987, num_pairs=10_000, samples_per_pair=1_000, loop_size=500)

{k:summary for k, (summary, _, _, _) in eval_results_independent.items()}

lr_bests = {}
for i in indices:
  ex = experiments[i]
  shortname = ex.name.split(" lr=")[0]
  if shortname not in lr_bests or eval_results_independent[ex.name][1] < lr_bests[shortname][2]:
    lr_bests[shortname] = (ex.name, i, eval_results_independent[ex.name][1])

lr_best_subset = [i for _, i, _ in lr_bests.values()]
lr_best_subset_independent = lr_best_subset
lr_bests

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_independent], [results[i] for i in lr_best_subset_independent], 1)

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_independent], [results[i] for i in lr_best_subset_independent], 2)

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_independent], [results[i] for i in lr_best_subset_independent], 3)

"""
Explanation: Extract subset for each mode
Independent
End of explanation
"""


indices = [i for i, ex in enumerate(experiments) if (ex.metadata["p_q_mode"] == "reverse")]
plt.figure(figsize=(5,5))
WINDOW_SIZE = 25
for i in indices:
  plt.plot(np.arange(0, experiments[i].num_steps, WINDOW_SIZE), np.reshape([m["loss"] for m in results[i].all_metrics], [-1, WINDOW_SIZE]).mean(-1), label=experiments[i].name)
plt.legend()

eval_results_reverse = experiment_util.evaluate_all([experiments[i] for i in indices], [results[i] for i in indices],
                                    seed=765987, num_pairs=10_000, samples_per_pair=1_000, loop_size=500)

{k:summary for k, (summary, _, _, _) in eval_results_reverse.items()}

lr_bests = {}
for i in indices:
  ex = experiments[i]
  shortname = ex.name.split(" lr=")[0]
  if shortname not in lr_bests or eval_results_reverse[ex.name][1] < lr_bests[shortname][2]:
    lr_bests[shortname] = (ex.name, i, eval_results_reverse[ex.name][1])

lr_best_subset = [i for _, i, _ in lr_bests.values()]
lr_best_subset_reverse = lr_best_subset
lr_bests

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_reverse], [results[i] for i in lr_best_subset_reverse], 1)

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_reverse], [results[i] for i in lr_best_subset_reverse], 2)

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_reverse], [results[i] for i in lr_best_subset_reverse], 3)

"""
Explanation: Reverse
End of explanation
"""


indices = [i for i, ex in enumerate(experiments) if (ex.metadata["p_q_mode"] == "fixed")]
plt.figure(figsize=(5,5))
WINDOW_SIZE = 25
for i in indices:
  plt.plot(np.arange(0, experiments[i].num_steps, WINDOW_SIZE), np.reshape([m["loss"] for m in results[i].all_metrics], [-1, WINDOW_SIZE]).mean(-1), label=experiments[i].name)
plt.legend()

eval_results_fixed = experiment_util.evaluate_all([experiments[i] for i in indices], [results[i] for i in indices],
                                    seed=765987, num_pairs=1, samples_per_pair=100_000, loop_size=500)

{k:summary for k, (summary, _, _, _) in eval_results_fixed.items()}

{k:(isd/np.sqrt(100_000)) for k, (_, _, _, isd) in eval_results_fixed.items()}

lr_bests = {}
for i in indices:
  ex = experiments[i]
  shortname = ex.name.split(" lr=")[0]
  if shortname not in lr_bests or eval_results_fixed[ex.name][1] < lr_bests[shortname][2]:
    lr_bests[shortname] = (ex.name, i, eval_results_fixed[ex.name][1])

lr_best_subset = [i for _, i, _ in lr_bests.values()]
lr_best_subset_fixed = lr_best_subset
lr_bests

experiment_util.compute_and_visualize_coupling_experiments([experiments[i] for i in lr_best_subset_fixed], [results[i] for i in lr_best_subset_fixed], 1)

"""
Explanation: Fixed
End of explanation
"""


best_lr_map = {}
for i in (lr_best_subset_fixed + lr_best_subset_independent + lr_best_subset_reverse):
  print(i, experiments[i].metadata)
  best_lr_map[experiments[i].metadata["is_gadget_2"], experiments[i].metadata["p_q_mode"]] = experiments[i].metadata["lr"]

"""
Explanation: Extract best learning rates
End of explanation
"""


S_dim = 10
experiments = []
for p_q_mode in ["reverse", "independent", "fixed"]:
  for is_gadget_2 in [True, False]:
    for train_seed in [1, 2, 3, 4, 5]:
      lr = best_lr_map[is_gadget_2, p_q_mode]
      ex = experiment_util.CouplingExperimentConfig(
        name=f"G{2 if is_gadget_2 else 1} relaxed {p_q_mode} lr={lr} seed={train_seed}",
        model=(
            gadget_2.GadgetTwoMLPPredictor(
                S_dim=10, Z_dim=20, hidden_features=[1024, 1024],
                relaxation_temperature=1.0, learn_prior=False)
            if is_gadget_2 else
            gadget_1.GadgetOneMLPPredictor(
                S_dim=10, hidden_features=[1024, 1024],
                relaxation_temperature=1.0)
        ),
        logit_pair_distribution_fn=functools.partial(
            softmax_uniform_logit_pair_distribution_fn,
            dim=S_dim,
            p_q_mode=p_q_mode),
        coupling_loss_matrix_fn=squared_loss_matrix_fn,
        inner_num_samples=16,
        batch_size=64,
        use_transpose=(not is_gadget_2),
        tx=optax.adam(lr),
        num_steps=50_000,
        print_every=2_500,
        metadata=dict(lr=lr, p_q_mode=p_q_mode, is_gadget_2=is_gadget_2, train_seed=train_seed),
      )
      experiments.append(ex)

results = []
for ex in experiments:
  print("=" * 20)
  print(f"Training: {ex.name}")
  res = ex.train(jax.random.PRNGKey(ex.metadata["train_seed"]))
  time.sleep(0.1)
  client = jax.lib.xla_bridge.get_backend()
  client.defragment()
  results.append(res)

plt.figure(figsize=(15,7))
WINDOW_SIZE = 25
for ex, res in zip(experiments, results):
  plt.plot(np.arange(0, ex.num_steps, WINDOW_SIZE),
           np.reshape([m["loss"] for m in res.all_metrics],
                    [-1, WINDOW_SIZE]).mean(-1),
           "-" if ex.metadata["is_gadget_2"] else "--",
           label=ex.name,
           color=mpl.cm.tab10(
               (3 if ex.metadata["is_gadget_2"] else 0)
               + {"independent": 0, "reverse":1, "fixed":2}[ex.metadata["p_q_mode"]]))
plt.legend()

"""
Explanation: Phase 2: Retrain for multiple seeds with tuned learning rate
End of explanation
"""


scores_per_mode = {}
# Each p_q_mode is a separate column we want data form
for p_q_mode in ["reverse", "independent", "fixed"]:
  scores_per_mode[p_q_mode] = []
  # For each column, run five random evaluations. Each evaluation will itself
  # average over some number of logit pairs, but we will average over enough
  # so that the variance in this is minimal within each evaluation. For our
  # gadgets, each evaluation represents a separate training run with a separate
  # evaluation; for the baselines, we just redo the evaluation with a different
  # test set.
  for seed in [1, 2, 3, 4, 5]:
    print(p_q_mode, seed)
    sub_exps = []
    sub_res = []
    for ex, res in zip(experiments, results):
      if ex.metadata["p_q_mode"] == p_q_mode and ex.metadata["train_seed"] == seed:
        sub_exps.append(ex)
        sub_res.append(res)
    
    num_pairs = 10_000 if p_q_mode is not "fixed" else 1
    samples_per_pair = 1_000 if p_q_mode is not "fixed" else 100_000
    eval_results = experiment_util.evaluate_all(
        sub_exps,
        sub_res,
        seed=1000 + seed,
        num_pairs=num_pairs,
        samples_per_pair=samples_per_pair,
        loop_size=500)
    scores_per_mode[p_q_mode].append(eval_results)

scores_per_mode

def fix_name(name):
  lookups = {
    "Independent":"Independent",
    "ICDF":"CDF$^{-1}$",
    "ICDF (permuted)":"Permuted CDF$^{-1}$",
    "Gumbel-max":"Gumbel-max",
    "G1":"Gadget 1",
    "G2":"Gadget 2",
  }
  return lookups[name.split(" relaxed")[0]]


all_scores = {}
for p_q_mode in ["reverse", "independent", "fixed"]:
  scores = collections.defaultdict(list)
  for eval_sample in scores_per_mode[p_q_mode]:
    for name, (_, eval_average_for_seed, _, _) in eval_sample.items():
      if " seed" in name:
        name = name.split(" seed")[0]
      scores[name].append(eval_average_for_seed)
  all_scores[p_q_mode] = {fix_name(k):np.array(v) for k,v in scores.items()}

all_scores

"""
Explanation: Evaluate each model for each seed
End of explanation
"""


name_order = [
  "Independent",
  "CDF$^{-1}$",
  "Permuted CDF$^{-1}$",
  "Gumbel-max",
  "Gadget 1",
  "Gadget 2",
]

for name in name_order:
  parts = [name]
  parts.append(" & ")
  for p_q_mode in ["independent", "reverse"]: # "fixed",
    avg = np.mean(all_scores[p_q_mode][name])
    std = np.std(all_scores[p_q_mode][name])
    parts.append(f"{avg:.2f} $\pm$ {std:.2f}")
    parts.append(" & ")
  parts.pop()
  parts.append(" \\\\")
  print("".join(parts))

vis_exs = []
vis_res = []
for ex, res in zip(experiments, results):
  if ex.metadata["p_q_mode"] == "reverse" and ex.metadata["train_seed"] == 1:
    vis_exs.append(ex)
    vis_res.append(res)

# vis_exs, vis_res = [experiments[i] for i in lr_best_subset_reverse], [results[i] for i in lr_best_subset_reverse]
(logits_1, logits_2), coupling_estimates = experiment_util.get_coupling_estimates(vis_exs, vis_res, 2)

loss_values = experiment_util.compute_coupling_losses(experiments, logits_1, logits_2, coupling_estimates)

ncouplings =  len(coupling_estimates)
ncols = 2 + ncouplings
fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(2 + 4*ncouplings, 4), gridspec_kw={'width_ratios': [2,2] + [11]*ncouplings})
axs[0].imshow(jnp.exp(logits_1)[:, None], vmin=0)
axs[0].set_xticks([])
axs[0].set_yticks(np.arange(10))
axs[0].set_title("$p$")
axs[1].imshow(jnp.exp(logits_2)[:, None], vmin=0)
axs[1].set_xticks([])
axs[1].set_yticks(np.arange(10))
axs[1].set_title("$q$")
for j, (name, coupling) in enumerate(sorted(coupling_estimates.items(), key=lambda t: fix_order(t[0]))):
  axs[j+2].imshow(coupling, vmin=0)
  axs[j+2].set_title(f"{fix_name(name)}: {loss_values[name]:.2f}")
  axs[j+2].set_xticks(np.arange(10))
  axs[j+2].set_yticks(np.arange(10))
fig.tight_layout()

"""
Explanation: Build the table
End of explanation
"""


