import numpy as np
import matplotlib.pyplot as plt
from ibllib import atlas
from brainbox.atlas import plot_atlas


def combine_layers_cortex(regions, delete_duplicates=False):
    remove = ["1", "2", "3", "4", "5", "6a", "6b", "/"]
    for i, region in enumerate(regions):
        for j, char in enumerate(remove):
            regions[i] = regions[i].replace(char, "")
    if delete_duplicates:
        regions = list(set(regions))
    return regions


# Coordinates of slices in mm
ML = -0.5
AP = 1
DV = -2

# Generate some mock data
ba = atlas.AllenAtlas(25)
all_regions = ba.regions.acronym
regions = np.random.choice(all_regions, size=500, replace=False)  # pick 500 random regions
values = np.random.uniform(-1, 1, 500)  # generate 500 random values

# Plot atlas
f, axs = plt.subplots(2, 3, figsize=(20, 10))
plot_atlas(regions, values, ML, AP, DV, color_palette="RdBu_r", minmax=[-1, 1], axs=axs[0])

# Now combine all layers of cortex
plot_regions = combine_layers_cortex(regions)
combined_cortex = combine_layers_cortex(all_regions)

# Plot atlas
plot_atlas(
    plot_regions,
    values,
    ML,
    AP,
    DV,
    color_palette="RdBu_r",
    minmax=[-1, 1],
    axs=axs[1],
    custom_region_list=combined_cortex,
)
