from pathlib import Path
from pynwb import NWBHDF5IO
import numpy as np
import matplotlib.pyplot as plt
from dandi import dandiapi, download
import os

output_dir = "./plots"
output_dir = Path(output_dir)
output_dir.mkdir(exist_ok=True)

dandiset_id = "001424"
# dandi_filepath = "sub-794237/sub-794237_ses-20250403T103830_image+ophys.nwb"
# dandi_filepath = "sub-794237/sub-794237_ses-20250424T142019_image+ophys.nwb"
dandi_filepath = "sub-794237/sub-794237_ses-20250508T145040_image+ophys.nwb"
download_loc = "."

def summarize_trial_counts_by_position_and_orientation(stim_table):
    # Convert orientations to degrees and normalize to [0, 180)
    orientations_deg = np.degrees(stim_table['orientation']) % 180
    orientations_deg = orientations_deg.round().astype(int)

    stim_table = stim_table.copy()
    stim_table['orientation_deg'] = orientations_deg

    # Group by x, y, and orientation
    summary = (
        stim_table
        .groupby(['x_position', 'y_position', 'orientation_deg'])
        .size()
        .reset_index(name='trial_count')
        .sort_values(['y_position', 'x_position', 'orientation_deg'])
    )

    return summary

def calculate_receptive_fields(stim_table, dff, timestamps, onset_delay, window_duration):
    x_positions = np.unique(stim_table.x_position)
    sorted_x_positions = sorted(x_positions, key=float)
    y_positions = np.unique(stim_table.y_position)
    sorted_y_positions = sorted(y_positions, key=float)

    height, width = len(sorted_y_positions), len(sorted_x_positions)
    n_timestamps, n_rois = dff.shape

    receptive_fields = np.zeros((n_rois, height, width))
    stim_counts = np.zeros((height, width))  # to normalize later

    for i, row in stim_table.iterrows():
        response_start = row['start_time'] + onset_delay
        response_end = row['start_time'] + onset_delay + window_duration

        response_baseline_start = row['start_time'] - 5.5
        response_baseline_end = row['start_time'] - 0.5

        response_baseline_start, response_baseline_end = np.searchsorted(timestamps, (response_baseline_start, response_baseline_end))

        response_start, response_end = np.searchsorted(timestamps, (response_start, response_end))
        
        if response_start >= response_end:
            continue

        response = (np.nanmean(dff[response_start:response_end], axis=0) - np.nanmean(dff[response_baseline_start:response_baseline_end], axis=0)) / np.nanstd(dff[response_baseline_start:response_baseline_end], axis=0)
        # mean response
        response[np.isnan(response)] = 0
        xi = sorted_x_positions.index(row['x_position'])
        yi = sorted_y_positions.index([row['y_position']])
        receptive_fields[:, yi, xi] += response
        stim_counts[yi, xi] += 1

    # Avoid division by zero
    stim_counts[stim_counts == 0] = 1
    receptive_fields = receptive_fields / stim_counts[None, :, :]

    return receptive_fields

