dandiset_id = "001637"
dandi_filepath = "sub-820459/sub-820459_ses-ecephys-820459-2025-11-12-14-34-33_ecephys.nwb"
download_loc = "."

client = dandiapi.DandiAPIClient()
dandiset = client.get_dandiset(dandiset_id)
file = dandiset.get_asset_by_path(dandi_filepath)
file_url = file.download_url
filename = dandi_filepath.split("/")[-1]
filepath = f"{download_loc}/{filename}"
if os.path.exists(filepath):
    print("File already exists")
else:
    # This can sometimes take a while depending on the size of the file
    download.download(file_url, output_dir=download_loc)
    print(f"Downloaded file to {filepath}")

# the download code above will set filepath according to your download location
# filepath = r"/root/capsule/data/ecephys_820459_2025-11-12_14-34-33_nwb_2025-11-23_02-55-05/ecephys_820459_2025-11-12_14-34-33_experiment1_recording1.nwb"

io = NWBHDF5IO(filepath, mode="r", load_namespaces=True)
nwb = io.read()

nwb

nwb.subject

for k in nwb.intervals.keys(): print(k)

stim_table = nwb.intervals["Sequence mismatch block_presentations"]
print(stim_table.colnames)
stim_table
# use this to get a pandas dataframe
# stim_table_df = stim_table.to_dataframe()

np.unique(stim_table["TrialType"])

lfp = nwb.processing['ecephys']['LFP'].electrical_series
for k in lfp.keys(): print(k)

# conversion to np.array required for plotting
probeC_lfp_times = np.array(lfp["ElectricalSeriesProbeC-LFP"].timestamps)
probeC_lfp_traces = np.array(lfp["ElectricalSeriesProbeC-LFP"].data)

print(probeC_lfp_times.shape)
print(probeC_lfp_traces.shape)


plt.plot(probeC_lfp_times[:100], probeC_lfp_traces[:100])
plt.xlabel("time")
plt.ylabel("mV")
plt.title("LFP over time")
plt.show()

units_table = nwb.units
print(units_table.colnames)
units_table
# use this to get a pandas dataframe (it may take a lot of time and memory)
# units_table_df = units_table.to_dataframe()

plt.hist(np.array(units_table["snr"]), bins=100)
plt.show()

# conversion to np.array required for plotting
unit0_spikes = np.array(nwb.units["spike_times"][0])
plt.vlines(unit0_spikes[:100], 0, 1, lw=0.5)
plt.title("Unit 0 spike times")
plt.xlabel("Time")
plt.ylabel("Spike?")
plt.show()

electrodes = nwb.electrodes
electrodes
# use this to get pandas dataframe
# electrodes_df = electrodes.to_dataframe()
print(electrodes.colnames)
electrodes

np.unique(nwb.electrodes['group_name'])

eye_tracking = nwb.processing['eye_tracking']
eye_tracking.keys()

pupil_tracking = eye_tracking["pupil"]
pupil_tracking

xs = pupil_tracking['data_x']
ys = pupil_tracking['data_y']
start, end = 20, 5000

# generate colors along viridis for the range
colors = mpl.cm.viridis(np.linspace(0, 1, end-start))

# plot the line
plt.plot(xs[start:end], ys[start:end], zorder=0, linewidth=0.25, color='gray')

# scatter with colormap
sc = plt.scatter(xs[start:end], ys[start:end], s=5, c=colors, zorder=1)

plt.xlabel("Relative X position")
plt.ylabel("Relative Y position")

# colorbar for the gradient
cbar = plt.colorbar(sc)
cbar.ax.get_yaxis().labelpad = 15
cbar.ax.set_ylabel('Time', rotation=270)
cbar.set_ticks([])

plt.show()

wheel_rotation = nwb.acquisition["raw_running_wheel_rotation"]
# conversion to np.array required for plotting
wheel_times = np.array(wheel_rotation.timestamps)
wheel_trace = np.array(wheel_rotation.data)

print(wheel_times.shape)
print(wheel_trace.shape)

plt.plot(wheel_times,wheel_trace)
plt.title("Wheel Rotation over Time")
plt.xlabel("Time")
plt.show()

plt.plot(wheel_times, np.cumsum(wheel_trace))