import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from dataclasses import dataclass
from typing import Tuple

@dataclass
class MandelbrotParams:
    max_iter: int = 100
    xmin: float = -2.0
    xmax: float = 1.0
    ymin: float = -1.5
    ymax: float = 1.5
    width: int = 1000
    height: int = 1000
    escape_radius: float = 2.0

@dataclass
class JuliaParams:
    c: complex = -0.4 + 0.6j
    max_iter: int = 100
    xmin: float = -1.5
    xmax: float = 1.5
    ymin: float = -1.5
    ymax: float = 1.5
    width: int = 1000
    height: int = 1000
    escape_radius: float = 2.0

@dataclass
class NewtonParams:
    power: int = 3
    max_iter: int = 50
    xmin: float = -1.5
    xmax: float = 1.5
    ymin: float = -1.5
    ymax: float = 1.5
    width: int = 1000
    height: int = 1000
    tolerance: float = 1e-5

class MandelbrotSet:
    def __init__(self, params: MandelbrotParams = None):
        self.params = params or MandelbrotParams()
        
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        c = x + y*1j
        z = c
        divtime = self.params.max_iter + np.zeros(z.shape, dtype=int)
        
        for i in range(self.params.max_iter):
            z = z**2 + c
            diverge = z*np.conj(z) > self.params.escape_radius**2
            div_now = diverge & (divtime == self.params.max_iter)
            divtime[div_now] = i
            z[diverge] = self.params.escape_radius
            
        return divtime
    
    def plot(self, data: np.ndarray, cmap: str = 'magma'):
        plt.figure(figsize=(10, 10))
        plt.imshow(data, cmap=cmap, extent=[
            self.params.xmin, self.params.xmax,
            self.params.ymin, self.params.ymax
        ])
        plt.colorbar(label='Iterations')
        plt.title('Mandelbrot Set')
        plt.xlabel('Re(c)')
        plt.ylabel('Im(c)')
        plt.show()

class JuliaSet:
    def __init__(self, params: JuliaParams = None):
        self.params = params or JuliaParams()
    
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        z = x + y*1j
        divtime = self.params.max_iter + np.zeros(z.shape, dtype=int)
        
        for i in range(self.params.max_iter):
            z = z**2 + self.params.c
            diverge = z*np.conj(z) > self.params.escape_radius**2
            div_now = diverge & (divtime == self.params.max_iter)
            divtime[div_now] = i
            z[diverge] = self.params.escape_radius
            
        return divtime
    
    def plot(self, data: np.ndarray, cmap: str = 'viridis'):
        plt.figure(figsize=(10, 10))
        plt.imshow(data, cmap=cmap, extent=[
            self.params.xmin, self.params.xmax,
            self.params.ymin, self.params.ymax
        ])
        plt.colorbar(label='Iterations')
        plt.title(f'Julia Set (c={self.params.c:.3f})')
        plt.xlabel('Re(z)')
        plt.ylabel('Im(z)')
        plt.show()

class NewtonFractal:
    def __init__(self, params: NewtonParams = None):
        self.params = params or NewtonParams()
        self.roots = np.array([
            np.exp(2j * np.pi * k / self.params.power) 
            for k in range(self.params.power)
        ])
    
    def _f(self, z: np.ndarray) -> np.ndarray:
        return z**self.params.power - 1
    
    def _df(self, z: np.ndarray) -> np.ndarray:
        return self.params.power * z**(self.params.power-1)
    
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        z = x + y*1j
        result = np.zeros(z.shape, dtype=int)
        
        for i in range(self.params.max_iter):
            dz = self._f(z) / self._df(z)
            z = z - dz
            for j, root in enumerate(self.roots):
                close = abs(z - root) < self.params.tolerance
                result[close & (result == 0)] = j + 1
                
        return result
    
    def plot(self, data: np.ndarray):
        colors = ['black'] + [plt.cm.hsv(i/self.params.power) 
                            for i in range(self.params.power)]
        cmap = LinearSegmentedColormap.from_list('newton', colors)
        
        plt.figure(figsize=(10, 10))
        plt.imshow(data, cmap=cmap, extent=[
            self.params.xmin, self.params.xmax,
            self.params.ymin, self.params.ymax
        ])
        plt.title(f'Newton Fractal (power={self.params.power})')
        plt.xlabel('Re(z)')
        plt.ylabel('Im(z)')
        plt.show()

params = MandelbrotParams(max_iter=150, xmin=-0.5, xmax=0.5)
fractal = MandelbrotSet(params)
data = fractal.generate()
fractal.plot(data)

params = JuliaParams(c=-0.8 + 0.156j)
fractal = JuliaSet(params)
data = fractal.generate()
fractal.plot(data)

params = NewtonParams(power=5)
fractal = NewtonFractal(params)
data = fractal.generate()
fractal.plot(data)

