import mne
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap


# lay object to plot
lay = mne.channels.read_layout("biosemi")

# info object to plot with
info = mne.create_info(lay.names, ch_types=['eeg']*len(lay.names), sfreq=256)
info.set_montage('biosemi64', match_case=False)

# get 3d locations
chs = info["chs"]
locs3d = np.array([ch["loc"][:3] for ch in chs])

# get associated colours from colour wheel
x, y, z = locs3d.T
colors = mne.viz.evoked._rgb(x, y, z)

# get inputs for plot_sensors
ch_grps = [[i] for i in range(len(info.ch_names))]
cmap = ListedColormap(colors)

fig, ax = plt.subplots()
pl = mne.viz.utils.plot_sensors(info, ch_groups=ch_grps, cmap=cmap, linewidth=0, pointsize=250, axes=ax)

topo_cols = pd.DataFrame(colors).rename(columns={0:'r', 1:'g', 2:'b'})
topo_cols['ch_name'] = info.ch_names
topo_cols = topo_cols[['ch_name', 'r', 'g', 'b']]

topo_cols.to_csv('Biosemi64_fiducial.csv', index=False)