import os
import json

import numpy as np
import torch
from tqdm.auto import tqdm

from models import CNN
from configs import CnnConfig
from utils import load_json
import cv2
import glob


def replace_nan(x, y):
    x = 0.0 if np.isnan(x) else x
    y = 0.0 if np.isnan(y) else y
    return x, y


def draw_hands(
    image,
    hand_x,
    hand_y,
    connections,
    connection_color,
    thickness,
    point_color,
    frame_length,
    frame_width,
):
    for connection in connections:
        x0 = hand_x[connection[0]] * frame_width
        y0 = hand_y[connection[0]] * frame_length
        x1 = hand_x[connection[1]] * frame_width
        y1 = hand_y[connection[1]] * frame_length

        x0, y0 = replace_nan(x0, y0)
        x1, y1 = replace_nan(x1, y1)

        cv2.line(
            image,
            (int(x0), int(y0)),
            (int(x1), int(y1)),
            connection_color,
            thickness,
        )

    for x, y in zip(hand_x, hand_y):
        x, y = replace_nan(x, y)
        cv2.circle(image, (int(x), int(y)), thickness, point_color, thickness)

    return image


def draw_pose(
    image,
    pose_x,
    pose_y,
    links,
    connection_color,
    thickness,
    point_color,
    frame_length,
    frame_width,
):
    for link in links:
        x0 = pose_x[link[0]] * frame_width
        y0 = pose_y[link[0]] * frame_length
        x1 = pose_x[link[1]] * frame_width
        y1 = pose_y[link[1]] * frame_length

        x0, y0 = replace_nan(x0, y0)
        x1, y1 = replace_nan(x1, y1)

        cv2.line(
            image,
            (int(x0), int(y0)),
            (int(x1), int(y1)),
            connection_color,
            thickness,
        )
        cv2.circle(image, (int(x0), int(y0)), thickness, point_color, thickness)
        cv2.circle(image, (int(x1), int(y1)), thickness, point_color, thickness)

    return image


def cnn_feat(file_path, save_dir):
    video_record = load_json(file_path)

    FRAME_LENGTH = 1080
    FRAME_WIDTH = 1920
    POINT_COLOR = (255, 0, 0)
    CONNECTION_COLOR = (0, 255, 0)
    THICKNESS = 2

    connections = [
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 4),
        (5, 6),
        (6, 7),
        (7, 8),
        (9, 10),
        (10, 11),
        (11, 12),
        (13, 14),
        (14, 15),
        (15, 16),
        (17, 18),
        (18, 19),
        (19, 20),
        (0, 5),
        (5, 9),
        (9, 13),
        (13, 17),
        (0, 17),
    ]

    links = [
        (11, 12),
        (11, 23),
        (12, 24),
        (23, 24),
        (11, 13),
        (13, 15),
        (12, 14),
        (14, 16),
        (15, 21),
        (15, 17),
        (17, 19),
        (19, 15),
        (22, 16),
        (16, 18),
        (18, 20),
        (16, 20),
    ]

    model = CNN(CnnConfig)
    features = np.empty((0, CnnConfig.output_dim))

    assert video_record["n_frames"] > 0, "Number of frames should be greater than zero"
    for i in range(video_record["n_frames"]):
        image = np.zeros((FRAME_LENGTH, FRAME_WIDTH, 3), np.uint8)
        pose_x = video_record["pose_x"][i]
        pose_y = video_record["pose_y"][i]
        hand1_x = video_record["hand1_x"][i]
        hand1_y = video_record["hand1_y"][i]
        hand2_x = video_record["hand2_x"][i]
        hand2_y = video_record["hand2_y"][i]

        if hand1_x[0] != 0:
            image = draw_hands(
                image,
                hand1_x,
                hand1_y,
                connections,
                CONNECTION_COLOR,
                THICKNESS,
                POINT_COLOR,
                FRAME_LENGTH,
                FRAME_WIDTH,
            )

        if hand2_x[0] != 0:
            image = draw_hands(
                image,
                hand2_x,
                hand2_y,
                connections,
                CONNECTION_COLOR,
                THICKNESS,
                POINT_COLOR,
                FRAME_LENGTH,
                FRAME_WIDTH,
            )

        image = draw_pose(
            image,
            pose_x,
            pose_y,
            links,
            CONNECTION_COLOR,
            THICKNESS,
            POINT_COLOR,
            FRAME_LENGTH,
            FRAME_WIDTH,
        )
        image = image.astype(np.float32) / 255
        image = cv2.resize(image, (224, 224))
        feat = model(torch.FloatTensor(image).permute(2, 0, 1).unsqueeze(0))
        features = np.vstack([features, feat.numpy()])

    save_path = os.path.join(save_dir, video_record["uid"] + ".npy")
    np.save(save_path, features)


def runner(args, mode):
    json_files = sorted(
        glob.glob(
            os.path.join(args.data_dir, f"{args.dataset}_{mode}_keypoints", "*.json")
        )
    )

    save_dir = os.path.join(args.data_dir, f"{args.dataset}_{mode}_features")
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    saved_files = glob.glob(os.path.join(save_dir, "*.npy"))
    if len(saved_files) != len(json_files):
        for file_path in tqdm(json_files, desc=f"Saving CNN features - {mode} files"):
            cnn_feat(file_path, save_dir)
    else:
        print(mode, "CNN features already exist!")


def save_cnn_features(args):
    runner(args, mode="train")
    runner(args, mode="val")
    runner(args, mode="test")