import os
import json
import random
import time
import streamlit as st
import re
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer

MODEL_FILE = r'bt_8_LAYERs_100_DATA_PCT_768_EMBD_DIM_epoch_10.pt' ##place model file in same directory as app.py

# Better Transformer Class –––––––––––––––––––––––––––––––––––––––––––––––

class MLP(nn.Module):
    def __init__(self, n_embd, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(), # replaced ReLU
            nn.Dropout(p=dropout),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        return self.net(x)

class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head, seq_length, dropout=0.1):
        super().__init__()

        self.n_embd = n_embd
        self.n_head = n_head
        self.head_dim = n_embd // n_head # Dimension of each head's key, query, and value
        assert self.head_dim * n_head == self.n_embd, "n_embd must be divisible by n_head"
        self.seq_length = seq_length
        self.drop = nn.Dropout(p=dropout)

        self.query = nn.Linear(n_embd, n_embd, bias=False)
        self.key = nn.Linear(n_embd, n_embd, bias=False)
        self.value = nn.Linear(n_embd, n_embd, bias=False)
        self.out = nn.Linear(n_embd, n_embd, bias=False) # multi-head combining weight matrix

    def split_heads(self, x):
        B, S, D = x.size()
        # split dimension into n_head * head_dim, then transpose the sequence length w/ n_head
        # output: [B, n_head, S, head_dim]
        return x.view(B, S, self.n_head, self.head_dim).transpose(1, 2)

    def combine_heads(self, x):
        # use permute or transpose to reverse
        # taking a view earlier may produce a non-contiguous tensor, so we convert back because view needs a contiguous input
        B, _, S, head_dim = x.size() # _ is n_head which we will merge
        # output: [B, S, n_embd]
        return x.transpose(1, 2).contiguous().view(B, S, self.n_embd)

    def scaled_dot_product(self, q, k, v, dropout, mask=None):
        # q,k,v are [B, n_head, S, head_dim]
        # the key transpose sets up batch multiplication s.t. wei = [B, n_head, S, S]
        wei = q @ k.transpose(-2,-1) / np.sqrt(self.head_dim)
        # mask is [B, 1, S, S], so simply broadcasted across each head and works as expected
        if mask is not None:
          wei = wei.masked_fill(mask, float('-inf'))
        wei = dropout(F.softmax(wei, dim=-1))
        out = wei @ v
        return out

    def forward(self, x, mask=None):
        # x: (B, S, n_embd)
        # Step 1 and 2: Project full query, key, value, then split via reshaping
        q = self.split_heads(self.query(x))
        k = self.split_heads(self.key(x))
        v = self.split_heads(self.value(x))

        # Step 3: Compute scaled dot-product attention with causal mask
        # not done. should use generate_mask
        attn = self.scaled_dot_product(q, k, v, self.drop, mask)

        # Step 4 and 5: Concatenate attention scores, return projected output matrix
        out = self.out(self.combine_heads(attn)) # (B, S, n_embd)
        return out

