# SPDX-FileCopyrightText: 2020 Lukas Schrangl <lukas.schrangl@tuwien.ac.at>
#
# SPDX-License-Identifier: BSD-3-Clause

"""Provide :py:class:`Analyzer` as a Jupyter notebook UI for smFRET analysis"""
import itertools
from io import BytesIO
from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Union
try:
    from typing import Literal
except ImportError:
    # Python < 3.8
    from .typing_extensions import Literal

import ipywidgets
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy.stats

from sdt import changepoint, flatfield, fret, image, multicolor, nbui, roi

from .data_store import DataStore

class Analyzer:
    """Jupyter notebook UI for analysis of smFRET tracks

    Analyze and filter smFRET tracks produced by :py:class:`Tracker`.
    """
    rois: Dict[str, roi.ROI]
    """channel name -> ROI"""
    frame_selector: multicolor.FrameSelector
    """Select frames of different excitation types"""
    analyzers: Dict[str, fret.SmFRETAnalyzer]
    """dataset key -> analyzer class instance"""
    special_analyzers: Dict[str, fret.SmFRETAnalyzer]
    """dataset key -> analyzer class instance for special-purpose (donor-only,
    acceptor-only) datasets
    """
    sources: Dict[str, Dict[int, Union[str, Sequence[str]]]]
    """Map of dataset name (see :py:meth:`add_dataset`) -> file id -> source
    image file name(s).
    """
    special_sources: Dict[str, Dict[int, Union[str, Sequence[str]]]]
    """Map of dataset name (see :py:meth:`add_dataset`) -> file id -> source
    image file name(s). This is for special purpose datasets. Allowed keys:
    - ``"registration"`` (fiducial markers for image registration)
    - ``"donor-profile"``, ``"acceptor-profile"`` (densly labeled samples for
        determination of excitation intensity profiles)
    - ``"donor-only"``, ``"acceptor-only"`` (samples for determination of
        leakage and direct excitation correction factors, respectively)
    """
    segment_images: Dict[str, Dict[str, np.ndarray]]
    """file name -> 3D array where the first index specifies which of possible
    multiple images to use.
    """
    flatfield: Dict[str, flatfield.Corrector]
    """channel name -> flatfield correction class instance"""

    def __init__(self, file_prefix: str = "tracking",
                 reset_filters: bool = True):
        """Parameters
        ----------
        file_prefix
            Prefix for :py:class:`Tracker` save files to load. Same as
            ``file_prefix`` argument passed to :py:meth:`Tracker.save`.
        reset_filters
            If `True`, reset filters in tracking data.
        """
        ds = DataStore.load(file_prefix, loc=False)

        self.rois = ds.rois
        self.frame_selector = ds.tracker.frame_selector
        self.analyzers = {
            k: fret.SmFRETAnalyzer(v) for k, v in ds.tracks.items()}
        self.special_analyzers = {
            k: fret.SmFRETAnalyzer(v) for k, v in ds.special_tracks.items()}
        if reset_filters:
            for a in itertools.chain(self.analyzers.values(),
                                     self.special_analyzers.values()):
                a.reset_filters()
        self.sources = ds.sources
        self.special_sources = ds.special_sources
        self.segment_images = ds.segment_images
        self.flatfield = ds.flatfield

        self._segment_fig = None
        self._brightness_fig = None
        self._thresholder = None
        self._beam_shape_fig = None
        self._beam_shape_artists = [None, None]
        self._population_fig = None
        self._eff_stoi_fig = None

    def present_at_start(self, frame: Optional[int] = None):
        """Remove tracks that are not present in the beginning

        Parameters
        ----------
        frame
            Start frame number. If `None`, use first donor excitation frame
            except for the acceptor-only sample, where the first acceptor
            excitation frame is used.
        """
        if frame is None:
            e_seq = self.frame_selector.eval_seq(-1)
            frame_d = np.nonzero(e_seq == "d")[0][0]
            frame_a = np.nonzero(e_seq == "a")[0]
            # If there was no acceptor excitation…
            frame_a = frame_a[0] if frame_a.size else -1
        else:
            frame_d = frame_a = frame

        query_pat = "donor_frame == {}"
        query_d = query_pat.format(frame_d)
        for ana in self.analyzers.values():
            ana.query_particles(query_d)
        for key, ana in self.special_analyzers.items():
            if key[0] == "a" and key.endswith("only"):
                if frame_a < 0:
                    raise ValueError(
                        "excitation seq does not contain \"a\"")
                ana.query_particles(query_pat.format(frame_a))
            else:
                ana.query_particles(query_d)

    def find_beam_shape_thresh(self) -> ipywidgets.Widget:
        """Display a widget to set the laser intensity threshold

        Use :py:meth:`filter_beam_shape_region` to select only datapoints
        within the region of bright laser excitation.
        """
        if self._beam_shape_fig is None:
            fig = plt.figure()
            fig.add_subplot(1, 2, 1)
            fig.add_subplot(1, 2, 2)

            self._beam_shape_fig = fig

        ax = self._beam_shape_fig.axes

        channel_sel = ipywidgets.Dropdown(options=list(self.flatfield),
                                          description="channel")
        thresh_sel = ipywidgets.BoundedIntText(value=50, min=0, max=100,
                                               description="threshold")

        def channel_changed(change=None):
            bs = self.flatfield[channel_sel.value]
            if self._beam_shape_artists[0] is not None:
                self._beam_shape_artists[0].remove()
            self._beam_shape_artists[0] = ax[0].imshow(bs.corr_img)
            update()

        def update(change=None):
            bs = self.flatfield[channel_sel.value]
            if self._beam_shape_artists[1] is not None:
                self._beam_shape_artists[1].remove()
            self._beam_shape_artists[1] = ax[1].imshow(
                bs.corr_img * 100 > thresh_sel.value)
            self._beam_shape_fig.tight_layout()
            self._beam_shape_fig.canvas.draw()

        channel_sel.observe(channel_changed, "value")
        thresh_sel.observe(update, "value")

        channel_changed()

        return ipywidgets.VBox([channel_sel, thresh_sel,
                                self._beam_shape_fig.canvas])

    def filter_beam_shape_region(self, channel: Literal["donor", "accetpor"],
                                 thresh: float):
        """Select only datapoints within the region of bright laser excitation

        Parameters
        ----------
        channel
            Which channel's laser profile to use.
        thresh
            Threshold in percent of the amplitude
        """
        mask = self.flatfield[channel].corr_img * 100 > thresh
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.image_mask(mask, channel, reason="beam_shape")

    def flatfield_correction(self):
        """Apply flatfield correction to brightness data

        See also
        --------
        fret.SmFretAnalyzer.flatfield_correction
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.flatfield_correction(self.flatfield["donor"],
                                   self.flatfield["acceptor"])

    def calc_fret_values(self, *args, **kwargs):
        """Calculate FRET-related quantities

        Compute apparent FRET efficiencies and stoichiometries,
        the total brightness (mass) upon donor excitation, and the acceptor
        brightness (mass) upon direct excitation, which is interpolated for
        donor excitation datapoints in order to allow for calculation of
        stoichiometries.

        See also
        --------
        fret.SmFretAnalyzer.calc_fret_values
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.calc_fret_values(*args, **kwargs)

    def _make_dataset_selector(self, state: Dict,
                               callback: Callable[[Optional[Any]], None]) \
            -> ipywidgets.Dropdown:
        """Create a drop-down for selecting the dataset in UIs

        Parameters
        ----------
        state
            Keep some state of the widget. This creates/overwrites "tr",
            "pnos", and "files" entries containing the dataset's tracking data,
            unique particle numbers, and source files, respectively.
        callback
            Gets called whenever the dataset was changed.

        Returns
        -------
        The selection UI element
        """
        d_sel = ipywidgets.Dropdown(options=list(self.analyzers.keys()),
                                    description="dataset")

        def change_dataset(key=None):
            trc = self.analyzers[d_sel.value].apply_filters()
            trc = trc[trc["fret", "exc_type"] == "d"]
            nan_mask = (np.isfinite(trc["fret", "a_mass"]) &
                        np.isfinite(trc["fret", "d_mass"]))
            all_p = trc["fret", "particle"].unique()
            bad_p = trc.loc[~nan_mask, ("fret", "particle")].unique()
            state["id"] = d_sel.value
            state["trc"] = trc
            state["pnos"] = np.setdiff1d(all_p, bad_p)

            callback()

        d_sel.observe(change_dataset, names="value")
        change_dataset()

        return d_sel

    def find_changepoint_options(self) -> ipywidgets.Widget:
        """Find options for :py:meth:`segment_mass`

        Returns
        -------
        The selection UI element
        """
        state = {}

        p_sel = ipywidgets.IntText(value=0, description="particle")
        p_label = ipywidgets.Label()
        pen_a_sel = ipywidgets.FloatText(value=2e7, description="acc. penalty")
        pen_d_sel = ipywidgets.FloatText(value=2e7, description="don. penalty")

        if self._segment_fig is None:
            self._segment_fig, ax = plt.subplots(1, 3, figsize=(8, 4))
            ax[1].twinx()

        ax_dm, ax_am, ax_eff, ax_hn = self._segment_fig.axes

        def show_track(change=None):
            for a in self._segment_fig.axes:
                a.cla()

            pi = state["pnos"][p_sel.value]
            p_label.value = f"Real particle number: {pi}"
            trc = state["trc"]
            t = trc[trc["fret", "particle"] == pi]

            fd = t["donor", "frame"].values
            dm = t["fret", "d_mass"].values
            am = t["fret", "a_mass"].values
            ef = t["fret", "eff_app"].values
            hn = t["fret", "has_neighbor"].values

            cp_a = self.analyzers[state["id"]].cp_detector.find_changepoints(
                am, pen_a_sel.value) - 1
            cp_d = self.analyzers[state["id"]].cp_detector.find_changepoints(
                dm, pen_d_sel.value) - 1
            cp_d_all = np.concatenate([np.searchsorted(fd, fd[cp_a]) - 1,
                                       cp_d])
            cp_d_all = np.unique(cp_d_all)
            cp_d_all.sort()

            if len(fd):
                changepoint.plot_changepoints(dm, cp_d, time=fd, ax=ax_dm)
                changepoint.plot_changepoints(ef, cp_d_all, time=fd, ax=ax_eff)
                changepoint.plot_changepoints(am, cp_a, time=fd, ax=ax_am)
            ax_hn.plot(fd, hn, c="C2", alpha=0.2)

            ax_dm.relim(visible_only=True)
            ax_am.relim(visible_only=True)
            ax_eff.set_ylim(-0.05, 1.05)
            ax_hn.set_ylim(-0.05, 1.05)

            ax_dm.set_title("donor ex.")
            ax_am.set_title("acceptor ex.")
            ax_eff.set_title("app. eff")

            self._segment_fig.tight_layout()
            self._segment_fig.canvas.draw()

        p_sel.observe(show_track, "value")
        pen_d_sel.observe(show_track, "value")
        pen_a_sel.observe(show_track, "value")

        d_sel = self._make_dataset_selector(state, show_track)

        return ipywidgets.VBox([d_sel, p_sel, pen_d_sel, pen_a_sel,
                                self._segment_fig.canvas, p_label])

    def mass_changepoints(self, channel: Literal["donor", "acceptor"],
                          stats: Union[Callable, str,
                                       Iterable[Union[Callable, str]]
                                       ] = "median",
                          stat_margin: int = 1,
                          **kwargs):
        """Segment tracks by changepoint detection in brightness time traces

        This adds a column with segment numbers for each track.

        Parameters
        ----------
        channel
            Whether to perform changepoint detection in the donor or acceptor
            emission channel.
        stats
            Statistics to calculate for each track segment. For each entry
            ``s``, a column named ``"{chan}_seg_{s}"`` is appendend, where
            ``chan`` is ``d`` for donor and ``a`` for acceptor.
            ``s`` can be the name of a numpy function or a callable returning
            a statistic, such as :py:func:`numpy.mean`.
        stat_margin
            Number of data points around a changepoint to exclude from
            statistics calculation. This can prevent bias in the statistics due
            to recording a bleaching event in progress.
        **kwargs
            Arguments passed to the changepoint detection algorithm
            (:py:meth:`sdt.changepoint.Pelt.find_changepoints`).

        See also
        --------
        fret.SmFretAnalyzer.mass_changepoints
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.mass_changepoints(channel, stats, stat_margin, **kwargs)

    def set_bleach_thresh(self, donor, acceptor):
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.bleach_thresh = (donor, acceptor)

    def filter_bleach_step(self,
                           condition: Literal[
                               "donor", "acceptor", "donor or acceptor",
                               "no partial"] = "donor or acceptor"):
        """Remove tracks that do not show expected bleaching steps

        Acceptor must bleach in a single step, while donor must not show more
        than one step.

        Parameters
        ----------
        condition
            If ``"donor"``, accept only tracks where the donor bleaches in a
            single step and the acceptor shows either no bleach step or
            completely bleaches in a single step.
            Likewise, ``"acceptor"`` will accept only tracks where the acceptor
            bleaches fully in one step and the donor shows no partial
            bleaching.
            ``donor or acceptor`` requires that one channel bleaches in a
            single step while the other either also bleaches in one step or not
            at all (no partial bleaching).
            If ``"no partial"``, there may be no partial bleaching, but
            bleaching is not required.

        See also
        --------
        fret.SmFretAnalyzer.bleach_step
        """
        for a in self.analyzers.values():
            a.bleach_step(condition)
        for k, a in self.special_analyzers.items():
            if k.endswith("only"):
                # This is the best we can do for donor-only and acceptor-only
                a.bleach_step("no partial")
            else:
                a.bleach_step(condition)

    def set_leakage(self, leakage: float):
        r"""Set the leakage correction factor for all datasets

        Parameters
        ----------
        leakage
            Leakage correction factor (often denoted :math:`\alpha`)

        See also
        --------
        fret.SmFretAnalyzer.leakage
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.leakage = leakage

    def calc_leakage(self):
        """Calculate leakage correction factor from donor-only sample

        This uses the ``"donor-only"`` special dataset (see
        :py:meth:`fret.SmFRETTracker.add_special_dataset`).

        See also
        --------
        calc_leakage_from_bleached
        fret.SmFRETAnalyzer.calc_leakage
        """
        a = self.special_analyzers["donor-only"]
        a.calc_leakage()
        self.set_leakage(a.leakage)

    # TODO: Move to sdt package
    def calc_leakage_from_bleached(
            self, datasets: Union[str, Sequence[str], None] = None,
            stat: str = "median", print_summary: bool = True):
        """Calculate leakage correction factor from bleached acceptor traces

        This takes those parts of traces where the acceptor is bleached, but
        the donor isn't.

        Parameters
        ----------
        datasets
            dataset(s) to use. If `None`, use all.
        stat
            Statistic to use to determine bleaching steps. Has to be one that
            was passed to via ``stats`` parameter to
            :py:meth:`mass_changepoints`.
        print_summary
            Print number of datapoints and result.

        See also
        --------
        calc_leakage
        fret.SmFretAnalyzer.calc_leakage
        """
        if datasets is None:
            datasets = list(self.analyzers)
        elif isinstance(datasets, str):
            datasets = [datasets]

        effs = []
        for d in datasets:
            trc = self.analyzers[d].tracks
            bleach_tresh = self.analyzers[d].bleach_thresh[1]
            # search for donor excitation frames where acceptor is bleached
            # but donor isn't
            sel = ((trc["fret", "exc_type"] == "d") &
                   (trc["fret", "has_neighbor"] == 0) &
                   (trc["fret", "a_seg"] > 1) &
                   (trc["fret", f"a_seg_{stat}"] < bleach_tresh) &
                   (trc["fret", "d_seg"] == 0))
            effs.append(trc.loc[sel, ("fret", "eff_app")].values)
        effs = np.concatenate(effs)

        m_eff = np.mean(effs)
        leakage = m_eff / (1 - m_eff)

        self.set_leakage(leakage)

        if print_summary:
            n_data = len(effs)
            e_eff = np.std(effs, ddof=1) / np.sqrt(n_data)  # SEM
            err = 1 / (1 - m_eff)  # derivative of x / (1-x) is 1 / (1-x)**2
            err = err * err * e_eff
            print(f"leakage: {leakage:.4f} ± {err:.4f} (from {n_data} "
                  "datapoints)")

    def set_direct_excitation(self, dir_exc: float):
        r"""Set the direct (cross-) excitation factor for all datasets

        Parameters
        ----------
        dir_exc
            Direct excitation correction factor (often denoted :math:`\delta`)

        See also
        --------
        fret.SmFretAnalyzer.direct_excitation
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.direct_excitation = dir_exc

    def calc_direct_excitation(self):
        """Calculate direct excitation correction factor from acc-only sample

        This uses the ``"acceptor-only"`` special dataset (see
        :py:meth:`fret.SmFRETTracker.add_special_dataset`).

        See also
        --------
        calc_direct_excitation_from_bleached
        fret.SmFRETAnalyzer.calc_direct_excitation
        """
        a = self.special_analyzers["acceptor_only"]
        a.calc_direct_excitation()
        self.set_direct_excitation(a.direct_excitation)

    # TODO: Move to sdt package
    def calc_direct_excitation_from_bleached(
            self, datasets: Union[str, Sequence[str], None] = None,
            stat: str = "median", print_summary: bool = True):
        """Calculate dir. exc. correction factor from bleached donor traces

        This takes those parts of traces where the donor is bleached, but
        the acceptor isn't.

        Parameters
        ----------
        datasets
            dataset(s) to use. If `None`, use all.
        stat
            Statistic to use to determine bleaching steps. Has to be one that
            was passed to via ``stats`` parameter to
            :py:meth:`mass_changepoints`.
        print_summary
            Print number of datapoints and result.

        See also
        --------
        calc_direct_excitation
        fret.SmFretAnalyzer.calc_direct_excitation
        """
        if datasets is None:
            datasets = list(self.analyzers)
        elif isinstance(datasets, str):
            datasets = [datasets]

        stois = []
        for d in datasets:
            trc = self.analyzers[d].tracks
            bleach_tresh = self.analyzers[d].bleach_thresh[0]
            # search for donor excitation frames where donor is bleached
            # but acceptor isn't
            sel = ((trc["fret", "exc_type"] == "d") &
                   (trc["fret", "has_neighbor"] == 0) &
                   (trc["fret", "a_seg"] == 0) &
                   (trc["fret", "d_seg"] > 0) &
                   (trc["fret", f"d_seg_{stat}"] < bleach_tresh))
            stois.append(trc.loc[sel, ("fret", "stoi_app")].values)
        stois = np.concatenate(stois)

        m_stoi = np.mean(stois)
        dir_exc = m_stoi / (1 - m_stoi)

        self.set_direct_excitation(dir_exc)

        if print_summary:
            n_data = len(stois)
            e_stoi = np.std(stois, ddof=1) / np.sqrt(n_data)  # SEM
            err = 1 / (1 - m_stoi)  # derivative of x / (1-x) is 1 / (1-x)**2
            err = err * err * e_stoi
            print(f"direct exc.: {dir_exc:.4f} ± {err:.4f} (from {n_data} "
                  "datapoints)")

    def set_detection_eff(self, eff: float):
        r"""Set the detection efficiency factor for all datasets

        Parameters
        ----------
        eff
            Detection efficiency correction factor (often denoted
            :math:`\gamma`)

        See also
        --------
        sdt.fret.SmFretAnalyzer.detection_eff
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.detection_eff = eff

    def calc_detection_eff(self, min_part_len: int = 5,
                           how: Union[Callable[[np.array], float],
                                      Literal["individual"]] = np.nanmedian,
                           aggregate: Literal["dataset", "all"] = "dataset",
                           dataset: Optional[str] = None):
        r"""Calculate detection efficieny correction factor

        The detection efficiency ratio is the ratio of decrease in acceptor
        brightness to the increase in donor brightness upon acceptor
        photobleaching.

        Parameters
        ----------
        min_part_len
            How many data points need to be present before and after the
            bleach step to ensure a reliable calculation of the mean
            intensities. If there are fewer data points, a value of NaN will be
            assigned.
        how
            If "individual", the :math:`\gamma` value for each track will be
            stored and used to correct the values individually when calling
            :py:meth:`fret_correction`. If a function, apply this function
            to the :math:`\gamma` array and its return value as a global
            correction factor. A sensible example for such a function would be
            :py:func:`numpy.nanmean`. Beware that some :math:`\gamma` may be
            NaN.
        aggregate
            Whether to use a callable `how` argument to calculate the
            correction for each dataset individually or a single factor for all
            datasets. Ignored if `how` is not callable or `dataset` is not
            `None`.
        dataset
            If `how` is callable, use the given dataset to compute a global
            correction factor for all datasets.

        See also
        --------
        sdt.fret.SmFretAnalyzer.calc_detection_eff
        """
        if callable(how):
            if dataset is not None:
                a = self.analyzers[dataset]
                a.calc_detection_eff(min_part_len, how)
                self.set_detection_eff(a.detection_eff)
            elif aggregate == "dataset":
                for a in self.analyzers.values():
                    a.calc_detection_eff(min_part_len, how)
            elif aggregate == "all":
                effs = []
                for a in self.analyzers.values():
                    a.calc_detection_eff(min_part_len, "individual")
                    effs.append(a.detection_eff)
                self.set_detection_eff(how(pd.concat(effs)))
            else:
                raise ValueError(
                    "`aggregate` must be in {\"dataset\", \"all\"}.")
        elif how == "individual":
            for a in self.analyzers.values():
                a.calc_detection_eff(min_part_len, how)
        else:
            raise ValueError(
                "`how` must be either callable or \"individual\".")

    def set_excitation_eff(self, eff: float):
        r"""Set the excitation efficiency factor for all datasets

        Parameters
        ----------
        eff
            Excitation efficiency correction factor (often denoted
            :math:`\beta`)

        See also
        --------
        sdt.fret.SmFretAnalyzer.excitation_eff
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.excitation_eff = eff

    def find_subpopulations(self) -> ipywidgets.Widget:
        """Interactively select component for calculation of excitation eff.

        The excitation efficiency factor is computer from data with known 1:1
        stoichiomtry. This displays datasets and allows for fitting Gaussian
        mixture models to select the right component in an E–S plot. Parameters
        can be passed to :py:meth:`calc_excitation_eff`.

        Returns
        -------
        UI element

        See also
        --------
        calc_excitation_eff
        """
        if self._population_fig is None:
            self._population_fig = plt.subplots()[0]

        state = {}

        n_comp_sel = ipywidgets.IntText(description="components", value=1)

        def update(change=None):
            d = state["trc"]
            d = d[(d["fret", "a_seg"] == 0) &
                  (d["fret", "d_seg"] == 0) &
                  np.isfinite(d["fret", "eff_app"]) &
                  np.isfinite(d["fret", "stoi_app"])].copy()
            labels = fret.gaussian_mixture_split(d, n_comp_sel.value)[0]

            ax = self._population_fig.axes[0]
            ax.cla()
            for lab in range(n_comp_sel.value):
                dl = d[labels == lab]
                sel = (np.isfinite(dl["fret", "eff_app"]) &
                       np.isfinite(dl["fret", "stoi_app"]))
                ax.scatter(dl.loc[sel, ("fret", "eff_app")],
                           dl.loc[sel, ("fret", "stoi_app")],
                           marker=".", label=str(lab), alpha=0.6)
            ax.legend(loc=0, title="index")
            ax.set_xlabel("apparent eff.")
            ax.set_ylabel("apparent stoi.")
            self._population_fig.canvas.draw()

        d_sel = self._make_dataset_selector(state, update)
        n_comp_sel.observe(update, "value")

        return ipywidgets.VBox([d_sel, n_comp_sel,
                                self._population_fig.canvas])

    def calc_excitation_eff(self, dataset: str, n_components: int = 1,
                            component: int = 0):
        """Calculate excitation efficieny correction factor

        Utilize data with known 1:1 stoichiometry to this end. To find the
        right parameters for this call, use
        :py:meth:`find_excitation_eff_component`.

        Parameters
        ----------
        dataset
            Dataset to use.
        n_components
            Perform an n-component Gaussian mixture model fit to separate
            subpopulations in E–S space.
        component
            Identifier of the component to use. They are sorted in descending
            order w.r.t. mean stoichiometry.

        See also
        --------
        sdt.fret.gaussian_mixture_split
        sdt.fret.SmFretAnalyzer.calc_excitation_eff
        """
        a = self.analyzers[dataset]
        a.calc_excitation_eff(n_components, component)
        self.set_excitation_eff(a.excitation_eff)

    def calc_detection_excitation_effs(
            self, n_components: int,
            components: Optional[Sequence[int]] = None,
            dataset: Optional[str] = None):
        r"""Get detection and excitation efficiency from multi-state sample

        States are found in efficiency-vs.-stoichiometry space using a
        Gaussian mixture fit. Detection efficiency factor :math:`\gamma` and
        excitation efficiency factor :math:`\delta` are found performing a
        linear fit to the results [Hell2018]_.

        Parameters
        ----------
        n_components
            Number of components for Gaussian mixture model
        components
            List of indices of components to use for the linear fit. If `None`,
            use all.
        dataset
            Name of the dataset to use. If `None` use the ``"multi-state"``
            special dataset.
        """
        if dataset is None:
            a = self.special_analyzers["multi-state"]
        else:
            a = self.analyzers[dataset]
        a.calc_detection_excitation_effs(n_components, components)
        self.set_detection_eff(a.detection_eff)
        self.set_excitation_eff(a.excitation_eff)

    def fret_correction(self, *args, **kwargs):
        """Apply corrections to calculate real FRET-related values

        Parameters
        ----------
        *args, **kwargs
            Passed to :py:meth:`sdt.fret.SmFretAnalyzer.fret_correction`

        See also
        --------
        sdt.fret.SmFretAnalyzer.fret_correction
        """
        for a in self.analyzers.values():
            a.fret_correction(*args, **kwargs)

    def query(self, expr: str, mi_sep: str = "_", reason: str = "query"):
        """Filter localizations

        Use :py:meth:`sdt.fret.SmFretAnalyzer.query` to select only
        datapoints (i.e., lines in tracking tables) that fulfill `expr`.

        Parameters
        ----------
        expr
            Filter expression. See :py:meth:`pandas.DataFrame.eval` for
            details.
        mi_sep
            Use this to separate levels when flattening the column
            MultiIndex. Defaults to "_".

        See also
        --------
        pandas.DataFrame.eval
        sdt.fret.SmFretAnalyzer.query
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.query(expr, mi_sep, reason)

    def query_particles(self, expr: str, min_abs: int = 1,
                        min_rel: float = 0.0, mi_sep: str = "_",
                        reason: str = "query_p"):
        """Filter tracks

        Use :py:meth:`sdt.fret.SmFretAnalyzer.query_particles` to select only
        trajectories that fulfill `expr` sufficiently many times.

        Parameters
        ----------
        expr
            Filter expression. See :py:meth:`pandas.DataFrame.eval` for
            details.
        min_abs
            Minimum number of times a particle has to fulfill `expr`. If
            negative, this means "all but ``abs(min_count)``". If 0, it has
            to be fulfilled in all frames.
        min_rel
            Minimum fraction of data points that have to fulfill `expr` for a
            particle not to be removed
        mi_sep
            Use this to separate levels when flattening the column
            MultiIndex. Defaults to "_".

        See also
        --------
        pandas.DataFrame.eval
        sdt.fret.SmFretAnalyzer.query_particles
        """
        for a in itertools.chain(self.analyzers.values(),
                                 self.special_analyzers.values()):
            a.query_particles(expr, min_abs, min_rel, mi_sep, reason)

    def find_segmentation_options(self) -> nbui.Thresholder:
        """UI for finding parameters for segmenting images

        Parameters can be passed to :py:meth:`apply_segmentation`.

        Returns
        -------
            Widget instance.
        """
        if self._thresholder is None:
            self._thresholder = nbui.Thresholder()

        self._thresholder.image_selector.images = {
            self.sources[did][fid]: v[0]
            for did, dset in self.segment_images.items()
            for fid, v in dset.items()}

        return self._thresholder

    def apply_segmentation(self, keys: Sequence[str],
                           thresh_algorithm: str = "adaptive",
                           channel: str = "donor", **kwargs):
        """Remove datapoints from regions not selected by segmentation

        Threshold cell images according to the parameters and use the resulting
        mask to discard datapoints outside of thresholded regions.

        Parameters
        ----------
        keys
            Which datasets to apply cell masks to.
        channel
            If "donor", apply filter to donor coordinates. If "acceptor", apply
            mask to acceptor coordinates.
        thresh_algorithm
            Use the ``thresh_algorithm + "_thresh"`` function from
            :py:mod:`sdt.image` for thresholding.
        **kwargs
            Passed to the thresholding function.
        """
        if isinstance(thresh_algorithm, str):
            thresh_algorithm = getattr(image, thresh_algorithm + "_thresh")

        for k in keys:
            # TODO: maybe move this to sdt-python?
            ana_ = self.analyzers[k]
            fids = np.unique(
                ana_.tracks.index.remove_unused_levels().levels[0])

            max_frame = int(ana_.tracks["donor", "frame"].max()) + 1
            c_frames = self.frame_selector.find_other_frames(
                max_frame, "da", "s", "previous")
            c_frames = self.frame_selector.renumber_frames(c_frames, "s")
            change_idx = np.nonzero(np.diff(c_frames))[0] + 1
            if len(change_idx) > 0:
                change_idx = self.frame_selector.renumber_frames(
                    change_idx, "da", restore=True)
            start_idx = np.empty(len(change_idx)+1, dtype=int)
            start_idx[0] = 0
            start_idx[1:] = change_idx
            stop_idx = np.empty_like(start_idx)
            stop_idx[:-1] = change_idx
            stop_idx[-1] = max_frame

            masks = []
            for f in fids:
                ci = self.segment_images[k][f]
                masks.extend(
                    [{"key": f, "mask": thresh_algorithm(c, **kwargs),
                      "start": s, "stop": e}
                     for c, s, e in zip(ci, start_idx, stop_idx)])
            ana_.image_mask(masks, channel, reason="image_seg")

    def save(self, file_prefix: str = "tracking"):
        """Save results to disk

        This will save filtered data to disk.

        Parameters
        ----------
        file_prefix
            Prefix for the file written by this method. It will be suffixed by
            the output format version (v{output_version}) and file extension.
        """
        DataStore(tracks={k: ana.tracks for k, ana in self.analyzers.items()},
                  special_tracks={k: ana.tracks
                                  for k, ana in self.special_analyzers.items()}
                  ).save(file_prefix, mode="update")

