from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
from typing_extensions import Literal

from types import MappingProxyType
from pathlib import Path

import scanpy as sc
import scvelo as scv
from anndata import Raw, AnnData
from cellrank import logging as logg
from cellrank._key import Key
from cellrank.ul._docs import d, inject_docs
from cellrank.tl._utils import RandomKeys, TestMethod, save_fig, _correlation_test
from cellrank.tl._colors import _create_categorical_colors
from cellrank.tl._lineage import Lineage
from cellrank.tl.estimators._utils import SafeGetter
from cellrank.tl.estimators.mixins._utils import BaseProtocol, logger, shadow
from cellrank.tl.estimators.mixins._absorption_probabilities import AbsProbsMixin

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from matplotlib import rc_context, patheffects
from matplotlib.axes import Axes
from matplotlib.patches import ArrowStyle

class LinDriversProtocol(BaseProtocol):  # noqa: D101
    def _write_lineage_drivers(
        self,
        drivers: Optional[pd.DataFrame],
        use_raw: bool,
        params: Mapping[str, Any] = MappingProxyType({}),
    ) -> None:
        ...

    @property
    def eigendecomposition(self) -> Dict[str, Any]:  # noqa: D102
        ...

    @property
    def absorption_probabilities(self) -> Optional[Lineage]:  # noqa: D102
        ...

    @property
    def lineage_drivers(self) -> Optional[pd.DataFrame]:  # noqa: D102
        ...

