import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

def mapInit(M, N, p_w, p_m, p_h, p_v, p_ii):
    """
    Initialize a 2D map for virus spread simulation.
    
    Parameters:
    M, N: dimensions of the map
    p_w: probability of placing water blocks
    p_m: probability of placing medical facilities  
    p_h: probability of placing humans
    p_v: probability of vaccinating humans
    p_ii: probability of infecting humans
    
    Returns:
    numpy array representing the initialized map
    """
    # Create empty map filled with zeros (empty cells)
    map_grid = np.zeros((M, N), dtype=int)
    
    # Place water blocks with probability p_w
    for i in range(M):
        for j in range(N):
            if np.random.random() < p_w:
                map_grid[i, j] = 1  # water block
    
    # Place medical facilities on remaining land with probability p_m
    for i in range(M):
        for j in range(N):
            if map_grid[i, j] == 0 and np.random.random() < p_m:
                map_grid[i, j] = 2  # medical facility
    
    # Place humans on remaining land with probability p_h
    for i in range(M):
        for j in range(N):
            if map_grid[i, j] == 0 and np.random.random() < p_h:
                # Place healthy human (5) initially
                map_grid[i, j] = 5
    
    # Vaccinate humans with probability p_v
    for i in range(M):
        for j in range(N):
            if map_grid[i, j] == 5 and np.random.random() < p_v:
                # Vaccinate the human - change to vaccinated state (6)
                map_grid[i, j] = 6
    
    # Infect humans with probability p_ii (both vaccinated and unvaccinated)
    for i in range(M):
        for j in range(N):
            if (map_grid[i, j] == 5 or map_grid[i, j] == 6) and np.random.random() < p_ii:
                # Infect the human - change to infected state (10)
                map_grid[i, j] = 10
    
    return map_grid

def map2txt(map_grid):
    """
    Print map as text using specified symbols.
    
    Symbols:
    - w for water
    - M for medical facility
    - . for empty space
    - h for healthy human
    - H for vaccinated human
    - i for infected human
    - x for dead human
    """
    symbol_map = {
        0: '.',  # empty space
        1: 'w',  # water
        2: 'M',  # medical facility
        5: 'h',  # healthy human
        6: 'H',  # vaccinated human
        10: 'i', # infected human
        20: 'x'  # dead human
    }
    
    for i in range(map_grid.shape[0]):
        row = ""
        for j in range(map_grid.shape[1]):
            row += symbol_map.get(map_grid[i, j], '?')
        print(row)

def map2fancy_image(map_grid, filename=None):
    """
    Show map as an image using specified colors.
    
    Colors:
    - blue for water
    - red bold + for medical facility
    - white for empty space
    - black circle for healthy human
    - green circle  for vaccinated human
    - red circle for infected human
    - black circle with cross inside for dead human
    """
    # Create figure and axis
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    
    # Set up the plot with equal aspect ratio
    ax.set_xlim(0, map_grid.shape[1])
    ax.set_ylim(0, map_grid.shape[0])
    ax.set_aspect('equal')  # This ensures 1:1 aspect ratio
    
    # Define colors for each cell type
    for i in range(map_grid.shape[0]):
        for j in range(map_grid.shape[1]):
            cell_value = map_grid[i, j]
            
            if cell_value == 0:  # Empty space
                rect = patches.Rectangle((j, map_grid.shape[0]-1-i), 1, 1, 
                                       facecolor='white', edgecolor='black', linewidth=0.1)
                ax.add_patch(rect)
            elif cell_value == 1:  # Water
                rect = patches.Rectangle((j, map_grid.shape[0]-1-i), 1, 1, 
                                       facecolor='blue', edgecolor='black', linewidth=0.1)
                ax.add_patch(rect)
            elif cell_value == 2:  # Medical facility
                rect = patches.Rectangle((j, map_grid.shape[0]-1-i), 1, 1, 
                                       facecolor='white', edgecolor='red', linewidth=0.1)
                ax.add_patch(rect)
                # Add red plus sign
                ax.plot([j+0.3, j+0.7], [map_grid.shape[0]-1-i+0.5, map_grid.shape[0]-1-i+0.5], 
                       color='red', linewidth=2)
                ax.plot([j+0.5, j+0.5], [map_grid.shape[0]-1-i+0.3, map_grid.shape[0]-1-i+0.7], 
                       color='red', linewidth=2)
            elif cell_value == 5:  # Healthy human
                circle = patches.Circle((j+0.5, map_grid.shape[0]-1-i+0.5), 0.3, 
                                       facecolor='black', edgecolor='black')
                ax.add_patch(circle)
            elif cell_value == 6:  # Vaccinated human
                circle = patches.Circle((j+0.5, map_grid.shape[0]-1-i+0.5), 0.3, 
                                       facecolor='green', edgecolor='black')
                ax.add_patch(circle)
            elif cell_value == 10:  # Infected human
                circle = patches.Circle((j+0.5, map_grid.shape[0]-1-i+0.5), 0.3, 
                                       facecolor='red', edgecolor='black')
                ax.add_patch(circle)
            elif cell_value == 20:  # Dead human
                circle = patches.Circle((j+0.5, map_grid.shape[0]-1-i+0.5), 0.3, 
                                       facecolor='black', edgecolor='black')
                ax.add_patch(circle)
                # Add cross inside
                ax.plot([j+0.3, j+0.7], [map_grid.shape[0]-1-i+0.3, map_grid.shape[0]-1-i+0.7], 
                       color='white', linewidth=1)
                ax.plot([j+0.3, j+0.7], [map_grid.shape[0]-1-i+0.7, map_grid.shape[0]-1-i+0.3], 
                       color='white', linewidth=1)
    
    # Remove axes
    ax.set_xticks([])
    ax.set_yticks([])
    
    # Save if filename provided, otherwise show
    if filename:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
    else:
        plt.show()

