import torch
import math


def compute_particle_xyz_from_internal(particle_1_xyz,
                                       particle_2_xyz,
                                       particle_3_xyz,
                                       bond, angle, dihedral):
    """ Compute the Cartesian coordinate of a particle based on its internal
    coordinate with respect to three other particles and the Cartesian 
    coordinates of these three particles. This function runs a batch mode, which
    mean that it compute the Cartensian coordinates of a batch of particles based
    on a batch of three other particles.

    Parameters:
    -----------
    particle_1_xyz: torch.Tensor 
        the Cartesian coordinate of atom 1.
    particle_2_xyz: torch.Tensor 
        the Cartesian coordinate of atom 2.
    particle_3_xyz: torch.Tensor 
        the Cartesian coordinate of atom 3.
        Note particle_1_xyz, particle_2_xyz and particle_3_xyz should the same size: [batch_size, 3].
    bond: torch.Tensor
        the length of the bond between atom 3 and atom 4
    angle: torch.Tensor
        the value of the angle between atom 2, atom 3 and atom 4
    dihedral: torch.Tensor
        the value of the dihedral angle between atom 1, atom 2, atom3 and atom 4
        Note bond, angle and dihedral have the same size: [batch_size]

    Returns:
    --------
    particle_4_xyz: torch.Tensor
        the Cartensian coordiante of atom 4. Its size is [batch_size, 3]
    logabsdet: the logarithm of the absolute value of the determinant of
        the transformation from the internal coordinates of particle 4 
        to its Cartesian coordinates
    """

    ## calculate the coordinate of the forth atom
    wi, wj, wk = particle_1_xyz, particle_2_xyz, particle_3_xyz

    e1 = wk - wj
    e1 = e1 / torch.norm(e1, dim = -1, keepdim = True)

    e3 = torch.cross(e1, wj - wi, dim = -1)
    e3 = e3 / torch.norm(e3, dim = -1, keepdim = True)

    e2 = torch.cross(e3, e1, dim = -1)

    bond = torch.unsqueeze(bond, -1)
    angle = torch.unsqueeze(angle, -1)
    dihedral = torch.unsqueeze(dihedral, -1)
    
    dw = torch.cos(math.pi-angle)*e1 + \
         (-1)*torch.sin(math.pi-angle)*torch.cos(dihedral)*e2 + \
         (-1)*torch.sin(math.pi-angle)*torch.sin(dihedral)*e3
    
    wl = wk + bond*dw

    particle_4_xyz = wl

    ## calculate the Jacobian of the transform
    logabsdet = torch.log(torch.abs(bond**2*torch.sin(math.pi - angle)))
    logabsdet = torch.squeeze(logabsdet)
    
    return particle_4_xyz, logabsdet