import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap


class Figure_Manager():

    def __init__(self, max_row, max_col, xlim, ylim):
        self.max_row = max_row
        self.max_col = max_col
        self.xlim = xlim
        self.ylim = ylim
        self.figure_num = 1

    def next(self):
        ax = plt.subplot(self.max_row, self.max_col, self.figure_num)
        self.figure_num += 1
        ax.set_xlim(self.xlim)
        ax.set_ylim(self.ylim)
        plt.xticks(())
        plt.yticks(())
        return ax


class Contour_Grid():

    def __init__(self, bound_square):
        x_min, x_max, y_min, y_max = bound_square
        X, Y = np.mgrid[x_min:x_max:1500j, y_min:y_max:1500j]
        self.x = 1500
        self.y = 1500
        self.X = X
        self.Y = Y
        self.xy = np.c_[X.ravel(), Y.ravel()]

    def get_xy(self):
        return self.xy

    def reshape_to_grid(self, v):
        return np.reshape(v, (self.x, self.y))


def plot_density(rs, bound_square, cmap_str):
    grid = Contour_Grid(bound_square)
    density = rs.eval(grid.xy)
    density = grid.reshape_to_grid(density)
    plt.contour(grid.X, grid.Y, density, 5, cmap=cmap_str, alpha=1, zorder=1)
    return True


def plot_boundary(f, bound_square):
    grid = Contour_Grid(bound_square)
    density = f.predict(grid.xy)
    density = grid.reshape_to_grid(density)
    colors = ['#C9283E'] * 4
    custom_cmap = LinearSegmentedColormap.from_list(
        'my_list', colors, N=len(colors))

    colors_background = ['#5b5b5b', '#999999', '#bcbcbc', '#ffffff']  # 灰色到白色
    cmap_background = LinearSegmentedColormap.from_list('my_cmap_background', colors_background, N=len(colors_background))
    # background_data = np.rot90(density, k=4)
    background_data = density.T
    plt.imshow(background_data,
               extent=(-2, 2, -2, 2),
               origin='lower',
               aspect='auto',
               cmap=cmap_background,
               alpha=1.0,  # 背景的透明度
               zorder=0)
    plt.contour(grid.X, grid.Y, density,  levels=[0, 1, 2, 3], cmap=custom_cmap, linewidths=1.5, zorder=1)
    return True


def plot_region(f, bound_square, mimic_gray_color):
    grid = Contour_Grid(bound_square)
    pred = f.predict(grid.xy)
    Z = grid.reshape_to_grid(pred)
    colors = mimic_gray_color
    custom_cmap = LinearSegmentedColormap.from_list(
        'my_list', colors, N=len(colors))
    plt.pcolormesh(grid.X, grid.Y, Z, cmap=custom_cmap)
    return True

def plot_circle(ax, Radius):
    circle = plt.Circle((0, 0), Radius, color='black', alpha=0.5, fill=False, linestyle='--', linewidth=2)
    ax.add_artist(circle)