"""Analysis of wetting / drying fronts."""

# Non-standard modules
import pandas as pd
import matplotlib.pyplot as plt

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


class Front1DViewer(AnalysisViewer):

    def _create_figure(self):
        self.fig = plt.figure(figsize=(5, 7))
        xmin = 0.05
        xmax = 0.95
        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."""
        self.ax_img.set_title(f"img #{data['num']}")
        self.imshow = self.analysis.img_series._imshow(data['image'],
                                                       ax=self.ax_img,
                                                       **self.kwargs)

        self.analysis_line, = self.ax_analysis.plot(data['analysis'])
        self.updated_artists = [self.analysis_line, self.imshow]

    def _update_plot(self, data):
        """What to do upon iterations of the plot after the first time."""
        self.ax_img.set_title(f"img #{data['num']}")
        self.imshow.set_array(data['image'])
        self.analysis_line.set_ydata(data['analysis'])

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


class Front1DFormatter_Pandas(PandasFormatter):

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

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

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

    def _regenerate_data(self, num):
        """How to go back to raw dict of data from self.data."""
        return {'analysis': self.analysis.results.data.loc[num, 0:]}


class Front1DResults_PandasTsv(PandasTsvResults):

    measurement_type = 'front1d'
    default_filename = 'Img_Front1D'


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


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

    measurement_type = 'front1d'

    DefaultViewer = Front1DViewer
    DefaultFormatter = Front1DFormatter_Pandas
    DefaultResults = Front1DResults_PandasTsv

    def __init__(
        self,
        img_series,
        savepath=None,
        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)

        - 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,
        )

    def _analyze(self, img):
        """Basic analysis function, to be threaded or multiprocessed.

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

        Output
        ------
        - data, handled by self.formatter._store_data()
        """
        front_data = img.mean(axis=0)
        data = {'analysis': front_data}
        return data
