import pathlib
from typing import List, Tuple, Union, Optional

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import attr
from attr.validators import instance_of, optional
import tifffile

from calcium_bflow_analysis.dff_analysis_and_plotting.dff_analysis import (
    locate_spikes_peakutils,
    scatter_spikes,
)

from calcium_bflow_analysis.dff_analysis_and_plotting.plot_cells_and_traces import (
    extract_cells_from_tif,
)

from calcium_bflow_analysis.colabeled_cells.find_colabeled_cells import TiffChannels


@attr.s
class FovSubsetData:
    """
    A dataclass-like object keeping tabs of data for a subset of data
    in a given FOV. Used, for example, when a FOV has both labeled
    and unlabeled cells. In this case, a FovData instance will contain
    two FovSubsetData instances.

    Parameters:
    :param pathlib.Path results_file: an .npz file generated by CaImAn
    :param Optional[bool] with_labeling: Controls whether the data
    was taken with a second channel containing morphological data. True
    means that this subset points to the labeled data, False means that
    this subset points to the unlabeled data, and None means that there was
    no colabeling involved with this data.
    :param pathlib.Path colabel_img: Path to the tif file containing the
    images of the cell as seen in the colabeled channel.
    """

    results_file = attr.ib(validator=instance_of(pathlib.Path))
    with_labeling = attr.ib(validator=optional(instance_of(bool)))
    colabel_img = attr.ib(default=None)
    tif_file = attr.ib(init=False)
    colabel_file = attr.ib(init=False)
    colabel_stack = attr.ib(init=False)
    dff = attr.ib(init=False, repr=False)
    indices = attr.ib(init=False, repr=False)
    loaded = attr.ib(init=False)

    def load_data(self):
        """ Main class method to populate its different
        attributes with the data and proper files """
        self.tif_file = self._find_tif_file()
        if self.with_labeling is not None:
            self.colabel_file = self._find_colabeled_file()
            self.colabel_stack = self._load_colabeled_img()
        self.dff, self.indices = self._populate_dff_data()
        self.loaded = True

    def _find_tif_file(self):
        """
        Finds and returns the associated tif file. Returns None if
        doesn't exist.
        """
        name = self.results_file.name[:-12] + ".tif"
        try:
            tif_file = next(self.results_file.parent.glob(name))
            return tif_file
        except StopIteration:
            print(f"Tif not found for {name}")
            return None

    def _find_colabeled_file(self) -> Union[pathlib.Path, None]:
        """
        Finds and returns the colabeled file. Returns None if
        doesn't exist.
        """
        name = self.results_file.name[:-11] + "colabeled_idx.npy"
        try:
            colabel_file = next(self.results_file.parent.glob(name))
            return colabel_file
        except StopIteration:
            return None

    def _load_colabeled_img(self) -> np.ndarray:
        """
        Loads a tif file containing the parallel data
        of the colabeled cells to memory.
        """
        return tifffile.imread(str(self.colabel_img))

    def _populate_dff_data(self):
        """
        Using the different found filenames, load the dF/F data into
        memory. If a subset of the rows should be loaded (since we're
        working with labeled data) the indices of the relevant
        rows are also returned.
        """
        all_data = np.load(self.results_file)["F_dff"]
        if self.with_labeling is None:
            return all_data, np.arange(all_data.shape[0])

        indices = np.load(self.colabel_file)
        if self.with_labeling:
            return all_data[indices], indices
        if not self.with_labeling:
            all_indices = np.arange(all_data.shape[0])
            remaining_indices = np.delete(all_indices, indices)
            remaining_traces = all_data[remaining_indices]
            return remaining_traces, remaining_indices


@attr.s
class FovDataContainer:
    """
    A dataclass-like object which holds data taken from a specific FOV
    in its two subsets, or variations, i.e. labeled and unlabeled
    cells.
    """

    results_file = attr.ib(validator=instance_of(pathlib.Path))
    labeled = attr.ib(validator=optional(instance_of(FovSubsetData)), repr=False)
    unlabeled = attr.ib(validator=instance_of(FovSubsetData), repr=False)
    fps = attr.ib(default=58.8, validator=instance_of(float))
    cell_radius = attr.ib(default=9, validator=instance_of(int))
    tif = attr.ib(init=False)
    all_data = attr.ib(init=False)

    def __attrs_post_init__(self):
        self.tif = self.unlabeled.tif_file
        self.all_data = np.load(self.results_file)["F_dff"]
        if self.labeled.loaded is None:
            self.labeled.load_data()
        if self.unlabeled.loaded in None:
            self.unlabeled.load_data()


