from keras.models import Model, Sequential
from keras.layers import Dense, Input, LeakyReLU
from keras.optimizers import Adam
from keras.datasets import mnist
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt

(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 데이터 정규화: 픽셀 값을 -1 ~ 1 범위로 조정
x_test = (x_test.astype(np.float32) - 127.5) / 127.5

# 데이터 형태 변환: (28, 28) 이미지를 784 크기의 1차원 벡터로 변환
mnist_data = x_test.reshape(10000, 784)

print(mnist_data.shape)
# (10000, 784)

def create_generator():
  generator = Sequential()

  # 1. 모델의 입력 형태를 정의하는 Input 레이어를 추가합니다.
  # 100은 100개의 픽셀을 넣을 거라서 큰 의미는 없으며, 이 100개의 픽셀은 노이즈값으로, 100개의 픽셀값이 랜덤한 값을 가진다.
  generator.add(Input(shape=(100,)))

  # 첫 번째 은닉층은 256개의 뉴런으로 구성됨
  generator.add(Dense(units=256))
  # LeakyReLU 함수 모양을 보면 음수값은 특정한 기울기를 보이는데 여기서는 그 기울기값을 0.2로 설정
  generator.add(LeakyReLU(0.2))

  # 두 번째 은닉층은 512개의 뉴런으로 구성
  generator.add(Dense(units=512))
  generator.add(LeakyReLU(0.2))

  # 출력층의 활성화 함수는 tanh 으로 사용하고, 뉴런 수는 784개
  # tanh 은 출력값을 -1~1 사이로 맞춰줌
  # 784개인 이유는 MNIST 데이터셋 모습이 바로 28*28 개의 픽셀로 구성되어 있기 때문임
  # 1*784 처럼 생성자가 만드는 데이터 모습도 이와 같은 모양인 784개의 픽셀을 나열한 모습으로 나타내야 함
  # 추후 이 모양을 다시 28*28 형태로 나타내면 숫자와 같은 모습으로 나타남
  generator.add(Dense(units=784, activation='tanh'))
  return generator
g = create_generator()
g.summary()

def create_discriminator():
  discriminator = Sequential()
  # 784차원의 이미지 벡터(진짜 혹은 가짜)를 입력 받습니다.
  생성자가 만든 손글씨(784개의 픽셀로 구성된)를 넣기 때문이다.
  discriminator.add(Input(shape=(784,)))
  
  # 첫 번째 은닉층 512개의 뉴런으로 구성되어있다.
  discriminator.add(Dense(units=512))
  discriminator.add(LeakyReLU(0.2))

  discriminator.add(Dense(units=256))
  discriminator.add(LeakyReLU(0.2))

  # 최종 출력값은 1개이다. 판별자가 하는 것은 데이터의 진위 여부를 판단하는 것이다.
  # 따라서 진짜면 1, 가짜면 0이라는 숫자를 보여준다.
  discriminator.add(Dense(units=1, activation='sigmoid'))
  
  # 모델 컴파일
  # !! 여기서 컴파일하면 정상적으로 학습이 되지 않음
  # discriminator.compile(loss='binary_crossentropy', optimizer=Adam(learning_rate=0.0002, beta_1=0.5))
  return discriminator
d = create_discriminator()
# d.summary()

from keras.models import Model
from keras.layers import Input

# GAN 생성 함수
def create_gan(discriminator, generator):
  # 핵심: GAN 모델에서는 판별자가 학습하지 못하도록 가중치를 freeze 함
  discriminator.trainable = False
  
  # 생성적 적대 신경망인 gan 에 입력할 데이터 모습을 정함
  # 입력하는 값은 100개의 값으로 구성된 데이터이고, 콤마 뒤의 값이 비어있는 이유는 총 데이터의 개수를 넣기 위함임
  # 이렇게 값을 비워두면 실제 데이터 개수(여기서는 1만개)를 자동으로 넣어줌
  gan_input = Input(shape=(100,))
  
  # 생성자 신경망에 바로 윗줄에서 작성할 픽셀 100개의 값과 데이터의 전체 수(Input(shape=(100,))만큼 데이터를 넣음
  # 이 픽셀 100개는 노이즈값으로, x에는 생성자가 노이즈를 입력받아 만든 가짜 이미지들이 저장됨
  x = generator(gan_input)

  # 생성적 적대 신경망 gan 의 결과값 데이터 정의
  # 생성자가 만든 그림(x)을 보고 판별자가 판단한 결과, 판단 결과는 진짜 혹은 가짜 둘 중 하나로 나옴
  gan_output = discriminator(x)
  
  # 생성자와 판별자를 연결한 최종 생성적 적대 신경망인 gan 모델 정의
  # 입력값은 생성자 신경망이 만든 그림이고, 출력값은 판별자 신경망이 판단한 결과
  gan = Model(inputs=gan_input, outputs=gan_output)

  # GAN 모델 컴파일
  gan.compile(loss='binary_crossentropy', optimizer='adam')
  return gan

# !! 판별자를 동결시키고 GAN 모델 컴파일(생성자 훈련용)
d.trainable = False
gan = create_gan(d, g)
gan.summary()

# !! 판별자를 다시 학습 가능하게 만들고, 판별자 모델만 따로 컴파일(판별자 훈련용)
d.trainable = True
d.compile(loss='binary_crossentropy', optimizer=Adam(learning_rate=0.0002, beta_1=0.5))
d.summary()

# 결과 확인 함수
def plot_generated_images(generator):
  # 생성자에 넣어줄 노이즈값 생성
  noise = np.random.normal(loc=0, scale=1, size=[100,100])

  # 노이즈를 생성자 모델에 넣어 100개의 이미지 생성
  generated_images = generator.predict(noise)

  # 생성된 1차원 이미지(784,)를 2차원(28,28) 형태(그림 형태)로 변환
  generated_images = generated_images.reshape(100, 28, 28)

  # 10X10 격자에 생성된 이미지 100개를 시각화한다.
  plt.figure(figsize=(10, 10))
  for i in range(generated_images.shape[0]):
    plt.subplot(10, 10, i+1)
    plt.imshow(generated_images[i], interpolation='nearest')
    # 그림 이름은 넣지 않음
    plt.axis('off')
  
  # 지금까지 만든 그림을 화면에 보여줌 
  plt.tight_layout()
  #plt.show() # 이미지를 실제로 표시

# 한 번에 학습할 이미지 개수 (배치 크기)
batch_size = 128
# 총 훈련 반복 횟수 (에포크)
epochs = 5000

# 신경망 학습
# tqdm 은 5000 반복하는데 그 중 몇 번째 반복인지 눈에 잘 보이게 시각화함
for e in tqdm(range(epochs)):
  # --- 1단계: 판별자(Discriminator) 훈련 ---
  
  # 1-1. 진짜 이미지와 가짜 이미지 준비
  noise = np.random.normal(0, 1, [batch_size, 100])
  # 생성자 모델에 노이즈를 입력하여 생성자 신경망이 그림을 그린 후 결과를 generated_images 에 저장
  generated_images = g.predict(noise)
  
  # 실제 MNIST 데이터셋(1만개)에서 128개만 랜덤으로 추출(학습할 때마다 다양한 모양의 손글씨를 인공지능이 학습하도록 하기 위함임)
  # 첫 번째부터(now=0) MNIST 데이터셋 개수(high=mnist_data.shape[0])까지, 즉 1만 개의 데이터 중에서 배치 사이즈만큼 랜덤으로 추출
  image_batch = mnist_data[np.random.randint(low=0, high=mnist_data.shape[0], size=batch_size)]
  
  # 진짜 이미지와 가짜 이미지를 하나로 합침 (총 256개)
  # 이 데이터는 총 256개의 데이터로 되어있으며, 각 데이터에는 -1~1 사이의 값이 784개씩 들어있다.
  X = np.concatenate([image_batch, generated_images])
  
  # 1-2. 정답 레이블 생성(=판별자에 전달할 결과값 생성)
  # 판별자는 이것으로 그림이 진짜인지 가짜인지 확인할 수 있음
  # 앞에서 만든 데이터 개수(256개, 진짜와 가짜 각각 128개)만큼 결과값을 256개 만드는 때 이 때 값은 0으로 채움
  y_dis = np.zeros(2*batch_size)
  # 진짜 이미지(앞쪽 128개)의 정답은 1로, 가짜 이미지(뒤쪽 128개)의 정답은 0으로 설정
  y_dis[:batch_size] = 1
  
  # 1-3. 판별자 학습 진행
  # 처음에는 판별자가 먼저 학습할 수 있어야 함
  # 정답을 보고 학습하여 생성자가 만든 그림과 진짜 그림을 구별할 수 있어야 하기 때문임
  d.trainable = True # 판별자의 가중치를 업데이트할 수 있도록 설정
  d.train_on_batch(X, y_dis) # 판별자를 학습시킴, 입력값은 X 이고, 출력값은 y_dis 임

  # --- 2단계: 생성자(Generator) 훈련 ---

  # 2-1. 생성자를 속이기 위한 데이터 준비
  noise = np.random.normal(0, 1, [batch_size, 100])
  # 생성자의 목표는 판별자가 '진짜(1)'라고 착각하게 만드는 것
  y_gen = np.ones(batch_size)
  
  # 2-2. 생성자 학습 진행 (판별자는 동결)
  # 판별자가 더 이상 학습할 수 없도록 한다. 판별자는 이제 학습하는 것이 아니라 생성자가 만든 그림이 진짜인지 판별하는 역할만 한다.
  d.trainable = False
  # gan 에 노이즈를 입력으로 넣고, 출력값으로 모두 다 진짜(y_gen의 값은 1인 상태)를 출력으로 넣어서 학습시킨다.
  gan.train_on_batch(noise, y_gen)

    # --- 중간 결과 확인 ---
  # 첫 번째 에포크와 1000번째, 2000번째.. 일 때 생성자가 만든 그림을 출력함
  if e % 1000 == 0:
    print(f"Epoch {e}")
    plot_generated_images(g)

from keras.models import Model, Sequential
from keras.layers import Dense, Input, LeakyReLU
from keras.optimizers import Adam
from keras.datasets import mnist
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt

(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 데이터 정규화: 픽셀 값을 -1 ~ 1 범위로 조정
x_test = (x_test.astype(np.float32) - 127.5) / 127.5

# 데이터 형태 변환: (28, 28) 이미지를 784 크기의 1차원 벡터로 변환
mnist_data = x_test.reshape(10000, 784)
print(mnist_data.shape)
len(mnist_data)

def create_generator():
  generator = Sequential()

  # 1. 모델의 입력 형태를 정의하는 Input 레이어를 추가합니다.
  # input_dim=100은 shape=(100,)과 동일합니다.
  # 100은 100개의 픽셀을 넣을 거라서 큰 의미는 없으며, 이 100개의 픽셀은 노이즈값으로, 100개의 픽셀값이 랜덤한 값을 가진다.
  generator.add(Input(shape=(100,)))

  # 첫 번째 은닉층은 256개의 뉴런으로 구성됨
  generator.add(Dense(units=256))
  # LeakyReLU 함수 모양을 보면 음수값은 특정한 기울기를 보이는데 여기서는 그 기울기값을 0.2로 설정
  generator.add(LeakyReLU(0.2))

  # 두 번째 은닉층은 512개의 뉴런으로 구성
  generator.add(Dense(units=512))
  generator.add(LeakyReLU(0.2))

  # 출력층의 활성화 함수는 tanh 으로 사용하고, 뉴런 수는 784개
  # 784개인 이유는 MNIST 데이터셋 모습이 바로 28*18 개의 픽셀로 구성되어 있기 때문임
  # 1*784 처럼 생성자가 만드는 데이터 모습도 이와 같은 모양인 784개의 픽셀을 나열한 모습으로 나타내야 함
  # 추후 이 모양을 다시 28*28 형태로 나타내면 숫자와 같은 모습으로 나타남
  generator.add(Dense(units=784, activation='tanh'))
  return generator
g = create_generator()
g.summary()

def create_discriminator():
  discriminator = Sequential()
  # 이 모델에 입력하는 값은 784이다. 생성자가 만든 손글씨(784개의 픽셀로 구성된)를 넣기 때문이다.
  discriminator.add(Input(shape=(784,)))
  # 첫 번재 은닉층 512개의 뉴런으로 구성되어있다.
  discriminator.add(Dense(units=512))
  discriminator.add(LeakyReLU(0.2))

  discriminator.add(Dense(units=256))
  discriminator.add(LeakyReLU(0.2))

  # 최종 출력값은 1개이다. 판별자가 하는 것은 데이터의 진위 여부를 판단하는 것이다.
  # 따라서 진짜면 1, 가짜면 0이라는 숫자를 보여준다.
  discriminator.add(Dense(units=1, activation='sigmoid'))

# !!컴파일
  #discriminator.compile(loss='binary_crossentropy', optimizer=Adam(learning_rate=0.0002, beta_1=0.5))
  return discriminator
d = create_discriminator()
#d.summary()


# GAN 생성 함수
def create_gan(discriminator, generator):
  # 판별자가 학습하지 못하도록 막음
  discriminator.trainable = False
  # 생성적 적대 신경망인 gan 에 입력할 데이터 모습을 정함
  # 입력하는 값은 100개의 값으로 구성된 데이터이고, 콤마 뒤의 값이 비어있는 이유는 총 데이터의 개수를 넣기 위함임
  # 이렇게 값을 비워두면 실제 데이터 개수(여기서는 1만개)를 자동으로 넣어줌
  gan_input = Input(shape=(100,))
  # 생성자 신경망에 바로 윗줄에서 작성할 픽셀 100개의 값과 데이터의 전체 수(Input(shape=(100,))만큼 데이터를 넣음
  # 이 픽셀 100개는 노이즈값임
  # x에는 생성자가 만든 새로운 그림들이 저장됨
  x = generator(gan_input)

  # 생성적 적대 신경망 gan 의 결과값 데이터 정의
  # 생성자가 만든 그림(x)을 보고 판별자가 판단한 결과, 판단 결과는 진짜 혹은 가짜 둘 중 하나로 나옴
  gan_output = discriminator(x)

  # 생성적 적대 신경망인 gan 모델 설계
  # 입력값은 생성자 신경망이 만든 그림이고, 출력값은 판별자 신경망이 판단한 결과
  gan = Model(inputs=gan_input, outputs=gan_output)

  # 신경망의 오차값을 줄이는 방법
  gan.compile(loss='binary_crossentropy', optimizer='adam')
  return gan
# !! 판별자를 동결시키고 GAN 모델 컴파일(생성자 훈련용)
d.trainable = False
gan = create_gan(d, g)
gan.summary()

# !! 판별자를 다시 학습 가능하게 만들고, 판별자 모델만 따로 컴파일(판별자 훈련용)
d.trainable = True
d.compile(loss='binary_crossentropy', optimizer=Adam(learning_rate=0.0002, beta_1=0.5))
d.summary()

# 결과 확인 함수
def plot_generated_images(generator):
  # 생성자에 넣어줄 노이즈값 생성
  # 이 때 균일한 값을 생성할 수 있도록 넘파이의 랜덤값 생성 라이브러리 중 정규 분포 함수 이용
  # 첫 번째 0은 평균이 0이라는 의미이고,
  # 두 번째 1은 평균에서 1만큼씩 떨어져있는 값(즉, -1~1사이의 값)
  # 세 번째 값은 노이즈 100개를 생성한다.
  # 각각의 노이즈는 숫자 100개씩 구성되어 있다는 것으로, 이 함수를 호출할 때마다 100개의 그림을 그려달라는 의미임
  noise = np.random.normal(loc=0, scale=1, size=[100,100])

  # generator 도 신경망 모델이므로 predict() 함수 이용 가능
  # 위의 노이즈를 신경망에 넣어서 값을 예측하는 기능
  generated_images = generator.predict(noise)

  # 이 형태를 그림 형태로 바꾸어준다.
  generated_images = generated_images.reshape(100, 28, 28)

  # 그림 크기를 정한다.
  plt.figure(figsize=(10, 10))

  # 100개의 그림을 그려준다.
  for i in range(generated_images.shape[0]):
    # 그림 위치를 정한다. ??? 각각의 인자값이 의미하는 거
    plt.subplot(10, 10, i+1)
    # imshow() 함수는 이미지를 출력함
    plt.imshow(generated_images[i], interpolation='nearest')
    # 그림 이름은 넣지 않음
    plt.axis('off')
  
  # 지금까지 만든 그림을 화면에 보여줌 
  plt.tight_layout()
  plt.show()

  # 한 번에 학습할 이미지 개수 (배치 크기)
batch_size = 128
# 총 훈련 반복 횟수 (에포크)
epochs = 5000

# 신경망 학습
# tqdm 은 5000 반복하는데 그 중 몇 번째 반복인지 눈에 잘 보이게 시각화함
# 신경망 학습
for e in tqdm(range(epochs)):
    # --- 1단계: 판별자(Discriminator) 훈련 ---
    
    # 1-1. 진짜 이미지와 가짜 이미지 준비
    noise = np.random.normal(0, 1, [batch_size, 100])
    generated_images = g.predict(noise)
    image_batch = mnist_data[np.random.randint(low=0, high=mnist_data.shape[0], size=batch_size)]
    X = np.concatenate([image_batch, generated_images])
    y_dis = np.zeros(2 * batch_size)
    y_dis[:batch_size] = 1
    
    # 1-2. 판별자 학습 진행
    # 판별자를 훈련시키기 직전에 학습 가능하도록 설정합니다.
    d.trainable = True
    d.train_on_batch(X, y_dis)

    # --- 2단계: 생성자(Generator) 훈련 ---

    # 2-1. 생성자를 속이기 위한 데이터 준비
    noise = np.random.normal(0, 1, [batch_size, 100])
    y_gen = np.ones(batch_size)
    
    # 2-2. 생성자 학습 진행
    # 생성자를 훈련시킬 때는 판별자의 가중치가 변하면 안 되므로 동결합니다.
    d.trainable = False
    gan.train_on_batch(noise, y_gen)

    # --- 중간 결과 확인 ---
    if e % 1000 == 0:
        print(f"Epoch {e}")
        plot_generated_images(g)