import os
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig


def get_model_and_tokenizer(device, CHECKPOINTS_DIR, WEIGHTS_NAME, model_type):
    config = AutoConfig.from_pretrained(model_type)
    
    tokenizer = AutoTokenizer.from_pretrained(model_type, use_fast=True)
    
    model = AutoModelForCausalLM.from_pretrained(
        model_type,
        from_tf=bool(".ckpt" in model_type),
        config=config,
    )
    model.resize_token_embeddings(len(tokenizer))
    model.gradient_checkpointing_enable()
    model.to(device) 
        
    model_file = os.path.join(CHECKPOINTS_DIR, WEIGHTS_NAME)
    try:
        checkpoint = torch.load(model_file)
    except:
        checkpoint = torch.load(model_file, map_location=device)
    model.load_state_dict(checkpoint)
    
    return tokenizer, model

def generate_prompt(title=""):
    return f"<BOS> TOPIC: {title} COMMENT: "

def extract_comment_from_prompt(text):
    if type(text) == str:
        starts = text.rfind("COMMENT: ") + len("COMMENT: ")
        ends = text.find(" <EOS>")
        result = text[starts:ends]
    else:
        text = pd.Series(text)
        starts = text.str.find("COMMENT: ") + len("COMMENT: ")
        ends = text.str.find(" <EOS>")
        result = [sentence[start:end] for (sentence, start, end) in zip(text, starts, ends)]
    return result

def generate_comment(prompt, tokenizer, model, entry_length=30, top_p=0.8, temperature=1.):
    filter_value = -float("Inf")
    model.eval()
    with torch.no_grad():
        generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
        for i in range(entry_length):
            outputs = model(generated, labels=generated)
            loss, logits = outputs[:2]
            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)

            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0

            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            logits[:, indices_to_remove] = filter_value

            next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
            generated = torch.cat((generated, next_token), dim=1)
            if next_token in tokenizer.encode("<EOS>"):
                break
        output_list = list(generated.squeeze().numpy())
        output_text = f"{tokenizer.decode(output_list)} <EOS>" 
    return output_text