"""Reflectance analysis: average grey level over time in img sequence."""

# Non-standard modules
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from imgbasics.cropping import _cropzone_draw

# Local imports
from .analysis_base import Analysis
from .formatters import PandasFormatter
from .results import PandasTsvResults
from ..parameters.analysis import Zones
from ..viewers import AnalysisViewer


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


class GreyLevelViewer(AnalysisViewer):

    def _create_figure(self):
        self.fig = plt.figure(figsize=(5, 7))
        xmin = 0.1
        xmax = 0.9
        w = xmax - xmin
        self.ax_img = self.fig.add_axes([xmin, 0.33, w, 0.65])
        self.ax_analysis = self.fig.add_axes([xmin, 0.09, w, 0.25])
        self.axs = self.ax_img, self.ax_analysis
        self.ax_img.axis('off')

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

        # image
        self.ax_img.set_title(f'img #{num}')
        self.imshow = self.analysis.img_series._imshow(
            img,
            ax=self.ax_img,
            **self.kwargs,
        )
        # curves
        self.curves = []
        self.pts = []

        for zone_name, glevel in zip(self.analysis.zones.data, glevels):

            if self.analysis.results.data is not None:
                # Existing or re-generated full data from analysis
                full_data = self.analysis.results.data[zone_name]

                curve, = self.ax_analysis.plot(full_data, label=zone_name)
                self.curves.append(curve)

                self.full_data_available = True

                color = curve.get_color()
                pt, = self.ax_analysis.plot(num, glevel, 'o', c=color)

            else:
                self.full_data_available = False
                pt, = self.ax_analysis.plot(num, glevel, 'o', label=zone_name)
                color = pt.get_color()

            self.pts.append(pt)

            zone = self.analysis.zones.data[zone_name]
            _cropzone_draw(self.ax_img, zone, c=color)

        self.ax_analysis.legend()
        self.ax_analysis.grid()

        self.updated_artists = self.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']
        glevels = data['glevels']

        self.ax_img.set_title(f'img #{num}')

        self.imshow.set_array(img)

        for pt, glevel in zip(self.pts, glevels):
            if self.full_data_available:
                pt.set_data((num, glevel))
            else:
                x, y = pt.get_data()
                new_x = np.append(x, num)
                new_y = np.append(y, glevel)
                pt.set_data((new_x, new_y))

        self.ax_analysis.relim()  # without this, axes limits change don't work
        self.ax_analysis.autoscale(axis='both')


class GreyLevelFormatter_Pandas(PandasFormatter):

    def _prepare_data_storage(self):
        """Prepare structure(s) that will hold the analyzed data."""
        self.glevel_data = []

    def _store_data(self, data):
        """How to store data generated by analysis on a single image."""
        self.glevel_data.append(data['glevels'])

    def _to_pandas(self):
        """How to convert data generated by _store_data() into a pandas table."""
        zone_names = self.analysis.zones.data.keys()  # 'zone 1', 'zone 2', etc.
        return pd.DataFrame(self.glevel_data, columns=zone_names)

    def _regenerate_data(self, num):
        """How to go back to raw dict of data from self.data."""
        data = {}
        try:
            data['glevels'] = list(self.analysis.results.data.filter(like='zone').loc[num])
        except AttributeError:  # if self.data not defined (analysis not made)
            pass
        return data


class GreyLevelResults_PandasTsv(PandasTsvResults):

    measurement_type = 'glevel'
    default_filename = 'Img_GreyLevel'


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


class GreyLevel(Analysis):
    """Class to perform analysis of average grey level on image series."""

    measurement_type = 'glevel'

    DefaultViewer = GreyLevelViewer
    DefaultFormatter = GreyLevelFormatter_Pandas
    DefaultResults = GreyLevelResults_PandasTsv

    def __init__(
        self,
        img_series,
        savepath=None,
        func=np.mean,
        Viewer=None,
        Formatter=None,
        Results=None,
    ):
        """Analysis of avg gray level on selected zone in series of images.

        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)

        - func: function to be applied on the image pixels in the defined
                analysis zones (default: np.mean). Other typical functions
                can be: np.sum, np.max, etc.

        - 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)
        """
        super().__init__(
            img_series=img_series,
            savepath=savepath,
            Viewer=Viewer,
            Formatter=Formatter,
            Results=Results,
        )

        # empty zones object, if not filled with zones.define() or
        # zones.load() prior to starting analysis with self.run(),
        # the whole image is considered
        self.zones = Zones(self)
        self.func = func

    def _analyze(self, img):
        """Analysis process on single image. Must return a dict.

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

        Output
        ------
        - dict of data, handled by formatter._store_data()
        """
        glevels = []

        for cropzone in self.zones.data.values():
            crop_func = self.img_series.img_transformer.img_manager.crop
            img_crop = crop_func(img, cropzone)
            glevel = self.func(img_crop)
            glevels.append(glevel)

        data = {'glevels': glevels}

        return data

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

    def _set_default_zone(self):
        print('Warning: no zones defined; taking full image as default.')
        default_crop = 0, 0, self.img_series.nx, self.img_series.ny
        self.zones.data = {'zone 1': default_crop}

    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['zones'] = self.zones.data

        func_info = f'{self.func.__module__}.{self.func.__name__}'
        self.results.metadata['function'] = func_info

    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 zones object
        self.zones.load(filename=filename)
