import sys, os
cur_file_path = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(cur_file_path, '..'))

import shutil, glob
import os.path as osp
import cv2
import numpy as np
import json
import torch

from body_model.body_model import BodyModel

from utils.transforms import rotation_matrix_to_angle_axis, batch_rodrigues, convert_to_rotmat
from utils.logging import mkdir, Logger

NSTAGES = 3 # number of stages in the optimization
DEFAULT_FOCAL_LEN = (1060.531764702488, 1060.3856705041237) # fx, fy

def read_keypoints(keypoint_fn):
    '''
    Only reads body keypoint data of first person.
    '''
    with open(keypoint_fn) as keypoint_file:
        data = json.load(keypoint_file)

    if len(data['people']) == 0:
        print('WARNING: Found no keypoints in %s! Returning zeros!' % (keypoint_fn))
        return np.zeros((OP_NUM_JOINTS, 3), dtype=np.float)

    person_data = data['people'][0]
    body_keypoints = np.array(person_data['pose_keypoints_2d'],
                                dtype=np.float)
    body_keypoints = body_keypoints.reshape([-1, 3])

    return body_keypoints

def resize_points(points_arr, num_pts):
    '''
    Either randomly subsamples or pads the given points_arr to be of the desired size.
    - points_arr : N x 3
    - num_pts : desired num point
    '''
    is_torch = isinstance(points_arr, torch.Tensor)
    N = points_arr.size(0) if is_torch else points_arr.shape[0]
    if N > num_pts:
        samp_inds = np.random.choice(np.arange(N), size=num_pts, replace=False)
        points_arr = points_arr[samp_inds]
    elif N < num_pts:
        while N < num_pts:
            pad_size = num_pts - N
            if is_torch:
                points_arr = torch.cat([points_arr, points_arr[:pad_size]], dim=0)
                N = points_arr.size(0)
            else:
                points_arr = np.concatenate([points_arr, points_arr[:pad_size]], axis=0)
                N = points_arr.shape[0]
    return points_arr

def compute_plane_intersection(point, direction, plane):
    '''
    Given a ray defined by a point in space and a direction, compute the intersection point with the given plane.
    Detect intersection in either direction or -direction so the given ray may not actually intersect with the plane.

    Returns the intersection point as well as s such that point + s*direction = intersection_point. if s < 0 it means
    -direction intersects.

    - point : B x 3
    - direction : B x 3
    - plane : B x 4 (a, b, c, d) where (a, b, c) is the normal and (d) the offset.
    '''
    plane_normal = plane[:,:3]
    plane_off = plane[:,3]
    s = (plane_off - bdot(plane_normal, point)) / bdot(plane_normal, direction)
    itsct_pt = point + s.reshape((-1, 1))*direction
    return itsct_pt, s

def bdot(A1, A2, keepdim=False):
    '''
    Batched dot product.
    - A1 : B x D
    - A2 : B x D.
    Returns B.
    '''
    return (A1*A2).sum(dim=-1, keepdim=keepdim)

def parse_floor_plane(floor_plane):
    '''
    Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
    (a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
    '''
    floor_offset = torch.norm(floor_plane, dim=1, keepdim=True)
    floor_normal = floor_plane / floor_offset

    # in camera system -y is up, so floor plane normal y component should never be positive
    #       (assuming the camera is not sideways or upside down)
    neg_mask = floor_normal[:,1:2] > 0.0
    floor_normal = torch.where(neg_mask.expand_as(floor_normal), -floor_normal, floor_normal)
    floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
    floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=1)

    return floor_plane_4d

def load_planercnn_res(res_path):
    '''
    Given a directory containing PlaneRCNN plane detection results, loads the first image result
    and heuristically finds and returns the floor plane.
    '''
    planes_param_path = glob.glob(res_path + '/*_plane_parameters_*.npy')[0]
    planes_mask_path = glob.glob(res_path + '/*_plane_masks_*.npy')[0]
    planes_params = np.load(planes_param_path)
    planes_masks = np.load(planes_mask_path)

    # heuristically determine the ground plane
    #   the plane with the most labeled pixels in the bottom N rows
    nrows = 10
    label_count = np.sum(planes_masks[:, -nrows:, :], axis=(1, 2))
    floor_idx = np.argmax(label_count)
    valid_floor = False
    floor_plane = None
    while not valid_floor:
        # loop until we find a plane with many pixels on the bottom
        #       and doesn't face in the complete wrong direction
        # we assume the y component is larger than any others
        # i.e. that the floor is not > 45 degrees relative rotation from the camera
        floor_plane = planes_params[floor_idx]
        # transform to our system
        floor_plane = np.array([floor_plane[0], -floor_plane[2], floor_plane[1]])
        # determine 4D parameterization
        # for this data we know y should always be negative
        floor_offset = np.linalg.norm(floor_plane)
        floor_normal = floor_plane / floor_offset
        if floor_normal[1] > 0.0:
            floor_offset *= -1.0
            floor_normal *= -1.0
        a, b, c = floor_normal
        d = floor_offset
        floor_plane = np.array([a, b, c, d])

        valid_floor = np.abs(b) > np.abs(a) and np.abs(b) > np.abs(c)
        if not valid_floor:
            label_count[floor_idx] = 0
            floor_idx = np.argmax(label_count)

    return floor_plane

def compute_cam2prior(floor_plane, trans, root_orient, joints):
    '''
    Computes rotation and translation from the camera frame to the canonical coordinate system
    used by the motion and initial state priors.
    - floor_plane : B x 3
    - trans : B x 3
    - root_orient : B x 3
    - joints : B x J x 3
    '''
    B = floor_plane.size(0)
    if floor_plane.size(1) == 3:
        floor_plane_4d = parse_floor_plane(floor_plane)
    else:
        floor_plane_4d = floor_plane
    floor_normal = floor_plane_4d[:,:3]
    floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)

    # compute prior frame axes within the camera frame
    # up is the floor_plane normal
    up_axis = floor_normal
    # right is body -x direction projected to floor plane
    root_orient_mat = batch_rodrigues(root_orient)
    body_right = -root_orient_mat[:, :, 0]
    floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
    right_axis = floor_body_right - floor_trans
    # body right may not actually intersect - in this case must negate axis because we have the -x
    right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
    right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
    # forward is their cross product
    fwd_axis = torch.cross(up_axis, right_axis)
    fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)

    prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
    cam2prior_R = prior_R.transpose(2, 1)

    # translation takes translation to origin plus offset to the floor
    cam2prior_t = -trans

    _, s_root = compute_plane_intersection(joints[:,0], -floor_normal, floor_plane_4d)
    root_height = s_root.reshape((B, 1))

    return cam2prior_R, cam2prior_t, root_height

def apply_robust_weighting(res, robust_loss_type='bisquare', robust_tuning_const=4.6851):
    '''
    Returns robustly weighted squared residuals.
    - res : torch.Tensor (B x N), take the MAD over each batch dimension independently.
    '''
    robust_choices = ['none', 'bisquare']
    if robust_loss_type not in robust_choices:
        print('Not a valid robust loss: %s. Please use %s' % (robust_loss_type, str(robust_choices)))

    w = None
    detach_res = res.clone().detach() # don't want gradients flowing through the weights to avoid degeneracy
    if robust_loss_type == 'none':
        w = torch.ones_like(detach_res)
    elif robust_loss_type == 'bisquare':
        w = bisquare_robust_weights(detach_res, tune_const=robust_tuning_const)

    # apply weights to squared residuals
    weighted_sqr_res = w * (res**2)
    return weighted_sqr_res, w

def robust_std(res):
    '''
    Compute robust estimate of standarad deviation using median absolute deviation (MAD)
    of the given residuals independently over each batch dimension.

    - res : (B x N)

    Returns:
    - std : B x 1
    '''
    B = res.size(0)
    med = torch.median(res, dim=-1)[0].reshape((B,1))
    abs_dev = torch.abs(res - med)
    MAD = torch.median(abs_dev, dim=-1)[0].reshape((B, 1))
    std = MAD / 0.67449
    return std

def bisquare_robust_weights(res, tune_const=4.6851):
    '''
    Bisquare (Tukey) loss.
    See https://www.mathworks.com/help/curvefit/least-squares-fitting.html

    - residuals
    '''
    # print(res.size())
    norm_res = res / (robust_std(res) * tune_const)
    # NOTE: this should use absolute value, it's ok right now since only used for 3d point cloud residuals
        #   which are guaranteed positive, but generally this won't work)
    outlier_mask = norm_res >= 1.0

    # print(torch.sum(outlier_mask))
    # print('Outlier frac: %f' % (float(torch.sum(outlier_mask)) / res.size(1)))

    w = (1.0 - norm_res**2)**2
    w[outlier_mask] = 0.0

    return w

def gmof(res, sigma):
    """
    Geman-McClure error function
    - residual
    - sigma scaling factor
    """
    x_squared = res ** 2
    sigma_squared = sigma ** 2
    return (sigma_squared * x_squared) / (sigma_squared + x_squared)