class DensityPlots(ipywidgets.Box):
    """A widget for displaying density plots

    With this, one can plot data after individual filtering steps to see
    the progress
    """
    def __init__(self, analyzer: Analyzer, datasets: Sequence[str],
                 columns: Sequence[Sequence[str]] = [("fret", "eff_app"),
                                                     ("fret", "stoi_app")],
                 bounds: Sequence[Sequence[float]] = [(-0.5, 1.5),
                                                      (-0.5, 1.5)],
                 direction: Literal["horizontal", "vertical"] = "vertical",
                 **kwargs):
        """Parameters
        ----------
        analyzer
            :py:class:`Analyzer` instance to get data from.
        datasets
            Which datasets to plot
        columns:
            Dataframe column names to plot
        bounds
            Plot only data within bounds
        direction
            In which direction to add new plots. Does not work (yet).
        **kwargs
            Passed to :py:class:`ipywidgets.Box` constructor.
        """
        super().__init__(**kwargs)
        la = self.layout or ipywidgets.Layout()
        la.overflow = la.overflow or "auto"
        la.display = la.display or "block"
        if direction.startswith("h"):
            la.flex_flow = "row nowrap"
            la.width = la.width or "100%"
        else:
            la.flex_flow = "column nowrap"
            la.height = la.height or "100%"
        self.layout = la

        self._ana = analyzer
        self._datasets = datasets
        self._cols = columns

        self._bounds = bounds

    @staticmethod
    def _fmt_axis_label(label: Union[Any, Sequence[Any]]) -> str:
        """Format (multi-index) column name for axis label display

        Parameters
        ----------
        label
            Single- or multi-index column name

        Returns
        -------
        Comma-separated parts of the name
        """
        if not isinstance(label, (tuple, list)):
            label = (label,)
        return ", ".join(str(la) for la in label)

    def _do_plot(self, title: str, plot_func: Callable):
        """Draw a plot

        This will get the current data from the Analyzer and create a plot,
        which is appended to the Box widget.

        Parameters
        ----------
        title
            Title of the plot
        plot_func
            Uses a :py:class:`matplotlib.axes.Axes` (first argument) to plot
            x–y data (second and third args).
        """
        n_plots = len(self._datasets)
        fig = plt.Figure(constrained_layout=True, figsize=(4 * n_plots, 4))
        grid = fig.add_gridspec(1, n_plots)
        ax = []
        for i, g in enumerate(grid):
            ax.append(fig.add_subplot(g, sharex=None if not ax else ax[0],
                                      sharey=None if not ax else ax[0]))

        for d, a in zip(self._datasets, ax):
            t = self._ana.analyzers[d].apply_filters()
            x = t[self._cols[0]].to_numpy()
            y = t[self._cols[1]].to_numpy()

            fin = np.isfinite(x) & np.isfinite(y)
            x = x[fin]
            y = y[fin]

            bounds = ((self._bounds[0][0] < x) & (x < self._bounds[0][1]) &
                      (self._bounds[1][0] < y) & (y < self._bounds[1][1]))
            x = x[bounds]
            y = y[bounds]

            plot_func(a, x, y)

            a.set_xlabel(self._fmt_axis_label(self._cols[0]))
            a.set_ylabel(self._fmt_axis_label(self._cols[1]))
            a.set_title(d)

        for a in ax[1:]:
            a.tick_params(left=False, labelleft=False)

        fig.suptitle(title)
        bio = BytesIO()
        fig.savefig(bio, format="png")

        self.children = [*self.children,
                         ipywidgets.Image(value=bio.getvalue())]

    def dscatter(self, title: Optional[str] = None):
        """Scatter plot with datapoints colored according to density

        Parameters
        ----------
        title
            Plot supertitle
        """
        self._do_plot(title, self._do_scatter)

    def contourf(self, title: Optional[str] = None, n_levels: int = 15):
        """Contour plot contours according to density

        Parameters
        ----------
        title
            Plot supertitle
        n_levels
            Number of contour levels
        """
        self._do_plot(title,
                      lambda a, x, y: self._do_contour(a, x, y, n_levels))

    def colormesh(self, title: Optional[str] = None):
        """Colored plot of density

        Parameters
        ----------
        title
            Plot supertitle
        """
        self._do_plot(title, self._do_mesh)

    @staticmethod
    def _calc_kde(kde_x: np.ndarray, kde_y: np.ndarray, x: np.ndarray,
                  y: np.ndarray, n_samp: int = 1000,
                  random_state: Optional[np.random.RandomState] = None):
        """Do a Gaussian KDE on data

        Density estimate is calculated from a subset of the data.

        Parameters
        ----------
        kde_x, kde_y
            x and y values to use for KDE
        x, y
            x and y values to evaluate KDE at
        n_samp
            Number datapoints to use from kde_x and kde_y
        random_state
            This instance is used to draw `n_samp` datapoints from ``kde_x``
            and ``kde_y``. If `None`, create a new instance.

        Returns
        -------
        Density estimate evaluated at each point (x[i], y[i]).
        """
        if random_state is None:
            random_state = np.random.RandomState()
        if n_samp < len(kde_x):
            idx = random_state.choice(len(kde_x), replace=False, size=n_samp)
        else:
            idx = slice(None, None)
        kde = scipy.stats.gaussian_kde(np.array([kde_x[idx], kde_y[idx]]))

        dens = np.empty(len(x), dtype=float)
        max_data = 100000  # kde crashes if too many data at once
        for s in range(0, len(x), max_data):
            e = s + max_data
            dens[s:e] = kde(np.array([x[s:e], y[s:e]]))
        return dens

    def _do_scatter(self, a: mpl.axes.Axes, x: np.ndarray, y: np.ndarray):
        """Create density scatter plot

        Parameters
        ----------
        a
            Axes to draw on
        x, y
            data
        """
        dens = self._calc_kde(x, y, x, y)
        sort_idx = np.argsort(dens)
        dens = dens[sort_idx]
        x = x[sort_idx]
        y = y[sort_idx]

        a.scatter(x, y, c=dens, cmap="viridis", marker=".")

    def _do_contour(self, a: mpl.axes.Axes, x: np.ndarray, y: np.ndarray,
                    n_levels: int):
        """Create density contour plot

        Parameters
        ----------
        a
            Axes to draw on
        x, y
            data
        n_levels
            Number of contour levels
        """
        contour_grid = np.array(
            np.meshgrid(np.linspace(*self._bounds[0], 100),
                        np.linspace(*self._bounds[1], 100)))
        contour_grid_flat = np.reshape(contour_grid, (2, -1))
        a.contourf(*contour_grid,
                   self._calc_kde(x, y, *contour_grid_flat).reshape(
                      contour_grid.shape[1:]),
                   cmap=mpl.cm.get_cmap("viridis", n_levels),
                   levels=n_levels)

    def _do_mesh(self, a: mpl.axes.Axes, x: np.ndarray, y: np.ndarray):
        """Create density plot

        Parameters
        ----------
        a
            Axes to draw on
        x, y
            data
        """
        grid = np.array(
            np.meshgrid(np.linspace(*self._bounds[0], 200),
                        np.linspace(*self._bounds[1], 200)))
        grid_flat = np.reshape(grid, (2, -1))
        c = self._calc_kde(x, y, *grid_flat).reshape(grid.shape[1:])
        a.pcolormesh(*grid, c)


