import tensorflow as tf
import tensorflow.linalg as tfl
import numpy as np
import matplotlib.pyplot as plt

from tqdm.notebook import tqdm

physical_devices = tf.config.list_physical_devices('GPU')
print('GPU devices:', physical_devices)

if physical_devices:
    for device in physical_devices:
        tf.config.experimental.set_memory_growth(device, True)

sigma_x = tfl.LinearOperatorFullMatrix([[0., 1.], [1., 0.]])
sigma_z = tfl.LinearOperatorFullMatrix([[1., 0.], [0., -1.]])
id      = tfl.LinearOperatorIdentity(2)

def H_perc(data, labels):
    n_data, n = data.shape

    # define a batch of ientities operators where batch is the #data
    id_batch = tfl.LinearOperatorIdentity(2, batch_shape=[n_data])

    # define a batch of sigma_z operators where batch is the #data
    sigma_z_batch = tf.tile(tf.expand_dims(tf.constant([[1., 0.], [0., -1.]]), axis=0), [n_data, 1, 1])
    sigma_z_op = tfl.LinearOperatorFullMatrix(sigma_z_batch)

    # define a batch of data operators where batch is the #data.
    # each coordinate of a single datum is casted into a 2x2 operator,
    # which will be composed with the corresponding sigma_z operator
    # result is a list of n [n_data, 2, 2] opearators
    data_ops = [tfl.LinearOperatorDiag(tf.repeat(data, [2], axis=1)[:,i:i+2]) for i in range(0, n*2, 2)]

    ops = [tfl.LinearOperatorKronecker([id_batch]*i+[tfl.LinearOperatorComposition([data_ops[i], sigma_z_op])]+[id_batch]*(n-i-1)).to_dense() for i in range(n)]
    
    # opearotrs are first stacked on the n dimension and reduced, then labels are applied,
    # then heaviside step function is applied (relu), and finally reduced over n_data
    return tf.reduce_sum(tf.nn.relu(-tf.reshape(labels, (-1, 1, 1))*tf.reduce_sum(tf.stack(ops), axis=0)), axis=0)# / tf.sqrt(n+0.)

# reimplementation of the previous function with an eye over ram usage
def H_perc_nobatch(data, labels):
    n_data, n = data.shape

    h_perc = tf.zeros((2**n, 2**n), dtype='float32')
    
    for i in tqdm(range(n_data), desc='Constructing H_perc'):

        op = tf.zeros((2**n, 2**n), dtype='float32')
        for j in range(n):

            data_op = tfl.LinearOperatorDiag(tf.repeat(data[i, :], [2], axis=0)[2*j:2*(j+1)])
            op     += tfl.LinearOperatorKronecker([id]*j+[tfl.LinearOperatorComposition([data_op, sigma_z])]+[id]*(n-j-1)).to_dense()
            del data_op

        h_perc += tf.nn.relu(-labels[i]*op)
        del op

    return h_perc #/ tf.sqrt(n+0.)

def H_x(n):
    sigma_xs = [tfl.LinearOperatorKronecker([id]*i+[sigma_x]+[id]*(n-i-1)).to_dense() for i in range(n)]
    return -tf.reduce_sum(tf.stack(sigma_xs), axis=0)

def H_z(n):
    sigma_zs = [tfl.LinearOperatorKronecker([id]*i+[sigma_z]+[id]*(n-i-1)).to_dense() for i in range(n)]
    return tf.reduce_sum(tf.stack(sigma_zs), axis=0)

def H_QA(p, P, Hz, Hx):
    frac = p/P
    return tf.cast(frac*Hz + (1-frac)*Hx, dtype='complex128')


def init_state(n):
    return tf.ones((2**n,), dtype='complex128')/tf.sqrt((2.**n+0.j))

data   = tf.constant([[1., -1.], [-1., -1.]])
labels = tf.constant([1., -1.])

h_perc = H_perc(data, labels)
h_perc

n_data, n = data.shape

# define a batch of ientities operators where batch is the #data
id_batch = tfl.LinearOperatorIdentity(2, batch_shape=[n_data])

# define a batch of sigma_z operators where batch is the #data
sigma_z_batch = tf.tile(tf.expand_dims(tf.constant([[1., 0.], [0., -1.]]), axis=0), [n_data, 1, 1])
sigma_z_op = tfl.LinearOperatorFullMatrix(sigma_z_batch)