def log_cur_stats(stats_dict, loss, iter=None):
    Logger.log('LOSS: %f' % (loss.cpu().item()))
    print('----')
    for k, v in stats_dict.items():
        if isinstance(v, float):
            Logger.log('%s: %f' % (k, v))
        else:
            Logger.log('%s: %f' % (k, v.cpu().item()))
    if iter is not None:
        print('======= iter %d =======' % (int(iter)))
    else:
        print('========')

def save_optim_result(cur_res_out_paths, optim_result, per_stage_results, gt_data, observed_data, data_type,
                      optim_floor=True,
                      obs_img_paths=None,
                      obs_mask_paths=None):
    # final optim results
    res_betas = optim_result['betas'].cpu().numpy()
    res_trans = optim_result['trans'].cpu().numpy()
    res_root_orient = optim_result['root_orient'].cpu().numpy()
    res_body_pose = optim_result['pose_body'].cpu().numpy()
    res_contacts = None
    res_floor_plane = None
    if 'contacts' in optim_result:
        res_contacts = optim_result['contacts'].cpu().numpy()
    if 'floor_plane' in optim_result:
        res_floor_plane = optim_result['floor_plane'].cpu().numpy()
    for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
        cur_res_out_path = os.path.join(cur_res_out_path, 'stage3_results.npz')
        save_dict = {
            'betas' : res_betas[bidx],
            'trans' : res_trans[bidx],
            'root_orient' : res_root_orient[bidx],
            'pose_body' : res_body_pose[bidx]
        }
        if res_contacts is not None:
            save_dict['contacts'] = res_contacts[bidx]
        if res_floor_plane is not None:
            save_dict['floor_plane'] = res_floor_plane[bidx]
        np.savez(cur_res_out_path, **save_dict)

    # in prior coordinate frame
    if 'stage3' in per_stage_results and optim_floor:
        res_trans = per_stage_results['stage3']['prior_trans'].detach().cpu().numpy()
        res_root_orient = per_stage_results['stage3']['prior_root_orient'].detach().cpu().numpy()
        for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
            cur_res_out_path = os.path.join(cur_res_out_path, 'stage3_results_prior.npz')
            save_dict = {
                'betas' : res_betas[bidx],
                'trans' : res_trans[bidx],
                'root_orient' : res_root_orient[bidx],
                'pose_body' : res_body_pose[bidx]
            }
            if res_contacts is not None:
                save_dict['contacts'] = res_contacts[bidx]
            np.savez(cur_res_out_path, **save_dict)

    # ground truth
    save_gt = 'betas' in gt_data and \
                'trans' in gt_data and \
                'root_orient' in gt_data and \
                'pose_body' in gt_data
    if save_gt:
        gt_betas = gt_data['betas'].cpu().numpy()
        if data_type not in ['PROX-RGB', 'PROX-RGBD']:
            gt_betas = gt_betas[:,0] # only need frame 1 for e.g. 3d data since it's the same over time.
        gt_trans = gt_data['trans'].cpu().numpy()
        gt_root_orient = gt_data['root_orient'].cpu().numpy()
        gt_body_pose = gt_data['pose_body'].cpu().numpy()
        gt_contacts = None
        if 'contacts' in gt_data:
            gt_contacts = gt_data['contacts'].cpu().numpy()
        cam_mat = None
        if 'cam_matx' in gt_data:
            cam_mat = gt_data['cam_matx'].cpu().numpy()
        for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
            gt_res_name = 'proxd_results.npz' if data_type in ['PROX-RGB', 'PROX-RGBD'] else 'gt_results.npz'
            cur_gt_out_path = os.path.join(cur_res_out_path, gt_res_name)
            save_dict = {
                'betas' : gt_betas[bidx],
                'trans' : gt_trans[bidx],
                'root_orient' : gt_root_orient[bidx],
                'pose_body' : gt_body_pose[bidx]
            }
            if gt_contacts is not None:
                save_dict['contacts'] = gt_contacts[bidx]
            if cam_mat is not None:
                save_dict['cam_mtx'] = cam_mat[bidx]
            np.savez(cur_gt_out_path, **save_dict)

            # if these are proxd results also need to save a GT with cam matrix
            if data_type in ['PROX-RGB', 'PROX-RGBD']:
                cur_gt_out_path = os.path.join(cur_res_out_path, 'gt_results.npz')
                np.savez(cur_gt_out_path, cam_mtx=cam_mat[bidx])

    elif 'joints3d' in gt_data:
        # don't have smpl params, but have 3D joints (e.g. imapper)
        gt_joints = gt_data['joints3d'].cpu().numpy()
        cam_mat = occlusions = None
        if 'cam_matx' in gt_data:
            cam_mat = gt_data['cam_matx'].cpu().numpy()
        if 'occlusions' in gt_data:
            occlusions = gt_data['occlusions'].cpu().numpy()
        for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
            cur_res_out_path = os.path.join(cur_res_out_path, 'gt_results.npz')
            save_dict = {
                'joints3d' : gt_joints[bidx]
            }
            if cam_mat is not None:
                save_dict['cam_mtx'] = cam_mat[bidx]
            if occlusions is not None:
                save_dict['occlusions'] = occlusions[bidx]
            np.savez(cur_res_out_path, **save_dict)
    elif 'cam_matx' in gt_data:
        # need the intrinsics even if we have nothing else
        cam_mat = gt_data['cam_matx'].cpu().numpy()
        for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
            cur_res_out_path = os.path.join(cur_res_out_path, 'gt_results.npz')
            save_dict = {
                'cam_mtx' : cam_mat[bidx]
            }
            np.savez(cur_res_out_path, **save_dict)

    # observations
    obs_out = {k : v.cpu().numpy() for k, v in observed_data.items() if k != 'prev_batch_overlap_res'}
    for bidx, cur_res_out_path in enumerate(cur_res_out_paths):
        obs_out_path = os.path.join(cur_res_out_path, 'observations.npz')
        cur_obs_out = {k : v[bidx] for k, v in obs_out.items() if k not in ['RGB']}
        if obs_img_paths is not None:
            cur_obs_out['img_paths'] = [frame_tup[bidx] for frame_tup in obs_img_paths]
            # print(cur_obs_out['img_paths'])
        if obs_mask_paths is not None:
            cur_obs_out['mask_paths'] = [frame_tup[bidx] for frame_tup in obs_mask_paths]
        np.savez(obs_out_path, **cur_obs_out)

