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

# In[1]:


import numpy as np
import sympy

import qsymm


# In[2]:


# C3 rotational symmetry - invariant under 2pi/3
C3 = qsymm.rotation(1/3, spin=1/2)

# Time reversal
TR = qsymm.time_reversal(2, spin=1/2)

# Mirror symmetry in x
Mx = qsymm.mirror([1, 0], spin=1/2)

symmetries = [C3, TR, Mx]


# In[3]:


dim = 2  # Momenta along x and y
total_power = 3  # Maximum power of momenta


# In[4]:


family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[5]:


qsymm.hamiltonian_generator.check_symmetry(family, symmetries)


# In[6]:


# Spatial inversion
pU = np.array([
    [1.0, 0.0, 0.0, 0.0],
    [0.0, -1.0, 0.0, 0.0],
    [0.0, 0.0, 1.0, 0.0],
    [0.0, 0.0, 0.0, -1.0],
])
pS = qsymm.inversion(2, U=pU)

# Time reversal
trU = np.array([
    [0.0, 0.0, -1.0, 0.0],
    [0.0, 0.0, 0.0, -1.0],
    [1.0, 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
])
trS = qsymm.time_reversal(2, U=trU)

# Rotation
phi = 2.0 * np.pi / 4.0  # Impose 4-fold rotational symmetry
rotU = np.array([
    [np.exp(-1j*phi/2), 0.0, 0.0, 0.0],
    [0.0, np.exp(-1j*3*phi/2), 0.0, 0.0],
    [0.0, 0.0, np.exp(1j*phi/2), 0.0],
    [0.0, 0.0, 0.0, np.exp(1j*3*phi/2)],
])
rotS = qsymm.rotation(1/4, U=rotU)

symmetries = [rotS, trS, pS]


# In[7]:


dim = 2
total_power = 2


# In[8]:


family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[9]:


# Spatial inversion
pU = np.diag([1, -1, 1, -1])
pS = qsymm.inversion(3, pU)

# Time reversal
trU = np.array([
    [0.0, 0.0, -1.0, 0.0],
    [0.0, 0.0, 0.0, -1.0],
    [1.0, 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
])
trS = qsymm.time_reversal(3, trU)

# Rotation
phi = 2.0 * np.pi / 3.0  # Impose 3-fold rotational symmetry
rotU = np.array([
    [np.exp(1j*phi/2), 0.0, 0.0, 0.0],
    [0.0, np.exp(1j*phi/2), 0.0, 0.0],
    [0.0, 0.0, np.exp(-1j*phi/2), 0.0],
    [0.0, 0.0, 0.0, np.exp(-1j*phi/2)],
])
rotS = qsymm.rotation(1/3, axis=[0, 0, 1], U=rotU)

symmetries = [rotS, trS, pS]


# In[10]:


dim = 3
total_power = 2


# In[11]:


family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[12]:


# 3D real space rotation generators
L = qsymm.groups.L_matrices(3)
# Spin-1/2 matrices
S = qsymm.groups.spin_matrices(1/2)
# Spin-3/2 spin matrices
J = qsymm.groups.spin_matrices(3/2)
# Continuous rotation generators
symmetries = [qsymm.ContinuousGroupGenerator(l, s) for l, s in zip(L, S)]

dim = 3
total_power = 2

family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[13]:


symmetries = [qsymm.ContinuousGroupGenerator(l, s) for l, s in zip(L, S)]
# Add inversion
symmetries.append(qsymm.PointGroupElement(-np.eye(3), False, False, np.eye(2)))
dim = 3
total_power = 2
family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[14]:


symmetries = [qsymm.ContinuousGroupGenerator(l, s) for l, s in zip(L, J)]
dim = 3
total_power = 2
family = qsymm.continuum_hamiltonian(symmetries, dim, total_power, prettify=True)
qsymm.display_family(family)


# In[15]:


# Double spin-1/2 representation
spin = [np.kron(s, np.eye(2)) for s in qsymm.groups.spin_matrices(1/2)]
# C3 rotational symmetry
C3 = qsymm.rotation(1/3, axis=[0, 0, 1], spin=spin)

# Time reversal
TR = qsymm.time_reversal(3, spin=spin)

# Mirror x
Mx = qsymm.mirror([1, 0, 0], spin=spin)

# Inversion
IU = np.kron(np.eye(2), np.diag([1, -1]))
I = qsymm.inversion(3, IU)


# In[16]:


dim = 3
total_power = 2
family = qsymm.continuum_hamiltonian([C3, TR, Mx, I], dim, total_power, prettify=True)
qsymm.display_family(family)


# In[17]:


identity_terms = qsymm.continuum_hamiltonian([qsymm.groups.identity(dim, 1)], dim, total_power)
identity_terms = [
    qsymm.Model({
        key: np.kron(val, np.eye(4))
        for key, val in term.items()
    })
    for term in identity_terms
]

family = qsymm.hamiltonian_generator.subtract_family(family, identity_terms, prettify=True)
qsymm.display_family(family)


# In[18]:


dim = 3
total_power = 2
no_c3_family = qsymm.continuum_hamiltonian([TR, Mx, I], dim, total_power, prettify=True)
no_c3_family = qsymm.hamiltonian_generator.subtract_family(no_c3_family, identity_terms, prettify=True)


# In[19]:


new_terms = qsymm.hamiltonian_generator.subtract_family(no_c3_family, family, prettify=True)
qsymm.display_family(new_terms)


# In[20]:


dim = 3
total_power = 2
broken_family = qsymm.continuum_hamiltonian([TR, Mx], dim, total_power, prettify=True)
broken_family = qsymm.hamiltonian_generator.subtract_family(broken_family, identity_terms, prettify=True)
new_terms = qsymm.hamiltonian_generator.subtract_family(broken_family, no_c3_family, prettify=True)
qsymm.display_family(new_terms)