#!/usr/bin/python3

import numpy as np
import matplotlib
import os

matplotlib.use("Agg")

import matplotlib.pyplot as plt

def process(fname):                        # read binary file
  with open(f'{fname}.bin', "rb") as f:
  
      width = np.fromfile(
          f, dtype=np.int32, count=1
      )[0]
  
      height = np.fromfile(
          f, dtype=np.int32, count=1
      )[0]
  
      max_iter = np.fromfile(
          f, dtype=np.int32, count=1
      )[0]
  
      data = np.fromfile(
          f,
          dtype=np.uint16,
          count=width * height
      )
  
  image = data.reshape((height, width))
  
  # Mathematical Mandelbrot coordinates
  
  xmin = -2.5
  xmax =  1.0
  
  ymin = -1.25
  ymax =  1.25
  
  # Select exactly the SAME mathematical zoom region
  # for every resolution
  
  zoom_xmin = -0.8
  zoom_xmax = -0.6
  
  zoom_ymin =  0.0
  zoom_ymax =  0.2
  
  # Convert mathematical coordinates to pixel indices
  
  x0 = int(
      (zoom_xmin - xmin)
      / (xmax - xmin)
      * width
  )
  
  x1 = int(
      (zoom_xmax - xmin)
      / (xmax - xmin)
      * width
  )
  
  y0 = int(
      (zoom_ymin - ymin)
      / (ymax - ymin)
      * height
  )
  
  y1 = int(
      (zoom_ymax - ymin)
      / (ymax - ymin)
      * height
  )
  
  zoom = image[y0:y1, x0:x1]
  
  # Plot
  
  plt.figure(figsize=(10, 7))
  
  plt.imshow(
      zoom,
      origin="lower",
  
      # Tell matplotlib the actual mathematical coordinates
      extent=[
          zoom_xmin,
          zoom_xmax,
          zoom_ymin,
          zoom_ymax
      ],
  
      cmap="magma",
      interpolation="nearest",
  
      aspect="equal"
  )
  
  plt.colorbar(
      label="Iterations"
  )
  
  plt.xlabel("Real part (x)")
  plt.ylabel("Imaginary part (y)")
  
  plt.title(
      f"Mandelbrot zoom\n"
      f"Resolution: {width} × {height}, "
      f"maximum iterations: {max_iter}"
  )
  
  plt.tight_layout()
  
  plt.savefig(
      f"{fname}.png",
      dpi=200
  )
  
  plt.close()
  
  print("Mandelbrot visualization")
  print("------------------------")
  print("Resolution :", width, "x", height)
  print("Iterations :", max_iter)
  
  print(
      "Displayed region:",
      f"x = [{zoom_xmin}, {zoom_xmax}],",
      f"y = [{zoom_ymin}, {zoom_ymax}]"
  )
  
  print("Zoom array   :", zoom.shape)
  print(f"Output       : {fname}.png")

for filename in os.listdir("."):
    if filename.endswith(".bin"):
        fname = filename.removesuffix(".bin")
        process(fname)