# 필요한 파일 임포트
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
%matplotlib inline

!wget https://raw.githubusercontent.com/karpathy/makemore/master/names.txt # 데이터

words = open('names.txt', 'r').read().splitlines()
len(words) # 32033

chars = sorted(list(set(''.join(words))))
stoi = {s:i+1 for i,s in enumerate(chars)}
stoi['.'] = 0
itos = {i:s for s,i in stoi.items()}
print(itos)
# {1: 'a', 2: 'b', 3: 'c', 4: 'd', 5: 'e', 6: 'f'...

# build the dataset
block_size = 3 # 다음 문자 예측하는 데 필요한 문자 수에 대한 context 길이

def build_dataset(words):
  X, Y = [], []
  # x는 입력
  # y는 레이블
  for w in words:

    #print(w)
    context = [0] * block_size
    for ch in w + '.':
      ix = stoi[ch]
      X.append(context)
      # 현재 context 입력에 추가
      Y.append(ix)
      # 현재 인덱스를 레이블에 추가
      #print(''.join(itos[i] for i in context), '--->', itos[ix])
      context = context[1:] + [ix] # crop and append

  X = torch.tensor(X)
  Y = torch.tensor(Y)
  print(X.shape, Y.shape)
  # torch.Size([22799, 3]) torch.Size([22799])
  return X, Y

import random
random.seed(42)
random.shuffle(words)
n1 = int(0.8*len(words))
n2 = int(0.9*len(words))

Xtr, Ytr = build_dataset(words[:n1])
Xdev, Ydev = build_dataset(words[n1:n2])
Xte, Yte = build_dataset(words[n2:])
"""
torch.Size([182580, 3]) torch.Size([182580])
torch.Size([22767, 3]) torch.Size([22767])
torch.Size([22799, 3]) torch.Size([22799])
"""

W1 = torch.randn((6, 100))
b1 = torch.randn(100)
#emb @ W 불가, (32,3,2) @ (6,100) 을 할 순 없음
#3개의 vector를 concatenate 

# x = torch.cat([emb[:, 0, :], emb[:, 1, :], emb[:, 2, :]], dim=1)
# 블록 크기를 변경할 경우 귀찮아짐

# x = torch.cat(torch.unbind(emb, dim=1), dim=1)
# 블록 크기 상관 x
# 새로운 메모리 공간 필요 -> 비효율적

x = emb.view(emb.shape[0], -1)  
# 가장 효율적임
# emb.shape[0] = 32

h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
h

"""tensor([[-0.9593, -0.9898,  0.9987,  ..., -0.6509, -0.9997,  0.9978],
        [-0.9995, -0.9882,  0.7896,  ...,  0.5566, -0.7900,  0.8329],
        [-0.9769,  0.9975, -0.4099,  ..., -0.0782, -0.9998,  0.0702],
        ...,
        [ 0.9977, -0.4416, -0.5036,  ..., -0.1040,  0.9661, -0.2288],
        [ 0.9991, -0.7121, -0.6554,  ..., -0.0184,  0.6605,  0.1505],
"""
h.shape # torch.Size([32, 100])

W2 = torch.randn((100, 27)) # 100: 은닉층 크기, 27: char(a~z)
b2 = torch.randn(27)

logits = h @ W2 + b2 
logits.shape
# torch.Size([32, 27])
counts = torch.exp(logits)
probs = counts / counts.sum(dim=1, keepdim=True)
# 소프트맥스 구현

g = torch.Generator().manual_seed(2147483647) # 재현 가능하게 생성기 사용
C = torch.randn((27, 10), generator=g)
W1 = torch.randn((30, 200), generator=g)
b1 = torch.randn(200, generator=g)
W2 = torch.randn((200, 27), generator=g)
b2 = torch.randn(27, generator=g)
parameters = [C, W1, b1, W2, b2] # 모든 매개변수를 단일 매개변수 목록으로 클러스터링

sum(p.nelement() for p in parameters) #  현재 파라미터 수: 11897

