import json
import os
import uuid
import pickle
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

import torch
import numpy as np

from PIL import Image
from datetime import date
from ultralytics import YOLO
from diffusers import DPMSolverSDEScheduler, UNet2DConditionModel
from transformers import CLIPProcessor, CLIPModel

from models.unet import FrameConditionalUNet2DModel
from pipelines.frame_cond_pipeline import FrameConditionalDDPMPipeline
from pipelines.text_cond_pipeline import TextConditionalDDPMPipeline
from layout_single_gen import plot_and_save_segments
from post_process import post_procss_furnitures, post_process_box_data, post_process_seg_data, post_process_frame_data
from trans_frame import SimpleHouse2Web, read_template, s3_bucket_inst, s3_file_url, s3_bucket


img_size = (256, 256)
cache_dir = 'D:\\hugging_face'
model_path = 'D:\\trains\\fp_cond_09111518'
text_control_model_path = 'D:\\trains\\text_control_v5'
obb_model_path = 'D:\\trains\\obbs\\best.pt'
seg_model_path = 'D:\\trains\\seg\\best.pt'
frame_data_path = 'D:\\data\\fp_json_split\\train'
run_path = 'D:\\data\\run'
device_name = 'cuda:2'
device = torch.device(device_name)
model = FrameConditionalUNet2DModel.from_pretrained(model_path, subfolder='unet')
model = model.to(device)
text_model = UNet2DConditionModel.from_pretrained(text_control_model_path, subfolder='unet')
text_model = text_model.to(device)
obb_model = YOLO(obb_model_path)
seg_model = YOLO(seg_model_path)
noise_scheduler = DPMSolverSDEScheduler.from_pretrained(model_path, subfolder='scheduler')
pipe = FrameConditionalDDPMPipeline(model, noise_scheduler)
pipe = pipe.to(device)
text_pipe = TextConditionalDDPMPipeline(text_model, noise_scheduler)
text_pipe = text_pipe.to(device)
generator = torch.Generator(device=device).manual_seed(0)
templates = read_template('web_template.csv')
processor = CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14',
                                          cache_dir=cache_dir, local_files_only=True)
clip_model = CLIPModel.from_pretrained('openai/clip-vit-large-patch14', cache_dir=cache_dir,
                                           local_files_only=True)
clip_model.requires_grad_(False)
clip_model = clip_model.to(device)


if not os.path.exists(run_path):
    os.makedirs(run_path)


def get_text_encoder(captions, batch_size):
    text_input_ids = processor(text=captions, return_tensors="pt", max_length=10, padding="max_length", truncation=True)['input_ids']
    encoder_hidden_states = clip_model.get_text_features(input_ids=text_input_ids.to(device))
    encoder_hidden_states = encoder_hidden_states.repeat(batch_size, 1, 1)
    return encoder_hidden_states


def frame_cond_gen_web(frame_id, batch_size=1, num_inference_steps=100):
    url_li, img_li = [], []
    minx, miny, maxx, maxy = plot_and_save_segments(f'{frame_id}.json', frame_data_path, run_path)
    with open(os.path.join(frame_data_path, f'{frame_id}.json'), 'r', encoding='utf8') as f:
        frame_data = json.load(f)
    frame_img_file = os.path.join(run_path, 'frame', f'{frame_id}.png')
    cond = Image.open(frame_img_file).resize(img_size).convert('RGB')
    cond = np.array(cond)
    cond = torch.from_numpy(cond).float()
    cond = cond.permute(2, 0, 1).repeat(batch_size, 1, 1, 1).to(device)
    pred_images = pipe(cond, batch_size=batch_size, generator=generator, num_inference_steps=num_inference_steps)
    obb_res_li = obb_model.predict(source=pred_images, conf=0.6, imgsz=512, device=device_name)
    for i in range(len(pred_images)):
        fuid = str(uuid.uuid4()).replace('-', '')
        box_data = post_process_box_data(obb_res_li[i])
        with open(os.path.join(run_path, f'{fuid}_box.json'), 'w', encoding='utf8') as f:
            json.dump(box_data, f)
        pretty_res = post_procss_furnitures(frame_data, box_data, (minx, miny, maxx, maxy))
        house = SimpleHouse2Web(pretty_res)
        house.fit(templates)
        filename = f'{date.today()}/{fuid}_bim.json'
        img_filename = os.path.join(run_path, f'{fuid}.png')
        s3_bucket_inst.put_object(Key=filename, Body=json.dumps(house.web_data, ensure_ascii=False))
        url_li.append('/'.join([s3_file_url, s3_bucket, filename]))
        pred_images[i].save(img_filename)
        img_li.append(img_filename)
    return url_li, img_li


def text_cond_gen_web(text_caption, batch_size=1, num_inference_steps=100):
    captions = text_caption.split(',')
    encoder_hidden_states = get_text_encoder(captions, batch_size)
    pred_images = text_pipe(encoder_hidden_states, batch_size=batch_size, generator=generator,
                            num_inference_steps=num_inference_steps)
    obb_res_li = obb_model.predict(source=pred_images, conf=0.3, imgsz=512, device=device_name)
    seg_res_li = seg_model.predict(source=pred_images, conf=0.8, imgsz=512, iou=0.7, device=device_name)
    url_li, img_li = [], []
    for i in range(len(pred_images)):
        fuid = str(uuid.uuid4()).replace('-', '')
        img_filename = os.path.join(run_path, f'{fuid}.png')
        pred_images[i].save(img_filename)
        img_li.append(img_filename)
        union_res_li = post_process_box_data(obb_res_li[i]) + post_process_seg_data(seg_res_li[i])
        with open(os.path.join(run_path, f'{fuid}_yolo.json'), 'w', encoding='utf8') as f:
            json.dump(union_res_li, f, ensure_ascii=False)
        pretty_res = post_process_frame_data(union_res_li, save_file=os.path.join(run_path, f'{fuid}_yolo.png'))
        house = SimpleHouse2Web(pretty_res)
        house.fit(templates)
        filename = f'{date.today()}/{fuid}.json'
        s3_bucket_inst.put_object(Key=filename, Body=json.dumps(house.web_data, ensure_ascii=False))
        url_li.append('/'.join([s3_file_url, s3_bucket, filename]))
    return url_li, img_li


if __name__ == '__main__':
    print(frame_cond_gen_web(1111257539))
