import os
import torch
import pickle

class Model:

    def __init__(self, words):
        self.words = words
        # Create the vocabulary and the necessary mapprings
        self.chars = sorted(set(''.join(words)))
        self.stoi = {s:i+1 for i,s in enumerate(self.chars)} # mapping from string to index
        self.stoi['.'] = 0
        self.itos = {i:s for s,i in self.stoi.items()} # mapping from index to string
        self.vocab_size = len(self.chars) + 1 # 26 characters + 1 special token

    def train(self):
        # Create the bigram probability matrix
        N = torch.zeros((self.vocab_size, self.vocab_size), dtype=torch.int32)
        for w in self.words:
            chs = ['.'] + list(w) + ['.']
            for ch1, ch2 in zip(chs, chs[1:]):
                ix1, ix2 = self.stoi[ch1], self.stoi[ch2]
                N[ix1, ix2] += 1

        self.P = (N+1).float()
        self.P /= self.P.sum(1, keepdims=True)
        loss = self._evaluate_likelihood()
        print(f"Average negative log likelihood: {loss}")
    
    def generate(self, n_names):
        names = []
        # Generate a series of n_names tokens
        for i in range(n_names):
            out = []
            ix = 0
            while True:
                p = self.P[ix].float()
                ix = torch.multinomial(p, num_samples=1, replacement=True).item()
                if ix == 0:
                    break
                else:
                    out.append(self.itos[ix])
            names.append(''.join(out))
        return ', '.join(names)
    
    def save_model(self, file_path='.'):
        with open(os.path.join(file_path, 'model.pkl'), 'wb') as f:
            pickle.dump(self, f)

    
    def _evaluate_likelihood(self):
        log_likelihood = 0.0
        n = 0
        for w in self.words:
            chs = ['.'] + list(w) + ['.']
            for ch1, ch2 in zip(chs, chs[1:]):
                ix1 = self.stoi[ch1]
                ix2 = self.stoi[ch2]
                prob = self.P[ix1, ix2]
                logprob = torch.log(prob)
                log_likelihood += logprob
                n += 1
        return -log_likelihood/n
    
if __name__ == '__main__':
    # Read in the words
    words = open('names.txt', 'r').read().splitlines()
    # Instantiate the model class
    model = Model(words)
    # Train the model
    model.train()
    model.save_model()