def save_rgb_stitched_result(seq_intervals, all_res_out_paths, res_out_path, device,
                                body_model_path, num_betas, use_joints2d):
    import cv2
    seq_overlaps = [0]
    for int_idx in range(len(seq_intervals)-1):
        prev_end = seq_intervals[int_idx][1]
        cur_start = seq_intervals[int_idx+1][0]
        seq_overlaps.append(prev_end - cur_start)

    # if arbitray RGB video data, stitch together to save full sequence output
    all_res_dirs = all_res_out_paths
    print(all_res_dirs)

    final_res_out_path = os.path.join(res_out_path, 'final_results')
    mkdir(final_res_out_path)

    concat_cam_res = None
    concat_contacts = None
    concat_ground_planes = None
    concat_joints2d = None
    concat_img_paths = None
    gt_cam_mtx = None
    for res_idx, res_dir in enumerate(all_res_dirs):
        # camera view
        cur_stage3_res = load_res(res_dir, 'stage3_results.npz')
        cur_contacts = torch.Tensor(cur_stage3_res['contacts']).to(device)
        if concat_ground_planes is None:
            concat_ground_planes = torch.Tensor(cur_stage3_res['floor_plane']).to(device).reshape((1, -1))
        else:
            concat_ground_planes = torch.cat([concat_ground_planes, torch.Tensor(cur_stage3_res['floor_plane']).to(device).reshape((1, -1))], dim=0)
        cur_stage3_res = {k : v for k, v in cur_stage3_res.items() if k in ['betas', 'trans', 'root_orient', 'pose_body']}
        cur_stage3_res = prep_res(cur_stage3_res, device, cur_stage3_res['trans'].shape[0])
        if concat_cam_res is None:
            concat_cam_res = cur_stage3_res
            concat_contacts = cur_contacts
        else:
            for k, v in concat_cam_res.items():
                concat_cam_res[k] = torch.cat([concat_cam_res[k], cur_stage3_res[k][seq_overlaps[res_idx]:]], dim=0)
            concat_contacts = torch.cat([concat_contacts, cur_contacts[seq_overlaps[res_idx]:]], dim=0)

        # gt
        if gt_cam_mtx is None:
            gt_res = load_res(res_dir, 'gt_results.npz')
            gt_cam_mtx = gt_res['cam_mtx']

        # obs
        cur_obs = load_res(res_dir, 'observations.npz')
        if concat_joints2d is None:
            concat_joints2d = cur_obs['joints2d']
        else:
            concat_joints2d = np.concatenate([concat_joints2d, cur_obs['joints2d'][seq_overlaps[res_idx]:]], axis=0)
        if concat_img_paths is None:
            concat_img_paths = list(cur_obs['img_paths'])
        else:
            concat_img_paths = concat_img_paths + list(cur_obs['img_paths'][seq_overlaps[res_idx]:])

        # ignore if we don't have an interval for this directory (was an extra due to even batching requirement)
        if res_idx >= len(seq_overlaps):
            break

    # copy meta
    src_meta_path = os.path.join(all_res_dirs[0], 'meta.txt')
    shutil.copyfile(src_meta_path, os.path.join(final_res_out_path, 'meta.txt'))

    #  gt results (cam matx)
    np.savez(os.path.join(final_res_out_path, 'gt_results.npz'), cam_mtx=gt_cam_mtx)

    # obs results (joints2d and img_paths)
    np.savez(os.path.join(final_res_out_path, 'observations.npz'), joints2d=concat_joints2d, img_paths=concat_img_paths)

    #  save the actual results npz for viz later
    concat_res_out_path = os.path.join(final_res_out_path, 'stage3_results.npz')
    res_betas = concat_cam_res['betas'].clone().detach().cpu().numpy()
    res_trans = concat_cam_res['trans'].clone().detach().cpu().numpy()
    res_root_orient = concat_cam_res['root_orient'].clone().detach().cpu().numpy()
    res_body_pose = concat_cam_res['pose_body'].clone().detach().cpu().numpy()
    res_floor_plane = concat_ground_planes[0].clone().detach().cpu().numpy() # NOTE: saves estimate from first subsequence
    res_contacts = concat_contacts.clone().detach().cpu().numpy()
    np.savez(concat_res_out_path, betas=res_betas,
                                trans=res_trans,
                                root_orient=res_root_orient,
                                pose_body=res_body_pose,
                                floor_plane=res_floor_plane,
                                contacts=res_contacts)

    # get body model
    num_viz_frames = concat_cam_res['trans'].size(0)
    viz_body_model = BodyModel(bm_path=body_model_path,
                            num_betas=num_betas,
                            batch_size=num_viz_frames,
                            use_vtx_selector=use_joints2d).to(device)
    viz_body = run_smpl(concat_cam_res, viz_body_model)

    # transform full camera-frame sequence into a shared prior frame based on a single ground plane
    viz_joints3d = viz_body.Jtr
    # compute the transformation based on t=0 and the first sequence floor plane
    cam2prior_R, cam2prior_t, cam2prior_root_height = compute_cam2prior(concat_ground_planes[0].unsqueeze(0),
                                                                        concat_cam_res['trans'][0].unsqueeze(0),
                                                                        concat_cam_res['root_orient'][0].unsqueeze(0),
                                                                        viz_joints3d[0].unsqueeze(0))
    # transform the whole sequence
    input_data_dict = {kb : vb.unsqueeze(0) for kb, vb in concat_cam_res.items() if kb in ['trans', 'root_orient', 'pose_body', 'betas']}
    viz_prior_data_dict = apply_cam2prior(input_data_dict, cam2prior_R, cam2prior_t, cam2prior_root_height,
                                            input_data_dict['pose_body'],
                                            input_data_dict['betas'],
                                            0,
                                            viz_body_model)
    concat_prior_res = {
        'trans' : viz_prior_data_dict['trans'][0],
        'root_orient' : viz_prior_data_dict['root_orient'][0],
        'pose_body' : concat_cam_res['pose_body'],
        'betas' : concat_cam_res['betas']
    }

    # save pose prior frame
    concat_prior_res_out_path = os.path.join(final_res_out_path, 'stage3_results_prior.npz')
    res_betas = concat_prior_res['betas'].clone().detach().cpu().numpy()
    res_trans = concat_prior_res['trans'].clone().detach().cpu().numpy()
    res_root_orient = concat_prior_res['root_orient'].clone().detach().cpu().numpy()
    res_body_pose = concat_prior_res['pose_body'].clone().detach().cpu().numpy()
    res_contacts = concat_contacts.clone().detach().cpu().numpy()
    np.savez(concat_prior_res_out_path, betas=res_betas,
                                trans=res_trans,
                                root_orient=res_root_orient,
                                pose_body=res_body_pose,
                                contacts=res_contacts)

def load_res(result_dir, file_name):
    '''
    Load np result from our model or GT
    '''
    res_path = os.path.join(result_dir, file_name)
    if not os.path.exists(res_path):
        return None
    res = np.load(res_path)
    res_dict = {k : res[k] for k in res.files}
    return res_dict

def prep_res(np_res, device, T):
    '''
    Load np result dict into dict of torch objects for use with SMPL body model.
    '''
    betas = np_res['betas']
    betas = torch.Tensor(betas).to(device)
    if len(betas.size()) == 1:
        num_betas = betas.size(0)
        betas = betas.reshape((1, num_betas)).expand((T, num_betas))
    else:
        num_betas = betas.size(1)
        assert(betas.size(0) == T)
    trans = np_res['trans']
    trans = torch.Tensor(trans).to(device)
    root_orient = np_res['root_orient']
    root_orient = torch.Tensor(root_orient).to(device)
    pose_body = np_res['pose_body']
    pose_body = torch.Tensor(pose_body).to(device)

    res_dict = {
        'betas' : betas,
        'trans' : trans,
        'root_orient' : root_orient,
        'pose_body' : pose_body
    }

    for k, v in np_res.items():
        if k not in ['betas', 'trans', 'root_orient', 'pose_body']:
            res_dict[k] = v
    return res_dict

def run_smpl(res_dict, body_model):
    smpl_body = body_model(pose_body=res_dict['pose_body'],
                            pose_hand=None,
                            betas=res_dict['betas'],
                            root_orient=res_dict['root_orient'],
                            trans=res_dict['trans'])
    return smpl_body

def apply_cam2prior(data_dict, R, t, root_height, body_pose, betas, key_frame_idx, body_model, inverse=False):
    '''
    Applies the camera2prior tranformation made up of R, t to the data in data dict and
    returns a new dictionary with the transformed data.
    Right now supports: trans, root_orient.

    NOTE: If the number of timesteps in trans/root_orient is 1, this function assumes they are at key_frame_idx.
            (othherwise the calculation of cur_root_height or trans_offset in inverse case is not correct)

    key_frame_idx : the timestep used to compute cam2prior size (B) tensor
    inverse : if true, applies the inverse transformation from prior space to camera
    '''
    prior_dict = dict()
    if 'root_orient' in data_dict:
        # B x T x 3
        root_orient = data_dict['root_orient']
        B, T, _ = root_orient.size()
        R_time = R.unsqueeze(1).expand((B, T, 3, 3))
        t_time = t.unsqueeze(1).expand((B, T, 3))
        root_orient_mat = batch_rodrigues(root_orient.reshape((-1, 3))).reshape((B, T, 3, 3))
        if inverse:
            prior_root_orient_mat = torch.matmul(R_time.transpose(3, 2), root_orient_mat)
        else:
            prior_root_orient_mat = torch.matmul(R_time, root_orient_mat)
        prior_root_orient = rotation_matrix_to_angle_axis(prior_root_orient_mat.reshape((B*T, 3, 3))).reshape((B, T, 3))
        prior_dict['root_orient'] = prior_root_orient

    if 'trans' in data_dict and 'root_orient' in data_dict:
        # B x T x 3
        trans = data_dict['trans']
        B, T, _ = trans.size()
        R_time = R.unsqueeze(1).expand((B, T, 3, 3))
        t_time = t.unsqueeze(1).expand((B, T, 3))
        if inverse:
            # transform so key frame at origin
            if T > 1:
                trans_offset = trans[np.arange(B),key_frame_idx,:].unsqueeze(1)
            else:
                trans_offset = trans[:,0:1,:]
            trans = trans - trans_offset
            # rotates to camera frame
            trans = torch.matmul(R_time.transpose(3, 2), trans.reshape((B, T, 3, 1)))[:,:,:,0]
            # translate to camera frame
            trans = trans - t_time
        else:
            # first transform so the trans of key frame is at origin
            trans = trans + t_time
            # then rotate to canonical frame
            trans = torch.matmul(R_time, trans.reshape((B, T, 3, 1)))[:,:,:,0]
            # then apply floor offset so the root joint is at the desired height
            cur_smpl_body = body_model(pose_body=body_pose.reshape((-1, body_pose.size(2))),
                                    pose_hand=None,
                                    betas=betas.reshape((-1, betas.size(2))),
                                    root_orient=prior_dict['root_orient'].reshape((-1, 3)),
                                    trans=trans.reshape((-1, 3)))
            smpl_joints3d = cur_smpl_body.Jtr.reshape((B, T, -1, 3))
            if T > 1:
                cur_root_height = smpl_joints3d[np.arange(B),key_frame_idx,0,2:3]
            else:
                cur_root_height = smpl_joints3d[:,0,0,2:3]
            height_diff = root_height - cur_root_height
            trans_offset = torch.cat([torch.zeros((B, 2)).to(height_diff), height_diff], axis=1)
            trans = trans + trans_offset.reshape((B, 1, 3))
        prior_dict['trans'] = trans
    elif 'trans' in data_dict:
        Logger.log('Cannot apply cam2prior on translation without root orient data!')
        exit()

    return prior_dict

