#!/usr/bin/env python3

import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np
import sympy as sp
import sympy.physics.quantum as qt

sp.init_printing(use_unicode=True)

def Main():

    L = sp.Symbol('L', positive=True)
    r = sp.Symbol('r', positive=True)

    a1 = sp.Rational(1, 3)
    a2 = sp.Rational(2, 3)

    Ap = sp.sqrt(a2) * sp.Matrix([[0, 1], [0, 0]])
    Ao = sp.sqrt(a1) * sp.Matrix([[-1, 0], [0, 1]])
    Am = sp.sqrt(a2) * sp.Matrix([[0, 0], [-1, 0]])

    T = (
        qt.TensorProduct(Ap, Ap) +
        qt.TensorProduct(Ao, Ao) +
        qt.TensorProduct(Am, Am)
    )

    Print('T =', T)


    S  = Diag([-a1, -a1, -a1, 1])

    V = sp.Matrix([
        [-1/sp.sqrt(2), 0, 0, 1/sp.sqrt(2)],
        [0,             1, 0, 0           ],
        [0,             0, 1, 0           ],
        [1/sp.sqrt(2),  0, 0, 1/sp.sqrt(2)]
    ])

    Print('T = V^T S V', 'S = ', S)
    Print('V = ', V)

    SL = Diag([0, 0, 0, 1])
    a0 = -a1**(r-1)
    Sr = Diag([a0, a0, a0, 1])

    Print('<Ψ|Ψ> = ', Trace(SL))

    U = (
        qt.TensorProduct(Ap, Ap) +
        0*qt.TensorProduct(Ao, Ao) +
        -qt.TensorProduct(Am, Am)
    )

    Print('U =', U)

    O = U * V * SL * V
    Print('<S^z> =', Trace(O) )

    O = U * V * Sr * V * U * V * SL * V
    Print('<S^z_0S^z_r> =', Trace(O) )

    return

def Trace(M):
    return sp.simplify(sp.Trace(M))

def Diag(S):
    return sp.Matrix(len(S), len(S), lambda i,j: S[i] if i == j else 0)

def Print(*O):
    for o in O:
        sp.pprint(o)
    print()
    return

def M(m):
    return np.array(m, dtype=np.float64)

Main()