# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import cv2
import time
# from numba import jit
# import numba

import torch
import torch.nn.functional as F

from .utils import rot_mat_to_euler

def find_dynamic_lmk_idx_and_bcoords_projective(vertices, neck_joint_loc, pose, dynamic_lmk_faces_idx,
                                     dynamic_lmk_b_coords,
                                     neck_kin_chain, dtype=torch.float32):
    batch_size = vertices.shape[0]
    # print('pose shape len')
    # print(len(pose.shape))
    if len(pose.shape) == 2:
        aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
                                     neck_kin_chain)
        rot_mats = batch_rodrigues(
            aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)
    else:
        rot_mats = torch.index_select(pose.view(batch_size, -1, 3, 3), 1,
                                     neck_kin_chain)

    b = rot_mats.shape[0]

    rel_rot_mat = torch.eye(3, device=vertices.device,
                            dtype=dtype).reshape(1, 3, 3).repeat(b, 1, 1)
    for idx in range(len(neck_kin_chain)):
        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
    neck_cam_loc = -torch.bmm(torch.transpose(rel_rot_mat, 1, 2), neck_joint_loc.reshape(-1, 3, 1)).reshape(-1, 3)
    neck_cam_loc_len = torch.norm(neck_cam_loc, dim=1).reshape(-1, 1).repeat(1, 3)
    neck_cam_dir = neck_cam_loc / neck_cam_loc_len
    y_ang = -torch.atan2(neck_cam_dir[:, 0], neck_cam_dir[:, 2])

    # print('neck cam dir')
    # print(neck_cam_dir)

    y_rot_angle = torch.round(
        torch.clamp(y_ang * 180.0 / np.pi,
                    max=39)).to(dtype=torch.long)

    # print('face cam ang')
    # print(y_rot_angle)

    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
    y_rot_angle = (neg_mask * neg_vals +
                   (1 - neg_mask) * y_rot_angle)
    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
                                           0, y_rot_angle)
    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
                                          0, y_rot_angle)

    return dyn_lmk_faces_idx, dyn_lmk_b_coords

def find_dynamic_lmk_idx_and_bcoords(vertices, pose, dynamic_lmk_faces_idx,
                                     dynamic_lmk_b_coords,
                                     neck_kin_chain, dtype=torch.float32):
    ''' Compute the faces, barycentric coordinates for the dynamic landmarks

        To do so, we first compute the rotation of the neck around the y-axis
        and then use a pre-computed look-up table to find the faces and the
        barycentric coordinates that will be used.

        Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de)
        for providing the original TensorFlow implementation and for the LUT.

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        pose: torch.tensor Bx(Jx3), dtype = torch.float32
            The current pose of the body model
        dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
            The look-up table from neck rotation to faces
        dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
            The look-up table from neck rotation to barycentric coordinates
        neck_kin_chain: list
            A python list that contains the indices of the joints that form the
            kinematic chain of the neck.
        dtype: torch.dtype, optional

        Returns
        -------
        dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
        dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
    '''

    batch_size = vertices.shape[0]

    if len(pose.shape) == 2:
        aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
                                     neck_kin_chain)
        rot_mats = batch_rodrigues(
            aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)
    else:
        rot_mats = torch.index_select(pose.view(batch_size, -1, 3, 3), 1,
                                      neck_kin_chain)

    b = rot_mats.shape[0]

    rel_rot_mat = torch.eye(3, device=vertices.device,
                            dtype=dtype).reshape(1, 3, 3).repeat(b, 1, 1)
    for idx in range(len(neck_kin_chain)):
        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)

    y_rot_angle = torch.round(
        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
                    max=39)).to(dtype=torch.long)
    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
    y_rot_angle = (neg_mask * neg_vals +
                   (1 - neg_mask) * y_rot_angle)
#     print('std yrot ang')
#     print(y_rot_angle)

    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
                                           0, y_rot_angle)
    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
                                          0, y_rot_angle)

    return dyn_lmk_faces_idx, dyn_lmk_b_coords

def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
    ''' Calculates landmarks by barycentric interpolation

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        faces: torch.tensor Fx3, dtype = torch.long
            The faces of the mesh
        lmk_faces_idx: torch.tensor L, dtype = torch.long
            The tensor with the indices of the faces used to calculate the
            landmarks.
        lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
            The tensor of barycentric coordinates that are used to interpolate
            the landmarks

        Returns
        -------
        landmarks: torch.tensor BxLx3, dtype = torch.float32
            The coordinates of the landmarks for each mesh in the batch
    '''
    # Extract the indices of the vertices for each face
    # BxLx3
    batch_size, num_verts = vertices.shape[:2]
    device = vertices.device

    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
        batch_size, -1, 3)

    lmk_faces += torch.arange(
        batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts

    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
        batch_size, -1, 3, 3)

    landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
    return landmarks

def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
        lbs_weights, pose2rot=True, dtype=torch.float32):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : torch.tensor Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template torch.tensor BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : torch.tensor Px(V * 3)
            The pose PCA coefficients
        J_regressor : torch.tensor JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: torch.tensor J
            The array that describes the kinematic tree for the model
        lbs_weights: torch.tensor N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        pose2rot: bool, optional
            Flag on whether to convert the input pose tensor to rotation
            matrices. The default value is True. If False, then the pose tensor
            should already contain rotation matrices and have a size of
            Bx(J + 1)x9
        dtype: torch.dtype, optional

        Returns
        -------
        verts: torch.tensor BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        joints: torch.tensor BxJx3
            The joints of the model
    '''

    batch_size = max(betas.shape[0], pose.shape[0])
    device = betas.device

    # Add shape contribution
    v_shaped = v_template + blend_shapes(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints(J_regressor, v_shaped)

    # print('j1t')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        rot_mats = batch_rodrigues(
            pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])

        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
        # (N x P) x (P, V * 3) -> N x V x 3
        pose_offsets = torch.matmul(pose_feature, posedirs) \
            .view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)

        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
                                    posedirs).view(batch_size, -1, 3)

    v_posed = pose_offsets + v_shaped
    # 4. Get the global joint location

    # print(rot_mats[0][1])

    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

    # 5. Do skinning:
    # W is N x V x (J + 1)
    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    # (N x V x (J + 1)) x (N x (J + 1) x 16)
    num_joints = J_regressor.shape[0]
    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
        .view(batch_size, -1, 4, 4)

    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
                               dtype=dtype, device=device)
    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))

    verts = v_homo[:, :, :3, 0]

    return verts, J_transformed, A, J

def vertices2joints(J_regressor, vertices):
    ''' Calculates the 3D joint locations from the vertices

    Parameters
    ----------
    J_regressor : torch.tensor JxV
        The regressor array that is used to calculate the joints from the
        position of the vertices
    vertices : torch.tensor BxVx3
        The tensor of mesh vertices

    Returns
    -------
    torch.tensor BxJx3
        The location of the joints
    '''

    return torch.einsum('bik,ji->bjk', [vertices, J_regressor])

def vertices2joints_np(J_regressor_np, vertices_np):
    ''' Calculates the 3D joint locations from the vertices

    Parameters
    ----------
    J_regressor : torch.tensor JxV
        The regressor array that is used to calculate the joints from the
        position of the vertices
    vertices : torch.tensor BxVx3
        The tensor of mesh vertices

    Returns
    -------
    torch.tensor BxJx3
        The location of the joints
    '''
    n_j, n_v = J_regressor_np.shape
    b = vertices_np.shape[0]
    return np.matmul(np.tile(J_regressor_np.reshape(1, n_j, n_v), (b, 1, 1)), vertices_np)

def blend_shapes(betas, shape_disps):
    ''' Calculates the per vertex displacement due to the blend shapes

    Parameters
    ----------
    betas : torch.tensor Bx(num_betas)
        Blend shape coefficients
    shape_disps: torch.tensor Vx3x(num_betas)
        Blend shapes

    Returns
    -------
    torch.tensor BxVx3
        The per-vertex displacement due to shape deformation
    '''

    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
    # i.e. Multiply each shape displacement by its corresponding beta and
    # then sum them.
    blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
    return blend_shape

def blend_shapes_np(betas, shape_disps):
    ''' Calculates the per vertex displacement due to the blend shapes

    Parameters
    ----------
    betas : torch.tensor Bx(num_betas)
        Blend shape coefficients
    shape_disps: torch.tensor Vx3x(num_betas)
        Blend shapes

    Returns
    -------
    torch.tensor BxVx3
        The per-vertex displacement due to shape deformation
    '''

    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
    # i.e. Multiply each shape displacement by its corresponding beta and
    # then sum them.
    # blend_shape = betas, shape_disps
    b = betas.shape[0]
    m, k, l = shape_disps.shape
    shape_disps_r = np.tile(shape_disps.reshape(1, m, k, l), (b, 1, 1, 1))
    betas_r = np.tile(betas.reshape(b, 1, -1, 1), (1, m, 1, 1))
    blend_shape = np.matmul(shape_disps_r, betas_r)

    return blend_shape.reshape((b, m, 3))

def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat

def rodnumba(p):
    rot, jac = cv2.Rodrigues(p)
    return rot, jac

def batch_rodrigues_np(pose_body):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: ndarray Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: ndarray Nx3x3
            The rotation matrices for the given axis-angle parameters
        R_jac: ndarray Nx3x3x3
            Jacobians of the rotation matrices
    '''
    batch_size = pose_body.shape[0]
    n_j = int(pose_body.shape[1]/3)
    dt = pose_body.dtype
    rot_mats = np.zeros((batch_size, n_j, 3, 3), dtype=dt)
    rot_mats_jac = np.zeros((batch_size, n_j, 3, 3, 3), dtype=dt)
    for b in range(0, pose_body.shape[0]):
        for i in range(0, n_j):
            # rot, jac = cv2.Rodrigues(pose_body[b][3 * (i): 3 * (i + 1)].reshape(-1))
            rot, jac = rodnumba(pose_body[b][3 * (i): 3 * (i + 1)].reshape(-1))
            # print(numba.typeof(rot))
            # print(numba.typeof(jac))
            rot_mats[0, i] = rot
            rot_mats_jac[0, i] = jac.reshape(3, 3, 3)

    return rot_mats, rot_mats_jac