# define a batch of data operators where batch is the #data.
# each coordinate of a single datum is casted into a 2x2 operator,
# which will be composed with the corresponding sigma_z operator
# result is a list of n [n_data, 2, 2] opearators
data_ops = [tfl.LinearOperatorDiag(tf.repeat(data, [2], axis=1)[:,i:i+2]) for i in range(0, n*2, 2)]

ops = [tfl.LinearOperatorKronecker([id_batch]*i+[tfl.LinearOperatorComposition([data_ops[i], sigma_z_op])]+[id_batch]*(n-i-1)).to_dense() for i in range(n)]


tfl.LinearOperatorKronecker([tfl.LinearOperatorFullMatrix(np.array([[0., 0.], [0., 1.]], 'float32')), id]).to_dense()

tfl.LinearOperatorKronecker([id, tfl.LinearOperatorFullMatrix(np.array([[0., 0.], [0., 1.]], 'float32'))]).to_dense()

n = 10
hadamard = tfl.LinearOperatorFullMatrix(np.array([[1., 1.], [1., -1.]], 'float32') / np.sqrt(2.))

ops = [tfl.LinearOperatorKronecker([id]*i+[hadamard]+[id]*(n-i-1)).to_dense() for i in range(n)]

np.diag(tfl.LinearOperatorKronecker([sigma_z]*10).to_dense())[:4]

tf.reduce_sum(tf.stack(ops), axis=0)[:,0]

# define a batch of sigma_z operators where batch is the #data
sigma_z_batch = tf.tile(tf.expand_dims(tf.constant([[1., 0.], [0., -1.]]), axis=0), [n_data, 1, 1])
sigma_z_op = tfl.LinearOperatorFullMatrix(sigma_z_batch)

# define a batch of data operators where batch is the #data.
# each coordinate of a single datum is casted into a 2x2 operator,
# which will be composed with the corresponding sigma_z operator
# result is a list of n [n_data, 2, 2] opearators
data_ops = [tfl.LinearOperatorDiag(tf.repeat(data, [2], axis=1)[:,i:i+2]) for i in range(0, n*2, 2)]

ops = [tfl.LinearOperatorKronecker([id_batch]*i+[tfl.LinearOperatorComposition([data_ops[i], sigma_z_op])]+[id_batch]*(n-i-1)).to_dense() for i in range(n)]


H_perc_nobatch(data, labels)

tfl.eigh(h_perc)

h_x = H_x(10)
h_x = tf.cast(h_x, 'complex64')

eigvals, eigvecs = tfl.eigh(h_x)

tfl.matmul(eigvecs, tfl.matmul(tfl.diag(tf.exp(-1.j*eigvals)), eigvecs, adjoint_b=True), adjoint_a=False)

datum_ops = [tfl.LinearOperatorDiag(tf.repeat(data[0,:], [2], axis=0)[i:i+2]) for i in range(0, 2*2, 2)]
datum_ops[1].to_dense()

init_state(2)

h_t = H_QA(1, 1_000, h_perc, h_x)
h_t

def ed_qa_step(state, Ht, dt):

    eigvals, eigvecs = tfl.eigh(Ht)
    # define time evolution operator
    U = tf.exp(-1.j*eigvals*dt)

    # rewrite state in eigenvector basis, apply time evolution operator, project back to computational basis
    evolved = tfl.matvec(eigvecs, U * tfl.matvec(eigvecs, state, adjoint_a=True))
    return evolved

def create_dataset(N : int, features : int):
    """Create dataset as described by ref. paper, i.e. random +-1 values."""
    x = np.random.randint(2, size=(N, features))
    x[ x == 0 ] = -1  # data is encoded as +- 1
    return x

tau = 1_000
P   = 1_000
dt  = tau/P
N_data = 8
N_feat = 10

data   = create_dataset(N_data, N_feat).astype('float32')
labels = tf.ones((N_data), 'float32')

data = np.load('patterns8_10.npy')
#h_perc = tf.cast(H_perc_nobatch(data, labels), 'complex128')


h_perc = tf.cast(H_perc_nobatch(data, labels), 'complex128')
E0     = tfl.eigh(h_perc)[0][0]
h_x    = tf.cast(H_x(data.shape[1]), 'complex128')
#h_z    = tf.cast(H_z(data.shape[1]), 'complex128')

state = init_state(data.shape[1])
loss  = []