import re
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo
import torch.optim as optim
import torch.utils.checkpoint as cp

import numpy as np
import PIL.Image as pimg

import layers
import losses
import utils
import evaluation
import data.data_utils as data_utils
import data.transform as transform

evaluate = evaluation.evaluate_semseg

use_pyl_in_spp = True
checkpoint_stem = False
checkpoint_upsample = False
use_dws_up = False
use_dws_down = False

first_stride = 2
avg_pooling_k = 2
upsample = lambda x, size: F.interpolate(x, size, mode='bilinear', align_corners=False)

checkpoint = lambda func, *inputs: cp.checkpoint(func, *inputs, preserve_rng_state=True)

def _batchnorm_factory(num_maps, momentum):
    return nn.BatchNorm2d(num_maps, eps=1e-5, momentum=momentum)

use_batchnorm = True
drop_rate = 0.0

def build(depth, pretrained, **kwargs):
    global imagenet_init
    imagenet_init = pretrained
    if depth == 32:
        return densenet32(**kwargs)
    if depth == 60:
        return densenet60(**kwargs)
    if depth == 121:
        return densenet121(**kwargs)
    elif depth == 169:
        return densenet169(**kwargs)
    elif depth == 201:
        return densenet201(**kwargs)
    elif depth == 161:
        return densenet161(**kwargs)
    else:
        raise ValueError('Unknown model variant')

