from copy import deepcopy

import numpy as np
import pandas as pd
from scipy.linalg import expm
import matplotlib.pyplot as plt
%matplotlib inline

from ph121c_lxvm import tensor, measure

# Build initial state
down = np.array([1., 0.]).reshape(2, 1)
up   = down[::-1].reshape((2, 1))
ξ = (up - np.sqrt(3) * down) / 2

def make_wave (L, d):
    """Create the inital wavefunctions."""
    ψ = tensor.mps(L=L, d=d)
    ψ.from_arr([ ξ for _ in range(L) ], center=-1)  
    return ψ

def eentropy (tren, i):
    """Calculate the entanglement entropy at a cut between quanta i, i+1."""
    tren.groupby_quanta_tag([i])
    u, s, vh = np.linalg.svd(tren.center.mat)
    return measure.entropy.entanglement(s)

sx = np.array([[0, 1], [1, 0]])
sy = np.array([[0, -1j], [1j, 0]])
sz = np.diag([1, -1])

# Build pieces of Hamiltonian in gate representation
def build_pieces_of_H (L, d, W, seed):
    """Build the field, odd, and even term Hamiltonians and also their union."""
    rng = np.random.default_rng(seed=seed)
    H_field = np.empty(L, dtype='O')
    for i in range(H_field.size):
        H_field[i] = tensor.mpo(L=L, d=d)
        H_field[i].set_local_oper(
            -(rng.uniform(-W, W) * sx + rng.uniform(-W, W) * sz), i + 1
        )
    H_odd = np.empty(L//2, dtype='O')
    for i in range(H_odd.size):
        H_odd[i] = tensor.mpo(L=L, d=d)
        H_odd[i].set_local_oper(-sz, 2*i + 1)
        H_odd[i].set_local_oper(sz, 2*i + 1 + 1)
    H_even = np.empty(L//2 - 1 + L%2, dtype='O')
    for i in range(H_even.size):
        H_even[i] = tensor.mpo(L=L, d=d)
        H_even[i].set_local_oper(-sz, 2*(i + 1))
        H_even[i].set_local_oper(sz, 2*(i + 1) + 1)
    H_full = np.array([*H_field, *H_odd, *H_even], dtype='O')
    return (H_field, H_odd, H_even, H_full)


# Construct propagators
def build_propagators (L, d, δτ, H_field, H_odd, H_even):
    """Exponentiate each non-commuting piece of the Hamiltonian"""
    U_field = tensor.mpo(L=L, d=d)
    for i, e in enumerate(H_field):
        U_field.set_local_oper(expm(- 1.j * δτ * e[0].mat), i+1)
    U_odd = tensor.mpo(L=L, d=d)
    for i, e in enumerate(H_odd):
        U_odd.set_local_oper(
            expm(- 1.j * δτ * np.kron(e[1].mat, e[0].mat)),
            2 * i + 1
        )
    U_even = tensor.mpo(L=L, d=d)
    for i, e in enumerate(H_even):
        U_even.set_local_oper(
            expm(- 1.j * δτ * np.kron(e[1].mat, e[0].mat)),
            2 * (i + 1)
        )
    return (U_field, U_odd, U_even)

def make_observables (L, d):
    """Create the observable operators of interest."""
    s1x = tensor.mpo(L=L, d=d)
    s1x.set_local_oper(sx, 1)
    sL2x = tensor.mpo(L=L, d=d)
    sL2x.set_local_oper(sx, L//2)
    s1z = tensor.mpo(L=L, d=d)
    s1z.set_local_oper(sz, 1)
    sL2z = tensor.mpo(L=L, d=d)
    sL2z.set_local_oper(sz, L//2)
    return (s1x, sL2x, s1z, sL2z)

W = 3
L = 24
d = 2
δτ = 0.1
N = 20
seeds = [935, 483, 102, 567, 304]

%%time
chi = 16
Nstp = 20
ψ = make_wave(L=L, d=d)
ψ_results = dict(t=[], S=[], s1x=[], sL2x=[], s1z=[], sL2z=[], seed=[])
for seed in seeds:
    H_field, H_odd, H_even, H_full = build_pieces_of_H(
        L=L, d=d, W=W, seed=seed
    )
    U_field, U_odd, U_even = build_propagators(
        L=L, d=d, δτ=δτ, H_field=H_field, H_odd=H_odd, H_even=H_even
    )
    wave = deepcopy(ψ)
    s1x, sL2x, s1z, sL2z = make_observables(L=L, d=d)
    ψ_results['t'].append(0)
    ψ_results['seed'].append(seed)
    ψ_results['S'].append(eentropy(ψ, L//2))
    ψ_results['s1x'].append(s1x.expval(ψ))
    ψ_results['sL2x'].append(sL2x.expval(ψ))
    ψ_results['s1z'].append(s1z.expval(ψ))
    ψ_results['sL2z'].append(sL2z.expval(ψ))
    # TEBD pattern
    for i in range(N):
        for e in [U_field, U_even, U_odd]:
            e.oper(wave, inplace=True)
        wave.trim_bonds(chi)
        wave.normalize()
        # measure entropy
        ψ_results['t'].append(δτ * (i + 1))
        ψ_results['seed'].append(seed)
        ψ_results['S'].append(eentropy(ψ, L//2))
        ψ_results['s1x'].append(s1x.expval(ψ))
        ψ_results['sL2x'].append(sL2x.expval(ψ))
        ψ_results['s1z'].append(s1z.expval(ψ))
        ψ_results['sL2z'].append(sL2z.expval(ψ))

df = pd.DataFrame(ψ_results)

fig, ax = plt.subplots()
# grp = df.groupby(['L', 't']).mean().reset_index()
for k, seed in enumerate(seeds):
    which = (df.seed == seed)
    ax.plot(
        df.t[which].values, df.S[which].values,
        label=f'$\\chi={chi}$'
    )
grp = df.groupby(['t']).mean().reset_index()
# ax.plot(np.arange(N) * δτ, np.mean(df.S))
ax.set_title('$S_{L/2}$')
ax.set_xlabel('time')
ax.set_ylabel('$S_{L/2}$')
ax.legend()
fig.tight_layout()
plt.show()

grp