def transform_mat(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    return torch.cat([F.pad(R, [0, 0, 0, 1]),
                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)

def transform_mat_np(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    b = R.shape[0]
    T = np.zeros((b, 4, 4), dtype=R.dtype)
    T[:, :3, :3] = R
    T[:, :3, 3] = t.reshape(b, 3)
    T[:, 3, 3] = 1
    return T

def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : torch.tensor BxNx3x3
        Tensor of rotation matrices
    joints : torch.tensor BxNx3
        Locations of joints
    parents : torch.tensor BxN
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    posed_joints : torch.tensor BxNx3
        The locations of the joints after applying the pose rotations
    rel_transforms : torch.tensor BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """

    joints = torch.unsqueeze(joints, dim=-1)

    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    # print(transforms_mat[0][0])
    # print(transforms_mat[0][1])

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        # Subtract the joint location at the rest pose
        # No need for rotation, since it's identity when at rest
        curr_res = torch.matmul(transform_chain[parents[i]],
                                transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)

    # print(transforms[0][1])

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])

    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])

    return posed_joints, rel_transforms

# @jit
def batch_rigid_transform_diff(rot_mats, rot_mats_jac, transform_jac_chain, joints, parents, is_jac=True):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : ndarray BxNx3x3
        Tensor of rotation matrices
    rot_mats_jac : ndarray BxNx3x3x3
        Tensor of rotation matrix Jacobians
    joints : ndarray BxNx3
        Locations of joints
    parents : ndarray BxN
        The kinematic tree of each object

    Returns
    -------
    posed_joints : ndarray BxNx3
        The locations of the joints after applying the pose rotations
    pose_joints_jac : ndarray BxNxNx3x3
        Jacobian of pose_joints w.r.t. pose
    rel_transforms : ndarray BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    rel_transforms_jac : jacobian w.r.t. pose
    """
    # joints = joints.reshape()
    # joints = torch.unsqueeze(joints, dim=-1)

    t_0 = time.time()

    rel_joints = joints.copy()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat_np(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    b = joints.shape[0]
    n_j = joints.shape[1]

    t_1 = time.time()

    if is_jac:
        ainds = np.arange(0, parents.shape[0])
        transform_jac_chain.fill(0)
        transform_jac_chain[:, ainds, ainds, :, 0:3, 0:3] = rot_mats_jac[:, ainds]

    transform_chain = np.zeros((b, parents.shape[0], 4, 4), dtype=rot_mats.dtype)
    transform_chain[:, 0] = transforms_mat[:, 0]
    for i in range(1, parents.shape[0]):
        transform_chain[:, i] = np.matmul(transform_chain[:, parents[i]],
                                transforms_mat[:, i])

    if is_jac:
        inds = np.arange(1, len(parents))
        trans_parent = transform_chain[:, parents[inds]].reshape(b, len(inds), 1, 1, 4, 4)
        trans_jac = transform_jac_chain[:, inds, inds, :].reshape(1, len(inds), 1, 3, 4, 4)
        transform_jac_chain[:, inds, inds, :] = np.matmul(trans_parent, trans_jac).reshape(len(inds), 3, 4, 4)

        m = np.eye(parents.shape[0], dtype=np.bool)
        b = transforms_mat.shape[0]
        for i in range(1, parents.shape[0]):
            m[i] = m[i] | m[parents[i]]
            # transform_jac_chain[:, i, :, :] += np.matmul(transform_jac_chain[:, parents[i], :, :], transforms_mat[:, i])
            tr_jac_ch_sel = transform_jac_chain[:, parents[i], m[i], :]
            transform_jac_chain[:, i, m[i], :] += np.matmul(tr_jac_ch_sel, transforms_mat[:, i]).reshape(b, -1, 3, 4, 4)

    t_2 = time.time()

    transforms = transform_chain
    posed_joints = transforms[:, :, :3, 3]
    joints_rot = np.matmul(transforms[:, :, 0:3, 0:3], joints.reshape(b, n_j, 3, 1)).reshape((b, n_j, 3))
    rel_transforms = transforms.copy()
    rel_transforms[:, :, 0:3, 3] = rel_transforms[:, :, 0:3, 3] - joints_rot
    if is_jac:
        transforms_jac = np.transpose(transform_jac_chain, (0, 2, 3, 1, 4, 5))
        posed_joints_jac = transforms_jac[:, :, :, :, :3, 3]
        tjhj = np.matmul(transforms_jac[:, :, :, :, 0:3, 0:3], joints.reshape((b, 1, 1, n_j, 3, 1))).reshape(
            (b, n_j, 3, n_j, 3))
        rel_transforms_jac = transforms_jac.copy()
        rel_transforms_jac[:, :, :, :, 0:3, 3] = rel_transforms_jac[:, :, :, :, 0:3, 3] - tjhj
    else:
        posed_joints_jac = None
        rel_transforms_jac = None
    t_3 = time.time()

    # print('brgd breakdown {} {} {} '.format(t_1-t_0, t_2-t_1, t_3-t_2))

    return posed_joints, posed_joints_jac, rel_transforms, rel_transforms_jac

def batch_rigid_transform_fast_diff(rot_mats, rot_mats_jac, joints, parents):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : ndarray BxNx3x3
        Tensor of rotation matrices
    rot_mats_jac : ndarray BxNx3x3x3
        Tensor of rotation matrix Jacobians
    joints : ndarray BxNx3
        Locations of joints
    parents : ndarray BxN
        The kinematic tree of each object

    Returns
    -------
    posed_joints : ndarray BxNx3
        The locations of the joints after applying the pose rotations
    pose_joints_jac : ndarray BxNxNx3x3
        Jacobian of pose_joints w.r.t. pose
    rel_transforms : ndarray BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    rel_transforms_jac : jacobian w.r.t. pose
    """
    # joints = joints.reshape()
    # joints = torch.unsqueeze(joints, dim=-1)

    t_0 = time.time()

    rel_joints = joints.copy()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat_np(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    # print(transforms_mat[0][0])
    # print(transforms_mat[0][1])

    b = joints.shape[0]
    n_j = joints.shape[1]

    t_1 = time.time()

    # transform_jac = np.zeros((b, n_j, 3, 4, 4))
    # transform_jac[:, 0, :, 0:3, 0:3] = rot_mats_jac[:, 0]
    transform_jac_chain = np.zeros((b, n_j, 3, parents.shape[0], 4, 4))
    for i in range(0, parents.shape[0]):
        transform_jac_chain[:, i, :, i, 0:3, 0:3] = rot_mats_jac[:, i]

    transform_chain = np.copy(transforms_mat)
    for i in range(1, parents.shape[0]):
        t_curr = np.matmul(transform_chain[:, parents[i], 0:3, 0:3],
                                transforms_mat[:, i, 0:3, 3].reshape(-1, 1, 3, 1)).reshape(-1, 1, 3)
        transform_chain[:, i, 0:3, 3] = t_curr + transform_chain[:, parents[i], 0:3, 3]
        t_i = np.tile(transforms_mat[:, i, 0:3, 3].reshape(-1, 1, 1, 3, 1), (1, n_j, 3, 1, 1))
        # print(transform_jac_chain[:, :, :, parents[i], 0:3, 0:3].shape)
        # print(t_i.shape)
        t_jac_curr = np.matmul(transform_jac_chain[:, :, :, parents[i], 0:3, 0:3], t_i)
        transform_jac_chain[:, :, :, i, 0:3, 3] = transform_jac_chain[:, :, :, parents[i], 0:3, 3] + \
                                                  t_jac_curr.reshape(-1, n_j, 3, 3)

    # transforms = np.stack(transform_chain, axis=1)
    transforms = transform_chain
    # transforms_jac = np.stack(transform_jac_chain, axis=3)
    transforms_jac = transform_jac_chain

    t_2 = time.time()
    # print(transforms[0][1])

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]
    posed_joints_jac = transforms_jac[:, :, :, :, :3, 3]

    # joints_homogen = np.zeros((b, n_j, 4, 1))
    # joints_homogen[:, :, 0:3, 0] = joints

    joints_rot = np.matmul(transforms[:,:,0:3,0:3], joints.reshape(b, n_j, 3, 1)).reshape((b, n_j, 3))
    # jht = np.pad(np.matmul(transforms, joints_homogen), [(0, 0), (0, 0), (0, 0), (3,0)])
    # rel_transforms = transforms - jht
    rel_transforms = transforms.copy()
    rel_transforms[:, :, 0:3, 3] = rel_transforms[:, :, 0:3, 3] - joints_rot

    # jhtj = np.matmul(transforms_jac, joints_homogen.reshape((b, 1, 1, n_j, 4, 1)))
    # jhtj = np.pad(jhtj, [(0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (3, 0)])

    tjhj = np.matmul(transforms_jac[:,:,:,:,0:3,0:3], joints.reshape((b, 1, 1, n_j, 3, 1))).reshape((b, n_j, 3, n_j, 3))

    rel_transforms_jac = transforms_jac.copy()
    rel_transforms_jac[:, :, :, :, 0:3, 3] = rel_transforms_jac[:, :, :, :, 0:3, 3] - tjhj
    # rel_transforms_jac = transforms_jac - jhtj

    t_3 = time.time()

    print('brgd breakdown {} {} {} '.format(t_1-t_0, t_2-t_1, t_3-t_2))

    return posed_joints, posed_joints_jac, rel_transforms, rel_transforms_jac

def prepare_J(betas, v_template, shapedirs, J_regressor, n_v):
    # Add shape contribution
    v_shaped = v_template + blend_shapes_np(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints_np(J_regressor, v_shaped)

    n_j = J_regressor.shape[0]

    batch_size = betas.shape[0]

    homogen_coord = np.ones((batch_size, n_v, 1), dtype=v_template.dtype)

    transform_jac_chain = np.zeros((batch_size, n_j, n_j, 3, 4, 4), dtype=v_template.dtype)

    return J, v_shaped, homogen_coord, transform_jac_chain

def rel_to_direct(pose, parents):

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)

    b = pose.shape[0]

    rot_chain = np.zeros((b, parents.shape[0], 3, 3))
    rot_chain[:, 0] = rot_mats[:, 0]
    for i in range(1, parents.shape[0]):
        rot_chain[:, i] = np.matmul(rot_chain[:, parents[i]], rot_mats[:, i])

    pose_dir = np.zeros_like(pose)
    n_j = int(pose.shape[1]/3)
    for b in range(0, pose.shape[0]):
        for i in range(0, n_j):
            rv, jac = cv2.Rodrigues(rot_chain[b, i])
            pose_dir[b, 3*i : 3*(i+1)] = rv.reshape(-1)

    return pose_dir

def lbs_diff_fast(pose, parents,
        J, v_shaped, W, W_j, homogen_coord, v_inds=None):
    t_0 = time.time()

    batch_size = pose.shape[0]

    # print('j1:')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)
    # rot_mats = rot_mats.reshape((batch_size, -1, 3, 3))
    n_j = rot_mats.shape[1]

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]

    n_v = v_shaped.shape[1]

    v_posed = v_shaped

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_fast_diff(rot_mats, rot_mat_jacs, J, parents)

    t_3 = time.time()

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))

    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))

    A_jact = A_jac  # .transpose(0, 2, 3, 1, 4, 5)
    T_jac = np.matmul(W_j, A_jact.reshape(batch_size, n_j, 3, n_j, -1))
    T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))
    # N x n_j x 3 x V x 16

    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)

    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))

    T_j = T.reshape((batch_size, 1, 1, n_v, 4, 4))

    v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
    v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)

    verts = v_homo[:, :, :3, 0]

    v_homo_jac = v_homo_jac2[:, :, :, :, :3, :]

    verts_jac = v_homo_jac[:, :, :, :, :3, 0]

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_4-t_3))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J  # , v_posed, v_posed_jac

# @jit
def lbs_diff(pose, posedirs, parents, J, v_shaped, lbs_weights, homogen_coord, transform_jac_chain, v_inds=None):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : ndarray Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template ndarray BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : ndarray Px(V * 3)
            The pose PCA coefficients
        J_regressor : ndarray JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: ndarray J
            The array that describes the kinematic tree for the model
        lbs_weights: ndarray N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        vinds: ndarray - list of required vertex indices (if None, all vertices will be processed)

        Returns
        -------
        verts: ndarray BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        verts_jac: ndarray BxVx(J+1)x3x3
        joints: ndarray BxJx3
            The joints of the model
        joints_jac: ndarray BxJxJx3x3
            Jacobian of joints' coordinates
    '''

    t_0 = time.time()

    batch_size = pose.shape[0]

    if v_inds is not None:
        lbs_weights = lbs_weights[v_inds]
        n_v = len(v_inds)
    else:
        n_v = 0

    n_j = J.shape[1]

    W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))
    W_j = W.reshape((batch_size, 1, 1, n_v, n_j))

    # print('j1:')
    # print(J[0][1])

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)
    # rot_mats = rot_mats.reshape((batch_size, -1, 3, 3))
    n_j = rot_mats.shape[1]
    pose_feature = (rot_mats[:, 1:, :, :] - np.tile(np.eye(3).reshape(1, 1, 3, 3), (batch_size, n_j-1, 1, 1)))\
        .reshape((batch_size, -1))

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]
        inds = np.stack([3*v_inds, 3*v_inds+1, 3*v_inds+2], axis=1).reshape(-1)
        posedirs = posedirs[:, inds]
        # lbs_weights = lbs_weights[v_inds]

        # (N x P) x (P, V * 3) -> N x V x 3
    pose_offsets = np.matmul(pose_feature, posedirs) \
        .reshape((batch_size, -1, 3))

    n_v = v_shaped.shape[1]

    pose_offset_jacs = np.zeros((batch_size, n_j, 3, n_v, 3), dtype=pose.dtype)
    for i in range (1, n_j):
        pdi = np.matmul(rot_mat_jacs[:, i].reshape(-1, 9), posedirs[(i-1)*9 : i*9, :]).reshape((batch_size, 3, -1, 3))
        pose_offset_jacs[:, i] = pdi

    v_posed = pose_offsets + v_shaped

    #NxVxJ+1x3x3
    v_posed_jac = pose_offset_jacs

    # print(rot_mats[0][1])

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_diff(rot_mats, rot_mat_jacs, transform_jac_chain, J, parents)

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))

    # W_j = np.tile(W.reshape((batch_size, 1, 1, n_v, n_j)), (1, n_j, 3, 1, 1))

    A_jact = A_jac #.transpose(0, 2, 3, 1, 4, 5)
    T_jac = np.matmul(W_j, A_jact.reshape(batch_size, n_j, 3, n_j, -1))
    T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))
    #N x n_j x 3 x V x 16

    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)

    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))

    # T_j = np.tile(T.reshape((batch_size, 1, 1, n_v, 4, 4)), (1, n_j, 3, 1, 1, 1))

    T_j = T.reshape((batch_size, 1, 1, n_v, 4, 4))

    # v_posed_jac_h = np.pad(v_posed_jac, ((0, 0), (0,0), (0,0), (0,0), (0,1))).reshape((batch_size, n_j, 3, n_v, 4, 1))
    # v_posed_jac_h = v_posed_jac_h.reshape((batch_size, n_j, 3, n_v, 4, 1))
    # v_homo_jac1 = np.matmul(T_j, v_posed_jac_h)

    v_homo_jac1 = np.matmul(T_j[:, :, :, :, 0:3, 0:3], v_posed_jac.reshape((batch_size, n_j, 3, n_v, 3, 1)))

    # v_posed_homo_j = np.tile(v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1)), (1, n_j, 3, 1, 1, 1))
    v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
    v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)
    # v_homo_jac2 = v_homo_jac2.transpose((0, 3, 1, 2, 4, 5))

    verts = v_homo[:, :, :3, 0]

    v_homo_jac = v_homo_jac1 + v_homo_jac2[:, :, :, :,  :3, :]

    verts_jac = v_homo_jac[:, :, :, :, :3, 0]

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_4-t_3))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J #, v_posed, v_posed_jac

def lbs_diff_nopd(pose, posedirs_face, parents, J, v_shaped, lbs_weights, homogen_coord, transform_jac_chain, v_inds=None,
                  is_jac=True, bpm=None, face_expression=None, w_thr=0):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : ndarray Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template ndarray BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : ndarray Px(V * 3)
            The pose PCA coefficients
        J_regressor : ndarray JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: ndarray J
            The array that describes the kinematic tree for the model
        lbs_weights: ndarray N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        vinds: ndarray - list of required vertex indices (if None, all vertices will be processed)

        Returns
        -------
        verts: ndarray BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        verts_jac: ndarray BxVx(J+1)x3x3
        joints: ndarray BxJx3
            The joints of the model
        joints_jac: ndarray BxJxJx3x3
            Jacobian of joints' coordinates
    '''

    t_0 = time.time()

    batch_size = pose.shape[0]
    if v_inds is not None:

        if bpm is None:
            w_arr = np.sum(lbs_weights[v_inds] ** 2, axis=0)/len(v_inds)
            bpm = w_arr > w_thr

        lbs_weights = lbs_weights[v_inds]

    # 3. Add pose blend shapes
    # N x J x 3 x 3

    t_1 = time.time()

    rot_mats, rot_mat_jacs = batch_rodrigues_np(pose)

    n_j = rot_mats.shape[1]

    if v_inds is not None:
        v_shaped = v_shaped[:, v_inds, :]

        if face_expression is not None:
            posedirs_face = posedirs_face[v_inds]
            n_v = len(v_inds)
            dv = np.matmul(posedirs_face.reshape(1, n_v, 3, 1, -1), np.tile(face_expression.reshape(1, 1, 1, -1, 1),
                                                                            (1, n_v, 3, 1, 1)))
            v_shaped += dv.reshape(1, -1, 3)
            dv_jac = np.transpose(posedirs_face.reshape(n_v, 3, -1), (2, 0, 1)).reshape(-1, 3 * n_v)

    n_v = v_shaped.shape[1]

    v_posed = v_shaped

    t_2 = time.time()

    J_transformed, J_transformed_jac, A, A_jac = batch_rigid_transform_diff(rot_mats, rot_mat_jacs, transform_jac_chain, J, parents, is_jac)

    t_3 = time.time()

    # 5. Do skinning:
    # W is N x V x (J + 1)
    # W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    # (N x V x (J + 1)) x (N x (J + 1) x 16) = N x V x 16
    num_joints = n_j
    W = np.tile(lbs_weights.reshape(1, n_v, n_j), (batch_size, 1, 1))
    T = np.matmul(W, A.reshape(batch_size, num_joints, 16)) \
        .reshape((batch_size, -1, 4, 4))
    v_posed_homo = np.concatenate([v_posed, homogen_coord], axis=2)
    v_homo = np.matmul(T, v_posed_homo.reshape((batch_size, n_v, 4, 1)))
    verts = v_homo[:, :, :3, 0]

    t_3_1 = time.time()

    if is_jac:

        Wred = W[:, :, bpm]
        n_bpm = Wred.shape[-1]
        W_j = Wred.reshape((batch_size, 1, 1, n_v, n_bpm))
        T_jac = np.matmul(W_j, A_jac.reshape(batch_size, n_j, 3, n_j, -1)[:, :, :, bpm])
        T_jac = T_jac.reshape((batch_size, n_j, 3, n_v, 4, 4))

        #N x n_j x 3 x V x 16

        t_3_2 = time.time()

        v_posed_homo_j = v_posed_homo.reshape((batch_size, 1, 1, n_v, 4, 1))
        v_homo_jac2 = np.matmul(T_jac, v_posed_homo_j)
        n_pars = pose.shape[1]
        verts_jac = v_homo_jac2[:, :, :, :,  :3, 0].reshape(n_pars, -1)
        if face_expression is not None:
            n_fp = 10
            dv_jac = np.matmul(T[:, :, 0:3, 0:3], dv_jac.reshape(n_fp, -1, 3, 1)).reshape(n_fp, -1)
            # dv_jac = np.zeros_like(dv_jac)
            verts_jac = np.concatenate([verts_jac, dv_jac], axis=0)
    else:
        verts_jac = None

    t_4 = time.time()

    # print('breakdown b {} a {} rt {} f1 {} f2 {} f3 {}'.format(t_1-t_0, t_2-t_1, t_3-t_2, t_3_1 - t_3, t_3_2 - t_3_1, t_4-t_3_2))

    return verts, verts_jac, J_transformed, J_transformed_jac, A, A_jac, J, bpm #, v_posed, v_posed_jac


from abc import abstractmethod
from collections import OrderedDict
from functools import partial
from typing import List

import numpy as np
import pymc3 as pm
import theano as th
import theano.tensor as tt
from pymc3.variational.updates import get_or_compute_grads

from .. import types
from ..io import io_commons
from ..models.fancy_model import GeneralizedContinuousModel

class FancyStochasticOptimizer:
    """The base class of stochastic optimizers equipped with the functionality of saving and loading the
    optimizer state to and from disk (e.g. for stateful optimizers such as ADAM and ADAMAX), and the
    possibility of utilizing the extra attributes of `GeneralizedContinuousModel` to perform structured
    parameter updates, e.g. updating only sample-specific variables while keeping global variables intact
    (see `FancyAdamax` for a concrete implementation).
    """
    @abstractmethod
    def get_optimizer(self,
                      model: GeneralizedContinuousModel=None,
                      approx: pm.MeanField=None):
        """

        Args:
            model: a generalized continuous PyMC3 model
            approx: an instance of PyMC3 mean-field approximation

        Returns:
            A callable function that upon providing `loss_or_grads` and `params`, returns an
            `OrderedDict` of shared theano tensor updates (for example, see `FancyAdamax.get_optimizer`).
        """
        raise NotImplementedError

    @staticmethod
    def get_call_kwargs(_locals_):
        _locals_ = _locals_.copy()
        _locals_.pop('loss_or_grads')
        _locals_.pop('params')
        return _locals_

    @abstractmethod
    def save(self, output_path: str):
        raise NotImplementedError

    @abstractmethod
    def load(self, input_path: str):
        raise NotImplementedError

class FancyAdamax(FancyStochasticOptimizer):
    """Adamax optimizer with saving/loading functionality and sample-specific-only update mode."""
    def __init__(self,
                 learning_rate: float = 0.002,
                 beta1: float = 0.9,
                 beta2: float = 0.999,
                 epsilon: float = 1e-8,
                 sample_specific_only: bool = False,
                 disable_bias_correction: bool = False):
        """Initializer.

        Args:
            learning_rate: learning rate
            beta1: first moment forgetting factor
            beta2: second moment forgetting factor
            epsilon: a small float for avoiding division-by-zero
            sample_specific_only: only update sample-specific variables (as specified in the generalized model)
            disable_bias_correction: disable moment estimation bias correction
        """
        self.learning_rate = learning_rate
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.sample_specific_only = sample_specific_only
        self.disable_bias_correction = disable_bias_correction

        # placeholder for first (m) and second (u) moments
        # in mean-field type approximation, ``mu`` and ``rho`` each have their own tensors
        # the list elements correspond to ``mu`` and ``rho`` moments, respectively
        self.m_tensors: List[types.TensorSharedVariable] = []
        self.u_tensors: List[types.TensorSharedVariable] = []

        # placeholder for the state of moment estimation bias corrector
        self.res_tensor: types.TensorSharedVariable = None

    def _assert_shared_tensors_available(self):
        m_u_available = len(self.m_tensors) == 2 and len(self.u_tensors) == 2
        res_available = self.disable_bias_correction or self.res_tensor is not None
        assert m_u_available and res_available, "Adamax tensors are not available yet"

    def get_mu_m(self):
        self._assert_shared_tensors_available()
        return self.m_tensors[0]

    def get_rho_m(self):
        self._assert_shared_tensors_available()
        return self.m_tensors[1]

    def get_mu_u(self):
        self._assert_shared_tensors_available()
        return self.u_tensors[0]

    def get_rho_u(self):
        self._assert_shared_tensors_available()
        return self.u_tensors[1]

    def get_res_tensor(self):
        return self.res_tensor

    @staticmethod
    def structured_adamax(loss_or_grads=None,
                          params=None,
                          model: GeneralizedContinuousModel=None,
                          approx: pm.MeanField=None,
                          learning_rate=0.002, beta1=0.9,
                          beta2=0.999, epsilon=1e-8,
                          sample_specific_only=False,
                          disable_bias_correction=False,
                          base_class: 'FancyAdamax' = None):
        """Adamax stochastic optimizer with partial sample-specific-only update functionality.

        Args:
            loss_or_grads: symbolic loss function or gradients
            params: variational parameter bundle
            model: an instance of generalized model
            approx: an instance of variational approximation for the model
            learning_rate: global learning rate
            beta1: first moment estimation forgetting factor
            beta2: second moment estimation forgetting factor
            epsilon: a small float to avoid division-by-zero
            sample_specific_only: only update parameters registered in the generalized model as sample-specific
            disable_bias_correction: disable moment estimation bias correction
            base_class: a reference to the base class to store a reference to the shared tensors (for I/O)

        Returns:
            returns the function itself if `loss_or_grads` and `params` are not given;
            otherwise, returns an ordered dict of shared tensor updates (to be used in pymc3 for compiling
            the step function)
        """
        if loss_or_grads is None and params is None:
            return partial(FancyAdamax.structured_adamax,
                           **FancyStochasticOptimizer.get_call_kwargs(locals()))
        elif loss_or_grads is None or params is None:
            raise ValueError('Please provide both `loss_or_grads` and `params` to get updates')
        assert model is not None, 'Please provide `model` to get updates'
        assert approx is not None, 'Please provide `approx` to get updates'

        all_grads = get_or_compute_grads(loss_or_grads, params)
        updates = OrderedDict()

        # indices of sample-specific vars
        if sample_specific_only:
            vmap_list = io_commons.get_var_map_list_from_meanfield_approx(approx)
            sample_specific_indices = []
            for vmap in vmap_list:
                if vmap.var in model.sample_specific_var_registry:
                    sample_specific_indices += [idx for idx in range(vmap.slc.start, vmap.slc.stop)]
            update_indices = th.shared(np.asarray(sample_specific_indices, dtype=np.int))
            num_dof = len(sample_specific_indices)

        # Using theano constant to prevent upcasting of float32
        one = tt.constant(1)

        if disable_bias_correction:
            a_t = learning_rate
        else:
            res_prev = th.shared(pm.theanof.floatX(beta1))
            res = beta1 * res_prev
            a_t = learning_rate / (one - res)
            updates[res_prev] = res
            if base_class is not None:
                base_class.res_tensor = res_prev

        for param, g_t in zip(params, all_grads):
            if sample_specific_only:
                g_t_view = g_t[update_indices]
                m_prev = th.shared(np.zeros((num_dof,), dtype=types.floatX),
                                   broadcastable=(False,))
                u_prev = th.shared(np.zeros((num_dof,), dtype=types.floatX),
                                   broadcastable=(False,))
            else:
                g_t_view = g_t
                value = param.get_value(borrow=True)
                m_prev = th.shared(np.zeros(value.shape, dtype=types.floatX),
                                   broadcastable=(False,))
                u_prev = th.shared(np.zeros(value.shape, dtype=types.floatX),
                                   broadcastable=(False,))

            # save a reference to m and u in the base class
            if base_class is not None:
                base_class.m_tensors.append(m_prev)
                base_class.u_tensors.append(u_prev)

            m_t = beta1 * m_prev + (one - beta1) * g_t_view
            u_t = tt.maximum(beta2 * u_prev, abs(g_t_view))
            step = a_t * m_t / (u_t + epsilon)

            if sample_specific_only:
                new_param = tt.inc_subtensor(param[update_indices], -step)
            else:
                new_param = param - step

            updates[m_prev] = m_t
            updates[u_prev] = u_t
            updates[param] = new_param

        return updates

    def get_optimizer(self,
                      model: GeneralizedContinuousModel=None,
                      approx: pm.MeanField=None):
        return FancyAdamax.structured_adamax(
            model=model,
            approx=approx,
            beta1=self.beta1,
            beta2=self.beta2,
            learning_rate=self.learning_rate,
            epsilon=self.epsilon,
            sample_specific_only=self.sample_specific_only,
            disable_bias_correction=self.disable_bias_correction,
            base_class=self)

    def save(self, output_path: str) -> None:
        """Saves the state of the optimizer to disk.

        Args:
            output_path: output path (must be writable directory)
        """
        from ..io import io_adamax  # lazy import to break import cycle
        io_adamax.AdamaxStateExporter(self, output_path)()

    def load(self, input_path: str):
        """Loads the state of the optimizer from disk.

        Args:
            input_path: input path (must be a readable directory)
        """
        from ..io import io_adamax  # lazy import to break import cycle
        io_adamax.AdamaxStateImporter(self, input_path)()

import numpy,import theano,import pymc3,from pymc3


#! /usr/bin/env python3

import numpy as np
import argparse
import json
import os
from scipy.interpolate import interp1d

parser = argparse.ArgumentParser(description="Compute log-likelihoods from light curve evaluations")
parser.add_argument("--interp-directory", help="Directory containing the interpolator evaluations and index files")
parser.add_argument("--output-file", help="Name of file to save posterior samples to")
parser.add_argument("--grid-file", help="Location of the grid file")
parser.add_argument("--lc-file", help="Location of the light curve JSON file")
parser.add_argument("--bands", nargs="+", help="Bands to use")
parser.add_argument("--distance", type=float, help="Luminosity distance (in Mpc)")
args = parser.parse_args()

# Load the files containing the grid indices corresponding to each angle
indices_fname_base = args.interp_directory + ("/" if args.interp_directory[-1] != "/" else "") + "indices_"
interp_angles = [0., 30., 45., 60., 75., 90.]
indices = {_theta_interp:np.loadtxt(indices_fname_base + str(int(_theta_interp)) + ".dat", ndmin=1).astype(int) for _theta_interp in interp_angles}

# Load the grid
grid = np.loadtxt(args.grid_file)
theta_grid = grid[:,-1] # angle is the fifth column

# Load the light curve data
with open(args.lc_file, "r") as fp:
    lc_data = json.load(fp)
for b in lc_data.keys():
    for key in lc_data[b].keys():
        lc_data[b][key] = np.array(lc_data[b][key]) # cast everything as a numpy array

interpolator_times = np.logspace(np.log10(0.125), np.log10(37.239195485411194), 264)
interpolator_times_str_to_float = {"{:.3f}".format(_t):_t for _t in interpolator_times} # this is useful since the times in the filenames are truncated to three decimal places

# Set up a data structure to hold all the interpolated magnitudes.
# This is a list with the following structure:
#   - Each element in the list corresponds to a row in the grid.
#   - This list element is a dictionary mapping (band, interpolator_angle) to *another* dictionary.
#       - This sub-dictionary maps the strings "time", "mag", and "mag_err" to lists containing those values
# The point is to get all the interpolated magnitudes in one place, since they currently exist in many different files,
# all while keeping track of which grid point each one corresponds to, its time, and its band, and to minimize lookup times going forward.
interp_data = [{} for _ in range(theta_grid.size)]

# Fill this data structure with data loaded from the interpolated magnitude files
fname_base = args.interp_directory + ("/" if args.interp_directory[-1] != "/" else "")
for fname in os.listdir(args.interp_directory):

    fname_split = fname.split("_")

    # Determine whether it's a file we want or not
    if fname_split[0] != "eval":
        continue

    fname_split = fname_split[2:]

    # Get the time, angle, and band from the filename
    t_interp = interpolator_times_str_to_float[fname_split[0]] # convert from the truncated time in the filename to the full time value
    theta_interp = float(fname_split[1])
    band = fname_split[2][0]

    if band not in args.bands:
        continue

    if indices[theta_interp].size == 0:
        continue

    # Load the interpolated magnitude data
    data = np.loadtxt(fname_base + fname)

    # "i" is the index within this data file
    # "grid_index" is the row of the grid corresponding to this particular interpolator evaluation
    for i, grid_index in enumerate(indices[theta_interp]):
        # initialize things if needed
        if (band, theta_interp) not in interp_data[grid_index].keys():
            interp_data[grid_index][(band, theta_interp)] = {"time":[], "mag":[], "mag_err":[]}
        interp_data[grid_index][(band, theta_interp)]["time"].append(t_interp)
        interp_data[grid_index][(band, theta_interp)]["mag"].append(data[i][0])
        interp_data[grid_index][(band, theta_interp)]["mag_err"].append(data[i][1])

# Go through and sort every array in interp_data by time (and also convert them to numpy arrays)
for i in range(theta_grid.size):
    for key in interp_data[i].keys():
        interp_data[i][key]["time"] = np.array(interp_data[i][key]["time"])
        interp_data[i][key]["mag"] = np.array(interp_data[i][key]["mag"])
        interp_data[i][key]["mag_err"] = np.array(interp_data[i][key]["mag_err"])
        sorted_indices = np.argsort(interp_data[i][key]["time"])
        interp_data[i][key]["time"] = interp_data[i][key]["time"][sorted_indices]
        interp_data[i][key]["mag"] = interp_data[i][key]["mag"][sorted_indices]
        interp_data[i][key]["mag_err"] = interp_data[i][key]["mag_err"][sorted_indices]

# Now build a list of interpolators corresponding to each row in the grid for each band, where the interpolation over angle has been done
lc_functions = []
for i, theta in enumerate(theta_grid):
    # Determine the angular bin
    for j in range(1, len(interp_angles)):
        if interp_angles[j] > theta:
            break
    theta_lower = interp_angles[j - 1]
    theta_upper = interp_angles[j]

    f_dict = {}
    for b in args.bands:
        # Get the interpolated values for this row in the grid
        t = interp_data[i][(b, theta_lower)]["time"]
        mag_lower = interp_data[i][(b, theta_lower)]["mag"]
        mag_err_lower = interp_data[i][(b, theta_lower)]["mag_err"]
        mag_upper = interp_data[i][(b, theta_upper)]["mag"]
        mag_err_upper = interp_data[i][(b, theta_upper)]["mag_err"]

        # Interpolate in angle
        mag = ((theta_upper - theta) / (theta_upper - theta_lower)) * mag_lower + ((theta - theta_lower) / (theta_upper - theta_lower)) * mag_upper
        mag_err = ((theta_upper - theta) / (theta_upper - theta_lower)) * mag_err_lower + ((theta - theta_lower) / (theta_upper - theta_lower)) * mag_err_upper

        # Interpolate in time and put the interpolators in the dictionary
        f_dict[b] = (interp1d(t, mag, fill_value="extrapolate"), interp1d(t, mag_err, fill_value="extrapolate"))

    lc_functions.append(f_dict)

# Now, finally, compute the likelihoods
lnL = np.zeros(theta_grid.size)
for i, f_dict in enumerate(lc_functions):
    for b in args.bands:
        mag_interpolator, mag_err_interpolator = f_dict[b]
        # make sure to account for the distance modulus when computing residuals
        residuals = mag_interpolator(lc_data[b]["time"]) + (5. * (np.log10(args.distance * 1.e6) - 1.)) - lc_data[b]["mag"]
        model_error = mag_err_interpolator(lc_data[b]["time"])
        lnL[i] += -0.5 * np.sum(residuals**2 / (model_error**2 + lc_data[b]["mag_err"]**2)
            + np.log(2. * np.pi * (model_error**2 + lc_data[b]["mag_err"]**2)))

grid[:,0] = lnL

# Get the header for the posterior sample file from the grid file, then save
with open(args.grid_file, "r") as fp:
    header = fp.readline().strip()[2:] # strip off the "# " at the beginning and the "\n" at the end
np.savetxt(args.output_file, grid, header=header)


import smart
import numpy as np
import sys, os, os.path, time
from astropy.table import Table
from numpy.linalg import inv, det
from ..utils.interpolations import trilinear_interpolation

##############################################################################################################

def InterpModel(teff, logg=4, metal=0, alpha=0, modelset='marcs-apogee-dr15', instrument='nirspec', order=33):

    FULL_PATH  = os.path.realpath(__file__)
    BASE, NAME = os.path.split(FULL_PATH)

    # Check the model set and instrument
    if instrument == 'nirspec':
        if modelset.lower() == 'btsettl08':
            path     = BASE + '/../libraries/btsettl08/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams.csv'
        elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
            path     = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams.csv'
        elif modelset.lower() == 'sonora-2018':
            path     = BASE + '/../libraries/SONORA_2018/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/SONORA_2018/SONORA_2018_gridparams.csv'

    elif instrument == 'apogee':
        if modelset.lower() == 'btsettl08':
            path     = BASE + '/../libraries/btsettl08/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams_apogee.csv'
        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            path     = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/PHOENIX_BTSETTL_CIFIST2011_2015_gridparams_apogee.csv'
        elif modelset.lower() == 'phoenix-aces-agss-cond-2011' :
            path     = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams_apogee.csv'
        elif modelset.lower() == 'marcs-apogee-dr15' :
            path     = BASE + '/../libraries/MARCS_APOGEE_DR15/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/MARCS_APOGEE_DR15/MARCS_APOGEE_DR15_gridparams_apogee.csv'

    # Read the grid file
    T1 = Table.read(Gridfile)

    ###################################################################################

    def GetModel(temp, wave=False, **kwargs):

        logg       = kwargs.get('logg', 4.5)
        metal      = kwargs.get('metal', 0)
        alpha      = kwargs.get('alpha', 0)
        gridfile   = kwargs.get('gridfile', None)
        instrument = kwargs.get('instrument', 'nirspec')
        order      = kwargs.get('order', None)

        if gridfile is None:
            raise ValueError('Model gridfile must be provided.')

        if instrument == 'nirspec':
            if modelset.lower() == 'btsettl08':
                filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(alpha)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
            elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
                filename = 'PHOENIX_ACES_AGSS_COND_2011_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_alpha{0:.2f}'.format(float(alpha)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
            elif modelset.lower() == 'sonora-2018':
                filename = 'SONORA_2018_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_FeH{0:.2f}'.format(0) + '_Y{0:.2f}'.format(0.28) + '_CO{0:.2f}'.format(1.00) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

        if instrument == 'apogee':
            filename = gridfile['File'][np.where( (gridfile['Temp']==temp) & (gridfile['Logg']==logg) & (gridfile['Metal']==metal) & (gridfile['Alpha']==alpha) )].data[0]

        Tab = Table.read(path+filename, format='ascii.tab', names=['wave', 'flux'])

        if wave:
            return Tab['wave']
        else:
            return Tab['flux']

    ###################################################################################

    # Check if the model already exists (grid point)
    if modelset.lower() == 'sonora-2018':
        if (teff, logg) in zip(T1['Temp'], T1['Logg']):
            metal, ys = 0, 0.28
            index0 = np.where( (T1['Temp'] == teff) & (T1['Logg'] == logg) & (T1['FeH'] == metal) & (T1['Y'] == ys) )
            #flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
            #waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
            flux2  = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['FeH'][index0], alpha=T1['Y'][index0], instrument=instrument, order=order, gridfile=T1)
            waves2 = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['FeH'][index0], alpha=T1['Y'][index0], instrument=instrument, order=order, gridfile=T1, wave=True)
            return waves2, flux2
    else:
        if (teff, logg, metal, alpha) in zip(T1['Temp'], T1['Logg'], T1['Metal'], T1['Alpha']):
            index0 = np.where( (T1['Temp'] == teff) & (T1['Logg'] == logg) & (T1['Metal'] == metal) & (T1['Alpha'] == alpha) )
            #flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
            #waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
            flux2  = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['Metal'][index0], alpha=T1['Alpha'][index0], instrument=instrument, order=order, gridfile=T1)
            waves2 = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['Metal'][index0], alpha=T1['Alpha'][index0], instrument=instrument, order=order, gridfile=T1, wave=True)
            return waves2, flux2

    try:
        if modelset.lower() == 'sonora-2018':
            metal, alpha = 0, 0.28
            # Get the nearest models to the gridpoint (Temp)
            x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= teff)])
            x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= teff)])
            #print(x0, x1)

            # Get the nearest grid point to Logg
            y0 = np.max(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] <= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] <= logg) )])))
            y1 = np.min(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] >= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] >= logg) )])))
            #print(y0, y1)

            # Get the nearest grid point to [M/H]
            z0 = np.max(list(set(T1['FeH'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] <= metal) )]) &
                             set(T1['FeH'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] <= metal) )])))
            z1 = np.min(list(set(T1['FeH'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] >= metal) )]) &
                             set(T1['FeH'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] >= metal) )])))
            #print(z0, z1)

            # Get the nearest grid point to Alpha
            t0 = np.max(list(set(T1['Y'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] <= alpha) )]) &
                             set(T1['Y'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] <= alpha) )])))
            t1 = np.min(list(set(T1['Y'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] >= alpha) )]) &
                             set(T1['Y'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] >= alpha) )])))
            #print(t0, t1)

        else:
            # Get the nearest models to the gridpoint (Temp)
            x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= teff)])
            x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= teff)])
            #print('teff:', x0, teff, x1)
            # Get the nearest grid point to Logg
            y0 = np.max(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] <= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] <= logg) )])))
            y1 = np.min(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] >= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] >= logg) )])))
            #print('logg:', y0, logg, y1)
            # Get the nearest grid point to [M/H]
            #print(metal)
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) )])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) )])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] <= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] >= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= metal))])))
            z0 = np.max(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] <= metal) )]) &
                             set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= metal) )])))
            z1 = np.min(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] >= metal) )]) &
                             set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= metal) )])))
            #print('metal:', z0, metal, z1)
            # Get the nearest grid point to Alpha
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] <= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] <= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] >= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] >= alpha) )])))
            t0 = np.max(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] <= alpha) )]) &
                             set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] <= alpha) )])))
            t1 = np.min(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] >= alpha) )]) &
                             set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] >= alpha) )])))
            #print('alpha:', z0, alpha, z1)
    except:
        raise ValueError('Model Parameters Teff: %0.3f, Logg: %0.3f, [M/H]: %0.3f, Alpha: %0.3f are outside the model grid.'%(teff, logg, metal, alpha))

    if modelset.lower() == 'sonora-2018':
        # Get the 16 points
        ind0000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 0000
        ind1000 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 1000
        ind0100 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 0100
        ind0010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 0010
        ind0001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 0001
        ind1001 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 1001
        ind0101 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 0101
        ind0011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 0011
        ind1011 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 1011
        ind0111 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 0111
        ind1111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 1111
        ind0110 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 0110
        ind1010 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 1010
        ind1100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 1100
        ind1101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 1101
        ind1110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 1110
        Points =  [ [np.log10(T1['Temp'][ind0000]), T1['Logg'][ind0000], T1['FeH'][ind0000], T1['Y'][ind0000],
                     np.log10(GetModel(T1['Temp'][ind0000], logg=T1['Logg'][ind0000], metal=T1['FeH'][ind0000], alpha=T1['Y'][ind0000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1000]), T1['Logg'][ind1000], T1['FeH'][ind1000], T1['Y'][ind1000],
                     np.log10(GetModel(T1['Temp'][ind1000], logg=T1['Logg'][ind1000], metal=T1['FeH'][ind1000], alpha=T1['Y'][ind1000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0100]), T1['Logg'][ind0100], T1['FeH'][ind0100], T1['Y'][ind0100],
                     np.log10(GetModel(T1['Temp'][ind0100], logg=T1['Logg'][ind0100], metal=T1['FeH'][ind0100], alpha=T1['Y'][ind0100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0010]), T1['Logg'][ind0010], T1['FeH'][ind0010], T1['Y'][ind0010],
                     np.log10(GetModel(T1['Temp'][ind0010], logg=T1['Logg'][ind0010], metal=T1['FeH'][ind0010], alpha=T1['Y'][ind0010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0001]), T1['Logg'][ind0001], T1['FeH'][ind0001], T1['Y'][ind0001],
                     np.log10(GetModel(T1['Temp'][ind0001], logg=T1['Logg'][ind0001], metal=T1['FeH'][ind0001], alpha=T1['Y'][ind0001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1001]), T1['Logg'][ind1001], T1['FeH'][ind1001], T1['Y'][ind1001],
                     np.log10(GetModel(T1['Temp'][ind1001], logg=T1['Logg'][ind1001], metal=T1['FeH'][ind1001], alpha=T1['Y'][ind1001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0101]), T1['Logg'][ind0101], T1['FeH'][ind0101], T1['Y'][ind0101],
                     np.log10(GetModel(T1['Temp'][ind0101], logg=T1['Logg'][ind0101], metal=T1['FeH'][ind0101], alpha=T1['Y'][ind0101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0011]), T1['Logg'][ind0011], T1['FeH'][ind0011], T1['Y'][ind0011],
                     np.log10(GetModel(T1['Temp'][ind0011], logg=T1['Logg'][ind0011], metal=T1['FeH'][ind0011], alpha=T1['Y'][ind0011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1011]), T1['Logg'][ind1011], T1['FeH'][ind1011], T1['Y'][ind1011],
                     np.log10(GetModel(T1['Temp'][ind1011], logg=T1['Logg'][ind1011], metal=T1['FeH'][ind1011], alpha=T1['Y'][ind1011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0111]), T1['Logg'][ind0111], T1['FeH'][ind0111], T1['Y'][ind0111],
                     np.log10(GetModel(T1['Temp'][ind0111], logg=T1['Logg'][ind0111], metal=T1['FeH'][ind0111], alpha=T1['Y'][ind0111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1111]), T1['Logg'][ind1111], T1['FeH'][ind1111], T1['Y'][ind1111],
                     np.log10(GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['FeH'][ind1111], alpha=T1['Y'][ind1111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0110]), T1['Logg'][ind0110], T1['FeH'][ind0110], T1['Y'][ind0110],
                     np.log10(GetModel(T1['Temp'][ind0110], logg=T1['Logg'][ind0110], metal=T1['FeH'][ind0110], alpha=T1['Y'][ind0110], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1010]), T1['Logg'][ind1010], T1['FeH'][ind1010], T1['Y'][ind1010],
                     np.log10(GetModel(T1['Temp'][ind1010], logg=T1['Logg'][ind1010], metal=T1['FeH'][ind1010], alpha=T1['Y'][ind1010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1100]), T1['Logg'][ind1100], T1['FeH'][ind1100], T1['Y'][ind1100],
                     np.log10(GetModel(T1['Temp'][ind1100], logg=T1['Logg'][ind1100], metal=T1['FeH'][ind1100], alpha=T1['Y'][ind1100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1101]), T1['Logg'][ind1101], T1['FeH'][ind1101], T1['Y'][ind1101],
                     np.log10(GetModel(T1['Temp'][ind1101], logg=T1['Logg'][ind1101], metal=T1['FeH'][ind1101], alpha=T1['Y'][ind1101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1110]), T1['Logg'][ind1110], T1['FeH'][ind1110], T1['Y'][ind1110],
                     np.log10(GetModel(T1['Temp'][ind1110], logg=T1['Logg'][ind1110], metal=T1['FeH'][ind1110], alpha=T1['Y'][ind1110], instrument=instrument, order=order, gridfile=T1))],
                  ]
        #print(Points)
        waves2 = GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['FeH'][ind1111], alpha=T1['Y'][ind1111], instrument=instrument, order=order, gridfile=T1, wave=True)
    else:
        # Get the 16 points
        ind0000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 0000
        ind1000 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 1000
        ind0100 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 0100
        ind0010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 0010
        ind0001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 0001
        ind1001 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 1001
        ind0101 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 0101
        ind0011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 0011
        ind1011 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 1011
        ind0111 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 0111
        ind1111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 1111
        ind0110 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 0110
        ind1010 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 1010
        ind1100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 1100
        ind1101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 1101
        ind1110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 1110
        Points =  [ [np.log10(T1['Temp'][ind0000]), T1['Logg'][ind0000], T1['Metal'][ind0000], T1['Alpha'][ind0000],
                     np.log10(GetModel(T1['Temp'][ind0000], logg=T1['Logg'][ind0000], metal=T1['Metal'][ind0000], alpha=T1['Alpha'][ind0000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1000]), T1['Logg'][ind1000], T1['Metal'][ind1000], T1['Alpha'][ind1000],
                     np.log10(GetModel(T1['Temp'][ind1000], logg=T1['Logg'][ind1000], metal=T1['Metal'][ind1000], alpha=T1['Alpha'][ind1000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0100]), T1['Logg'][ind0100], T1['Metal'][ind0100], T1['Alpha'][ind0100],
                     np.log10(GetModel(T1['Temp'][ind0100], logg=T1['Logg'][ind0100], metal=T1['Metal'][ind0100], alpha=T1['Alpha'][ind0100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0010]), T1['Logg'][ind0010], T1['Metal'][ind0010], T1['Alpha'][ind0010],
                     np.log10(GetModel(T1['Temp'][ind0010], logg=T1['Logg'][ind0010], metal=T1['Metal'][ind0010], alpha=T1['Alpha'][ind0010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0001]), T1['Logg'][ind0001], T1['Metal'][ind0001], T1['Alpha'][ind0001],
                     np.log10(GetModel(T1['Temp'][ind0001], logg=T1['Logg'][ind0001], metal=T1['Metal'][ind0001], alpha=T1['Alpha'][ind0001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1001]), T1['Logg'][ind1001], T1['Metal'][ind1001], T1['Alpha'][ind1001],
                     np.log10(GetModel(T1['Temp'][ind1001], logg=T1['Logg'][ind1001], metal=T1['Metal'][ind1001], alpha=T1['Alpha'][ind1001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0101]), T1['Logg'][ind0101], T1['Metal'][ind0101], T1['Alpha'][ind0101],
                     np.log10(GetModel(T1['Temp'][ind0101], logg=T1['Logg'][ind0101], metal=T1['Metal'][ind0101], alpha=T1['Alpha'][ind0101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0011]), T1['Logg'][ind0011], T1['Metal'][ind0011], T1['Alpha'][ind0011],
                     np.log10(GetModel(T1['Temp'][ind0011], logg=T1['Logg'][ind0011], metal=T1['Metal'][ind0011], alpha=T1['Alpha'][ind0011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1011]), T1['Logg'][ind1011], T1['Metal'][ind1011], T1['Alpha'][ind1011],
                     np.log10(GetModel(T1['Temp'][ind1011], logg=T1['Logg'][ind1011], metal=T1['Metal'][ind1011], alpha=T1['Alpha'][ind1011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0111]), T1['Logg'][ind0111], T1['Metal'][ind0111], T1['Alpha'][ind0111],
                     np.log10(GetModel(T1['Temp'][ind0111], logg=T1['Logg'][ind0111], metal=T1['Metal'][ind0111], alpha=T1['Alpha'][ind0111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1111]), T1['Logg'][ind1111], T1['Metal'][ind1111], T1['Alpha'][ind1111],
                     np.log10(GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['Metal'][ind1111], alpha=T1['Alpha'][ind1111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0110]), T1['Logg'][ind0110], T1['Metal'][ind0110], T1['Alpha'][ind0110],
                     np.log10(GetModel(T1['Temp'][ind0110], logg=T1['Logg'][ind0110], metal=T1['Metal'][ind0110], alpha=T1['Alpha'][ind0110], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1010]), T1['Logg'][ind1010], T1['Metal'][ind1010], T1['Alpha'][ind1010],
                     np.log10(GetModel(T1['Temp'][ind1010], logg=T1['Logg'][ind1010], metal=T1['Metal'][ind1010], alpha=T1['Alpha'][ind1010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1100]), T1['Logg'][ind1100], T1['Metal'][ind1100], T1['Alpha'][ind1100],
                     np.log10(GetModel(T1['Temp'][ind1100], logg=T1['Logg'][ind1100], metal=T1['Metal'][ind1100], alpha=T1['Alpha'][ind1100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1101]), T1['Logg'][ind1101], T1['Metal'][ind1101], T1['Alpha'][ind1101],
                     np.log10(GetModel(T1['Temp'][ind1101], logg=T1['Logg'][ind1101], metal=T1['Metal'][ind1101], alpha=T1['Alpha'][ind1101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1110]), T1['Logg'][ind1110], T1['Metal'][ind1110], T1['Alpha'][ind1110],
                     np.log10(GetModel(T1['Temp'][ind1110], logg=T1['Logg'][ind1110], metal=T1['Metal'][ind1110], alpha=T1['Alpha'][ind1110], instrument=instrument, order=order, gridfile=T1))],
                  ]
        #print(Points)
        waves2 = GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['Metal'][ind1111], alpha=T1['Alpha'][ind1111], instrument=instrument, order=order, gridfile=T1, wave=True)

    return waves2, smart.utils.interpolations.quadlinear_interpolation(np.log10(teff), logg, metal, alpha, Points)

################################################################################################################################################################################################

def InterpModel_3D(Teff, Logg, Metal, modelset='marcs-apogee-dr15', instrument='nirspec', order=33):
    Alpha = 0.0
    FULL_PATH  = os.path.realpath(__file__)
    BASE, NAME = os.path.split(FULL_PATH)

    # Check the model set
    if instrument == 'nirspec':

        if modelset.lower() == 'btsettl08':
            path = BASE + '/../libraries/btsettl08/NIRSPEC-O%s-RAW/'%order

        elif modelset.lower() == 'phoenixaces' :
            path = BASE + '/../libraries/phoenixaces/NIRSPEC-O%s-RAW/'%order

    elif instrument == 'apogee':

        if modelset.lower() == 'btsettl08':
            path = BASE + '/../libraries/btsettl08/APOGEE-RAW/'

        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            path = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/APOGEE-RAW/'

        elif modelset.lower() == 'phoenix-aces-agss-cond-2011' :
            path = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/APOGEE-RAW/'

        elif modelset.lower() == 'marcs-apogee-dr15' :
            path = BASE + '/../libraries/MARCS_APOGEE_DR15/APOGEE-RAW_3D/'

    def GetModel(temp, logg, metal, modelset='marcs-apogee-dr15', wave=False):
        en = 0.00
        if instrument == 'nirspec':

            if modelset.lower() == 'btsettl08':
                # alpha enhancement correction
                if metal == -0.5:
                    en = 0.2
                elif (metal == -1.0) or (metal == -1.5) or (metal == -2.0) or (metal == -2.5):
                    en = 0.4

                if metal == 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
                elif metal == 0.5:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
                else:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

            elif modelset.lower() == 'phoenixaces':
                filename = 'phoenixaces_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z-{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

        elif instrument == 'apogee':
            if modelset.lower() == 'btsettl08':
                # alpha enhancement correction
                if metal == -0.5:
                    en = 0.2
                elif (metal == -1.0) or (metal == -1.5) or (metal == -2.0) or (metal == -2.5):
                    en = 0.4

                if metal == 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'
                elif metal > 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'
                else:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
                filename = 'PHOENIX_BTSETTL_CIFIST2011_2015_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
                filename = 'PHOENIX_ACES_AGSS_COND_2011_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'marcs-apogee-dr15':
                filename = 'MARCS_APOGEE_DR15_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

        Tab = Table.read(path+filename, format='ascii.tab', names=['wave', 'flux'])

        if wave:
            return Tab['wave']
        else:
            return Tab['flux']

    if instrument == 'nirspec':

        if modelset.lower() == 'btsettl08':
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams.csv'

        elif modelset.lower() == 'phoenixaces':
            Gridfile = BASE + '/../libraries/phoenixaces/phoenixaces_gridparams.csv'

    elif instrument == 'apogee':

        if modelset.lower() == 'btsettl08':
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams_apogee.csv'

        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            Gridfile = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/PHOENIX_BTSETTL_CIFIST2011_2015_gridparams_apogee.csv'

        elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams_apogee.csv'

        elif modelset.lower() == 'marcs-apogee-dr15':
            Gridfile = BASE + '/../libraries/MARCS_APOGEE_DR15/MARCS_APOGEE_DR15_3D_gridparams_apogee.csv'

    T1 = Table.read(Gridfile)
    # Check if the model already exists (grid point)
    if (Teff, Logg, Metal) in zip(T1['Temp'], T1['Logg'], T1['Metal']):
        if modelset.lower() == 'btsettl08':
            if Metal == -0.5:
                Alpha = 0.2
            elif (Metal == -1.0) or (Metal == -1.5) or (Metal == -2.0) or (Metal == -2.5):
                Alpha = 0.4
        index0 = np.where( (T1['Temp'] == Teff) & (T1['Logg'] == Logg) & (T1['Metal'] == Metal) & (T1['Alpha'] == Alpha) )
        flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
        waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
        return waves2, flux2

    # Get the nearest models to the gridpoint (Temp)
    x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= Teff)])
    x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= Teff)])
    #print(x0, Teff, x1)
    #y0 = T1['Logg'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & (T1['Logg'] <= Logg) )][-1]
    #y1 = T1['Logg'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & (T1['Logg'] >= Logg) )][0]
    #print(x0, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] <= Logg) ) )])))
    #print(x1, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] <= Logg) ) )])))
    #print(x0, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] >= Logg) ) )])))
    #print(x1, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] >= Logg) ) )])))
    y0 = np.max(list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] <= Logg) ) )]) & set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] <= Logg) ) )])))
    y1 = np.min(list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] >= Logg) ) )]) & set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] >= Logg) ) )])))
    #print(y0, Logg, y1)
    #z0 = T1['pgs'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & ( (T1['Logg'] == y0) | (T1['Logg'] == y1) ) & (T1['pgs'] <= PGS) )][-1]
    #z1 = T1['pgs'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & ( (T1['Logg'] == y0) | (T1['Logg'] == y1) ) & (T1['pgs'] >= PGS) )][0]
    #print(x0, y0, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['pgs'] <= PGS))])))
    #print(x1, y1, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) ) & (T1['pgs'] <= PGS))])))
    #print(x0, y0, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['pgs'] >= PGS))])))
    #print(x1, y1, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) ) & (T1['pgs'] >= PGS))])))
    z0 = np.max(list(set(T1['Metal'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['Metal'] <= Metal) )]) & set(T1['Metal'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= Metal)) )])))
    z1 = np.min(list(set(T1['Metal'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['Metal'] >= Metal) )]) & set(T1['Metal'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= Metal)) )])))
    #print(z0, PGS, z1)

    # Check if the gridpoint exists within the model ranges
    for x in [x0, x1]:
        for y in [y0, y1]:
            for z in [z0, z1]:
                if (x, y, z) not in zip(T1['Temp'], T1['Logg'], T1['Metal']):
                    print('No Model', x, y, z)
                    return 1
    '''
    print(np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1)))
    print(np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2)))
    print(np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1)))
    print(np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2)))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2))]))
    '''
    # Get the 16 points
    if modelset.lower() == 'btsettl08':
        if z0 == -0.5:
            Alpha0 = 0.2
        elif (z0 == -1.0) or (z0 == -1.5) or (z0 == -2.0) or (z0 == -2.5):
            Alpha0 = 0.4
        else:
            Alpha0 = 0.0
        if z1 == -0.5:
            Alpha1 = 0.2
        elif (z1 == -1.0) or (z1 == -1.5) or (z1 == -2.0) or (z1 == -2.5):
            Alpha1 = 0.4
        else:
            Alpha1 = 0.0
        ind000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 000
        ind100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 100
        ind010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 010
        ind110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 110
        ind001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 001
        ind101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 101
        ind011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 011
        ind111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 111

    else:
        ind000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 000
        ind100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 100
        ind010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 010
        ind110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 110
        ind001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 001
        ind101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 101
        ind011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 011
        ind111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 111

    Points =  [ [np.log10(T1['Temp'][ind000]), T1['Logg'][ind000], T1['Metal'][ind000],
                 np.log10(GetModel(T1['Temp'][ind000], T1['Logg'][ind000], T1['Metal'][ind000], modelset=modelset))],
                [np.log10(T1['Temp'][ind100]), T1['Logg'][ind100], T1['Metal'][ind100],
                 np.log10(GetModel(T1['Temp'][ind100], T1['Logg'][ind100], T1['Metal'][ind100], modelset=modelset))],
                [np.log10(T1['Temp'][ind010]), T1['Logg'][ind010], T1['Metal'][ind010],
                 np.log10(GetModel(T1['Temp'][ind010], T1['Logg'][ind010], T1['Metal'][ind010], modelset=modelset))],
                [np.log10(T1['Temp'][ind110]), T1['Logg'][ind110], T1['Metal'][ind110],
                 np.log10(GetModel(T1['Temp'][ind110], T1['Logg'][ind110], T1['Metal'][ind110], modelset=modelset))],
                [np.log10(T1['Temp'][ind001]), T1['Logg'][ind001], T1['Metal'][ind001],
                 np.log10(GetModel(T1['Temp'][ind001], T1['Logg'][ind001], T1['Metal'][ind001], modelset=modelset))],
                [np.log10(T1['Temp'][ind101]), T1['Logg'][ind101], T1['Metal'][ind101],
                 np.log10(GetModel(T1['Temp'][ind101], T1['Logg'][ind101], T1['Metal'][ind101], modelset=modelset))],
                [np.log10(T1['Temp'][ind011]), T1['Logg'][ind011], T1['Metal'][ind011],
                 np.log10(GetModel(T1['Temp'][ind011], T1['Logg'][ind011], T1['Metal'][ind011], modelset=modelset))],
                [np.log10(T1['Temp'][ind111]), T1['Logg'][ind111], T1['Metal'][ind111],
                 np.log10(GetModel(T1['Temp'][ind111], T1['Logg'][ind111], T1['Metal'][ind111], modelset=modelset))],
              ]
    #print(Points)
    waves2 = GetModel(T1['Temp'][ind000], T1['Logg'][ind000], T1['Metal'][ind000], wave=True, modelset=modelset)

    return waves2, trilinear_interpolation(np.log10(Teff), Logg, Metal, Points)

    ################################################################################################