import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt

from lib import read_load, read_xy_alt


def boxcount(data, size):
    np.sign(blockshaped(data, 2).sum(axis=(1, 2))).sum()
    blocked_arrays = blockshaped(data, size)
    counter = 0
    a = blockshaped(data, 2).sum(axis=(1, 2)).signum()
    for i in range(len(blocked_arrays)):
        for j in range(len(blocked_arrays[i])):
            if (blocked_arrays[i][j].any()):
                counter += 1
                break
    return counter


def blockshaped(data, square_size):
    h, w = data.shape
    assert h % square_size == 0, f"Array is not evenly divisible {h} {square_size} {h % square_size}"
    return (data.reshape(h // square_size, square_size, -1, square_size).swapaxes(1, 2).reshape(-1,
                                                                                                square_size,
                                                                                                square_size))


df = read_xy_alt(
    "/Users/joshuacoles/Developer/checkouts/jc3091/CompB DLA/data-analysis/data/alpha/1.csv"
)

max_actual_radius = max(df.x.abs().max(), df.y.abs().max())
max_radius = int(np.exp2(np.floor(np.log2(max_actual_radius)) + 1))
data = np.zeros((max_radius * 2, max_radius * 2))
data[df.x + max_radius, df.y + max_radius] = 1

iterations = int(np.floor(np.log2(max_actual_radius)) + 1)
sizes = 2 ** np.arange(1, iterations)

fb = []
dims = []

for size in sizes:
    filled_box = np.sign(blockshaped(data, size).sum(axis=(1, 2))).sum()
    fb.append(filled_box)
    dims.append(size / np.exp2(iterations))



class ImageFractalDimension:
    def __init__(self, image_name, SIZE):
        self.SIZE = SIZE
        image = Image.open(image_name)

        assert image.size[0] == image.size[1] and image.size[
            0] == self.SIZE, "Height and Width of the image must be equal."

        image = np.asarray(image)
        self.img_px_array = np.copy(image)

        self.convertImg()
        self.fractal_dim = self.calculate_fractal_dim()

    def convertImg(self):
        for i in range(len(self.img_px_array)):
            for j in range(len(self.img_px_array[i])):
                for k in range(len(self.img_px_array[i][j])):
                    if (self.img_px_array[i][j][k] == 255):
                        self.img_px_array[i][j][k] = 0
                    else:
                        self.img_px_array[i][j][k] = 1

    def calculate_fractal_dim(self):
        self.dimensions = []
        self.filled_boxes = []

        self.img_px_array = self.img_px_array[:, :, 0]

        2 ** np.range(1, self.SIZE)

        size = 1
        while size != self.SIZE:
            size *= 2
            filled_box = self.boxcount(size)
            self.filled_boxes.append(filled_box)
            self.dimensions.append(size / self.SIZE)

        return -np.polyfit(np.log(self.dimensions), np.log(self.filled_boxes), 1)[0]

    def graph(self):
        plt.plot(-np.log(self.dimensions), np.log(self.filled_boxes))
        plt.title("Fractal Dimension : " + str(self.fractal_dim))
        plt.show()
        plt.clf()


ImageFractalDimension("/Users/joshuacoles/Developer/checkouts/jc3091/CompB DLA/data-analysis/figures/dla-eg.png", 1024)