def perspective_projection(points, rotation, translation,
                           focal_length, camera_center):
    """
    Adapted from https://github.com/mkocabas/VIBE/blob/master/lib/models/spin.py
    This function computes the perspective projection of a set of points.
    Input:
        points (bs, N, 3): 3D points
        rotation (bs, 3, 3): Camera rotation
        translation (bs, 3): Camera translation
        focal_length (bs, 2): Focal length
        camera_center (bs, 2): Camera center
    """
    batch_size = points.shape[0]
    K = torch.zeros([batch_size, 3, 3], device=points.device)
    K[:,0,0] = focal_length[:,0]
    K[:,1,1] = focal_length[:,1]
    K[:,2,2] = 1.
    K[:,:-1, -1] = camera_center

    # Transform points
    points = torch.einsum('bij,bkj->bki', rotation, points)
    points = points + translation.unsqueeze(1)

    # Apply perspective distortion
    projected_points = points / points[:,:,-1].unsqueeze(-1)

    # Apply camera intrinsics
    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)

    return projected_points[:, :, :-1]

OP_NUM_JOINTS = 25
OP_IGNORE_JOINTS = [1, 9, 12] # neck and left/right hip
OP_EDGE_LIST = [[1,8], [1,2], [1,5], [2,3], [3,4], [5,6], [6,7], [8,9], [9,10], [10,11], [8,12], [12,13], [13,14], [1,0], [0,15], [15,17], [0,16], [16,18], [14,19], [19,20], [14,21], [11,22], [22,23], [11,24]]
# indices to map an openpose detection to its flipped version
OP_FLIP_MAP = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]

#
# The following 2 functions are borrowed from VPoser (https://github.com/nghorbani/human_body_prior).
# See their license for usage restrictions.
#
def expid2model(expr_dir):
    from configer import Configer

    if not os.path.exists(expr_dir): raise ValueError('Could not find the experiment directory: %s' % expr_dir)

    best_model_fname = sorted(glob.glob(os.path.join(expr_dir, 'snapshots', '*.pt')), key=os.path.getmtime)[-1]
    try_num = os.path.basename(best_model_fname).split('_')[0]

    print(('Found Trained Model: %s' % best_model_fname))

    default_ps_fname = glob.glob(os.path.join(expr_dir,'*.ini'))[0]
    if not os.path.exists(
        default_ps_fname): raise ValueError('Could not find the appropriate vposer_settings: %s' % default_ps_fname)
    ps = Configer(default_ps_fname=default_ps_fname, work_dir = expr_dir, best_model_fname=best_model_fname)

    return ps, best_model_fname

