import os
import sys
import time
import math
import threading


import cudaq
import numpy as np
import cupy as cp
from scipy import interpolate

from pathlib import Path

from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter

##########################################################################################################

# INPUTS HERE

num_reg_qubits=8   # Total number of qubits is 2*num_reg_qubits+3;

T=1000             # Number of timesteps to run;

video_length=T       # Length of the video you eventually wish to render in seconds;
fps=0.1              # fps of that video. This and video_length will tell this program how many
                     # timesteps (frames) need to be written to file;
frames=int(fps*video_length)

folder_name=None     # Name of folder to save results in. If None, it will create a name based
                     # on the other inputs specified.

downsampling_factor=2**5

# Initial state as a function of x and y
initial_state_function = lambda x,y : np.sin(x*2*np.pi/N)*(1-0.5*x/N)*np.sin(y*4*np.pi/N)*(1-0.5*y/N)+1

##########################################################################################################

cudaq.set_target('nvidia', option='mgpu,fp64')

@cudaq.kernel
def alloc_kernel(num_qubits:int):
    qubits = cudaq.qvector(num_qubits)

from cupy.cuda.memory import MemoryPointer, UnownedMemory

def to_cupy_array(state):
    tensor = state.getTensor()
    pDevice = tensor.data()
    sizeByte = tensor.get_num_elements() * tensor.get_element_size()
    mem = UnownedMemory(pDevice, sizeByte, owner=state)
    memptr = MemoryPointer(mem, 0)
    cupy_array = cp.ndarray(tensor.get_num_elements(),
                            dtype=cp.complex128,
                            memptr=memptr)
    return cupy_array

num_anc=3
num_qubits=2*num_reg_qubits+num_anc

