"""Defines the structure of individual analysis jobs."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_processing.ipynb.

# %% auto 0
__all__ = ['process_analysis_rois', 'AnalysisJob']

# %% ../nbs/04_processing.ipynb 3
import multiprocessing
import numpy as np
import pandas as pd
from pathlib import Path
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes._axes import Axes 
from datetime import datetime
from shapely import get_coordinates


from dateutil.tz import tzlocal
from uuid import uuid4
from pynwb import NWBHDF5IO, NWBFile
from pynwb.ophys import (
    Fluorescence,
    ImageSegmentation,
    OnePhotonSeries,
    OpticalChannel,
    RoiResponseSeries
)
from typing import Any

# %% ../nbs/04_processing.ipynb 4
from .input import DataLoader, ROI, GridWrapperROILoader
from .datamodels import Config
from .analysis import AnalysisROI
from . import results, __version__

# %% ../nbs/04_processing.ipynb 5
def process_analysis_rois(analysis_roi: AnalysisROI, config: Config) -> AnalysisROI:
    analysis_roi.compute_mean_intensity_timeseries(config.use_frame_range, config.start_frame_idx, config.end_frame_idx)
    if np.mean(analysis_roi.mean_intensity_over_time) >= config.mean_signal_threshold:
        analysis_roi.detect_peaks(config.signal_to_noise_ratio, config.min_octave_span, config.noise_window_size)
        analysis_roi.estimate_baseline(config.baseline_estimation_method)
        analysis_roi.compute_area_under_curve()
        analysis_roi.compute_amplitude_and_delta_f_over_f()
        analysis_roi.compute_variance_area(config.variance_window_size)
    return analysis_roi

# %% ../nbs/04_processing.ipynb 6
class AnalysisJob:

    def __init__(self, 
                 number_of_parallel_processes: int,
                 data_loaders: dict[str, DataLoader | list[DataLoader]],
                 results_dir: Path | None = None
                ) -> None:
        self.number_of_parallel_processes = number_of_parallel_processes
        self.recording_loader = data_loaders['recording']
        self.results_dir = results_dir
        self.results_dir_path: Path | None = None
        self.nwb_metadata: dict | None = None
        self.analysis_start_datetime: datetime | None = None
        if 'rois' in data_loaders.keys():
            self.rois_source = 'file'
            self.roi_loaders = data_loaders['rois']
        else:
            self.rois_source = 'grid'
            self.roi_loaders = [GridWrapperROILoader(self.recording_loader.filepath)]
        if 'focus_area' in data_loaders.keys():
            self.focus_area_enabled = True
            self.focus_area_loader = data_loaders['focus_area']
        else:
            self.focus_area_enabled = False
            self.focus_area = None

    
    def preview_window_size(self, window_size: int) -> tuple[Figure, Axes]:
        self.load_data_into_memory(window_size)
        grid_configs = self.roi_loaders[0].configs
        fig, ax = results.plot_window_size_preview(self.recording.preview, grid_configs, self.focus_area)
        return fig, ax

    
    def load_data_into_memory(self, window_size: int) -> None:
        self.recording = self.recording_loader.load_and_parse_file_content()
        if not hasattr(self, 'all_analysis_rois'):
            self.all_analysis_rois = self._create_all_analysis_rois(window_size)
        elif self.rois_source == 'grid':
            if self.roi_loaders[0].window_size != window_size:
                self.all_analysis_rois = self._create_all_analysis_rois(window_size)


    def _create_all_analysis_rois(self, window_size: int) -> list[AnalysisROI]:     
        self.all_rois = self._load_data_from_all_roi_loaders(window_size)
        if self.focus_area_enabled:
            self.focus_area = self.focus_area_loader.load_and_parse_file_content()[0]
            self.all_rois = self._filter_rois_by_focus_area()
        return self._create_analysis_rois()

    
    def _load_data_from_all_roi_loaders(self, window_size: int) -> list[ROI]:
        all_rois = []
        for roi_loader in self.roi_loaders:
            if type(roi_loader) == GridWrapperROILoader:
                roi_loader.set_configs_for_grid_creation(self.recording.preview.shape[1], self.recording.preview.shape[0], window_size)
            all_rois += roi_loader.load_and_parse_file_content()
        if self.rois_source == 'grid':
            self.grid_configs = self.roi_loaders[0].configs
        else:
            self.grid_configs = None
        return all_rois
    

    def _filter_rois_by_focus_area(self) -> list[ROI]:
        filtered_rois = [roi for roi in self.all_rois if roi.as_polygon.within(self.focus_area.as_polygon)]
        return filtered_rois
        

    def _create_analysis_rois(self) -> list[AnalysisROI]:
        all_analysis_rois = []
        if self.rois_source == 'file':
            self._create_and_add_label_ids_to_all_rois_from_file()
        for roi in self.all_rois:
            roi_bounding_box_row_col_coords = get_coordinates(roi.as_polygon.envelope).astype('int')
            row_min, col_min = roi_bounding_box_row_col_coords.min(axis=0)
            row_max, col_max = roi_bounding_box_row_col_coords.max(axis=0)   
            zstack = self.recording.zstack[:, row_min:row_max, col_min:col_max, :]
            analysis_roi = AnalysisROI(roi, (row_min, col_min), zstack)
            all_analysis_rois.append(analysis_roi)
        return all_analysis_rois


    def _create_and_add_label_ids_to_all_rois_from_file(self) -> None:
        roi_count = len(self.all_rois)
        zfill_factor = int(np.log10(roi_count)) + 1
        for idx, roi in enumerate(self.all_rois):
            label_id = str(idx + 1).zfill(zfill_factor)
            roi.add_label_id(label_id)

    def _create_results_dir(self) -> Path:
        """
        Creates and returns the directory path where analysis results will be saved.

        The directory name is constructed using the analysis start datetime, the recording filename, and optionally the focus area filename if focus area is enabled. If a custom results directory was provided during initialization, it is used instead. The method ensures the directory exists (creates it if necessary) and returns the Path object for the results directory.
        """
        prefix_with_datetime = self.analysis_start_datetime.strftime('%Y_%m_%d_%H-%M-%S_results_for')
        recording_filename = self.recording.filepath.stem
        if self.focus_area_enabled:
            focus_area_filename = self.focus_area.filepath.stem
            results_dir_name = f'{prefix_with_datetime}_{recording_filename}_with_{focus_area_filename}'
        else:
            results_dir_name = f'{prefix_with_datetime}_{recording_filename}'
        if self.results_dir is not None:
            results_dir_path = self.results_dir / results_dir_name
        else:
            results_dir_path = self.recording.filepath.parent / results_dir_name
        results_dir_path.mkdir(exist_ok=True, parents=True)
        return results_dir_path

    
    def run_analysis(self, config: Config) -> None:
        self.analysis_start_datetime = datetime.now(tzlocal())
        self.load_data_into_memory(config.grid_size)
        self.results_dir_path = self._create_results_dir()
        copy_of_all_analysis_rois = self.all_analysis_rois.copy()
        with multiprocessing.Pool(processes = self.number_of_parallel_processes) as pool:
            processed_analysis_rois = pool.starmap(process_analysis_rois, [(analysis_roi, config) for analysis_roi in copy_of_all_analysis_rois])
        self.all_analysis_rois = processed_analysis_rois


    def create_results(self, config: Config, nwb_metadata: dict[str, Any] = None) -> None:
        activity_filtered_analysis_rois = [roi for roi in self.all_analysis_rois if roi.peaks_count >= config.min_peak_count]
        self.activity_overview_plot = results.plot_activity_overview(analysis_rois_with_sufficient_activity = activity_filtered_analysis_rois,
                                                                     preview_image = self.recording.preview,
                                                                     indicate_activity = True,
                                                                     focus_area = self.focus_area,
                                                                     grid_configs = self.grid_configs)
        if config.save_overview_png:
            self.activity_overview_plot[0].savefig(self.results_dir_path.joinpath('activity_overview.png'), dpi=300)
            label_id_overview_fig, label_id_overview_ax = results.plot_rois_with_label_id_overview(
                analysis_rois_with_sufficient_activity=activity_filtered_analysis_rois,
                preview_image=self.recording.preview,
                focus_area=self.focus_area,
                grid_configs=self.grid_configs)
            label_id_overview_fig.savefig(self.results_dir_path.joinpath('ROI_label_IDs_overview.png'), dpi=300)
            plt.close()
        if config.save_summary_results:
            self._create_and_save_csv_result_files(activity_filtered_analysis_rois)
            self._create_and_save_individual_traces_pdf_result_file(activity_filtered_analysis_rois)
        if config.save_single_trace_results:
            self._create_and_save_single_trace_results_as_csv(activity_filtered_analysis_rois)
        if config.export_to_nwb:
            self._export_all_data_to_nwb(nwb_metadata)
 
    def _export_all_data_to_nwb(self, nwb_metadata: dict[str, Any] = None) -> None:
        if nwb_metadata is None:
            nwb_metadata = {}
        nwb_file = self._initialize_nwb_file(nwb_metadata)
        nwb_file = self._add_source_data(nwb_file, nwb_metadata)
        nwb_file = self._add_na3_results(nwb_file)
        self._write_nwb_file(nwb_file)


    def _initialize_nwb_file(self, nwb_metadata: dict[str, Any]) -> NWBFile:
        nwb_file = NWBFile(
            session_description=f'Autogenerated NWB file by neuralactivitycubic version {__version__}',
            identifier=str(uuid4()),
            session_start_time=datetime.now(tzlocal()),
            experimenter=nwb_metadata.get('experimenter', ''),
            experiment_description=nwb_metadata.get('experiment_description', ''),
            lab=nwb_metadata.get('lab', ''),
            institution=nwb_metadata.get('institution', ''),
            notes=nwb_metadata.get('notes', ''),
        )
        return nwb_file


    def _add_source_data(self, nwb_file: NWBFile, nwb_metadata: dict[str, Any]) -> NWBFile:
        # Add recording source data
        microscope = nwb_file.create_device(
            name = 'Microscope',
            model_number=nwb_metadata.get('model_number', ''),
            manufacturer=nwb_metadata.get('manufacturer', ''),
            description=nwb_metadata.get('description', ''),
        )
        optical_channel = OpticalChannel(
            name = 'channel_01',
            description = nwb_metadata.get('imaging_description', ''),
            emission_lambda = nwb_metadata.get('emission_wavelength', 0.0),
        )
        imaging_plane = nwb_file.create_imaging_plane(
            name = 'ImagingPlane',
            optical_channel = optical_channel,
            description = 'In vitro Calcium Imaging of neuronal cell culture',
            device = microscope,
            excitation_lambda = nwb_metadata.get('excitation_wavelength', 0.0),
            indicator = nwb_metadata.get('calcium_indicator', ''),
            location = nwb_metadata.get('location', ''),
        )
        cai_series = OnePhotonSeries(
            name = 'OnePhotonSeries',
            imaging_plane = imaging_plane,
            unit = 'a.u.',
            data = self.recording.zstack[:, :, :, 0], 
            rate = 10.0 #tbd
        )
        nwb_file.add_acquisition(cai_series)
        # Add ROI source data:
        ophys_module = nwb_file.create_processing_module(
            name="ophys", description="optical physiology processed data"
        )
        img_seg = ImageSegmentation()
        roi_plane_seg = img_seg.create_plane_segmentation(
            name="PlaneSegmentation",
            description="ROIs for analysis; created manually using ImageJ2",
            imaging_plane=imaging_plane,
            reference_images=cai_series,
        )
        ophys_module.add(img_seg)
        roi_plane_seg.add_column(name = 'ROI_label_ID', description = 'ROI label ID assigned by neuralactivitycubic', index = False)
        for roi in self.all_rois:
            roi.create_nwb_compliant_pixel_mask()
            roi_plane_seg.add_roi(pixel_mask = roi.pixel_mask, ROI_label_ID = roi.label_id)
        return nwb_file


    def _add_na3_results(self, nwb_file: NWBFile) -> NWBFile:
        # Add mean intensity time series per ROI:
        plane_segmentation = nwb_file.processing['ophys']['ImageSegmentation']['PlaneSegmentation']
        all_rois_table_region = plane_segmentation.create_roi_table_region(description = 'All ROIs')
        stacked_arrays = np.stack([analysis_roi.mean_intensity_over_time for analysis_roi in self.all_analysis_rois], axis=1)
        mean_int_roi_response_series = RoiResponseSeries(
            name="RoiResponseSeries",
            description="Mean intensity time series for all ROIs",
            data=stacked_arrays,
            rois=all_rois_table_region,
            unit="bit-values",
            rate=10.0, # use rate of image series in acquisition of NWB file
        )
        mean_int_fluorescence_series = Fluorescence(mean_int_roi_response_series)
        nwb_file.processing['ophys'].add(mean_int_fluorescence_series)
        # Add all analysis results
        custom_unit_column_infos = [
            ('ROI_label_ID', 'ROI label ID assigned by neuralactivitycubic', False),
            ('total_peaks_count', 'Total number of detected peaks', False),
            ('variance_area', 'Computed variance area', False),
            ('peaks_frame_idxs', 'The frame indices of all peaks', True),
            ('peaks_bit_value', 'The fluorescence intensities (bit values) of all peak', True),
            ('peaks_amplitude', 'The amplitudes of all peaks (above the computed baseline)', True),
            ('peaks_df_over_f', 'The delta-F over F values of all peaks', True),
            ('peaks_auc', 'The area under curve values of all peaks (np.NaN for individual peaks if no AUC could be computed)', True),
            ('peaks_classification', 'The classification of all peaks as assigned by neuralactivitycubic', True)
        ]
        for (name, description, index) in custom_unit_column_infos:
            nwb_file.add_unit_column(name = name, description = description, index = index)
        for analysis_roi in self.all_analysis_rois:
            info = self._get_info_for_nwb_units_table_from_analysis_roi(analysis_roi)
            nwb_file.add_unit(**info)
        return nwb_file


    def _write_nwb_file(self, nwb_file: NWBFile) -> None:
        filepath = self.results_dir_path.joinpath('autogenerated_nwb_file.nwb')
        with NWBHDF5IO(filepath, 'w') as io:
            io.write(nwb_file)


    def _get_info_for_nwb_units_table_from_analysis_roi(self, analysis_roi) -> dict:
        info = {
            'ROI_label_ID': analysis_roi.label_id, 
            'total_peaks_count': analysis_roi.peaks_count, 
            'variance_area': np.nan, 
            'peaks_frame_idxs': [], 
            'peaks_bit_value': [], 
            'peaks_amplitude': [], 
            'peaks_df_over_f': [], 
            'peaks_auc': [], 
            'peaks_classification': []
        }
        if hasattr(analysis_roi, 'variance_area'):
            info['variance_area'] = analysis_roi.variance_area
        if analysis_roi.peaks_count > 0:
            info['peaks_frame_idxs'] = [peak.frame_idx for peak in analysis_roi.peaks.values()]
            info['peaks_bit_value'] = [peak.intensity for peak in analysis_roi.peaks.values()]
            if hasattr(analysis_roi, 'baseline'):
                info['peaks_amplitude'] = [peak.amplitude for peak in analysis_roi.peaks.values()]
                info['peaks_df_over_f'] = [peak.delta_f_over_f for peak in analysis_roi.peaks.values()]
                info['peaks_auc'] = [peak.area_under_curve for peak in analysis_roi.peaks.values()]
                info['peaks_classification'] = [peak.peak_type for peak in analysis_roi.peaks.values()]
        return info


    def _create_variance_area_dataframe(self, filtered_rois: list[AnalysisROI]) -> pd.DataFrame:
        data = {'ROI label ID': [],
                'Variance Area': []}
        for roi in filtered_rois:
            data['ROI label ID'].append(roi.label_id)
            data['Variance Area'].append(roi.variance_area)
        df = pd.DataFrame(data = data)
        return df
    
    
    def _create_and_save_csv_result_files(self, filtered_rois: list[AnalysisROI]) -> None:
        if len(filtered_rois) > 0:
            df_variance_areas = self._create_variance_area_dataframe(filtered_rois)
            df_variance_areas.to_csv(self.results_dir_path.joinpath('Variance_area_results.csv'), index = False)
            peak_results_per_roi = [results.export_peak_results_df_from_analysis_roi(roi) for roi in filtered_rois]
            df_all_peak_results = pd.concat(peak_results_per_roi, ignore_index = True)
            max_peak_count_across_all_rois = df_all_peak_results.groupby('ROI label ID').count()['peak frame index'].max()
            zfill_factor = int(np.log10(max_peak_count_across_all_rois)) + 1
            amplitude_and_delta_f_over_f_results_all_rois = []
            auc_results_all_rois = []
            for roi_label_id in df_all_peak_results['ROI label ID'].unique():
                tmp_df_single_roi = df_all_peak_results[df_all_peak_results['ROI label ID'] == roi_label_id].copy()
                amplitude_and_delta_f_over_f_results_all_rois.append(results.create_single_roi_amplitude_and_delta_f_over_f_results(tmp_df_single_roi, zfill_factor))
                auc_results_all_rois.append(results.create_single_roi_auc_results(tmp_df_single_roi, zfill_factor))
            df_all_amplitude_and_delta_f_over_f_results = pd.concat(amplitude_and_delta_f_over_f_results_all_rois, ignore_index = True)
            df_all_auc_results = pd.concat(auc_results_all_rois, ignore_index = True)
            # Once all DataFrames are created successfully, write them to disk 
            df_all_peak_results.to_csv(self.results_dir_path.joinpath('all_peak_results.csv'), index = False)
            df_all_amplitude_and_delta_f_over_f_results.to_csv(self.results_dir_path.joinpath('Amplitude_and_dF_over_F_results.csv'), index = False)
            df_all_auc_results.to_csv(self.results_dir_path.joinpath('AUC_results.csv'), index = False)

    
    def _create_and_save_individual_traces_pdf_result_file(self, filtered_rois: list[AnalysisROI]) -> None:
            filepath = self.results_dir_path.joinpath('Individual_traces_with_identified_events.pdf')
            with PdfPages(filepath) as pdf:
                for indicate_activity in [True, False]:
                    overview_fig, ax = results.plot_activity_overview(filtered_rois, self.recording.preview, indicate_activity, self.focus_area, self.grid_configs)
                    pdf.savefig(overview_fig)
                    plt.close()
                label_ids_overview_fig, ax = results.plot_rois_with_label_id_overview(filtered_rois, self.recording.preview, self.focus_area, self.grid_configs)
                pdf.savefig(label_ids_overview_fig)
                plt.close()
                for roi in filtered_rois:
                    fig = results.plot_intensity_trace_with_identified_peaks_for_individual_roi(roi)
                    pdf.savefig(fig)
                    plt.close()


    def _create_and_save_single_trace_results_as_csv(self, filtered_rois: list[AnalysisROI]) -> None:
        single_trace_subdir_path = self.results_dir_path.joinpath('single_traces')
        single_trace_subdir_path.mkdir(parents = True, exist_ok = True)
        for analysis_roi in filtered_rois:
            df_single_trace = self._initialize_single_trace_df(analysis_roi)
            for peak_idx in df_single_trace[df_single_trace['is_peak'] == True].index:
                peak = analysis_roi.peaks[peak_idx]
                df_single_trace.at[peak_idx, 'peak_type'] = peak.peak_type
                df_single_trace.at[peak_idx, 'amplitude'] = peak.amplitude
                df_single_trace.at[peak_idx, 'delta_f_over_f'] = peak.delta_f_over_f
                df_single_trace.at[peak_idx, 'has_baseline_intersections'] = peak.has_neighboring_intersections
                if peak.has_neighboring_intersections == True:
                    df_single_trace.at[peak_idx, 'pre_peak_intersection_idx'] = peak.frame_idxs_of_neighboring_intersections[0]
                    df_single_trace.at[peak_idx, 'post_peak_intersection_idx'] = peak.frame_idxs_of_neighboring_intersections[1]
                    df_single_trace.at[peak_idx, 'area_under_curve'] = peak.area_under_curve
            self._save_single_trace_df(df_single_trace, analysis_roi, single_trace_subdir_path)
            

    def _initialize_single_trace_df(self, analysis_roi: AnalysisROI) -> pd.DataFrame:
        is_peak_mask = np.zeros(analysis_roi.mean_intensity_over_time.shape[0], dtype=bool)
        is_peak_mask[analysis_roi.frame_idxs_of_peaks] = True
        data = {'intensity': analysis_roi.mean_intensity_over_time,
                'estimated_baseline': analysis_roi.baseline,
                'is_peak': is_peak_mask}
        df = pd.DataFrame(data = data)
        df.index.name = 'frame_idx'
        new_columns = ['peak_type', 
                       'amplitude', 
                       'delta_f_over_f', 
                       'has_baseline_intersections', 
                       'pre_peak_intersection_idx', 
                       'post_peak_intersection_idx', 
                       'area_under_curve']
        df[new_columns] = np.nan
        df = df.astype({'peak_type': 'object',
                        'has_baseline_intersections': 'object'})
        return df


    def _save_single_trace_df(self, 
                              df_single_trace: pd.DataFrame, 
                              analysis_roi: AnalysisROI,
                              single_trace_subdir_path: Path) -> None:
        filepath_friendly_roi_label_id = analysis_roi.label_id.replace('/', '-')
        filepath = single_trace_subdir_path.joinpath(f'data_of_ROI_{filepath_friendly_roi_label_id}.csv')
        df_single_trace.to_csv(filepath)
