# -*- 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 torch
import torch.nn.functional as F

def rot_mat_to_euler(rot_mats):
    # Calculates rotation matrix to euler angles
    # Careful for extreme cases of eular angles like [0.0, pi, 0.0]

    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
                    rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
    return torch.atan2(-rot_mats[:, 2, 0], sy)

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]

    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)

    rel_rot_mat = torch.eye(
        3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0).repeat(
        batch_size, 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)

    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 joints2bones(joints, parents):
    ''' Decompose joints location to bone length and direction.

        Parameters
        ----------
        joints: torch.tensor Bx24x3
    '''
    assert joints.shape[1] == parents.shape[0]
    bone_dirs = torch.zeros_like(joints)
    bone_lens = torch.zeros_like(joints[:, :, :1])

    for c_id in range(parents.shape[0]):
        p_id = parents[c_id]
        if p_id == -1:
            # Parent node
            bone_dirs[:, c_id] = joints[:, c_id]
        else:
            # Child node
            # (B, 3)
            diff = joints[:, c_id] - joints[:, p_id]
            length = torch.norm(diff, dim=1, keepdim=True) + 1e-8
            direct = diff / length

            bone_dirs[:, c_id] = direct
            bone_lens[:, c_id] = length

    return bone_dirs, bone_lens

def bones2joints(bone_dirs, bone_lens, parents):
    ''' Recover bone length and direction to joints location.

        Parameters
        ----------
        bone_dirs: torch.tensor 1x24x3
        bone_lens: torch.tensor Bx24x1
    '''
    batch_size = bone_lens.shape[0]
    joints = torch.zeros_like(bone_dirs).expand(batch_size, 24, 3)

    for c_id in range(parents.shape[0]):
        p_id = parents[c_id]
        if p_id == -1:
            # Parent node
            joints[:, c_id] = bone_dirs[:, c_id]
        else:
            # Child node
            joints[:, c_id] = joints[:, p_id] + bone_dirs[:, c_id] * bone_lens[:, c_id]

    return joints

