import enum
import itertools
import time
from math import floor, ceil

import numba as nb
import numpy as np
import matplotlib
from PIL import Image, PyAccess


# Amount of times to print the total progress
PROGRESS_STEPS: int = 20

# Set to `True` to plot the shortcut version of the fractal
SHORTCUT: bool = True

# Make all integers critical points
FIX_CRITICAL_POINTS: bool = True

# Width of the image in pixels and aspect ratio
RESOLUTION: int = 1920*1080//4
ASPECT_RATIO: float = 21/9 if FIX_CRITICAL_POINTS else 16/9

# Value of the center pixel
CENTER: complex = 0 + 0j

# Value range of the real part (width of the horizontal axis)
RE_RANGE: float = 10 if FIX_CRITICAL_POINTS else 5

# Show grid lines for integer real and imaginary parts
SHOW_GRID: bool = False
GRID_COLOR: tuple[int, int, int] = (125, 125, 125)

# Matplotlib named colormap
COLORMAP_NAME: str = 'inferno'

# Plot range of the axes
X_MIN = CENTER.real - RE_RANGE/2  # min Re(z)
X_MAX = CENTER.real + RE_RANGE/2  # max Re(z)
Y_MIN = CENTER.imag - RE_RANGE/(2*ASPECT_RATIO)  # min Im(z)
Y_MAX = CENTER.imag + RE_RANGE/(2*ASPECT_RATIO)  # max Im(z)

x_range = X_MAX - X_MIN
y_range = Y_MAX - Y_MIN
pixels_per_unit = np.sqrt(RESOLUTION/(x_range*y_range))

# Width and height of the image in pixels
WIDTH = round(pixels_per_unit*x_range)
HEIGHT = round(pixels_per_unit*y_range)


# Maximum iterations for the divergence test (recommended >= 60)
MAX_ITER: int = 60


# Max value of Re(z) and Im(z) for which the recursion doesn't overflow
CUTOFF_RE = 7.564545572282618e+153
CUTOFF_IM = 112.10398935569289 if SHORTCUT else 111.95836403625282

# Smallest positive real fixed point
INNER_FIXED_POINT = 0.277733766171606 if SHORTCUT else 0.150108511304474


# Precompute the colormap
CMAP_LEN: int = 2000
cmap_mpl = matplotlib.colormaps[COLORMAP_NAME]
# Start away from 0 (discard black values for the 'inferno' colormap)
# Matplotlib's colormaps have 256 discrete color points
n_cmap = round(256*0.98)
CMAP = [cmap_mpl(k/256) for k in range(256 - n_cmap, 256)]
# Interpolate
x = np.linspace(0, 1, num=CMAP_LEN)
xp = np.linspace(0, 1, num=n_cmap)
c0, c1, c2 = tuple(np.interp(x, xp, [c[k] for c in CMAP]) for k in range(3))
CMAP = []
for x0, x1, x2 in zip(c0, c1, c2):
    CMAP.append(tuple(round(255*x) for x in (x0, x1, x2)))


class DivType(enum.Enum):
    """Divergence type."""

    CONVERGED = -1  # Converged
    MAX_ITER = 0  # Maximum iterations reached
    CUTOFF_RE = 1  # Diverged by exceeding the real part cutoff
    CUTOFF_IM = 2  # Diverged by exceeding the imaginary part cutoff


@nb.jit(nb.float64(nb.float64, nb.int64), nopython=True)
def smooth(x, k=1):
    """Recursive exponential smoothing function."""

    y = np.expm1(np.pi*x)/np.expm1(np.pi)
    if k <= 1:
        return y
    return smooth(y, np.fmin(6, k - 1))


@nb.jit(nb.float64(nb.float64, nb.float64), nopython=True)
def get_delta(x, cutoff):
    """Get the fractional part of the smoothed divergence count."""

    nu = np.log(np.abs(x)/cutoff)/(np.pi*cutoff - np.log(cutoff))
    nu = np.fmax(0, np.fmin(nu, 1))
    return smooth(1 - nu, 2)