@attr.s
class ShowLabeledAndUnlabeled:
    """
    Plots a simple comparison of the dF/F traces that
    originated from the unlabeled cells and the labeled
    cells.
    """

    fovlist = attr.ib(validator=instance_of(List))

    def run(self):
        """ Main pipeline """

        max_shape = self._find_max_shape(self.fovlist)
        # labeled, unlabeled = self._stack_dff_arrays(self.file_pairs, max_shape)
        fig = plt.figure(figsize=(20, 16))
        gs = gridspec.GridSpec(max_shape[0], 16, figure=fig)
        self._plot_traces_with_cell_img(self.fovlist, gs, max_shape)
        # self._populate_fig_rows(
        #     gs, labeled, unlabeled, self.fps, self.indices_per_fov, max_shape
        # )

        # self._plot_against(labeled, unlabeled, self.fps, self.indices_per_fov)

    def _find_max_shape(self, fovlist):
        """
        Iterate over the found files and decide upon the shape of
        the array that will hold the stacked data. This is useful
        when the number of measurements in each FOV was unequal.
        """
        shapes = []
        num_of_labeled = 0
        for fov in fovlist:
            shapes.append(fov.all_data.shape)
            try:
                num_of_labeled += fov.labeled.indices.shape[0]
            except AttributeError:  # no labeled data present
                pass

        shapes = np.array(shapes)
        num_of_rows = shapes[:, 0].sum() - num_of_labeled
        max_cols = shapes[:, 1].max()
        return num_of_rows, max_cols

    def _plot_traces_with_cell_img(self, fovlist, gs, maxshape):
        """
        Create a plot built from rows, each containing an small
        image of an identified component and it corresponding trace.
        The final plot contains two columns, one for each
        type of neuron shown in the image., as well as the framerate of the acquisition
        """
        starting_row_idx = 0
        ending_row_idx = 0
        for fov in fovlist:
            time_vec = np.arange(maxshape[1]) / fov.fps
            starting_row_idx = ending_row_idx
            if fov.labeled:
                ending_row_idx += max(
                    fov.unlabeled.dff.shape[0], fov.labeled.dff.shape[0]
                )
            else:
                ending_row_idx += fov.all_data.shape[0]
            midpoint = gs.get_geometry()[1] // 2
            self._show_cell_excerpts(
                fovsubset=fov.unlabeled,
                radius=fov.cell_radius,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_col=0,
            )
            self._show_cell_excerpts(
                fovsubset=fov.labeled,
                radius=fov.cell_radius,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_col=midpoint,
            )
            self._show_traces(
                fovsubset=fov.unlabeled,
                time_vec=time_vec,
                fps=fov.fps,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_cols=slice(1, midpoint),
            )
            self._show_traces(
                fovsubset=fov.labeled,
                time_vec=time_vec,
                fps=fov.fps,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_cols=slice(midpoint+1, None),
            )

    def _show_cell_excerpts(self, fovsubset, radius, gs, gs_rows, gs_col):
        """
        Plots a column of cell excerpts in the given ax
        """
        excerpts = extract_cells_from_tif(
            fovsubset.results_file,
            str(fovsubset.tif_file),
            fovsubset.indices,
            num=100,
            cell_radius=radius,
            data_channel=TiffChannels.ONE,
            number_of_channels=1,
        )
        cell_means = np.nanmean(excerpts, axis=1)
        if fovsubset.colabel_stack is not None:
            colabeled_excerpts = extract_cells_from_tif(
                fovsubset.results_file,
                str(fovsubset.colabel_img),
                fovsubset.indices,
                num=100,
                cell_radius=radius,
                data_channel=TiffChannels.ONE,
                number_of_channels=1,
            )
            colabeled_means = np.nanmean(colabeled_excerpts, axis=1)
        else:
            colabeled_means = [np.array([]) for item in cell_means]
        axis_rows = range(gs_rows.start, gs_rows.stop)[::-1]
        for cell, colabel, ax_row in zip(cell_means, colabeled_means, axis_rows):
            ax_img = plt.subplot(gs[ax_row, gs_col])
            ax_img.imshow(cell, cmap='gray')
            ax_img.imshow(colabel, cmap='cool', alpha=0.4)
            ax_img.axis('off')

    def _show_traces(self, fovsubset, time_vec, fps, gs, gs_rows, gs_cols):
        """
        For the given axes shows the trace of the dF/F of the
        cell over time.
        """
        detected_spikes = locate_spikes_peakutils(fovsubset.dff, fps=fps)
        ax = plt.subplot(gs[gs_rows, gs_cols])
        if time_vec.shape[0] != fovsubset.dff.shape[1]:
            dff = np.zeros((fovsubset.dff.shape[0], time_vec.shape[0]))
            spikes = np.zeros((detected_spikes.shape[0], time_vec.shape[0]))
            spikes[:, :detected_spikes.shape[1]] = detected_spikes
            dff[:, :fovsubset.dff.shape[1]] = fovsubset.dff
        else:
            dff = fovsubset.dff
            spikes = detected_spikes
        scatter_spikes(dff, spikes, downsample_display=1, time_vec=time_vec, ax=ax)


