import numpy as np

def place_atom_nerf(a, b, c, bond_length, bond_angle_deg, torsion_angle_deg):
    """
    Calculates the coordinates of atom D given atoms A, B, C and internal coords.
    
    Parameters:
    - a, b, c: (3,) numpy arrays of xyz coordinates of the previous atoms.
    - bond_length: distance between c and d (Angstroms).
    - bond_angle_deg: angle b-c-d (Degrees).
    - torsion_angle_deg: dihedral angle a-b-c-d (Degrees).
    
    Returns:
    - d: (3,) numpy array of coordinates for the new atom.
    """
    # 1. Convert to rad 
    bond_angle_rad = np.radians(bond_angle_deg)
    torsion_angle_rad = np.radians(torsion_angle_deg)
    
    # 2. Calculate the position of D in the local reference frame
    # We align the local X-axis with the bond B->C.
    # Note: The geometric calculation uses the complement of the bond angle 
    #  so we use pi - bond_angle.
    
    # D_local represents the vector C->D in the local frame
    d_local = np.array([
        bond_length * np.cos(np.pi - bond_angle_rad),
        bond_length * np.sin(np.pi - bond_angle_rad) * np.cos(torsion_angle_rad),
        bond_length * np.sin(np.pi - bond_angle_rad) * np.sin(torsion_angle_rad)
    ])
    
    # 3. Build the transformation matrix from Local to Global Frame
    
    # Vector from B to C (normalized) - this is our local X-axis
    bc = c - b
    bc_u = bc / np.linalg.norm(bc)
    
    # Vector from A to B
    ab = b - a
    
    # Normal to the plane A-B-C (normalized) - this is our local Z-axis
    # We use Cross Product to find the perpendicular vector
    n = np.cross(ab, bc_u)
    n_u = n / np.linalg.norm(n)
    
    # The "Up" vector in the plane (normalized) - this is our local Y-axis
    # It is perpendicular to both the bond BC and the normal N
    m_u = np.cross(n_u, bc_u)
    
    # Create the Rotation Matrix columns [x_axis, y_axis, z_axis]
    # Shape: (3, 3)
    rotation_matrix = np.column_stack((bc_u, m_u, n_u))
    
    # 4. Transform D_local to Global Coordinates
    # Global_Pos = Origin(C) + (Rotation * Local_Pos)
    d_global = c + np.dot(rotation_matrix, d_local)
    
    return d_global

# --- Simple Test Case ---
if __name__ == "__main__":
    # Define 3 arbitrary starting atoms (A, B, C)
    atom_a = np.array([0.0, 1.0, 0.0])
    atom_b = np.array([0.0, 0.0, 0.0]) # Origin
    atom_c = np.array([1.5, 0.0, 0.0]) # Along X-axis
    
    # Define internal coordinates for the next atom (D)
    length = 1.5       # Angstroms
    b_angle = 120.0    # Degrees
    torsion = 90.0     # Degrees (Should point "out" of the screen/plane)

    new_atom = place_atom_nerf(atom_a, atom_b, atom_c, length, b_angle, torsion)
    
    print(f"Atom A: {atom_a}")
    print(f"Atom B: {atom_b}")
    print(f"Atom C: {atom_c}")
    print(f"Atom D (Calculated): {np.round(new_atom, 3)}")
