import numpy as np
import matplotlib.pyplot as plt


def deg_2_rad(winkel):
    return winkel/180.0 * np.pi


# all units in angstrom
base_a_m = 5.75
base_b_m = 4.5
base_c_m = 5.38

base_c_r = 2.856
base_b_r = 4.554
base_a_r = base_b_r

alpha_m = 122.64  # degree


def mono_2_rutile(c_m, a_m):
    a_r = np.cos(deg_2_rad(alpha_m - 90)) * c_m * base_c_m
    c_r = (a_m) * base_a_m + np.sin(deg_2_rad(alpha_m - 90)) * c_m * base_c_m
    return a_r, c_r


def main():
    offset_a_m = 0.25 - 0.23947
    offset_c_m = 0.02646

    offset_a_r, offset_c_r = mono_2_rutile(offset_c_m, offset_a_m)

    print("A_r: ", offset_a_r, "C_r: ", offset_c_r)

    x = np.arange(10)
    y = np.arange(10)
    X, Y = np.meshgrid(x, y)

    atom_x = offset_a_r + X * base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
    atom_x[np.mod(X, 2) == 0] -= 2 * offset_a_r

    atom_y = offset_c_r + 0.5 * Y * base_a_r
    atom_y[np.mod(X, 2) == 0] -= 2 * offset_c_r

    fig, ax = plt.subplots(1)
    plt.scatter(atom_x, atom_y)

    atom_x_rut = X * base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
    atom_y_rut = Y * 0.5 * base_a_r

    plt.scatter(atom_x_rut, atom_y_rut)
    print(np.square(atom_x-atom_x_rut) + np.square(atom_y-atom_y_rut))

    cryst = np.loadtxt("./crystal_V.xyz", delimiter=" ")
    print(cryst.shape)
    plt.scatter(-cryst[:, 0]+5.5*base_c_r, cryst[:, 2])
    ax.set_aspect(1)
    plt.show()


if __name__ == "__main__":
    main()