def plot_rfs(rfs_matrix, output_dir):
    n_cols = 8
    n_rows = (len(rfs_matrix) // n_cols)
    if len(rfs_matrix) % n_cols != 0:
        n_rows += 1

    fig, axs = plt.subplots(n_rows,n_cols, figsize=(n_cols*3, n_rows*3))
    for i, ax in enumerate(axs.flatten()):
        if i >= len(rfs_matrix):
            ax.axis('off')
            continue
        ax.imshow(rfs_matrix[i], clim=[-0.1, 0.5])
        ax.set_title(f'ROI {i}')
    fig.savefig(output_dir / "all_rfs.png", format="png", dpi=300)
    plt.close(fig)

    fig2 = plt.figure()
    plt.imshow(np.mean(rfs_matrix, axis=0), clim=[-0.1, 0.5])
    plt.savefig(output_dir / "avg_rf.png", format="png", dpi=300)
    plt.close(fig2)

def plot_rf_traces_max_only(stim_table, dff, timestamps, onset_delay, window_duration, output_dir, rfs_matrix):
    x_positions = np.unique(stim_table.x_position)
    sorted_x_positions = sorted(x_positions, key=float)
    y_positions = np.unique(stim_table.y_position)
    sorted_y_positions = sorted(y_positions, key=float)

    full_window_duration = 1.0 + window_duration
    trace_window = np.linspace(-1.0, window_duration, int(np.round(full_window_duration / np.median(np.diff(timestamps)))))
    n_rois = dff.shape[1]

    location_traces = {}
    orientation_by_trial = {}
    stop_offsets = []

    for i, row in stim_table.iterrows():
        x, y = row['x_position'], row['y_position']
        xi = sorted_x_positions.index(x)
        yi = sorted_y_positions.index(y)

        orientation_rad = row.orientation
        orientation_deg = int(round(np.degrees(orientation_rad))) % 180  # Normalize to [0, 180)
        
        start_time = row['start_time']
        stop_time = row['stop_time']

        base_start = start_time - 1.0
        base_end = start_time
        stim_start = start_time + onset_delay
        stim_end = stim_start + window_duration

        base_idx_start, base_idx_end = np.searchsorted(timestamps, (base_start, base_end))
        stim_idx_start, stim_idx_end = np.searchsorted(timestamps, (stim_start, stim_end))

        if stim_idx_end > stim_idx_start and base_idx_end > base_idx_start:
            baseline = dff[base_idx_start:base_idx_end]
            trace = dff[stim_idx_start:stim_idx_end]

            mean_base = np.nanmean(baseline, axis=0)
            std_base = np.nanstd(baseline, axis=0)
            z_baseline = (baseline - mean_base) / (std_base + 1e-10)
            z_trace = (trace - mean_base) / (std_base + 1e-10)

            if z_baseline.shape[0] + z_trace.shape[0] == len(trace_window):
                z_full = np.concatenate((z_baseline, z_trace), axis=0)
                location_traces.setdefault((xi, yi), []).append((z_full, orientation_deg))

                stop_offset = stop_time - stim_start
                stop_offsets.append(stop_offset)

    avg_stop_offset = np.mean(stop_offsets) if stop_offsets else window_duration

    # Assign colors per orientation
    all_orientations = sorted(set(o for traces in location_traces.values() for _, o in traces))
    color_map = plt.cm.get_cmap('tab10', len(all_orientations))
    orientation_colors = {ori: color_map(i) for i, ori in enumerate(all_orientations)}

    for roi in range(n_rois):
        best_response = -np.inf
        best_loc = None

        for (xi, yi), trials in location_traces.items():
            traces = np.stack([t[0] for t in trials])
            mean_resp = np.nanmean(traces[:, :, roi])
            if mean_resp > best_response:
                best_response = mean_resp
                best_loc = (xi, yi)

        if best_loc is None:
            continue

        trials = location_traces[best_loc]
        trace_plot_dir = output_dir / "traces_max" / f"roi_{roi}"
        trace_plot_dir.mkdir(parents=True, exist_ok=True)

        plt.figure(figsize=(8, 4))
        for z_full, orientation_deg in trials:
            color = orientation_colors[orientation_deg]
            plt.plot(trace_window, z_full[:, roi], color=color, alpha=0.4)

        plt.plot(trace_window, np.nanmean([z[:, roi] for z, _ in trials], axis=0), color='black', linewidth=2, label='Mean')
        plt.axvline(0, color='black', linestyle='--', linewidth=1)
        plt.axvline(avg_stop_offset, color='black', linestyle='--', linewidth=1)
        plt.title(f'ROI {roi} - strongest response at x={sorted_x_positions[best_loc[0]]}, y={sorted_y_positions[best_loc[1]]}')
        plt.xlabel('Time (s)')
        plt.ylabel('z-scored dF/F')

        # Add orientation legend
        handles = [plt.Line2D([0], [0], color=color, lw=2, label=f'{ori}°') for ori, color in orientation_colors.items()]
        plt.legend(handles=handles, title="Orientation", loc='upper right')

        plt.tight_layout()
        plt.savefig(trace_plot_dir / f"roi_{roi}_traces.png", dpi=150)
        plt.close()

        # RF map
        plt.figure()
        plt.imshow(rfs_matrix[roi], clim=[-0.1, 0.5])
        plt.title(f"RF map - ROI {roi}")
        plt.axis('off')
        plt.tight_layout()
        plt.savefig(trace_plot_dir / f"roi_{roi}_rf.png", dpi=150)
        plt.close()

onset_delay = 0.0
window_duration = 0.5