__all__ = ['DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161']

model_urls = {
    'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth',
    'densenet169': 'https://download.pytorch.org/models/densenet169-b2777c0a.pth',
    'densenet201': 'https://download.pytorch.org/models/densenet201-c1103571.pth',
    'densenet161': 'https://download.pytorch.org/models/densenet161-8d451a50.pth',
}

def densenet(**kwargs):
    r"""Densenet-121 model from
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = DenseNet(num_init_features=64, **kwargs)
    return model

def _load_imagenet_weights(name):
    params = model_zoo.load_url(model_urls[name])
    delete_keys = ['features.norm5.weight', 'features.norm5.bias',
    'features.norm5.running_mean', 'features.norm5.running_var',
    'classifier.weight', 'classifier.bias']
    for key in delete_keys:
        del params[key]
    return params

def densenet_small(**kwargs):
    assert imagenet_init == 0
    model = DenseNet(num_init_features=64, growth_rate=64, bn_size=2,
                     block_config=(3, 6, 12, 8), **kwargs)
    return model

def densenet32(**kwargs):
    assert imagenet_init == 0
    model = DenseNet(num_init_features=64, growth_rate=64, bn_size=2,
                     block_config=(2, 4, 6, 4), **kwargs)
    return model

def densenet60(**kwargs):
    assert imagenet_init == 0
    model = DenseNet(num_init_features=64, growth_rate=64, bn_size=2,
                     block_config=(3, 6, 12, 8), **kwargs)
    return model

def densenet121(**kwargs):
    r"""Densenet-121 model from
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
            imagenet_init (bool): If True, returns a model pre-trained on ImageNet
    """
    if imagenet_init:
        model = DenseNet(num_init_features=64, growth_rate=32,
                         block_config=(6, 12, 24, 16), **kwargs)
        params = _load_imagenet_weights('densenet121')
        model.load_state_dict(params)
    else:
        model = DenseNet(num_init_features=64, growth_rate=32,
                         block_config=(6, 12, 24, 16), drop_rate=drop_rate, **kwargs)
    return model

def densenet169(**kwargs):
    r"""Densenet-169 model from
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
            imagenet_init (bool): If True, returns a model pre-trained on ImageNet
    """
    if imagenet_init:
        model = DenseNet(num_init_features=64, growth_rate=32,
                         block_config=(6, 12, 32, 32), **kwargs)
        params = _load_imagenet_weights('densenet169')
        model.load_state_dict(params)
    else:
        model = DenseNet(num_init_features=64, growth_rate=32,
                         block_config=(6, 12, 32, 32), drop_rate=0.2, **kwargs)
    return model

def densenet201(**kwargs):
    r"""Densenet-201 model from
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
            imagenet_init (bool): If True, returns a model pre-trained on ImageNet
    """
    assert imagenet_init
    if imagenet_init:
        model = DenseNet(num_init_features=64, growth_rate=32,
                         block_config=(6, 12, 48, 32), **kwargs)
        params = _load_imagenet_weights('densenet201')
        model.load_state_dict(params)
    return model

def densenet161(**kwargs):
    r"""Densenet-161 model from
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
            imagenet_init (bool): If True, returns a model pre-trained on ImageNet
    """
    if imagenet_init:
        model = DenseNet(num_init_features=96, growth_rate=48,
                         block_config=(6, 12, 36, 24), **kwargs)
        params = _load_imagenet_weights('densenet161')
        model.load_state_dict(params)
    else:
        model = DenseNet(num_init_features=96, growth_rate=48,
                         block_config=(6, 12, 36, 24), drop_rate=0.2, **kwargs)

    return model

def _checkpoint_unit_nobt(bn, relu, conv):
    def func(*x):
        x = torch.cat(x, 1)
        return conv(relu(bn(x)))
    return func

def _checkpoint_unit(bn1, relu1, conv1, bn2, relu2, conv2):
    def func(*x):
        x = torch.cat(x, 1)
        x = conv1(relu1(bn1(x)))
        return conv2(relu2(bn2(x)))
    return func

def _checkpoint_block_func(block):
    def func(x):
        return block(x)
    return func

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 padding=1, dilation=1, bias=False):
        super(SeparableConv2d, self).__init__()
        assert kernel_size > 1
        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size, stride,
                              padding, dilation, groups=in_channels, bias=False)
        self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias)

    def forward(self,x):
        x = self.conv(x)
        x = self.pointwise(x)
        return x

class _DenseLayer(nn.Sequential):
    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate,
                 dilation=1, checkpointing=False):
        super(_DenseLayer, self).__init__()
        m = batchnorm_momentum_ckpt if checkpointing else batchnorm_momentum

        if not use_dws_down:
            bottleneck_size = bn_size * growth_rate
            if imagenet_init or num_input_features > bottleneck_size:
                if use_batchnorm:
                    self.add_module('norm1', _batchnorm_factory(num_input_features, m))
                self.add_module('relu1', nn.ReLU(inplace=True)),
                self.add_module('conv1', nn.Conv2d(num_input_features, bottleneck_size,
                                kernel_size=1, stride=1, bias=False))
                num_feats = bottleneck_size
            else:
                num_feats = num_input_features
            if use_batchnorm:
                self.add_module('norm2', _batchnorm_factory(num_feats, m))
            self.add_module('relu2', nn.ReLU(inplace=True)),
            self.add_module('conv2', nn.Conv2d(num_feats, growth_rate, kernel_size=3,
                            stride=1, padding=dilation, bias=False, dilation=dilation))
        else:
            num_feats = num_input_features
            if use_batchnorm:
                self.add_module('norm2', _batchnorm_factory(num_feats, m))
            self.add_module('relu2', nn.ReLU(inplace=True)),
            self.add_module('conv2', SeparableConv2d(num_feats, growth_rate, kernel_size=3,
                            stride=1, padding=dilation, bias=False, dilation=dilation))

        self.drop_rate = drop_rate
        self.checkpointing = checkpointing
        if checkpointing:
            if len(self) == 6:
                self.conv_func = _checkpoint_unit(self.norm1, self.relu1, self.conv1,
                                                  self.norm2, self.relu2, self.conv2)
            else:
                self.conv_func = _checkpoint_unit_nobt(self.norm2, self.relu2, self.conv2)

    def forward(self, x):
        if self.checkpointing:
            if self.training:
                x = checkpoint(self.conv_func, *x)
            else:
                x = self.conv_func(*x)
        else:
            x = super(_DenseLayer, self).forward(x)

        if self.drop_rate > 0:
            x = F.dropout(x, p=self.drop_rate, training=self.training, inplace=True)

        return x

class _DenseBlock(nn.Sequential):
    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate,
                 split=False, dilation=1, checkpointing=True):
        super(_DenseBlock, self).__init__()
        self.checkpointing = checkpointing
        self.split = split

        for i in range(num_layers):
            layer = _DenseLayer(
                num_input_features + i * growth_rate, growth_rate=growth_rate, bn_size=bn_size,
                drop_rate=drop_rate, dilation=dilation, checkpointing=checkpointing)
            self.add_module('denselayer%d' % (i + 1), layer)
        if split:
            self.split_size = num_input_features + (num_layers // 2) * growth_rate
            k = avg_pooling_k
            pad = (k-1) // 2
            self.pool_func = lambda x: F.avg_pool2d(x, k, 2, padding=pad, ceil_mode=False,
                                                    count_include_pad=False)

    def forward(self, x):
        if self.checkpointing:
            x = [x]
        for i, layer in enumerate(self.children()):
            if self.split and len(self) // 2 == i:
                if self.checkpointing:
                    split = torch.cat(x, 1)
                    x = [self.pool_func(split)]
                else:
                    split = x
                    x = self.pool_func(split)
            if self.checkpointing:
                x.append(layer(x))
            else:
                x = torch.cat([x, layer(x)], 1)
        if self.checkpointing:
            x = torch.cat(x, 1)
        if self.split:
            return x, split
        else:
            return x

class _Transition(nn.Sequential):
    @staticmethod
    def _checkpoint_function(bn, relu, conv, pool):
        def func(inputs):
            return pool(conv(relu(bn(inputs))))
        return func

    def __init__(self, num_input_features, num_output_features, stride=2, checkpointing=False):
        super(_Transition, self).__init__()
        self.stride = stride
        if use_batchnorm:
            m = batchnorm_momentum_ckpt if checkpointing else batchnorm_momentum
            self.add_module('norm', _batchnorm_factory(num_input_features, m))
        self.add_module('relu', nn.ReLU(inplace=True))
        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,
                                          kernel_size=1, stride=1, bias=False))
        if stride > 1:
            if avg_pooling_k == 2:
                self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=stride))
            elif avg_pooling_k == 3:
                self.add_module('pool', nn.AvgPool2d(kernel_size=3, stride=stride,
            					padding=1, ceil_mode=False, count_include_pad=False))
        else:
            self.pool = lambda x: x
        self.checkpointing = checkpointing
        if checkpointing:
            self.func = _Transition._checkpoint_function(self.norm, self.relu, self.conv, self.pool)

    def forward(self, x):
        if self.checkpointing and self.training:
            return checkpoint(self.func, x)
        else:
            return super(_Transition, self).forward(x)

class _BNReluConv(nn.Sequential):
    @staticmethod
    def _checkpoint_function(bn, relu, conv):
        def func(inputs):
            return conv(relu(bn(inputs)))
        return func

    def __init__(self, num_maps_in, num_maps_out, k=3, output_conv=False,
                 dilation=1, drop_rate=0, checkpointing=False):
        super(_BNReluConv, self).__init__()
        self.drop_rate = drop_rate
        if use_batchnorm:
            m = batchnorm_momentum_ckpt if checkpointing else batchnorm_momentum
            self.add_module('norm', _batchnorm_factory(num_maps_in, m))
        self.add_module('relu', nn.ReLU(inplace=True))
        padding = ((k-1) // 2) * dilation
        if k >= 3 and use_dws_up:
            self.add_module('conv', SeparableConv2d(num_maps_in, num_maps_out, kernel_size=k,
                padding=padding, bias=output_conv, dilation=dilation))
        else:
            self.add_module('conv', nn.Conv2d(num_maps_in, num_maps_out, kernel_size=k,
                            padding=padding, bias=output_conv, dilation=dilation))
        self.checkpointing = checkpointing
        if checkpointing:
            self.func = _BNReluConv._checkpoint_function(self.norm, self.relu, self.conv)

    def forward(self, x):
        if self.checkpointing and self.training:
            x = checkpoint(self.func, x)
        else:
            x = super(_BNReluConv, self).forward(x)
        return x

class SequentialPrint(nn.Sequential):
    def forward(self, input):
        for module in self._modules.values():
            input = module(input)
            print(module)
            print(input.size())
        return input

class DenseNet(nn.Module):
    r"""Densenet-BC model class, based on
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Args:
            growth_rate (int) - how many filters to add each layer (`k` in paper)
            block_config (list of 4 ints) - how many layers in each pooling block
            num_init_features (int) - the number of filters to learn in the first convolution layer
            bn_size (int) - multiplicative factor for number of bottle neck layers
                (i.e. bn_size * k features in the bottleneck layer)
            drop_rate (float) - dropout rate after each dense layer
            num_classes (int) - number of classification classes
    """
    def __init__(self, dataset, args, growth_rate=32, block_config=(6, 12, 24, 16),
                 num_init_features=64, bn_size=4, drop_rate=0):

        super(DenseNet, self).__init__()
        global batchnorm_momentum, batchnorm_momentum_ckpt
        batchnorm_momentum = dataset.batchnorm_momentum
        print('BatchNorm momentum =', batchnorm_momentum)
        self.checkpointing = args.checkpointing
        self.checkpoint_stem = checkpoint_stem
        if self.checkpointing or self.checkpoint_stem:
            batchnorm_momentum_ckpt = min(np.roots([-1, 2, -batchnorm_momentum]))
            print('BatchNorm momentum ckpt =', batchnorm_momentum_ckpt)
        self.dataset = dataset
        self.num_classes = dataset.num_classes
        self.num_logits = dataset.num_logits
        self.args = args
        self.block_config = block_config
        self.num_blocks = len(block_config)
        self.growth_rate = growth_rate

        # First convolution
        self.features = nn.Sequential()
        self.features.add_module('conv0', nn.Conv2d(3, num_init_features, kernel_size=7,
                                 stride=first_stride, padding=3, bias=False))
        if imagenet_init and use_batchnorm:
            m = batchnorm_momentum_ckpt if self.checkpoint_stem else batchnorm_momentum
            self.features.add_module('norm0', _batchnorm_factory(num_init_features, m))
            self.features.add_module('relu0', nn.ReLU(inplace=True))

        self.features.add_module('pool0', nn.MaxPool2d(kernel_size=2, stride=2))
        self.first_block_idx = len(self.features)
        if self.checkpoint_stem:
            self.first_ckpt_func = self._checkpoint_segment(0, self.first_block_idx)

        self.random_init = []
        self.fine_tune = []
        if args.pretrained:
            self.fine_tune.append(self.features)
        else:
            self.random_init.append(self.features)

        splits = [False, False, True, False]
        up_sizes = [256, 256, 128, 128]

        # for cityscapes/vistas
        spp_square_grid = False
        spp_grids = [8,4,2,1]
        # spp_grids = [4,2,1]

        # for VOC2012
        # spp_square_grid = True
        # spp_grids = [6,3,2,1]

        self.spp_size = 512
        bt_size = 512
        level_size = 128
        dilations = [1] * len(block_config)
        strides = [2] * (len(block_config) - 1)

        # dilate last block
        # dilations = [1, 1, 1, 2]
        # strides = [2, 2, 1]
        # up_sizes = up_sizes[1:]

        # dilate last two blocks
        # dilations = [1, 1, 2, 4]
        # strides = [2, 1, 1]
        # up_sizes = up_sizes[2:]

        num_downs = first_stride + strides.count(2) + sum(splits)
        num_ups = len(up_sizes)
        self.downsampling_factor = 2**num_downs
        self.upsampling_factor = 2**(num_downs-num_ups)
        print(self.downsampling_factor)
        print(self.upsampling_factor)
        args.downsampling_factor = self.downsampling_factor
        args.upsampling_factor = self.upsampling_factor

        self.use_upsampling_path = True

        skip_sizes = []
        num_features = num_init_features

        for i, num_layers in enumerate(block_config):
            block = _DenseBlock(
                num_layers=num_layers, num_input_features=num_features, bn_size=bn_size,
                growth_rate=growth_rate, drop_rate=drop_rate, split=splits[i],
                dilation=dilations[i], checkpointing=self.checkpointing)
            self.features.add_module('denseblock%d' % (i + 1), block)
            num_features = num_features + num_layers * growth_rate
            if block.split and self.use_upsampling_path:
                skip_sizes.append(block.split_size)
            if i != len(block_config) - 1:
                if strides[i] > 1 and self.use_upsampling_path:
                    skip_sizes.append(num_features)
                trans = _Transition(
                    num_input_features=num_features, num_output_features=num_features // 2,
                    stride=strides[i], checkpointing=self.checkpointing)
                self.features.add_module('transition%d' % (i + 1), trans)
                num_features = num_features // 2

        self.use_aux = args.aux_loss
        self.spp = layers.SpatialPyramidPooling(
            _BNReluConv, upsample, num_features, bt_size, level_size,
            self.spp_size, spp_grids, spp_square_grid)

        self.random_init.append(self.spp)
        num_features = self.spp_size

        if self.use_aux:
            self.pyramid_loss_scales = data_utils.get_pyramid_loss_scales(
                    args.downsampling_factor, args.upsampling_factor)
            if use_pyl_in_spp:
                spp_scales = []
                for scale in reversed(spp_grids):
                    assert args.crop_size % scale == 0
                    spp_scales.append(args.crop_size // scale)
                self.pyramid_loss_scales = spp_scales + self.pyramid_loss_scales

        if self.use_upsampling_path:
            up_convs = [3] * len(up_sizes)
            self.upsample_layers = nn.Sequential()
            self.random_init.append(self.upsample_layers)
            print(up_sizes, skip_sizes)
            assert len(up_sizes) == len(skip_sizes)
            for i in range(num_ups):
                upsample_unit = layers.UpsampleResidual(
                    _BNReluConv, upsample, num_features, skip_sizes[-1-i], up_sizes[i],
                    up_convs[i], self.args.aux_loss, self.num_classes)
                num_features = upsample_unit.num_maps_out
                self.upsample_layers.add_module('upsample_'+str(i), upsample_unit)
            print()

        self.logits = _BNReluConv(num_features, self.num_logits, k=1, output_conv=True,
                                  checkpointing=self.checkpointing)
        self.random_init.append(self.logits)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        nn.init.xavier_normal_(self.logits.conv.weight.data)
        if self.use_aux and self.use_upsampling_path:
            for module in self.upsample_layers:
                nn.init.xavier_normal_(module.aux_logits.conv.weight.data)

    def _checkpoint_segment(self, start, end):
        def func(x):
            for i in range(start, end):
                x = self.features[i](x)
            return x
        return func

    def forward(self, x, target_size=None):
        skip_layers = []
        if target_size is None:
            target_size = x.size()[2:4]

        if not self.training or not self.checkpoint_stem:
            for i in range(self.first_block_idx+1):
                x = self.features[i].forward(x)
        else:
            x = checkpoint(self.first_ckpt_func, x)
            x = self.features[self.first_block_idx].forward(x)
        for i in range(self.first_block_idx+1, len(self.features), 2):
            if self.features[i].stride > 1 and self.use_upsampling_path:
                skip_layers.append(x)
            x = self.features[i].forward(x)
            x = self.features[i+1].forward(x)
            if isinstance(x, tuple) and self.use_upsampling_path:
                x, split = x
                skip_layers.append(split)

        x = self.spp(x)

        aux_logits = []
        if self.use_upsampling_path:
            for i, up in enumerate(self.upsample_layers):
                x = up(x, skip_layers[-1-i])
                if self.use_aux:
                    x, aux = x
                    aux_logits.append(aux)

        x = self.logits(x)
        x = upsample(x, target_size)

        return x, aux_logits

    def send_to_gpu(self, batch):
        x = batch['image'].cuda(non_blocking=True)
        batch['image'] = x
        batch['labels_cpu'] = batch['labels']
        batch['labels'] = batch['labels'].cuda(non_blocking=True)
        if self.use_aux and self.training:
            transform.build_pyramid_labels_th(batch, self.pyramid_loss_scales, self.num_classes)

    def forward_loss(self, batch, return_outputs=False):
        x = batch['image']
        logits, aux_logits = self.forward(x)
        if not self.training:
            aux_logits = []
        self.output = logits
        loss = losses.segmentation_loss(logits, aux_logits, batch, self.args.aux_loss_weight,
                                        self.dataset.ignore_id, equal_level_weights=use_pyl_in_spp)
        loss, self.aux_losses = loss
        if return_outputs:
            return loss, (logits, aux_logits)
        return loss

    def load_state_dict(self, state_dict, strict=False):
        """Copies parameters and buffers from :attr:`state_dict` into
        this module and its descendants. The keys of :attr:`state_dict` must
        exactly match the keys returned by this module's :func:`state_dict()`
        function.

        Arguments:
            state_dict (dict): A dict containing parameters and persistent buffers.
        """
        pattern = re.compile(
            r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
        for key in list(state_dict.keys()):
            res = pattern.match(key)
            if res:
                new_key = res.group(1) + res.group(2)
                state_dict[new_key] = state_dict[key]
                del state_dict[key]

        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                # print('Variable "{}" missing in model'.format(name))
                continue
            if isinstance(param, torch.nn.parameter.Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            try:
                own_state[name].copy_(param)
            except:
                print('While copying the parameter named {}, whose dimensions in the model are'
                      ' {} and whose dimensions in the checkpoint are {}, ...'.format(
                      name, own_state[name].size(), param.size()))
                raise
        # missing = set(own_state.keys()) - set(state_dict.keys())
        # if len(missing) > 0:
        #         print('missing keys in state_dict: "{}"'.format(missing))


#!/usr/bin/env python3
# encoding: utf-8

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.keras import Model as M
from tensorflow.keras import Input as I
from tensorflow.keras import Sequential
from skimage.util.shape import view_as_windows
from tensorflow.keras.layers import (Dense,
                                     Flatten,
                                     LayerNormalization)

from rls.utils.tf2_utils import (clip_nn_log_std,
                                 squash_rsample,
                                 gaussian_entropy,
                                 update_target_net_weights)
from rls.algos.base.off_policy import Off_Policy
from rls.utils.sundry_utils import LinearAnnealing
from rls.nn.networks import get_visual_network_from_type
from rls.nn.initializers import initKernelAndBias
from rls.utils.indexs import VisualNetworkType
from rls.utils.build_networks import (ValueNetwork,
                                      DoubleValueNetwork)
from rls.utils.indexs import OutputNetworkType

class VisualEncoder(M):

    def __init__(self, img_dim, fc_dim):
        super().__init__()
        self.net = Sequential([
            get_visual_network_from_type(VisualNetworkType.NATURE)(),
            Dense(fc_dim, **initKernelAndBias),
            LayerNormalization()
        ])
        self(I(shape=img_dim))

    def call(self, vis):
        return self.net(vis)

def random_crop(imgs, output_size):
    """
    Vectorized way to do random crop using sliding windows
    and picking out random ones

    args:
        imgs, batch images with shape (B,C,H,W)
    """
    # batch size
    n = imgs.shape[0]
    img_size = imgs.shape[-1]
    crop_max = img_size - output_size
    imgs = np.transpose(imgs, (0, 2, 3, 1))
    w1 = np.random.randint(0, crop_max, n)
    h1 = np.random.randint(0, crop_max, n)
    # creates all sliding windows combinations of size (output_size)
    windows = view_as_windows(
        imgs, (1, output_size, output_size, 1))[..., 0, :, :, 0]
    # selects a random window for each batch element
    cropped_imgs = windows[np.arange(n), w1, h1]
    return cropped_imgs

def center_crop_image(image, output_size):
    h, w = image.shape[1:3]
    new_h, new_w = output_size, output_size

    top = (h - new_h) // 2
    left = (w - new_w) // 2

    image = image[:, top:top + new_h, left:left + new_w]
    return image

class CURL(Off_Policy):
    """
    CURL: Contrastive Unsupervised Representations for Reinforcement Learning, http://arxiv.org/abs/2004.04136
    """

    def __init__(self,
                 envspec,

                 alpha=0.2,
                 annealing=True,
                 last_alpha=0.01,
                 ployak=0.995,
                 discrete_tau=1.0,
                 log_std_bound=[-20, 2],
                 network_settings={
                     'actor_continuous': {
                         'share': [128, 128],
                         'mu': [64],
                         'log_std': [64]
                     },
                     'actor_discrete': [64, 32],
                     'q': [128, 128],
                     'encoder': 128
                 },
                 auto_adaption=True,
                 actor_lr=5.0e-4,
                 critic_lr=1.0e-3,
                 alpha_lr=5.0e-4,
                 curl_lr=5.0e-4,
                 img_size=64,
                 **kwargs):
        super().__init__(envspec=envspec, **kwargs)
        assert self.visual_sources == 1
        self.ployak = ployak
        self.discrete_tau = discrete_tau
        self.log_std_min, self.log_std_max = log_std_bound[:]
        self.auto_adaption = auto_adaption
        self.annealing = annealing
        self.img_size = img_size
        self.img_dim = [img_size, img_size, self.visual_dim[-1]]
        self.vis_feat_size = network_settings['encoder']

        if self.auto_adaption:
            self.log_alpha = tf.Variable(initial_value=0.0, name='log_alpha', dtype=tf.float32, trainable=True)
        else:
            self.log_alpha = tf.Variable(initial_value=tf.math.log(alpha), name='log_alpha', dtype=tf.float32, trainable=False)
            if self.annealing:
                self.alpha_annealing = LinearAnnealing(alpha, last_alpha, 1.0e6)

        def _create_net(name): return DoubleValueNetwork(
            name=name,
            value_net_type=OutputNetworkType.CRITIC_QVALUE_ONE,
            value_net_kwargs=dict(vector_dim=self.s_dim + self.vis_feat_size,
                                  action_dim=self.a_dim,
                                  network_settings=network_settings['q'])
        )

        self.critic_net = _create_net('critic_net')
        self.critic_target_net = _create_net('critic_target_net')

        if self.is_continuous:
            self.actor_net = ValueNetwork(
                name='actor_net',
                value_net_type=OutputNetworkType.ACTOR_CTS,
                value_net_kwargs=dict(vector_dim=self.s_dim + self.vis_feat_size,
                                      output_shape=self.a_dim,
                                      network_settings=network_settings['actor_continuous'])
            )
        else:
            self.actor_net = ValueNetwork(
                name='actor_net',
                value_net_type=OutputNetworkType.ACTOR_DCT,
                value_net_kwargs=dict(vector_dim=self.s_dim + self.vis_feat_size,
                                      output_shape=self.a_dim,
                                      network_settings=network_settings['actor_discrete'])
            )
            self.gumbel_dist = tfp.distributions.Gumbel(0, 1)

        # entropy = -log(1/|A|) = log |A|
        self.target_entropy = 0.98 * (-self.a_dim if self.is_continuous else np.log(self.a_dim))

        self.encoder = VisualEncoder(self.img_dim, self.vis_feat_size)
        self.encoder_target = VisualEncoder(self.img_dim, self.vis_feat_size)

        self.curl_w = tf.Variable(initial_value=tf.random.normal(shape=(self.vis_feat_size, self.vis_feat_size)), name='curl_w', dtype=tf.float32, trainable=True)

        self.critic_tv = self.critic_net.trainable_variables + self.encoder.trainable_variables

        update_target_net_weights(
            self.critic_target_net.weights + self.encoder_target.trainable_variables,
            self.critic_net.weights + self.encoder.trainable_variables
        )
        self.actor_lr, self.critic_lr, self.alpha_lr, self.curl_lr = map(self.init_lr, [actor_lr, critic_lr, alpha_lr, curl_lr])
        self.optimizer_actor, self.optimizer_critic, self.optimizer_alpha, self.optimizer_curl = map(self.init_optimizer, [self.actor_lr, self.critic_lr, self.alpha_lr, self.curl_lr])

        self._worker_params_dict.update(self.actor_net._policy_models)
        self._worker_params_dict.update(encoder=self.encoder)

        self._all_params_dict.update(self.actor_net._all_models)
        self._all_params_dict.update(self.critic_net._all_models)
        self._all_params_dict.update(curl_w=self.curl_w,
                                     encoder=self.encoder,
                                     optimizer_actor=self.optimizer_actor,
                                     optimizer_critic=self.optimizer_critic,
                                     optimizer_alpha=self.optimizer_alpha,
                                     optimizer_curl=self.optimizer_curl)
        self._model_post_process()

    def choose_action(self, s, visual_s, evaluation=False):
        visual_s = center_crop_image(visual_s[:, 0], self.img_size)
        mu, pi = self._get_action(s, visual_s)
        a = mu.numpy() if evaluation else pi.numpy()
        return a

    @tf.function
    def _get_action(self, s, visual_s):
        with tf.device(self.device):
            feat = tf.concat([self.encoder(visual_s), s], axis=-1)
            if self.is_continuous:
                mu, log_std = self.actor_net.value_net(feat)
                log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
                pi, _ = squash_rsample(mu, log_std)
                mu = tf.tanh(mu)  # squash mu
            else:
                logits = self.actor_net.value_net(feat)
                mu = tf.argmax(logits, axis=1)
                cate_dist = tfp.distributions.Categorical(logits=tf.nn.log_softmax(logits))
                pi = cate_dist.sample()
            return mu, pi

    def _process_before_train(self, data):
        data['visual_s'] = np.transpose(data['visual_s'][:, 0].numpy(), (0, 3, 1, 2))
        data['visual_s_'] = np.transpose(data['visual_s_'][:, 0].numpy(), (0, 3, 1, 2))
        data['pos'] = self.data_convert(
            np.transpose(random_crop(data['visual_s'], self.img_size), (0, 2, 3, 1))
        )
        data['visual_s'] = self.data_convert(
            np.transpose(random_crop(data['visual_s'], self.img_size), (0, 2, 3, 1))
        )
        data['visual_s_'] = self.data_convert(
            np.transpose(random_crop(data['visual_s_'], self.img_size), (0, 2, 3, 1))
        )
        return (data,)

    def _target_params_update(self):
        update_target_net_weights(self.critic_target_net.weights + self.encoder_target.trainable_variables,
                                  self.critic_net.weights + self.encoder.trainable_variables,
                                  self.ployak)

    def learn(self, **kwargs):
        self.train_step = kwargs.get('train_step')

        for i in range(self.train_times_per_step):
            self._learn(function_dict={
                'summary_dict': dict([
                    ['LEARNING_RATE/actor_lr', self.actor_lr(self.train_step)],
                    ['LEARNING_RATE/critic_lr', self.critic_lr(self.train_step)],
                    ['LEARNING_RATE/alpha_lr', self.alpha_lr(self.train_step)]
                ]),
                'train_data_list': ['s', 'visual_s', 'a', 'r', 's_', 'visual_s_', 'done', 'pos'],
            })

    @property
    def alpha(self):
        return tf.exp(self.log_alpha)

    def _train(self, memories, isw, cell_state):
        td_error, summaries = self.train(memories, isw, cell_state)
        if self.annealing and not self.auto_adaption:
            self.log_alpha.assign(tf.math.log(tf.cast(self.alpha_annealing(self.global_step.numpy()), tf.float32)))
    return td_error, summaries

    @tf.function(experimental_relax_shapes=True)
    def train(self, memories, isw, cell_state):
        s, visual_s, a, r, s_, visual_s_, done, pos = memories
        with tf.device(self.device):
            with tf.GradientTape(persistent=True) as tape:
                vis_feat = self.encoder(visual_s)
                vis_feat_ = self.encoder(visual_s_)
                target_vis_feat_ = self.encoder_target(visual_s_)
                feat = tf.concat([vis_feat, s], axis=-1)
                feat_ = tf.concat([vis_feat_, s_], axis=-1)
                target_feat_ = tf.concat([target_vis_feat_, s_], axis=-1)
                if self.is_continuous:
                    target_mu, target_log_std = self.actor_net.value_net(feat_)
                    target_log_std = clip_nn_log_std(target_log_std, self.log_std_min, self.log_std_max)
                    target_pi, target_log_pi = squash_rsample(target_mu, target_log_std)
                else:
                    target_logits = self.actor_net.value_net(feat_)
                    target_cate_dist = tfp.distributions.Categorical(logits=tf.nn.log_softmax(target_logits))
                    target_pi = target_cate_dist.sample()
                    target_log_pi = target_cate_dist.log_prob(target_pi)
                    target_pi = tf.one_hot(target_pi, self.a_dim, dtype=tf.float32)
                q1, q2 = self.critic_net.value_net(feat, a)
                q1_target, q2_target = self.critic_target_net.value_net(feat_, target_pi)
                q_target = tf.minimum(q1_target, q2_target)
                dc_r = tf.stop_gradient(r + self.gamma * (1 - done) * (q_target - self.alpha * target_log_pi))
                td_error1 = q1 - dc_r
                td_error2 = q2 - dc_r
                q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
                q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
                critic_loss = 0.5 * q1_loss + 0.5 * q2_loss

                z_a = vis_feat  # [B, N]
                z_out = self.encoder_target(pos)
                logits = tf.matmul(z_a, tf.matmul(self.curl_w, tf.transpose(z_out, [1, 0])))
                logits -= tf.reduce_max(logits, axis=-1, keepdims=True)
                curl_loss = tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(tf.range(self.batch_size), logits))
            critic_grads = tape.gradient(critic_loss, self.critic_tv)
            self.optimizer_critic.apply_gradients(
                zip(critic_grads, self.critic_tv)
            )
            curl_grads = tape.gradient(curl_loss, [self.curl_w] + self.encoder.trainable_variables)
            self.optimizer_curl.apply_gradients(
                zip(curl_grads, [self.curl_w] + self.encoder.trainable_variables)
            )

            with tf.GradientTape() as tape:
                if self.is_continuous:
                    mu, log_std = self.actor_net.value_net(feat)
                    log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
                    pi, log_pi = squash_rsample(mu, log_std)
                    entropy = gaussian_entropy(log_std)
                else:
                    logits = self.actor_net.value_net(feat)
                    logp_all = tf.nn.log_softmax(logits)
                    gumbel_noise = tf.cast(self.gumbel_dist.sample(a.shape), dtype=tf.float32)
                    _pi = tf.nn.softmax((logp_all + gumbel_noise) / self.discrete_tau)
                    _pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1), self.a_dim)
                    _pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
                    pi = _pi_diff + _pi
                    log_pi = tf.reduce_sum(tf.multiply(logp_all, pi), axis=1, keepdims=True)
                    entropy = -tf.reduce_mean(tf.reduce_sum(tf.exp(logp_all) * logp_all, axis=1, keepdims=True))
                q_s_pi = self.critic_net.get_min(feat, pi)
                actor_loss = -tf.reduce_mean(q_s_pi - self.alpha * log_pi)
            actor_grads = tape.gradient(actor_loss, self.actor_net.trainable_variables)
            self.optimizer_actor.apply_gradients(
                zip(actor_grads, self.actor_net.trainable_variables)
            )

            if self.auto_adaption:
                with tf.GradientTape() as tape:
                    if self.is_continuous:
                        mu, log_std = self.actor_net.value_net(feat)
                        log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
                        norm_dist = tfp.distributions.Normal(loc=mu, scale=tf.exp(log_std))
                        log_pi = tf.reduce_sum(norm_dist.log_prob(norm_dist.sample()), axis=-1)
                    else:
                        logits = self.actor_net.value_net(feat)
                        cate_dist = tfp.distributions.Categorical(logits=tf.nn.log_softmax(logits))
                        log_pi = cate_dist.log_prob(cate_dist.sample())
                    alpha_loss = -tf.reduce_mean(self.alpha * tf.stop_gradient(log_pi + self.target_entropy))
                alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
                self.optimizer_alpha.apply_gradients(
                    [(alpha_grad, self.log_alpha)]
                )
            self.global_step.assign_add(1)
            summaries = dict([
                ['LOSS/actor_loss', actor_loss],
                ['LOSS/q1_loss', q1_loss],
                ['LOSS/q2_loss', q2_loss],
                ['LOSS/critic_loss', critic_loss],
                ['LOSS/curl_loss', curl_loss],
                ['Statistics/log_alpha', self.log_alpha],
                ['Statistics/alpha', self.alpha],
                ['Statistics/entropy', entropy],
                ['Statistics/q_min', tf.reduce_min(tf.minimum(q1, q2))],
                ['Statistics/q_mean', tf.reduce_mean(tf.minimum(q1, q2))],
                ['Statistics/q_max', tf.reduce_max(tf.maximum(q1, q2))]
            ])
            if self.auto_adaption:
                summaries.update({
                    'LOSS/alpha_loss': alpha_loss
                })
            return (td_error1 + td_error2) / 2., summaries


# from __future__ import absolute_import
# from __future__ import division
# from __future__ import print_function
import random
#random.seed(1017)

import time

import numpy as np
import tensorflow as tf
from data_io import tpn_iterator, tpn_raw_data
import cPickle
import os
import os.path as osp
import glog as log
from model import BiTPNModel

flags = tf.flags
logging = tf.logging

flags.DEFINE_string(
    "model", "default",
    "A type of model. [default]")
flags.DEFINE_string("data_path", None, "data_path")
flags.DEFINE_string("log_path", '.', "log_path")
flags.DEFINE_string("save_path", '.', "save_path")
flags.DEFINE_integer("num_layers", 1, "number of LSTM layers. [1]")
flags.DEFINE_string("type", "residual,basic", "type of LSTM cells {residual, basic}. [residual]")
flags.DEFINE_string("combine", "max,ave,concat", "Bidirectional feature combination scheme.")

FLAGS = flags.FLAGS

if not osp.isdir(FLAGS.save_path):
  os.makedirs(FLAGS.save_path)

class DefaultConfig(object):
  """Default config."""
  init_scale = 0.01
  learning_rate = 0.001
  momentum = 0.9
  max_grad_norm = 1.5
  num_layers = FLAGS.num_layers
  num_steps = 20
  input_size = 1024
  hidden_size = 1024
  max_epoch = 5
  iter_epoch = 2000
  keep_prob = 1.0
  lr_decay = 0.5
  batch_size = 128
  num_classes = 31
  cls_weight = 1.0
  bbox_weight = 0.0
  ending_weight = 1.0
  vid_per_batch = 4
  cls_init = 'cls_score_vid_params.pkl'
  bbox_init = 'bbox_pred_vid_params.pkl'
  combine = FLAGS.combine
  type = FLAGS.type

class TestConfig(object):
  """Tiny config, for testing."""
  init_scale = 0.1
  learning_rate = 1.0
  max_grad_norm = 1
  num_layers = 1
  num_steps = 2
  hidden_size = 2
  max_epoch = 1
  max_max_epoch = 1
  keep_prob = 1.0
  lr_decay = 0.5
  batch_size = 20
  vocab_size = 10000
  combine = FLAGS.combine

def run_epoch(session, m, data, eval_op, init_state, epoch_idx, verbose=False):
  """Runs the model on the given data."""
  start_time = time.time()
  tot_costs = 0.0
  costs = 0.0
  cls_costs = 0.0
  bbox_costs = 0.0
  end_costs = 0.0
  display_iter = 20.
  state, b_state = init_state
  for step in xrange(m.iter_epoch):
    x, cls_t, end_t, bbox_t, bbox_weights = tpn_iterator(data, m.batch_size, m.num_steps, m.num_classes, m.vid_per_batch)
    data_time = time.time()
    cost, cls_cost, bbox_cost, end_cost, state, b_state, _, global_norm = session.run(
        [m.cost, m.cls_cost, m.bbox_cost, m.end_cost,
        m.final_state, m.final_backward_state,
        eval_op,  m.global_norm],
         {m.input_data: x,
          m.cls_targets: cls_t,
          m.bbox_targets: bbox_t,
          m.bbox_weights: bbox_weights,
          m.end_targets: end_t,
          m.initial_state: state,
          m.initial_backward_state: b_state})
    costs += cost
    tot_costs += cost
    cls_costs += cls_cost
    bbox_costs += bbox_cost
    end_costs += end_cost

    if verbose and (step + 1) % display_iter == 0:
      costs /= display_iter
      cls_costs /= display_iter
      bbox_costs /= display_iter
      end_costs /= display_iter
      log.info("Iter {:06d} {:.03f} s/iter data time: {:.03f} s: cost {:.03f} = cls_cost {:.03f} * {:.02f} + end_cost {:.03f} * {:.02f} + bbox_cost {:.03f} * {:.02f}. Global norm: {:.03f}".format(
        step+1+epoch_idx * m.iter_epoch,
        (time.time() - start_time) / display_iter,
        (data_time - start_time) / display_iter,
        costs, cls_costs, m.cls_weight,
        end_costs, m.ending_weight,
        bbox_costs, m.bbox_weight, global_norm))
      costs = 0.0
      cls_costs = 0.0
      bbox_costs = 0.0
      end_costs = 0.0
      start_time = time.time()

  return tot_costs / m.iter_epoch, state, b_state

def get_config():
  if FLAGS.model == "default":
    return DefaultConfig()
  else:
    raise ValueError("Invalid model: %s", FLAGS.model)

def main(_):
  if not FLAGS.data_path:
    raise ValueError("Must set --data_path to TPN data directory")

  log.info("Processing data...")
  raw_data = tpn_raw_data(FLAGS.data_path)
  train_data, valid_data = raw_data

  config = get_config()
  eval_config = get_config()
  eval_config.batch_size = 1
  eval_config.num_steps = 1

  #tf.set_random_seed(1017)
  with tf.Graph().as_default(), tf.Session() as session:
    initializer = tf.random_uniform_initializer(-config.init_scale,
                                                config.init_scale, seed=1017)
    with tf.variable_scope("model", reuse=None, initializer=initializer):
      m = BiTPNModel(is_training=True, config=config)
    with tf.variable_scope("model", reuse=True, initializer=initializer):
      mvalid = BiTPNModel(is_training=False, config=config)

    tf.initialize_all_variables().run()

    saver = tf.train.Saver()
    for i in range(config.max_epoch):
      lr_decay = config.lr_decay ** i
      m.assign_lr(session, config.learning_rate * lr_decay)

      if i == 0:
        state = m.initial_state.eval()
        b_state = m.initial_backward_state.eval()
      log.info("Epoch: {} Learning rate {:.03e}.".format(i+1,session.run(m.lr)))
      train_cost, state, b_state = run_epoch(session, m, train_data, m.train_op,
                             [state, b_state], i, verbose=True)
      log.info("Epoch: %d Train Cost: %.3f" % (i + 1, train_cost))
      save_path = osp.join(FLAGS.save_path, '{}LSTM_{}'.format(FLAGS.type, FLAGS.num_layers))
      log.info("Save to {}_{}".format(save_path, (i+1)*m.iter_epoch))
      saver.save(session, save_path, global_step=(i+1) * m.iter_epoch)

if __name__ == "__main__":
  tf.app.run()