import torch
import numpy as np
import hashlib
from torch.autograd import Variable
import os
    
def deterministic_random(min_value, max_value, data):
    digest = hashlib.sha256(data.encode()).digest()
    raw_value = int.from_bytes(digest[:4], byteorder='little', signed=False)
    return int(raw_value / (2 ** 32 - 1) * (max_value - min_value)) + min_value


def mpjpe_cal(predicted, target):
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1))


def test_calculation(predicted, target, action, error_sum, data_type, subject, MAE=False):
    error_sum = mpjpe_by_action_p1(predicted, target, action, error_sum)
    if not MAE:
        error_sum = mpjpe_by_action_p2(predicted, target, action, error_sum)

    return error_sum


def mpjpe_by_action_p1(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    batch_num = predicted.size(0)
    frame_num = predicted.size(1)
    dist = torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 2)

    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name]['p1'].update(torch.mean(dist).item()*batch_num*frame_num, batch_num*frame_num)
    else:
        for i in range(batch_num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name]['p1'].update(torch.mean(dist[i]).item()*frame_num, frame_num)
            
    return action_error_sum


def mpjpe_by_action_p2(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    pred = predicted.detach().cpu().numpy().reshape(-1, predicted.shape[-2], predicted.shape[-1])
    gt = target.detach().cpu().numpy().reshape(-1, target.shape[-2], target.shape[-1])
    dist = p_mpjpe(pred, gt)
    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]
        action_error_sum[action_name]['p2'].update(np.mean(dist) * num, num)
    else:
        for i in range(num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]
            action_error_sum[action_name]['p2'].update(np.mean(dist), 1)
            
    return action_error_sum


def p_mpjpe(predicted, target):
    assert predicted.shape == target.shape

    muX = np.mean(target, axis=1, keepdims=True)
    muY = np.mean(predicted, axis=1, keepdims=True)

    X0 = target - muX
    Y0 = predicted - muY

    normX = np.sqrt(np.sum(X0 ** 2, axis=(1, 2), keepdims=True))
    normY = np.sqrt(np.sum(Y0 ** 2, axis=(1, 2), keepdims=True))

    X0 /= normX
    Y0 /= normY

    H = np.matmul(X0.transpose(0, 2, 1), Y0)
    U, s, Vt = np.linalg.svd(H)
    V = Vt.transpose(0, 2, 1)
    R = np.matmul(V, U.transpose(0, 2, 1))

    sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))
    V[:, :, -1] *= sign_detR
    s[:, -1] *= sign_detR.flatten()
    R = np.matmul(V, U.transpose(0, 2, 1))

    tr = np.expand_dims(np.sum(s, axis=1, keepdims=True), axis=2)

    a = tr * normX / normY
    t = muX - a * np.matmul(muY, R)

    predicted_aligned = a * np.matmul(predicted, R) + t

    return np.mean(np.linalg.norm(predicted_aligned - target, axis=len(target.shape) - 1), axis=len(target.shape) - 2)


def define_actions( action ):

  actions = ["Directions","Discussion","Eating","Greeting",
           "Phoning","Photo","Posing","Purchases",
           "Sitting","SittingDown","Smoking","Waiting",
           "WalkDog","Walking","WalkTogether"]

  if action == "All" or action == "all" or action == '*':
    return actions

  if not action in actions:
    raise( ValueError, "Unrecognized action: %s" % action )

  return [action]


def define_error_list(actions):
    error_sum = {}
    error_sum.update({actions[i]: {'p1':AccumLoss(), 'p2':AccumLoss()} for i in range(len(actions))})
    return error_sum


class AccumLoss(object):
    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val
        self.count += n
        self.avg = self.sum / self.count
        

def get_varialbe(split, target):
    num = len(target)
    var = []
    if split == 'train':
        for i in range(num):
            temp = Variable(target[i], requires_grad=False).contiguous().type(torch.cuda.FloatTensor)
            var.append(temp)
    else:
        for i in range(num):
            temp = Variable(target[i]).contiguous().cuda().type(torch.cuda.FloatTensor)
            var.append(temp)

    return var


def print_error(data_type, action_error_sum, is_train):
    mean_error_p1, mean_error_p2 = print_error_action(action_error_sum, is_train)

    return mean_error_p1, mean_error_p2