def load_vposer(expr_dir, vp_model='snapshot'):
    '''
    :param expr_dir:
    :param vp_model: either 'snapshot' to use the experiment folder's code or a VPoser imported module, e.g.
    from human_body_prior.train.vposer_smpl import VPoser, then pass VPoser to this function
    :param if True will load the model definition used for training, and not the one in current repository
    :return:
    '''
    import importlib
    import os
    import torch

    ps, trained_model_fname = expid2model(expr_dir)
    if vp_model == 'snapshot':

        vposer_path = sorted(glob.glob(os.path.join(expr_dir, 'vposer_*.py')), key=os.path.getmtime)[-1]

        spec = importlib.util.spec_from_file_location('VPoser', vposer_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        vposer_pt = getattr(module, 'VPoser')(num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape)
    else:
        vposer_pt = vp_model(num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape)

    vposer_pt.load_state_dict(torch.load(trained_model_fname, map_location='cpu'))
    vposer_pt.eval()

    return vposer_pt, ps


# -*- coding: utf-8 -*-
"""
Created on Wed Dec  4 11:35:36 2019

@author: aaussel
"""

from cells.RE_mdPul import *
from cells.TC_mdPul import *
#from cells.HTC_mdPul import *
from cells.HTC_buffer_mdPul_Destxhe_tests import *
from scipy import signal

prefs.codegen.target = 'numpy'

defaultclock.dt = 0.01*ms
runtime=2*second

start_scope()

def generate_syn(source,target,syntype,connection_pattern,g_i,taur_i,taud_i,V_i):
    eq_syn='''_post=s_i*g_i*(V_post-V_i) : amp * meter ** -2 (summed)
        ds_i/dt=-s_i/taud_i+(1-s_i)/taur_i*0.5*(1+tanh(V_pre/10/mV)) : 1
        g_i : siemens * meter**-2
        V_i : volt
        taud_i : second
        taur_i : second
    '''

    S=Synapses(source,target,model=syntype+eq_syn,method='exact')
    if connection_pattern=='':
        S.connect()
    else :
        S.connect(j=connection_pattern, skip_if_invalid=True)
    S.g_i=g_i
    S.taur_i=taur_i
    S.taud_i=taud_i
    S.V_i=V_i
    return S

def mon_delay(mon,i,t_delay):
    return mon.V[i][-t_delay]

def generate_syn_delay(source,target,syntype,connection_pattern,g_i,taur_i,taud_i,V_i):
    eq_syn_delay='''_post=s_i*g_i*(V_post-V_i) : amp * meter ** -2 (summed)
        ds_i/dt=-s_i/taud_i+(1-s_i)/taur_i*0.5*(1+tanh(voltage_delayed_pre/10/mV)) : 1
        g_i : siemens * meter**-2
        V_i : volt
        taud_i : second
        taur_i : second
    '''

    S=Synapses(source,target,model=syntype+eq_syn_delay,method='exact')
    if connection_pattern=='':
        S.connect()
    else :
        S.connect(j=connection_pattern, skip_if_invalid=True)
    S.g_i=g_i
    S.taur_i=taur_i
    S.taud_i=taud_i
    S.V_i=V_i
    return S

def generate_spike_timing(N,f,start_time,end_time=runtime):
    list_time_and_i=[]
    for i in range(N):
        list_time=[(start_time,i)]
        next_spike=list_time[-1][0]+(1+0.1*rand())/f
        while next_spike<end_time:
            list_time.append((next_spike,i))
            next_spike=list_time[-1][0]+(1+0.1*rand())/f
        list_time_and_i+=list_time
    return array(list_time_and_i)

def create_mdPul_column(N_HTC,N_TC,N_RE,condition,in_mode,theta_phase):
    HTC=NeuronGroup(N_HTC,eq_HTC_buffer_mdPul,threshold='V>0*mvolt',refractory=3*ms,method='rk4')
    HTC.V = '-25*mvolt'
    HTC.Ca_iTLT = '1e-7 * mole * metre**-3'
    HTC.Ca_iTHT = '0.01*mmolar'
    HTC.J='0 * uA * cmeter ** -2'
    HTC.mAHP='0.3'

    TC=NeuronGroup(N_TC,eq_TC_mdPul,threshold='V>0*mvolt',refractory=3*ms,method='rk4')
    TC.V = '-45*mvolt+20*mvolt*rand()'
    TC.J='0 * nA * cmeter ** -2'
    TC.Ca_i = '1e-7 * mole * metre**-3'

    RE=NeuronGroup(N_RE,eq_RE_mdPul,threshold='V>0*mvolt',refractory=3*ms,method='rk4')
    RE.V = '-70*mvolt+20*mvolt*rand()'
    RE.J = '0 * nA * cmeter ** -2'
    RE.Ca_i = '1e-7 * mole * metre**-3'

    ##Define monitors
    R1=SpikeMonitor(HTC,record=True)
    R2=SpikeMonitor(TC,record=True)
    R3=SpikeMonitor(RE,record=True)

    V1=StateMonitor(HTC,'V',record=True)
    V2=StateMonitor(TC,'V',record=True)
    V3=StateMonitor(RE,'V',record=True)

    I2=StateMonitor(HTC,'ITHT',record=[0])
    I3=StateMonitor(HTC,'ITLT',record=[0])

    ##Synapses

    gAMPA_HTC_RE= 0.1 * msiemens * cm **-2 #0.42 - 0.70
    gGABA_A_RE_RE = 0.1 * msiemens * cm **-2 #0.42 - 0.70
    gGABA_A_RE_TC = 0.4 * msiemens * cm **-2
#    gGABA_A_RE_TC = 0 * msiemens * cm **-2
#    gGABA_A_HTC_TC = 0.2 * msiemens * cm **-2
    gGABA_A_HTC_TC = 0.2 * msiemens * cm **-2
    gGABA_B_RE_TC= gGABA_A_RE_TC*0.1
    gGABA_A_RE_HTC = 1 * msiemens * cm **-2
    gGABA_B_RE_HTC= 1 * msiemens * cm **-2
#    gGABA_A_RE_HTC = 0 * msiemens * cm **-2
#    gGABA_B_RE_HTC= 0 * msiemens * cm **-2

    synHTCRE_AMPA=None
    if condition=='mGluR1':
        synHTCRE_AMPA = generate_syn(HTC,RE,'IsynHTC','',gAMPA_HTC_RE,0.125*ms,1*ms,0*mV)
        synHTCRE_AMPA.connect() #all-to-all connection

    synRERE_GABA = generate_syn(RE,RE,'IsynREA','',gGABA_A_RE_RE,0.25*ms,5*ms,-80*mV)
    synRETC_A=generate_syn(RE,TC,'IsynREA','',gGABA_A_RE_TC,0.25*ms,5*ms,-80*mV)
    synRETC_B=generate_syn(RE,TC,'IsynREB','',gGABA_B_RE_TC,20*ms,100*ms,-80*mV)
    synREHTC_A=generate_syn(RE,HTC,'IsynREintGABAA','',gGABA_A_RE_HTC,0.25*ms,5*ms,-80*mV)
    synREHTC_B=generate_syn(RE,HTC,'IsynREintGABAB','',gGABA_B_RE_HTC,20*ms,100*ms,-80*mV)

    synHTCTC=generate_syn_delay(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    synHTCTC=generate_syn(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    if in_mode=='burst':
#        if condition=='mAChR':
#            synHTCTC=generate_syn_delay(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    #        synHTCTC.delay=delayTC
#
#        if condition=='mGluR1':
#            synHTCTC=generate_syn_delay(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    #        synHTCTC.delay=delayTC
#    else :
#        if condition=='mAChR':
#            synHTCTC=generate_syn(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    #        synHTCTC.delay=delayTC
#
#        if condition=='mGluR1':
#            synHTCTC=generate_syn(HTC,TC,'IsynHTC','',gGABA_A_HTC_TC,0.25*ms,5*ms,-80*mV)
#    #        synHTCTC.delay=delayTC

    GJ_HTC=Synapses(HTC,HTC,'''
                 w : siemens * meter **-2 # gap junction conductance
                 IGJ_post = w * (V_post - V_pre) : amp * meter ** -2 (summed)
                 ''')
    GJ_HTC.connect()
    #GJ_HTC.w = 0.003e-3 * siemens * cm **-2
    GJ_HTC.w = 0.08e-3 * siemens * cm **-2

    gEPSP_RE = 20 * msiemens * cm **-2 #0.42 - 0.70
#    gEPSP_TC = 0.5 * msiemens * cm **-2 #0.42 - 0.70
    gEPSP_TC = 0 * msiemens * cm **-2 #0.42 - 0.70
    gEPSP_TC2 = 10 * msiemens * cm **-2 #0.42 - 0.70
    gIPSP_TC = -1.5 * gEPSP_TC
    #Inputs :
    inputs_TC=generate_spike_timing(N_TC,100*Hz,0*ms,end_time=3000*ms)
    G_in = SpikeGeneratorGroup(N_TC, inputs_TC[:,1], inputs_TC[:,0]*second)
    Syn_in=Synapses(G_in,TC,on_pre='Vinp=Vhigh')
    Syn_in.connect(j='i')
    TC.ginp_TC=gEPSP_TC

    inputs_RE=generate_spike_timing(N_RE,13*Hz,0*ms,end_time=3000*ms)
    G_in2 = SpikeGeneratorGroup(N_RE, inputs_RE[:,1], inputs_RE[:,0]*second)
    Syn_in2=Synapses(G_in2,RE,on_pre='Vinp=Vhigh')
    Syn_in2.connect(j='i')

    inputs_TC3=generate_spike_timing(N_TC,13*Hz,0*ms,end_time=3000*ms)
    G_in3 = SpikeGeneratorGroup(N_TC, inputs_TC3[:,1], inputs_TC3[:,0]*second)
    Syn_in3=Synapses(G_in3,TC,on_pre='Vinp=Vhigh')
    Syn_in3.connect(j='i')

#    inputs_HTC=generate_spike_timing(N_HTC,13*Hz,0*ms,end_time=3000*ms)
#    G_in4 = SpikeGeneratorGroup(N_HTC, inputs_HTC[:,1], inputs_HTC[:,0]*second)
#    Syn_in4=Synapses(G_in4,HTC,on_pre='Vinp=Vhigh')
#    Syn_in4.connect(j='i')
#    HTC.ginp_HTC=0e-3 * siemens * cm **-2

    if theta_phase=='bad':
        RE.ginp_RE=gEPSP_RE
        TC.ginp_TC2=gEPSP_TC2

    all_neurons=HTC,TC,RE,G_in,G_in2,G_in3
    all_synapses=synHTCRE_AMPA,synRERE_GABA,synRETC_A,synRETC_B,synREHTC_A,synREHTC_B,synHTCTC,Syn_in,Syn_in2,Syn_in3
    all_synapses=tuple([y for y in all_synapses if y])
    all_monitors=R1,R2,R3,V1,V2,V3,I2,I3
    all_gap_junctions=(GJ_HTC,)

#    if in_mode=='single_spike':
#        HTC.delay_steps = [1]  # delay in time steps per neuron
#        buffer_size = 2  # 1+Maximum delay (in time steps)
#    else :
#        HTC.delay_steps = [3999]  # delay in time steps per neuron
#        buffer_size = 4000  # 1+Maximum delay (in time steps)
#
#    print(HTC.delay_steps)
#    HTC.variables.add_array('voltage_buffer', dimensions=volt.dim, size=(buffer_size, len(HTC)))
#
#    update_code = '''buffer_pointer = (buffer_pointer + 1) % buffer_size
#                     voltage_delayed = update_voltage_buffer(V, voltage_buffer, buffer_pointer, delay_steps, buffer_size)'''
#
#    buffer_updater = HTC.run_regularly(update_code, codeobj_class=NumpyCodeObject)
#
#    @check_units(V=volt, voltage_buffer=volt, buffer_pointer=1, delay_steps=1, buffer_size=1, result=volt)
#    def update_voltage_buffer(V, voltage_buffer, buffer_pointer, delay_steps, buffer_size):
#        # Write current rate into the buffer
#        voltage_buffer[buffer_pointer, :] = V
#        # Get delayed rates
#        rows = (buffer_pointer - delay_steps) % buffer_size
#        return voltage_buffer[rows, arange(len(rows))]

    return all_neurons,all_synapses,all_gap_junctions,all_monitors

if __name__=='__main__':
    close('all')
    runtime=1*second
    f=13*Hz #rythmic input frequency
    input_on=False
    N_HTC,N_TC,N_RE= 20,80,100 #Number of neurons of RE, TC, and HTC type
#    N_HTC,N_TC,N_RE= 1,80,100 #Number of neurons of RE, TC, and HTC type

    Vrev_inp=0*mV
    taurinp=0.1*ms
    taudinp=0.5*ms
    tauinp=taudinp
    Vhigh=0*mV
    Vlow=-80*mV
    ginp_IB=0* msiemens * cm **-2
    ginp=0* msiemens * cm **-2

    Vrev_inp2=0*mV
    taurinp2=0.1*ms
    taudinp2=0.5*ms
    tauinp2=taudinp2
    Vhigh2=0*mV
    Vlow2=-80*mV

    #condition='mGluR1'
    condition='mAChR'
#    in_mode='single_spike'
    in_mode='burst'
    theta_phase='good'

    if condition=='mGluR1':
        gKL_TC=0.0028e-3 * siemens * cm **-2
        gKL_HTC=0.0069e-3 * siemens * cm **-2
        gKL_RE=0.05e-3 * siemens * cm **-2
    elif condition=='mAChR':
        gKL_TC=0.0028e-3 * siemens * cm **-2
        gKL_HTC=0.0069e-3 * siemens * cm **-2
        gKL_RE=0.08e-3 * siemens * cm **-2

#    Cm_HTC = 2.5* ufarad * cm ** -2
#    gNa_HTC=90e-3 * siemens * cm **-2
#    ENa_HTC=50*mV
#    gK_HTC=10e-3 * siemens * cm **-2
#    EK_HTC=-100*mV
#    gL_HTC=0.001e-3 * siemens * cm **-2
#    EL_HTC=-70*mV
    gKL_HTC=0.001e-3 * siemens * cm **-2
#    EKL_HTC=-100*mV
#    gTLT_HTC= 2.1e-3 * siemens * cm **-2
#    gTHT_HTC= 15e-3 * siemens * cm **-2
#    gAHP_HTC= 45e-3 * siemens * cm **-2
#    EAHP=-95*mV
#    gH_HTC = 0.36e-3 * siemens * cm **-2
#    EH_HTC = -40 * mV
    gapp=0.1*mamp * cmeter ** -2 # in HTC cells
#    gapp=0.001*mamp * cmeter ** -2 # in HTC cells

    net=Network()
    all_neurons,all_synapses,all_gap_junctions,all_monitors=create_mdPul_column(N_HTC,N_TC,N_RE,condition,in_mode,theta_phase)
#    R1,R2,R3,V1,V2,V3=all_monitors
    R1,R2,R3,V1,V2,V3,I2,I3=all_monitors

    HTC=all_neurons[0]
    if in_mode=='single_spike':
        HTC.delay_steps = [1]  # delay in time steps per neuron
        buffer_size = 2  # 1+Maximum delay (in time steps)
    else :
        HTC.delay_steps = [3999]  # delay in time steps per neuron
        buffer_size = 4000  # 1+Maximum delay (in time steps)

    init_array=-70*mV*ones((buffer_size, len(HTC)))
    HTC.variables.add_array('voltage_buffer', dimensions=volt.dim, size=(buffer_size, len(HTC)))#,values=init_array)
    HTC.voltage_buffer=init_array

    update_code = '''buffer_pointer = (buffer_pointer + 1) % buffer_size
                     voltage_delayed = update_voltage_buffer(V, voltage_buffer, buffer_pointer, delay_steps, buffer_size)'''

    buffer_updater = HTC.run_regularly(update_code, codeobj_class=NumpyCodeObject)

    @check_units(V=volt, voltage_buffer=volt, buffer_pointer=1, delay_steps=1, buffer_size=1, result=volt)
    def update_voltage_buffer(V, voltage_buffer, buffer_pointer, delay_steps, buffer_size):
        # Write current rate into the buffer
        voltage_buffer[buffer_pointer, :] = V
        # Get delayed rates
        rows = (buffer_pointer - delay_steps) % buffer_size
        return voltage_buffer[rows, arange(len(rows))]

#    print(HTC.delay_steps)

    SynHTCTC=all_synapses[5]

    net.add(all_neurons)
    net.add(all_synapses)
    net.add(all_gap_junctions)
    net.add(all_monitors)

    TC=all_neurons[1]
    M=StateMonitor(TC,'Isyn', record=True)
    M2=StateMonitor(HTC,'voltage_delayed', record=True)
    M3=StateMonitor(SynHTCTC,'s_i',record=[0])
    net.add(M,M2,M3)

    prefs.codegen.target = 'cython'
    net.run(runtime,report='text',report_period=300*second)

    figure()
    plot(R1.t,R1.i+0,'r.',label='HTC')
    plot(R2.t,R2.i+20,'b.',label='TC')
    plot(R3.t,R3.i+100,'g.',label='RE')
    xlim(0,runtime/second)
    legend()

    figure()
    plot(V1.t,V1.V[0],label='HTC V')
    plot(V2.t,V2.V[0],label='TC V')
    plot(V3.t,V3.V[0],label='RE V')
    legend()

#    figure()
#    plot(I1.t,I1.ISK[0],label='I_SK HTC')
#    legend()

#    figure()
#    plot(I2.t,I2.ITHT[0],label='I_THT HTC')
#    legend()
#
#    figure()
#    plot(I3.t,I3.ITLT[0],label='I_TLT HTC')
#    legend()

#    f,Spectrum_LFP_V1=signal.periodogram(V1.V[0], 100000,'flattop', scaling='spectrum')
#    figure()
#    plot(f,Spectrum_LFP_V1)
#    xlim(0,100)
    Isyn_calc=(V2.V[0]+80*mV)*0.2 * msiemens * cm **-2*M3.s_i[0]

    figure()
    plot(M.t,M.Isyn[0],label='I_syn HTC-TC')
    plot(M2.t,M2.voltage_delayed[0],label='HTC voltage_delayed')
    plot(V2.t,V2.V[0],label='TC V')
    plot(V1.t,V1.V[0],label='HTC V')
    plot(M3.t,M3.s_i[0],label='s_i syn HTC-TC')
    plot(M3.t,Isyn_calc,label='Isyn HTC-TC calc')
    legend()

    clear_cache('cython')


"""
Please cite this paper if you re-use parts of the code
S. Koehler, T. Hussain, Z. Blair, T. Huffaker, F. Ritzmann, A. Tandon, T. Pickardt, S. Sarikouch, H. Latus, G. Greil, I. Wolf, S. Engelhardt,
"Unsupervised Domain Adaptation from Axial to Short-Axis Multi-Slice Cardiac MR Images by Incorporating Pretrained Task Networks,"
in IEEE Transactions on Medical Imaging (TMI), early access, doi: 10.1109/TMI.2021.3052972.

The spatial transformer layer is based on the neuron project:
Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
MICCAI 2018.
"""

# main imports
import logging
import sys

# third party
import numpy as np
import tensorflow
import tensorflow as tf
import tensorflow.keras.initializers
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

import src.utils.Loss_and_metrics as metr
from src.models.KerasLayers import Euler2Matrix, UnetWrapper, ConvEncoder, Inverse3DMatrix
from src.models.ModelUtils import get_optimizer

sys.path.append('src/models/ext/neuron')
sys.path.append('src/models/ext/pynd-lib')
sys.path.append('src/models/ext/pytools-lib')
import src.models.ext.neuron.neuron.layers as nrn_layers

def create_affine_cycle_transformer_model(config, networkname='affine_cycle_transformer', unet=None):
    """
    Create a compiled Domain adaption (AX2SAX) spatial transformer model with three loss functions.

    :param config: Key value pairs for image size and other network parameters
    :param metrics: list of tensorflow or keras compatible metrics
    :param networkname: string, name of this model scope
    :param unet: tf.keras.Model, pre-trained 2D U-net
    :return: compiled tf.keras.Model

    The returned tf.keras.Model expects the following input during training:
    [ax_cmr,  sax_cmr]
    and returns the following elements:
    [ax2sax_cmr, sax2ax_cmr, ax2sax_mod_cmr,sax_msk, ax_msk, m, m_mod]
    During inference sax_cmr could be None or a zero initialised ndarray

    This model has the following flow:
    inputs = [AX, SAX]
    m, m_mod = Encoder(AX)
    m_inv = Inverse(m)
    ax2sax = SpatialTransformer(AX, m)
    ax2sax = SpatialTransformer(AX, m_mod)
    sax2ax = SpatialTransformer(SAX, m_inv)
    sax_msk = Unet(SAX)
    ax_msk = SpatialTransformer(sax_msk, m_mod)
    outputs = [ax2sax, sax2ax, ax2sax_mod, sax_msk, ax_msk, m, m_mod]

    This model calculates the gradients according to the following loss functions:
    MSE_mod(ax2sax_gt, ax2sax_pred) # learn an affine transformation that fits to our gt
    MSE_mod(sax2ax_gt, sax2ax_pred) # apply the opposite transformation to align the cycle consistency
    Loss_focus(ax_msk) or Loss_focus(sax_msk)

    - MSE_mod = MSE(y_true, y_pred) * weighting[None, None,: , : ],
    Example in-plane weighting with an in-plane resolution of 10 x 10 pixels:
    weighting =
    [[0.   0.   0.   0.   0.   0.   0.   0.   0.   0.  ]
     [0.   0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.  ]
     [0.   0.25 0.5  0.5  0.5  0.5  0.5  0.5  0.25 0.  ]
     [0.   0.25 0.5  0.75 0.75 0.75 0.75 0.5  0.25 0.  ]
     [0.   0.25 0.5  0.75 1.   1.   0.75 0.5  0.25 0.  ]
     [0.   0.25 0.5  0.75 1.   1.   0.75 0.5  0.25 0.  ]
     [0.   0.25 0.5  0.75 0.75 0.75 0.75 0.5  0.25 0.  ]
     [0.   0.25 0.5  0.5  0.5  0.5  0.5  0.5  0.25 0.  ]
     [0.   0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.  ]
     [0.   0.   0.   0.   0.   0.   0.   0.   0.   0.  ]]

    - Loss_focus =
        # ignore background, we want to maximize the number of captured ventricle voxel
        y_pred = y_pred[...,1:] if background given by the U-Net predictions
        y_pred = tf.cast(y_pred, dtype=tf.float32)
        # keep only the highest prob
        sum_bigger_than = tf.reduce_max(y_pred, axis=-1)
        # create a mask of voxels greater than the threshold
        mask_bigger_than = tf.cast(sum_bigger_than > min_probabillity, tf.float32)
        # we cant use the mask directly as this creates no gradients, this keeps also the prob value
        sum_bigger_than = sum_bigger_than * mask_bigger_than
        # return a scalar between 0 and 1, the loss is typically close to 1, as the background class is overrepresented
        return 1- tf.reduce_mean(sum_bigger_than)
    """

    if tf.distribute.has_strategy():
        strategy = tf.distribute.get_strategy()
    else:
        # distribute the training with the "mirrored data"-paradigm across multiple gpus if available, if not use gpu 0
        strategy = tf.distribute.MirroredStrategy(devices=config.get('GPUS', ["/gpu:0"]))
    with strategy.scope():

        input_shape = config.get('DIM', [10, 224, 224])
        inputs_ax = Input((*input_shape, config.get('IMG_CHANNELS', 1)))
        inputs_sax = Input((*input_shape, config.get('IMG_CHANNELS', 1)))
        # define standard values according to the convention over configuration paradigm
        activation = config.get('ACTIVATION', 'elu')
        batch_norm = config.get('BATCH_NORMALISATION', False)
        pad = config.get('PAD', 'same')
        kernel_init = config.get('KERNEL_INIT', 'he_normal')
        m_pool = config.get('M_POOL', (1, 2, 2))
        f_size = config.get('F_SIZE', (3, 3, 3))
        filters = config.get('FILTERS', 16)
        drop_1 = config.get('DROPOUT_min', 0.3)
        drop_3 = config.get('DROPOUT_max', 0.5)
        bn_first = config.get('BN_FIRST', False)
        ndims = len(config.get('DIM', [10, 224, 224]))
        depth = config.get('DEPTH', 4)
        dense_weights = config.get('DENSE_WEIGHTS', 256)
        indexing = config.get('INDEXING', 'ij')

        weight_mse_inplane = config.get('WEIGHT_MSE_INPLANE',
                                        True)  # weight the MSE loss pixels in the center have greater weights
        mask_smaller_than_threshold = config.get('MASK_SMALLER_THAN_THRESHOLD',
                                                 0.01)  # calc the MSe loss only where our image has values greater than
        ax_weight = config.get('AX_LOSS_WEIGHT', 2)

        cycle_loss = config.get('CYCLE_LOSS', False)
        sax_weight = config.get('SAX_LOSS_WEIGHT', 2)

        focus_loss = config.get('FOCUS_LOSS', False)
        focus_weight = config.get('FOCUS_LOSS_WEIGHT', 1)
        min_unet_probability = config.get('MIN_UNET_PROBABILITY',
                                          0.9)  # sum the foreground voxels with a prob higher than
        use_mask2ax_prob = config.get('USE_SAX2AX_PROB', True)  # otherwise use the SAX probability

        # increase the dropout through the layer depth
        dropouts = list(np.linspace(drop_1, drop_3, depth))
        dropouts = [round(i, 1) for i in dropouts]

        enc, _ = ConvEncoder(activation=activation,
                             batch_norm=batch_norm,
                             bn_first=bn_first,
                             depth=depth,
                             drop_3=drop_3,
                             dropouts=dropouts,
                             f_size=f_size,
                             filters=filters,
                             kernel_init=kernel_init,
                             m_pool=m_pool,
                             ndims=ndims,
                             pad=pad)(inputs_ax)

        # Shrink the encoding towards the euler angles and translation params,
        # no additional dense layers before the GAP layer
        m_raw = tensorflow.keras.layers.GlobalAveragePooling3D()(enc)  # m.shape --> b, 512
        m_raw = tensorflow.keras.layers.Dense(dense_weights, kernel_initializer=kernel_init, activation=activation,
                                              name='dense1')(m_raw)
        m_raw = tensorflow.keras.layers.Dense(9, kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
                                              activation=activation, name='dense2')(m_raw)
        m = Euler2Matrix(name='ax2sax_matrix')(m_raw[:, 0:6])

        # Cycle flow - use M and the inverse M to transform the SAX and AX input
        ax2sax = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing, ident=False, fill_value=0,
                                               name='ax2sax')([inputs_ax, m])
        m_inv = Inverse3DMatrix()(m)
        sax2ax = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing, ident=False, fill_value=0,
                                               name='sax2ax')([inputs_sax, m_inv])

        if unet:
            logging.info('unet given, use it to max probability')

            # concat the rotation parameters with the second set of translation params
            m_mod = tf.keras.layers.Concatenate(axis=-1)([m_raw[:, 0:3], m_raw[:, 6:9]])  # rot + translation
            m_mod = Euler2Matrix(name='ax2sax_mod_matrix')(m_mod)  #
            ax2sax_mod = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing, ident=False,
                                                       fill_value=0, name='ax2sax_mod_st')([inputs_ax, m_mod])

            # we use the probabilities of a pre-trained U-net with fixed weights
            # to learn a second set of translation parameters which maximize the Unet probability
            mask_prob = UnetWrapper(unet, name='mask_prob')(ax2sax_mod)  #
            m_mod_inv = Inverse3DMatrix()(m_mod)  #
            mask2ax = nrn_layers.SpatialTransformer(interp_method='nearest', indexing=indexing, ident=False,
                                                    fill_value=0, name='mask2ax')([mask_prob, m_mod_inv])
            # Define the model output
            outputs = [ax2sax, sax2ax, ax2sax_mod, mask_prob, mask2ax, m, m_mod]

            # baseline loss
            logging.info('adding ax2sax MSE loss with a weighting of {}'.format(ax_weight))
            losses = {'ax2sax': metr.loss_with_zero_mask(mask_smaller_than=mask_smaller_than_threshold,
                                                         weight_inplane=weight_mse_inplane, xy_shape=input_shape[-2])}
            loss_w = {'ax2sax': ax_weight}

            # extend losses by cycle MSE loss
            if cycle_loss:
                logging.info('adding cycle loss with a weighting of {}'.format(sax_weight))
                losses['sax2ax'] = metr.loss_with_zero_mask(mask_smaller_than=mask_smaller_than_threshold,
                                                            weight_inplane=weight_mse_inplane, xy_shape=input_shape[-2])
                loss_w['sax2ax'] = sax_weight

            # extend losses by probability loss
            if focus_loss:
                # Use the SAX predictions or the SAX2AX predictions to maximise the unet probability
                # probability_object must fit a output-layer name
                if use_mask2ax_prob:
                    probability_object = 'mask2ax'
                else:
                    probability_object = 'mask_prob'
                logging.info('adding focus loss on {} with a weighting of {}'.format(probability_object, focus_weight))
                losses[probability_object] = metr.max_volume_loss(min_probability=min_unet_probability)
                loss_w[probability_object] = focus_weight

        else:  # no u-net given
            outputs = [ax2sax, sax2ax, m]
            # baseline loss
            logging.info('adding ax2sax MSE loss with a weighting of {}'.format(ax_weight))
            losses = {'ax2sax': metr.loss_with_zero_mask(mask_smaller_than=mask_smaller_than_threshold,
                                                         weight_inplane=weight_mse_inplane,
                                                         xy_shape=input_shape[-2])}
            loss_w = {'ax2sax': ax_weight}

            # extend losses by cycle MSE loss
            if cycle_loss:
                logging.info('adding cycle MSE loss with a weighting of {}'.format(sax_weight))
                losses['sax2ax'] = metr.loss_with_zero_mask(mask_smaller_than=mask_smaller_than_threshold,
                                                            weight_inplane=weight_mse_inplane,
                                                            xy_shape=input_shape[-2])
                loss_w['sax2ax'] = sax_weight

        model = Model(inputs=[inputs_ax, inputs_sax], outputs=outputs, name=networkname)
        model.compile(optimizer=get_optimizer(config, networkname), loss=losses, loss_weights=loss_w)

        return model

