#encoding = utf8
import copy
import glob
import logging
import os
import pickle

import cv2
import numpy as np
import torch
from tqdm import tqdm
from PIL import Image

import face_detection
from models import Wav2Lip, Wav2LipV2

logger = logging.getLogger(__name__)


def mirror_index(size, index):
    # size = len(self.coord_list_cycle)
    turn = index // size
    res = index % size
    if turn % 2 == 0:
        return res
    else:
        return size - res - 1


def read_image(path):
    image = Image.open(path)
    return image


def read_images(img_list):
    frames = []
    print('reading images...')
    for img_path in tqdm(img_list):
        # frame = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
        frame = Image.open(img_path)
        frame = np.array(frame)
        frames.append(frame)
    return frames


def read_files_path(path):
    file_paths = []
    files = os.listdir(path)
    for file in files:
        if not os.path.isdir(file) and file.endswith('.png') or file.endswith('.jpg'):
            file_paths.append(os.path.join(path, file))
    return file_paths


def get_smoothened_boxes(boxes, t):
    for i in range(len(boxes)):
        if i + t > len(boxes):
            window = boxes[len(boxes) - t:]
        else:
            window = boxes[i: i + t]
        boxes[i] = np.mean(window, axis=0)
    return boxes


def datagen_signal(frame, mel, face_det_results, img_size, wav2lip_batch_size=128):
    img_batch, mel_batch, frame_batch, coord_batch = [], [], [], []

    idx = 0
    frame_to_save = frame.copy()
    face, coord = face_det_results[idx].copy()

    face = cv2.resize(face, (img_size, img_size))

    for i, m in enumerate(mel):
        img_batch.append(face)
        mel_batch.append(m)
        frame_batch.append(frame_to_save)
        coord_batch.append(coord)

    if len(img_batch) >= wav2lip_batch_size:
        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

        img_masked = img_batch.copy()
        img_masked[:, img_size // 2:] = 0
        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

        return img_batch, mel_batch, frame_batch, coord_batch

    if len(img_batch) > 0:
        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
        img_masked = img_batch.copy()
        img_masked[:, img_size//2:] = 0

        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

        return img_batch, mel_batch, frame_batch, coord_batch


def face_detect(images, device):
    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
                                            flip_input=False, device=device)
    batch_size = 16

    while 1:
        predictions = []
        try:
            for i in range(0, len(images), batch_size):
                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
        except RuntimeError:
            if batch_size == 1:
                raise RuntimeError(
                    'Image too big to run face detection on GPU. Please use the --resize_factor argument')
            batch_size //= 2
            print('Recovering from OOM error; New batch size: {}'.format(batch_size))
            continue
        break

    results = []
    pad_y1, pad_y2, pad_x1, pad_x2 = [0, 10, 0, 0]
    for rect, image in zip(predictions, images):
        if rect is None:
            cv2.imwrite('temp/faulty_frame.jpg', image)  # check this frame where the face was not detected.
            raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')

        y1 = max(0, rect[1] - pad_y1)
        y2 = min(image.shape[0], rect[3] + pad_y2)
        x1 = max(0, rect[0] - pad_x1)
        x2 = min(image.shape[1], rect[2] + pad_x2)

        results.append([x1, y1, x2, y2])

    boxes = np.array(results)
    if not False:
        boxes = get_smoothened_boxes(boxes, t=5)
    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]

    del detector
    return results


def get_device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


def _load(checkpoint_path):
    device = get_device()
    if device == 'cuda':
        checkpoint = torch.load(checkpoint_path)
    else:
        checkpoint = torch.load(checkpoint_path,
                                map_location=lambda storage, loc: storage)
    return checkpoint


def load_model(path):
    model = Wav2LipV2()
    print("Load checkpoint from: {}".format(path))
    logging.info(f'Load checkpoint from {path}')
    checkpoint = _load(path)
    s = checkpoint["state_dict"]
    new_s = {}
    for k, v in s.items():
        new_s[k.replace('module.', '')] = v
    model.load_state_dict(new_s)
    device = get_device()
    model = model.to(device)
    return model.eval()


def load_avatar(path, img_size, device):
    print(f'load avatar:{path}')
    face_images_path = os.path.join(path, 'face')
    face_images_path = read_files_path(face_images_path)
    full_list_cycle = read_images(face_images_path)

    face_det_results = face_detect(full_list_cycle, device)

    face_frames = []
    coord_frames = []
    for face, coord in face_det_results:
        resized_crop_frame = cv2.resize(face[:, :, :3], (img_size, img_size))
        face_frames.append(resized_crop_frame)
        coord_frames.append(coord)

    return full_list_cycle, face_frames, coord_frames


