import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

class LatticeQuantumWalk:
    """
    Unified class for 1D, 2D, and 3D Quantum Walk Simulations on regular lattices.
    """
    def __init__(self, dimension: int, network_size: int = 100):
        if dimension not in [1, 2, 3]:
            raise ValueError("Dimension must be 1, 2, or 3.")
        
        self.dimension = dimension
        self.N = network_size
        self.LSIZE = 2 * self.N + 1
        self.results = None
        self.meta = {}

    def run(self, steps: int = 100, initial_dist: str = 'uniform', compare_classical: bool = False) -> dict:
        self.meta = {
            "steps": steps,
            "initial_dist": initial_dist,
            "compare_classical": compare_classical,
            "network_size": self.N
        }
        
        if self.dimension == 1:
            self.results = self._run_1d(steps, initial_dist, compare_classical)
        elif self.dimension == 2:
            self.results = self._run_2d(steps, initial_dist, compare_classical)
        elif self.dimension == 3:
            self.results = self._run_3d(steps, initial_dist)
        
        return self.results

    def _run_1d(self, steps, initial_dist, compare_classical):
        N = self.N
        LSIZE = self.LSIZE
        COINS = 1
        TOTAL_DIM = LSIZE * COINS * 2
        
        TOSS = np.zeros(shape=(TOTAL_DIM, TOTAL_DIM), dtype=complex)
        STEP = np.zeros(shape=(TOTAL_DIM, TOTAL_DIM), dtype=complex)
        
        for row in range(TOTAL_DIM):
            for col in range(TOTAL_DIM):
                if (row == col) and (row < LSIZE):     TOSS[row, col] = 1/np.sqrt(2)
                elif (row == col) and (row >= LSIZE):  TOSS[row, col] = -1/np.sqrt(2)
                elif (row == col + LSIZE):             TOSS[row, col] = 1/np.sqrt(2)
                elif (col == row + LSIZE):             TOSS[row, col] = 1/np.sqrt(2)
                
                if (row == col + 1) and (row < LSIZE):   STEP[row, col] = 1
                elif (col == row + 1) and (row >= LSIZE): STEP[row, col] = 1
        
        state0 = np.zeros(shape=(TOTAL_DIM, 1), dtype=complex)
        def get_indices(lattice_x):
            base_idx = int(lattice_x + N)
            return [base_idx, base_idx + LSIZE]
        center_indices = get_indices(0)
        
        if initial_dist == 'uniform':
            state0[center_indices[0]] = 1/np.sqrt(2)
            state0[center_indices[1]] = (0+1j)/np.sqrt(2)
        elif initial_dist == 'biased':
            state0[center_indices[0]] = 1.0
        
        current_state = state0
        for t in range(steps):
            current_state = np.dot(TOSS, current_state)
            current_state = np.dot(STEP, current_state)

        lattice_grid = np.arange(-N, N + 1)
        quantum_probs = np.zeros(LSIZE)
        for i, x in enumerate(lattice_grid):
            locs = get_indices(x)
            quantum_probs[i] = (np.abs(current_state[locs[0]])**2 + np.abs(current_state[locs[1]])**2).item()

        classical_probs = None
        if compare_classical:
            walkers = 5000
            positions = np.zeros(walkers)
            for _ in range(steps):
                positions += np.random.choice([-1, 1], size=walkers)
            c_counts, _ = np.histogram(positions, bins=np.arange(-N, N + 2) - 0.5, density=True)
            classical_probs = c_counts

        return {"type": "1D", "grid": lattice_grid, "quantum_probs": quantum_probs, "classical_probs": classical_probs}

    def _run_2d(self, steps, initial_dist, compare_classical):
        N = self.N
        LSIZE = self.LSIZE
        state = np.zeros((2, 2, LSIZE, LSIZE), dtype=complex)
        if initial_dist == 'uniform':
            state[0, 0, N, N], state[0, 1, N, N] = 0.5, 0.5j
            state[1, 0, N, N], state[1, 1, N, N] = -0.5j, 0.5
        else:
            state[0, 0, N, N] = 1.0

        for _ in range(steps):
            s = state
            new_s = np.empty_like(s)
            new_s[0, 0] = 0.5 * (s[0, 0] + s[0, 1] + s[1, 0] + s[1, 1])
            new_s[0, 1] = 0.5 * (s[0, 0] - s[0, 1] + s[1, 0] - s[1, 1])
            new_s[1, 0] = 0.5 * (s[0, 0] + s[0, 1] - s[1, 0] - s[1, 1])
            new_s[1, 1] = 0.5 * (s[0, 0] - s[0, 1] - s[1, 0] + s[1, 1])
            state[0, 0] = np.roll(new_s[0, 0], shift=(1, 1), axis=(0, 1))
            state[0, 1] = np.roll(new_s[0, 1], shift=(1, -1), axis=(0, 1))
            state[1, 0] = np.roll(new_s[1, 0], shift=(-1, 1), axis=(0, 1))
            state[1, 1] = np.roll(new_s[1, 1], shift=(-1, -1), axis=(0, 1))

        q_probs = np.sum(np.abs(state) ** 2, axis=(0, 1))
        c_probs = None
        if compare_classical:
            walkers = 10000
            x, y = np.zeros(walkers, dtype=int), np.zeros(walkers, dtype=int)
            for _ in range(steps):
                x += np.random.choice([-1, 1], size=walkers)
                y += np.random.choice([-1, 1], size=walkers)
            c_probs, _, _ = np.histogram2d(x, y, bins=[np.arange(-N, N + 2) - 0.5] * 2, density=True)
            c_probs = c_probs.T
        return {"type": "2D", "q_data": q_probs, "c_data": c_probs}

    def _run_3d(self, steps, initial_dist):
        N = self.N
        LSIZE = self.LSIZE
        COIN_DIM = 6
        state = np.zeros((COIN_DIM, LSIZE, LSIZE, LSIZE), dtype=np.complex128)
        center = N
        if initial_dist == 'uniform':
            state[:, center, center, center] = 1.0 / np.sqrt(COIN_DIM)
        else:
            state[0, center, center, center] = 1.0
        grover_coin = (2.0 / COIN_DIM) * np.ones((COIN_DIM, COIN_DIM)) - np.eye(COIN_DIM)
        
        for _ in range(steps):
            state = np.einsum('ij,jxyz->ixyz', grover_coin, state)
            new_state = np.zeros_like(state)
            new_state[0] = np.roll(state[0], 1, 0)
            new_state[1] = np.roll(state[1], -1, 0)
            new_state[2] = np.roll(state[2], 1, 1)
            new_state[3] = np.roll(state[3], -1, 1)
            new_state[4] = np.roll(state[4], 1, 2)
            new_state[5] = np.roll(state[5], -1, 2)
            state = new_state
        q_probs = np.sum(np.abs(state) ** 2, axis=0)
        return {"type": "3D", "q_data": q_probs}

    def visualize(self):
        if self.results is None: raise ValueError("No results. Run simulation first.")
        if self.dimension == 1: self._visualize_1d()
        elif self.dimension == 2: self._visualize_2d()
        elif self.dimension == 3: self._visualize_3d()

    def _visualize_1d(self):
        result = self.results
        plt.figure()
        plt.plot(result['grid'], result['quantum_probs'], label='Quantum')
        if result['classical_probs'] is not None:
            plt.plot(result['grid'], result['classical_probs'], label='Classical', linestyle='--')
        plt.legend()
        plt.title(f"1D Quantum Walk (N={self.meta['network_size']}, T={self.meta['steps']})")
        plt.show()

    def _visualize_2d(self):
        result = self.results
        fig, axes = plt.subplots(1, 2 if self.meta['compare_classical'] else 1, figsize=(12, 5))
        ax_q = axes[0] if self.meta['compare_classical'] else axes
        im_q = ax_q.imshow(result['q_data'], cmap='Blues', origin='lower', extent=[-self.N, self.N, -self.N, self.N])
        ax_q.set_title(f"Quantum Walk (Steps={self.meta['steps']})")
        plt.colorbar(im_q, ax=ax_q)
        if self.meta['compare_classical']:
            ax_c = axes[1]
            im_c = ax_c.imshow(result['c_data'], cmap='Blues', origin='lower', extent=[-self.N, self.N, -self.N, self.N])
            ax_c.set_title("Classical Random Walk")
            plt.colorbar(im_c, ax=ax_c)
        plt.tight_layout()
        plt.show()

    def _visualize_3d(self):
        q_probs = self.results['q_data']
        N = self.N
        center = N
        xy = q_probs[:, :, center]
        yz = q_probs[center, :, :]
        xz = q_probs[:, center, :]
        fig, axes = plt.subplots(1, 3, figsize=(18, 5), dpi=100)
        cmap = 'YlGnBu'
        norm = LogNorm(vmin=max(q_probs.min(), 1e-12), vmax=q_probs.max())
        
        for ax, data, title, xlabel, ylabel in zip(axes, [xy, yz, xz], ["XY Slice", "YZ Slice", "XZ Slice"], ["x","y","x"], ["y","z","z"]):
            im = ax.imshow(data.T, cmap=cmap, norm=norm, origin='lower', extent=[-N, N, -N, N])
            ax.set_title(title); ax.set_xlabel(xlabel); ax.set_ylabel(ylabel)
        fig.colorbar(axes[2].images[0], ax=axes.ravel().tolist(), pad=0.02).set_label('Raw Probability Density')
        plt.suptitle(f"3D Quantum Walk Analysis (Steps={self.meta['steps']})")
        plt.show()
