import torch
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
from datasets import load_dataset
from PIL import Image
import numpy as np
import paddlehub as hub
import random
from PIL import ImageDraw,ImageFont

import streamlit as st

@st.experimental_singleton
def load_bg_model():
    bg_model = hub.Module(name='U2NetP', directory='assets/models/')
    return bg_model


bg_model = load_bg_model()
def remove_bg(img):
    result = bg_model.Segmentation(
        images=[np.array(img)[:,:,::-1]],
        paths=None,
        batch_size=1,
        input_size=320,
        output_dir=None,
        visualization=False)
    output = result[0]
    mask=Image.fromarray(output['mask'])
    front=Image.fromarray(output['front'][:,:,::-1]).convert("RGBA")
    front.putalpha(mask)
    return front

meme_template=Image.open("./assets/pigeon_meme.jpg").convert("RGBA")
def make_meme(pigeon,text="Is this a pigeon?",show_text=True,remove_background=True):
    
    meme=meme_template.copy()
    approx_butterfly_center=(850,30)

    if remove_background:
        pigeon=remove_bg(pigeon)
        
    else:
        pigeon=Image.fromarray(pigeon).convert("RGBA")

    random_rotate=random.randint(-30,30)
    random_size=random.randint(150,200)
    pigeon=pigeon.resize((random_size,random_size)).rotate(random_rotate,expand=True)

    meme.alpha_composite(pigeon, approx_butterfly_center)

    #ref: https://blog.lipsumarium.com/caption-memes-in-python/
    def drawTextWithOutline(text, x, y):
        draw.text((x-2, y-2), text,(0,0,0),font=font)
        draw.text((x+2, y-2), text,(0,0,0),font=font)
        draw.text((x+2, y+2), text,(0,0,0),font=font)
        draw.text((x-2, y+2), text,(0,0,0),font=font)
        draw.text((x, y), text, (255,255,255), font=font)

    if show_text:
        draw = ImageDraw.Draw(meme)
        font_size=52
        font = ImageFont.truetype("assets/impact.ttf", font_size)
        w, h = draw.textsize(text, font) # measure the size the text will take
        drawTextWithOutline(text, meme.width/2 - w/2, meme.height - font_size*2)
    meme = meme.convert("RGB")    
    return meme

def get_train_data(dataset_name="huggan/smithsonian_butterflies_subset"):
    dataset=load_dataset(dataset_name)
    dataset=dataset.sort("sim_score")
    return dataset["train"]

from transformers import BeitFeatureExtractor, BeitForImageClassification
emb_feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-base-patch16-224')
emb_model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224')
def embed(images):
    inputs = emb_feature_extractor(images=images, return_tensors="pt")
    outputs = emb_model(**inputs,output_hidden_states= True)
    last_hidden=outputs.hidden_states[-1]
    pooler=emb_model.base_model.pooler
    final_emb=pooler(last_hidden).detach().numpy()
    return final_emb    
    
def build_index():    
    dataset=get_train_data()
    ds_with_embeddings = dataset.map(lambda x: {"beit_embeddings":embed(x["image"])},batched=True,batch_size=20)
    ds_with_embeddings.add_faiss_index(column='beit_embeddings')
    ds_with_embeddings.save_faiss_index('beit_embeddings', 'beit_index.faiss')

def get_dataset():
    dataset=get_train_data()
    dataset.load_faiss_index('beit_embeddings', 'beit_index.faiss')
    return dataset

def load_model(model_name='ceyda/butterfly_cropped_uniq1K_512',model_version=None):
    gan = LightweightGAN.from_pretrained(model_name,version=model_version)
    gan.eval()
    return gan
    
def generate(gan,batch_size=1):
    with torch.no_grad():
        ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0., 1.)*255
        ims = ims.permute(0,2,3,1).detach().cpu().numpy().astype(np.uint8)
    return ims

def interpolate():
    pass