if __name__ == "__main__":
    foldername = pathlib.Path("/data/Amit_QNAP/ForHagai/FOV3")
    fovs = []
    resfile = foldername / '355_GCaMP6-Ch2_WFA-590-Ch1_X25_mag3_act3b-940nm_256px_20180313_00001_CHANNEL_2_results.npz'
    colabel_img = foldername / ''
    for res_file in foldername.rglob("*act3a*results.npz"):
        subset_with = FovSubsetData(res_file, with_labeling=True)
        subset_with.load_data()
        subset_without = FovSubsetData(res_file, with_labeling=False)
        subset_without.load_data()
        fovs.append(FovDataContainer(res_file, subset_with, subset_without))

    showl = ShowLabeledAndUnlabeled(fovs)
    showl.run()
    plt.show()

import pathlib
from typing import List, Tuple, Union, Optional

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import attr
from attr.validators import instance_of, optional
import tifffile

from calcium_bflow_analysis.dff_analysis_and_plotting.dff_analysis import (
    locate_spikes_peakutils,
    scatter_spikes,
)

from calcium_bflow_analysis.dff_analysis_and_plotting.plot_cells_and_traces import (
    extract_cells_from_tif,
)

from calcium_bflow_analysis.colabeled_cells.find_colabeled_cells import TiffChannels


@attr.s
class FovSubsetData:
    """
    A dataclass-like object keeping tabs of data for a subset of data
    in a given FOV. Used, for example, when a FOV has both labeled
    and unlabeled cells. In this case, a FovData instance will contain
    two FovSubsetData instances.

    Parameters:
    :param pathlib.Path results_file: an .npz file generated by CaImAn
    :param Optional[bool] with_labeling: Controls whether the data
    was taken with a second channel containing morphological data. True
    means that this subset points to the labeled data, False means that
    this subset points to the unlabeled data, and None means that there was
    no colabeling involved with this data.
    :param pathlib.Path colabel_img: Path to the tif file containing the
    images of the cell as seen in the colabeled channel.
    """

    results_file = attr.ib(validator=instance_of(pathlib.Path))
    with_labeling = attr.ib(validator=optional(instance_of(bool)))
    colabel_img = attr.ib(default=None)
    tif_file = attr.ib(init=False)
    colabel_file = attr.ib(init=False)
    colabel_stack = attr.ib(init=False)
    dff = attr.ib(init=False, repr=False)
    indices = attr.ib(init=False, repr=False)
    loaded = attr.ib(init=False)

    def load_data(self):
        """ Main class method to populate its different
        attributes with the data and proper files """
        self.tif_file = self._find_tif_file()
        if self.with_labeling is not None:
            self.colabel_file = self._find_colabeled_file()
            self.colabel_stack = self._load_colabeled_img()
        self.dff, self.indices = self._populate_dff_data()
        self.loaded = True

    def _find_tif_file(self):
        """
        Finds and returns the associated tif file. Returns None if
        doesn't exist.
        """
        name = self.results_file.name[:-12] + ".tif"
        try:
            tif_file = next(self.results_file.parent.glob(name))
            return tif_file
        except StopIteration:
            print(f"Tif not found for {name}")
            return None

    def _find_colabeled_file(self) -> Union[pathlib.Path, None]:
        """
        Finds and returns the colabeled file. Returns None if
        doesn't exist.
        """
        name = self.results_file.name[:-11] + "colabeled_idx.npy"
        try:
            colabel_file = next(self.results_file.parent.glob(name))
            return colabel_file
        except StopIteration:
            return None

    def _load_colabeled_img(self) -> np.ndarray:
        """
        Loads a tif file containing the parallel data
        of the colabeled cells to memory.
        """
        return tifffile.imread(str(self.colabel_img))

    def _populate_dff_data(self):
        """
        Using the different found filenames, load the dF/F data into
        memory. If a subset of the rows should be loaded (since we're
        working with labeled data) the indices of the relevant
        rows are also returned.
        """
        all_data = np.load(self.results_file)["F_dff"]
        if self.with_labeling is None:
            return all_data, np.arange(all_data.shape[0])

        indices = np.load(self.colabel_file)
        if self.with_labeling:
            return all_data[indices], indices
        if not self.with_labeling:
            all_indices = np.arange(all_data.shape[0])
            remaining_indices = np.delete(all_indices, indices)
            remaining_traces = all_data[remaining_indices]
            return remaining_traces, remaining_indices


