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


# Load sample data from JSON file
def load_sample_data(filename, num_samples=1000):
    with open(filename, "r") as f:
        boards = [json.loads(line) for line in f.readlines()[:num_samples]]
    return boards


# Plot 1: Estimated days to complete the full run with different numbers of processes
def plot_extended_estimated_days_vs_processes():
    processes = [1, 2, 4, 8, 16, 32, 64, 128, 256]
    estimated_days = [166.64, 83.32, 20.83, 10.42, 5.21, 2.6, 1.3, 0.65, 0.325]

    plt.figure(figsize=(10, 6))
    plt.plot(processes, estimated_days, marker="o")
    plt.title("Extended Estimated Days to Complete Full Run vs. Number of Processes")
    plt.xlabel("Number of Processes")
    plt.ylabel("Estimated Days")
    plt.xscale("log", base=2)
    plt.xticks(processes, labels=[str(p) for p in processes])
    plt.grid(True)
    plt.savefig("extended_estimated_days_vs_processes.png")
    plt.close()


# Plot 2: Size of All Battleship Generations
def plot_size_of_all_battleship_generations():
    labels = ["1M Original", "1M Compressed", "30B Original", "30B Compressed"]
    sizes_mb = [517, 19, 15_510_000, 570_000]  # Estimated sizes in MB
    sizes_tb = [size / 1_024_000 for size in sizes_mb]  # Convert MB to TB

    plt.figure(figsize=(10, 6))
    plt.bar(labels, sizes_tb, color=["blue", "green", "blue", "green"])
    plt.title("Size of All Battleship Generations")
    plt.ylabel("Size (TB)")
    plt.ylim(0, 20)
    for i, size in enumerate(sizes_tb):
        plt.text(i, size + 0.5, f"{sizes_mb[i]} MB\n{size:.2f} TB", ha="center")
    plt.savefig("file_size_comparison.png")
    plt.close()


# Plot 3: Distribution of ship placements on the grid
def plot_ship_placement_distribution(sample_file):
    boards = load_sample_data(sample_file)
    grid_size = 100
    ship_counts = np.zeros(grid_size)

    for board in boards:
        for i, cell in enumerate(board):
            if cell:  # If the cell is occupied by a ship
                ship_counts[i] += 1

    plt.figure(figsize=(12, 6))
    plt.bar(range(grid_size), ship_counts, color="purple")
    plt.title("Distribution of Ship Placements on the Grid")
    plt.xlabel("Grid Position")
    plt.ylabel("Number of Ships")
    plt.savefig("ship_placement_distribution.png")
    plt.close()


# Generate the plots
plot_extended_estimated_days_vs_processes()
plot_size_of_all_battleship_generations()
plot_ship_placement_distribution("battleship_boards.json")
