import numpy as np
from scipy.stats import norm
import tensorflow as tf
from matplotlib import pyplot as plt


class Generator(tf.keras.Model):

  def __init__(self):
    super(Generator, self).__init__()
    self.w = tf.Variable(1, dtype=tf.float32)
    self.b = tf.Variable(0, dtype=tf.float32)

  def call(self, z):
    x = self.w*z + self.b
    return x

class Discriminator(tf.keras.Model):

  def __init__(self, hidden_units=8):
    super(Discriminator, self).__init__()
    self.dense = tf.keras.layers.Dense(hidden_units)
    self.logits = tf.keras.layers.Dense(1, kernel_initializer=tf.keras.initializers.zeros())

  def call(self, x):
    x = tf.expand_dims(x, axis=-1)
    logits = self.logits(tf.nn.relu(self.dense(x)))
    logits = tf.squeeze(logits, axis=-1)
    p = 1 / (1 + tf.math.exp(-logits))
    return p

# parameters true distribution
mu = -4
sigma = 2

def visualize(G, D):
  interval = np.linspace(-10, 10, 100).astype(np.float32)
  d_values = D(interval)
  g_dist = norm.pdf(interval, loc=G.b.numpy(), scale=G.w.numpy())
  true_dist = norm.pdf(interval, loc=mu, scale=sigma)
  plt.plot(interval, true_dist, label="true_dist")
  plt.plot(interval, g_dist, label="G_dist")
  plt.plot(interval, d_values, label="D, p_true_data(x)")
  plt.legend()
  plt.show()

N = 32
x = np.random.normal(loc=mu, scale=sigma, size=N).astype(np.float32)
indices = np.array(range(N))

G = Generator()
D = Discriminator()

learning_rate = 0.1
D_optimizer = tf.keras.optimizers.Adam(learning_rate)
G_optimizer = tf.keras.optimizers.Adam(learning_rate)

batch_size = 16
critic_iters = 2
iterations = 100
plot_interval = 1

for iteration in range(iterations):

  if iteration % plot_interval == 0:
    visualize(G, D)

  # discriminator update
  for _ in range(critic_iters):
    # sample real data (from data distribution)
    np.random.shuffle(indices)
    real = x[indices[:batch_size]]
    z = tf.random.normal(shape=[batch_size])
    fake = G(z)
    with tf.GradientTape() as tape:
      loss_real = tf.reduce_mean(-tf.math.log(D(real)))
      loss_fake = tf.reduce_mean(-tf.math.log(1-D(fake)))
      D_loss = loss_real + loss_fake

    grads = tape.gradient(D_loss, D.trainable_variables)
    D_optimizer.apply_gradients(zip(grads, D.trainable_variables))

  # generator update
  z = tf.random.normal(shape=[batch_size])
  with tf.GradientTape() as tape:
    G_loss = tf.reduce_mean(tf.math.log(1-D(G(z))))

  grads = tape.gradient(G_loss, G.trainable_variables)
  G_optimizer.apply_gradients(zip(grads, G.trainable_variables))