N=2**num_reg_qubits
N_tot=2**num_qubits
num_ranks=1
rank=0
N_sub=int(N_tot//num_ranks)

timesteps_per_frame=1
if frames<T:
    timesteps_per_frame=int(T/frames)

if folder_name is None:
    folder_name="d2q5_nq_"+str(N)+"x"+str(N)+"_T_"+str(T)+"_frames_"+str(frames)

Path("Results/"+folder_name+"/").mkdir(parents=True, exist_ok=True)

class QLBMAdvecDiffD2Q5_new:
    def __init__(self,ux=0.2,uy=0.15) -> None:
        self.dim = 2
        self.ndir = 5
        self.nq_dir = math.ceil(np.log2(self.ndir))
        self.dirs=[] 
        for dir_int in range(self.ndir):
            dir_bin = f"{dir_int:b}".zfill(self.nq_dir)
            self.dirs.append(dir_bin)
        self.e_unitvec = np.array([0, 1, -1, 1, -1])
        self.wts = np.array([2/6, 1/6, 1/6, 1/6, 1/6])
        self.cs = 1/np.sqrt(3)
        self.ux = ux
        self.uy = uy
        self.u = np.array([0, self.ux, self.ux, self.uy, self.uy])
        self.wtcoeffs = np.multiply(self.wts, 1+self.e_unitvec*self.u/self.cs**2)
        self.lambdas = np.arccos(self.wtcoeffs)
        self.create_circuit()
  
    def create_circuit(self):
        v=np.pad(self.wtcoeffs,(0,2**num_anc - self.ndir))
        v=v**0.5
        v[0]+=1
        v=v/np.linalg.norm(v)
        U_prep=2*np.outer(v,v)-np.eye(len(v))
        cudaq.register_operation("prep_op", U_prep)

        def collisionOp(dirs):
            dirs_i_list=[]
            for dir_ in dirs:
                dirs_i=[(int(c)) for c in dir_]
                dirs_i_list+=dirs_i[::-1]
            return dirs_i_list

        self.dirs_i_list=collisionOp(self.dirs)

        @cudaq.kernel
        def rshift(q: cudaq.qview, n: int):
            for i in range(n):
                if i == n-1:
                    x(q[n-1-i])
                elif i == n-2:
                    x.ctrl(q[n-1-(i+1)], q[n-1-i])
                else:
                    x.ctrl(q[0:n-1-i], q[n-1-i])

        @cudaq.kernel
        def lshift(q: cudaq.qview, n: int):
            for i in range(n):
                if i == 0:
                    x(q[0])
                elif i == 1:
                    x.ctrl(q[0], q[1])
                else:
                    x.ctrl(q[0:i], q[i])

        @cudaq.kernel
        def d2q5_tstep(q:cudaq.qview,nqx:int,nqy:int,nq_dir:int,dirs_i:list[int]):
            qx=q[0:nqx]
            qy=q[nqx:nqx+nqy]
            qdir=q[nqx+nqy:nqx+nqy+nq_dir]
            i=2
            b_list=dirs_i[i*nq_dir:(i+1)*nq_dir]
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            cudaq.control(lshift,qdir,qx,nqx)
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            i=1
            b_list=dirs_i[i*nq_dir:(i+1)*nq_dir]
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            cudaq.control(rshift,qdir,qx,nqx)
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            i=4
            b_list=dirs_i[i*nq_dir:(i+1)*nq_dir]
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            cudaq.control(lshift,qdir,qy,nqy)
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            i=3
            b_list=dirs_i[i*nq_dir:(i+1)*nq_dir]
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])
            cudaq.control(rshift,qdir,qy,nqy)
            for j in range(nq_dir):
                b=b_list[j]
                if b==0:
                    x(qdir[j])

        @cudaq.kernel
        def d2q5_tstep_wrapper(state: cudaq.State,nqx:int,nqy:int,nq_dir:int,dirs_i:list[int]):
            q=cudaq.qvector(state)
            qdir=q[nqx+nqy:nqx+nqy+nq_dir]
            prep_op(qdir[2],qdir[1],qdir[0])
            d2q5_tstep(q,nqx,nqy,nq_dir,dirs_i)
            prep_op(qdir[2],qdir[1],qdir[0])
    
        @cudaq.kernel
        def d2q5_tstep_wrapper_hadamard(vec:list[complex],nqx:int,nqy:int,nq_dir:int,dirs_i:list[int]):
            q=cudaq.qvector(vec)
            qdir=q[nqx+nqy:nqx+nqy+nq_dir]
            qy=q[nqx:nqx+nqy]
            prep_op(qdir[2],qdir[1],qdir[0])
            d2q5_tstep(q,nqx,nqy,nq_dir,dirs_i)
            prep_op(qdir[2],qdir[1],qdir[0])
            for i in range(nqy):
                h(qy[i])

        def run_timestep(vec,hadamard=False):
            if hadamard:
                result=cudaq.get_state(d2q5_tstep_wrapper_hadamard,vec,num_reg_qubits,num_reg_qubits,self.nq_dir,self.dirs_i_list)
            else:
                result=cudaq.get_state(d2q5_tstep_wrapper,vec,num_reg_qubits,num_reg_qubits,self.nq_dir,self.dirs_i_list)
            num_nonzero_ranks=num_ranks/(2**num_anc)
            rank_slice = to_cupy_array(result)
            len_sub_sv=int(N_tot/num_ranks)
            if rank>=num_nonzero_ranks:
                sub_sv=np.zeros(len_sub_sv,dtype=np.complex128)
                cp.cuda.runtime.memcpy(rank_slice.data.ptr, sub_sv.ctypes.data, sub_sv.nbytes, cp.cuda.runtime.memcpyHostToDevice)
            if rank==0 and num_nonzero_ranks<1:
                sub_sv=(rank_slice).get()
                sub_sv[int(N_tot/(2**num_anc)):]=0
                cp.cuda.runtime.memcpy(rank_slice.data.ptr, sub_sv.ctypes.data, sub_sv.nbytes, cp.cuda.runtime.memcpyHostToDevice)
            return result

        self.run_timestep = run_timestep

    def write_state(self,state,t):
        rank_slice = to_cupy_array(state)
        num_nonzero_ranks=num_ranks/(2**num_anc)
        num_ranks_per_row=num_ranks/N
        if int(num_ranks_per_row*downsampling_factor)>0:
            if rank%int(num_ranks_per_row*downsampling_factor)>=num_ranks_per_row:
                return
        if rank<num_nonzero_ranks:
            with open('Results/'+folder_name+'/'+str(t)+'_'+str(rank)+'.npy','wb') as f:
                if num_nonzero_ranks<1:
                    arr=cp.real(rank_slice)[:int(N_tot/(2**num_anc))]
                else:
                    arr=cp.real(rank_slice)
                if len(arr)>N:
                    arr=arr.reshape((-1,N))
                    arr=arr[::downsampling_factor,::downsampling_factor]
                    arr=arr.flatten()
                else:
                    arr=arr[::downsampling_factor]
                np.save(f,arr)

    def run_evolution(self,initial_state,timesteps,observable=False):
        state=initial_state
        for t in range(timesteps):
            if t==timesteps-1 and observable:
                next_state=self.run_timestep(state,True)
                self.write_state(next_state,str(t+1)+"_h")
            else:
                next_state=self.run_timestep(state)
                if (t+1)%timesteps_per_frame==0 and (timesteps-t)>timesteps_per_frame:
                    self.write_state(next_state,t+1)
                    if rank==0:
                        print("Timestep: ",t+1)
                cp.get_default_memory_pool().free_all_blocks()
            state=next_state
        self.write_state(next_state,timesteps)
        if rank==0:
            print("Timestep: ",timesteps)
        cp.get_default_memory_pool().free_all_blocks()
        state=next_state
        self.final_state=state

