from pymatgen.core.structure import IStructure
from pymatgen.core.structure import Lattice
import numpy as np

file_name = '/home/sure/Document/vasptoolkit/examples/Lattice_Rotate_for_pDOS/POSCAR'
struct = IStructure.from_file(file_name)
sites = struct._sites
Mn_list = []
for n,site in enumerate(sites):
    if site.specie.symbol == 'Mn':
        Mn_list.append(n)

manual_pick_list = [1]
origin_Mn_list = Mn_list
Mn_list = []
for i in manual_pick_list:
    Mn_list.append(origin_Mn_list[i-1])

Mn_z_axis_list = []
Mn_x_axis_list = []
Mn_y_axis_list = []
nbr = struct.get_all_neighbors(2.5)
for pick_atom_index, Mn_site in enumerate(Mn_list):
    Mn_coords = np.array(sites[Mn_site].coords)
    nbr_vectors = np.array([np.array(nbr_site.coords)-Mn_coords for nbr_site in nbr[Mn_site]])  

    # following cell is to part the vectors in couple that are opposite to each other
    # dot product of neighbor vectors to identify the aligned ones
    nbr_vector_product = np.dot(nbr_vectors,nbr_vectors.T)
    # couple index print the smallest dot product of vectors
    couple_index = np.argmin(nbr_vector_product,axis=0)
    couple_list = []
    for a,b in zip(range(0,6),couple_index):
        if a < b:
            couple_list.append((a,b))

    # create axis_list
    axis_list = []
    for couple in couple_list:
        axis_list.append((nbr_vectors[couple[0]]-nbr_vectors[couple[1]])/2)
    axis_vectors = np.array(axis_list)
    axis_length = np.linalg.norm(axis_vectors, axis=1)
    


    # set axis for the first atom only, and deal with others according to the first one
    if Mn_z_axis_list:
        axis_vectors = np.append(axis_vectors,-axis_vectors,axis=0)
        z_product = np.dot(axis_vectors, Mn_z_axis_list[0].T)
        z_axis = axis_vectors[np.argmax(z_product)]
        x_product = np.dot(axis_vectors, Mn_x_axis_list[0].T)
        x_axis = axis_vectors[np.argmax(x_product)]
        y_product = np.dot(axis_vectors, Mn_y_axis_list[0].T)
        y_axis = axis_vectors[np.argmax(y_product)]


    # if Mn_z_axis_list is empty we initialize the first axis
    else:
        # find z axis
        z_axis_index = np.argmax(axis_length)
        z_axis = axis_vectors[z_axis_index]

        # set x y axis according to right hand rule
        xy_axis_list = np.delete(axis_vectors,z_axis_index,axis=0)
        x_axis = xy_axis_list[0]
        y_axis = xy_axis_list[1]

        # check the triple product
        triple_product = np.dot(np.cross(x_axis,y_axis),(z_axis).T)
        if triple_product < 0:
            x_axis, y_axis = y_axis, x_axis


    print('Mn{}_x_axis'.format(manual_pick_list[pick_atom_index]),x_axis,np.linalg.norm(x_axis))
    print('Mn{}_y_axis'.format(manual_pick_list[pick_atom_index]),y_axis,np.linalg.norm(y_axis))
    print('Mn{}_z_axis'.format(manual_pick_list[pick_atom_index]),z_axis,np.linalg.norm(z_axis))

    Mn_x_axis_list.append(x_axis)
    Mn_y_axis_list.append(y_axis)
    Mn_z_axis_list.append(z_axis)

    
Mn_z_axis_list = np.array(Mn_z_axis_list)
Mn_z_axis = np.mean(Mn_z_axis_list,axis=0)

Mn_x_axis_list = np.array(Mn_x_axis_list)
Mn_x_axis = np.mean(Mn_x_axis_list,axis=0)

Mn_y_axis_list = np.array(Mn_y_axis_list)
Mn_y_axis = np.mean(Mn_y_axis_list,axis=0)


print('Averaged Mn_x_axis',Mn_x_axis,np.linalg.norm(Mn_x_axis))
print('Averaged Mn_y_axis',Mn_y_axis,np.linalg.norm(Mn_y_axis))
print('Averaged Mn_z_axis',Mn_z_axis,np.linalg.norm(Mn_z_axis))


# Schmidt orthogonalization
orth_z_axis = Mn_z_axis
orth_x_axis = Mn_x_axis - np.dot(Mn_x_axis,orth_z_axis)/np.dot(orth_z_axis,orth_z_axis)*orth_z_axis
orth_y_axis = Mn_y_axis - np.dot(Mn_y_axis,orth_z_axis)/np.dot(orth_z_axis,orth_z_axis)*orth_z_axis\
     - np.dot(Mn_y_axis,orth_x_axis)/np.dot(orth_x_axis,orth_x_axis)*orth_x_axis

orth_z_axis = orth_z_axis/np.linalg.norm(orth_z_axis)
orth_x_axis = orth_x_axis/np.linalg.norm(orth_x_axis)
orth_y_axis = orth_y_axis/np.linalg.norm(orth_y_axis)

Mn_orth_xyz_coord = np.matrix([orth_x_axis,orth_y_axis,orth_z_axis])


print('Mn_orth_xyz_coord')
print(Mn_orth_xyz_coord)


print('original lattice')
print(struct.lattice.matrix)
target_lattice = np.dot(struct.lattice.matrix,Mn_orth_xyz_coord.T)

# orthogonalization check
#print(np.dot(target_lattice[0],target_lattice[1].T))
#print(np.dot(target_lattice[0],target_lattice[2].T))
print('lattice after rotation')
print(Lattice(target_lattice))