import re, collections
def get_stats(vocab):
	pairs = collections.defaultdict(int)
	for word, freq in vocab.items():
		symbols = word.split()
	for i in range(len(symbols)-1):
		pairs[symbols[i],symbols[i+1]] += freq
		return pairs

def merge_vocab(pair, v_in):
	v_out = {}
	bigram = re.escape(' '.join(pair))
	p = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
	for word in v_in:
		w_out = p.sub(''.join(pair), word)
		v_out[w_out] = v_in[word]
		return v_out

vocab = {'l o w </w>' : 5, 'l o w e r </w>' : 2,
'n e w e s t </w>':6, 'w i d e s t </w>':3}
num_merges = 10
for i in range(num_merges):
	pairs = get_stats(vocab)
	best = max(pairs, key=pairs.get)
	vocab = merge_vocab(best, vocab)
	print(best)

from typing import List

import collections, re

# 단어 끝을 나타내는 문자
WORD_END = '_'

def build_bpe(
    corpus: List[str],
    max_vocab_size: int
) -> List[int]:
  
    #단어끝에 언더스코어 붙여주기
    #단어 각 알파벳으로 바꿔주기, 단어의 등장빈도를 Counter로 세어서 가중치로 이용하기
    words_freq = collections.Counter(corpus)
    words = {}
    for word, freq in words_freq.items():
      words[' '.join(list(word)) + ' ' + WORD_END] = freq
     #words에 corpus의 각 단어가 나온 횟수를 Count해서 dict로 저장
    print(words)

    #각 단어에서 나올 수 있는 쌍에 대해서 freq를 가중치로 주는 dict인 pair를 리턴
    def get_stats(voca):
      pairs = collections.defaultdict(int)
      for word, freq in voca.items():
        symbols = word.split()
        for i in range(len(symbols)-1):
        #pairs에 freq만큼 가중치를 더해줌
          pairs[symbols[i],symbols[i+1]] += freq
      return pairs
		
    def merge_vocab(pair, v_in):
      v_out = {}
      #token 두개를 공백으로 연결
      bigram = re.escape(' '.join(pair))
      #정규표현식 패턴을 컴파일해서 p에 저장
      #단어 중간에 있는 경우를 배재하기 위해 단어의 시작과 끝에서만 탐색
      p = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
      for word in v_in:
      #v_in의 모든 단어에 대해서 두 token을 합쳐서 w_out에 저장
        w_out = p.sub(''.join(pair), word)
        v_out[w_out] = v_in[word]
      return v_out

    # id2token: List[str] = None

    id2token = set()
    for word in words:
        id2token.update(word.split())
    id2token = list(id2token)
#len(id2token)이 max_vocab_size보다 작을때까지 반복
    while len(id2token) < max_vocab_size:
      pairs = get_stats(words)
      if not pairs:
        break
      best = max(pairs, key = pairs.get)
      words = merge_vocab(best, words)
      # id2token = list(words.keys())
      id2token.append(''.join(best))
      print(id2token)
    #함수 조건에 맞게 길이가 긴 token순으로 정렬
      id2token.sort(key=len, reverse=True)
    ### END YOUR CODE

    return id2token

'''Vocab: bcde ab cd bc de a b c d e _
abcde ==> a bcde _ '''

def encode(
    sentence: str,
    id2token: List[str]
) -> List[int]:
    '''
    1. 문장에서 공백을 기준으로 단어를 하나씩 나눠서 리스트로 만들기
    2. 리스트에서 for문으로 단어 하나씩 가져오기
    3. id2token에서 인덱스 0부터 값 하나씩  가져와서 있는지 확인
      일치하면 해당 단어에서 voca부분 없애고 3번 반복
      일치하지 않으면 다음 voca에서 찾기
    4. 해당단어 끝났으면 len(id2token)-1인 WORD_END를 token_ids에 넣기
    5.return token_ids
    '''
    #sentence에서 공백을 기준으로 단어 나눠서 리스트로 만들기
    token_ids = []
    words = sentence.split(" ")
    print(words)

    for word in words:
      #word에 알파벳이 모두 없어지고, 모두 id와 comma로 바뀔때까지 while문 돌기
      while any(char.isalpha() for char in word) == True:
        #가장 길이가 긴 토큰부터 검사
        for token in id2token:
          token_length=len(token)
          if token in word:
            word = word.replace(token,',' + str(id2token.index(token)))
            # token_ids.append(id2token.index(token))
            print(word)
            break
          else: #일치하지 않으면 다음 token에서 찾기
            continue
            print(word)
      word = word.split(",")
      word.pop(0)
      word = list(map(int,word))
      word.append(len(id2token)-1)
      print(word)
      token_ids.extend(word) #extend는 append와 다르게 요소를 개별적으로 추가함.
    print(token_ids)
    ### END YOUR CODE

    return token_ids

def decode(
    token_ids: List[int],
    id2token: List[str]
) -> str:
    sentence = ""
    # token을 공백을 기준으로 나눠서 리스트로 만들기
    # token_ids = token_ids.split(" ")
    #token(숫자) 하나씩 불러와서 token으로 바꾼 후 sentence에 추가해주기
    for id in token_ids:
      sentence+=id2token[id]
    #언더스코어는 공백으로 바꾸기
    #언더스코어를 기준으로 리스트로 만들고, 마지막 언더스코어는 공백으로 남았기 때문에 pop()으로 빼준 뒤 다시 join으로 합치기
    sentence = sentence.split(WORD_END)
    sentence.pop()
    sentence = ' '.join(sentence)

    print(sentence)
    ### END YOUR CODE
    return sentence

'''[ 196 62 20 6 ] ==> [ I_ li ke_ it_ ]==> "I_like_it_" ==> "I like it " ==> "I like it" '''