pbar = tqdm(range(P), desc='ED QA')
for i in pbar:

    h_t   = H_QA(i+1, P, h_perc, h_x)
    state = ed_qa_step(state, h_t, dt)

    loss.append(tf.cast((tf.tensordot(tf.math.conj(state), tfl.matvec(tf.cast(h_perc, 'complex128'), state), axes=1)-E0)/N_feat, 'float32').numpy())

    pbar.set_postfix({'loss':loss[-1], 'norm':tf.cast(tf.norm(state), 'float32').numpy()})

state

np.save('result', state.numpy())

which_down = tf.where(tf.cast(tfl.eigh(h_perc)[1][0], 'float32') > 0.).numpy()
which_down

state.numpy()[which_down[0]]

plt.plot((np.array(range(P))+1)/P, abs(np.asarray(loss)))
plt.yscale('log')
plt.ylabel(r'$\epsilon(s)$')
plt.xlabel(r'$s$')
plt.grid(alpha=0.2)

# comparison with dQA approach

psi_dqa = np.load('stato2.npy')
psi_qa  = np.load('result.npy')


tf.cast(tf.norm(tfl.matvec(tf.cast(psi_dqa, 'complex128'), psi_qa, adjoint_a=True )), 'float32')

tau = 500
P   = 1_000
dt  = tau/P
N_data = 3
N_feat = 4

data   = create_dataset(N_data, N_feat)
labels = np.ones((N_data), 'float32')

h_perc_diag = H_perc_diag(data, labels)
E0     = np.sort(h_perc_diag)[0]
h_x    = H_x(data.shape[1])
h_perc = np.diag(h_perc_diag)

eigvals_array = np.empty((P, 2**N_feat))

pbar = tqdm(range(P), desc='Spectral analysis')
for i in pbar:
    
    h_t   = H_QA(i+1, P, h_perc, h_x)
    eigvals, eigvecs = ncp.linalg.eigh(h_t)

    eigvals_array[i] = np.real_if_close(eigvals.get())

np.save('eigvals_3-4', eigvals_array)

plt.plot(np.linspace(0, 1, P), eigvals_array[:,::1], lw=1.8)
plt.xlabel(r'$s(t)$', fontsize=14)
plt.ylabel(r'Energy [$a.u.$]', fontsize=14)
plt.tick_params('both', which='major', labelsize=12)
plt.title('Quantum Annealing\neigenvalues evolution', fontsize=16)
#plt.grid(alpha=0.3)
plt.savefig('qa_eigevolution.svg')

plt.savefig('qa_eigevolution.svg')

state = tf.cast(psi_dqa.squeeze(), 'complex128')
epsilon = (tf.cast(tf.tensordot(tf.math.conj(state), tfl.matvec(h_perc, state), axes=1), 'float32').numpy()-E0)/N_data
epsilon

state

import matplotlib.pyplot as plt
import numpy as np
import cupy
import math
from datetime import datetime
from tqdm.notebook import tqdm
from importlib.util import find_spec
GPU = True

if GPU:
    if find_spec('cupy') is not None:
        import cupy as ncp
    else:
        print('Selected device is GPU but cupy is not installed, falling back to numpy')
        import numpy as ncp
else:
    import numpy as ncp

def kronecker_prod(operators):

    result = operators[0]
    for op in operators[1:]:

        result = ncp.kron(result, op)

    return result

def ReLU(x):
    return x * (x > 0)

def H_perc_diag(data, labels):
    ## ASSUMING H_PERC IS DIAGONAL, WHICH IS IN THE COMPUTATIONAL BASIS ##
    n_data, n = data.shape
    identity = ncp.ones((2,), 'float64')
    sigma_z  = ncp.array([1., -1.])

    h_perc = ncp.zeros((2**n,), dtype='float64')
    
    for i in range(n_data):

        op = ncp.zeros((2**n,), dtype='float64')
        for j in range(n):

            op += kronecker_prod([identity]*j+[data[i, j] * sigma_z]+[identity]*(n-j-1))

        h_perc += ReLU(-labels[i]*op)
        del op

    return h_perc / ncp.sqrt(n)

def H_x(n):
    identity = ncp.diag([1., 1.])
    sigma_x  = ncp.array([[0., 1.], [1., 0.]])

    op = ncp.zeros((2**n, 2**n), dtype='float64')
    for j in range(n):
        op += kronecker_prod([identity]*j+[sigma_x]+[identity]*(n-j-1))

    return -op

def H_QA(p, P, Hz, Hx):
    frac = p/P
    return (frac*Hz + (1-frac)*Hx).astype('complex128')