# добавляем палитру

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class FractalParams:
    max_iter: int = 100
    width: int = 1000
    height: int = 1000
    xmin: float = -2.0
    xmax: float = 1.0
    ymin: float = -1.5
    ymax: float = 1.5
    escape_radius: float = 2.0
    palette_levels: int = 8
    palette_colors: List[Tuple[float, float, float]] = None

    def __post_init__(self):
        if self.palette_colors is None:
            # Палитра по умолчанию: черный -> красный -> желтый -> белый
            self.palette_colors = [
                (0, 0, 0),
                (1, 0, 0),
                (1, 1, 0),
                (1, 1, 1)
            ]

class RecursivePalette:
    def __init__(self, base_colors: List[Tuple[float, float, float]], depth: int = 8):
        self.base_colors = base_colors
        self.depth = depth
        
    def get_color(self, value: float) -> Tuple[float, float, float]:
        for d in range(self.depth):
            threshold = 1.0 / (2 ** (d + 1))
            if value > threshold:
                color_idx = d % len(self.base_colors)
                return self.base_colors[color_idx]
        return self.base_colors[-1]
    
    def create_colormap(self, name: str = 'recursive') -> LinearSegmentedColormap:
        colors = []
        for i in range(256):
            colors.append(self.get_color(i / 255.0))
        return LinearSegmentedColormap.from_list(name, colors)

class MandelbrotSet:
    def __init__(self, params: FractalParams = None):
        self.params = params or FractalParams()
        self.palette = RecursivePalette(
            self.params.palette_colors, 
            self.params.palette_levels
        )
        
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        c = x + y*1j
        z = c
        divtime = self.params.max_iter + np.zeros(z.shape, dtype=float)
        
        for i in range(self.params.max_iter):
            z = z**2 + c
            diverge = z*np.conj(z) > self.params.escape_radius**2
            div_now = diverge & (divtime == self.params.max_iter)
            divtime[div_now] = i + 1 - np.log(np.log(abs(z[div_now]))) / np.log(2)
            z[diverge] = self.params.escape_radius
            
        return divtime / self.params.max_iter
    
    def plot(self, data: np.ndarray, title: str = "Mandelbrot Set"):
        plt.figure(figsize=(12, 12))
        plt.imshow(data.T, cmap=self.palette.create_colormap(),
                  extent=[self.params.xmin, self.params.xmax,
                         self.params.ymin, self.params.ymax])
        plt.axis('equal')
        plt.title(title)
        plt.show()

class JuliaSet:
    def __init__(self, params: FractalParams = None, c: complex = -0.4 + 0.6j):
        self.params = params or FractalParams()
        self.c = c
        self.palette = RecursivePalette(
            self.params.palette_colors, 
            self.params.palette_levels
        )
        
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        z = x + y*1j
        divtime = self.params.max_iter + np.zeros(z.shape, dtype=float)
        
        for i in range(self.params.max_iter):
            z = z**2 + self.c
            diverge = z*np.conj(z) > self.params.escape_radius**2
            div_now = diverge & (divtime == self.params.max_iter)
            divtime[div_now] = i + 1 - np.log(np.log(abs(z[div_now]))) / np.log(2)
            z[diverge] = self.params.escape_radius
            
        return divtime / self.params.max_iter
    
    def plot(self, data: np.ndarray, title: str = None):
        if title is None:
            title = f"Julia Set (c={self.c:.3f})"
        plt.figure(figsize=(12, 12))
        plt.imshow(data.T, cmap=self.palette.create_colormap(),
                  extent=[self.params.xmin, self.params.xmax,
                         self.params.ymin, self.params.ymax])
        plt.axis('equal')
        plt.title(title)
        plt.show()

class NewtonFractal:
    def __init__(self, params: FractalParams = None, power: int = 3):
        self.params = params or FractalParams()
        self.power = power
        self.roots = np.array([
            np.exp(2j * np.pi * k / self.power) 
            for k in range(self.power)
        ])
        self.palette = RecursivePalette(
            self.params.palette_colors, 
            self.params.palette_levels
        )
    
    def _f(self, z: np.ndarray) -> np.ndarray:
        return z**self.power - 1
    
    def _df(self, z: np.ndarray) -> np.ndarray:
        return self.power * z**(self.power-1)
    
    def generate(self) -> np.ndarray:
        y, x = np.ogrid[
            self.params.ymin:self.params.ymax:self.params.height*1j, 
            self.params.xmin:self.params.xmax:self.params.width*1j
        ]
        z = x + y*1j
        
        for i in range(self.params.max_iter):
            dz = self._f(z) / self._df(z)
            z = z - dz
            if np.all(abs(dz) < 1e-7):
                break
                
        result = np.zeros_like(z, dtype=float)
        for i, root in enumerate(self.roots, 1):
            result[abs(z - root) < 1e-3] = i / self.power
            
        return result
    
    def plot(self, data: np.ndarray, title: str = None):
        if title is None:
            title = f"Newton Fractal (power={self.power})"
        plt.figure(figsize=(12, 12))
        plt.imshow(data.T, cmap=self.palette.create_colormap(),
                  extent=[self.params.xmin, self.params.xmax,
                         self.params.ymin, self.params.ymax])
        plt.axis('equal')
        plt.title(title)
        plt.show()