@attr.s
class FovDataContainer:
    """
    A dataclass-like object which holds data taken from a specific FOV
    in its two subsets, or variations, i.e. labeled and unlabeled
    cells.
    """

    results_file = attr.ib(validator=instance_of(pathlib.Path))
    labeled = attr.ib(validator=optional(instance_of(FovSubsetData)), repr=False)
    unlabeled = attr.ib(validator=instance_of(FovSubsetData), repr=False)
    fps = attr.ib(default=58.8, validator=instance_of(float))
    cell_radius = attr.ib(default=9, validator=instance_of(int))
    tif = attr.ib(init=False)
    all_data = attr.ib(init=False)

    def __attrs_post_init__(self):
        self.tif = self.unlabeled.tif_file
        self.all_data = np.load(self.results_file)["F_dff"]
        if self.labeled.loaded is None:
            self.labeled.load_data()
        if self.unlabeled.loaded in None:
            self.unlabeled.load_data()


@attr.s
class ShowLabeledAndUnlabeled:
    """
    Plots a simple comparison of the dF/F traces that
    originated from the unlabeled cells and the labeled
    cells.
    """

    fovlist = attr.ib(validator=instance_of(List))

    def run(self):
        """ Main pipeline """

        max_shape = self._find_max_shape(self.fovlist)
        # labeled, unlabeled = self._stack_dff_arrays(self.file_pairs, max_shape)
        fig = plt.figure(figsize=(20, 16))
        gs = gridspec.GridSpec(max_shape[0], 16, figure=fig)
        self._plot_traces_with_cell_img(self.fovlist, gs, max_shape)
        # self._populate_fig_rows(
        #     gs, labeled, unlabeled, self.fps, self.indices_per_fov, max_shape
        # )

        # self._plot_against(labeled, unlabeled, self.fps, self.indices_per_fov)

    def _find_max_shape(self, fovlist):
        """
        Iterate over the found files and decide upon the shape of
        the array that will hold the stacked data. This is useful
        when the number of measurements in each FOV was unequal.
        """
        shapes = []
        num_of_labeled = 0
        for fov in fovlist:
            shapes.append(fov.all_data.shape)
            try:
                num_of_labeled += fov.labeled.indices.shape[0]
            except AttributeError:  # no labeled data present
                pass

        shapes = np.array(shapes)
        num_of_rows = shapes[:, 0].sum() - num_of_labeled
        max_cols = shapes[:, 1].max()
        return num_of_rows, max_cols

    def _plot_traces_with_cell_img(self, fovlist, gs, maxshape):
        """
        Create a plot built from rows, each containing an small
        image of an identified component and it corresponding trace.
        The final plot contains two columns, one for each
        type of neuron shown in the image., as well as the framerate of the acquisition
        """
        starting_row_idx = 0
        ending_row_idx = 0
        for fov in fovlist:
            time_vec = np.arange(maxshape[1]) / fov.fps
            starting_row_idx = ending_row_idx
            if fov.labeled:
                ending_row_idx += max(
                    fov.unlabeled.dff.shape[0], fov.labeled.dff.shape[0]
                )
            else:
                ending_row_idx += fov.all_data.shape[0]
            midpoint = gs.get_geometry()[1] // 2
            self._show_cell_excerpts(
                fovsubset=fov.unlabeled,
                radius=fov.cell_radius,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_col=0,
            )
            self._show_cell_excerpts(
                fovsubset=fov.labeled,
                radius=fov.cell_radius,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_col=midpoint,
            )
            self._show_traces(
                fovsubset=fov.unlabeled,
                time_vec=time_vec,
                fps=fov.fps,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_cols=slice(1, midpoint),
            )
            self._show_traces(
                fovsubset=fov.labeled,
                time_vec=time_vec,
                fps=fov.fps,
                gs=gs,
                gs_rows=slice(starting_row_idx, ending_row_idx),
                gs_cols=slice(midpoint+1, None),
            )

    def _show_cell_excerpts(self, fovsubset, radius, gs, gs_rows, gs_col):
        """
        Plots a column of cell excerpts in the given ax
        """
        excerpts = extract_cells_from_tif(
            fovsubset.results_file,
            str(fovsubset.tif_file),
            fovsubset.indices,
            num=100,
            cell_radius=radius,
            data_channel=TiffChannels.ONE,
            number_of_channels=1,
        )
        cell_means = np.nanmean(excerpts, axis=1)
        if fovsubset.colabel_stack is not None:
            colabeled_excerpts = extract_cells_from_tif(
                fovsubset.results_file,
                str(fovsubset.colabel_img),
                fovsubset.indices,
                num=100,
                cell_radius=radius,
                data_channel=TiffChannels.ONE,
                number_of_channels=1,
            )
            colabeled_means = np.nanmean(colabeled_excerpts, axis=1)
        else:
            colabeled_means = [np.array([]) for item in cell_means]
        axis_rows = range(gs_rows.start, gs_rows.stop)[::-1]
        for cell, colabel, ax_row in zip(cell_means, colabeled_means, axis_rows):
            ax_img = plt.subplot(gs[ax_row, gs_col])
            ax_img.imshow(cell, cmap='gray')
            ax_img.imshow(colabel, cmap='cool', alpha=0.4)
            ax_img.axis('off')

    def _show_traces(self, fovsubset, time_vec, fps, gs, gs_rows, gs_cols):
        """
        For the given axes shows the trace of the dF/F of the
        cell over time.
        """
        detected_spikes = locate_spikes_peakutils(fovsubset.dff, fps=fps)
        ax = plt.subplot(gs[gs_rows, gs_cols])
        if time_vec.shape[0] != fovsubset.dff.shape[1]:
            dff = np.zeros((fovsubset.dff.shape[0], time_vec.shape[0]))
            spikes = np.zeros((detected_spikes.shape[0], time_vec.shape[0]))
            spikes[:, :detected_spikes.shape[1]] = detected_spikes
            dff[:, :fovsubset.dff.shape[1]] = fovsubset.dff
        else:
            dff = fovsubset.dff
            spikes = detected_spikes
        scatter_spikes(dff, spikes, downsample_display=1, time_vec=time_vec, ax=ax)


