"""Contour tracking on image series."""

# Misc. package imports
from skimage import measure
import pandas as pd
from numpy import nan as NaN
import matplotlib.pyplot as plt
import imgbasics

# Local imports
from .analysis_base import Analysis
from .formatters import PandasFormatter
from .results import PandasTsvResults
from ..managers import FileManager
from ..parameters.analysis import Contours, Threshold
from ..viewers import AnalysisViewer


# ======================= Plotting / Animation classes =======================


class ContourTrackingViewer(AnalysisViewer):

    def _create_figure(self):
        self.fig, self.ax = plt.subplots()
        self.axs = self.ax,

    def _first_plot(self, data):
        """What to do the first time data arrives on the plot."""
        img = data['image']
        num = data['num']

        self.ax.set_title(f'img #{num}, grey level {self.analysis.threshold.value}')
        self.imshow = self.analysis.img_series._imshow(img, ax=self.ax, **self.kwargs)

        self.ax.axis('off')
        self.fig.tight_layout()

        # Plot full contour lines if data is present -------------------------

        try:
            raw_contours = data['raw contours']
        except KeyError:
            pass
        else:
            self.contour_lines = []
            for raw_contour in raw_contours:
                contour_line, = self.ax.plot(*raw_contour, '-r')
                self.contour_lines.append(contour_line)

        # Plot position of centroids of contours -----------------------------

        self.centroid_pts = []
        for contour_ppty in data['contour properties']:
            centroid_pt, = self.ax.plot(*contour_ppty[:2], '+b')
            self.centroid_pts.append(centroid_pt)

        self.updated_artists = self.contour_lines + self.centroid_pts + [self.imshow]

    def _update_plot(self, data):
        """What to do upon iterations of the plot after the first time."""
        img = data['image']
        num = data['num']

        self.ax.set_title(f'img #{num}, grey level {self.analysis.threshold.value}')

        self.imshow.set_array(img)

        for contour, analysis, line, pt in zip(data['raw contours'],
                                               data['contour properties'],
                                               self.contour_lines,
                                               self.centroid_pts):

            if contour is not None:
                line.set_data(*contour)
                pt.set_data(*analysis[:2])
            else:
                line.set_data(None, None)
                pt.set_data(None, None)


class ContourTrackingFormatter_Pandas(PandasFormatter):
    """Formatting of results spit out by analysis methods"""

    def _prepare_data_storage(self):
        """Prepare structure(s) that will hold the analyzed data."""
        n = len(self.analysis.reference_positions)

        # Initiate dict to store all contour data (for json saving later) ----
        if self.analysis.save_raw_contours:
            self._raw_contour_data = {str(k + 1): {} for k in range(n)}

        # Initiate pandas table to store data (for tsv saving later) ---------

        names = 'x', 'y', 'p', 'a'  # measurement names (p, a perimeter, area)
        cols = [name + str(k + 1) for k in range(n) for name in names]

        self._data = pd.DataFrame(index=self.analysis.nums, columns=cols)

    def _store_raw_contours(self, data):
        """ # Save raw contour data into dict """
        num = data['num']
        n = len(data['contour properties'])
        for k in range(n):
            contour = data['raw contours'][k]
            if contour is not None:
                x, y = contour
                coords = {'x': list(x), 'y': list(y)}
                # The str is because JSON converts to str, and so this makes
                # live data compatible with reloaded data from JSON
            else:
                coords = None
            self._raw_contour_data[str(k + 1)][str(num)] = coords

    def _store_data(self, data):
        """How to store data generated by analysis on a single image."""
        line = sum(data['contour properties'], start=())  # "Flatten" list of tuples
        self._data.loc[data['num']] = line
        if self.analysis.save_raw_contours:
            self._store_raw_contours(data)

    def _to_pandas(self):
        """How to convert data generated by _store_data() into a pandas table."""
        return self._data

    def _to_results(self):
        super()._to_results()
        self.analysis.results.raw_contour_data = self._raw_contour_data

    def _regenerate_data(self, num):
        """How to go back to raw dict of data from self.data.

        Useful for plotting / animating results again after analysis, among
        other things.
        """
        data = {}
        data['contour properties'] = []
        data['raw contours'] = []

        n = len(self.analysis.contours.data['position'])

        for k in range(n):

            # contour positions and perimeters
            lim1 = 'x' + str(k + 1)
            lim2 = 'a' + str(k + 1)
            xc, yc, perimeter, area = self.analysis.results.data.loc[num, lim1:lim2]
            data['contour properties'].append((xc, yc, perimeter, area))

            # full contour data
            try:
                raw_contour_data = self.analysis.results.raw_contour_data
            except AttributeError:
                data['raw contours'].append(None)
            else:
                contour = raw_contour_data[str(k + 1)][str(num)]
                if contour is not None:
                    x = raw_contour_data[str(k + 1)][str(num)]['x']
                    y = raw_contour_data[str(k + 1)][str(num)]['y']
                    data['raw contours'].append((x, y))
                else:
                    data['raw contours'].append(None)

        return data


class ContourTrackingResults_PandasTsv(PandasTsvResults):

    measurement_type = 'ctrack'
    default_filename = 'Img_ContourTracking'

    def _save_raw_contour_data(self, filename=None):
        name = self._set_filename(filename)
        raw_data_filename = name + '_RawContourData'
        FileManager.to_json(
            self.raw_contour_data,
            self.savepath,
            raw_data_filename,
        )

    def _load_raw_contour_data(self, filename=None):
        name = self._set_filename(filename)
        raw_data_filename = name + '_RawContourData'
        return FileManager.from_json(
            self.savepath,
            raw_data_filename,
        )

    def load(self, filename=None):
        """Load data and metadata from tsv/json files."""
        super().load(filename=filename)
        try:
            self.raw_contour_data = self._load_raw_contour_data(
                filename=filename,
            )
        except FileNotFoundError:
            pass

    def save(self, filename=None):
        """Save data and metadata into tsv/json files."""
        super().save(filename=filename)
        try:
            self.raw_contour_data
        except AttributeError:
            pass
        else:
            self._save_raw_contour_data(filename=filename)


