# Path: enn/experiments/neurips_2021/base.py
# class Data(NamedTuple):
# class PriorKnowledge:
# class ENNQuality:
# class EpistemicSampler(typing_extensions.Protocol):
# class TestbedAgent(typing_extensions.Protocol):
# class TestbedProblem(abc.ABC):
#   def __call__(self, x: chex.Array, seed: int = 0) -> chex.Array:
#   def __call__(self,
#                data: Data,
#                prior: Optional[PriorKnowledge] = None) -> EpistemicSampler:
#   def train_data(self) -> Data:
#   def evaluate_quality(self, enn_sampler: EpistemicSampler) -> ENNQuality:
#   def prior_knowledge(self) -> PriorKnowledge:
# 
# Path: enn/experiments/neurips_2021/testbed.py
# class GPRegression:
# class TestbedGPRegression(testbed_base.TestbedProblem):
#   def __init__(self,
#                kernel_fn: nt_types.KernelFn,
#                x_train: chex.Array,
#                x_test: chex.Array,
#                key: chex.PRNGKey,
#                tau: int = 1,
#                noise_std: float = 1,
#                kernel_ridge: float = 1e-6):
#   def x_test(self) -> chex.Array:
#   def test_mean(self) -> chex.Array:
#   def test_cov(self) -> chex.Array:
#   def train_data(self) -> testbed_base.Data:
#   def train_data(self) -> testbed_base.Data:
#   def prior_knowledge(self) -> testbed_base.PriorKnowledge:
#   def evaluate_quality(
#       self,
#       enn_sampler: testbed_base.EpistemicSampler) -> testbed_base.ENNQuality:
# def _kl_gaussian(
#     mean_1: float, std_1: float, mean_2: float, std_2: float) -> float:


import chex
import haiku as hk
import numpy as np
import pandas as pd
import plotnine as gg
from enn.experiments.neurips_2021 import base as testbed_base
from enn.experiments.neurips_2021 import testbed

# Setting plot theme
gg.theme_set(gg.theme_bw(base_size=16, base_family='serif'))
gg.theme_update(figure_size=(12, 8), panel_spacing=0.5)


def sanity_1d(true_model: testbed_base.TestbedProblem,
              enn_sampler: testbed_base.EpistemicSampler) -> gg.ggplot:
  """Sanity check to plot 1D representation of the GP testbed output."""
  if hasattr(true_model, 'problem'):
    true_model = true_model.problem  # Removing logging wrappers
  if not hasattr(true_model, 'data_sampler'):
    return gg.ggplot()
  gp_model = true_model.data_sampler
  return plot_1d_regression(gp_model, enn_sampler)


def _gen_samples(enn_sampler: testbed_base.EpistemicSampler,
                 x: chex.Array,
                 num_samples: int) -> pd.DataFrame:
  """Generate posterior samples at x (not implemented for all posterior)."""
  data = []
  rng = hk.PRNGSequence(0)
  for seed in range(num_samples):
    net_out = enn_sampler(x, next(rng))
    y = net_out[:, 0]
    data.append(pd.DataFrame({'x': x[:, 0], 'y': y, 'seed': seed}))
  return pd.concat(data)



# Path: enn/experiments/neurips_2021/base.py
# class Data(NamedTuple):
# class PriorKnowledge:
# class ENNQuality:
# class EpistemicSampler(typing_extensions.Protocol):
# class TestbedAgent(typing_extensions.Protocol):
# class TestbedProblem(abc.ABC):
#   def __call__(self, x: chex.Array, seed: int = 0) -> chex.Array:
#   def __call__(self,
#                data: Data,
#                prior: Optional[PriorKnowledge] = None) -> EpistemicSampler:
#   def train_data(self) -> Data:
#   def evaluate_quality(self, enn_sampler: EpistemicSampler) -> ENNQuality:
#   def prior_knowledge(self) -> PriorKnowledge:
# 
# Path: enn/experiments/neurips_2021/testbed.py
# class GPRegression:
# class TestbedGPRegression(testbed_base.TestbedProblem):
#   def __init__(self,
#                kernel_fn: nt_types.KernelFn,
#                x_train: chex.Array,
#                x_test: chex.Array,
#                key: chex.PRNGKey,
#                tau: int = 1,
#                noise_std: float = 1,
#                kernel_ridge: float = 1e-6):
#   def x_test(self) -> chex.Array:
#   def test_mean(self) -> chex.Array:
#   def test_cov(self) -> chex.Array:
#   def train_data(self) -> testbed_base.Data:
#   def train_data(self) -> testbed_base.Data:
#   def prior_knowledge(self) -> testbed_base.PriorKnowledge:
#   def evaluate_quality(
#       self,
#       enn_sampler: testbed_base.EpistemicSampler) -> testbed_base.ENNQuality:
# def _kl_gaussian(
#     mean_1: float, std_1: float, mean_2: float, std_2: float) -> float:


# Path: enn/experiments/neurips_2021/plotting.py
import chex
import haiku as hk
import numpy as np
import pandas as pd
import plotnine as gg
from enn.experiments.neurips_2021 import base as testbed_base
from enn.experiments.neurips_2021 import testbed


# Setting plot theme
gg.theme_set(gg.theme_bw(base_size=16, base_family='serif'))
gg.theme_update(figure_size=(12, 8), panel_spacing=0.5)


def sanity_1d(true_model: testbed_base.TestbedProblem,
              enn_sampler: testbed_base.EpistemicSampler) -> gg.ggplot:
  """Sanity check to plot 1D representation of the GP testbed output."""
  if hasattr(true_model, 'problem'):
    true_model = true_model.problem  # Removing logging wrappers
  if not hasattr(true_model, 'data_sampler'):
    return gg.ggplot()
  gp_model = true_model.data_sampler
  return plot_1d_regression(gp_model, enn_sampler)


def _gen_samples(enn_sampler: testbed_base.EpistemicSampler,
                 x: chex.Array,
                 num_samples: int) -> pd.DataFrame:
  """Generate posterior samples at x (not implemented for all posterior)."""
  data = []
  rng = hk.PRNGSequence(0)
  for seed in range(num_samples):
    net_out = enn_sampler(x, next(rng))
    y = net_out[:, 0]
    data.append(pd.DataFrame({'x': x[:, 0], 'y': y, 'seed': seed}))
  return pd.concat(data)



def plot_1d_regression(gp_model: testbed.GPRegression,