# ST to apply m to an volume
def create_affine_transformer_fixed(config, metrics=None, networkname='affine_transformer_fixed', fill_value=0,
                                    interp_method='linear'):
    """
    Apply a learned transformation matrix to an input image, no training possible
    :param config:  Key value pairs for image size and other network parameters
    :param metrics: list of tensorflow or keras compatible metrics
    :param networkname: string, name of this model scope
    :param fill_value:
    :return: compiled tf.keras model
    """
    if tf.distribute.has_strategy():
        strategy = tf.distribute.get_strategy()
    else:
        # distribute the training with the mirrored data paradigm across multiple gpus if available, if not use gpu 0
        strategy = tf.distribute.MirroredStrategy(devices=config.get('GPUS', ["/gpu:0"]))

    with strategy.scope():

        inputs = Input((*config.get('DIM', [10, 224, 224]), config.get('IMG_CHANNELS', 1)))
        input_matrix = Input((12), dtype=np.float32)
        indexing = config.get('INDEXING', 'ij')

        # warp the source with the flow
        y = nrn_layers.SpatialTransformer(interp_method=interp_method, indexing=indexing, ident=False,
                                          fill_value=fill_value)([inputs, input_matrix])

        model = Model(inputs=[inputs, input_matrix], outputs=[y, input_matrix], name=networkname)

        return model