def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, J_regressor_h36m, 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
        rot_mats: torch.tensor BxJx3x3
            The rotation matrics of each joints
    '''
    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)

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        if pose.numel() == batch_size * 24 * 4:
            rot_mats = quat_to_rotmat(pose.reshape(batch_size * 24, 4)).reshape(batch_size, 24, 3, 3)
        else:
            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
    J_transformed, A = batch_rigid_transform(rot_mats, J, parents[:24], 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]

    J_from_verts = vertices2joints(J_regressor_h36m, verts)

    return verts, J_transformed, rot_mats, J_from_verts

def hybrik(betas, global_orient, pose_skeleton, phis,
           v_template, shapedirs, posedirs, J_regressor, J_regressor_h36m, parents, children,
           lbs_weights, dtype=torch.float32, train=False, leaf_thetas=None):
    ''' Performs Linear Blend Skinning with the given shape and skeleton joints

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        global_orient : torch.tensor Bx3
            The tensor of global orientation
        pose_skeleton : torch.tensor BxJ*3
            The pose skeleton in (X, Y, Z) format
        phis : torch.tensor BxJx2
            The rotation on bone axis parameters
        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
        J_regressor_h36m : torch.tensor 17xV
            The regressor array that is used to calculate the 17 Human3.6M joints from
            the position of the vertices
        parents: torch.tensor J
            The array that describes the kinematic parents for the model
        children: dict
            The dictionary that describes the kinematic chidrens 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
        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
        rot_mats: torch.tensor BxJx3x3
            The rotation matrics of each joints
    '''
    batch_size = max(betas.shape[0], pose_skeleton.shape[0])
    device = betas.device

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

    # 2. Get the rest joints
    # NxJx3 array
    if leaf_thetas is not None:
        rest_J = vertices2joints(J_regressor, v_shaped)
    else:
        rest_J = torch.zeros((v_shaped.shape[0], 29, 3), dtype=dtype, device=device)
        rest_J[:, :24] = vertices2joints(J_regressor, v_shaped)

        leaf_number = [411, 2445, 5905, 3216, 6617]
        leaf_vertices = v_shaped[:, leaf_number].clone()
        rest_J[:, 24:] = leaf_vertices

    # 3. Get the rotation matrics
    rot_mats, rotate_rest_pose = batch_inverse_kinematics_transform(
        pose_skeleton, global_orient, phis,
        rest_J.clone(), children, parents, dtype=dtype, train=train,
        leaf_thetas=leaf_thetas)

    test_joints = True
    if test_joints:
        J_transformed, A = batch_rigid_transform(rot_mats, rest_J[:, :24].clone(), parents[:24], dtype=dtype)
    else:
        J_transformed = None

    # assert torch.mean(torch.abs(rotate_rest_pose - J_transformed)) < 1e-5
    # 4. Add pose blend shapes
    # rot_mats: N x (J + 1) x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    pose_feature = (rot_mats[:, 1:] - ident).view([batch_size, -1])
    pose_offsets = torch.matmul(pose_feature, posedirs) \
        .view(batch_size, -1, 3)

    v_posed = pose_offsets + v_shaped

    # 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]
    J_from_verts_h36m = vertices2joints(J_regressor_h36m, verts)

    return verts, J_transformed, rot_mats, J_from_verts_h36m

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 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 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 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 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. (Template Pose)
    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:]].clone()

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

    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
        # (B, 4, 4) x (B, 4, 4)
        curr_res = torch.matmul(transform_chain[parents[i]],
                                transforms_mat[:, i])
        transform_chain.append(curr_res)

    # (B, K + 1, 4, 4)
    transforms = torch.stack(transform_chain, dim=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

def batch_inverse_kinematics_transform(
        pose_skeleton, global_orient,
        phis,
        rest_pose,
        children, parents, dtype=torch.float32, train=False,
        leaf_thetas=None):
    """
    Applies a batch of inverse kinematics transfoirm to the joints

    Parameters
    ----------
    pose_skeleton : torch.tensor BxNx3
        Locations of estimated pose skeleton.
    global_orient : torch.tensor Bx1x3x3
        Tensor of global rotation matrices
    phis : torch.tensor BxNx2
        The rotation on bone axis parameters
    rest_pose : torch.tensor Bx(N+1)x3
        Locations of rest_pose. (Template Pose)
    children: dict
        The dictionary that describes the kinematic chidrens for the model
    parents : torch.tensor Bx(N+1)
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    rot_mats: torch.tensor Bx(N+1)x3x3
        The rotation matrics of each joints
    rel_transforms : torch.tensor Bx(N+1)x4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """
    batch_size = pose_skeleton.shape[0]
    device = pose_skeleton.device

    rel_rest_pose = rest_pose.clone()
    rel_rest_pose[:, 1:] -= rest_pose[:, parents[1:]].clone()
    rel_rest_pose = torch.unsqueeze(rel_rest_pose, dim=-1)

    # rotate the T pose
    rotate_rest_pose = torch.zeros_like(rel_rest_pose)
    # set up the root
    rotate_rest_pose[:, 0] = rel_rest_pose[:, 0]

    rel_pose_skeleton = torch.unsqueeze(pose_skeleton.clone(), dim=-1).detach()
    rel_pose_skeleton[:, 1:] = rel_pose_skeleton[:, 1:] - rel_pose_skeleton[:, parents[1:]].clone()
    rel_pose_skeleton[:, 0] = rel_rest_pose[:, 0]

    # the predicted final pose
    final_pose_skeleton = torch.unsqueeze(pose_skeleton.clone(), dim=-1)
    final_pose_skeleton = final_pose_skeleton - final_pose_skeleton[:, 0:1] + rel_rest_pose[:, 0:1]

    rel_rest_pose = rel_rest_pose
    rel_pose_skeleton = rel_pose_skeleton
    final_pose_skeleton = final_pose_skeleton
    rotate_rest_pose = rotate_rest_pose

    assert phis.dim() == 3
    phis = phis / (torch.norm(phis, dim=2, keepdim=True) + 1e-8)

    # TODO
    if train:
        global_orient_mat = batch_get_pelvis_orient(
            rel_pose_skeleton.clone(), rel_rest_pose.clone(), parents, children, dtype)
    else:
        global_orient_mat = batch_get_pelvis_orient_svd(
            rel_pose_skeleton.clone(), rel_rest_pose.clone(), parents, children, dtype)

    rot_mat_chain = [global_orient_mat]
    rot_mat_local = [global_orient_mat]
    # leaf nodes rot_mats
    if leaf_thetas is not None:
        leaf_cnt = 0
        leaf_rot_mats = leaf_thetas.view([batch_size, 5, 3, 3])

    for i in range(1, parents.shape[0]):
        if children[i] == -1:
            # leaf nodes
            if leaf_thetas is not None:
                rot_mat = leaf_rot_mats[:, leaf_cnt, :, :]
                leaf_cnt += 1

                rotate_rest_pose[:, i] = rotate_rest_pose[:, parents[i]] + torch.matmul(
                    rot_mat_chain[parents[i]],
                    rel_rest_pose[:, i]
                )

                rot_mat_chain.append(torch.matmul(
                    rot_mat_chain[parents[i]],
                    rot_mat))
                rot_mat_local.append(rot_mat)
        elif children[i] == -3:
            # three children
            rotate_rest_pose[:, i] = rotate_rest_pose[:, parents[i]] + torch.matmul(
                rot_mat_chain[parents[i]],
                rel_rest_pose[:, i]
            )

            spine_child = []
            for c in range(1, parents.shape[0]):
                if parents[c] == i and c not in spine_child:
                    spine_child.append(c)

            # original
            spine_child = []
            for c in range(1, parents.shape[0]):
                if parents[c] == i and c not in spine_child:
                    spine_child.append(c)

            children_final_loc = []
            children_rest_loc = []
            for c in spine_child:
                temp = final_pose_skeleton[:, c] - rotate_rest_pose[:, i]
                children_final_loc.append(temp)

                children_rest_loc.append(rel_rest_pose[:, c].clone())

            rot_mat = batch_get_3children_orient_svd(
                children_final_loc, children_rest_loc,
                rot_mat_chain[parents[i]], spine_child, dtype)

            rot_mat_chain.append(
                torch.matmul(
                    rot_mat_chain[parents[i]],
                    rot_mat)
            )
            rot_mat_local.append(rot_mat)
        else:
            # (B, 3, 1)
            rotate_rest_pose[:, i] = rotate_rest_pose[:, parents[i]] + torch.matmul(
                rot_mat_chain[parents[i]],
                rel_rest_pose[:, i]
            )
            # (B, 3, 1)
            child_final_loc = final_pose_skeleton[:, children[i]] - rotate_rest_pose[:, i]

            if not train:
                orig_vec = rel_pose_skeleton[:, children[i]]
                template_vec = rel_rest_pose[:, children[i]]
                norm_t = torch.norm(template_vec, dim=1, keepdim=True)
                orig_vec = orig_vec * norm_t / torch.norm(orig_vec, dim=1, keepdim=True)

                diff = torch.norm(child_final_loc - orig_vec, dim=1, keepdim=True)
                big_diff_idx = torch.where(diff > 15 / 1000)[0]

                child_final_loc[big_diff_idx] = orig_vec[big_diff_idx]

            child_final_loc = torch.matmul(
                rot_mat_chain[parents[i]].transpose(1, 2),
                child_final_loc)

            child_rest_loc = rel_rest_pose[:, children[i]]
            # (B, 1, 1)
            child_final_norm = torch.norm(child_final_loc, dim=1, keepdim=True)
            child_rest_norm = torch.norm(child_rest_loc, dim=1, keepdim=True)

            child_final_norm = torch.norm(child_final_loc, dim=1, keepdim=True)

            # (B, 3, 1)
            axis = torch.cross(child_rest_loc, child_final_loc, dim=1)
            axis_norm = torch.norm(axis, dim=1, keepdim=True)

            # (B, 1, 1)
            cos = torch.sum(child_rest_loc * child_final_loc, dim=1, keepdim=True) / (child_rest_norm * child_final_norm + 1e-8)
            sin = axis_norm / (child_rest_norm * child_final_norm + 1e-8)

            # (B, 3, 1)
            axis = axis / (axis_norm + 1e-8)

            # Convert location revolve to rot_mat by rodrigues
            # (B, 1, 1)
            rx, ry, rz = torch.split(axis, 1, dim=1)
            zeros = torch.zeros((batch_size, 1, 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_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

            # Convert spin to rot_mat
            # (B, 3, 1)
            spin_axis = child_rest_loc / child_rest_norm
            # (B, 1, 1)
            rx, ry, rz = torch.split(spin_axis, 1, dim=1)
            zeros = torch.zeros((batch_size, 1, 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)
            # (B, 1, 1)
            cos, sin = torch.split(phis[:, i - 1], 1, dim=1)
            cos = torch.unsqueeze(cos, dim=2)
            sin = torch.unsqueeze(sin, dim=2)
            rot_mat_spin = ident + sin * K + (1 - cos) * torch.bmm(K, K)
            rot_mat = torch.matmul(rot_mat_loc, rot_mat_spin)

            rot_mat_chain.append(torch.matmul(
                rot_mat_chain[parents[i]],
                rot_mat))
            rot_mat_local.append(rot_mat)

    # (B, K + 1, 3, 3)
    rot_mats = torch.stack(rot_mat_local, dim=1)

    return rot_mats, rotate_rest_pose.squeeze(-1)

def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, _, V = torch.svd(S_non_zero)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat

def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    assert children[0] == 3
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, int(children[0])].clone()
    spine_rest_loc = rel_rest_pose[:, int(children[0])].clone()
    spine_norm = torch.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine)
    center_final_loc = 0
    center_rest_loc = 0
    for child in pelvis_child:
        if child == int(children[0]):
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()
    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.norm(center_rest_loc, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat

def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, _, V = torch.svd(S)

    rot_mat = torch.bmm(V, U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat

def vectors2rotmat(vec_rest, vec_final, dtype):
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc

def rotmat_to_quat(rotmat):
    """Convert rotation matrix to quaternion coefficients.
    Args:
        rotmat: size is [B, 3, 3]
    Returns:
        Quaternion: size is [B, 4] <===> (w, x, y, z)
    """
    quaternion = torch.zeros([rotmat.size(0), 4], device=rotmat.device)
    trace = rotmat[:, 0, 0] + rotmat[:, 1, 1] + rotmat[:, 2, 2]
    flag = 1 + trace > 0
    s = torch.zeros_like(trace)

    # pos
    s[flag] = 2 * torch.sqrt(1 + trace[flag]) + 1e-16
    s_pos = s[flag]
    quaternion[flag, 0] = s_pos / 4
    quaternion[flag, 1] = (rotmat[flag, 2, 1] - rotmat[flag, 1, 2]) / s_pos
    quaternion[flag, 2] = (rotmat[flag, 0, 2] - rotmat[flag, 2, 0]) / s_pos
    quaternion[flag, 3] = (rotmat[flag, 1, 0] - rotmat[flag, 0, 1]) / s_pos

    # neg
    diag = torch.stack([rotmat[:, 0, 0], rotmat[:, 1, 1], rotmat[:, 2, 2]])
    max_val, max_ind = torch.max(diag, dim=0)

    s[~flag] = 2 * torch.sqrt(1 - trace[~flag] + 2 * max_val[~flag]) + 1e-16

    f0 = ~flag * (max_ind == 0)
    s0 = s[f0]
    quaternion[f0, 0] = (rotmat[f0, 2, 1] - rotmat[f0, 1, 2]) / s0
    quaternion[f0, 1] = s0 / 4
    quaternion[f0, 2] = (rotmat[f0, 0, 1] + rotmat[f0, 1, 0]) / s0
    quaternion[f0, 3] = (rotmat[f0, 0, 2] + rotmat[f0, 2, 0]) / s0

    f1 = ~flag * (max_ind == 1)
    s1 = s[f1]
    quaternion[f1, 0] = (rotmat[f1, 0, 2] - rotmat[f1, 2, 0]) / s1
    quaternion[f1, 1] = (rotmat[f1, 0, 1] + rotmat[f1, 1, 0]) / s1
    quaternion[f1, 2] = s1 / 4
    quaternion[f1, 3] = (rotmat[f1, 1, 2] + rotmat[f1, 2, 1]) / s1

    f2 = ~flag * (max_ind == 2)
    s2 = s[f2]
    quaternion[f2, 0] = (rotmat[f2, 1, 0] - rotmat[f2, 0, 1]) / s2
    quaternion[f2, 1] = (rotmat[f2, 0, 2] + rotmat[f2, 2, 0]) / s2
    quaternion[f2, 2] = (rotmat[f2, 1, 2] + rotmat[f2, 2, 1]) / s2
    quaternion[f2, 3] = s2 / 4

    return quaternion

def quat_to_rotmat(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    norm_quat = quat
    norm_quat = norm_quat / (norm_quat.norm(p=2, dim=1, keepdim=True) + 1e-8)
    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w * x, w * y, w * z
    xy, xz, yz = x * y, x * z, y * z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz,
                          2 * wz + 2 * xy, w2 - x2 + y2 - z2, 2 * yz - 2 * wx,
                          2 * xz - 2 * wy, 2 * wx + 2 * yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat


#!/usr/bin/python
"""
Do MCMC runs to fit FG models to simulated data, over a grid of
(nu_min, nu_max) values.
"""
import numpy as np
import models
import model_list
import fitting
from utils import rj2cmb, bands_log
import sys, time, os, copy
#from multiprocessing import Pool
from mpi4py import MPI

# Set-up MPI
comm = MPI.COMM_WORLD
myid = comm.Get_rank()
nproc = comm.Get_size()

# Prefix for output files
PREFIX = "paper01"
NBURN = 500
NSTEPS = 10000
NWALKERS = 100

# Reference noise curve (assumes noise file contains sigma_P=sigma_Q=sigma_U
# in uK_CMB.deg for a given experiment)
NOISE_FILE = "data/noise_coreplus_extended.dat"
#NOISE_FILE = "data/core_plus_extended_noise.dat"

# Band parameter definitions
nbands = 7
nthreads = 2

# Set random seed
SEED = 10
if len(sys.argv) > 1:
    SEED = int(sys.argv[1])
print "SEED =", SEED
np.random.seed(SEED)

# Lists of input and output models
in_list = ['cmb', ]; fit_list = ['cmb', ]

# Define input models and their amplitudes/parameters
allowed_comps = model_list.model_dict
cmb_model = model_list.cmb_model

# Parse args to define input and output models
if len(sys.argv) > 2:
    in_list += sys.argv[2].split(",")
    fit_list += sys.argv[3].split(",")
else:
    in_list = ['cmb', 'synch', 'mbb']
    fit_list = ['cmb', 'synch', 'mbb']

# Make sure models are of known types
for item in in_list:
    if item not in allowed_comps.keys():
        raise ValueError("Unknown component type '%s'" % item)
for item in fit_list:
    if item not in allowed_comps.keys():
        raise ValueError("Unknown component type '%s'" % item)

# Print recognised models and specify name
print "Input components:", in_list
print "Fitting components:", fit_list
name_in = "-".join(in_list)
name_fit = "-".join(fit_list)

# Frequency ranges
numin_vals = [20., 30., 40.]
numax_vals = [300., 400., 500., 600., 700., 800.]
#numin_vals = [20., ]
#numax_vals = [300., 700.]

# Temperature/polarisation noise rms for all bands, as a fraction of T_cmb
fsigma_T = 1e4 #1. / np.sqrt(2.)
fsigma_P = 1.

# Collect components into lists and set input amplitudes
mods_in = [allowed_comps[comp] for comp in in_list]
mods_fit = [allowed_comps[comp] for comp in fit_list]
amps_in = np.array([m.amps() for m in mods_in])
params_in = np.array([m.params() for m in mods_in])

# Expand into all combinations of nu_min,max
nu_min, nu_max = np.meshgrid(numin_vals, numax_vals)
nu_params = np.column_stack((nu_min.flatten(), nu_max.flatten()))

# Prepare output files for writing
filename = "output/%s_summary_%s.%s_nb%d_seed%d" \
             % (PREFIX, name_in, name_fit, nbands, SEED)
cut_range = np.arange(NSTEPS, step=200)
for cut in cut_range:
    fname = filename + "_cut%d.dat" % cut
    f = open(fname, 'w')
    f.close()

def model_test(nu, D_vec, Ninv, models_fit, initial_vals=None, burn=500,
               steps=1000, nwalkers=100, cmb_amp_in=None, sample_file=None):
    """
    Generate simulated data given an input model, and perform MCMC fit using
    another model.
    """
    # Collect together data and noise/instrument model
    beam_mat = np.identity(3*len(nu)) # Beam model
    data_spec = (nu, D_vec, Ninv, beam_mat)

    # Get a list of amplitude/parameter names and initial values
    amp_names = []; amp_vals = []; param_names = []; param_vals = []
    amp_parent_model = []; param_parent_model = []
    for mod in models_fit:
        # Parameter names
        amp_names += ["%s_%s" % (mod.model, pol) for pol in "IQU"]
        param_names += mod.param_names

        # Parameter values
        amp_vals = np.concatenate( (amp_vals, mod.amps()) )
        param_vals = np.concatenate( (param_vals, mod.params()) )

        # Parent model list
        amp_parent_model.append(mod)
        param_parent_model.append(mod)

    # Concatenate parameter lists
    pnames = amp_names + param_names
    pvals = np.concatenate((amp_vals, param_vals))
    parent_model = amp_parent_model + param_parent_model

    # Use 'guess' as the initial point for the MCMC if specified
    if initial_vals is None: initial_vals = pvals

    # Collect names, initial values, and parent components for the parameters
    param_spec = (pnames, initial_vals, parent_model)

    # Run MCMC sampler on this model
    t0 = time.time()
    pnames, samples, logp = fitting.joint_mcmc(data_spec, models_fit, param_spec,
                                               burn=burn, steps=steps,
                                               nwalkers=nwalkers,
                                               nthreads=nthreads,
                                               sample_file=sample_file)
    print "MCMC run in %d sec." % (time.time() - t0)

    # Return parameter names and samples
    return pnames, samples, logp, initial_vals

def run_model(nu_params):
    # Get band definition
    nu_min, nu_max = nu_params
    print "nu_min = %d GHz, nu_max = %d GHz" % (nu_min, nu_max)
    nu = bands_log(nu_min, nu_max, nbands)
    label = str(nu_min) + '_' + str(nu_max)

    # Make copies of models
    my_mods_in = mods_in
    my_mods_fit = mods_fit

    # Name of sample file
    #fname_samples = "output/%s_samples_%s.%s_nb%d_seed%d_%s.dat" \
    #              % (PREFIX, name_in, name_fit, nbands, SEED, label)
    fname_samples = None

    # Simulate data and run MCMC fit
    D_vec, Ninv = fitting.generate_data(nu, fsigma_T, fsigma_P,
                                        components=my_mods_in,
                                        noise_file=NOISE_FILE)

    pnames, samples, logp, ini = model_test(nu, D_vec, Ninv, my_mods_fit,
                                            burn=NBURN, steps=NSTEPS,
                                            nwalkers=NWALKERS,
                                            cmb_amp_in=cmb_model.amps(),
                                            sample_file=fname_samples)
    # Calculate best-fit chisq.
    chisq = -2.*logp
    dof = D_vec.size - len(pnames)

    # Get best-fit (max. prob.) parameter values
    maxl_idx = np.argmax(logp)
    bf_params = samples[:, maxl_idx]

    # Reshape sample array into (Nparams, Nwalkers, Nsamples)
    samples = samples.reshape((samples.shape[0],
                               NWALKERS,
                               samples.shape[1]/NWALKERS))

    # Loop over different burn-in cuts to produce summary stats
    for n, cut in enumerate(cut_range):

        # Set output filename
        fname = filename + "_cut%d.dat" % cut

        # Output mean and bias
        summary_data = [nu_min, nu_max, np.min(chisq), maxl_idx]
        header = "nu_min nu_max chi2_min maxlike_idx "

        # Loop over parameter names
        for i in range(len(pnames)):

            # Mean, std. dev., and fractional shift from true value
            # (for mean, median, and max. likelihood param values)
            _mean = np.mean(samples[i,:,cut:])
            _std = np.std(samples[i,:,cut:])
            _fracbias = (np.mean(samples[i,:,cut:]) - ini[i]) \
                      / np.std(samples[i,:,cut:])
            _med_fracbias = (np.median(samples[i,:,cut:]) - ini[i]) \
                          / np.std(samples[i,:,cut:])
            _ml_fracbias = (bf_params[i] - ini[i]) / np.std(samples[i,:,cut:])

            stats = [_mean, _std, _fracbias, _med_fracbias, _ml_fracbias]

            # Keep summary stats, to be written to file
            summary_data += stats
            header += "mean_%s std_%s Delta_%s MedDelta_%s MLDelta_%s " \
                    % (pnames[i], pnames[i], pnames[i], pnames[i], pnames[i])

            # Only output summary stats once
            if n == 0:
                print "%14s: %+3.3e +/- %3.3e [Delta = %+3.3f, MLDelta = %+3.3f]" \
                      % (pnames[i], stats[0], stats[1], stats[2], stats[3])

        # If file is empty, set flag to write header when saving output
        has_header = False if os.stat(fname).st_size == 0 else True

        # Append summary statistics to file
        f = open(fname, 'a')
        if has_header:
            np.savetxt(f, np.atleast_2d(summary_data))
        else:
            np.savetxt(f, np.atleast_2d(summary_data), header=header[:-1])
        f.close()

# Run pool of processes
#pool = Pool(NPROC)
#pool.map(run_model, nu_params)

for i in range(len(nu_params)):
    if i % nproc != myid: continue

    # Run the model for this set of params
    run_model(nu_params[i])


from typing import List, Tuple, Optional, Any
import time
import numba
import numpy as np
import cv2
import reeds_shepp
from .debugger import Debugger

class RRTStar(object):
    def __init__(self):
        self.debug = False
        self.vertices = None  # type: Optional[List[RRTStar.StateNode]]
        self.root = None  # type: Optional[RRTStar.StateNode]
        self.gain = None  # type: Optional[RRTStar.StateNode]
        self.x_best = None  # type: Optional[RRTStar.StateNode]
        self.epsilon = 1e-6

        self.check_res = 0.3  # type: Optional[float]
        self.check_poly = None  # type: Optional[np.ndarray]
        self.exist_res = 0.1
        self.maximum_curvature = 0.2  # type: Optional[float]

        self.start = None  # type: Optional[RRTStar.StateNode]
        self.goal = None  # type: Optional[RRTStar.StateNode]
        self.grid_map = None  # type: Optional[np.ndarray]
        self.grid_res = None  # type: Optional[float]
        self.grid_ori = None  # type: Optional[RRTStar.StateNode]
        self.obstacle = None  # type: Optional[int]
        self.heuristic = None  # type: Optional[List[(Tuple[float], Tuple[float], Tuple[float])]]

    def set_vehicle(self, check_poly, check_res, maximum_curvature):
        # type: (np.ndarray, float, float) -> RRTStar
        """
        set parameter of the vehicle
        :param check_poly: contour of the vehicle for collision check. cv2.contour.
        :param check_res: the resolution of curve interpolation.
        :param maximum_curvature: equal to 1/minimum_turning_radius of the vehicle.
        """
        self.check_poly, self.check_res, self.maximum_curvature = check_poly, check_res, maximum_curvature
        return self

    def preset(self, start, goal, grid_map, grid_res, grid_ori, obstacle, heuristic):
        # type: (StateNode, StateNode, np.ndarray, float, StateNode, int, Any) -> RRTStar
        """
        initialize the parameters for planning: Start State, Goal State and other needs.
        :param start: the start state.
        :param goal: the goal state.
        :param grid_map: occupancy grid map.
        :param grid_res: resolution of grid map.
        :param grid_ori: the center point of the occupancy grid map.
        :param obstacle: the value of pixels of the obstacles region.
        :param heuristic: [(state, biasing)], sampling heuristic path.
            state: state (x_o, y_o, a_o) of the point of the path.
            biasing = (x_mu, x_sigma), (y_mu, y_sigma), (a_mu, a_sigma).
        :return: RRTStar object
        """
        self.grid_map, self.grid_res, self.grid_ori, self.obstacle = grid_map, grid_res, grid_ori, obstacle
        self.heuristic = heuristic
        self.start, self.goal = start, goal
        self.start.g, self.start.hl = 0, self.cost(start, goal)
        self.start.hu = self.start.hl if self.collision_free(start, goal) else np.inf
        self.start.fl, self.start.fu = self.start.g + self.start.hl, self.start.g + self.start.hu
        self.root, self.gain = self.start, self.goal
        self.vertices, self.x_best = [self.root], self.root
        return self

    def planning(self, times, repeat=10, optimize=False, debug=False):
        """main flow."""
        self.debug = debug
        past = time.time()
        for i in range(times):
            x_new = self.sample_free(i, repeat)
            x_nearest = self.nearest(x_new) if not optimize else self.least(x_new)
            if x_nearest and self.benefit(x_new) and self.collision_free(x_nearest, x_new):
                self.attach(x_nearest, x_new)
                self.rewire(x_new)
                self.x_best = self.best()
                self.branch_and_bound()
            Debugger().debug_planned_path(self, i, switch=self.debug)
            Debugger().debug_planning_hist(self, i, (time.time() - past), switch=True)

    def branch_and_bound(self, space=None):
        def out(x):
            vertices.remove(x)
            x.remove()
        vertices = space if space else self.vertices
        vs = filter(lambda x: x.fl > self.x_best.fu + self.epsilon, vertices)
        map(out, vs)
        Debugger().debug_branch_and_bound(vs, switch=self.debug)

    def sample_free(self, n, repeat=10, default=((2., .5), (0., np.pi / 4.), (0, np.pi / 6.))):
        """sample a state from free configuration space."""

        def is_free(state):
            contour = self.transform(self.check_poly, state[0], state[1], state[2])
            contour = np.floor(contour / self.grid_res + self.grid_map.shape[0] / 2.).astype(int)
            mask = np.zeros_like(self.grid_map, dtype=np.uint8)
            cv2.fillPoly(mask, [contour], 255)
            result = np.bitwise_and(mask, self.grid_map)
            Debugger().breaker('sample free: {}'.format(np.all(result < self.obstacle)), self.debug)
            return np.all(result < self.obstacle)

        def exist(state):
            def key(y):
                dxy = np.fabs(y.state[:-1] - s[:-1])
                da = ((y.state[-1] + np.pi) % (2 * np.pi) - np.pi) - ((s[-1] + np.pi) % (2 * np.pi) - np.pi)
                return dxy[0] < self.exist_res and dxy[1] < self.exist_res and da < self.exist_res
            s = np.array(state)
            result = filter(key, self.vertices)
            Debugger.breaker('sample exclusive: {}'.format(result == []), switch=self.debug)
            return result

        def emerge():
            if self.heuristic:
                i = n % len(self.heuristic)
                state, biasing = self.heuristic[i]
                rand = [state[0], state[1], state[2]]  # [x_o, y_o, a_o]
                (x_mu, x_sigma), (y_mu, y_sigma), (a_mu, a_sigma) = biasing
                rand[0] += np.random.normal(x_mu, x_sigma)
                rand[1] += np.random.normal(y_mu, y_sigma)
                rand[2] += np.random.normal(a_mu, a_sigma)
                return rand
            else:
                Debugger().debug_no_heuristic(vertex.state, default, self.debug)
                rand = [vertex.state[0], vertex.state[1], vertex.state[2]]
                (r_mu, r_sigma), (t_mu, t_sigma), (a_mu, a_sigma) = default
                r, theta = np.random.normal(r_mu, r_sigma), np.random.normal(t_mu, t_sigma) + rand[2]
                rand[0] += r * np.cos(theta)
                rand[1] += r * np.sin(theta)
                rand[2] += np.random.normal(a_mu, a_sigma)
                return rand

        vertex = np.random.choice(self.vertices)
        for i in range(repeat):
            x_rand = emerge()
            Debugger().debug_sampling(x_rand, self.check_poly, switch=self.debug)
            if is_free(x_rand):
                if not exist(x_rand):
                    return self.StateNode(tuple(x_rand))
        return self.StateNode(tuple(x_rand))

    def nearest(self, x_rand):  # type: (StateNode) -> StateNode
        """find the state in the tree which is nearest to the sampled state.
        And fill the g, hl and fl properties of the sampled state.
        """

        def replenish(x_n, x_r):
            x_r.g, x_r.hl = x_n.g + self.cost(x_n, x_r), self.cost(x_r, self.gain)
            x_r.fl = x_r.g + x_r.hl

        # quick shot
        if self.collision_free(self.root, x_rand):
            x_nearest = self.root
        else:
            costs = list(map(lambda x: self.cost(x, x_rand), self.vertices))
            x_nearest = self.vertices[int(np.argmin(costs))]
        replenish(x_nearest, x_rand)
        Debugger().debug_nearest_searching(x_nearest.state, switch=self.debug)
        return x_nearest

    def least(self, x_rand):  # type: (StateNode) -> StateNode
        def replenish(x_n, x_r):
            x_r.g, x_r.hl = x_n.g + self.cost(x_n, x_r), self.cost(x_r, self.gain)
            x_r.fl = x_r.g + x_r.hl

        # quick shot
        if self.collision_free(self.root, x_rand):
            x_least = self.root
        else:
            nodes = filter(lambda x: self.collision_free(x, x_rand), self.vertices)
            if nodes:
                costs = list(map(lambda x: x.g + self.cost(x, x_rand), nodes))
                x_least = nodes[int(np.argmin(costs))]
            else:
                x_least = None
        if x_least:
            replenish(x_least, x_rand)
            Debugger().debug_nearest_searching(x_least.state, switch=self.debug)
        return x_least

    def benefit(self, x_new):
        words = 'Benefit: {}/ ({}, {})'.format(x_new.fl <= self.x_best.fu, x_new.fl, self.x_best.fu)
        Debugger.breaker(words, switch=self.debug)
        return x_new.fl < self.x_best.fu

    def collision_free(self, x_from, x_to):  # type: (StateNode, StateNode) -> bool
        """check if the path from one state to another state collides with any obstacles or not."""
        # making contours of the curve
        states = reeds_shepp.path_sample(x_from.state, x_to.state, 1. / self.maximum_curvature, 0.3)
        # states.append(tuple(x_to.state))  # include the end point
        contours = [self.transform(self.check_poly, s[0], s[1], s[2]) for s in states]
        contours = [np.floor(con / self.grid_res + self.grid_map.shape[0] / 2.).astype(int) for con in contours]
        # making mask
        mask = np.zeros_like(self.grid_map, dtype=np.uint8)
        [cv2.fillPoly(mask, [con], 255) for con in contours]
        # checking
        result = np.bitwise_and(mask, self.grid_map)
        Debugger().debug_collision_checking(states, self.check_poly, np.all(result < self.obstacle), switch=self.debug)
        return np.all(result < self.obstacle)

    @staticmethod
    @numba.njit
    def transform(poly, x, y, a):
        pts = poly.transpose()
        xyo = np.array([[x], [y]])
        rot = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
        return (np.dot(rot, pts) + xyo).transpose()

    def attach(self, x_nearest, x_new):  # type: (StateNode, StateNode) -> None
        """add the new state to the tree and complement other values.
        And fill the hu and fu properties of x_new.
        """
        x_new.match(x_nearest)
        available = self.collision_free(x_new, self.gain)
        x_new.hu = x_new.hl if available else np.inf
        x_new.fu = x_new.g + x_new.hu
        x_new.status = 0 if available else 1
        self.vertices.append(x_new)
        Debugger().debug_attaching(x_nearest, x_new, 1. / self.maximum_curvature, switch=self.debug)

    def rewire(self, x_new, gamma=0.2):  # type: (StateNode, float) -> None
        """rewiring tree by the new state."""

        def recheck(x):
            available, cost = self.collision_free(x_new, x), self.cost(x_new, x)
            if available and x.g > x_new.g + cost:
                Debugger().debug_rewiring(x, x_new.g + cost, switch=self.debug)
                x.g = x_new.g + cost
                x.fu, x.fl = x.g + x.hu, x.g + x.hl
                x.rematch(x_new)

        xs = filter(lambda x: x.g > x_new.g + gamma, self.vertices)
        Debugger().debug_rewiring_check(xs, x_new, switch=self.debug)
        map(recheck, xs)

    def cost(self, x_from, x_to):  # type: (StateNode, StateNode) -> float
        """calculate the cost from one state to another state"""
        return reeds_shepp.path_length(x_from.state, x_to.state, 1. / self.maximum_curvature)

    def best(self):
        return sorted(self.vertices, key=lambda x: x.fu)[0]

    def spare(self, space=None):
        return sorted(self.vertices if not space else space, key=lambda x: x.hl)[0]

    def path(self):  # type: () -> List[RRTStar.StateNode]
        """extract the planning result. including the goal state if the the goal is available"""
        x_best = self.best()
        if x_best.fu < np.inf:
            p = x_best.trace()
            self.gain.g = self.gain.fu = x_best.fu
            p.append(self.gain)
            return p
        else:
            x_spare = self.spare()
            return x_spare.trace()

    def trajectory(self, a_cc=4, v_max=10, res=0.1):
        """
        planning velocity for a path to generate a trajectory.
        """

        def interpolate(q_ori, segment_type, length, radius):
            x0, y0, a0 = q_ori[0], q_ori[1], q_ori[2]
            transfer = np.array([[np.cos(a0), -np.sin(a0), 0., x0], [np.sin(a0), np.cos(a0), 0., y0], [0., 0., 1., a0]])
            sign, phi = np.sign(length), np.fabs(length) / radius
            if segment_type == 1:
                r = 2 * radius * np.sin(phi / 2.)
                x_lcs = np.array([[r * sign * np.cos(phi / 2.)], [r * np.sin(phi / 2.)], [sign * phi], [1]])
            elif segment_type == 3:
                r = 2 * radius * np.sin(phi / 2.)
                x_lcs = np.array([[r * sign * np.cos(phi / 2.)], [- r * np.sin(phi / 2.)], [- sign * phi], [1]])
            else:
                x_lcs = np.array([[length], [0], [0], [1]])
            x_tar = np.dot(transfer, x_lcs)
            return x_tar[0, 0], x_tar[1, 0], x_tar[2, 0]

        def extract_segments(q_from, q_to):
            segments.extend(reeds_shepp.path_type(q_from, q_to, 1. / self.maximum_curvature))
            return q_to

        def extract_discontinuities(q0, sgs):
            sg0, sg1 = sgs[0], sgs[1]
            q1 = interpolate(q0, sg0[0], sg0[1], 1. / self.maximum_curvature)
            if sg0[1] * sg1[1] < 0:
                discontinuities2.append(self.Configuration(q1, v=0))
            else:
                discontinuities2.append(self.Configuration(q1))
            return q1

        def plan_motions(index_sec):
            motion = []
            v0, v1 = 0, 0
            seg = discontinuities2[index_sec[0]:index_sec[1]+1]
            seg = list(zip(seg[:-1], seg[1:]))
            extent = 0
            for sec in seg:
                s0, s1 = sec[0].state, sec[1].state
                extent += np.abs(reeds_shepp.path_length(s0, s1, 1. / self.maximum_curvature))
            samples = []
            for sec in seg:
                s0, s1 = sec[0].state, sec[1].state
                samples.extend(reeds_shepp.path_sample(s0, s1, 1. / self.maximum_curvature, res))

            acc = min([(v_max ** 2 - v1 ** 2) / extent, a_cc])
            vcc = np.sqrt(v1 ** 2 + acc * extent)
            for i, sample in enumerate(samples):
                if i * res < extent / 2.:
                    vt = min([np.sqrt(v0 ** 2 + 2 * acc * (i * res)), vcc])
                else:
                    vt = min([np.sqrt(v1 ** 2 + 2 * acc * np.abs(extent - i * res)), vcc])
                motion.append(self.Configuration(sample[:3], k=sample[3], v=np.sign(sample[4]) * vt))
            return motion

        segments = []  # type: List[(float, float)]
        path = [tuple(node.state) for node in self.path()]
        reduce(extract_segments, path)
        segments = list(zip(segments[:-1], segments[1:]))  # type: List[(Tuple[float], Tuple[float])]

        discontinuities2 = []  # type: List[(Tuple[float], float)]
        segments.insert(0, tuple(self.root.state))
        reduce(extract_discontinuities, segments)

        discontinuities2.append(self.Configuration().from_state_node(self.gain))
        discontinuities2.insert(0, self.Configuration().from_state_node(self.root))
        discontinuity_indexes = []
        for j, dis in enumerate(discontinuities2):
            if dis.v == 0:
                discontinuity_indexes.append(j)

        motions = []
        sector_indexes = list(zip(discontinuity_indexes[:-1], discontinuity_indexes[1:]))
        for sector_index in sector_indexes:
            motions.extend(plan_motions(sector_index))
        motions.append(self.Configuration(state=self.gain.state, v=self.gain.v, k=motions[-1].k))
        return motions

    class Configuration(object):
        def __init__(self, state=(), v=None, k=None):
            self.state = np.array(state)
            self.v, self.k = v, k

        def from_state_node(self, state_node):
            # type: (RRTStar.StateNode) -> RRTStar.Configuration
            self.state = state_node.state
            self.v, self.k = state_node.v, state_node.k
            return self

    class StateNode(object):
        def __init__(self, state=()):
            # type: (tuple) -> None
            self.state = np.array(state)  # state of the Node, a tuple (x, y, orientation)
            self.g = np.inf  # cost from root to here.
            self.hu = np.inf  # cost from here to goal if available.
            self.hl = np.inf  # cost from here to goal if not available.
            self.fu = self.g + self.hu
            self.fl = self.g + self.hl
            self.parent = None  # type: Optional[RRTStar.StateNode]
            self.children = []  # type: List[RRTStar.StateNode]
            self.status = 0  # 0 for safe, 1 for dangerous.
            self.v, self.k = 0, None  # velocity of the state, curvature of the state (related to the steering angle)

        def match(self, x_parent):
            # type: (RRTStar.StateNode) -> None
            """
            add a state as parent.
            """
            self.parent = x_parent
            x_parent.children.append(self)

        def remove(self):
            if self.parent:
                self.parent.children.remove(self)
                self.parent = None

        def rematch(self, x_new_parent):
            if self.parent:
                self.parent.children.remove(self)
                self.match(x_new_parent)

        def trace(self):  # type: ()->List[RRTStar.StateNode]
            p, ptr = [self], self.parent
            while ptr:
                p.append(ptr)
                ptr = ptr.parent
            p.reverse()
            return p

        def lcs2gcs(self, origin):
            # type: (RRTStar.StateNode) -> RRTStar.StateNode
            """
            transform self's coordinate from local coordinate system (LCS) to global coordinate system (GCS)
            :param origin: the tuple the coordinate (in GCS) of the origin of LCS.
            """
            xo, yo, ao = origin[0], origin[1], origin[2]
            x = self.state[0] * np.cos(ao) - self.state[1] * np.sin(ao) + xo
            y = self.state[0] * np.sin(ao) + self.state[1] * np.cos(ao) + yo
            a = self.state[2] + ao
            self.state = np.array((x, y, a))
            return self

        def gcs2lcs(self, origin):
            # type: (RRTStar.StateNode) -> RRTStar.StateNode
            """
            transform self's coordinate from global coordinate system (LCS) to local coordinate system (GCS)
            :param origin: the circle-node contains the coordinate (in GCS) of the origin of LCS.
            """
            xo, yo, ao = origin[0], origin[1], origin[2]
            x = (self.state[0] - xo) * np.cos(ao) + (self.state[1] - yo) * np.sin(ao)
            y = -(self.state[0] - xo) * np.sin(ao) + (self.state[1] - yo) * np.cos(ao)
            a = self.state[2] - ao
            self.state = np.array((x, y, a))
            return self

class BiRRTStar(RRTStar):
    def __init__(self):
        super(BiRRTStar, self).__init__()
        self.s_vertices = None
        self.g_vertices = None

    def preset(self, start, goal, grid_map, grid_res, grid_ori, obstacle, heuristic):
        self.grid_map, self.grid_res, self.grid_ori, self.obstacle = grid_map, grid_res, grid_ori, obstacle
        self.heuristic = heuristic
        self.start, self.goal = start, goal

        self.start.g, self.start.hl = 0, self.cost(start, goal)
        self.start.hu = self.start.hl if self.collision_free(start, goal) else np.inf
        self.start.fl, self.start.fu = self.start.g + self.start.hl, self.start.g + self.start.hu

        self.goal.g, self.goal.hl = 0, self.cost(goal, start)
        self.goal.hu = self.goal.hl if self.collision_free(goal, start) else np.inf
        self.goal.fl, self.goal.fu = self.goal.g + self.goal.hl, self.goal.g + self.goal.hu

        self.root, self.gain = self.start, self.goal
        self.s_vertices = [self.start]
        self.g_vertices = [self.goal]
        self.x_best = self.root
        return self

    def swap(self, i):
        self.branch_and_bound(self.g_vertices)
        self.branch_and_bound(self.s_vertices)
        if self.root is self.start:
            self.root = self.goal
            self.gain = self.start
            self.vertices = self.g_vertices
            n = -((i/2) % len(self.heuristic)) - 1 if self.heuristic else i
            Debugger.breaker('swap: goal -> start, {}, {}'.format(i, n), self.debug)
            return n
        else:
            self.root = self.start
            self.gain = self.goal
            self.vertices = self.s_vertices
            n = (i/2) % len(self.heuristic) if self.heuristic else i
            Debugger.breaker('swap: start -> goal, {}, {}'.format(i, n), self.debug)
            return n

    def planning(self, times, repeat=10, optimize=True, debug=False):
        """main flow."""
        self.debug = debug
        past = time.time()
        for i in range(times):
            n = self.swap(i)
            x_new = self.sample_free(n, repeat)
            x_nearest = self.nearest(x_new) if not optimize else self.least(x_new)
            if x_nearest and self.benefit(x_new) and self.collision_free(x_nearest, x_new):
                self.attach(x_nearest, x_new)
                self.rewire(x_new)
                self.x_best = self.best_of_all()
                self.connect_graphs(x_new)
            Debugger().debug_planned_path(self, i, switch=self.debug)
            Debugger().debug_planning_hist(self, i, (time.time() - past), switch=True)

    def best_of_all(self):
        x_new_best = self.best()
        return x_new_best if x_new_best.fu < self.x_best.fu else self.x_best

    def connect_graphs(self, x_new):
        if self.root is self.start:
            vs = self.g_vertices
        else:
            vs = self.s_vertices
        costs = map(lambda x: self.cost(x_new, x), vs)
        x_nearest, cost = vs[int(np.argmin(costs))], np.min(costs)
        Debugger().debug_connect_graphs(x_nearest.state, x_new.g+cost+x_nearest.g, self.x_best.fu, switch=self.debug)
        if x_new.g + cost + x_nearest.g < self.x_best.fu:
            if self.collision_free(x_new, x_nearest):
                x_nearest.fu = x_new.fu = x_new.g + cost + x_nearest.g
                x_new.hu = x_new.fu - x_new.g
                x_nearest.hu = x_nearest.fu - x_nearest.g
                x_new.neighbor = x_nearest
                x_nearest.neighbor = x_new
                self.x_best = x_new

    def path(self):
        x_best = self.x_best
        if x_best.fu < np.inf:
            p = x_best.trace()
            if hasattr(x_best, 'neighbor') and x_best.neighbor:
                p1 = x_best.neighbor.trace()
                p1.reverse()
                p.extend(p1)
                if p[-1] is self.start:
                    p.reverse()
                return p
            else:
                if p[0] is self.start:
                    p.append(self.gain)
                else:
                    p.append(self.start)
                    p.reverse()
                return p
        else:
            x_spare = self.spare(space=self.s_vertices)
            return x_spare.trace()


r""":mod:`mirgecom.inviscid` provides helper functions for inviscid flow.

Inviscid Flux Calculation
^^^^^^^^^^^^^^^^^^^^^^^^^

.. autofunction:: inviscid_flux
.. autofunction:: inviscid_facial_flux

Inviscid Time Step Computation
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. autofunction:: get_inviscid_timestep
.. autofunction:: get_inviscid_cfl
"""

__copyright__ = """
Copyright (C) 2020 University of Illinois Board of Trustees
"""

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import numpy as np
from meshmode.dof_array import thaw
from grudge.trace_pair import TracePair
from mirgecom.flux import divergence_flux_lfr
from mirgecom.fluid import make_conserved

def inviscid_flux(state):
    r"""Compute the inviscid flux vectors from fluid conserved vars *cv*.

    The inviscid fluxes are
    $(\rho\vec{V},(\rho{E}+p)\vec{V},\rho(\vec{V}\otimes\vec{V})
    +p\mathbf{I}, \rho{Y_s}\vec{V})$

    .. note::

        The fluxes are returned as a :class:`mirgecom.fluid.ConservedVars`
        object with a *dim-vector* for each conservation equation. See
        :class:`mirgecom.fluid.ConservedVars` for more information about
        how the fluxes are represented.

    Parameters
    ----------
    state: :class:`~mirgecom.gas_model.FluidState`

        Full fluid conserved and thermal state.

    Returns
    -------
    :class:`~mirgecom.fluid.ConservedVars`

        A CV object containing the inviscid flux vector for each
        conservation equation.
    """
    mass_flux = state.momentum_density
    energy_flux = state.velocity * (state.energy_density + state.pressure)
    mom_flux = (
        state.mass_density * np.outer(state.velocity, state.velocity)
        + np.eye(state.dim)*state.pressure
    )
    species_mass_flux = \
        state.velocity*state.species_mass_density.reshape(-1, 1)

    return make_conserved(state.dim, mass=mass_flux, energy=energy_flux,
                          momentum=mom_flux, species_mass=species_mass_flux)

def inviscid_facial_flux(discr, state_tpair, local=False):
    r"""Return the flux across a face given the solution on both sides *q_tpair*.

    This flux is currently hard-coded to use a Rusanov-type  local Lax-Friedrichs
    (LFR) numerical flux at element boundaries. The numerical inviscid flux $F^*$ is
    calculated as:

    .. math::

        \mathbf{F}^{*}_{\mathtt{LFR}} = \frac{1}{2}(\mathbf{F}(q^-)
        +\mathbf{F}(q^+)) \cdot \hat{n} + \frac{\lambda}{2}(q^{-} - q^{+}),

    where $q^-, q^+$ are the fluid solution state on the interior and the
    exterior of the face on which the LFR flux is to be calculated, $\mathbf{F}$ is
    the inviscid fluid flux, $\hat{n}$ is the face normal, and $\lambda$ is the
    *local* maximum fluid wavespeed.

    Parameters
    ----------
    discr: :class:`~grudge.eager.EagerDGDiscretization`

        The discretization collection to use

    state_tpair: :class:`~grudge.trace_pair.TracePair`

        Trace pair of :class:`~mirgecom.gas_model.FluidState` for the face upon
        which the flux calculation is to be performed

    local: bool

        Indicates whether to skip projection of fluxes to "all_faces" or not. If
        set to *False* (the default), the returned fluxes are projected to
        "all_faces."  If set to *True*, the returned fluxes are not projected to
        "all_faces"; remaining instead on the boundary restriction.

    Returns
    -------
    :class:`~mirgecom.fluid.ConservedVars`

        A CV object containing the scalar numerical fluxes at the input faces.
        The returned fluxes are scalar because they've already been dotted with
        the face normals as required by the divergence operator for which they
        are being computed.
    """
    actx = state_tpair.int.array_context
    dd = state_tpair.dd
    dd_all_faces = dd.with_dtag("all_faces")

    flux_tpair = TracePair(dd,
                           interior=inviscid_flux(state_tpair.int),
                           exterior=inviscid_flux(state_tpair.ext))

    # This calculates the local maximum eigenvalue of the flux Jacobian
    # for a single component gas, i.e. the element-local max wavespeed |v| + c.
    w_int = state_tpair.int.speed_of_sound + state_tpair.int.speed
    w_ext = state_tpair.ext.speed_of_sound + state_tpair.ext.speed
    lam = actx.np.maximum(w_int, w_ext)

    normal = thaw(actx, discr.normal(dd))
    cv_tpair = TracePair(dd,
                         interior=state_tpair.int.cv,
                         exterior=state_tpair.ext.cv)

    # todo: user-supplied flux routine
    flux_weak = divergence_flux_lfr(cv_tpair, flux_tpair, normal=normal, lam=lam)

    if local is False:
        return discr.project(dd, dd_all_faces, flux_weak)

    return flux_weak

def get_inviscid_timestep(discr, state):
    """Return node-local stable timestep estimate for an inviscid fluid.

    The maximum stable timestep is computed from the acoustic wavespeed.

    Parameters
    ----------
    discr: grudge.eager.EagerDGDiscretization

        the discretization to use

    state: :class:`~mirgecom.gas_model.FluidState`

        Full fluid conserved and thermal state

    Returns
    -------
    class:`~meshmode.dof_array.DOFArray`

        The maximum stable timestep at each node.
    """
    from grudge.dt_utils import characteristic_lengthscales
    return (
        characteristic_lengthscales(state.array_context, discr)
        / state.wavespeed
    )

def get_inviscid_cfl(discr, state, dt):
    """Return node-local CFL based on current state and timestep.

    Parameters
    ----------
    discr: :class:`~grudge.eager.EagerDGDiscretization`

        the discretization to use

    dt: float or :class:`~meshmode.dof_array.DOFArray`

        A constant scalar dt or node-local dt

    state: :class:`~mirgecom.gas_model.FluidState`

        The full fluid conserved and thermal state

    Returns
    -------
    :class:`~meshmode.dof_array.DOFArray`

        The CFL at each node.
    """
    return dt / get_inviscid_timestep(discr, state=state)


# -*- coding: utf-8 -*-
"""
Created on Thu Nov  1 14:27:54 2018

@author: dkorff
"""

import numpy as np
import cantera as ct

from li_ion_battery_p2d_inputs import Inputs

# Import Cantera objects:
anode_obj = ct.Solution(Inputs.ctifile,Inputs.anode_phase)
elyte_obj = ct.Solution(Inputs.ctifile,Inputs.elyte_phase)
cathode_obj = ct.Solution(Inputs.ctifile,Inputs.cathode_phase)
conductor_obj = ct.Solution(Inputs.ctifile,Inputs.metal_phase)

anode_surf_obj = ct.Interface(Inputs.ctifile,Inputs.anode_surf_phase,
    [anode_obj,elyte_obj,conductor_obj])
cathode_surf_obj = ct.Interface(Inputs.ctifile,Inputs.cathode_surf_phase,
    [cathode_obj,elyte_obj,conductor_obj])

# Anode initial conditions:
LiC6_0 = Inputs.SOC_0; C6_0 = 1 - LiC6_0
X_an_init = '{}:{}, {}:{}'.format(Inputs.Li_species_anode, LiC6_0,
    Inputs.Vac_species_anode, C6_0)

# Set Cantera object states
anode_obj.X = X_an_init
anode_obj.TP = Inputs.T, ct.one_atm
elyte_obj.TP = Inputs.T, ct.one_atm
anode_surf_obj.TP = Inputs.T, ct.one_atm

X_an_0 = anode_obj.X

if hasattr(Inputs, 'X_elyte_init'):
    elyte_obj.X = X_elyte_init

# Initial conditions: UPDATE TO GENERALIZE FOR MULT. SPECIES - DK 8/31/18
LiCoO2_0 = 1- Inputs.SOC_0; CoO2_0 = 1 - LiCoO2_0
X_ca_init = '{}:{}, {}:{}'.format(Inputs.Li_species_cathode, LiCoO2_0,
        Inputs.Vac_species_cathode, CoO2_0)

cathode_obj.X = X_ca_init
cathode_obj.TP = Inputs.T, ct.one_atm
cathode_surf_obj.TP = Inputs.T, ct.one_atm

X_ca_0 = cathode_obj.X
X_elyte_0 = elyte_obj.X

"""========================================================================="""
"""========================================================================="""
"""========================================================================="""

class anode():

    # Set flag so solver knows whether to implement the anode:
    flag = Inputs.flag_anode

    # Number of nodes in the y-direction:
    npoints = Inputs.npoints_anode

    # Number of "shells" in anode particle:
    nshells = Inputs.nshells_anode

    # Initial conditions: UPDATE TO GENERALIZE FOR MULT. SPEC - DK 8/31/18

    # Anode variables for a given volume include X_Li for each shell, Phi_an,
    #       Phi_elyte, and rho_k_elyte for all elyte species.
    nVars = nshells + 2 + elyte_obj.n_species

    # Pointers
    ptr = {}
    ptr['iFar'] = 1
    ptr['X_ed'] = np.arange(0, Inputs.nshells_anode)
    ptr['rho_k_elyte'] = nshells + np.arange(0,elyte_obj.n_species)
    ptr['Phi_ed'] = nshells + elyte_obj.n_species
    ptr['Phi_dl'] = nshells + elyte_obj.n_species + 1

    # Anode/elyte interface area per unit volume
    #   [m^2 interface / m_3 total electrode volume]
    # For spherical particles, the total surface area per unit volume can be
    #   calculated from the geometry. Since some particles will overlap, we
    #   multiply by an 'overlap' input parameter.
    A_surf = (1-Inputs.overlap_an)*6*Inputs.eps_solid_an/Inputs.d_part_an

    # Set up solution vector
    nSV = npoints*nVars

    # Array of offsets to point to each node's variable:
    offsets = np.arange(0,int(nSV),int(nVars))

    # Store parameters as class object attributes:
    T = Inputs.T
    C_dl = Inputs.C_dl_an
    X_Li_max = Inputs.SOC_max
    X_Li_min = Inputs.SOC_min
    D_Li_ed = Inputs.D_Li_an

    # Geometric parameters:
    eps_ed = Inputs.eps_solid_an
    eps_elyte = 1 - Inputs.eps_solid_an
    tau_ed = Inputs.tau_an
    r_pore = Inputs.r_p_an
    d_part = Inputs.d_part_an
    dyInv = npoints/Inputs.H_an
    dr = d_part*0.5/nshells

    # Calculate the current density [A/m^2] corresponding to a C_rate of 1:
    oneC = eps_ed*anode_obj.density_mole*Inputs.H_an*ct.faraday/3600

    # Calculate the percent volume of a single graphite particle that exists in
    #   each 'shell'. I.e. for shell j, what is the volume of that shell,
    #   relative to the total particle volume? The radius of the volume is
    #   currently discretized evenly (i.e. 'dr' the differential radius is
    #   constant). Certainly other discretizations (such as constant
    #   differential volume) are possible (and perhaps better).
    #
    #   Because the volume is 4/3 pi*r^3, the volume of the shell relative to
    #   the total volume is (r_shell/r_particle)^3, and the differential volume
    #   relative to the total, for shell 'j' is:
    #       (r_shell(j+1)^3 - r_shell(j)^3)/r_particle^3
    #   Because the radius is discretized evenly, the radius of shell j, r_j,
    #   relative to the total radius r_particle, is:
    #       r_j/r_particle = j/nshells

    V_shell = np.zeros([nshells])
    nsr3 = 1/nshells/nshells/nshells
    for j in np.arange(0, nshells, 1):
        V_shell[j] = ((j+1)**3 - (j)**3)*nsr3

    # Electronic conductivity of the electrode phase:
    sigma_eff_ed = Inputs.sigma_an*eps_ed/tau_ed**3
    # Species mobilities of the electrolyte phase.  Converted from user input
    #   diffusion coefficients:
    u_Li_elyte = (Inputs.D_Li_elyte*eps_elyte/ct.gas_constant
          /Inputs.T/tau_ed**3)

    t_flag = []

    """========================================================================="""
    """========================================================================="""
    """========================================================================="""

class separator():
    # Set a flag to let the solver know whether to implement this class:
    flag = Inputs.flag_sep

    # Number of nodes in the y-direction
    npoints = Inputs.npoints_elyte

    # Number of variables per node:
    nVars = 1 + elyte_obj.n_species

    H = Inputs.H_elyte  # Separator thickness [m]

    tau_sep = Inputs.tau_sep  # Tortuosity of separator

    # Geometric parameters:
    eps_elyte = Inputs.eps_elyte_sep
    dyInv = npoints/H
    tau_sep = tau_sep

    ptr = {}
    ptr['rho_k_elyte'] = np.arange(0,elyte_obj.n_species)
    ptr['Phi'] = elyte_obj.n_species

    # Set up the solution vector
    nSV = npoints*nVars

    # Array of offsets to point to each node's variable:
    offsets = np.arange(int(anode.nSV),int(anode.nSV)+int(nSV),int(nVars))

"""========================================================================="""
"""========================================================================="""
"""========================================================================="""

class cathode():
    # Set flag so solver knows it's looking at the anode class
    flag = Inputs.flag_cathode

    # Number of nodes in the y-direction:
    npoints = Inputs.npoints_cathode

    # Number of "shells" in the cathode particle
    nshells = Inputs.n_shells_cathode

    # Number of state variables per node:
    nVars = nshells + 2 + elyte_obj.n_species

    # Pointers
    ptr = {}
    ptr['iFar'] = 0
    ptr['X_ed'] = np.arange(0, nshells)
    ptr['rho_k_elyte'] = nshells + np.arange(0,elyte_obj.n_species)
    ptr['Phi_ed'] = nshells + elyte_obj.n_species
    ptr['Phi_dl'] = nshells + elyte_obj.n_species + 1

    # Cathode/elyte interface area per unit volume
    #   [m^2 interface / m_3 total electrode volume]
    # For spherical particles, the total surface area per unit volume can be
    #   calculated from the geometry. Since some particles will overlap, we
    #   multiply by an 'overlap' input parameter.
    A_surf = (1-Inputs.overlap_ca)*6*Inputs.eps_solid_ca/Inputs.d_part_ca

    # Set up the solution vector
    nSV = npoints*nVars

    # Array of offsets to point to each node's variable:
    offsets = np.arange(int(anode.nSV + separator.nSV), \
                    int(anode.nSV + separator.nSV)+int(nSV),int(nVars))

    #  Store Parameters in the Class object:
    T = Inputs.T
    C_dl = Inputs.C_dl_ca
    X_Li_max = Inputs.SOC_max
    X_Li_min = Inputs.SOC_min
    D_Li_ed = Inputs.D_Li_ca

    # Geometric parameters:
    eps_ed = Inputs.eps_solid_ca
    eps_elyte = 1 - Inputs.eps_solid_ca
    tau_ed = Inputs.tau_ca
    r_p = Inputs.r_p_ca
    d_part = Inputs.d_part_ca
    dyInv = npoints/Inputs.H_ca
    dr = d_part*0.5/nshells

    # Calculate the current density [A/m^2] corresponding to a C_rate of 1:
    oneC = eps_ed*cathode_obj.density_mole*Inputs.H_ca*ct.faraday/3600

    # Calculate the percent volume of a single graphite particle that exists in
    #   each 'shell'. I.e. for shell j, what is the volume of that shell,
    #   relative to the total particle volume? The radius of the volume is
    #   currently discretized evenly (i.e. 'dr' the differential radius is
    #   constant). Certainly other discretizations (such as constant
    #   differential volume) are possible (and perhaps better).
    #
    #   Because the volume is 4/3 pi*r^3, the volume of the shell relative to
    #   the total volume is (r_shell/r_particle)^3, and the differential volume
    #   relative to the total, for shell 'j' is:
    #       (r_shell(j+1)^3 - r_shell(j)^3)/r_particle^3
    #   Because the radius is discretized evenly, the radius of shell j, r_j,
    #   relative to the total radius r_particle, is:
    #       r_j/r_particle = j/nshells

    V_shell = np.zeros([nshells])
    nsr3 = 1/nshells/nshells/nshells
    for j in np.arange(0, nshells, 1):
        V_shell[j] = ((j+1)**3 - (j)**3)*nsr3

    sigma_eff_ed = Inputs.sigma_ca*eps_ed/tau_ed**3
    u_Li_elyte = (Inputs.D_Li_elyte*eps_elyte/ct.gas_constant
          /Inputs.T/tau_ed**3)

    t_flag = []

"""========================================================================="""
"""========================================================================="""
"""========================================================================="""

# Calculate the actual current density:
# The minus sign is because we begin with the charging reaction, which
#   delivers negative charge to the anode:
i_ext = -Inputs.C_rate*min(anode.oneC,cathode.oneC)

SV_0 = np.zeros([anode.nSV+separator.nSV+cathode.nSV])

# Set up algebraic variable vector:
algvar = np.zeros_like(SV_0)

offsets = anode.offsets
ptr = anode.ptr
for j in range(anode.npoints):
    SV_0[offsets[j] + ptr['X_ed']] = \
        np.ones([anode.nshells])*X_an_0[0]
    algvar[offsets[j] + ptr['X_ed']] = 1

    SV_0[offsets[j] + ptr['rho_k_elyte']] = \
        elyte_obj.Y*elyte_obj.density_mass
    algvar[offsets[j] + ptr['rho_k_elyte']] = 1

    SV_0[offsets[j] + ptr['Phi_ed']] = \
        Inputs.Phi_anode_init

    SV_0[offsets[j] + ptr['Phi_dl']] = \
        Inputs.Phi_elyte_init - Inputs.Phi_anode_init
    algvar[offsets[j] + ptr['Phi_dl']] = 1

offsets = separator.offsets
ptr = separator.ptr
for j in np.arange(0, separator.npoints):
    SV_0[offsets[j] + ptr['rho_k_elyte']] = \
        elyte_obj.Y*elyte_obj.density_mass
    algvar[offsets[j] + ptr['rho_k_elyte']] = 1

    SV_0[offsets[j] + ptr['Phi']] = \
        Inputs.Phi_elyte_init

offsets = cathode.offsets
ptr = cathode.ptr
for j in range(cathode.npoints):
    SV_0[offsets[j] + ptr['X_ed']] = \
        np.ones([cathode.nshells])*X_ca_0[0]
    algvar[offsets[j] + ptr['X_ed']] = 1

    SV_0[offsets[j] + ptr['rho_k_elyte']] = \
        elyte_obj.Y*elyte_obj.density_mass
    algvar[offsets[j] + ptr['rho_k_elyte']] = 1

    SV_0[offsets[j] + ptr['Phi_ed']] = \
        Inputs.Phi_anode_init + Inputs.Delta_Phi_init

    SV_0[offsets[j] + ptr['Phi_dl']] = \
        Inputs.Phi_elyte_init - (Inputs.Phi_anode_init + Inputs.Delta_Phi_init)
    algvar[offsets[j] + ptr['Phi_dl']] = 1


import numpy as np
from argparse import ArgumentParser
from collections import defaultdict

import torch
import torch.nn as nn
import torch.nn.functional as F

from utils import CEVAE_mod, load_data

parser = ArgumentParser()
parser.add_argument('-lr', type=float, default=0.0005)
parser.add_argument('-hDim', type=int, default=50)
parser.add_argument('-nTest', type=int, default=1000)
parser.add_argument('-nTrain', type=int, default=1000)
parser.add_argument('-zDim', type=int, default=5)
parser.add_argument('-rep', type=int, default=20)
parser.add_argument('-nIter', type=int, default=10001)
parser.add_argument('-batchSize', type=int, default=512)
parser.add_argument('-nSamplesZ', type=int, default=1)
parser.add_argument('-evalIter', type=int, default=500)
parser.add_argument('-device', type=str, default='cpu')
parser.add_argument('-comment', type=str, default='')
args = parser.parse_args()

# Model class for auxiliary model
class aux_m(nn.Module):
    def __init__(self, input_dim=20, h_dim_f=100, out_dim=1):
        super().__init__()
        # init network
        self.lin1 = nn.Linear(input_dim, h_dim_f)
        self.lin2 = nn.Linear(h_dim_f, out_dim)

    def forward(self, z_in):
        h1 = F.elu(self.lin1(z_in))
        out = torch.sigmoid(self.lin2(h1))
        return out

if __name__ == "__main__":

    acc_list = defaultdict(list)
    stat_par_y1a0 = defaultdict(list)
    stat_par_y1a1 = defaultdict(list)

    # repeat whole training process to observe stability
    for rep_i in range(args.rep):
        print('rep: %i' % rep_i)
        # data loader:
        # order of packed data: data = [y, x_con, x_bin, r, b, a]
        train_data, test_data, cat_bin_dict = load_data(n_test=args.nTest)

        # unpack train x in order to init in right shapes
        train_data = [data[:args.nTrain] for data in train_data]
        y_tr, x_tr_con, x_tr_bin, r_tr, b_tr, a_tr = train_data
        x_tr = np.hstack((x_tr_bin, x_tr_con))

        CEVAE = CEVAE_mod(args=args, dim_x=x_tr.shape[1], dim_b=b_tr.shape[1], dim_a=1, dim_x_con=x_tr_con.shape[1],
                          dim_x_bin=cat_bin_dict, dim_z=args.zDim, dim_r=r_tr.shape[1], dim_q_h=args.hDim,
                          dim_p_h=args.hDim).to(args.device)

        model = './model_path.pt'
        CEVAE.load_state_dict(torch.load(model))

        # init Causal Path Enabler (auxiliary -fair- models)
        path_combinations = ['z', 'zb', 'zbr', 'zbrx', 'zbrxa']
        AUX = dict()
        AUX['z'] = aux_m(input_dim=args.zDim, h_dim_f=100)
        AUX['zb'] = aux_m(input_dim=args.zDim + b_tr.shape[1], h_dim_f=100)
        AUX['zbr'] = aux_m(input_dim=args.zDim + b_tr.shape[1] + r_tr.shape[1], h_dim_f=100)
        AUX['zbrx'] = aux_m(input_dim=args.zDim + b_tr.shape[1] + r_tr.shape[1] + x_tr.shape[1], h_dim_f=100)
        AUX['zbrxa'] = aux_m(input_dim=args.zDim + b_tr.shape[1] + r_tr.shape[1] + x_tr.shape[1] + a_tr.shape[1],
                             h_dim_f=100)

        # init optimizer
        optimizer = dict()
        for combination in path_combinations:
            optimizer[combination] = torch.optim.RMSprop(AUX[combination].parameters(), lr=args.lr)

        # Maintain loss development for monitoring
        loss_dict = defaultdict(list)

        # training loop
        for i in range(args.nIter):
            # select random batch
            batch_idx = np.random.choice(a=range(x_tr.shape[0]), size=args.batchSize, replace=False)
            batch_data = [torch.Tensor(g[batch_idx]).to(args.device) for g in train_data]

            y_batch, x_batch_con, x_batch_bin, r_batch, b_batch, a_batch = batch_data
            # INFER distribution over z, using inference network of CEVAE
            z_infer = CEVAE.q_z.forward(
                observations=torch.cat((x_batch_con, x_batch_bin, r_batch, b_batch, a_batch), 1))
            # No need to store derivative to CEVAE
            z_sample = z_infer.sample().detach()

            aux_output = dict()
            aux_output['z'] = AUX['z'].forward(z_sample)
            aux_output['zb'] = AUX['zb'].forward(torch.cat((z_sample, b_batch), 1))
            # use R(x,Z,B,A) below instead in order to exclude A->X->R->Y
            aux_output['zbr'] = AUX['zbr'].forward(torch.cat((z_sample, b_batch, r_batch), 1))
            aux_output['zbrx'] = AUX['zbrx'].forward(
                torch.cat((z_sample, b_batch, r_batch, x_batch_con, x_batch_bin), 1))
            aux_output['zbrxa'] = AUX['zbrxa'].forward(
                torch.cat((z_sample, b_batch, r_batch, x_batch_con, x_batch_bin, a_batch), 1))

            # calculate loss and update step
            objective = dict()
            for combination in path_combinations:
                objective[combination] = torch.mean((aux_output[combination] - y_batch) ** 2)
                optimizer[combination].zero_grad()
                objective[combination].backward()
                optimizer[combination].step()

                loss_dict[combination].append(float(objective[combination].cpu().detach().numpy()))

        # test on test set
        batch_data = [torch.Tensor(g).to(args.device) for g in test_data]

        y_batch, x_batch_con, x_batch_bin, r_batch, b_batch, a_batch = batch_data
        # INFER distribution over z, using inference network of CEVAE
        z_infer = CEVAE.q_z.forward(observations=torch.cat((x_batch_con, x_batch_bin, r_batch, b_batch, a_batch), 1))
        # No need to store derivative to CEVAE
        z_sample = z_infer.sample().detach()

        aux_output = dict()
        aux_output['z'] = AUX['z'].forward(z_sample)
        aux_output['zb'] = AUX['zb'].forward(torch.cat((z_sample, b_batch), 1))
        aux_output['zbr'] = AUX['zbr'].forward(torch.cat((z_sample, b_batch, r_batch), 1))
        aux_output['zbrx'] = AUX['zbrx'].forward(torch.cat((z_sample, b_batch, r_batch, x_batch_con, x_batch_bin), 1))
        aux_output['zbrxa'] = AUX['zbrxa'].forward(
            torch.cat((z_sample, b_batch, r_batch, x_batch_con, x_batch_bin, a_batch), 1))

        # ------ Test Statistical parity -----
        mask_a0 = (a_batch == 0).squeeze()
        mask_a1 = (a_batch == 1).squeeze()

        for combination in path_combinations:
            y_predict = torch.round(aux_output[combination])
            accuracy = torch.sum(y_predict == y_batch).cpu().detach().numpy() / y_batch.shape[0]
            acc_list[combination].append(accuracy)
            p_y1_a0 = y_predict[mask_a0].sum() / mask_a0.sum()
            p_y1_a1 = y_predict[mask_a1].sum() / mask_a1.sum()
            stat_par_y1a0[combination].append(p_y1_a0.cpu().detach().numpy())
            stat_par_y1a1[combination].append(p_y1_a1.cpu().detach().numpy())

    np.save('accuracy_dict.npy', acc_list)
    np.save('stat_par_y1a0_dict.npy', stat_par_y1a0)
    np.save('stat_par_y1a1_dict.npy', stat_par_y1a1)