import numpy as np
import matplotlib.pyplot as plt
import virus_model
import visualizer_3d


def map_init():
    """Test the map initialization function"""
    print("=== Testing mapInit ===")
    M, N = 100, 200
    p_w = 0.01
    p_m = 0.05
    p_h = 0.8
    p_v = 0.02
    p_ii = 0.01
    
    result = virus_model.mapInit(M, N, p_w, p_m, p_h, p_v, p_ii)
    print(f"Map shape: {result.shape}")
    
    # Count different cell types
    unique, counts = np.unique(result, return_counts=True)
    print("Initial cell type counts:")
    for u, c in zip(unique, counts):
        if u == 0:
            print(f"  Empty cells: {c}")
        elif u == 1:
            print(f"  Water blocks: {c}")
        elif u == 2:
            print(f"  Medical facilities: {c}")
        elif u == 5:
            print(f"  Healthy humans: {c}")
        elif u == 6:
            print(f"  Vaccinated humans: {c}")
        elif u == 10:
            print(f"  Infected humans: {c}")
    
    return result

map0 = map_init()
[Nr, Nc] = map0.shape
Ns = 1000 # number of evolution steps

evolution_map = np.zeros( (Ns, Nr, Nc), dtype = int)

evolution_map[0,:,:] = np.copy(map0)

# virus property
p_d=0.03
p_i=0.03
d_e=2.0

for s in range(1, Ns):
    evolution_map[s,:,:] = virus_model.virusMove(evolution_map[s-1,:,:], p_d=p_d, p_i=p_i, d_e=d_e)

# Save the evolution map as a numpy array at the end
np.save('evolution_map.npy', evolution_map)
print("Evolution map saved as 'evolution_map.npy'")

# visualizer_3d.visualize_3d_array(evolution_map)