def load_avatar_from_processed(base_path, avatar_name):
    avatar_path = os.path.join(base_path, 'data', 'avatars', avatar_name)
    print(f'load avatar from processed:{avatar_path}')
    coord_path = os.path.join(avatar_path, 'coords.pkl')
    print(f'load avatar_path from processed:{avatar_path}')
    face_image_path = os.path.join(avatar_path, 'face_imgs')
    print(f'load face_image_path from processed:{face_image_path}')
    full_image_path = os.path.join(avatar_path, 'full_imgs')
    print(f'load full_image_path from processed:{full_image_path}')

    with open(coord_path, 'rb') as f:
        coord_list_frames = pickle.load(f)

    face_image_list = glob.glob(os.path.join(face_image_path, '*.[jpJP][pnPN]*[gG]'))
    face_image_list = sorted(face_image_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
    face_list_cycle = read_images(face_image_list)

    full_image_list = glob.glob(os.path.join(full_image_path, '*.[jpJP][pnPN]*[gG]'))
    full_image_list = sorted(full_image_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
    frame_list_cycle = read_images(full_image_list)

    return frame_list_cycle, face_list_cycle, coord_list_frames


def jpeg_to_png(image):
    min_green = np.array([50, 100, 100])
    max_green = np.array([70, 255, 255])

    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    mask = cv2.inRange(hsv, min_green, max_green)
    mask_not = cv2.bitwise_not(mask)
    green_not = cv2.bitwise_and(image, image, mask=mask_not)
    b, g, r = cv2.split(green_not)

    # todo 合成四通道
    image = cv2.merge([b, g, r, mask_not])
    return image


def load_avatar_from_256_processed(base_path, avatar_name, pkl):
    avatar_path = os.path.join(base_path, 'data', 'avatars', avatar_name, pkl)
    print(f'load avatar from processed:{avatar_path}')

    with open(avatar_path, "rb") as f:
        avatar_data = pickle.load(f)

    face_list_cycle = []
    frame_list_cycle = []
    coord_list_frames = []
    align_frames = []
    m_frames = []
    inv_m_frames = []

    frame_info_list = avatar_data['frame_info_list']
    for frame_info in tqdm(frame_info_list):
        face_list_cycle.append(frame_info['img'])
        frame_list_cycle.append(jpeg_to_png(frame_info['frame']))
        coord_list_frames.append(frame_info['coords'])
        align_frames.append(frame_info['align_frame'])
        m_frames.append(frame_info['m'])
        inv_m_frames.append(frame_info['inv_m'])

    return frame_list_cycle, face_list_cycle, coord_list_frames, align_frames, m_frames, inv_m_frames


def config_logging(file_name: str, console_level: int = logging.INFO, file_level: int = logging.DEBUG):
    file_handler = logging.FileHandler(file_name, mode='a', encoding="utf8")
    file_handler.setFormatter(logging.Formatter(
        '%(asctime)s [%(levelname)s] %(module)s.%(lineno)d %(name)s:\t%(message)s'
        ))
    file_handler.setLevel(file_level)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(logging.Formatter(
        '[%(asctime)s.%(msecs)03d %(levelname)s] %(message)s',
        datefmt="%Y/%m/%d %H:%M:%S"
        ))
    console_handler.setLevel(console_level)

    logging.basicConfig(
        level=min(console_level, file_level),
        handlers=[file_handler, console_handler],
    )


def object_stop(obj):
    if obj is not None:
        obj.stop()


def img_warp_back_inv_m(img, img_to, inv_m):
    h_up, w_up, c = img_to.shape
    mask = np.ones_like(img).astype(np.float32)
    inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
    inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
    mask_indices = inv_mask == 1
    if 4 == c:
        img_to[:, :, :3][mask_indices] = inv_img[mask_indices]
    else:
        img_to[inv_mask == 1] = inv_img[inv_mask == 1]
    return img_to


def render_image(context, frame):
    res_frame, idx, type_ = frame

    if type_ == 0:
        combine_frame = context.frame_list_cycle[idx]
    else:
        bbox = context.coord_list_cycle[idx]
        combine_frame = copy.deepcopy(context.frame_list_cycle[idx])
        af = context.align_frames[idx]
        inv_m = context.inv_m_frames[idx]
        y1, y2, x1, x2 = bbox
        try:
            res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
            af[y1:y2, x1:x2] = res_frame
            combine_frame = img_warp_back_inv_m(af, combine_frame, inv_m)
        except Exception as e:
            logging.error(f'resize error:{e}')
            return None

    return combine_frame