import numpy as np
import sys
import dynaphopy.displacements as disp

def progress_bar(progress):
    bar_length = 30
    status = ""
    if isinstance(progress, int):
        progress = float(progress)
    if not isinstance(progress, float):
        progress = 0
        status = "Progress error\r\n"
    if progress < 0:
        progress = 0
        status = "Halt ...\r\n"
    if progress >= 1:
        progress = 1
        status = "Done...\r\n"
    block = int(round(bar_length*progress))
    text = "\rTrajectory: [{0}] {1:.2f}% {2}".format("#"*block + "-"*(bar_length-block),
                                                      progress*100, status)
    sys.stdout.write(text)
    sys.stdout.flush()

#print(disp.relative_trajectory(cell, traj ,pos))

def relativize_trajectory(dynamic):

    cell = dynamic.structure.get_cell()
    number_of_atoms = dynamic.trajectory.shape[1]
    super_cell = dynamic.get_super_cell_matrix()
    position = dynamic.structure.get_positions(super_cell=super_cell)
    normalized_trajectory = np.zeros_like(dynamic.trajectory.real)
    trajectory = dynamic.trajectory.real

    progress_bar(0)



    for j in range(number_of_atoms):

    #    print(trajectory, position, cell)
    #    print (disp.relative_trajectory(cell, trajectory[:,j,:], position[j]))
        for i in range(0, normalized_trajectory.shape[0]):
            difference = trajectory[i, j, :] - position[j]
            difference_matrix = np.dot(np.linalg.inv(cell), difference)
#            print(difference_matrix)

        normalized_trajectory[:, j, :] = disp.relative_trajectory(cell, trajectory[:, j, :], position[j])

        progress_bar(float(j+1)/number_of_atoms)

  #  print(normalized_trajectory)

    return normalized_trajectory

def relativize_trajectory_py(dynamic):

    cell = dynamic.structure.get_cell()
    number_of_atoms = dynamic.trajectory.shape[1]
    super_cell = dynamic.get_super_cell_matrix()
    position = dynamic.structure.get_positions(super_cell=super_cell)
    normalized_trajectory = dynamic.trajectory.real.copy()

    progress_bar(0)

    for j in range(number_of_atoms):
        for i in range(0, normalized_trajectory.shape[0]):

            difference = normalized_trajectory[i, j, :] - position[j]

         #   difference_matrix = np.array(np.dot(np.linalg.inv(cell),(IniSep)),dtype=int)
            difference_matrix = np.around(np.dot(np.linalg.inv(cell), difference), decimals=0)

            normalized_trajectory[i, j, :] -= np.dot(difference_matrix, cell) + position[j]

        progress_bar(float(j+1)/number_of_atoms)

  #  print(normalized_trajectory)
    return normalized_trajectory

def trajectory_projection(dynamic, direction):

    direction = np.array(direction)/np.linalg.norm(direction)

    super_cell = dynamic.get_super_cell_matrix()
    trajectory = dynamic.get_relative_trajectory()

    atom_type_index = dynamic.structure.get_atom_type_index(super_cell=super_cell)
    number_of_atom_types = dynamic.structure.get_number_of_atom_types()

    projections = []

    for j in range(number_of_atom_types):
        projection = np.array([])
        for i in range(0, trajectory.shape[1]):
            if atom_type_index[i] == j:
           #     print('atom:', i, 'type:', atom_type_index[i])
                projection = np.append(projection, np.dot(trajectory[:, i, :], direction/np.dot(direction, direction)),axis=1)
        projections.append(projection)

    return np.array(projections)

import numpy as np
import sys
import dynaphopy.displacements as disp

def progress_bar(progress):
    bar_length = 30
    status = ""
    if isinstance(progress, int):
        progress = float(progress)
    if not isinstance(progress, float):
        progress = 0
        status = "Progress error\r\n"
    if progress < 0:
        progress = 0
        status = "Halt ...\r\n"
    if progress >= 1:
        progress = 1
        status = "Done...\r\n"
    block = int(round(bar_length*progress))
    text = "\rTrajectory: [{0}] {1:.2f}% {2}".format("#"*block + "-"*(bar_length-block),
                                                      progress*100, status)
    sys.stdout.write(text)
    sys.stdout.flush()

#print(disp.relative_trajectory(cell, traj ,pos))

def relativize_trajectory(dynamic):

    cell = dynamic.structure.get_cell()
    number_of_atoms = dynamic.trajectory.shape[1]
    super_cell = dynamic.get_super_cell_matrix()
    position = dynamic.structure.get_positions(super_cell=super_cell)
    normalized_trajectory = np.zeros_like(dynamic.trajectory.real)
    trajectory = dynamic.trajectory.real

    progress_bar(0)



    for j in range(number_of_atoms):

    #    print(trajectory, position, cell)
    #    print (disp.relative_trajectory(cell, trajectory[:,j,:], position[j]))
        for i in range(0, normalized_trajectory.shape[0]):
            difference = trajectory[i, j, :] - position[j]
            difference_matrix = np.dot(np.linalg.inv(cell), difference)
#            print(difference_matrix)

        normalized_trajectory[:, j, :] = disp.relative_trajectory(cell, trajectory[:, j, :], position[j])

        progress_bar(float(j+1)/number_of_atoms)

  #  print(normalized_trajectory)

    return normalized_trajectory

def relativize_trajectory_py(dynamic):

    cell = dynamic.structure.get_cell()
    number_of_atoms = dynamic.trajectory.shape[1]
    super_cell = dynamic.get_super_cell_matrix()
    position = dynamic.structure.get_positions(super_cell=super_cell)
    normalized_trajectory = dynamic.trajectory.real.copy()

    progress_bar(0)

    for j in range(number_of_atoms):
        for i in range(0, normalized_trajectory.shape[0]):

            difference = normalized_trajectory[i, j, :] - position[j]

         #   difference_matrix = np.array(np.dot(np.linalg.inv(cell),(IniSep)),dtype=int)
            difference_matrix = np.around(np.dot(np.linalg.inv(cell), difference), decimals=0)

            normalized_trajectory[i, j, :] -= np.dot(difference_matrix, cell) + position[j]

        progress_bar(float(j+1)/number_of_atoms)

  #  print(normalized_trajectory)
    return normalized_trajectory

def trajectory_projection(dynamic, direction):

    direction = np.array(direction)/np.linalg.norm(direction)

    super_cell = dynamic.get_super_cell_matrix()
    trajectory = dynamic.get_relative_trajectory()

    atom_type_index = dynamic.structure.get_atom_type_index(super_cell=super_cell)
    number_of_atom_types = dynamic.structure.get_number_of_atom_types()

    projections = []

    for j in range(number_of_atom_types):
        projection = np.array([])
        for i in range(0, trajectory.shape[1]):
            if atom_type_index[i] == j:
           #     print('atom:', i, 'type:', atom_type_index[i])
                projection = np.append(projection, np.dot(trajectory[:, i, :], direction/np.dot(direction, direction)),axis=1)
        projections.append(projection)

    return np.array(projections)

#print(disp.relative_trajectory(cell, traj ,pos)) #    print(trajectory, position, cell) #    print (disp.relative_trajectory(cell, trajectory[:,j,:], position[j])) #            print(difference_matrix) #  print(normalized_trajectory) #   difference_matrix = np.array(np.dot(np.linalg.inv(cell),(IniSep)),dtype=int) #  print(normalized_trajectory) #     print('atom:', i, 'type:', atom_type_index[i])