def map2colormap(map_grid, filename=None):
    """
    Show map as a colored image using specified color scheme.
    
    Colors:
    - blue for water
    - magenta for medical facility
    - white for empty space
    - grey for healthy human
    - green for vaccinated human
    - red for infected human
    - black for dead human
    """
    # Create a color map based on cell values
    color_map = np.zeros((map_grid.shape[0], map_grid.shape[1], 3))
    
    # Define color mappings (RGB values)
    colors = {
        0: [1, 1, 1],      # white for empty space
        1: [0, 0, 1],      # blue for water
        2: [1, 0, 1],      # magenta for medical facility (FIXED)
        5: [0.5, 0.5, 0.5], # grey for healthy human
        6: [0, 1, 0],      # green for vaccinated human
        10: [1, 0, 0],     # red for infected human
        20: [0, 0, 0]      # black for dead human
    }
    
    # Fill the color map
    for i in range(map_grid.shape[0]):
        for j in range(map_grid.shape[1]):
            cell_value = map_grid[i, j]
            if cell_value in colors:
                color_map[i, j] = colors[cell_value]
            else:
                color_map[i, j] = [0, 0, 0]  # default to black for unknown values
    
    # Create figure and axis
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    
    # Display the color map
    ax.imshow(color_map, aspect='equal')
    
    # Remove axes
    ax.set_xticks([])
    ax.set_yticks([])
    
    # Save if filename provided, otherwise show
    if filename:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
    else:
        plt.show()

def virusMove(map_grid, p_d=0.1, p_r=0.2, p_i=0.3, d_e=2.0):
    """
    Apply one step of virus spread to the map.
    
    Rules:
    - Every infected human can become dead with probability p_d
    - Every infected human can recover with probability p_r and become immune
    - Every infected human can infect neighbors at distance d according to 
      probability p(d) = p_i * exp(-d^2 / d_e^2)
    
    Parameters:
    map_grid: current map state
    p_d: probability of infected human becoming dead
    p_r: probability of infected human to recover and become immune
    p_i: probability scale for infection
    d_e: effective distance of infection
    
    Returns:
    Updated map after virus spread step
    """
    # Create a copy of the map to avoid modifying during iteration
    new_map = map_grid.copy()
    
    # Get map dimensions
    M, N = map_grid.shape
    
    # Process each cell
    for i in range(M):
        for j in range(N):
            current_cell = map_grid[i, j]
            
            # Rule 1: Infected humans can become dead
            if current_cell == 10:  # Infected human
                if np.random.random() < p_d:
                    new_map[i, j] = 20  # Become dead
            
            # Rule 2: Infected humans can recover
            if current_cell == 10:  # Infected human
                if np.random.random() < p_d:
                    new_map[i, j] = 6  # Become health and immune
            
            # Rule 3: Infected humans can infect neighbors
            if current_cell == 10:  # Only infected humans can spread
                # Check all neighboring cells (including diagonals)
                for di in [-1, 0, 1]:
                    for dj in [-1, 0, 1]:
                        if di == 0 and dj == 0:
                            continue  # Skip the cell itself
                        
                        ni, nj = i + di, j + dj
                        
                        # Check if neighbor is within bounds
                        if 0 <= ni < M and 0 <= nj < N:
                            neighbor_cell = map_grid[ni, nj]
                            
                            # Can only infect healthy or vaccinated humans
                            if neighbor_cell == 5 or neighbor_cell == 6:
                                # Calculate distance
                                d = np.sqrt(di**2 + dj**2)
                                
                                # Calculate infection probability
                                prob = p_i * np.exp(-d**2 / d_e**2)
                                
                                # Try to infect
                                if np.random.random() < prob:
                                    # Infect the neighbor (change to infected)
                                    new_map[ni, nj] = 10
    
    return new_map