import os
import argparse
from collections import defaultdict
import logging
import pickle
import json

import numpy as np
import torch
from torch import nn
from scipy.io import loadmat

from configs.default import get_cfg_defaults
import dlib
import cv2


def _reset_parameters(model):
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)


def get_video_style(video_name, style_type):
    person_id, direction, emotion, level, *_ = video_name.split("_")
    if style_type == "id_dir_emo_level":
        style = "_".join([person_id, direction, emotion, level])
    elif style_type == "emotion":
        style = emotion
    elif style_type == "id":
        style = person_id
    else:
        raise ValueError("Unknown style type")

    return style


def get_style_video_lists(video_list, style_type):
    style2video_list = defaultdict(list)
    for video in video_list:
        style = get_video_style(video, style_type)
        style2video_list[style].append(video)

    return style2video_list


def get_face3d_clip(
    video_name, video_root_dir, num_frames, start_idx, dtype=torch.float32
):
    """_summary_

    Args:
        video_name (_type_): _description_
        video_root_dir (_type_): _description_
        num_frames (_type_): _description_
        start_idx (_type_): "random" , middle, int
        dtype (_type_, optional): _description_. Defaults to torch.float32.

    Raises:
        ValueError: _description_
        ValueError: _description_

    Returns:
        _type_: _description_
    """
    video_path = os.path.join(video_root_dir, video_name)
    if video_path[-3:] == "mat":
        face3d_all = loadmat(video_path)["coeff"]
        face3d_exp = face3d_all[:, 80:144]  # expression 3DMM range
    elif video_path[-3:] == "txt":
        face3d_exp = np.loadtxt(video_path)
    else:
        raise ValueError("Invalid 3DMM file extension")

    length = face3d_exp.shape[0]
    clip_num_frames = num_frames
    if start_idx == "random":
        clip_start_idx = np.random.randint(low=0, high=length - clip_num_frames + 1)
    elif start_idx == "middle":
        clip_start_idx = (length - clip_num_frames + 1) // 2
    elif isinstance(start_idx, int):
        clip_start_idx = start_idx
    else:
        raise ValueError(f"Invalid start_idx {start_idx}")

    face3d_clip = face3d_exp[clip_start_idx : clip_start_idx + clip_num_frames]
    face3d_clip = torch.tensor(face3d_clip, dtype=dtype)

    return face3d_clip


def get_video_style_clip(
    video_name,
    video_root_dir,
    style_max_len,
    start_idx="random",
    dtype=torch.float32,
    return_start_idx=False,
):
    video_path = os.path.join(video_root_dir, video_name)
    if video_path[-3:] == "mat":
        face3d_all = loadmat(video_path)["coeff"]
        face3d_exp = face3d_all[:, 80:144]  # expression 3DMM range
    elif video_path[-3:] == "txt":
        face3d_exp = np.loadtxt(video_path)
    else:
        raise ValueError("Invalid 3DMM file extension")

    face3d_exp = torch.tensor(face3d_exp, dtype=dtype)

    length = face3d_exp.shape[0]
    if length >= style_max_len:
        clip_num_frames = style_max_len
        if start_idx == "random":
            clip_start_idx = np.random.randint(low=0, high=length - clip_num_frames + 1)
        elif start_idx == "middle":
            clip_start_idx = (length - clip_num_frames + 1) // 2
        elif isinstance(start_idx, int):
            clip_start_idx = start_idx
        else:
            raise ValueError(f"Invalid start_idx {start_idx}")

        face3d_clip = face3d_exp[clip_start_idx : clip_start_idx + clip_num_frames]
        pad_mask = torch.tensor([False] * style_max_len)
    else:
        clip_start_idx = None
        padding = torch.zeros(style_max_len - length, face3d_exp.shape[1])
        face3d_clip = torch.cat((face3d_exp, padding), dim=0)
        pad_mask = torch.tensor([False] * length + [True] * (style_max_len - length))

    if return_start_idx:
        return face3d_clip, pad_mask, clip_start_idx
    else:
        return face3d_clip, pad_mask


