import sys
import traceback
import os
from saicinpainting.evaluation.utils import move_to_device
from torchvision import transforms
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
from PIL import Image
from PIL import ImageColor
import cv2
import hydra
import numpy as np
import torch
import tqdm
import yaml
from omegaconf import OmegaConf
from torch.utils.data._utils.collate import default_collate
import torch.fft
from saicinpainting.training.data.datasets import make_default_val_dataset
from saicinpainting.training.trainers import load_checkpoint
from saicinpainting.utils import register_debug_signal_handlers
import json
import io
import pickle
from torchvision.transforms import ToPILImage
import boto3
import re
from PIL import Image
import itertools
from itertools import chain
import random
import shutil
import tempfile
from contextlib import redirect_stdout
import io
import os
import glob
import traceback
import copy
from itertools import product
# Based on https://github.com/pytorch/examples/blob/master/mnist/main.py

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
s3_client = boto3.client('s3')

##########################Inference###################################

def model_fn(model_dir):
    with open(r'code/default.yaml') as file: 
        predict_config = yaml.safe_load(file)
   
    register_debug_signal_handlers()  
    predict_config['model']['path'] = 'big-lama' 
    train_config_path = os.path.join(predict_config['model']['path'], 'config.yaml') 
    with open(train_config_path, 'r') as f: 
        train_config = OmegaConf.create(yaml.safe_load(f)) 
    train_config.training_model.predict_only = True 
    train_config.visualizer.kind = 'noop' 

    out_ext = '.png' 
    checkpoint_path = os.path.join(predict_config['model']['path'],'models',predict_config['model']['checkpoint'])
    print('------------load checkpoint!!!!-----------------------')
    with redirect_stdout(io.StringIO()) as f:
    
        model = load_checkpoint(train_config, checkpoint_path, strict=False, map_location='cpu')        
        model.freeze() 
        model.to(device) 
    print('------------finish loading checkpoint!!!!-----------------------')
    
    return model


