matplotlib inline

from numpy import *
from matplotlib import *
from matplotlib.pyplot import *

#todo:
#now make these drawings for the 
#hexagon on diagonal multiplied by [1, 1.21, 1.31, 1.46]
#and the other way, turned 90 degrees and multiplied by [1, 0.83, 0.75, 0.69]
#and squeezed in one way

def rot(rotation):
    al=rotation*2*pi/360
    return np.array([[cos(al), sin(al)],[-sin(al), cos(al)]]).T

def unit_cell_matrix(a, rotation, extension_faction_11):
    hexagonal_cell = np.array([[a, 0], [-a/2, a*sqrt(3)/2]])
    rotated = np.dot(hexagonal_cell, rot(-rotation))

    extension_11 = np.dot(rot(45), np.dot([[extension_faction_11,0],[0,1]], rot(-45).T))
    return np.dot(rotated, extension_11)


#parameters of the unit cell: a, rotation angle, extension along diagonal
params = [[a, 15, 1],
          [1.21*a, 15, 1.1],
          [1.31*a, 15, 1.1],
          [1.46*a, 15, 1.1],
          [a, -15, 1],
          [0.9*a, -15, 1/1.1],
          [0.8*a, -15, 1/1.1],
          [0.7*a, -15, 1/1.1]]

param=params[0]

import numpy as np
from numpy.linalg import norm
from numpy.fft import fft2, fftshift
import matplotlib.pyplot as plt

def gaussian_2d(x, y, sigma=1.0):
    return np.exp(-(x**2 + y**2)/(2*sigma**2))

def draw_cell(cell, drawing_no):
    # Parameters
    N = 1024  # image size
    a = 32   # lattice constant
    r = 5    # atom radius
    sigma_atom = 0.0  # atom gaussian width
    sigma_window = N/8  # window function width
    
    # cell = np.array([[a, 0],[-a/2, a]])
    
    # Create coordinate grids
    x = np.linspace(-N/2, N/2, N)
    y = np.linspace(-N/2, N/2, N)
    X, Y = np.meshgrid(x, y)
    
    # Initialize image
    image = np.zeros((N, N))
    
    # Generate hexagonal lattice positions with random displacements
    n_cells = N//a
    for i in range(-n_cells//2, n_cells//2):
        for j in range(-n_cells//2, n_cells//2):
            x_pos,y_pos = dot(cell.T, [i,j])
            
            # Add random displacement
            dx = np.random.normal(0, sigma_atom)
            dy = np.random.normal(0, sigma_atom)
            
            # Add gaussian at each lattice point
            image += gaussian_2d(X - (x_pos + dx), Y - (y_pos + dy), r)
    
    # Apply window function. This removes ugly artifacts which arrise due to trunctation on the edge of the reconstructed plane
    window = gaussian_2d(X, Y, sigma_window)
    image *= window
    
    # Calculate Fourier transform
    ft = fftshift(np.abs(fft2(image)))
    
    # Plot results
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 5))

    
    ax1.imshow(image>0.6*amax(image), cmap=cm.Greys, extent=[-N/2,N/2-1,-N/2,N/2-1])
    ax1.set_title('Real Space')
    ax1.axis('equal')
    rlim = 80
    ax1.axis([-rlim,rlim,-rlim,rlim])
    
    clim = amax(ft)/3
    ax2.imshow(ft, cmap=cm.Greys_r, extent=[-N/2,N/2-1,-N/2,N/2-1], clim=[0,clim])
    
    lim = 80
    
    ax2.axis('equal')
    ax2.axis([-lim,lim,-lim,lim])
    
    ax2.set_title('Fourier Transform')
    
    for ax in [ax1,ax2]:
        ax.set_xticks([])
        ax.set_yticks([])

    savefig(f"drawing_{drawing_no}.png")


for i,param in list(enumerate(params)):
    cell = unit_cell_matrix(*param)
    draw_cell(cell, i)