#! -*- coding: utf-8 -*-
# 自定义层

import numpy as np
import tensorflow as tf
from .backend import keras, K, get_all_attributes

# 等价于 from keras.layers import *
locals().update(get_all_attributes(keras.layers))
# 等价于 from keras.models import Model
locals()['Model'] = keras.models.Model
# 等价于 from keras.utils import get_custom_objects
locals()['get_custom_objects'] = keras.utils.get_custom_objects

def gelu_erf(x):
    # 基于Erf直接计算的gelu函数
    return 0.5 * x * (1.0 + tf.math.erf(x / np.sqrt(2.0)))

def gelu_tanh(x):
    # 基于Tanh近似计算的gelu函数
    cdf = 0.5 * (1.0 + K.tanh(
        (np.sqrt(2 / np.pi) * (x + 0.044715 * K.pow(x, 3)))))
    return x * cdf

def add_seq_mask(x, mask, mode=0, axis=None, heads=1):
    """为序列条件mask的函数
    mask: 形如(batch_size, seq_len)的0-1矩阵；
    mode: 如果是0，则直接乘以mask；
          如果是1，则在padding部分减去一个大正数。
    axis: 序列所在轴，默认为1；
    heads: 相当于batch这一维要被重复的次数。
    """
    if mask is None or mode not in [0, 1]:
        return x
    else:
        if heads is not 1:
            mask = K.expand_dims(mask, 1)
            mask = K.tile(mask, (1, heads, 1))
            mask = K.reshape(mask, (-1, K.shape(mask)[2]))
        if axis is None:
            axis = 1
        if axis == -1:
            axis = K.ndim(x) - 1
        assert axis > 0, 'axis muse be greater than 0'
        for _ in range(axis - 1):
            mask = K.expand_dims(mask, 1)
        for _ in range(K.ndim(x) - K.ndim(mask) - axis + 1):
            mask = K.expand_dims(mask, K.ndim(mask))
        if mode == 0:
            return x * mask
        else:
            return x - (1 - mask) * 1e12