class Block(nn.Module):
    def __init__(self, n_embd, n_head, seq_length, dropout=0.1):
        super().__init__()
        self.sa = MultiHeadAttention(n_embd, n_head, seq_length, dropout)
        self.mlp = MLP(n_embd, dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        # experimentally, apply layer norm before attention/MLP
        self.drop = nn.Dropout(p=dropout)

    def forward(self, x, mask):
        # residual connection (stream)
        x = x + self.drop(self.sa(self.ln1(x), mask))
        x = x + self.drop(self.mlp(self.ln2(x)))
        return x

class PositionalEncoding(nn.Module):
  """
  Formula taken from the original Transformer paper:
  PE(pos, 2i (even)) = sin(pos/(10000^{2i/d_model}))
  PE(pos, 2i+1 (odd)) = cos(pos/(10000^{2i/d_model}))

  See reference for more details:
  https://kikaben.com/transformers-positional-encoding/
  """
  def __init__(self, d_model, max_len):
      # just set d_model = n_embd and max_len = seq_len
      super().__init__()

      position = torch.arange(max_len).unsqueeze(1) # [max_len, 1]
      divisor = torch.exp(torch.arange(0, d_model, 2) * (-np.log(10000.0) / d_model)) # [d_model / 2, half for each of sin and cos]
      pe = torch.zeros(max_len, d_model)
      pe[:, 0::2] = torch.sin(position * divisor) # 0 for second dim or :?
      pe[:, 1::2] = torch.cos(position * divisor)
      self.register_buffer('pe', pe) # result: self.pe = [max_len, d_model], mapping each token index to a vector of length d_model as desired

  def forward(self, x):
      # x = torch.arange(seq_length) has shape [seq_length], so x.size(0) extracts it, then we index self.pe for the first seq_length mappings
      # note we do not add the positional embeddings to x itself yet, we simply return them
      # output = (seq_length, d_model=n_embd)
      return self.pe[:x.size(0)]

class BetterTransformer(nn.Module):
    def __init__(self, vocab_size, seq_length, n_embd, n_head, n_layer, pad_idx, eos_token_id, device, dropout=0.1):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=pad_idx)
        # we need to make sure the embedding ignores the padding token right?
        self.position_embedding = PositionalEncoding(n_embd, seq_length)
        self.blocks = nn.Sequential(*[Block(n_embd,
                                            n_head,
                                            seq_length,
                                            dropout) for _ in range(n_layer)])
        self.lm_head = nn.Linear(n_embd, vocab_size)
        self.drop = nn.Dropout(dropout)
        self.seq_length = seq_length
        self.pad_idx = pad_idx
        self.eos_token_id = eos_token_id
        self.device = device
        self.init_params()

    # optional weight initialization (e.g. Xavier uniform)
    def init_params(self, default_initialization=False):
        if not default_initialization:
            for name, p in self.named_parameters():
                if p.dim() > 1:
                    nn.init.xavier_uniform_(p)

    def get_causal_mask(self, x):
        """
        Generates causal mask for decoding
        """
        seq_len = x.size(-1) # x = (batch_size x seq_len)
        attn_shape = (1, seq_len, seq_len)
        subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8') # k = 1 shifts the diagonal, so that the main diagonal gets 0's
        return (torch.from_numpy(subsequent_mask) == 0).to(self.device) # (1, seq_len x seq_len)
        # True along main diagonal + below, False elsewhere

    def get_pad_mask(self, x, pad_idx):
        """
        Generates padding mask
        """
        return (x != pad_idx).unsqueeze(1).unsqueeze(-2).to(self.device)
        # (batch_size x 1 x 1 x seq_len)

    def forward(self, x, targets=None):

        # should alr be int64 tokens but explicit cast in case
        x = x.to(torch.int64)
        B, S = x.shape

        # get mask
        mask = self.get_pad_mask(x, self.pad_idx) & self.get_causal_mask(x).to(self.device)
        # mask = (batch_size x 1 x seq_len x seq_len)

        tok_emb = self.token_embedding(x)
        pos_emb = self.position_embedding(torch.arange(S))
        x = self.drop(tok_emb + pos_emb)
        # (B, S, n_embd)
        for block in self.blocks:
            x = block(x, ~mask) # (batch_size, seq_length, n_embd)
        # negate mask to fill originally False values with -inf later
        logits = self.lm_head(x) # (batch_size, seq_length, vocab_size)

        # this code assumes teacher forcing——for each text of seq length S we have S autoregressive predictions,
        # thus we have B*S logits and B*S targets
        if targets is None:
            loss = None
        else:
            B, S, C = logits.shape
            logits = logits.view(B*S, C)
            targets = targets.view(B*S)
            loss = F.cross_entropy(logits, targets, ignore_index=self.pad_idx)
            # we need to make sure loss ignores the padding token right?
            # this helps it avoid wasting compute on learning PAD -> PAD, etc.

        return logits, loss


    def generate(self, input_ids, method='multinomial',
                 max_new_tokens=1000, temp=None,
                 num_beams=None, p_nucleus=None, k=None):

        # TODO: see Huggingface's .generate() function
        # https://huggingface.co/transformers/v3.4.0/_modules/transformers/generation_utils.html

        if method == 'temperature':
            assert (temp is not None) and (0 < temp) and (temp <= 1)
        # if method == 'num_beams':
        #     assert isinstance(num_beams, int) and (num_beams) > 0 and (num_beams) < 100
        if method == 'top-k':
            assert isinstance(k, int) and (k > 0)

        # input_ids begins as (batch_size, seq_length)

        for _ in range(max_new_tokens):
            if method in ['multinomial', 'temperature', 'greedy', 'nucleus', 'top-k']:
                # i) Truncate to the most recent `max length` tokens
                text_cond = input_ids[:, -self.seq_length:]
                # ii) Retrieve predictions
                logits, loss = self(text_cond) # no loss because no targets ofc
                # model output: (batch_size, seq_length, vocab_size)
                # iii) Find last token logits of each
                logits = logits[:, -1, :] # (batch_size, vocab_size)

                # aside: if temperature sampling, divide logits by temp before applying softmax
                if method == 'temperature':
                    logits = logits / temp

                # iv) Take softmax along each
                probs = F.softmax(logits, dim=-1)

                # v) Sample next token depending on method
                if method == 'greedy':
                    next_idx = probs.argmax(dim=-1).unsqueeze(-1)

                elif method in ['multinomial', 'temperature', 'nucleus', 'top-k']:
                    if method == 'nucleus':
                        assert p_nucleus is not None and (0 < p_nucleus) and (p_nucleus <= 1)

                        sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
                        prob_cumsum = sorted_probs.cumsum(dim=-1)
                        idx_remove = prob_cumsum > p_nucleus
                        # shift one right to ensure the first token is above the threshold
                        idx_remove[..., 1:] = idx_remove[..., :-1].clone()
                        idx_remove[..., 0] = False
                        # retrieve original indices by reverse-sorting
                        remove_mask = idx_remove.gather(dim=-1,
                                          index=sorted_idx.argsort(dim=-1))
                        # ^ specifically, we do this by first argsorting the indices which were returned from argsort. this is crazy y'all
                        # you can show that this returns indices that when used to subset a sorted array, returns the original array in unsorted order
                        # https://stackoverflow.com/questions/52127723/pytorch-better-way-to-get-back-original-tensor-order-after-torch-sort
                        # torch.gather is how we apply a multi-dimensional index
                        # https://stackoverflow.com/questions/50999977/what-does-the-gather-function-do-in-pytorch-in-layman-terms
                        probs[remove_mask] = 0

                    if method == 'top-k':
                        remove_mask = probs < torch.topk(probs, k).values[..., -1, None] # the topk returns (B, 1), leaving only the
                        # kth largest probs (i.e. the cutoff value for each). Then mask is same size as probs (B, vocab_size)
                        probs[remove_mask] = 0

                    # Sample probabilistically via scores
                    next_idx = torch.multinomial(probs, num_samples=1) # (batch_size, 1)

                # vi) Autoregressively append to input_text
                input_ids = torch.cat((input_ids, next_idx), dim=-1)
                # end prematurely if <EOS> generated
                if next_idx == self.eos_token_id:
                  break
                # now input_text = (batch_size, seq_length + 1)

        return input_ids