# data preprocessing
def input_fn(request_body, request_content_type):
    
    tempdir = tempfile.mkdtemp(dir='/tmp')
    print("TEMPDIR IS: ",tempdir)
    directory_im=f'{tempdir}/Image' 
    # Create the directory
    try:
        os.makedirs(directory_im, exist_ok = True)
        print("Directory '%s' created successfully" % directory_im)
    except OSError as error:
        print("Directory '%s' can not be created" % directory_im)

    # 2 Create directory for masks
    dir_UIMask=f'{tempdir}/UIMask' 
    try:
        os.makedirs(dir_UIMask, exist_ok = True)
        print("Directory '%s' created successfully" % dir_UIMask)
    except OSError as error:
        print("Directory '%s' can not be created" % dir_UIMask)

    # 3 Create directory for transformed masks
    dir_input_mask=f'{tempdir}/InputMask' 
    try:
        os.makedirs(dir_input_mask, exist_ok = True)
        print("Directory '%s' created successfully" % dir_input_mask)
    except OSError as error:
        print("Directory '%s' can not be created" % dir_input_mask)   

    # 4 Create directory for training dataset
    dir_train=f'{tempdir}/Train'  
    try:
        os.makedirs(dir_train, exist_ok = True)
        print("Directory '%s' created successfully" % dir_train)
    except OSError as error:
        print("Directory '%s' can not be created" % dir_train)

    # 5 Create directory for training dataset
    dir_test=f'{tempdir}/Test' 
    try:
        os.makedirs(dir_test, exist_ok = True)
        print("Directory '%s' created successfully" % dir_test)
    except OSError as error:
        print("Directory '%s' can not be created" % dir_test)
    
    directories=[tempdir,directory_im,dir_UIMask,dir_input_mask,dir_train,dir_test]

    input_json = json.loads(request_body)
    s3_image_list=input_json['input-image-location']
    mask_meta_list=input_json['mask-info']
    output_bucket=input_json['output-bucket']
    output_project=input_json['output-project']
    
    #Download images
    print("Downloading images")
    for i in range(len(input_json['input-image-location'])):
        image_location=input_json['input-image-location'][i]
        s1=image_location.split('s3://')[1]
        image_bucket=s1.split('/')[0]
        im_s3=s1.split(image_bucket+'/')[1]
        s3 = boto3.client('s3')
        with open(directory_im+'/'+str(i)+'_im.jpg', 'wb') as f:
            s3.download_fileobj(image_bucket, im_s3, f)
    #Image list
    imglist= [file for file in os.listdir(directory_im) if file.endswith('.jpg')]
    
    #Download masks
    print("Downloading masks")
    colors = []
    single_mask_list = []
    defect_name_list=[]
    hex_color_list=[]
    
    mask_location=input_json['mask-location']
    print('mask location:',mask_location)
    s1=mask_location.split('s3://')[1]
    mask_bucket=s1.split('/')[0]
    mask_s3=s1.split(mask_bucket+'/')[1]
    mask_raw_name=mask_s3.split('mask/')[1]
    print('mask s3:',mask_s3)
    print('mask name:',mask_raw_name)
    
    for i in range(len(input_json['mask-info'])-1):
        
        color_tmp = ImageColor.getcolor(input_json['mask-info'][f'{i+1}']['hex-color'], "RGB")
        
        #change that to BGR
        color_tmp2 = (color_tmp[2],color_tmp[1],color_tmp[0])
        colors.append(color_tmp2)
        
        defect_name=input_json['mask-info'][f'{i+1}']['class-name']
        defect_name_list.append(defect_name)
        hex_color_list.append(input_json['mask-info'][f'{i+1}']['hex-color'])
        
        
        s3_client.download_file(mask_bucket,mask_s3,f'{dir_UIMask}/mask_{defect_name}.png')
        single_mask_list.append(f'mask_{defect_name}.png')
   
    #Transform single masks for ML inference
    for i in range(len(single_mask_list)):
        mask= cv2.imread(dir_UIMask+'/'+single_mask_list[i])
        mask[np.logical_not(np.all(mask == colors[i], axis=-1))] = (255, 255, 255)
        #added this so i can create the mask as is expected by the service
        cv2.imwrite(dir_UIMask+'/'+single_mask_list[i], mask)
        mask[np.logical_not(np.all(mask == colors[i], axis=-1))] = (0, 0, 0)
        mask[np.all(mask == colors[i], axis=-1)] = (255, 255,  255)        
        cv2.imwrite(dir_input_mask+'/'+single_mask_list[i], mask)
    
    #number of mask
    mask_num=len(input_json['mask-info'])-1
    #Generate all combinations of masks
    all_comb=[]
    for i in range(2,mask_num+1):
        comb_list=list(itertools.combinations(range(mask_num), i))
        all_comb.append(comb_list)
    all_comb=list(chain(*all_comb))
    
    #Combine masks and save them in UIMask
    for i in range(len(all_comb)):
        class1=defect_name_list[all_comb[i][0]]
        class2=defect_name_list[all_comb[i][1]]
        mask_a = cv2.imread(dir_UIMask+'/'+'mask_'+class1+'.png')
        mask_b = cv2.imread(dir_UIMask+'/'+'mask_'+class2+'.png')
        mask_bwa = cv2.bitwise_and(mask_a,mask_b)
        
        class_name=str()
        for j in range(len(all_comb[i])-2):
            class_new=defect_name_list[all_comb[i][j+2]]
            mask_c = cv2.imread(dir_UIMask+'/'+'mask_'+class_new+'.png')
            mask_bwa = cv2.bitwise_and(mask_bwa,mask_c)
            class_name=class_name+'_'+class_new
        cv2.imwrite(dir_UIMask+'/'+"mask_"+class1+'_'+class2+class_name+'.png', mask_bwa)
        
    
    #combined mask list
    masklist= [file for file in os.listdir(dir_UIMask) if file.endswith('.png')]
    #upload mask to s3
    for i in range(len(masklist)):
        s3_client.upload_file(dir_UIMask+'/'+masklist[i],output_bucket,output_project+'/mask/MixMask/'+masklist[i])
        
    #Transform combined masks and save them in InputMask directory
    for i in range(len(all_comb)):
        class1=defect_name_list[all_comb[i][0]]
        class2=defect_name_list[all_comb[i][1]]

        mask_a = cv2.imread(dir_UIMask+'/'+'mask_'+class1+'.png')
        mask_a[np.all(mask_a == (255,255,255), axis=-1)] = (0, 0, 0)
        mask_a[np.all(mask_a != (0,0,0), axis=-1)] = (255, 255,  255)

        mask_b = cv2.imread(dir_UIMask+'/'+'mask_'+class2+'.png')
        mask_b[np.all(mask_b == (255,255,255), axis=-1)] = (0, 0, 0)
        mask_b[np.all(mask_b != (0,0,0), axis=-1)] = (255, 255,  255)

        mask_bwa = cv2.bitwise_or(mask_a,mask_b)


        class_name=str()
        for j in range(len(all_comb[i])-2):
            class_new=defect_name_list[all_comb[i][j+2]]
            mask_c = cv2.imread(dir_UIMask+'/'+'mask_'+class_new+'.png')
            mask_c[np.all(mask_c == (255,255,255), axis=-1)] = (0, 0, 0)
            mask_c[np.all(mask_c != (0,0,0), axis=-1)] = (255, 255,  255)
            mask_bwa = cv2.bitwise_or(mask_bwa,mask_c)
            class_name=class_name+'_'+class_new
        cv2.imwrite(dir_input_mask+'/'+"mask_"+class1+'_'+class2+class_name+'.png', mask_bwa)
    return(imglist,single_mask_list,masklist,mask_meta_list,output_bucket,output_project,mask_num,s3_image_list,directories,defect_name_list,hex_color_list)
    