def init_state(n):
    return (ncp.ones((2**n,))/ncp.sqrt(2**n)).astype('complex128')


def ed_qa_step(state, Ht, dt):

    eigvals, eigvecs = ncp.linalg.eigh(Ht)
    # define time evolution operator
    U = ncp.exp(-1.j*eigvals*dt)

    # rewrite state in eigenvector basis, apply time evolution operator, project back to computational basis
    evolved = ncp.dot(eigvecs, U * ncp.dot(eigvecs.transpose().conjugate(), state))
    return evolved

def create_dataset(N : int, features : int):
    """Create dataset as described by ref. paper, i.e. random +-1 values."""
    x = np.random.randint(2, size=(N, features))
    x[ x == 0 ] = -1  # data is encoded as +- 1
    return x

tau = 1_000
P   = 5
dt  = 0.5
 
N_feats = np.array([i for i in range(1, 14)])
N_datas = np.arange(1, np.floor(0.83*N_feats[[-1]]), dtype='int64')

memory = cupy.get_default_memory_pool()

#data   = np.load('../data/patterns8_10.npy')


N_datas

used_ram_tot = np.empty((N_feats.shape[0], N_datas.shape[0], P, 4))
times = np.empty((N_feats.shape[0], N_datas.shape[0], P+1))

data_pbar = tqdm(total=len(N_datas), desc='testing N_data', leave=True)
p_pbar = tqdm(total=P, desc='ED QA', leave=True)
for m, N_feat in enumerate(tqdm(N_feats, desc='testing N_feat')):
    for n, N_data in enumerate(N_datas):
        start = datetime.now()
        data   = create_dataset(N_data, N_feat)
        labels = np.ones((N_data), 'float32')
        h_perc_diag = H_perc_diag(data, labels)
        E0     = np.sort(h_perc_diag)[0]
        h_x    = H_x(data.shape[1])
        h_perc = np.diag(h_perc_diag)
        #h_z    = tf.cast(H_z(data.shape[1]), 'complex128')

        state = init_state(data.shape[1])
        end = datetime.now()
        times[m, n, -1] = (end-start).microseconds
        loss  = []
        
        for i in range(P):
            start = datetime.now()
            h_t   = H_QA(i+1, P, h_perc, h_x)
            used_ram_tot[m, n, i, 0] = memory.used_bytes()
            state = ed_qa_step(state, h_t, dt)
            used_ram_tot[m, n, i, 1] = memory.used_bytes()
            loss.append(ncp.real_if_close((ncp.tensordot(state.conjugate(), ncp.dot(h_perc, state), axes=1)-E0)/N_feat))

            end = datetime.now()
            times[m, n, i] = (end-start).microseconds

            pbar.set_postfix({'loss':ncp.real_if_close(loss[-1]), 'norm':ncp.real_if_close(ncp.linalg.norm(state))})
            used_ram_tot[m, n, i, 2] = memory.used_bytes()
            del h_t
            used_ram_tot[m, n, i, 3] = memory.used_bytes()
            p_pbar.update()
        del state
        p_pbar.refresh()
        p_pbar.reset()
        data_pbar.update()
        memory.free_all_blocks()
    data_pbar.refresh()
    data_pbar.reset()

used_ram_tot = np.load('perf_ram.npy')
times = np.load('perf_time.npy')

FONTSIZE=12
fig, (ax_ram, ax_time) = plt.subplots(2, 1, figsize=(8, 7), tight_layout=True, sharex=True)

ram_cmap = ax_ram.imshow(used_ram_tot[:,:,:,2].mean(axis=2).T/1e3, cmap='plasma', norm='log', origin='lower')
time_cmap = ax_time.imshow((times[:,:,:-1].mean(axis=2) + times[:,:,-1]).T/1e3, cmap='viridis', norm='log', origin='lower')


ax_ram.set_ylabel('#patterns', fontsize=FONTSIZE+2)
ax_ram.tick_params(axis='y', which='major', labelsize=FONTSIZE)
ax_ram.set_yticks(range(0,len(N_datas),2), labels=N_datas[::2])
ax_ram.set_title(r'RAM usage [$KB$]', fontsize=FONTSIZE+4)

ax_time.set_xlabel('#qbits', fontsize=FONTSIZE+2)
ax_time.set_ylabel('#patterns', fontsize=FONTSIZE+2)
ax_time.tick_params(axis='both', which='major', labelsize=FONTSIZE)
ax_time.set_xticks(range(0,len(N_feats),2), labels=N_feats[::2])
ax_time.set_yticks(range(0,len(N_datas),2), labels=N_datas[::2])
ax_time.set_title(r'Execution time [$ms$]', fontsize=FONTSIZE+4)

