import numpy as np


def get_trajectory_txt(trajectory):

    cell = trajectory.get_cells()[0]

    a = np.linalg.norm(cell[0])
    b = np.linalg.norm(cell[1])
    c = np.linalg.norm(cell[2])

    alpha = np.arccos(np.dot(cell[1], cell[2])/(c*b))
    gamma = np.arccos(np.dot(cell[1], cell[0])/(a*b))
    beta = np.arccos(np.dot(cell[2], cell[0])/(a*c))

    xhi = a
    xy = b * np.cos(gamma)
    xz = c * np.cos(beta)
    yhi = np.sqrt(pow(b,2)- pow(xy,2))
    yz = (b*c*np.cos(alpha)-xy * xz)/yhi
    zhi = np.sqrt(pow(c,2)-pow(xz,2)-pow(yz,2))

    xhi = xhi + max(0,0, xy, xz, xy+xz)
    yhi = yhi + max(0,0, yz)

    xlo_bound = np.min([0.0, xy, xz, xy+xz])
    xhi_bound = xhi + np.max([0.0, xy, xz, xy+xz])
    ylo_bound = np.min([0.0, yz])
    yhi_bound = yhi + np.max([0.0, yz])
    zlo_bound = 0
    zhi_bound = zhi

    ind = trajectory.get_array('steps')
    lammps_data_file = ''
    for i, position_step in enumerate(trajectory.get_positions()):
        lammps_data_file += 'ITEM: TIMESTEP\n'
        lammps_data_file += '{}\n'.format(ind[i])
        lammps_data_file += 'ITEM: NUMBER OF ATOMS\n'
        lammps_data_file += '{}\n'.format(len(position_step))
        lammps_data_file += 'ITEM: BOX BOUNDS xy xz yz pp pp pp\n'
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(xlo_bound, xhi_bound, xy)
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(ylo_bound, yhi_bound, xz)
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(zlo_bound, zhi_bound, yz)
        lammps_data_file += ('ITEM: ATOMS x y z\n')
        for position in position_step:
            lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(*position)
    return lammps_data_file


def parameters_to_input_file(parameters_object):

    parameters = parameters_object.get_dict()
    input_file = ('STRUCTURE FILE POSCAR\nPOSCAR\n\n')
    input_file += ('FORCE CONSTANTS\nFORCE_CONSTANTS\n\n')
    input_file += ('PRIMITIVE MATRIX\n')
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[0])
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[1])
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[2])
    input_file += ('\n')
    input_file += ('SUPERCELL MATRIX PHONOPY\n')
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[0])
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[1])
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[2])
    input_file += ('\n')

    return input_file

import numpy as np


def get_trajectory_txt(trajectory):

    cell = trajectory.get_cells()[0]

    a = np.linalg.norm(cell[0])
    b = np.linalg.norm(cell[1])
    c = np.linalg.norm(cell[2])

    alpha = np.arccos(np.dot(cell[1], cell[2])/(c*b))
    gamma = np.arccos(np.dot(cell[1], cell[0])/(a*b))
    beta = np.arccos(np.dot(cell[2], cell[0])/(a*c))

    xhi = a
    xy = b * np.cos(gamma)
    xz = c * np.cos(beta)
    yhi = np.sqrt(pow(b,2)- pow(xy,2))
    yz = (b*c*np.cos(alpha)-xy * xz)/yhi
    zhi = np.sqrt(pow(c,2)-pow(xz,2)-pow(yz,2))

    xhi = xhi + max(0,0, xy, xz, xy+xz)
    yhi = yhi + max(0,0, yz)

    xlo_bound = np.min([0.0, xy, xz, xy+xz])
    xhi_bound = xhi + np.max([0.0, xy, xz, xy+xz])
    ylo_bound = np.min([0.0, yz])
    yhi_bound = yhi + np.max([0.0, yz])
    zlo_bound = 0
    zhi_bound = zhi

    ind = trajectory.get_array('steps')
    lammps_data_file = ''
    for i, position_step in enumerate(trajectory.get_positions()):
        lammps_data_file += 'ITEM: TIMESTEP\n'
        lammps_data_file += '{}\n'.format(ind[i])
        lammps_data_file += 'ITEM: NUMBER OF ATOMS\n'
        lammps_data_file += '{}\n'.format(len(position_step))
        lammps_data_file += 'ITEM: BOX BOUNDS xy xz yz pp pp pp\n'
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(xlo_bound, xhi_bound, xy)
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(ylo_bound, yhi_bound, xz)
        lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(zlo_bound, zhi_bound, yz)
        lammps_data_file += ('ITEM: ATOMS x y z\n')
        for position in position_step:
            lammps_data_file += '{0:20.10f} {1:20.10f} {2:20.10f}\n'.format(*position)
    return lammps_data_file


def parameters_to_input_file(parameters_object):

    parameters = parameters_object.get_dict()
    input_file = ('STRUCTURE FILE POSCAR\nPOSCAR\n\n')
    input_file += ('FORCE CONSTANTS\nFORCE_CONSTANTS\n\n')
    input_file += ('PRIMITIVE MATRIX\n')
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[0])
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[1])
    input_file += ('{} {} {} \n').format(*np.array(parameters['primitive'])[2])
    input_file += ('\n')
    input_file += ('SUPERCELL MATRIX PHONOPY\n')
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[0])
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[1])
    input_file += ('{} {} {} \n').format(*np.array(parameters['supercell'])[2])
    input_file += ('\n')

    return input_file
