# Experiment E2: Image generation with a genetic algorithm (numpy + matplotlib, offline)
# Requires: numpy, matplotlib | no dataset, no internet
# ===========================================================================

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)

# 1. Target image defined mathematically (no dataset): an 8x8 cross
SIZE = 8
target = np.zeros((SIZE, SIZE), dtype=int)
target[SIZE // 2 - 1:SIZE // 2 + 1, :] = 1      # horizontal bar
target[:, SIZE // 2 - 1:SIZE // 2 + 1] = 1      # vertical bar
# (for a square instead: target = np.zeros((SIZE, SIZE), int); target[2:6, 2:6] = 1)

# 2. Genetic-algorithm settings
POP_SIZE = 60          # number of images in each generation
MUTATION_RATE = 0.02   # probability of flipping each pixel
GENERATIONS = 300
ELITE = 2              # best images copied unchanged to the next generation

# 3. Fitness = fraction of pixels that match the target (1.0 = perfect)
def fitness(img):
    return np.mean(img == target)

# 4. Initial population: completely random black/white images
population = np.random.randint(0, 2, size=(POP_SIZE, SIZE, SIZE))

def select_parent(population, scores):
    # tournament selection: pick 3 random images, keep the fittest
    idx = np.random.choice(len(population), 3, replace=False)
    return population[idx[np.argmax(scores[idx])]]

def crossover(p1, p2):
    # uniform crossover: each pixel comes from parent 1 or parent 2
    mask = np.random.rand(SIZE, SIZE) < 0.5
    return np.where(mask, p1, p2)

def mutate(img):
    flip = np.random.rand(SIZE, SIZE) < MUTATION_RATE
    return np.where(flip, 1 - img, img)

best_history, mean_history, snapshots = [], [], {}
for gen in range(GENERATIONS):
    scores = np.array([fitness(ind) for ind in population])
    best = population[np.argmax(scores)]
    best_history.append(scores.max()); mean_history.append(scores.mean())
    if gen in (0, 5, 15, 30) or scores.max() == 1.0:
        snapshots[gen] = best.copy()
    if gen % 20 == 0:
        print(f"generation {gen:3d}  best fitness {scores.max():.3f}  mean {scores.mean():.3f}")
    if scores.max() == 1.0:
        print(f"Target reached at generation {gen}")
        break
    # next generation: elitism + selection + crossover + mutation
    order = np.argsort(scores)[::-1]
    new_population = [population[i] for i in order[:ELITE]]
    while len(new_population) < POP_SIZE:
        child = crossover(select_parent(population, scores), select_parent(population, scores))
        new_population.append(mutate(child))
    population = np.array(new_population)

# 5. Results
fig, ax = plt.subplots(1, len(snapshots) + 1, figsize=(2.2 * (len(snapshots) + 1), 2.6))
ax[0].imshow(target, cmap="gray"); ax[0].set_title("target"); ax[0].axis("off")
for a, (g, img) in zip(ax[1:], snapshots.items()):
    a.imshow(img, cmap="gray"); a.set_title(f"gen {g}\nfit {fitness(img):.2f}"); a.axis("off")
plt.tight_layout(); plt.show()

plt.plot(best_history, label="best fitness"); plt.plot(mean_history, label="mean fitness")
plt.xlabel("generation"); plt.ylabel("fitness (fraction of correct pixels)")
plt.title("Genetic algorithm progress"); plt.legend(); plt.grid(alpha=0.3); plt.show()