def load_image(fname, mode='RGB', return_orig=False, auto_resize = 2000):
    img_ori = Image.open(fname).convert(mode)
    h=img_ori.size[0]
    w=img_ori.size[1]
    img_ori_2 = copy.deepcopy(img_ori)
    if h > auto_resize or w > auto_resize:
        ratio1=h/auto_resize
        ratio2=w/auto_resize
        ratio=max(ratio1,ratio2)
        h_resize=int(h/ratio)
        w_resize=int(w/ratio)

        print(h_resize,w_resize)
        img_ori_2 = img_ori.resize((h_resize, w_resize))
    else:
        h_resize=h
        w_resize=w
    img = np.array(img_ori_2)
    if img.ndim == 3:
        img = np.transpose(img, (2, 0, 1))
    out_img = img.astype('float32') / 255
    if return_orig:
        return out_img, h, w, h_resize, w_resize, img_ori #img
    else:
        return out_img, h, w, h_resize, w_resize

def syn_gen(model,dataset_type,dir_dataset,directory_im,dir_input_mask,imname,maskname,output_bucket,output_project):
    im_short_name = (imname.split('.jpg'))[0]
    imlocation = directory_im+'/'+imname
    mask_short_name = (maskname.split('.png'))[0]
    masklocation = dir_input_mask + '/' +  maskname
    #load image and mask 
    image, h, w,h_resize,w_resize, im_ori = load_image(imlocation, mode='RGB', return_orig=True)
    im_ori = np.array(im_ori)
    mask,_,_,_,_,mask_ori =  load_image(masklocation, mode='L', return_orig=True)
    mask_ori = np.array(mask_ori)
    result = dict(image=image, mask=mask[None, ...])
    #perform ML inference 
    with torch.no_grad():
        img = torch.from_numpy(result['image']).unsqueeze(0)
        mask = torch.from_numpy(result['mask']).unsqueeze(0)
        mask = (mask > 0) * 1
        masked_img = img * (1 - mask)
        masked_img = torch.cat([masked_img, mask], dim=1)
        out1 = model.generator(masked_img.to(device))
        out1 = out1.squeeze_(0)
        outimg = ToPILImage()(out1)
        outimg = outimg.crop((0,0,h_resize,w_resize))
        outimg = outimg.resize((h,w))

        new_im_name = im_short_name+'_'+mask_short_name+'.jpg'
        print('Generate synthetic defect:')
        print(new_im_name)
        print('is done!')
        pixels = np.where(mask_ori == 255)
        im_ori[pixels[0],pixels[1],:] = np.array(outimg)[pixels[0],pixels[1],:]
        im_ori_pil = Image.fromarray(im_ori)
        im_ori_pil.save(dir_dataset+'/'+new_im_name)
        # upload generated synthetic defects to s3
        s3_client.upload_file(dir_dataset+'/'+new_im_name,output_bucket,f'{output_project}/anomaly/{dataset_type}/{new_im_name}')
        #remove defects 
        os.remove(dir_dataset+'/'+new_im_name)
    return new_im_name