def print_error_action(action_error_sum, is_train):
    mean_error_each = {'p1': 0.0, 'p2': 0.0}
    mean_error_all = {'p1': AccumLoss(), 'p2': AccumLoss()}

    if is_train == 0:
        print("{0:=^12} {1:=^10} {2:=^8}".format("Action", "p#1 mm", "p#2 mm"))

    for action, value in action_error_sum.items():
        if is_train == 0:
            print("{0:<12} ".format(action), end="")
            
        mean_error_each['p1'] = action_error_sum[action]['p1'].avg * 1000.0
        mean_error_all['p1'].update(mean_error_each['p1'], 1)

        mean_error_each['p2'] = action_error_sum[action]['p2'].avg * 1000.0
        mean_error_all['p2'].update(mean_error_each['p2'], 1)

        if is_train == 0:
            print("{0:>6.2f} {1:>10.2f}".format(mean_error_each['p1'], mean_error_each['p2']))

    if is_train == 0:
        print("{0:<12} {1:>6.2f} {2:>10.2f}".format("Average", mean_error_all['p1'].avg, \
            mean_error_all['p2'].avg))
    
    return mean_error_all['p1'].avg, mean_error_all['p2'].avg


def save_model(previous_name, save_dir,epoch, data_threshold, model, model_name):
    # if os.path.exists(previous_name):
    #     os.remove(previous_name)

    torch.save(model.state_dict(),
               '%s/%s_%d_%d.pth' % (save_dir, model_name, epoch, data_threshold * 100))

    previous_name = '%s/%s_%d_%d.pth' % (save_dir, model_name, epoch, data_threshold * 100)
    return previous_name
    
def save_model_new(save_dir,epoch, data_threshold, lr, optimizer, model, model_name):
    # if os.path.exists(previous_name):
    #     os.remove(previous_name)

    # torch.save(model.state_dict(),
    #            '%s/%s_%d_%d.pth' % (save_dir, model_name, epoch, data_threshold * 100))
    torch.save({
                'epoch': epoch,
                'lr': lr,
                'optimizer': optimizer.state_dict(),
                'model_pos': model.state_dict(),
            },
               '%s/%s_%d_%d.pth' % (save_dir, model_name, epoch, data_threshold * 100))









[((81, 10, 81, 48), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((82, 10, 82, 51), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((94, 15, 94, 31), 'numpy.linalg.svd', 'np.linalg.svd', ({(94, 29, 94, 30): 'H'}, {}), '(H)', True, 'import numpy as np\n'), ((87, 20, 87, 63), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((88, 20, 88, 63), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((103, 24, 103, 56), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((98, 39, 98, 55), 'numpy.linalg.det', 'np.linalg.det', ({(98, 53, 98, 54): 'R'}, {}), '(R)', True, 'import numpy as np\n'), ((106, 18, 106, 35), 'numpy.matmul', 'np.matmul', ({(106, 28, 106, 31): 'muY', (106, 33, 106, 34): 'R'}, {}), '(muY, R)', True, 'import numpy as np\n'), ((108, 28, 108, 51), 'numpy.matmul', 'np.matmul', ({(108, 38, 108, 47): 'predicted', (108, 49, 108, 50): 'R'}, {}), '(predicted, R)', True, 'import numpy as np\n'), ((65, 51, 65, 64), 'numpy.mean', 'np.mean', ({(65, 59, 65, 63): 'dist'}, {}), '(dist)', True, 'import numpy as np\n'), ((73, 55, 73, 68), 'numpy.mean', 'np.mean', ({(73, 63, 73, 67): 'dist'}, {}), '(dist)', True, 'import numpy as np\n'), ((39, 51, 39, 67), 'torch.mean', 'torch.mean', ({(39, 62, 39, 66): 'dist'}, {}), '(dist)', False, 'import torch\n'), ((48, 55, 48, 74), 'torch.mean', 'torch.mean', ({(48, 66, 48, 73): 'dist[i]'}, {}), '(dist[i])', False, 'import torch\n'), ((154, 19, 154, 59), 'torch.autograd.Variable', 'Variable', (), '', False, 'from torch.autograd import Variable\n'), ((158, 19, 158, 38), 'torch.autograd.Variable', 'Variable', ({(158, 28, 158, 37): 'target[i]'}, {}), '(target[i])', False, 'from torch.autograd import Variable\n')]