from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
from diffusers import DiffusionPipeline
from transformers import AutoModelForSeq2SeqLM
from samplings import top_p_sampling, temperature_sampling
import torch

class AIAssistant:
    def __init__(self):
        pass

    def entity_pos_tagger(self, example):
        tokenizer = AutoTokenizer.from_pretrained("Davlan/bert-base-multilingual-cased-ner-hrl")
        model = AutoModelForTokenClassification.from_pretrained("Davlan/bert-base-multilingual-cased-ner-hrl")
        nlp = pipeline("ner", model=model, tokenizer=tokenizer)
        ner_results = nlp(example)
        return ner_results

    def text_to_image_generation(self, prompt, n_steps=40, high_noise_frac=0.8):
        base = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
        )
        base.to("cuda")
        refiner = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-refiner-1.0",
            text_encoder_2=base.text_encoder_2,
            vae=base.vae,
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16",
        )
        refiner.to("cuda")

        image = base(
            prompt=prompt,
            num_inference_steps=n_steps,
            denoising_end=high_noise_frac,
            output_type="latent",
        ).images
        image = refiner(
            prompt=prompt,
            num_inference_steps=n_steps,
            denoising_start=high_noise_frac,
            image=image,
        ).images[0]
        return image

    def grammatical_pos_tagger(self, text):
        nlp_pos = pipeline(
            "ner",
            model="mrm8488/bert-spanish-cased-finetuned-pos",
            tokenizer=(
                'mrm8488/bert-spanish-cased-finetuned-pos',
                {"use_fast": False}
        ))
        return nlp_pos(text)

    def text_to_music(self, text, max_length=1024, top_p=0.9, temperature=1.0):
        tokenizer = AutoTokenizer.from_pretrained('sander-wood/text-to-music')
        model = AutoModelForSeq2SeqLM.from_pretrained('sander-wood/text-to-music')

        input_ids = tokenizer(text,
                              return_tensors='pt',
                              truncation=True,
                              max_length=max_length)['input_ids']

        decoder_start_token_id = model.config.decoder_start_token_id
        eos_token_id = model.config.eos_token_id

        decoder_input_ids = torch.tensor([[decoder_start_token_id]])

        for t_idx in range(max_length):
            outputs = model(input_ids=input_ids,
                            decoder_input_ids=decoder_input_ids)
            probs = outputs.logits[0][-1]
            probs = torch.nn.Softmax(dim=-1)(probs).detach().numpy()
            sampled_id = temperature_sampling(probs=top_p_sampling(probs,
                                                                top_p=top_p,
                                                                return_probs=True),
                                            temperature=temperature)
            decoder_input_ids = torch.cat((decoder_input_ids, torch.tensor([[sampled_id]])), 1)
            if sampled_id!=eos_token_id:
                continue
            else:
                tune = "X:1\n"
                tune += tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
                return tune
                break

# Ejemplo de uso
assistant = AIAssistant()
ner_results = assistant.entity_pos_tagger("Nader Jokhadar had given Syria the lead with a well-struck header in the seventh minute.")
print(ner_results)

image = assistant.text_to_image_generation("A majestic lion jumping from a big stone at night")
print(image)

pos_tags = assistant.grammatical_pos_tagger('Mis amigos están pensando en viajar a Londres este verano')
print(pos_tags)

tune = assistant.text_to_music("This is a traditional Irish dance music.")
print(tune)