def predict_fn(input_object, model):
    
    imglist = input_object[0]
    single_mask_list = input_object[1]
    masklist = input_object[2] 
    mask_meta_list = input_object[3]
    output_bucket = input_object[4]
    output_project = input_object[5]
    mask_num = input_object[6]
    s3_image_list=input_object[7]
    directories = input_object[8]
    defect_name_list=input_object[9]
    hex_color_list=input_object[10]
        
    directory_im = directories[1]
    dir_UIMask = directories[2]
    dir_input_mask = directories[3]
    dir_train = directories[4]
    dir_test = directories[5]
    
  
    # Split the train and test image dataset
    if len(imglist)<20:
        print('There should be at least 20 normal images!')
    else:
        if int(len(imglist)*0.2)>10:
            len_train=len(imglist)*0.8
            len_test=len(imglist)*0.2
        else:
            len_test=10
            len_train=len(imglist)-10
        if len_train>40:
            len_train=40
            len_test=len(imglist)-40
        if len_test>40:
            len_test=40

    # Split the image dataset
    train_normal_list=random.sample(imglist, len_train)
    test_normal_list_all=list(set(imglist) - set(train_normal_list))
    test_normal_list=random.sample(test_normal_list_all, len_test)
    
    # Upload the normal images to s3
    for i in range(len(train_normal_list)):
        s3_client.upload_file(directory_im+'/'+train_normal_list[i],output_bucket,output_project+'/normal/train/'+train_normal_list[i])
    
    for j in range(len(test_normal_list)):
        s3_client.upload_file(directory_im+'/'+test_normal_list[j],output_bucket,output_project+'/normal/test/'+test_normal_list[j])
    
    print("Generating training Images")
    train_anomaly_list=[]
    for i in range(len(single_mask_list)):
        for j in range(len(train_normal_list)):
            try: 
                # get image name 
                train_imname = train_normal_list[j]
                # get mask name 
                train_maskname = single_mask_list[i]
               
                # ML inference
                train_defect=syn_gen(model,'train',dir_train,directory_im,dir_input_mask,train_imname,train_maskname,output_bucket,output_project)
                train_anomaly_list.append(train_defect)    
                print('Generating training synthetic defects is done!')
            except Exception:
                print(traceback.format_exc())
                # or
                print(sys.exc_info()[2])
                print('Error for train dataset!')          

    print('Train anomaly list',train_anomaly_list)
    
    # Generate test dataset
    print("Generating testing Images")
    test_anomaly_list=[]
    
    produc_list_test = list(product(test_normal_list,single_mask_list))    
    test_perm_list = np.random.permutation(produc_list_test)
    
    # Get the combinatino of masks
    comb_mask_list=list(set(masklist) - set(single_mask_list))
    
    produc_list = list(product(imglist,comb_mask_list))        
    rand_perm_list = np.random.permutation(produc_list)
    num_test_images = min(20,len(rand_perm_list))
    test_comb_list = rand_perm_list[0:num_test_images]        
    
    test_im_mask_list=list(test_perm_list) +list(test_comb_list)
 
    for i in range(len(test_im_mask_list)):
        try:
            # get image name 
            test_imname =test_im_mask_list[i][0]
            # get mask name      
            test_maskname=test_im_mask_list[i][1]
            # ML inference
            test_defect=syn_gen(model,'test',dir_test,directory_im,dir_input_mask,test_imname,test_maskname,output_bucket,output_project)
            test_anomaly_list.append(test_defect)
            print('Generating test synthetic defects is done!')
        except:
            print(traceback.format_exc())
            # or
            print(sys.exc_info()[2])
            print('Error for test dataset!') 

    print('Test anomaly list',test_anomaly_list)
    return (train_anomaly_list,test_anomaly_list,mask_meta_list,output_bucket,output_project,s3_image_list,directories,train_normal_list,test_normal_list,defect_name_list,hex_color_list)