class MultiHeadAttention(Layer):
    """多头注意力机制
    """
    def __init__(self, heads, head_size, key_size=None, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.out_dim = heads * head_size
        self.key_size = key_size if key_size else head_size

    def build(self, input_shape):
        super(MultiHeadAttention, self).build(input_shape)
        self.q_dense = Dense(self.key_size * self.heads)
        self.k_dense = Dense(self.key_size * self.heads)
        self.v_dense = Dense(self.out_dim)
        self.o_dense = Dense(self.out_dim)

    def call(self, inputs, q_mask=False, v_mask=False, a_mask=False):
        """实现多头注意力
        q_mask: 对输入的query序列的mask。
                主要是将输出结果的padding部分置0。
        v_mask: 对输入的value序列的mask。
                主要是防止attention读取到padding信息。
        a_mask: 对attention矩阵的mask。
                不同的attention mask对应不同的应用。
        """
        q, k, v = inputs[:3]
        # 处理mask
        idx = 3
        if q_mask:
            q_mask = inputs[idx]
            idx += 1
        else:
            q_mask = None
        if v_mask:
            v_mask = inputs[idx]
            idx += 1
        else:
            v_mask = None
        if a_mask:
            if len(inputs) > idx:
                a_mask = inputs[idx]
            else:
                a_mask = 'history_only'
        else:
            a_mask = None
        # 线性变换
        qw = self.q_dense(q)
        kw = self.k_dense(k)
        vw = self.v_dense(v)
        # 形状变换
        qw = K.reshape(qw, (-1, K.shape(q)[1], self.heads, self.key_size))
        kw = K.reshape(kw, (-1, K.shape(k)[1], self.heads, self.key_size))
        vw = K.reshape(vw, (-1, K.shape(v)[1], self.heads, self.head_size))
        # 维度置换
        qw = K.permute_dimensions(qw, (0, 2, 1, 3))
        kw = K.permute_dimensions(kw, (0, 2, 1, 3))
        vw = K.permute_dimensions(vw, (0, 2, 1, 3))
        # 转为三阶张量
        qw = K.reshape(qw, (-1, K.shape(q)[1], self.key_size))
        kw = K.reshape(kw, (-1, K.shape(k)[1], self.key_size))
        vw = K.reshape(vw, (-1, K.shape(v)[1], self.head_size))
        # Attention
        a = K.batch_dot(qw, kw, [2, 2]) / np.sqrt(self.key_size)
        a = add_seq_mask(a, v_mask, 1, -1, self.heads)
        if a_mask is not None:
            if a_mask == 'history_only':
                ones = K.ones_like(a[:1])
                a_mask = (ones - tf.matrix_band_part(ones, -1, 0)) * 1e12
                a = a - a_mask
            else:
                a = a - (1 - a_mask) * 1e12
        a = K.softmax(a)
        # 完成输出
        o = K.batch_dot(a, vw, [2, 1])
        o = K.reshape(o, (-1, self.heads, K.shape(q)[1], self.head_size))
        o = K.permute_dimensions(o, (0, 2, 1, 3))
        o = K.reshape(o, (-1, K.shape(o)[1], self.out_dim))
        o = self.o_dense(o)
        o = add_seq_mask(o, q_mask, 0)
        return o

    def compute_output_shape(self, input_shape):
        return (input_shape[0][0], input_shape[0][1], self.out_dim)

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'key_size': self.key_size
        }
        base_config = super(MultiHeadAttention, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class LayerNormalization(Layer):
    """实现基本的Layer Norm，只保留核心运算部分
    """
    def __init__(self, **kwargs):
        super(LayerNormalization, self).__init__(**kwargs)
        self.epsilon = K.epsilon() * K.epsilon()

    def build(self, input_shape):
        super(LayerNormalization, self).build(input_shape)
        shape = (input_shape[-1], )
        self.gamma = self.add_weight(shape=shape,
                                     initializer='ones',
                                     name='gamma')
        self.beta = self.add_weight(shape=shape,
                                    initializer='zeros',
                                    name='beta')

    def call(self, inputs):
        mean = K.mean(inputs, axis=-1, keepdims=True)
        variance = K.mean(K.square(inputs - mean), axis=-1, keepdims=True)
        std = K.sqrt(variance + self.epsilon)
        outputs = (inputs - mean) / std
        outputs *= self.gamma
        outputs += self.beta
        return outputs

class FactorizedEmbedding(Layer):
    """基于低秩分解的Embedding层
    """
    def __init__(self, input_dim, output_dim, hidden_dim=None, **kwargs):
        super(FactorizedEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        if hidden_dim is None:
            self.hidden_dim = output_dim
        else:
            self.hidden_dim = hidden_dim

    def build(self, input_shape):
        super(FactorizedEmbedding, self).build(input_shape)
        self._embeddings = self.add_weight(name='embeddings',
                                           shape=(self.input_dim,
                                                  self.hidden_dim),
                                           initializer='uniform')
        self._project_kernel = self.add_weight(name='project_kernel',
                                               shape=(self.hidden_dim,
                                                      self.output_dim),
                                               initializer='glorot_uniform')
        self.embeddings = K.dot(self._embeddings, self._project_kernel)

    def call(self, inputs):
        if K.dtype(inputs) != 'int32':
            inputs = K.cast(inputs, 'int32')
        outputs = K.gather(self._embeddings, inputs)
        outputs = K.dot(outputs, self._project_kernel)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape + (self.output_dim, )

    def get_config(self):
        config = {
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'hidden_dim': self.hidden_dim
        }
        base_config = super(FactorizedEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class PositionEmbedding(Layer):
    """定义位置Embedding，这里的Embedding是可训练的。
    """
    def __init__(self, input_dim, output_dim, merge_mode='add', **kwargs):
        super(PositionEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.merge_mode = merge_mode

    def build(self, input_shape):
        super(PositionEmbedding, self).build(input_shape)
        self.embeddings = self.add_weight(name='embeddings',
                                          shape=(self.input_dim,
                                                 self.output_dim),
                                          initializer='zeros')

    def call(self, inputs):
        input_shape = K.shape(inputs)
        batch_size, seq_len = input_shape[0], input_shape[1]
        pos_embeddings = self.embeddings[:seq_len]
        pos_embeddings = K.expand_dims(pos_embeddings, 0)
        pos_embeddings = K.tile(pos_embeddings, [batch_size, 1, 1])
        if self.merge_mode == 'add':
            return inputs + pos_embeddings
        else:
            return K.concatenate([inputs, pos_embeddings])

    def compute_output_shape(self, input_shape):
        if self.merge_mode == 'add':
            return input_shape
        else:
            return input_shape[:2] + (input_shape[2] + self.v_dim, )

    def get_config(self):
        config = {
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'merge_mode': self.merge_mode
        }
        base_config = super(PositionEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class FeedForward(Layer):
    """FeedForward层，其实就是两个Dense层的叠加
    """
    def __init__(self, units, activation='relu', **kwargs):
        super(FeedForward, self).__init__(**kwargs)
        self.units = units
        self.activation = activation

    def build(self, input_shape):
        super(FeedForward, self).build(input_shape)
        output_dim = input_shape[-1]
        self.dense_1 = Dense(self.units, activation=self.activation)
        self.dense_2 = Dense(output_dim)

    def call(self, inputs):
        x = self.dense_1(inputs)
        x = self.dense_2(x)
        return x

    def get_config(self):
        config = {'units': self.units, 'activation': self.activation}
        base_config = super(FeedForward, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class EmbeddingDense(Layer):
    """运算跟Dense一致，但kernel用Embedding层的embeddings矩阵。
    根据Embedding层的名字来搜索定位Embedding层。
    """
    def __init__(self, embedding_name, activation='softmax', **kwargs):
        super(EmbeddingDense, self).__init__(**kwargs)
        self.embedding_name = embedding_name
        self.activation = activation

    def call(self, inputs):
        if not hasattr(self, 'kernel'):
            embedding_layer = inputs._keras_history[0]

            if embedding_layer.name != self.embedding_name:

                def recursive_search(layer):
                    """递归向上搜索，根据名字找Embedding层
                    """
                    last_layer = layer._inbound_nodes[0].inbound_layers
                    if isinstance(last_layer, list):
                        if len(last_layer) == 0:
                            return None
                        else:
                            last_layer = last_layer[0]
                    if last_layer.name == self.embedding_name:
                        return last_layer
                    else:
                        return recursive_search(last_layer)

                embedding_layer = recursive_search(embedding_layer)
                if embedding_layer is None:
                    raise Exception('Embedding layer not found')

                self.kernel = K.transpose(embedding_layer.embeddings)
                self.units = K.int_shape(self.kernel)[1]
                self.bias = self.add_weight(name='bias',
                                            shape=(self.units, ),
                                            initializer='zeros')

        outputs = K.dot(inputs, self.kernel)
        outputs = K.bias_add(outputs, self.bias)
        outputs = Activation(self.activation).call(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.units, )

    def get_config(self):
        config = {
            'embedding_name': self.embedding_name,
            'activation': self.activation
        }
        base_config = super(EmbeddingDense, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

custom_objects = {
    'gelu_erf': gelu_erf,
    'gelu_tanh': gelu_tanh,
    'MultiHeadAttention': MultiHeadAttention,
    'LayerNormalization': LayerNormalization,
    'FactorizedEmbedding': FactorizedEmbedding,
    'PositionEmbedding': PositionEmbedding,
    'FeedForward': FeedForward,
    'EmbeddingDense': EmbeddingDense
}

get_custom_objects().update(custom_objects)