if __name__ == "__main__":
    foldername = pathlib.Path("/data/Amit_QNAP/ForHagai/FOV3")
    fovs = []
    resfile = foldername / '355_GCaMP6-Ch2_WFA-590-Ch1_X25_mag3_act3b-940nm_256px_20180313_00001_CHANNEL_2_results.npz'
    colabel_img = foldername / ''
    for res_file in foldername.rglob("*act3a*results.npz"):
        subset_with = FovSubsetData(res_file, with_labeling=True)
        subset_with.load_data()
        subset_without = FovSubsetData(res_file, with_labeling=False)
        subset_without.load_data()
        fovs.append(FovDataContainer(res_file, subset_with, subset_without))

    showl = ShowLabeledAndUnlabeled(fovs)
    showl.run()
    plt.show()

A dataclass-like object keeping tabs of data for a subset of data     in a given FOV. Used, for example, when a FOV has both labeled     and unlabeled cells. In this case, a FovData instance will contain     two FovSubsetData instances.      Parameters:     :param pathlib.Path results_file: an .npz file generated by CaImAn     :param Optional[bool] with_labeling: Controls whether the data     was taken with a second channel containing morphological data. True     means that this subset points to the labeled data, False means that     this subset points to the unlabeled data, and None means that there was     no colabeling involved with this data.     :param pathlib.Path colabel_img: Path to the tif file containing the     images of the cell as seen in the colabeled channel. Main class method to populate its different         attributes with the data and proper files Finds and returns the associated tif file. Returns None if         doesn't exist. Finds and returns the colabeled file. Returns None if         doesn't exist. Loads a tif file containing the parallel data         of the colabeled cells to memory. Using the different found filenames, load the dF/F data into         memory. If a subset of the rows should be loaded (since we're         working with labeled data) the indices of the relevant         rows are also returned. A dataclass-like object which holds data taken from a specific FOV     in its two subsets, or variations, i.e. labeled and unlabeled     cells. Plots a simple comparison of the dF/F traces that     originated from the unlabeled cells and the labeled     cells. Main pipeline # labeled, unlabeled = self._stack_dff_arrays(self.file_pairs, max_shape) # self._populate_fig_rows( #     gs, labeled, unlabeled, self.fps, self.indices_per_fov, max_shape # ) # self._plot_against(labeled, unlabeled, self.fps, self.indices_per_fov) Iterate over the found files and decide upon the shape of         the array that will hold the stacked data. This is useful         when the number of measurements in each FOV was unequal. # no labeled data present Create a plot built from rows, each containing an small         image of an identified component and it corresponding trace.         The final plot contains two columns, one for each         type of neuron shown in the image., as well as the framerate of the acquisition Plots a column of cell excerpts in the given ax For the given axes shows the trace of the dF/F of the         cell over time.