from enum import Enum

import numpy as np
import plotly.express as px


class FractalType(Enum):
    Julia = 1
    Mandelbrot = 2


class FractalGenerator:
    """Creates a single fractal object and either returns it as as a numpy array, plot it or persists it as an pgn
    image. The output of this class is used by FractalTrainingValidationSet to generate training/val sets
    Args:
        complex_function -- complex function to make a Julia fractal
        n -- fractal size will ne n*n
        xlim,ylim -- tuples with the plotting region on the complex plane
        thr -- once a function grows larger that this number is considered to be divergent to infinity
        max_iter -- number of compositions of the complex function with itself
        type_ -- fractal type
        fractal -- numpy array with the fractal
    """

    def __init__(self, n=256, xlim=(-2, 2), ylim=(-2, 2), thr=2, max_iter=10):
        self.type_ = None
        self.fractal = None
        self.n = n
        self.xlim = xlim
        self.ylim = ylim
        self.thr = thr
        self.max_iter = max_iter

    def create_julia(self, complex_function=lambda z: np.sin(z ** 4 + 1.41)):
        """Creates a fractal of the Julia family, the fractal is stored inside self.fractal"""
        fractal = np.zeros((self.n, self.n), dtype="complex")
        x_space = np.linspace(self.xlim[0], self.xlim[1], self.n)
        y_space = np.linspace(self.ylim[0], self.ylim[1], self.n)
        for ix, x in enumerate(x_space):
            for iy, y in enumerate(y_space):
                for i in range(self.max_iter):
                    if i == 0:
                        z = complex(x, y)
                    z = complex_function(z)
                    if np.abs(z) >= self.thr:
                        z = self.thr
                        break
                fractal[ix, iy] = z
        self.fractal = np.abs(fractal)
        self.type_ = FractalType.Julia
        return self

    def create_mandelbrot(self):
        """Creates a fractal of the Mandelbrot family, the fractal is stored inside self.fractal"""
        fractal = np.zeros((self.n, self.n), dtype="complex")
        x_space = np.linspace(self.xlim[0], self.xlim[1], self.n)
        y_space = np.linspace(self.ylim[0], self.ylim[1], self.n)
        for ix, x in enumerate(x_space):
            for iy, y in enumerate(y_space):
                for i in range(self.max_iter):
                    if i == 0:
                        z = 0
                    z = z ** 2 + complex(x, y)
                    if np.abs(z) >= self.thr:
                        z = self.thr
                        break
                fractal[ix, iy] = z
        self.fractal = np.abs(fractal.transpose())
        self.type_ = FractalType.Mandelbrot
        return self

    def plot(self, **kwargs):
        if self.fractal is None:
            print("Nothing to plot. Generate a fractal first.")
            return None
        random_colormap = np.random.choice(
            ["orrd", "inferno_r", "hot_r", "jet_r", "purples", "agsunset_r"]
        )
        fig = px.imshow(
            img=self.fractal, color_continuous_scale=random_colormap, **kwargs
        )
        return fig