# END OF Better Transformer Class –––––––––––––––––––––––––––––––––––––––––––––––

def set_seed(seed = 42):
    random.seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    # torch.cuda.manual_seed_all(seed) # if multi-GPU
    torch.backends.cudnn.deterministic=True # only applies to CUDA convolution operations
    torch.backends.cudnn.benchmark = False
    # usually CuDNN has heuristics as to which algorithm to pick. cudnn.benchmark benchmarks several algorithms and picks the fastest
    # often helpful if your input shapes are fixed and not changing a lot during training
    # however, this means it may pick a different algorithm even when the deterministic flag is set.
    # As such it is good practice to turn off cudnn.benchmark when turning on cudnn.deterministic

def load_tokenizer(device):
    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    EMPTY_TOKENS = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)
    return tokenizer, EMPTY_TOKENS


def load_big_model(tokenizer, device):
    ## Model architecture
    set_seed(42)
    N_HEAD = 16
    N_LAYER = 8
    N_EMBD = 768
    VOCAB_SIZE = 50258
    SEQ_LENGTH = 384

    model = BetterTransformer(VOCAB_SIZE, SEQ_LENGTH, N_EMBD, N_HEAD, N_LAYER, tokenizer.pad_token_id, tokenizer.eos_token_id, device=device)
    model.init_params()
    path = MODEL_FILE
    model.load_state_dict(torch.load(path, map_location=device)["model_state_dict"])

    return model
    