# Океанская палитра
ocean_params = FractalParams(
    palette_colors=[
        (0, 0, 0.5),    # темно-синий
        (0, 0.5, 1),    # голубой
        (0, 1, 1),      # циан
        (1, 1, 1)       # белый
    ],
    palette_levels=8
)

# Космическая палитра
space_params = FractalParams(
    palette_colors=[
        (0, 0, 0),      # черный
        (0.5, 0, 1),    # пурпурный
        (1, 0, 1),      # фиолетовый
        (1, 1, 1)       # белый
    ],
    palette_levels=8
)

def create_palettes():
    # "Черная медуза"
    black_medusa = FractalParams(
        max_iter=500,
        palette_colors=[
            (0, 0, 0),      # черный
            (0.2, 0.2, 0.2), # темно-серый
            (0.7, 0.7, 0.7), # светло-серый
            (1, 1, 1)       # белый
        ],
        palette_levels=16
    )
    
    # "Циклон"
    cyclone = FractalParams(
        max_iter=500,
        palette_colors=[
            (0.1, 0.1, 0.3),  # темно-синий
            (0.2, 0.2, 0.5),  # синий
            (0.4, 0.4, 0.8),  # светло-синий
            (0.8, 0.8, 1.0),  # бледно-голубой
            (1, 1, 1)         # белый
        ],
        palette_levels=12
    )
    
    # "Огненный корабль"
    burning_ship = FractalParams(
        max_iter=500,
        palette_colors=[
            (0, 0, 0),        # черный
            (0.5, 0, 0),      # темно-красный
            (1, 0, 0),        # красный
            (1, 0.5, 0),      # оранжевый
            (1, 1, 0),        # желтый
            (1, 1, 1)         # белый
        ],
        palette_levels=10
    )
    
    # "Электрическая плесень"
    electric_mold = FractalParams(
        max_iter=500,
        palette_colors=[
            (0, 0, 0),        # черный
            (0.2, 0, 0.4),    # темно-фиолетовый
            (0.4, 0, 0.8),    # фиолетовый
            (0.6, 0.3, 1),    # светло-фиолетовый
            (0.8, 0.6, 1),    # лавандовый
            (1, 1, 1)         # белый
        ],
        palette_levels=12
    )
    
    # "Долина вилок"
    fork_valley = FractalParams(
        max_iter=500,
        palette_colors=[
            (0, 0.1, 0),      # темно-зеленый
            (0, 0.3, 0),      # зеленый
            (0.4, 0.6, 0),    # светло-зеленый
            (0.8, 1, 0.2),    # желто-зеленый
            (1, 1, 1)         # белый
        ],
        palette_levels=14
    )

    return [
        ("Black Medusa", black_medusa),
        ("Cyclone", cyclone),
        ("Burning Ship", burning_ship),
        ("Electric Mold", electric_mold),
        ("Fork Valley", fork_valley)
    ]

def generate_fractal_showcase():
    palettes = create_palettes()
    
    mandelbrot_areas = [
        {
            'title': 'Full Mandelbrot',
            'xmin': -2.0, 'xmax': 0.5,
            'ymin': -1.25, 'ymax': 1.25
        },
        {
            'title': 'Valley of Forks',
            'xmin': -0.7, 'xmax': -0.4,
            'ymin': -0.2, 'ymax': 0.1
        },
        {
            'title': 'Spiral Detail',
            'xmin': -0.744, 'xmax': -0.735,
            'ymin': 0.1, 'ymax': 0.109
        }
    ]
    
    julia_constants = [
        (-0.4 + 0.6j, "Classic Julia"),
        (-0.8 + 0.156j, "Dragon Julia"),
        (-0.5251993 + 0.5251993j, "Spiral Julia")
    ]
    
    newton_powers = [3, 4, 5]
    
    for palette_name, params in palettes:
        print(f"\nGenerating fractals with {palette_name} palette...")
        
        for area in mandelbrot_areas:
            params.xmin = area['xmin']
            params.xmax = area['xmax']
            params.ymin = area['ymin']
            params.ymax = area['ymax']
            
            mandel = MandelbrotSet(params)
            data = mandel.generate()
            mandel.plot(data, f"Mandelbrot Set - {palette_name} - {area['title']}")
        
        for c, desc in julia_constants:
            params.xmin = -1.5
            params.xmax = 1.5
            params.ymin = -1.5
            params.ymax = 1.5
            
            julia = JuliaSet(params, c)
            data = julia.generate()
            julia.plot(data, f"Julia Set - {palette_name} - {desc}")
        
        for power in newton_powers:
            params.xmin = -1.5
            params.xmax = 1.5
            params.ymin = -1.5
            params.ymax = 1.5
            
            newton = NewtonFractal(params, power)
            data = newton.generate()
            newton.plot(data, f"Newton Fractal - {palette_name} - Power {power}")

generate_fractal_showcase()

Generating fractals with Black Medusa palette...

Generating fractals with Cyclone palette...

Generating fractals with Burning Ship palette...

Generating fractals with Electric Mold palette...