# =========================== Main ANALYSIS class ============================


class ContourTracking(Analysis):
    """Class to track contours on image series."""

    measurement_type = 'ctrack'

    DefaultViewer = ContourTrackingViewer
    DefaultFormatter = ContourTrackingFormatter_Pandas
    DefaultResults = ContourTrackingResults_PandasTsv

    def __init__(
        self,
        img_series,
        savepath=None,
        save_raw_contours=True,
        Viewer=None,
        Formatter=None,
        Results=None,
    ):
        """Analysis of iso-grey-level contours and their evolution in series.

        Parameters
        ----------
        - img_series: image series from the ImgSeries class or subclasses

        - savepath: folder in which to save analysis data & metadata
                    (if not specified, the img_series savepath is used)

        - save_raw_contours: if True (default), all coordinates of contour
                             lines are saved in addition to contour positions,
                             centroids and areas.

        - Viewer: Viewer class/subclasses that is used to display and inspect
                  analysis data (is used by ViewerTools)
                  (if not specified, use default Viewer)

        - Formatter: class/subclass of Formatter to format results spit out
                     by the raw analysis into something storable/saveable
                     by the Results class.
                     (if not specified, use default Formatter)

        - Results: Results class/subclasses that is used to store, save and
                   load analysis data and metadata.
                   (if not specified, use default Results)
        """
        self.save_raw_contours = save_raw_contours

        super().__init__(
            img_series=img_series,
            Viewer=Viewer,
            Formatter=Formatter,
            Results=Results,
            savepath=savepath,
        )

        # empty contour param object, needs to be filled with contours.define()
        # or contours.load() prior to starting analysis with self.run()
        self.contours = Contours(self)
        self.threshold = Threshold(self)

    def _find_contours(self, img, level):
        """Define how contours are found on an image."""
        if img.ndim == 2:
            image = img
        else:
            image = self.img_series.imge_processor.grayscale(img)
        return measure.find_contours(image, level)

    def _update_reference_positions(self, data):
        """Next iteration will look for contours close to the current ones."""
        for i, contour_analysis in enumerate(data['contour properties']):
            if any(qty is NaN for qty in contour_analysis):
                # There has been a problem in detecting the contour
                pass
            else:
                # if position correctly detected, update where to look next
                xc, yc, *_ = contour_analysis
                self.reference_positions[i] = (xc, yc)

    def _analyze(self, img):
        """Find contours at level in file i closest to the reference positions.

        Parameters
        ----------
        - img: image array to be analyzed (e.g. numpy array).

        Output
        ------
        [(x1, y1, p1, a1), (y2, y2, p2, a1), ..., (xn, yn, pn, an)] where n is the
        number of contours followed and (x, y), p, a is position, perimeter, area
        """
        contours = self._find_contours(img, self.threshold.value)

        data = {'contour properties': []}     # Stores analysis data (centroid etc.)
        data['raw contours'] = []       # Stores full (x, y) contour data

        for refpos in self.reference_positions:

            try:
                # this time edge=false, because trying to find contour closest
                # to the recorded centroid position, not edges
                contour = imgbasics.closest_contour(
                    contours=contours,
                    position=refpos,
                    edge=True,
                )

            except imgbasics.ContourError:
                # No contour at all detected on image --> return NaN
                xc, yc, perimeter, area = (NaN,) * 4
                data['raw contours'].append(None)

            else:

                x, y = imgbasics.contour_coords(contour, source='scikit')

                contprops = imgbasics.contour_properties(x, y)

                xc, yc = contprops['centroid']
                perimeter = contprops['perimeter']
                area = contprops['area']
                data['raw contours'].append((x, y))

            data['contour properties'].append((xc, yc, perimeter, area))

        self._update_reference_positions(data)

        return data

    def _initialize(self):
        """Check everything OK before starting analysis & initialize params."""

        if self.contours.is_empty:
            msg = "Contours not defined yet. Use self.contours.define(), "\
                  "or self.contours.load() if contours have been previously saved."
            raise AttributeError(msg)

        self.reference_positions = list(self.contours.data['position'].values())

    def _add_metadata(self):
        """Add useful analysis parameters etc. to the self.metadata dict.

        (later saved in the metadata json file)
        Define in subclasses
        """
        self.results.metadata['contours'] = self.contours.data

    def regenerate(self, filename=None):
        """Load saved data, metadata and regenerate objects from them.

        Is used to reset the system in a state similar to the end of the
        analysis that was made before saving the results.

        Parameters
        ----------
        - filename: name of the analysis results file (if None, use default)

        More or less equivalent to:
        analysis.results.load(filename=filename)
        image_series.load_transforms()
        (except that transforms are loaded from the metadata file of the
        analysis, not from a file generated by image_series.save_transforms())
        """

        # Load data
        super().regenerate(filename=filename)

        # regenerate internal threshold / contours object
        self.contours.load(filename=filename)

        # at the moment, this is already done by contours.load(), but I'm
        # putting this there to be sure in case contours are modified to not
        # include threshold level information
        self.threshold.load(filename=filename)