def get_video_style_clip_from_np(
    face3d_exp,
    style_max_len,
    start_idx="random",
    dtype=torch.float32,
    return_start_idx=False,
):
    face3d_exp = torch.tensor(face3d_exp, dtype=dtype)

    length = face3d_exp.shape[0]
    if length >= style_max_len:
        clip_num_frames = style_max_len
        if start_idx == "random":
            clip_start_idx = np.random.randint(low=0, high=length - clip_num_frames + 1)
        elif start_idx == "middle":
            clip_start_idx = (length - clip_num_frames + 1) // 2
        elif isinstance(start_idx, int):
            clip_start_idx = start_idx
        else:
            raise ValueError(f"Invalid start_idx {start_idx}")

        face3d_clip = face3d_exp[clip_start_idx : clip_start_idx + clip_num_frames]
        pad_mask = torch.tensor([False] * style_max_len)
    else:
        clip_start_idx = None
        padding = torch.zeros(style_max_len - length, face3d_exp.shape[1])
        face3d_clip = torch.cat((face3d_exp, padding), dim=0)
        pad_mask = torch.tensor([False] * length + [True] * (style_max_len - length))

    if return_start_idx:
        return face3d_clip, pad_mask, clip_start_idx
    else:
        return face3d_clip, pad_mask


def get_wav2vec_audio_window(audio_feat, start_idx, num_frames, win_size):
    """

    Args:
        audio_feat (np.ndarray): (N, 1024)
        start_idx (_type_): _description_
        num_frames (_type_): _description_
    """
    center_idx_list = [2 * idx for idx in range(start_idx, start_idx + num_frames)]
    audio_window_list = []
    padding = np.zeros(audio_feat.shape[1], dtype=np.float32)
    for center_idx in center_idx_list:
        cur_audio_window = []
        for i in range(center_idx - win_size, center_idx + win_size + 1):
            if i < 0:
                cur_audio_window.append(padding)
            elif i >= len(audio_feat):
                cur_audio_window.append(padding)
            else:
                cur_audio_window.append(audio_feat[i])
        cur_audio_win_array = np.stack(cur_audio_window, axis=0)
        audio_window_list.append(cur_audio_win_array)

    audio_window_array = np.stack(audio_window_list, axis=0)
    return audio_window_array


def setup_config():
    parser = argparse.ArgumentParser(description="voice2pose main program")
    parser.add_argument(
        "--config_file", default="", metavar="FILE", help="path to config file"
    )
    parser.add_argument(
        "--resume_from", type=str, default=None, help="the checkpoint to resume from"
    )
    parser.add_argument(
        "--test_only", action="store_true", help="perform testing and evaluation only"
    )
    parser.add_argument(
        "--demo_input", type=str, default=None, help="path to input for demo"
    )
    parser.add_argument(
        "--checkpoint", type=str, default=None, help="the checkpoint to test with"
    )
    parser.add_argument("--tag", type=str, default="", help="tag for the experiment")
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    parser.add_argument(
        "--local_rank",
        type=int,
        help="local rank for DistributedDataParallel",
    )
    parser.add_argument(
        "--master_port",
        type=str,
        default="12345",
    )
    parser.add_argument(
        "--max_audio_len",
        type=int,
        default=450,
        help="max_audio_len for inference",
    )
    parser.add_argument(
        "--ddim_num_step",
        type=int,
        default=10,
    )
    parser.add_argument(
        "--inference_seed",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--inference_sample_method",
        type=str,
        default="ddim",
    )
    args = parser.parse_args()

    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    return args, cfg


def setup_logger(base_path, exp_name):
    rootLogger = logging.getLogger()
    rootLogger.setLevel(logging.INFO)

    logFormatter = logging.Formatter("%(asctime)s [%(levelname)-0.5s] %(message)s")

    log_path = "{0}/{1}.log".format(base_path, exp_name)
    fileHandler = logging.FileHandler(log_path)
    fileHandler.setFormatter(logFormatter)
    rootLogger.addHandler(fileHandler)

    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    rootLogger.addHandler(consoleHandler)
    rootLogger.handlers[0].setLevel(logging.INFO)

    logging.info("log path: %s" % log_path)


def cosine_loss(a, v, y, logloss=nn.BCELoss()):
    d = nn.functional.cosine_similarity(a, v)
    loss = logloss(d.unsqueeze(1), y)
    return loss


def get_pose_params(mat_path):
    """Get pose parameters from mat file

    Args:
        mat_path (str): path of mat file

    Returns:
        pose_params (numpy.ndarray): shape (L_video, 9), angle, translation, crop paramters
    """
    mat_dict = loadmat(mat_path)

    np_3dmm = mat_dict["coeff"]
    angles = np_3dmm[:, 224:227]
    translations = np_3dmm[:, 254:257]

    np_trans_params = mat_dict["transform_params"]
    crop = np_trans_params[:, -3:]

    pose_params = np.concatenate((angles, translations, crop), axis=1)

    return pose_params