def write_manifest(tempdir,dataset_type,dataset_anomaly_list,dataset_normal_list,output_bucket,output_project,defect_name_list,hex_color_list):
    manifest_name=f'{tempdir}/l4v_{dataset_type}.manifest' 
    with open(manifest_name, 'w') as the_new_file:
        for i in range(len(dataset_anomaly_list)):
            new_json = {}
            new_json['source-ref'] =f's3://{output_bucket}/{output_project}/anomaly/{dataset_type}/{dataset_anomaly_list[i]}'
            new_json['auto-label']=11
            new_json['auto-label-metadata']={"class-name": "anomaly","type": "groundtruth/image-classification"}
            
            #get defect class from mask name
            s=dataset_anomaly_list[i]
            t1=s.split('.jpg')[0]
            t2=t1.split('im_')[1]
            mask_name=t2+'.png' 
            new_json['anomaly-mask-ref']=f's3://{output_bucket}/{output_project}/mask/MixMask/'+mask_name

            s1=s.split('.jpg')[0]
            s2=s1.split('mask')[1]
            s3=s2.split('_')
            s4=s3[1:]
            metadata={}
            metadata[0]={'class-name': 'BACKGROUND', 'hex-color': '#ffffff', 'confidence': 0}

            for j in range(len(s4)):
                class_index=defect_name_list.index(s4[j])
                newdefect={}
                newdefect['class-name']=s4[j]
                newdefect['hex-color']=hex_color_list[class_index]
                newdefect['confidence']=0
                metadata[class_index]=newdefect

            new_json['anomaly-mask-ref-metadata']={}
            new_json['anomaly-mask-ref-metadata']['internal-color-map']=metadata
            new_json['anomaly-mask-ref-metadata']['type']="groundtruth/semantic-segmentation"
            
            the_new_file.write(json.dumps(new_json))
            the_new_file.write('\n')
        
        #add normal data for train dataset
        normal_json = {}
        
        for i in range(len(dataset_normal_list)):
            normal_json['source-ref']=f's3://{output_bucket}/{output_project}/normal/{dataset_type}/{dataset_normal_list[i]}'
            normal_json["auto-label"]=12
            normal_json["auto-label-metadata"]={"class-name": "normal","type": "groundtruth/image-classification"}
            the_new_file.write(json.dumps(normal_json))
            the_new_file.write('\n')
                                     
    s3_client.upload_file(manifest_name,output_bucket,output_project+f'/manifest/{dataset_type}/l4v_{dataset_type}.manifest')
                       
def output_fn(output, content_type):
    print("On the output function")
   
    train_anomaly_list=output[0]
    test_anomaly_list=output[1]
    mask_meta_list=output[2]
    output_bucket=output[3]
    output_project=output[4]
    s3_image_list=output[5]
    directories = output[6]
    train_normal_list=output[7]
    test_normal_list=output[8]
    defect_name_list=output[9]
    hex_color_list=output[10]
        
    tempdir = directories[0]
    directory_im = directories[1]
    dir_UIMask = directories[2]
    dir_input_mask = directories[3]
    dir_train = directories[4]
    dir_test = directories[5]
 
    #clean up all the directories
    shutil.rmtree(directory_im)
    shutil.rmtree(dir_UIMask)
    shutil.rmtree(dir_input_mask)
    
    # Create manifest for train dataset
    write_manifest(tempdir,'train',train_anomaly_list,train_normal_list,output_bucket,output_project,defect_name_list,hex_color_list)
    print('Training manifest is generated!')
    
    # Create manifest for test dataset
    write_manifest(tempdir,'test',test_anomaly_list,test_normal_list,output_bucket,output_project,defect_name_list,hex_color_list)
    print('Test manifest is generated!')
    
    # Remove tempdir
    shutil.rmtree(tempdir)
      
    return "done"