def generate(model, tokenizer, device, method=None, k=None, 
    p_nucleus=None, temp=None, max_new_tokens=None, cond="", deterministic=None):
    """
    Wrapper for generating text using the specified model. Generates unconditionally if cond=None.

    Inputs:
      -model: Decoder model to be used for text generation
      -tokenizer: Compatible tokenizer
      -device: Device of model (CPU/CUDA)
      -method (str): Decoding method for text generation ('multinomial', 'temperature', 'greedy', 'nucleus', or 'top-k')
      -k (int): Positive integer for top-k logits to sample if top-k decoding
      -p_nucleus (float/int): Cumulative probability cutoff if nucleus/top-p decoding
      -temp (float/int): Temperature if temperature decoding
      -max_new_tokens (int): Maximum number of tokens to generate
      -cond (str=None): If provided, will serve as conditional prompt for text generation
      -deterministic (int): If deterministic, uses the specified seed for model generation
    Returns:
      -res (str): Generated text string
    """
    assert method in ['multinomial', 'temperature', 'greedy', 'nucleus', 'top-k'], \
        "method must be 'multinomial', 'temperature', 'greedy', 'nucleus', or 'top-k'"

    #if method == 'temperature':
    #    assert (temp is not None) and isinstance(temp, (int, float)) and (0 < temp) and (temp <= 1), \
    #    "temp must be defined as a number between (0, 1]"
    #if method == 'nucleus':
    #    assert (p_nucleus is not None) and isinstance(p_nucleus, (int, float)) and (0 < p_nucleus) and (p_nucleus <= 1), \
    #    "p_nucleus must be defined as a number between (0, 1]"
    ## if method == 'num_beams':
    ##     assert isinstance(num_beams, int) and (num_beams) > 0 and (num_beams) < 100
    #if method == 'top-k':
    #    assert (k is not None) and isinstance(k, int) and (k > 0) and (k < SEQ_LENGTH), \
    #    "k must be defined as an integer greater than 0 and less than the model sequence length"

    #if max_new_tokens is None:
    #    print('No max_new_tokens provided, using a default value of 250\n')
    #    max_new_tokens = 250

    #assert (max_new_tokens is not None) and isinstance(max_new_tokens, int) and (max_new_tokens) > 0 and (max_new_tokens) <= 1000, \
    #"max_new_tokens must be an integer between (0, 1000]"

    if deterministic is not None:
        set_seed(deterministic)

    if cond != "":

        cond_tokens = tokenizer(cond).input_ids

        gen_tokens = model.generate(torch.tensor(cond_tokens).unsqueeze(0).long().to(device),
                                    method=method, k=k, p_nucleus=p_nucleus, temp=temp,
                                    max_new_tokens=max_new_tokens)[0]

        # Insert delimiter to indicate where prompt ends
        gen_prep = torch.zeros(len(gen_tokens)+2).long() # make space for two more tokens for delimiter
        gen_prep -= 1
        gen_prep[:len(cond_tokens)] = gen_tokens[:len(cond_tokens)]
        gen_prep[-(len(gen_tokens)-len(cond_tokens)):] = gen_tokens[-(len(gen_tokens)-len(cond_tokens)):]
        gen_prep[gen_prep == -1] = torch.tensor(tokenizer.encode(' || ')) # insert tokens for || in between

        res = tokenizer.decode(gen_prep)
        res = re.sub(re.escape(tokenizer.bos_token), '', res, count=1) ## Remove end token
        

    else:
        empty_tokens = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)

        res = tokenizer.batch_decode(model.generate(empty_tokens,
                                                    method=method, k=k,
                                                    p_nucleus=p_nucleus, temp=temp,
                                                    max_new_tokens=max_new_tokens))[0]

        res = re.sub(re.escape(tokenizer.bos_token), '', res, count=2) ## Remove start and end tokens

    # Clean up Unicode character issues
    # '“' then 'â€' = opening and closing double quotes
    # '’' = apostrophe
    res = re.sub(r'“', '"', res)
    res = re.sub(r'’', "'", res)
    res = re.sub(r'â€', '"', res)
    res = res + " [END]" ## better end token
    return res