def sinusoidal_embedding(timesteps, dim):
    """

    Args:
        timesteps (_type_): (B,)
        dim (_type_): (C_embed)

    Returns:
        _type_: (B, C_embed)
    """
    # check input
    half = dim // 2
    timesteps = timesteps.float()

    # compute sinusoidal embedding
    sinusoid = torch.outer(
        timesteps, torch.pow(10000, -torch.arange(half).to(timesteps).div(half))
    )
    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
    if dim % 2 != 0:
        x = torch.cat([x, torch.zeros_like(x[:, :1])], dim=1)
    return x


def get_wav2vec_audio_window(audio_feat, start_idx, num_frames, win_size):
    """

    Args:
        audio_feat (np.ndarray): (250, 1024)
        start_idx (_type_): _description_
        num_frames (_type_): _description_
    """
    center_idx_list = [2 * idx for idx in range(start_idx, start_idx + num_frames)]
    audio_window_list = []
    padding = np.zeros(audio_feat.shape[1], dtype=np.float32)
    for center_idx in center_idx_list:
        cur_audio_window = []
        for i in range(center_idx - win_size, center_idx + win_size + 1):
            if i < 0:
                cur_audio_window.append(padding)
            elif i >= len(audio_feat):
                cur_audio_window.append(padding)
            else:
                cur_audio_window.append(audio_feat[i])
        cur_audio_win_array = np.stack(cur_audio_window, axis=0)
        audio_window_list.append(cur_audio_win_array)

    audio_window_array = np.stack(audio_window_list, axis=0)
    return audio_window_array


def reshape_audio_feat(style_audio_all_raw, stride):
    """_summary_

    Args:
        style_audio_all_raw (_type_): (stride * L, C)
        stride (_type_): int

    Returns:
        _type_: (L, C * stride)
    """
    style_audio_all_raw = style_audio_all_raw[
        : style_audio_all_raw.shape[0] // stride * stride
    ]
    style_audio_all_raw = style_audio_all_raw.reshape(
        style_audio_all_raw.shape[0] // stride, stride, style_audio_all_raw.shape[1]
    )
    style_audio_all = style_audio_all_raw.reshape(style_audio_all_raw.shape[0], -1)
    return style_audio_all


import random


def get_derangement_tuple(n):
    while True:
        v = [i for i in range(n)]
        for j in range(n - 1, -1, -1):
            p = random.randint(0, j)
            if v[p] == j:
                break
            else:
                v[j], v[p] = v[p], v[j]
        else:
            if v[0] != 0:
                return tuple(v)


def compute_aspect_preserved_bbox(bbox, increase_area, h, w):
    left, top, right, bot = bbox
    width = right - left
    height = bot - top

    width_increase = max(
        increase_area, ((1 + 2 * increase_area) * height - width) / (2 * width)
    )
    height_increase = max(
        increase_area, ((1 + 2 * increase_area) * width - height) / (2 * height)
    )

    left_t = int(left - width_increase * width)
    top_t = int(top - height_increase * height)
    right_t = int(right + width_increase * width)
    bot_t = int(bot + height_increase * height)

    left_oob = -min(0, left_t)
    right_oob = right - min(right_t, w)
    top_oob = -min(0, top_t)
    bot_oob = bot - min(bot_t, h)

    if max(left_oob, right_oob, top_oob, bot_oob) > 0:
        max_w = max(left_oob, right_oob)
        max_h = max(top_oob, bot_oob)
        if max_w > max_h:
            return left_t + max_w, top_t + max_w, right_t - max_w, bot_t - max_w
        else:
            return left_t + max_h, top_t + max_h, right_t - max_h, bot_t - max_h

    else:
        return (left_t, top_t, right_t, bot_t)


def crop_src_image(src_img, save_img, increase_ratio, detector=None):
    if detector is None:
        detector = dlib.get_frontal_face_detector()

    img = cv2.imread(src_img)
    faces = detector(img, 0)
    h, width, _ = img.shape
    if len(faces) > 0:
        bbox = [faces[0].left(), faces[0].top(), faces[0].right(), faces[0].bottom()]
        l = bbox[3] - bbox[1]
        bbox[1] = bbox[1] - l * 0.1
        bbox[3] = bbox[3] - l * 0.1
        bbox[1] = max(0, bbox[1])
        bbox[3] = min(h, bbox[3])
        bbox = compute_aspect_preserved_bbox(
            tuple(bbox), increase_ratio, img.shape[0], img.shape[1]
        )
        img = img[bbox[1] : bbox[3], bbox[0] : bbox[2]]
        img = cv2.resize(img, (256, 256))
        cv2.imwrite(save_img, img)
    else:
        raise ValueError("No face detected in the input image")
        # img = cv2.resize(img, (256, 256))
        # cv2.imwrite(save_img, img)