class LinDriversMixin(AbsProbsMixin):
    """Mixin that computes potential driver genes for lineages."""

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        self._lineage_drivers: Optional[pd.DataFrame] = None

    @d.get_full_description(base="lineage_drivers")
    @d.get_sections(base="lineage_drivers", sections=["Parameters", "Returns"])
    @d.dedent
    @inject_docs(tm=TestMethod)
    def compute_lineage_drivers(
        self: LinDriversProtocol,
        lineages: Optional[Union[str, Sequence]] = None,
        method: Literal["fischer", "perm_test"] = TestMethod.FISCHER,
        cluster_key: Optional[str] = None,
        clusters: Optional[Union[str, Sequence]] = None,
        layer: Optional[str] = None,
        use_raw: bool = False,
        confidence_level: float = 0.95,
        n_perms: int = 1000,
        seed: Optional[int] = None,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """
        Compute driver genes per lineage.

        Correlates gene expression with lineage probabilities, for a given lineage and set of clusters.
        Often, it makes sense to restrict this to a set of clusters which are relevant for the specified lineages.

        Parameters
        ----------
        lineages
            Lineage names from :attr:`absorption_probabilities`. If `None`, use all lineages.
        method
            Mode to use when calculating p-values and confidence intervals. Valid options are:

                - `{tm.FISCHER!r}` - use Fischer transformation :cite:`fischer:21`.
                - `{tm.PERM_TEST!r}` - use permutation test.
        cluster_key
            Key from :attr:`anndata.AnnData.obs` to obtain cluster annotations. These are considered for ``clusters``.
        clusters
            Restrict the correlations to these clusters.
        layer
            Key from :attr:`anndata.AnnData.layers` from which to get the expression.
            If `None` or `'X'`, use :attr:`anndata.AnnData.X`.
        use_raw
            Whether or not to use :attr:`anndata.AnnData.raw` to correlate gene expression.
        confidence_level
            Confidence level for the confidence interval calculation. Must be in interval `[0, 1]`.
        n_perms
            Number of permutations to use when ``method = {tm.PERM_TEST!r}``.
        seed
            Random seed when ``method = {tm.PERM_TEST!r}``.
        %(parallel)s

        Returns
        -------
        %(correlation_test.returns)s

        Also updates the following field:

            - :attr:`lineage_drivers` - the same :class:`pandas.DataFrame` as described above.
        """

        # check that lineage probs have been computed
        method = TestMethod(method)
        abs_probs = self.absorption_probabilities
        if abs_probs is None:
            raise RuntimeError(
                "Compute `.absorption_probabilities` first as `.compute_absorption_probabilities()`."
            )

        if abs_probs.shape[1] == 1:
            logg.warning(
                "There is only 1 lineage present. Using stationary distribution instead"
            )
            stat_dist = self.eigendecomposition.get("stationary_dist", None)
            if stat_dist is None:
                raise RuntimeError(
                    "No stationary distribution found in `.eigendecomposition['stationary_dist']`."
                )
            abs_probs = Lineage(
                stat_dist,
                names=abs_probs.names,
                colors=abs_probs.colors,
            )

        # check all lin_keys exist in self.lin_names
        if isinstance(lineages, str):
            lineages = [lineages]
        if lineages is not None:
            _ = abs_probs[lineages]
        else:
            lineages = abs_probs.names
        lineages = list(lineages)

        if not len(lineages):
            raise ValueError("No lineages have been selected.")

        # use `cluster_key` and clusters to subset the data
        if clusters is not None:
            if cluster_key not in self.adata.obs.keys():
                raise KeyError(
                    f"Unable to find clusters in `adata.obs[{cluster_key!r}]`."
                )
            if isinstance(clusters, str):
                clusters = [clusters]

            all_clusters = np.array(self.adata.obs[cluster_key].cat.categories)
            cluster_mask = np.array([name not in all_clusters for name in clusters])

            if any(cluster_mask):
                raise KeyError(
                    f"Clusters `{list(np.array(clusters)[cluster_mask])}` not found in "
                    f"`adata.obs[{cluster_key!r}]`."
                )
            subset_mask = np.in1d(self.adata.obs[cluster_key], clusters)
            adata_comp = self.adata[subset_mask]
            lin_probs = abs_probs[subset_mask, :]
        else:
            adata_comp = self.adata
            lin_probs = abs_probs

        # check that the layer exists, and that use raw is only used with layer X
        if layer in (None, "X"):
            if use_raw and self.adata.raw is None:
                logg.warning("No raw attribute set. Using `.X` instead")
                use_raw = False
            data = adata_comp.raw.X if use_raw else adata_comp.X
            var_names = adata_comp.raw.var_names if use_raw else adata_comp.var_names
        else:
            if layer not in self.adata.layers:
                raise KeyError(f"Layer `{layer!r}` not found in `adata.layers`.")
            if use_raw:
                logg.warning("If `use_raw=True`, layer must be `None`.")
                use_raw = False
            data = adata_comp.layers[layer]
            var_names = adata_comp.var_names

        start = logg.debug(
            f"Computing correlations for lineages `{sorted(lineages)}` restricted to clusters `{clusters}` in "
            f"layer `{'X' if layer is None else layer}` with `use_raw={use_raw}`"
        )

        lin_probs = lin_probs[lineages]
        drivers = _correlation_test(
            data,
            lin_probs,
            gene_names=var_names,
            method=method,
            n_perms=n_perms,
            seed=seed,
            confidence_level=confidence_level,
            **kwargs,
        )
        params = self._create_params()
        self._write_lineage_drivers(
            drivers.loc[var_names], use_raw=use_raw, params=params, time=start
        )

        return drivers

    @d.get_sections(base="plot_lineage_drivers", sections=["Parameters"])
    @d.dedent
    def plot_lineage_drivers(
        self: LinDriversProtocol,
        lineage: str,
        n_genes: int = 8,
        ncols: Optional[int] = None,
        use_raw: bool = False,
        title_fmt: str = "{gene} qval={qval:.4e}",
        figsize: Optional[Tuple[float, float]] = None,
        dpi: Optional[int] = None,
        save: Optional[Union[str, Path]] = None,
        **kwargs: Any,
    ) -> None:
        """
        Plot lineage drivers discovered by :meth:`compute_lineage_drivers`.

        Parameters
        ----------
        lineage
            Lineage for which to plot the driver genes.
        n_genes
            Top most correlated genes to plot.
        ncols
            Number of columns.
        use_raw
            Whether to access in :attr:`anndata.AnnData.raw` or not.
        title_fmt
            Title format. Can include `{gene}`, `{pval}`, `{qval}` or `{corr}`, which will be substituted
            with the actual values.
        %(plotting)s
        kwargs
            Keyword arguments for :func:`scvelo.pl.scatter`.

        Returns
        -------
        %(just_plots)s
        """

        def prepare_format(
            gene: str,
            *,
            pval: Optional[float],
            qval: Optional[float],
            corr: Optional[float],
        ) -> Dict[str, Any]:
            kwargs: Dict[str, Optional[Union[str, float]]] = {}
            if "{gene" in title_fmt:
                kwargs["gene"] = gene
            if "{pval" in title_fmt:
                kwargs["pval"] = float(pval) if pval is not None else np.nan
            if "{qval" in title_fmt:
                kwargs["qval"] = float(qval) if qval is not None else np.nan
            if "{corr" in title_fmt:
                kwargs["corr"] = float(corr) if corr is not None else np.nan

            return kwargs

        lin_drivers = self.lineage_drivers
        if lin_drivers is None:
            raise RuntimeError(
                "Compute `.lineage_drivers` first as `.compute_lineage_drivers()`."
            )

        key = f"{lineage}_corr"
        if key not in lin_drivers:
            raise KeyError(
                f"Lineage `{key!r}` not found in `{list(lin_drivers.columns)}`."
            )

        if n_genes <= 0:
            raise ValueError(f"Expected `n_genes` to be positive, found `{n_genes}`.")

        _ = kwargs.pop("save", None)
        genes = lin_drivers.sort_values(by=key, ascending=False).head(n_genes)

        ncols = 4 if ncols is None else ncols
        nrows = int(np.ceil(len(genes) / ncols))

        fig, axes = plt.subplots(
            ncols=ncols,
            nrows=nrows,
            dpi=dpi,
            figsize=(ncols * 6, nrows * 4) if figsize is None else figsize,
        )
        axes = np.ravel([axes])

        _i = 0
        for _i, (gene, ax) in enumerate(zip(genes.index, axes)):
            data = genes.loc[gene]
            scv.pl.scatter(
                self.adata,
                color=gene,
                ncols=ncols,
                use_raw=use_raw,
                ax=ax,
                show=False,
                title=title_fmt.format(
                    **prepare_format(
                        gene,
                        pval=data.get(f"{lineage}_pval", None),
                        qval=data.get(f"{lineage}_qval", None),
                        corr=data.get(f"{lineage}_corr", None),
                    )
                ),
                **kwargs,
            )

        for j in range(_i + 1, len(axes)):
            axes[j].remove()

        if save is not None:
            save_fig(fig, save)

    @d.dedent
    def plot_lineage_drivers_correlation(
        self: LinDriversProtocol,
        lineage_x: str,
        lineage_y: str,
        color: Optional[str] = None,
        gene_sets: Optional[Dict[str, Sequence[str]]] = None,
        gene_sets_colors: Optional[Sequence[str]] = None,
        use_raw: bool = False,
        cmap: str = "RdYlBu_r",
        fontsize: int = 12,
        adjust_text: bool = False,
        legend_loc: Optional[str] = "best",
        figsize: Optional[Tuple[float, float]] = (4, 4),
        dpi: Optional[int] = None,
        save: Optional[Union[str, Path]] = None,
        show: bool = True,
        **kwargs: Any,
    ) -> Optional[Axes]:
        """
        Show scatter plot of gene-correlations between two lineages.

        Optionally, you can pass a :class:`dict` of gene names that will be annotated in the plot.

        Parameters
        ----------
        lineage_x
            Name of the lineage on the x-axis.
        lineage_y
            Name of the lineage on the y-axis.
        color
            Key in :attr:`anndata.AnnData.var` or :attr:`anndata.AnnData.varm`, preferring for the former.
        gene_sets
            Gene sets annotations of the form `{'gene_set_name': ['gene_1', 'gene_2'], ...}`.
        gene_sets_colors
            List of colors where each entry corresponds to a gene set from ``genes_sets``.
            If `None` and keys in ``gene_sets`` correspond to lineage names, use the lineage colors.
            Otherwise, use default colors.
        use_raw
            Whether to access :attr:`anndata.AnnData.raw` or not.
        cmap
            Colormap to use.
        fontsize
            Size of the text when plotting ``gene_sets``.
        adjust_text
            Whether to automatically adjust text in order to reduce overlap.
        legend_loc
            Position of the legend. If `None`, don't show the legend. Only used when ``gene_sets != None``.
        %(plotting)s
        show
            If `False`, return :class:`matplotlib.pyplot.Axes`.
        kwargs
            Keyword arguments for :func:`scanpy.pl.scatter`.

        Returns
        -------
        The axes object, if ``show = False``.
        %(just_plots)s

        Notes
        -----
        This plot is based on the following
        `notebook <https://github.com/theislab/gastrulation_analysis/blob/main/6_cellrank.ipynb>`_ by Maren Büttner.
        """
        from cellrank.pl._utils import _position_legend

        if use_raw and self.adata.raw is None:
            logg.warning("No raw attribute set. Setting `use_raw=False`")
            use_raw = False
        adata = self.adata.raw if use_raw else self.adata
        dkey = Key.varm.lineage_drivers(self.backward)
        # this makes a useless copy, wish there was another way
        if isinstance(adata, Raw):
            adata = adata.to_adata()

        if color is not None:
            if not isinstance(color, str):
                raise TypeError(
                    "Expected `color` to be either of type `str` or `None`, "
                    f"found `{type(color).__name__}`."
                )
            if color not in adata.var and color not in adata.varm[dkey]:
                raise KeyError(
                    f"Unable to find key `{color!r}` in `adata{'.raw' if use_raw else ''}.var` "
                    f"or `adata.{'.raw' if use_raw else ''}varm[{dkey!r}]`."
                )

        # silent assumption: `.compute_lineage_drivers()` always writes to AnnData
        key1, key2 = f"{lineage_x}_corr", f"{lineage_y}_corr"
        if key1 not in adata.varm[dkey] or key2 not in adata.varm[dkey]:
            haystack = "adata.raw.varm" if use_raw else "adata.varm"
            raise RuntimeError(
                f"Unable to find correlations in `{haystack}[{key1!r}]` or `{haystack}[{key2!r}]`."
                f"Compute `.lineage_drivers` first as "
                f"`.compute_lineage_drivers([{lineage_x!r}, {lineage_y!r}], use_raw={use_raw})`."
            )
        # produce the actual scatter plot
        ctx = {"figure.figsize": figsize, "figure.dpi": dpi}
        for key in list(ctx.keys()):
            if ctx[key] is None:
                del ctx[key]
        with rc_context(ctx), RandomKeys(adata, n=3, where="var") as keys:
            adata.var[keys[0]] = adata.varm[dkey][key1]
            adata.var[keys[1]] = adata.varm[dkey][key2]
            if color is not None:
                if color in adata.var:
                    ckey = color
                else:
                    adata.var[keys[2]] = adata.varm[dkey][color]
                    ckey = keys[2]

            ax = sc.pl.scatter(
                adata,
                x=keys[0],
                y=keys[1],
                color_map=cmap,
                use_raw=False,
                show=False,
                color=None if color is None else ckey,
                **kwargs,
            )

            # add some lines to highlight the origin
            xmin, xmax = np.nanmin(adata.var[keys[0]]), np.nanmax(adata.var[keys[0]])
            ymin, ymax = np.nanmin(adata.var[keys[1]]), np.nanmax(adata.var[keys[1]])
            ax.hlines(0, xmin=xmin, xmax=xmax, color="grey", alpha=0.5, zorder=-1)
            ax.vlines(0, ymin=ymin, ymax=ymax, color="grey", alpha=0.5, zorder=-1)

        fig = ax.figure
        ax.set_xlabel(key1)
        ax.set_ylabel(key2)
        ax.set_title(color)

        # annotate the passed set of genes
        if gene_sets is not None:
            if gene_sets_colors is None:
                try:
                    # fmt: off
                    sets = list(gene_sets.keys())
                    gene_sets_colors = self.absorption_probabilities[sets].colors
                    # fmt: on
                except KeyError:
                    logg.warning(
                        "Unable to determine gene sets colors from lineages. Using default colors"
                    )
                    gene_sets_colors = _create_categorical_colors(len(gene_sets))
            if len(gene_sets_colors) != len(gene_sets):
                raise ValueError(
                    f"Expected `gene_sets_colors` to be of length `{len(gene_sets)}`, "
                    f"found `{len(gene_sets_colors)}`."
                )

            path_effect = [
                patheffects.Stroke(linewidth=2, foreground="w", alpha=0.8),
                patheffects.Normal(),
            ]
            annots = []
            for (key, values), color in zip(gene_sets.items(), gene_sets_colors):
                arrowprops = (
                    {
                        "arrowstyle": ArrowStyle.CurveFilledB(
                            head_width=0.1, head_length=0.2
                        ),
                        "fc": color,
                        "ec": color,
                    }
                    if adjust_text
                    else None
                )
                values = set(values) & set(adata.var_names)
                for value in values:
                    x = adata.varm[dkey].loc[value, key1]
                    y = adata.varm[dkey].loc[value, key2]

                    annot = ax.annotate(
                        value,
                        xy=(x, y),
                        va="top",
                        ha="left",
                        path_effects=path_effect,
                        arrowprops=arrowprops,
                        size=fontsize,
                        c=color,
                    )
                    annots.append(annot)
                if values:
                    ax.scatter([], [], color=color, label=key)

            if adjust_text:
                try:
                    import adjustText

                    start = logg.info("Adjusting text position")
                    adjustText.adjust_text(
                        annots,
                        x=adata.varm[dkey][key1].values,
                        y=adata.varm[dkey][key2].values,
                        ax=ax,
                    )
                    logg.info("    Finish", time=start)
                except ImportError:
                    logg.error(
                        "Please install `adjustText` first as `pip install adjustText`"
                    )
            if len(annots) and legend_loc not in (None, "none"):
                _position_legend(ax, legend_loc=legend_loc)

        if save is not None:
            save_fig(fig, path=save)

        if not show:
            return ax

    @logger
    @shadow
    def _write_absorption_probabilities(
        self: LinDriversProtocol,
        abs_probs: Optional[Lineage],
        abs_times: Optional[pd.DataFrame],
        params: Mapping[str, Any] = MappingProxyType({}),
    ) -> str:
        self._write_lineage_drivers(None, use_raw=False, log=False)
        try:
            self._write_lineage_drivers(None, use_raw=True, log=False)
        except AttributeError:
            pass
        return super()._write_absorption_probabilities(
            abs_probs, abs_times, params=params, log=False
        )

    @logger
    @shadow
    def _write_lineage_drivers(
        self: LinDriversProtocol,
        drivers: Optional[pd.DataFrame],
        use_raw: bool,
        params: Mapping[str, Any] = MappingProxyType({}),
    ) -> str:
        self._lineage_drivers = drivers

        # fmt: off
        key = Key.varm.lineage_drivers(self.backward)
        self._set("_lineage_drivers", self.adata.raw.varm if use_raw else self.adata.varm, key=key, value=drivers)
        self.params[key] = dict(params)
        # fmt: on

        return (
            f"Adding `adata.{'raw.' if use_raw else ''}varm[{key!r}]`\n"
            f"       `.lineage_drivers`\n"
            "    Finish"
        )

    def _read_absorption_probabilities(
        self: LinDriversProtocol, adata: AnnData
    ) -> bool:
        ok = super()._read_absorption_probabilities(adata)
        if not ok:
            return False

        # fmt: off
        key = Key.varm.lineage_drivers(self.backward)
        with SafeGetter(self, allowed=(KeyError, AttributeError)) as sg:
            self.params[key] = self._read_params(key)
            if self.params[key].get("use_raw", False):
                self._get("_lineage_drivers", self.adata.raw.varm, key=key, where="varm", dtype=pd.DataFrame)
            else:
                self._get("_lineage_drivers", self.adata.varm, key=key, where="varm", dtype=pd.DataFrame)
        # fmt: on

        return sg.ok

    @property
    @d.dedent
    def lineage_drivers(self) -> Optional[pd.DataFrame]:
        """
        Potential lineage drivers.

        Computes Pearson correlation of each gene with fate probabilities for every terminal state. High Pearson
        correlation indicates potential lineage drivers. Also computes p-values and confidence intervals.

        Returns
        -------
        %(correlation_test.returns)s
        """
        return self._lineage_drivers


import sys,os
sys.path.append("..")
from model_io import model_io
import numpy as np
import tensorflow as tf
from bunch import Bunch
from example import feature_writer, write_to_tfrecords, classifier_processor
from data_generator import tokenization
from data_generator import hvd_distributed_tf_data_utils as tf_data_utils

from example import hvd_bert_order_classifier as bert_order_classifier
import horovod.tensorflow as hvd
from example import esim_bert

flags = tf.flags

FLAGS = flags.FLAGS

## Required parameters
flags.DEFINE_string(
    "eval_data_file", None,
    "The config json file corresponding to the pre-trained BERT model. "
    "This specifies the model architecture.")

flags.DEFINE_string(
    "output_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "config_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "init_checkpoint", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "result_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "vocab_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "label_id", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_integer(
    "max_length", 128,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "train_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "dev_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "model_output", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "gpu_id", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_integer(
    "epoch", 5,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_integer(
    "num_classes", 3,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_integer(
    "train_size", 256434,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_integer(
    "batch_size", 32,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "model_type", None,
    "Input TF example files (can be a glob or comma separated).")

def main(_):

    hvd.init()

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.visible_device_list = str(hvd.local_rank())

    graph = tf.Graph()
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
    with graph.as_default():
        import json

        # config = json.load(open("/data/xuht/bert/chinese_L-12_H-768_A-12/bert_config.json", "r"))

        config = json.load(open(FLAGS.config_file, "r"))

        init_checkpoint = FLAGS.init_checkpoint
        print("===init checkoutpoint==={}".format(init_checkpoint))

        # init_checkpoint = "/data/xuht/bert/chinese_L-12_H-768_A-12/bert_model.ckpt"
        # init_checkpoint = "/data/xuht/concat/model_1/oqmrc.ckpt"
        config = Bunch(config)
        config.use_one_hot_embeddings = True
        config.scope = "esim/bert"
        config.dropout_prob = 0.1
        config.label_type = "single_label"
        config.lstm_dim = 128
        config.num_heads = 12
        config.num_units = 768

        import json
        label_dict = json.load(open(FLAGS.label_id))

        # label_tensor = np.asarray(label_dict["class_ratio"]).astype(np.float32)
        label_tensor = None
        # config.loss = "focal_loss"

        json.dump(config, open(FLAGS.model_output+"/config.json", "w"))

        # os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu_id
        sess = tf.Session(config=sess_config)

        train_size = int(FLAGS.train_size/hvd.size())

        num_train_steps = int(
            train_size / FLAGS.batch_size * FLAGS.epoch)
        num_warmup_steps = int(num_train_steps * 0.1)

        num_storage_steps = int(train_size / FLAGS.batch_size)

        print(num_train_steps, num_warmup_steps, "=============")

        opt_config = Bunch({"init_lr":(2e-5/hvd.size()),
                            "num_train_steps":num_train_steps,
                            "num_warmup_steps":num_warmup_steps,
                            "train_op":"adam"})

        model_io_config = Bunch({"fix_lm":False})

        model_io_fn = model_io.ModelIO(model_io_config)

        num_choice = FLAGS.num_classes
        max_seq_length = FLAGS.max_length

        if FLAGS.model_type == "original":
            model_function = bert_order_classifier.classifier_model_fn_builder
        elif FLAGS.model_type == "attn":
            model_function = bert_order_classifier.classifier_attn_model_fn_builder
        elif FLAGS.model_type == "orignal_nonlinear":
            model_function = bert_order_classifier.classifier_model_fn_builder_v1
        elif FLAGS.model_type == "esim_bert":
            model_function = esim_bert.classifier_attn_model_fn_builder

        model_train_fn = model_function(
                                config,
                                num_choice,
                                init_checkpoint,
                                model_reuse=None,
                                load_pretrained=True,
                                model_io_fn=model_io_fn,
                                model_io_config=model_io_config,
                                opt_config=opt_config,
                                input_name=["a", "b"],
                                label_tensor=label_tensor,
                                not_storage_params=["adam", "adam_1"],
                                exclude_scope_dict={"task":"esim"})

        model_eval_fn = model_function(
                                    config,
                                    num_choice,
                                    init_checkpoint,
                                    model_reuse=True,
                                    load_pretrained=True,
                                    model_io_fn=model_io_fn,
                                    model_io_config=model_io_config,
                                    opt_config=opt_config,
                                    input_name=["a", "b"],
                                    label_tensor=label_tensor,
                                    not_storage_params=["adam", "adam_1"],
                                    exclude_scope_dict={"task":"esim"})

        def metric_fn(features, logits, loss):
            print(logits.get_shape(), "===logits shape===")
            pred_label = tf.argmax(logits, axis=-1, output_type=tf.int32)
            prob = tf.nn.softmax(logits)
            accuracy = correct = tf.equal(
                tf.cast(pred_label, tf.int32),
                tf.cast(features["label_ids"], tf.int32)
            )
            accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
            return {"accuracy":accuracy, "loss":loss, "pred_label":pred_label, "label_ids":features["label_ids"]}

        name_to_features = {
                "input_ids_a":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "input_mask_a":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "segment_ids_a":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "input_ids_b":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "input_mask_b":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "segment_ids_b":
                        tf.FixedLenFeature([max_seq_length], tf.int64),
                "label_ids":
                        tf.FixedLenFeature([], tf.int64),
        }

        def _decode_record(record, name_to_features):
            """Decodes a record to a TensorFlow example.
            """
            example = tf.parse_single_example(record, name_to_features)

            # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
            # So cast all int64 to int32.
            for name in list(example.keys()):
                t = example[name]
                if t.dtype == tf.int64:
                    t = tf.to_int32(t)
                example[name] = t
            return example

        params = Bunch({})
        params.epoch = FLAGS.epoch
        params.batch_size = FLAGS.batch_size
        # train_features = tf_data_utils.train_input_fn("/data/xuht/wsdm19/data/train.tfrecords",
        #                             _decode_record, name_to_features, params)
        # eval_features = tf_data_utils.eval_input_fn("/data/xuht/wsdm19/data/dev.tfrecords",
        #                             _decode_record, name_to_features, params)

        train_features = tf_data_utils.train_input_fn(FLAGS.train_file,
                                    _decode_record, name_to_features, params)
        eval_features = tf_data_utils.eval_input_fn(FLAGS.dev_file,
                                    _decode_record, name_to_features, params)

        [train_op, train_loss, train_per_example_loss, train_logits] = model_train_fn(train_features, [], tf.estimator.ModeKeys.TRAIN)
        [_, eval_loss, eval_per_example_loss, eval_logits] = model_eval_fn(eval_features, [], tf.estimator.ModeKeys.EVAL)
        result = metric_fn(eval_features, eval_logits, eval_loss)

        init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
        sess.run(init_op)

        sess.run(hvd.broadcast_global_variables(0))

        def eval_fn(result):
            i = 0
            total_accuracy = 0
            label, label_id = [], []
            # label_weight = []
            while True:
                try:
                    eval_result = sess.run(result)
                    total_accuracy += eval_result["accuracy"]
                    label_id.extend(eval_result["label_ids"])
                    label.extend(eval_result["pred_label"])
                    # for item in eval_result["label_ids"]:
                    #     label_weight.append(label_tensor[item])
                    i += 1
                except tf.errors.OutOfRangeError:
                    print("End of dataset")
                    break
            # f1 = f1_score(label_id, label, average="macro", sample_weight=label_weight)
            # accuracy = accuracy_score(label_id, label, sample_weight=label_weight)
            f1 = f1_score(label_id, label, average="macro")
            accuracy = accuracy_score(label_id, label)
            print("test accuracy accuracy {} {} f1 {}".format(total_accuracy/i,
                accuracy, f1))
            return total_accuracy/ i, f1

        def train_fn(op, loss):
            i = 0
            cnt = 0
            total_loss = 0.0
            while True:
                try:
                    [_, train_loss] = sess.run([op, loss])
                    total_loss += train_loss
                    i += 1
                    cnt += 1
                    if np.mod(i, num_storage_steps) == 0:
                        print(total_loss/cnt)
                        # model_io_fn.save_model(sess, "/data/xuht/wsdm19/data/model_11_15_focal_loss/oqmrc_{}.ckpt".format(int(i/8000)))
                        if hvd.rank() == 0:
                            model_io_fn.save_model(sess, FLAGS.model_output+"/oqmrc_{}.ckpt".format(int(i/num_storage_steps)))
                            print("==successful storing model=={}".format(int(i/num_storage_steps)))
                        total_loss = 0
                        cnt = 0
                except tf.errors.OutOfRangeError:
                    break
        print("===========begin to train============")
        train_fn(train_op, train_loss)
        print("===========begin to eval============")
        accuracy, f1 = eval_fn(result)
        print("==accuracy {} f1 {}==".format(accuracy, f1))
        # model_io_fn.save_model(sess, "/data/xuht/wsdm19/data/model_11_15_focal_loss/oqmrc.ckpt")
        if hvd.rank() == 0:
            model_io_fn.save_model(sess, FLAGS.model_output+"/oqmrc.ckpt")

if __name__ == "__main__":
    flags.mark_flag_as_required("eval_data_file")
    flags.mark_flag_as_required("output_file")
    flags.mark_flag_as_required("config_file")
    flags.mark_flag_as_required("init_checkpoint")
    flags.mark_flag_as_required("result_file")
    flags.mark_flag_as_required("vocab_file")
    flags.mark_flag_as_required("train_file")
    flags.mark_flag_as_required("dev_file")
    flags.mark_flag_as_required("max_length")
    flags.mark_flag_as_required("model_output")
    flags.mark_flag_as_required("gpu_id")
    flags.mark_flag_as_required("epoch")
    flags.mark_flag_as_required("num_classes")
    tf.app.run()


__author__ = 'Mehmet Mert Yildiran, mert.yildiran@bil.omu.edu.tr'

import datetime # Supplies classes for manipulating dates and times in both simple and complex ways.
import imutils # A series of convenience functions to make basic image processing functions such as translation, rotation, resizing, skeletonization etc.
import time # Provides various time-related functions.
import cv2 # (Open Source Computer Vision) is a library of programming functions mainly aimed at real-time computer vision.
import numpy # The fundamental package for scientific computing with Python.
import os # Provides a portable way of using operating system dependent functionality.
import multiprocessing # A package that supports spawning processes using an API similar to the threading module.
from cerebrum.vision.utilities import VisionMemoryUtil # BUILT-IN Memory operations package
import random # Pseudo-random number generators for various distributions.
import Tkinter

STABILIZATION_DETECTION = 5 # Number of frames to detect stabilization
NON_STATIONARY_PERCENTAGE = 70 # Percentage of frame for detecting NON-STATIONARY CAMERA. Like: ( height * width * float(X) / float(100) )
NON_ZERO_PERCENTAGE = 0 #  Percentage of frame(threshold) for detecting unnecessary movement
TARGET_HEIGHT = 360 # Number of horizontal lines for target video and processing. Like 720p, 360p etc.
MIN_AREA = 500 # Minimum area in square pixels to detect a motion
root = Tkinter.Tk()
SCREEN_WIDTH = root.winfo_screenwidth()
SCREEN_HEIGHT = root.winfo_screenheight()

class VisionPerception():

	# MAIN CODE BLOCK
	@staticmethod
	def start(video_input, vision_perception_stimulated):

		if video_input == "0": # If the video_input is None, then we are reading from webcam
			camera = cv2.VideoCapture(0)
			time.sleep(0.25)
		else: # Otherwise, we are reading from a video file
			time.sleep(0.25)
			camera = cv2.VideoCapture(video_input)

		referenceFrame = None # Initialize the reference frame in the video stream

		starting_time = None
		memory_data_thresh = []
		memory_data_frameDeltaColored = []

		(grabbed, first_frame) = camera.read() # Grab the first frame

		height, width = first_frame.shape[:2] # Get video height and width  from first frame(size)
		#if not height == 720 or not width == 1280:
		if float(width) / float(height) != float(16) / float(9):
			if video_input == "0":
				# There is a STUPIDTY in here
				pass
			else:
				raise ValueError('Aspect ratio of input stream must be [16:9]')
			#warnings.warn("Aspect ratio of input stream must be [16:9]")

		frame_counter = 1 # Define frame counter variable
		motion_detected = 0 # Delta situation checking variable
		delta_value_stack = [] # List of delta values
		non_stationary_camera = 0
		motion_counter = 0
		nonzero_toolow = 0

		beginning_of_stream = datetime.datetime.now()
		while True: # Loop over the frames of the video

			(grabbed, frame) = camera.read() # Grab the current frame and initialize the occupied/unoccupied
			if not grabbed: # If the frame could not be grabbed, then we have reached the end of the video
				break
			frame_counter += 1 # Increase frame counter's value

			if video_input == "0":
				# If we are capturing from camera fuck Time Correction, there is also a STUPIDTY in here
				pass
			else:
				# -------------------- TIME CORRECTION --------------------
				time_delta = datetime.datetime.now() - beginning_of_stream
				current_time_of_realworld = time_delta.seconds + time_delta.microseconds / float(1000000)
				current_time_of_stream = frame_counter / camera.get(cv2.cv.CV_CAP_PROP_FPS)
				diff_of_time = current_time_of_stream - current_time_of_realworld
				if abs(diff_of_time) > (1 / camera.get(cv2.cv.CV_CAP_PROP_FPS)):
					if diff_of_time > 0:
						time.sleep(1 / camera.get(cv2.cv.CV_CAP_PROP_FPS))
					else:
						(grabbed, frame) = camera.read() # Grab the current frame and initialize the occupied/unoccupied
						if not grabbed: # If the frame could not be grabbed, then we have reached the end of the video
							break
						frame_counter += 1 # Increase frame counter's value
						continue
				# -------------------- TIME CORRECTION --------------------

			delta_value = 0 # Delta Value for storing max continuous contour area for current frame

			frame = imutils.resize(frame, height=TARGET_HEIGHT) # Resize frame to 360p. Alternative resizing method:
			height, width = frame.shape[:2] # Get video height and width  from first frame(size)

			gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # Convert frame to grayscale

			gray = cv2.bilateralFilter(gray,9,75,75) # Blur current frame with Bilateral Filter for noise reduction

			if referenceFrame is None: # If Reference Frame is None, initialize it
				referenceFrame = gray
				continue

			frameDelta = cv2.absdiff(referenceFrame, gray) # Compute the absolute difference between the current frame and reference frame
			thresh = cv2.threshold(frameDelta, 12, 255, cv2.THRESH_BINARY)[1] # Apply OpenCV's threshold function to get binary frame

			thresh = cv2.dilate(thresh, None, iterations=1) # Dilate the thresholded image to fill in holes
			frameDeltaColored = cv2.bitwise_and(frame,frame, mask= thresh) # Bitwise and - to get delta frame

			# Find contours on thresholded image
			(cnts, _) = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL,
				cv2.CHAIN_APPROX_SIMPLE)

			contour_area_stack = [] # List of contour areas's values

			# Loop over the contours
			if cnts:
				for c in cnts: # Contour in Contours
					contour_area_stack.append(cv2.contourArea(c)) # Calculate contour area and append to contour stack
					if cv2.contourArea(c) > MIN_AREA: # If contour area greater than min area
						(x, y, w, h) = cv2.boundingRect(c) # Compute the bounding box for this contour
						cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2) # Draw it on the frame
				delta_value = max(contour_area_stack) # Assign max contour area to delta value

				if delta_value > MIN_AREA: # If max contour area (delta value) greater than min area
					motion_detected = 1 # Initialize delta situation

				if delta_value > (height * width * float(NON_STATIONARY_PERCENTAGE) / float(100)): # If delta value is too much
					non_stationary_camera = 1
					status_text = "WARNING: NON-STATIONARY CAMERA"
					frameDeltaColored = numpy.zeros_like(frame)
				else:
					non_stationary_camera = 0

				if cv2.countNonZero(thresh) < (height * width * float(NON_ZERO_PERCENTAGE) / float(100)): # If Non Zero count is too low
					nonzero_toolow = 1
					status_text = "WARNING: NON-ZERO TOO LOW"
					frameDeltaColored = numpy.zeros_like(frame)
				else:
					nonzero_toolow = 0

			if motion_detected: # If we are on delta situation

				if starting_time is None:
					starting_time = datetime.datetime.now() # Starting time of the memory
					vision_perception_stimulated.value = 1 # Vision perception stimulated

				if random.randint(0,2) == 1: # IMPORTANT
					memory_data_thresh.append(thresh.tostring())
					memory_data_frameDeltaColored.append(frameDeltaColored.tostring())
					#print type(memory_data_thresh[0])

				if not non_stationary_camera:
					status_text = "MOTION DETECTED"
				delta_value_stack.append(delta_value) # Append max contour area (delta value) to delta value stack

				if len(delta_value_stack) >= STABILIZATION_DETECTION: # If length of delta value stack is greater than or equal to STABILIZATION_DETECTION constant
					delta_value_stack.pop(0) # Pop first element of delta value stack
					# If minimum delta value is greater than (mean of last 5 frame - minimum area / 2) and maximum delta value is less than (mean of last 5 frame + minimum area / 2)
					if min(delta_value_stack) > (numpy.mean(delta_value_stack) - MIN_AREA / 2) and max(delta_value_stack) < (numpy.mean(delta_value_stack) + MIN_AREA / 2):
						ending_time = datetime.datetime.now() # Ending time of the memory
						vision_perception_stimulated.value = 0 # Vision perception NOT stimulated

						if memory_data_thresh and memory_data_frameDeltaColored:
							process4 = multiprocessing.Process(target=VisionMemoryUtil.add_memory, args=(memory_data_thresh, memory_data_frameDeltaColored, starting_time, ending_time)) # Define write memory process
							process4.start() # Start write memory process
						memory_data_thresh = []
						memory_data_frameDeltaColored = []
						starting_time = None

						motion_detected = 0 # Then video STABILIZED
						delta_value_stack = [] # Empty delta value stack
						referenceFrame = None  # Clear reference frame
						if not non_stationary_camera and not nonzero_toolow:
							motion_counter += 1
			else:
				if not non_stationary_camera and not nonzero_toolow:
					status_text = "MOTION UNDETECTED"
					frameDeltaColored = numpy.zeros_like(frame)

			# Draw the text and timestamp on the frame
			cv2.putText(frame, "Diff    : {}".format(delta_value), (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
			cv2.putText(frame, "Thresh : {}".format(MIN_AREA), (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
			cv2.putText(frame, "Frame : {}".format(frame_counter), (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
			cv2.putText(frame, "Status  : {}".format(status_text), (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
			cv2.putText(frame, datetime.datetime.now().strftime("%A %d %B %Y %I:%M:%S%p"), (10, frame.shape[0] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

			# Show the frames and record if the user presses ESC or q
			cv2.imshow("Original Frame", frame)
			cv2.moveWindow("Original Frame",50 * SCREEN_WIDTH / 1920,100 * SCREEN_HEIGHT / 1080)
			cv2.imshow("Frame Threshhold", thresh)
			cv2.moveWindow("Frame Threshhold",50 * SCREEN_WIDTH / 1920,550 * SCREEN_HEIGHT / 1080)
			cv2.imshow("Frame Delta", frameDelta)
			cv2.moveWindow("Frame Delta",1200 * SCREEN_WIDTH / 1920,550 * SCREEN_HEIGHT / 1080)
			cv2.imshow("Frame Delta Colored", frameDeltaColored)
			cv2.moveWindow("Frame Delta Colored",1200 * SCREEN_WIDTH / 1920,100 * SCREEN_HEIGHT / 1080)
			key = cv2.waitKey(1) & 0xFF

			# if the `ESC` or `q` key is pressed, break the loop
			if key == ord("q") or key == ord("\x1b"):
				os.system("killall python") # Temporary line for practicality in DEVELOPMENT
				break

		cv2.destroyAllWindows() # Close any open windows
		camera.release() # Release the capture device


import tensorflow as tf
import numpy as np
import cv2,os
from blending import Blending_test
from load_imgs import load_imgs
from model import Flow
import argparse
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.get_logger().setLevel('ERROR')

import wandb

parser = argparse.ArgumentParser()
parser.add_argument('--dataset',  type=str,
                    help='path to dataset',        default = 'dataset/view_light_time/pomegranate')
parser.add_argument('--type',     type=str, nargs= "*",
                    help='xfield type',            default = ['light','view','time'])
parser.add_argument('--dim',      type=int, nargs= "*",
                    help='dimension of input Xfields',   default = [3,3,3])
parser.add_argument('--factor',   type=int,
                    help='downsampling factor',      default = 6)
parser.add_argument('--nfg',      type=int,
                    help='capacity multiplier',    default = 4)
parser.add_argument('--num_n',    type=int,
                    help='number of neighbors',    default = 8)
parser.add_argument('--sigma',    type=float,
                    help='bandwidth parameter',    default = 0.1)
parser.add_argument('--br',      type=float,
                    help='baseline ratio',         default = 1)
parser.add_argument('--savepath', type=str,
                    help='saving path',      default = 'results/')
parser.add_argument('--scale',      type=int,
                    help='number of intermediate points',     default = 90)
parser.add_argument('--fps',      type=float,
                    help='output video frame rate',  default = 90)
parser.add_argument('--wandb_entity', type=str,
                    help='provide entity name for wandb init', default='wandb')
parser.add_argument('--wandb_project', type=str,
                    help='provide project name for wandb init', default='xfields')

args = parser.parse_args()

def run_test(args):
    print('---------- Initialize W&B run for experiment tracking----------\n')
    run = wandb.init(entity=args.wandb_entity, project=args.wandb_project, job_type='train')
    wandb.config.update(args)

    print('---------- Perform Testing ----------')

    savedir = args.savepath
    if not os.path.exists(savedir):
        os.mkdir( savedir )
    head_tail = os.path.split(args.dataset)
    savedir = os.path.join(savedir,head_tail[1])

    if not os.path.exists(savedir):
        raise NameError('There is no directory:\n %s'%(savedir))

    if not os.path.exists(os.path.join(savedir,"rendered videos") ):
        os.mkdir( os.path.join(savedir,"rendered videos") )
        print('creating directory %s'%(os.path.join(savedir,"rendered videos")))

    print('XField type: %s'%(args.type))
    print( 'Dimension of input xfield: %s'%(args.dim))
    print('output video fps: %d'%(args.fps))
    print('number of intermediate points for interpolation: %d'%(args.scale))

    images,coordinates,all_pairs,h_res,w_res = load_imgs(args)
    min_ = np.min(coordinates)
    max_ = np.max(coordinates)

    dims = args.dim
    num_n = args.num_n

    if num_n > np.prod(dims):
            num_n = np.prod(dims)

    input = tf.placeholder(tf.float32,shape=[1,1,1,len(dims)])

    num_output = len(args.type)*2
    with tf.variable_scope("gen_flows"):
        flows = Flow(input,h_res,w_res,num_output,args.nfg,min_,max_)

    if args.type == ['light','view','time']:

        with tf.variable_scope("gen_flows"):
                albedos = tf.Variable(tf.constant(1.0, shape=[dims[1]*dims[2],h_res, w_res,3]), name='albedo')
                index_albedo   = tf.placeholder(tf.int32,shape=(num_n,))
                albedo         = tf.gather(albedos,index_albedo,0)

    elif args.type == ['light']:

        with tf.variable_scope("gen_flows"):
            albedo = tf.Variable(tf.constant(1.0, shape=[1,h_res, w_res,3]), name='albedo')

    else:
        albedo =   tf.constant(1.0, shape=[1,h_res, w_res,3])

    input_N        = tf.placeholder(tf.float32,shape=[num_n,1,1,len(dims)])
    Neighbors_img  = tf.placeholder(tf.float32,shape=[num_n,h_res,w_res,3])
    Neighbors_flow = tf.placeholder(tf.float32,shape=[num_n,h_res,w_res,len(args.type)*2])

    interpolated = Blending_test(input,input_N,Neighbors_img,Neighbors_flow,flows,albedo,h_res,w_res,args)

    saver = tf.train.Saver(max_to_keep=1000)
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    ckpt=tf.train.get_checkpoint_state("%s/trained model/"%(savedir))
    if ckpt:
        print('\n loading pretrained model  '+ckpt.model_checkpoint_path)
        saver.restore(sess,ckpt.model_checkpoint_path)
    else:
        raise NameError('There is no pretrained model located at dir:\n %s/trained model/'%(savedir))

    precomputed_flows = []

    for i in range(len(coordinates)):
        flows_out = sess.run(flows,feed_dict={input:coordinates[[i],::]})
        precomputed_flows.append(flows_out[0,::])

    precomputed_flows = np.stack(precomputed_flows,0)

    if args.type == ['view'] or args.type == ['light'] or args.type == ['time']:

        theta = [np.pi/args.scale*i for i in range(args.scale+1)]

        X1 = 1 - np.cos(theta);
        X2 = 1 + np.cos(theta);
        Y1 = 1 + np.sqrt(1-(X1-1)**2)
        Y2 = 1 - np.sqrt(1-(X2-1)**2)

        X = np.append(X1,X2)
        Y = np.append(Y1,Y2)
        X = X/2
        Y = Y/2

        if args.type == ['view'] or args.type == ['light'] :

            X = X*(dims[1]-1)
            Y = Y*(dims[0]-1)
            rendering_path = np.transpose([X,Y])

        if args.type == ['time']:

            rendering_path = np.transpose([X*(dims[0]-1)])

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter('%s/rendered videos/rendered.mp4'%(savedir),fourcc, args.fps, (w_res,h_res))
        for id in range(len(X)):

                input_coord = np.array([[[rendering_path[id,:]]]])
                indices = np.argsort(np.sum(np.square(input_coord[0,0,0,:]-coordinates[:,0,0,:]),-1))[:num_n]

                input_coord_N   = coordinates[indices,::]
                input_Neighbors = images[indices,::]
                input_flows     = precomputed_flows[indices,::]

                im_out = sess.run(interpolated,feed_dict={        input         :input_coord,
                                                                  input_N       :input_coord_N,
                                                                  Neighbors_img :input_Neighbors,
                                                                  Neighbors_flow:input_flows,
                                                                          })
                im_out = np.minimum(np.maximum(im_out[0,::],0.0),1.0)
                out.write(np.uint8(im_out*255))

                print('\r interpolated image %d of %d'%(id+1,len(rendering_path)),end=" " )

        out.release()
        wandb.log({"rendered": wandb.Video('%s/rendered videos/rendered.mp4'%(savedir), fps=4, format="mp4")})

    if args.type == ['light','view','time']:

        print('\n number of neighbors for interpolation: %d'%(num_n))
        max_L = dims[0]-1
        max_V = dims[1]-1
        max_T = dims[2]-1

        X_L = np.linspace(0,max_L,max_L*args.scale)
        X_L = np.append(X_L,np.flip(X_L))
        X_V = np.linspace(0,max_V,max_V*args.scale)
        X_V = np.append(X_V,np.flip(X_V))
        X_T = np.linspace(0,max_T,max_T*args.scale)
        X_T = np.append(X_T,np.flip(X_T))
        middle_X_L = max_L*0.5*np.ones_like(X_L)
        middle_X_V = max_V*0.5*np.ones_like(X_V)
        middle_X_T = max_T*0.5*np.ones_like(X_T)

        all_dimensions = {'light'      :np.stack([   X_L,    middle_X_V,middle_X_T],1),
                          'view'       :np.stack([middle_X_L,    X_V,   middle_X_T],1),
                          'time'       :np.stack([middle_X_L,middle_X_V       ,X_T],1),
                          'light_view' :np.stack([X_L,            X_V,   middle_X_T],1),
                          'light_time' :np.stack([X_L,        middle_X_V,       X_T],1),
                          'view_time'  :np.stack([middle_X_L,    X_V,           X_T],1),
                          'light_view_time': np.stack([X_L,       X_V,           X_T],1) }

        for case, rendering_path in all_dimensions.items():

            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter('%s/rendered videos/rendered_%s.mp4'%(savedir,case),fourcc, args.fps, (w_res,h_res))

            print('\n --------- %s interpolation ---------'%(case))

            for id in range(len(rendering_path)):

                input_coord = np.array([[[rendering_path[id,:]]]])
                indices = np.argsort(np.sum(np.square(input_coord[0,0,0,:]-coordinates[:,0,0,:]),-1))[:num_n]

                input_coord_N   = coordinates[indices,::]
                input_Neighbors = images[indices,::]
                input_flows     = precomputed_flows[indices,::]

                time_idx  = indices //(dims[0]*dims[1])
                rest  = indices % (dims[0]*dims[1])
                view_idx  = rest % dims[1]
                albedo_index    = view_idx*dims[1] + time_idx

                im_out = sess.run(interpolated,feed_dict={ input         :input_coord,
                                                           input_N       :input_coord_N,
                                                           Neighbors_img :input_Neighbors,
                                                           Neighbors_flow:input_flows,
                                                           index_albedo  :albedo_index,
                                                                          })

                im_out = np.minimum(np.maximum(im_out[0,::],0.0),1.0)
                out.write(np.uint8(im_out*255))

                print('\r interpolated image %d of %d'%(id+1,len(rendering_path)),end=" " )

            out.release()
            wandb.log({"rendered": wandb.Video('%s/rendered videos/rendered_%s.mp4'%(savedir,case), fps=4, format="mp4")})

if __name__=='__main__':

    run_test(args)
    print('\n The interpolation result is located in the \'rendered videos\' folder.')


import torch
import torch.nn as nn
import numpy as np

from mmcv.cnn import ConvModule

from mmdet.core import (auto_fp16, build_bbox_coder, force_fp32, multi_apply,
                        multiclass_arb_nms, get_bbox_dim, bbox2type, arb2roi)
from mmdet.models.builder import HEADS, build_loss
from .obboxa_head import OBBoxAHead
from mmdet.models.losses import accuracy

@HEADS.register_module()
class OBBAConvFCBBoxHead(OBBoxAHead):
    r"""More general bbox head, with shared conv and fc layers and two optional
    separated branches.

    .. code-block:: none

                                    /-> cls convs -> cls fcs -> cls
        shared convs -> shared fcs   -> attribution convs -> cls fcs -> attribution cls
                                    \-> reg convs -> reg fcs -> reg
    """  # noqa: W605

    def __init__(self,
                 num_shared_convs=0,
                 num_shared_fcs=0,
                 num_cls_convs=0,
                 num_cls_fcs=0,
                 num_reg_convs=0,
                 num_reg_fcs=0,
                 num_attr_convs=0,
                 num_attr_fcs=0,
                 conv_out_channels=256,
                 fc_out_channels=1024,
                 conv_cfg=None,
                 norm_cfg=None,
                 *args,
                 **kwargs):
        super(OBBAConvFCBBoxHead, self).__init__(*args, **kwargs)
        assert (num_shared_convs + num_shared_fcs + num_cls_convs +
                num_cls_fcs + num_reg_convs + num_reg_fcs > 0)
        if num_cls_convs > 0 or num_reg_convs > 0:
            assert num_shared_fcs == 0
        if not self.with_cls:
            assert num_cls_convs == 0 and num_cls_fcs == 0
        if not self.with_reg:
            assert num_reg_convs == 0 and num_reg_fcs == 0
        if not self.with_attr:
            assert num_attr_convs == 0 and num_attr_fcs == 0

        self.num_shared_convs = num_shared_convs
        self.num_shared_fcs = num_shared_fcs
        self.num_cls_convs = num_cls_convs
        self.num_cls_fcs = num_cls_fcs
        self.num_attr_convs = num_attr_convs  # attributes
        self.num_attr_fcs = num_attr_fcs      # attributes
        self.num_reg_convs = num_reg_convs
        self.num_reg_fcs = num_reg_fcs
        self.conv_out_channels = conv_out_channels
        self.fc_out_channels = fc_out_channels
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg

        # add shared convs and fcs
        self.shared_convs, self.shared_fcs, last_layer_dim = \
            self._add_conv_fc_branch(
                self.num_shared_convs, self.num_shared_fcs, self.in_channels,
                True)
        self.shared_out_channels = last_layer_dim

        # add cls specific branch
        self.cls_convs, self.cls_fcs, self.cls_last_dim = \
            self._add_conv_fc_branch(
                self.num_cls_convs, self.num_cls_fcs, self.shared_out_channels)

        # add attribute cls specific branch, mixiaoxin
        self.attr_convs, self.attr_fcs, self.attr_last_dim = \
            self._add_conv_fc_branch(
                self.num_attr_convs, self.num_attr_fcs, self.shared_out_channels)

        # add reg specific branch
        self.reg_convs, self.reg_fcs, self.reg_last_dim = \
            self._add_conv_fc_branch(
                self.num_reg_convs, self.num_reg_fcs, self.shared_out_channels)

        if self.num_shared_fcs == 0 and not self.with_avg_pool:
            if self.num_cls_fcs == 0:
                self.cls_last_dim *= self.roi_feat_area
            if self.num_reg_fcs == 0:
                self.reg_last_dim *= self.roi_feat_area
            if self.num_attr_fcs == 0:
                self.attr_last_dim *= self.roi_feat_area

        self.relu = nn.ReLU(inplace=True)
        # reconstruct fc_cls and fc_reg since input channels are changed
        if self.with_cls:
            self.fc_cls = nn.Linear(self.cls_last_dim, self.num_classes + 1)
        if self.with_reg:
            out_dim_reg = self.reg_dim if self.reg_class_agnostic else \
                    self.reg_dim * self.num_classes
            self.fc_reg = nn.Linear(self.reg_last_dim, out_dim_reg)
        if self.with_attr:
            self.fc_attr = nn.Linear(self.attr_last_dim, self.num_attr)

    def _add_conv_fc_branch(self,
                            num_branch_convs,
                            num_branch_fcs,
                            in_channels,
                            is_shared=False):
        """Add shared or separable branch

        convs -> avg pool (optional) -> fcs
        """
        last_layer_dim = in_channels
        # add branch specific conv layers
        branch_convs = nn.ModuleList()
        if num_branch_convs > 0:
            for i in range(num_branch_convs):
                conv_in_channels = (
                    last_layer_dim if i == 0 else self.conv_out_channels)
                branch_convs.append(
                    ConvModule(
                        conv_in_channels,
                        self.conv_out_channels,
                        3,
                        padding=1,
                        conv_cfg=self.conv_cfg,
                        norm_cfg=self.norm_cfg))
            last_layer_dim = self.conv_out_channels
        # add branch specific fc layers
        branch_fcs = nn.ModuleList()
        if num_branch_fcs > 0:
            # for shared branch, only consider self.with_avg_pool
            # for separated branches, also consider self.num_shared_fcs
            if (is_shared
                    or self.num_shared_fcs == 0) and not self.with_avg_pool:
                last_layer_dim *= self.roi_feat_area
            for i in range(num_branch_fcs):
                fc_in_channels = (
                    last_layer_dim if i == 0 else self.fc_out_channels)
                branch_fcs.append(
                    nn.Linear(fc_in_channels, self.fc_out_channels))
            last_layer_dim = self.fc_out_channels
        return branch_convs, branch_fcs, last_layer_dim

    def init_weights(self):
        super(OBBAConvFCBBoxHead, self).init_weights()
        # conv layers are already initialized by ConvModule
        for module_list in [self.shared_fcs, self.cls_fcs, self.reg_fcs, self.attr_fcs]:
            for m in module_list.modules():
                if isinstance(m, nn.Linear):
                    nn.init.xavier_uniform_(m.weight)
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        # shared part
        if self.num_shared_convs > 0:
            for conv in self.shared_convs:
                x = conv(x)

        if self.num_shared_fcs > 0:
            if self.with_avg_pool:
                x = self.avg_pool(x)

            x = x.flatten(1)

            for fc in self.shared_fcs:
                x = self.relu(fc(x))
        # separate branches
        x_cls = x
        x_reg = x
        x_attr = x

        for conv in self.cls_convs:
            x_cls = conv(x_cls)
        if x_cls.dim() > 2:
            if self.with_avg_pool:
                x_cls = self.avg_pool(x_cls)
            x_cls = x_cls.flatten(1)
        for fc in self.cls_fcs:
            x_cls = self.relu(fc(x_cls))

        for conv in self.reg_convs:
            x_reg = conv(x_reg)
        if x_reg.dim() > 2:  # x_reg.dim() == 2
            if self.with_avg_pool:
                x_reg = self.avg_pool(x_reg)
            x_reg = x_reg.flatten(1)
        for fc in self.reg_fcs:
            x_reg = self.relu(fc(x_reg))

        # add attribution classification:
        for conv in self.attr_convs:
            x_attr = conv(x_attr)
        if x_attr.dim() > 2:  # x_attr.dim() == 2
            if self.with_avg_pool:
                x_attr = self.avg_pool(x_attr)
            x_attr = x_attr.flatten(1)
        for fc in self.attr_fcs:
            x_attr = self.relu(fc(x_attr))

        cls_score = self.fc_cls(x_cls) if self.with_cls else None
        bbox_pred = self.fc_reg(x_reg) if self.with_reg else None
        attr_score = self.fc_attr(x_attr) if self.with_attr else None
        #print('acls score size after fully convolution, ', acls_score)
        #acls_score = torch.sigmoid(acls_score) if self.with_acls else None # sigmoid 操作将在计算cross entropy loss时实现
        #print('acls score size after sigmoid, ', acls_score)
        return cls_score, bbox_pred, attr_score

    def _get_target_single(self, pos_bboxes, neg_bboxes, pos_gt_bboxes,
                           pos_gt_labels, pos_bboxes_attr, cfg):
        num_pos = pos_bboxes.size(0)
        num_neg = neg_bboxes.size(0)
        num_samples = num_pos + num_neg

        # original implementation uses new_zeros since BG are set to be 0
        # now use empty & fill because BG cat_id = num_classes,
        # FG cat_id = [0, num_classes-1]
        labels = pos_bboxes.new_full((num_samples, ),
                                     self.num_classes,
                                     dtype=torch.long)
        label_weights = pos_bboxes.new_zeros(num_samples)
        target_dim = self.reg_dim if not self.reg_decoded_bbox \
                else get_bbox_dim(self.end_bbox_type)
        bbox_targets = pos_bboxes.new_zeros(num_samples, target_dim)
        bbox_weights = pos_bboxes.new_zeros(num_samples, target_dim)
        attr_targets = pos_bboxes.new_full((num_samples, self.num_attr), 0, dtype=torch.int8)  # 属性targets
        attr_weights = pos_bboxes.new_zeros(num_samples)                                              # 属性targets的权重
        # pos_gt_masks_data = [np.array(mask[0]) for mask in pos_gt_masks.data.masks]
        # pos_gt_masks_data = np.array(pos_gt_masks_data)
        # print('raw mask size: ', len(pos_gt_masks.data.masks))

        if num_pos > 0:
            labels[:num_pos] = pos_gt_labels     # 类别标签赋值
            pos_weight = 1.0 if cfg.pos_weight <= 0 else cfg.pos_weight
            label_weights[:num_pos] = pos_weight
            if not self.reg_decoded_bbox:
                pos_bbox_targets = self.bbox_coder.encode(
                    pos_bboxes, pos_gt_bboxes)
            else:
                pos_bbox_targets = pos_gt_bboxes
            bbox_targets[:num_pos, :] = pos_bbox_targets
            bbox_weights[:num_pos, :] = 1
            # enlarge angle weight
            bbox_weights[:, -1] *= 2
            # orient = pos_gt_masks_data[:, 1] - pos_gt_masks_data[:, 5]
            # orient = (orient > 0.0)
            # print('orient: ', orient.shape)
            # print('num_positive ', num_pos)
            # print('number gt bbox: ', pos_gt_bboxes.shape)
            # attr_targets[:num_pos] = torch.from_numpy(orient)
            attr_targets[:num_pos] = pos_bboxes_attr
            attr_weights[:num_pos] = 1.0
            # for k in range(num_pos):
            #     if pos_gt_labels[k] == 6 or pos_gt_labels[k] == 7 or pos_gt_labels[k] == 10 or pos_gt_labels[k] == 11 or pos_gt_labels[k] == 12:
            #         attr_targets[k] = 1  # 这几个类别的正方向都初始化为1, 后续预测的值也变为1,避免计算loss时这个分支没有值导致网络不稳定

        if num_neg > 0:
            label_weights[-num_neg:] = 1.0
        #print('bbox_weights in get_targets_single_func: ', bbox_weights)
        return labels, label_weights, bbox_targets, bbox_weights, attr_targets, attr_weights

    def get_targets(self,
                    sampling_results,
                    gt_bboxes,
                    gt_labels,
                    rcnn_train_cfg,
                    concat=True):
        pos_bboxes_list = [res.pos_bboxes for res in sampling_results]
        neg_bboxes_list = [res.neg_bboxes for res in sampling_results]
        pos_gt_bboxes_list = [res.pos_gt_bboxes for res in sampling_results]
        pos_gt_labels_list = [res.pos_gt_labels for res in sampling_results]
        pos_bboxes_attr_list = [res.pos_bboxes_attr for res in sampling_results]

        labels, label_weights, bbox_targets, bbox_weights, attr_targets, attr_weights = multi_apply(
            self._get_target_single,
            pos_bboxes_list,
            neg_bboxes_list,
            pos_gt_bboxes_list,
            pos_gt_labels_list,
            pos_bboxes_attr_list,
            cfg=rcnn_train_cfg)

        if concat:
            labels = torch.cat(labels, 0)
            label_weights = torch.cat(label_weights, 0)
            bbox_targets = torch.cat(bbox_targets, 0)
            bbox_weights = torch.cat(bbox_weights, 0)
            attr_targets = torch.cat(attr_targets, 0)
            attr_weights = torch.cat(attr_weights, 0)
        return labels, label_weights, bbox_targets, bbox_weights, attr_targets, attr_weights

    @force_fp32(apply_to=('cls_score', 'bbox_pred', 'attr_score'))
    def loss(self,
             cls_score,
             bbox_pred,
             attr_score,
             rois,
             labels,
             label_weights,
             bbox_targets,
             bbox_weights,
             attr_targets,
             attr_weights,
             reduction_override=None):
        losses = dict()
        if cls_score is not None:
            avg_factor = max(torch.sum(label_weights > 0).float().item(), 1.)
            if cls_score.numel() > 0:
                losses['loss_cls'] = self.loss_cls(
                    cls_score,
                    labels,
                    label_weights,
                    avg_factor=avg_factor,
                    reduction_override=reduction_override)
                losses['acc'] = accuracy(cls_score, labels)
        if bbox_pred is not None:
            bg_class_ind = self.num_classes
            # 0~self.num_classes-1 are FG, self.num_classes is BG
            pos_inds = (labels >= 0) & (labels < bg_class_ind)
            # do not perform bounding box regression for BG anymore.
            target_dim = self.reg_dim
            if pos_inds.any():
                if self.reg_decoded_bbox:
                    bbox_pred = self.bbox_coder.decode(rois[:, 1:], bbox_pred)
                    target_dim = get_bbox_dim(self.end_bbox_type)
                if self.reg_class_agnostic:
                    pos_bbox_pred = bbox_pred.view(
                        bbox_pred.size(0), target_dim)[pos_inds.type(torch.bool)]
                else:
                    pos_bbox_pred = bbox_pred.view(
                        bbox_pred.size(0), -1,
                        target_dim)[pos_inds.type(torch.bool),
                                    labels[pos_inds.type(torch.bool)]]
                losses['loss_bbox'] = self.loss_bbox(
                    pos_bbox_pred,
                    bbox_targets[pos_inds.type(torch.bool)],
                    bbox_weights[pos_inds.type(torch.bool)],
                    avg_factor=bbox_targets.size(0),
                    reduction_override=reduction_override)
                # print('bbox_weights in second_stage_loss_func: ', bbox_weights, '\n',
                #       'pos_bbox_pred: ', pos_bbox_pred, '\n',
                #       'bbox_targets: ', bbox_targets[pos_inds.type(torch.bool)])

            else:
                losses['loss_bbox'] = bbox_pred.sum() * 0
        if attr_score is not None:
            avg_factor = max(torch.sum(attr_weights > 0).float().item(), 1.)  # 参与分母计算的数值
            if self.use_sigmoid_attr:
                attr_score = attr_score.sigmoid()

            if attr_score.numel() > 0:
                losses['loss_attr'] = self.loss_attr(
                    attr_score.reshape(-1),
                    attr_targets.reshape(-1),
                    attr_weights,
                    avg_factor=avg_factor,
                    reduction_override='mean')
                # losses['num_pos_targets'] = torch.sum(attr_weights > 0).float()
                # loss_iou=(self.loss_attr(
                #     attr_score.reshape(-1, 1),
                #     attr_targets.reshape(-1, 1),
                #     reduction_override='none'
                # )*attr_weights.reshape(-1, 1).float()
                # ).sum() / avg_factor
                #
                # losses['test_cross_entropy_loss'] = loss_iou

            else:
                losses['loss_attr'] = 0

        return losses

    def simple_test_bboxes_attr(self,
                           x,
                           img_metas,
                           proposals,
                           rcnn_test_cfg,
                           rescale=False):
        """Test only det bboxes without augmentation."""
        rois = arb2roi(proposals, bbox_type=self.bbox_head.start_bbox_type)
        bbox_results = self._bbox_forward(x, rois)
        img_shape = img_metas[0]['img_shape']
        scale_factor = img_metas[0]['scale_factor']
        det_bboxes, det_labels, det_attr = self.bbox_head.get_bboxes( # 跳转到obbxa_head.get_bboxes
            rois,
            bbox_results['cls_score'],
            bbox_results['bbox_pred'],
            bbox_results['attr_score'],
            img_shape,
            scale_factor,
            rescale=rescale,
            cfg=rcnn_test_cfg)
        return det_bboxes, det_labels, det_attr

@HEADS.register_module()
class OBBAShared2FCBBoxHead(OBBAConvFCBBoxHead):

    def __init__(self, fc_out_channels=1024, *args, **kwargs):
        super(OBBAShared2FCBBoxHead, self).__init__(
            num_shared_convs=0,
            num_shared_fcs=2,
            num_cls_convs=0,
            num_cls_fcs=0,
            num_reg_convs=0,
            num_reg_fcs=0,
            fc_out_channels=fc_out_channels,
            *args,
            **kwargs)


from model import Encoder, Decoder, Discriminator, Generator
from load_data import TrainDataProvider
import matplotlib.pyplot as plt
import tensorflow as tf
import pandas as pd
import numpy as np
import argparse
import datetime
import glob
import os
import time

class Trainer:

    def __init__(self, opt):

        self.data_dir = opt.data_dir
        self.batch_size = opt.batch_size
        self.img_size = opt.img_size
        self.train_name = opt.train_name
        self.embedding_name = opt.embedding_name

        self.embeddings = pd.read_pickle(os.path.join(self.data_dir, self.embedding_name))
        self.embedding_num = self.embeddings.shape[0]
        self.embedding_dim = self.embeddings.shape[3]

        self.fonts_num = len(self.embeddings)

        self.data_provider = TrainDataProvider(self.data_dir, self.train_name)
        self.total_batches = self.data_provider.compute_total_batch_num(self.batch_size)
        print("total batches:", self.total_batches)

    def gd_loss(self, real_source, real_target, fake_target, encoded_source, embedding_ids,
                En, Dis, L1_penalty, Lconst_penalty, Lcat_penalty):

        l1_criterion = tf.keras.losses.MeanAbsoluteError()
        bce_criterion = tf.keras.losses.BinaryCrossentropy(from_logits=True)
        mse_criterion = tf.keras.losses.MeanSquaredError()

        mae = tf.keras.losses.MeanAbsoluteError()

        real_TS = tf.concat([real_source, real_target], axis=3)
        fake_TS = tf.concat([real_source, fake_target], axis=3)

        real_score, real_score_logit, real_cat_logit = Dis(real_TS)
        fake_score, fake_score_logit, fake_cat_logit = Dis(fake_TS)

        # G_1. constant loss
        encoded_fake = En(fake_target)[0]
        const_loss = Lconst_penalty * mse_criterion(encoded_source, encoded_fake)

        # G_2, D_1. category loss
        real_category = tf.convert_to_tensor(np.eye(self.fonts_num)[embedding_ids])
        real_category_loss = Lcat_penalty * bce_criterion(real_category, real_cat_logit)
        fake_category_loss = Lcat_penalty * bce_criterion(real_category, fake_cat_logit)
        category_loss = 0.5 * (real_category_loss + fake_category_loss)

        # D_2. binary loss - T/F
        one_labels = tf.ones([self.batch_size, 1])
        zero_labels = tf.zeros([self.batch_size, 1])

        real_binary_loss = bce_criterion(real_score, one_labels)
        fake_binary_loss = bce_criterion(fake_score, zero_labels)
        binary_loss = 10 * (real_binary_loss + fake_binary_loss)

        # G_3. L1 loss between real and fake images
        l1_loss = L1_penalty * l1_criterion(real_target, fake_target)

        # G_4. cheat loss for generator to fool discriminator
        cheat_loss = 10 * bce_criterion(fake_score, one_labels)

        # g_loss, d_loss
        g_loss = const_loss + fake_category_loss + l1_loss + cheat_loss
        d_loss = category_loss + binary_loss
        return g_loss, d_loss, [cheat_loss, l1_loss, fake_category_loss, const_loss, binary_loss, category_loss]

    def generate_images(self, real_src, real_tar, En, De, embeddings, embedding_ids):

        fake_tar, encoded_source, encode_layers = Generator(real_src, En, De, embeddings, embedding_ids)

        plt.figure(figsize=(15, 15))
        display_list = [real_src[0], real_tar[0], fake_tar[0]]
        title = ['Input Image', 'Ground Truth', 'Predicted Image']

        for i in range(3):
            plt.subplot(1, 3, i + 1)
            plt.title(title[i])
            img = tf.concat([display_list[i]] * 3, axis=2)
            plt.imshow(img * 0.5 + 0.5)
            plt.axis('off')
        plt.show()

    def train_step(self, source, target, epoch, embeddings, embedding_ids, L1_penalty, Lconst_penalty, Lcat_penalty,
                   En, De, Dis, g_optimizer, d_optimizer, summary_writer):

        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
            fake_target, encoded_source, encode_layers = Generator(source, En, De, embeddings, embedding_ids)
            gen_loss, disc_loss, loss_list = self.gd_loss(source, target, fake_target, encoded_source, embedding_ids,
                                                          En, Dis, L1_penalty, Lconst_penalty, Lcat_penalty)

        ge_trainable_variables = En.trainable_variables + De.trainable_variables

        generator_gradients = gen_tape.gradient(gen_loss,
                                                ge_trainable_variables)
        discriminator_gradients = disc_tape.gradient(disc_loss,
                                                     Dis.trainable_variables)

        g_optimizer.apply_gradients(zip(generator_gradients,
                                        ge_trainable_variables))
        d_optimizer.apply_gradients(zip(discriminator_gradients,
                                        Dis.trainable_variables))

        with summary_writer.as_default():
            tf.summary.scalar('G_gen_loss', gen_loss, step=epoch)
            tf.summary.scalar('D_disc_loss', disc_loss, step=epoch)
            tf.summary.scalar('G1_const_loss', loss_list[3], step=epoch)
            tf.summary.scalar('G2_fake_category_loss', loss_list[2], step=epoch)
            tf.summary.scalar('G3_l1_loss', loss_list[1], step=epoch)
            tf.summary.scalar('G4_cheat_loss', loss_list[0], step=epoch)
            tf.summary.scalar('D1_category_loss', loss_list[5], step=epoch)
            tf.summary.scalar('D2_binary_loss', loss_list[4], step=epoch)

    def train(self, max_epoch, schedule, lr=0.001, fine_tune=False, flip_labels=False,
              restore=None, ckpt_dir=None, with_charid=True, model_save_step=1):
        max_epoch = opt.max_epoch
        schedule = opt.schedule
        lr = opt.learning_rate

        if not fine_tune:
            L1_penalty, Lconst_penalty, Lcat_penalty = 100, 50, 100
        else:
            L1_penalty, Lconst_penalty, Lcat_penalty = 500, 1000, 500

        En = Encoder()
        De = Decoder()
        Dis = Discriminator(cat_num=self.fonts_num)

        g_optimizer = tf.keras.optimizers.Adam(lr, beta_1=0.5)
        d_optimizer = tf.keras.optimizers.Adam(lr, beta_1=0.5)

        log_dir = "./logs/"
        summary_writer = tf.summary.create_file_writer(log_dir + datetime.datetime.now().strftime("%Y_%m_%d-%H_%M_%S"))

        checkpoint_dir = './ckpt'
        checkpoint = tf.train.Checkpoint(generator_optimizer=g_optimizer, discriminator_optimizer=d_optimizer,
                                         encoder=En, decoder=De, discriminator=Dis)

        if restore:
            files = glob.glob(os.path.join(ckpt_dir, '*'))
            ckpt_list = [file[:file.rfind('.')] for file in files if file.endswith(".index")]

            prev_epoch = int(ckpt_list[-1].split('-')[0][-5:-2])

            g_optimizer = tf.keras.optimizers.Adam()
            d_optimizer = tf.keras.optimizers.Adam()

            En = Encoder()
            De = Decoder()
            Dis = Discriminator(cat_num=self.fonts_num)

            checkpoint = tf.train.Checkpoint(generator_optimizer=g_optimizer, discriminator_optimizer=d_optimizer,
                                       encoder=En, decoder=De, discriminator=Dis)

            checkpoint.restore(tf.train.latest_checkpoint(ckpt_dir))
            print("%d epoch trained model has restored" % prev_epoch)
        else:
            prev_epoch = 0
            print("New model training start!!")

        start = time.time()
        for epoch in range(max_epoch - prev_epoch):

            epoch += prev_epoch
            train_batch_iter = self.data_provider.get_train_iter(self.batch_size, with_charid=with_charid)

            if (epoch + 1) % schedule == 0:
                updated_lr = max(lr / 2, 0.0001)
                g_optimizer = tf.keras.optimizers.Adam(updated_lr, beta_1=0.5)
                d_optimizer = tf.keras.optimizers.Adam(updated_lr, beta_1=0.5)
                if lr != updated_lr:
                    print("decay learning rate from %.5f to %.5f" % (lr, updated_lr))
                lr = updated_lr

                updated_Lcat = min(10000, 2 * Lcat_penalty)
                if Lcat_penalty != updated_Lcat:
                    print("increase category penalty from %.5f to %.5f" % (Lcat_penalty, updated_Lcat))
                Lcat_penalty = updated_Lcat

            perc = 0
            for i, batch in enumerate(train_batch_iter):

                if i == 0:
                    val_embedding_ids = batch[0]
                    val_source = tf.reshape(batch[2][:, :, :, 0], [self.batch_size, self.img_size, self.img_size, 1])
                    val_target = tf.reshape(batch[2][:, :, :, 1], [self.batch_size, self.img_size, self.img_size, 1])

                    print('\n %d EPOCHS :' % (epoch + 1))
                    self.generate_images(val_source, val_target, En, De, self.embeddings, val_embedding_ids)

                    print("Learning Rate : {0}, L1_penalty : {1}, Lconst_penalty : {2}, Lcat_penalty : {3} \n".format(
                        lr, L1_penalty, Lconst_penalty, Lcat_penalty))

                perc_check = perc
                perc = (i / self.total_batches) * 100
                if (perc // 3) - (perc_check // 3) != 0:
                    print('....', end='')

                if with_charid:
                    font_ids, char_ids, batch_images = batch
                else:
                    font_ids, batch_images = batch

                embedding_ids = font_ids

                if flip_labels:
                    np.random.shuffle(embedding_ids)

                real_target = tf.reshape(batch_images[:, :, :, 1], [self.batch_size, self.img_size, self.img_size, 1])
                real_source = tf.reshape(batch_images[:, :, :, 0], [self.batch_size, self.img_size, self.img_size, 1])

                self.train_step(real_source, real_target, epoch, self.embeddings, embedding_ids,
                                L1_penalty, Lconst_penalty, Lcat_penalty, En, De, Dis,
                                g_optimizer, d_optimizer, summary_writer)

            if (epoch + 1) % model_save_step == 0:
                checkpoint_prefix = os.path.join(checkpoint_dir, f"ckpt_{epoch + 1}EP")
                checkpoint.save(file_prefix=checkpoint_prefix)

            if epoch == 0:
                print("\n It takes about %.0fs for 1 EPOCH" % float(time.time() - start))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', required=True, help='the path Where pickledata exist')
    parser.add_argument('--embedding_name', required=True, help='the name of embedding file')
    parser.add_argument('--train_name', required=True, help='the name of training file')
    parser.add_argument('--batch_size', default=128, type=int, help='batch size')
    parser.add_argument('--img_size', default=128, type=int, help='image size')
    parser.add_argument('--max_epoch', default=1000, type=int, help='maximum of training epoch')
    parser.add_argument('--schedule', default=50, type=int, help='the term of weight updating')
    parser.add_argument('--learning_rate', default=0.001, type=float, help='learning rate')
    parser.add_argument('--fine_tune', action='store_true', help='Whether fine tuning')

    opt = parser.parse_args()

    trainer = Trainer(opt)
    trainer.train(opt)


import numpy as np
from play_sound import play_sound
import time

def compliment_machine(name,N):

    adjectives = [
        'cute',
        'strong',
        'funny',
        'fg',
        'schmooga',
        'lovely',
        'five years old',
        'mnbvcxzertyui',
        'a really good piano player',
    ]

    time.sleep(0.1)
    for i in range(int(N)):
        text = '{} is {}!'.format(name,np.random.choice(adjectives))

        # time.sleep(0.25)
        print(text)
        # time.sleep(0.1)
        play_sound(text)
        play_sound("I'm going to compliment {} {} more times".format(name,int(N)-i-1))
        # time.sleep(0.25)

if __name__ == '__main__':
    print('')
    name = 'Elliott'
    N = input('how many times do you want to say {}? '.format(name))
    print('\n')
    compliment_machine(name,N)


#! /usr/bin/env python

import tensorflow as tf
import yaml
from model import AttnNet

import numpy as np

#import memory_saving_gradients
# monkey patch tf.gradients to point to our custom version, with automatic checkpoint selection
#tf.__dict__["gradients"] = memory_saving_gradients.gradients_memory

def _parse_function(example, max_atom):
    features = {
        "na"    : tf.FixedLenFeature([], tf.int64),
        "energy": tf.FixedLenFeature([], tf.float32),
        "serial": tf.FixedLenFeature([max_atom], tf.int64),
        "latt"  : tf.FixedLenFeature([3, 3], tf.float32),
        "coord":  tf.FixedLenFeature([max_atom, 3], tf.float32),
        "stress": tf.FixedLenFeature([3, 3], tf.float32) ,
        "force" : tf.FixedLenFeature([max_atom, 3], tf.float32), }
    parsed_features = tf.parse_single_example(example, features)
    #parsed_features["coords"] = tf.reshape(parsed_features["coords"], (max_atom, 3))
    #parsed_features["force"] = tf.reshape(parsed_features["force"], (max_atom, 3))
    #parsed_features["latt"] = tf.reshape(parsed_features["latt"], (3, 3))
    #parsed_features["stress"] = tf.reshape(parsed_features["stress"], (3, 3))
    return parsed_features

class ModelTrainer(object):
    """
        This class is aimed at training DensityNet model
    """

    def __init__(self, config):
        self.config = yaml.load(open(config))
        self.max_atom = self.config["max_atom_num"]

        self.n_atom_embed = self.config["n_atom_embed"]
        self.n_kmesh = self.config["n_kmesh"]
        self.n_trans = self.config["n_trans"]
        self.n_heads = self.config["n_heads"]
        self.n_ff = self.config["n_ff"]
        self.n_zetas = self.config["n_zetas"]

        self.train_path = self.config["train_file"]
        self.valid_path = self.config["valid_file"]
        self.nepoch = self.config["nepoch"]
        self.bs = self.config["batch_size"]
        self.lr = self.config["learning_rate"]
        self.model_path = self.config.get("model_path")

        self.attn_net = AttnNet(self.max_atom, self.n_atom_embed, self.n_kmesh, self.n_trans, self.n_heads, self.n_ff, self.n_zetas)
        self.valid_ops = self._get_valid_ops()
        self.train_ops = self._get_train_ops()

        #config = tf.ConfigProto()
        #config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

        self.train_session = tf.Session(graph=self.attn_net.train_graph) #, config=config)

        saver = self.train_saver
        s = self.train_session
        s.run(self.train_init)

        if self.model_path:
            print("Before training, loading model at:", self.model_path)
            saver.restore(s, self.model_path)
        self.valid_session = tf.Session(graph=self.attn_net.valid_graph) #, config=config)

    def _get_valid_ops(self):

        with self.attn_net.valid_graph.as_default():

            dataset = tf.data.TFRecordDataset(tf.constant([self.valid_path]))
            dataset = dataset.map(lambda x: _parse_function(x, self.max_atom))
            dataset = dataset.batch(self.bs)
            dataset = dataset.repeat(1)
            #iterator = dataset.make_one_shot_iterator()
            iterator = dataset.make_initializable_iterator()
            features = iterator.get_next()

            loss_t, energy_loss_t, force_loss_t, stress_loss_t = \
                self.attn_net.validate(features["serial"], features["coord"], features["latt"],
                    features["force"], features["energy"], features["stress"])
            global_step = tf.Variable(0, name='global_step',trainable=False)

            tf.summary.scalar("0.loss", loss_t)
            tf.summary.scalar("1.energy_loss", energy_loss_t)
            tf.summary.scalar("2.force_loss", force_loss_t)
            tf.summary.scalar("3.stress_loss", stress_loss_t)

            ms_op = tf.summary.merge_all()
            self.valid_init = tf.global_variables_initializer()
            self.valid_saver = tf.train.Saver()
            self.valid_iterator = iterator
            self.valid_writer = tf.summary.FileWriter("./valid_logs", self.attn_net.valid_graph)
            return ms_op, global_step, [loss_t, energy_loss_t, force_loss_t, stress_loss_t]

    def valid(self, model_path = "./models/model.ckpt"):

        writer = self.valid_writer
        saver = self.valid_saver
        s = self.valid_session
        s.run(self.valid_init)
        s.run(self.valid_iterator.initializer)
        pdata = [0.0, 0.0, 0.0, 0.0]
        cnt = 0

        #if self.model_path:
        #    print("Loading model at:", self.model_path)
        #    saver.restore(s, self.model_path)
        # model_path = "./models/model.ckpt"
        saver.restore(s, model_path)

        while True:
            try:
                ms, step, pdatat = s.run(self.valid_ops)
                writer.add_summary(ms, step)
                pdata[0] += pdatat[0].sum()
                pdata[1] += pdatat[1].sum()
                pdata[2] += pdatat[2].sum()
                pdata[3] += pdatat[3].sum()
                cnt += 1
            except tf.errors.OutOfRangeError as e:
                print("Avg Loss: {:.3f}, Avg Energy RMS: {:.3f}, Avg Force RMS: {:.3f}, Avg Stress RMS: {:.3f}".format(
                    pdata[0]/cnt, pdata[1]/cnt, pdata[2]/cnt, pdata[3]/cnt))
                #print("Fininshed training...")
                break

    def _get_train_ops(self):

        with self.attn_net.train_graph.as_default():

            dataset = tf.data.TFRecordDataset(tf.constant([self.train_path]))
            dataset = dataset.map(lambda x: _parse_function(x, self.max_atom))
            dataset = dataset.shuffle(buffer_size=200000)
            dataset = dataset.batch(self.bs)
            #dataset = dataset.repeat(self.nepoch)
            dataset = dataset.repeat(1)
            #iterator = dataset.make_one_shot_iterator()
            iterator = dataset.make_initializable_iterator()
            features = iterator.get_next()

            loss_t, energy_loss_t, force_loss_t, stress_loss_t = \
                self.attn_net.train(features["serial"], features["coord"], features["latt"],
                    features["force"], features["energy"], features["stress"])

            global_step = tf.Variable(0, name='global_step',trainable=False)
            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                #optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=0.99, beta2=0.999)
                optimizer = tf.train.AdamOptimizer(learning_rate=self.lr)
                optimizer = tf.contrib.estimator.clip_gradients_by_norm(optimizer, clip_norm=1.0)
                optim = optimizer.minimize(loss_t, global_step=global_step)
            #print(tf.trainable_variables()); import sys; sys.exit()
            tf.summary.scalar("0.loss", loss_t)
            tf.summary.scalar("1.energy_loss", energy_loss_t)
            tf.summary.scalar("2.force_loss", force_loss_t)
            tf.summary.scalar("3.stress_loss", stress_loss_t)

            for var in tf.trainable_variables():
                tf.summary.histogram(var.op.name, var)

            #print(dir(optim.inputs))
            #print(optim.inputs())
            #import sys; sys.exit()
            #grads_vars = tf.gradients(loss_t, tf.trainable_variables())
            #for var in grads_vars:
            #    tf.summary.histogram(var.op.name, var)
            #for grad, var in grads_vars:
            #    if grad is not None:
            #        tf.summary.histogram(var.op.name + '/gradients', grad)

            ms_op = tf.summary.merge_all()
            #check_op = tf.add_check_numerics_ops()
            self.train_init = tf.global_variables_initializer()
            self.train_saver = tf.train.Saver()
            self.train_iterator = iterator
            self.train_writer = tf.summary.FileWriter("./logs", self.attn_net.train_graph)
            return loss_t, optim, ms_op, global_step, self.attn_net.avg_na, self.attn_net.check_ops
            #return check_op, loss_t, optim, ms_op, global_step, self.attn_net.avg_na, self.attn_net.check_ops

    def train(self):

        writer = self.train_writer
        saver = self.train_saver

        s = self.train_session
        #s.run(self.train_init)
        s.run(self.train_iterator.initializer)

        while True:
            try:
                loss, _, ms, step, avg_na, ret = s.run(self.train_ops)
                #_, loss, _, ms, step, avg_na, ret = s.run(self.train_ops)
                writer.add_summary(ms, step)
                print("Current iteration: {:5d}, Current loss: {:.2f}, avg_na: {:.2f}".format(step, loss, avg_na), end="\r")
                #import sys; sys.exit()
            except tf.errors.OutOfRangeError as e:
                print("")
                #print("\nFininshed training...")
                break
        #print("")
        print("Save model at: {}".format(saver.save(s, "./models/model.ckpt")))

    def train_and_eval(self):
        tf.logging.set_verbosity(tf.logging.ERROR)
        print("###")

        for i in range(self.nepoch):
            print("In epoch #{}, ".format(i+1), end="\n")
            self.train()
            #self.valid()
            if (i+1)%10==0: self.valid()

        self.train_session.close()
        self.valid_session.close()

if __name__ == "__main__":
    mt = ModelTrainer("./config.yml")
    mt.train_and_eval()
    #mt.train()
    #mt.valid()
    #mt.valid()
    #for i in range(100):
    #    print("In epoch #{}, ".format(i+1), end="")
    #    mt.train()