from datetime import timedelta
import time
starttime = time.perf_counter()

qlbm_obj=QLBMAdvecDiffD2Q5_new()

print(num_qubits, num_reg_qubits)

initial_state = cudaq.get_state(alloc_kernel,num_qubits)
rank_slice = to_cupy_array(initial_state)

initial_state_function_=lambda x,y : initial_state_function(x,y)*(y<N).astype(int)
xv=np.arange((rank*N_sub)%N,((rank+1)*N_sub-1)%N + 1)
if N_sub>N:
    yv=np.arange(int((rank*N_sub)//N),int(((rank+1)*N_sub)//N))
else:
    yv=np.array([int((rank*N_sub)//N)])

print('Start initializing state')
sub_sv=initial_state_function_(*np.meshgrid(xv,yv)).flatten().astype(np.complex128)
print('Rank ',rank,': Sub-state initialized')

cp.cuda.runtime.memcpy(rank_slice.data.ptr, sub_sv.ctypes.data, sub_sv.nbytes, cp.cuda.runtime.memcpyHostToDevice)

qlbm_obj.run_evolution(initial_state,T)

# Animation Generation
print("Simulation complete. Generating animation...")

# Calculate downsampled grid size
downsampled_N = 2**num_reg_qubits // downsampling_factor  # 256 // 32 = 8
frames_to_plot = list(range(10, T + 1, 10))  # Timesteps: 10, 20, ..., 1000

# Preload data from saved .npy files
data = []
for i in frames_to_plot:
    try:
        sol = np.load(f"Results/{folder_name}/{i}_0.npy")
        Z = np.reshape(sol, (downsampled_N, downsampled_N))
        data.append(Z)
    except FileNotFoundError:
        print(f"Warning: File not found for timestep {i}. Skipping this frame.")
        continue

if not data:
    print("Error: No data available to create animation.")
    sys.exit(1)

# Set up the figure and axes
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Create meshgrid for plotting
x = np.linspace(-10, 10, downsampled_N)
y = np.linspace(-10, 10, downsampled_N)
X, Y = np.meshgrid(x, y)

# Compute normalization factor from the first frame
norm_factor = np.linalg.norm(data[0])

# Define update function for animation
def update_frame(frame_idx):
    ax.clear()
    Z = data[frame_idx]
    surf = ax.plot_surface(X, Y, Z / norm_factor, cmap='viridis', linewidth=0, antialiased=False)
    ax.set_title('Quantum Simulation Evolution')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlim(0, 0.6)
    return surf,

# Generate the animation
ani = FuncAnimation(fig, update_frame, frames=len(data), blit=False)

# Save the animation as a GIF
gif_path = f"Results/{folder_name}/animation.gif"

# output_dir = "/app/results"  # Directory inside the container
# os.makedirs(output_dir, exist_ok=True)  # Ensure the directory exists
# gif_path = os.path.join(output_dir, "animation.gif")
# ani.save(gif_path, writer='pillow', fps=10)
# print(f"Animation saved to {gif_path}")

# Save the animation as a GIF
output_dir = f"Results/{folder_name}"  # Use the results folder in the workspace
os.makedirs(output_dir, exist_ok=True)  # Ensure the directory exists
gif_path = os.path.join(output_dir, "animation.gif")
ani.save(gif_path, writer='pillow', fps=10)
print(f"Animation saved to {gif_path}")

# Clean up
plt.close(fig)