fig.colorbar(ram_cmap, ax=ax_ram)
fig.colorbar(time_cmap, ax=ax_time)
fig.suptitle('Exact diagonalization performances', fontsize=FONTSIZE+6, x=0.5845)

fig.savefig('ed_performances.svg')

plt.imshow(times[:,:,-1].T, cmap='plasma', origin='lower')

plt.plot(((ncp.array(range(P))+1)/P).get(), ncp.asarray(loss).get())
plt.yscale('log')
plt.ylabel(r'$\epsilon(s)$')
plt.xlabel(r'$s$')
plt.grid(alpha=0.2)

FONTSIZE=12
fig, ax = plt.subplots(figsize=(5, 4), tight_layout=True)

handle1 = ax.plot(N_feats, used_ram_tot[:,:,2].mean(axis=1)/1e3, label='RAM', lw=2)
ax.set_yscale('log')
ax.set_xlabel('#qbits', fontsize=FONTSIZE+2)
ax.set_xticks(N_feats[1::2])
ax.tick_params(axis='both', which='major', labelsize=FONTSIZE)
ax.set_ylabel(r'RAM usage [$KB$]', fontsize=FONTSIZE+2)
ax2 = ax.twinx()
handle2 = ax2.plot(N_feats, times[:,:-1].mean(axis=1)/1e3, c='tab:orange', label='Time', lw=2)
ax2.set_yscale('log')
ax2.tick_params(axis='y', which='major', labelsize=FONTSIZE)
ax2.set_ylabel(r'Execution time [$ms$]', fontsize=FONTSIZE+2)
ax.legend(handles=handle1+handle2, fontsize=FONTSIZE+2)

ax.set_title('Exact diagonalization performances', fontsize=FONTSIZE+4)
ax.grid(alpha=0.3)
ax2.grid(alpha=0.3)

plt.plot(N_feats, times[:,-1])

FONTSIZE=12
fig, ax = plt.subplots(figsize=(5, 4), tight_layout=True)

handle1 = ax.plot(N_feats, used_ram_tot[:,:,2].mean(axis=1)/1e3, label='RAM', lw=2)
ax.set_yscale('log')
ax.set_xlabel('#qbits', fontsize=FONTSIZE+2)
ax.set_xticks(N_feats[1::2])
ax.tick_params(axis='both', which='major', labelsize=FONTSIZE)
ax.set_ylabel(r'RAM usage [$KB$]', fontsize=FONTSIZE+2)
ax2 = ax.twinx()
handle2 = ax2.plot(N_feats, times[:,:-1].mean(axis=1)/1e3, c='tab:orange', label='Time', lw=2)
ax2.set_yscale('log')
ax2.tick_params(axis='y', which='major', labelsize=FONTSIZE)
ax2.set_ylabel(r'Execution time [$ms$]', fontsize=FONTSIZE+2)
ax.legend(handles=handle1+handle2, fontsize=FONTSIZE+2)

ax.set_title('Exact diagonalization performances', fontsize=FONTSIZE+4)
ax.grid(alpha=0.3)
ax2.grid(alpha=0.3)

plt.plot(N_feats, times[:,-1])

def H_perc2(state, data, labels):
    n_data, n = data.shape
    h_perc = tf.zeros((2**n, 2**n), dtype='complex128')
    
    id      = tfl.LinearOperatorIdentity(2) 
    sigma_z = tfl.LinearOperatorFullMatrix(tf.constant([[1., 0.], [0., -1.]]))

    for i in range(n_data):
        datum_ops = [tfl.LinearOperatorDiag(tf.repeat(data[i,:], [2], axis=0)[j:j+2]) for j in range(0, n*2, 2)]
        ops = [tfl.LinearOperatorKronecker([id]*i+[tfl.LinearOperatorComposition([sigma_z, datum_ops[i]])]+[id]*(n-i-1)).to_dense() for j in range(n)]
        result_op = tf.cast(-tf.reshape(labels[i], (1, 1))*tf.reduce_sum(tf.stack(ops), axis=0), 'complex128')
        if tf.cast(tf.tensordot(tf.math.conj(state), tfl.matvec(result_op, state), axes=1), 'float32') < 0.:
            continue
        else:
            h_perc = h_perc + result_op

    return h_perc

