import qsymm

import numpy as np
import sympy

ham = ("hbar^2 / (2 * m) * (k_x**2 + k_y**2 + k_z**2) * eye(2) +" +
        "alpha * sigma_x * k_x + alpha * sigma_y * k_y + alpha * sigma_z * k_z")

H = qsymm.Model(ham)

print(H)

H.tosympy(nsimplify=True)

H.momenta

ham2D = ("hbar^2 / (2 * m) * (k_x**2 + k_z**2) * eye(2) +" +
         "alpha * sigma_x * k_x + alpha * sigma_y * k_z")
H2D = qsymm.Model(ham2D, momenta=['k_x', 'k_z'])

H2D.tosympy(nsimplify=True)

H2D.momenta

# Identity in 3D
E = qsymm.identity(3)
# Inversion in 3D
I = qsymm.inversion(3)
# 4-fold rotation around the x-axis
C4 = qsymm.rotation(1/4, [1, 0, 0])
# 3-fold rotation around the [1, 1, 1] axis
C3 = qsymm.rotation(1/3, [1, 1, 1])
# Time reversal
TR = qsymm.time_reversal(3)
# Particle-hole
PH = qsymm.particle_hole(3)

C4

TR

cubic_gens = {I, C4, C3, TR, PH}
cubic_group = qsymm.groups.generate_group(cubic_gens)
print(len(cubic_group))

C3 * C4

C3**-1

H_with_TR = TR.apply(H)
H_with_TR.tosympy(nsimplify=True)

sz = qsymm.ContinuousGroupGenerator(None, np.array([[1, 0], [0, -1]]))

sz.apply(H).tosympy(nsimplify=True)

discrete_symm, continuous_symm = qsymm.symmetries(H, cubic_group)
print(len(discrete_symm), len(continuous_symm))

family = qsymm.continuum_hamiltonian(discrete_symm, dim=3, total_power=2, prettify=True)
qsymm.display_family(family)