#!/usr/bin/env python
# coding: utf-8

# In[1]:


import qsymm


# In[2]:


import numpy as np
import sympy


# In[3]:


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")


# In[4]:


H = qsymm.Model(ham)


# In[5]:


print(H)


# In[6]:


H.tosympy(nsimplify=True)


# In[7]:


H.momenta


# In[8]:


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'])


# In[9]:


H2D.tosympy(nsimplify=True)


# In[10]:


H2D.momenta


# In[11]:


# 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)


# In[12]:


qsymm.PointGroupElement(R=np.diag([1, -1]), U='kron(sigma_y, sigma_z)')


# In[13]:


C4


# In[14]:


TR


# In[15]:


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


# In[16]:


C3 * C4


# In[17]:


C3**-1


# In[18]:


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


# In[19]:


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


# In[20]:


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


# In[21]:


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


# In[22]:


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


# In[23]:


from pathlib import Path

h2d_path = Path("docs/build/H2D.txt")
H2D_sympy = H2D.tosympy()

with h2d_path.open("w") as f:
    f.write(repr(H2D_sympy))


# In[24]:


with h2d_path.open() as f:
    data = f.read()

loaded_H2D = qsymm.Model(
    sympy.parsing.sympy_parser.parse_expr(data),
    momenta=['k_x', 'k_z']
)