from manipulate import Manipulator
import tensorflow as tf
import numpy as np 
import torch
import clip
from MapTS import GetBoundary,GetDt

class StyleCLIP():
    
    def __init__(self,dataset_name='ffhq'):
        print('load clip')
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model, preprocess = clip.load("ViT-B/32", device=device)
        self.LoadData(dataset_name)
        
    def LoadData(self, dataset_name):
        tf.keras.backend.clear_session()
        M=Manipulator(dataset_name=dataset_name)
        np.set_printoptions(suppress=True)
        fs3=np.load('./npy/'+dataset_name+'/fs3.npy')
        
        self.M=M
        self.fs3=fs3
        
        w_plus=np.load('./data/'+dataset_name+'/w_plus.npy')
        self.M.dlatents=M.W2S(w_plus)
        
        if dataset_name=='ffhq':
            self.c_threshold=20
        else:
            self.c_threshold=100
        self.SetInitP()
    
    def SetInitP(self):
        self.M.alpha=[3]
        self.M.num_images=1
        
        self.target=''
        self.neutral=''
        self.GetDt2()
        img_index=0
        self.M.dlatent_tmp=[tmp[img_index:(img_index+1)] for tmp in self.M.dlatents]
        
        
    def GetDt2(self):
        classnames=[self.target,self.neutral]
        dt=GetDt(classnames,self.model)
        
        self.dt=dt
        num_cs=[]
        betas=np.arange(0.1,0.3,0.01)
        for i in range(len(betas)):
            boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=betas[i])
            print(betas[i])
            num_cs.append(num_c)
        
        num_cs=np.array(num_cs)
        select=num_cs>self.c_threshold
        
        if sum(select)==0:
            self.beta=0.1
        else:
            self.beta=betas[select][-1]
        
    
    def GetCode(self):
        boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=self.beta)
        codes=self.M.MSCode(self.M.dlatent_tmp,boundary_tmp2)
        return codes
    
    def GetImg(self):
        
        codes=self.GetCode()
        out=self.M.GenerateImg(codes)
        img=out[0,0]
        return img
    



#%%
if __name__ == "__main__":
    style_clip=StyleCLIP()
    self=style_clip