@nb.jit(
    nb.types.containers.Tuple((
        nb.float64,
        nb.types.EnumMember(DivType, nb.int64)
    ))(nb.complex128),
    nopython=True
)
def divergence_count(z):
    """Return a smoothed divergence count and the type of divergence."""

    z_fix = 0 + 0j
    for k in range(MAX_ITER):
        c = np.cos(np.pi*z)
        if SHORTCUT:
            if FIX_CRITICAL_POINTS:
                z_fix = (0.5 - c)*np.sin(np.pi*z)/np.pi
            z = 0.25 + z - (0.25 + 0.5*z)*c + z_fix
        else:  # Regular
            if FIX_CRITICAL_POINTS:
                z_fix = (1.25 - 1.75*c)*np.sin(np.pi*z)/np.pi
            z = 0.5 + 1.75*z - (0.5 + 1.25*z)*c + z_fix

        if np.abs(z.imag) > CUTOFF_IM:
            # Diverged by exceeding the imaginary part cutoff
            return k + get_delta(z.imag, CUTOFF_IM), DivType.CUTOFF_IM
        if np.abs(z.real) > CUTOFF_RE:
            # Diverged by exceeding the real part cutoff
            return k + get_delta(z.real, CUTOFF_RE), DivType.CUTOFF_RE
        if np.abs(z) < INNER_FIXED_POINT:
            # Converged to a fixed point
            return -1, DivType.CONVERGED

    # Maximum iterations reached
    return -1, DivType.MAX_ITER


@nb.jit(nb.float64(nb.float64), nopython=True)
def cyclic_map(g):
    """A continuous function that cycles back and forth from 0 to 1."""

    # This can be any continuous function.
    # Log scale removes high-frequency color cycles.
    freq_div = 12
    g = np.log1p(np.fmax(0, (g - 1)/freq_div))

    # Beyond this value for float64, decimals are truncated
    if g >= 2**51:
        return -1

    # Normalize and cycle
    # g += 0.5  # phase from 0 to 1
    return 1 - np.abs(2*(g - np.floor(g)) - 1)


@nb.jit(nb.complex128(nb.types.containers.UniTuple(nb.float64, 2)),
        nopython=True)
def pixel_to_z(p):
    """Convert pixel coordinates to its corresponding complex value."""

    re = X_MIN + (X_MAX - X_MIN)*p[0]/WIDTH
    im = Y_MAX - (Y_MAX - Y_MIN)*p[1]/HEIGHT
    return re + 1j*im


class Progress:
    """Simple progress check helper class."""

    def __init__(self, n: int, steps: int = 10):
        self.n = n
        self.k = 0
        self.steps = steps
        self.step = 1
        self.progress = 0

    def check(self) -> bool:
        self.k += 1
        self.progress = self.k/self.n
        if self.steps*self.k >= self.step*self.n:
            self.step += 1
            return True
        return self.progress == 1


def create_image():
    img = Image.new('RGB', (WIDTH, HEIGHT))
    pix = img.load()
    pix: PyAccess
    n_pix = WIDTH*HEIGHT

    prog = Progress(n_pix, steps=PROGRESS_STEPS)
    for p in itertools.product(range(WIDTH), range(HEIGHT)):
        c = pixel_to_z(p)
        g, div_type = divergence_count(c)
        if g >= 0:
            pix[p] = CMAP[round(cyclic_map(g)*(CMAP_LEN - 1))]
        else:
            # Color of the interior of the fractal
            pix[p] = (0, 0, 0)
        if prog.check():
            print(f'{prog.progress:<7.1%}')

    if SHOW_GRID:
        for x in range(ceil(X_MIN), floor(X_MAX) + 1):
            px = round((x - X_MIN)*(WIDTH - 1)/(X_MAX - X_MIN))
            for py in range(HEIGHT):
                pix[(px, py)] = GRID_COLOR
        for y in range(ceil(Y_MIN), floor(Y_MAX) + 1):
            py = round((Y_MAX - y)*(HEIGHT - 1)/(Y_MAX - Y_MIN))
            for px in range(WIDTH):
                pix[(px, py)] = GRID_COLOR

    return img


img = create_image()
strtime = time.strftime('%Y%m%d-%H%M%S')
fractal_type = 'Shortcut' if SHORTCUT else 'Regular'
filename = f'Collatz_{fractal_type}_{strtime}.png'
img.save(filename)