emb = C[X]
h = torch.tanh(emb.view(-1,6) @ W1 + b1) # (32,100)
logits = h @ W2 + b2 # (32,27)

#counts = logits.exp()
#prob = counts/ counts.sum(1, keepdims=True)
#loss = -prob[torch.arange(32),Y].log().mean()
#loss 
# tensor(17,7697)

F.cross_entropy(logits,Y) # tensor(17,7697)

for p in parameters:
  p.requires_grad = True
lre = torch.linspace(-3, 0, 1000)
lrs = 10**lre

lri = []
lossi = []
stepi = []

for i in range(200000):
  
  # minibatch construct
  ix = torch.randint(0, Xtr.shape[0], (32,)) # 228146 개 임으로 이중에 32개만 minibatch로 고르는것 -> 훨씬 빠르게 반복
  
  # forward pass
  emb = C[Xtr[ix]] # (32, 3, 2)
  h = torch.tanh(emb.view(-1, 30) @ W1 + b1) # (32, 100)
  logits = h @ W2 + b2 # (32, 27)
  loss = F.cross_entropy(logits, Ytr[ix])
  #print(loss.item())
  
  # backward pass
  for p in parameters:
    p.grad = None
  loss.backward()
  
  # update
  #lr = lrs[i]
  lr = 0.1 if i < 100000 else 0.01
  for p in parameters:
    p.data += -lr * p.grad
  # learning rate 또한 training 이 진행될수록 점점 감소

  # track stats
  #lri.append(lre[i])
  stepi.append(i)
  lossi.append(loss.log10().item()) # lossi에 log를 씌워서 hockey stick 을 방지

#print(loss.item())

plt.plot(stepi, lossi)

# training split(매개변수 훈련), dev/validation split(하이퍼파라미터 훈련), test split(모델 성능 평가) 분할

# training loss 
emb = C[Xtr] # (32, 3, 2)
h = torch.tanh(emb.view(-1, 30) @ W1 + b1) # (32, 100)
logits = h @ W2 + b2 # (32, 27)
loss = F.cross_entropy(logits, Ytr)
loss # tensor(2.1335, grad_fn=<NllLossBackward0>)

# validation loss
emb = C[Xdev] # (32, 3, 2)
h = torch.tanh(emb.view(-1, 30) @ W1 + b1) # (32, 100)
logits = h @ W2 + b2 # (32, 27)
loss = F.cross_entropy(logits, Ydev)
loss # tensor(2.1635, grad_fn=<NllLossBackward0>)

# test loss
emb = C[Xte] # (32, 3, 2)
h = torch.tanh(emb.view(-1, 30) @ W1 + b1) # (32, 100)
logits = h @ W2 + b2 # (32, 27)
loss = F.cross_entropy(logits, Yte)
loss # tensor(2.1637, grad_fn=<NllLossBackward0>)

# visualize dimensions 0 and 1 of the embedding matrix C for all characters
plt.figure(figsize=(8,8))
plt.scatter(C[:,0].data, C[:,1].data, s=200)
for i in range(C.shape[0]):
    plt.text(C[i,0].item(), C[i,1].item(), itos[i], ha="center", va="center", color='white')
plt.grid('minor')

# sample from the model
g = torch.Generator().manual_seed(2147483647 + 10)

for _ in range(20):
    
    out = []
    context = [0] * block_size # initialize with all ...
    while True:
      emb = C[torch.tensor([context])] # (1,block_size,d)
      h = torch.tanh(emb.view(1, -1) @ W1 + b1)
      logits = h @ W2 + b2
      probs = F.softmax(logits, dim=1)
      ix = torch.multinomial(probs, num_samples=1, generator=g).item()
      context = context[1:] + [ix]
      out.append(ix)
      if ix == 0:
        break
    
    print(''.join(itos[i] for i in out))
    """
    carmah.
    ambril.
    khi.
    mili.
    thylahnanden.
    jazhnen.
    deliah.
    jareei.
    ner.
    kentzeriiv.
    kaleigh.
    ham.
    jory.
    quint.
    shon.
    marian.
    quinteron.
    jarynn.
    jaxen.
    dus.
    """