from NL.Math.Utilities import rotationMatrix
import math
from QuantumATK import *

def twister_displacement(x, rotation_angle_per_z, rotation_axis,
                 rotation_axis_center, z_start, z_end):
    """
       Function for twisting a 1-d structure
       @param x : Coordinates of 1-d structure
       @param rotation_angle_per_z : size of twist in angle/length
       @param rotation_axis : axis to apply twist along
       @param rotation_axis_center : center of the rotation axis
       @param z_start : z value for starting the twist
       @param z_end : z value for ending the twist
    """

    # do not twist for z > z_end
    z = x[2]
    z = min(z,z_end)
    # do not twist for z < z_start
    z = z - z_start
    z = max (z,0.0)
    # find twist angle
    theta = z*rotation_angle_per_z
    # calculate the rotation matrix
    rotation_matrix = rotationMatrix(theta, *rotation_axis)
    # apply rotation
    return rotation_axis_center+numpy.dot(rotation_matrix, x-rotation_axis_center)

def wrapping_displacement(x, width, wrapping_angle):
    """
       Function for converting a nanosheet coordinate into a partly wrapped nanotube
       @param x : Coordinates of nanosheet atom
       @param width : Width of the nano-sheet
       @param wrapping_angle : maximum wrapping angle of the nanotube in radians
    """
    # calculate the average radius of the incomplete wrapped tube
    radius = width/wrapping_angle
    # find the angle of the current atom
    angle = (x[2]-width/2.)/radius
    # calculate the radius of the current atom
    atom_radius = radius+x[1]

    # return atom position of the wrapped atom
    return numpy.array([x[0], atom_radius*math.cos(angle),atom_radius*math.sin(angle)])

def Moebius(ribbon, n, m, repetition):
    """
       Function for generating a moebius molecule
       @param n : Chiral vector index
       @param m : Chiral vector index
       @param repetition : Repetition along z
    """

    # build n,m ribbon
    #ribbon = NanoRibbon(n,m)
    ribbon = ribbon.repeat(1,1,repetition)

    # get properties of the ribbon
    lattice = ribbon.bravaisLattice()
    elements = ribbon.elements()
    cartesian_coordinates=ribbon.cartesianCoordinates().inUnitsOf(Angstrom)

    # calculate the length of the 1-d structure
    z_length = numpy.linalg.norm(lattice.primitiveVectors()[2].inUnitsOf(Angstrom))

    # calculate twist parameters
    rotation_angle_per_z =  math.pi /z_length
    rotation_axis = numpy.array([0,0,1])
    rotation_axis_center = numpy.sum(cartesian_coordinates,axis=0)/len(cartesian_coordinates)

    # define a function of one variable, f(c),  for displacing the atoms
    f = lambda c : twister_displacement(c, rotation_angle_per_z, rotation_axis,
                                rotation_axis_center, 0.,z_length)
    # apply the function to find new displaced coordinates
    cartesian_coordinates = numpy.apply_along_axis(f, 1, cartesian_coordinates)
    cartesian_center = numpy.sum(cartesian_coordinates,axis=0)/len(cartesian_coordinates)
    cartesian_coordinates = cartesian_coordinates - cartesian_center


    # define a function of one variable, f(c),  for displacing the atoms
    f = lambda c : wrapping_displacement(c, z_length,2.0*math.pi)
    # apply the function to find new displaced coordinates
    cartesian_coordinates = numpy.apply_along_axis(f, 1, cartesian_coordinates)

    return MoleculeConfiguration(
            elements=elements,
            cartesian_coordinates=cartesian_coordinates * Angstrom
            )

ribbon = nlread('ribbon.nc', BulkConfiguration)[-1]
moebius = Moebius(ribbon,1,1,100)
nlsave('moebius.nc', moebius)