import numpy as np
from pyfaidx import Fasta
import offsetbasedgraph as obg
import logging
from collections import namedtuple, deque
from itertools import chain, dropwhile, takewhile

Variant = namedtuple("Variant", ["offset", "ref", "alt"])
FullVariant = namedtuple("FullVariant", ["offset", "ref", "alt", "precence"])

def find_valid_haplotypes(variants, combination):
    combination = np.asanyarray(combination)
    if not np.any(combination):
        return np.arange(1136)
    nonzero = np.flatnonzero(combination)
    informative_haplotypes = []
    informative_variants = [variant for i, variant in enumerate(variants)
                            if i in nonzero]
    for i, variant in enumerate(informative_variants):
        informative_haplotypes.extend(
            variant.precence.get_samples(combination[nonzero[i]]))
    f = np.unique(informative_haplotypes)
    if not f.size:
        return np.array([])
    for code, variant in zip(combination, variants):
        f = variant.precence.get_samples(code, f)
    return f

class VariantPrecence:
    strict = False
    accept_ref = False

    def __init__(self, precence):
        self._precence = precence

    def get_samples(self, variant, f=None):
        s = self._precence if f is None else self._precence[f, :]
        if self.accept_ref and variant == 0:
            return f if f is not None else np.arange(self._precence.shape[0])
        match = np.any((s == variant) | (s == -1),
                       axis=-1)
        res = np.flatnonzero(match)
        if f is not None:
            return f[res]
        return res

    @classmethod
    def from_line(cls, line):
        if cls.strict:
            parts = line.replace(".", "0").replace("/", "|").split("\t")
        else:
            parts = line.replace(".", "-1").replace("/", "|").split("\t")
        haplo_type = [part.split(":", 2)[0].split("|") for part in parts]
        precence = np.array([[int(a), int(b)] for a, b in haplo_type])
        return cls(precence)

    @classmethod
    def join(cls, precences, counts):
        offsets = [0] + list(np.cumsum(counts))[:-1]
        combined = np.zeros_like(precences[0]._precence)
        for precence in precences:
            combined[precence._precence < 0] = -1
        for precence, offset in zip(precences, offsets):
            combined[precence._precence > 0] = precence._precence[precence._precence > 0] + offset
        return cls(combined)

    def __repr__(self):
        return "P(%s)" % np.count_nonzero(self._precence > 0)

class VariantList:
    def __init__(self, start, end):
        self._start = int(start)
        self._end = int(end)
        self._variants = []

    def append(self, variant):
        self._variants.append(variant)

    def finalize(self):
        return [FullVariant(variant.offset-self._start,
                            variant.ref, variant.alt,
                            variant.precence) for variant in self._variants]

        final_list = []
        cur_end = 0
        var_buffer = []

        for variant in self._variants:
            transformed_var = FullVariant(
                variant.offset-self._start,
                variant.ref, variant.alt,
                variant.precence)
            assert transformed_var.offset <= self._end-self._start, (variant, self._start, transformed_var.offset, self._end-self._start)
            if transformed_var.offset >= cur_end:
                if var_buffer:
                    final_list.append(self._join_variants(var_buffer))
                var_buffer = []
            var_buffer.append(transformed_var)
            cur_end = max(transformed_var.offset + len(transformed_var.ref),
                          cur_end)
        if var_buffer:
            final_list.append(self._join_variants(var_buffer))
        return final_list

    def _join_variants(self, variants):
        if len(variants) == 1:
            return variants[0]
        start = min(variant.offset for variant in variants)
        end = max(variant.offset + len(variant.ref) for variant in variants)
        ref = []*(end-start)
        for variant in variants:
            ref[variant.offset-start:variant.offset-start+len(variant.ref)] = variant.ref
        ref = "".join(ref)
        alts = [ref[:variant.offset-start] + v_alt + ref[variant.offset+len(variant.ref)-start:]
                for variant in variants for v_alt in variant.alt]
        counts = [len(variant.alt) for variant in variants]
        precences = [variant.precence for variant in variants]
        combined_precence = VariantPrecence.join(precences, counts)
        return FullVariant(start, ref, alts, combined_precence)
def prune_seqs(ref, alts):
    offset = 0
    for cs in zip(*([ref]+alts)):
        if all(c == cs[0] for c in cs):
            offset += 1
        else:
            break
    return ref[offset:], [alt[offset:] for alt in alts], offset

class VCF:
    def __init__(self, vcf_file_name):
        self.file_name = vcf_file_name
        self.f = open(self.file_name)
        self.cur_lines = []

    def _join_variants(self, variants):
        if len(variants) == 1:
            return variants[0]
        start = min(variant.offset for variant in variants)
        end = max(variant.offset + len(variant.ref) for variant in variants)
        ref = []*(end-start)
        for variant in variants:
            ref[variant.offset-start:variant.offset-start+len(variant.ref)] = variant.ref
        ref = "".join(ref)
        alts = [ref[:variant.offset-start] + v_alt + ref[variant.offset+len(variant.ref)-start:]
                for variant in variants for v_alt in variant.alt]
        return Variant(start, ref, alts)

    def get_haplotype_sequences(self, haplotype, interval, ref_seq, lines=None):
        interval = (int(interval[0]), int(interval[1]))
        if lines is None:
            lines = self.f
        line_parts = (line.split("\t", 9) for line in lines
                      if not line.startswith("#"))
        interval_line_parts = dropwhile(
            lambda parts: int(parts[1])-1 < interval[0],
            line_parts)
        interval_line_parts = takewhile(
            lambda parts: int(parts[1])-1 <= interval[1],
            interval_line_parts)

        seqs = [[], []]
        cur_positions = [0, 0]
        for parts in interval_line_parts:
            variant = VariantPrecence.from_line(parts[-1])._precence[haplotype]
            variant_start = int(parts[1])-1-int(interval[0])
            variant_end = variant_start + len(parts[3])
            alts = parts[4].split(",")
            for i, (seq, allele) in enumerate(zip(seqs, variant)):
                if allele > 0:
                    seq.append(ref_seq[cur_positions[i]:variant_start])
                    seq.append(alts[allele-1])
                    cur_positions[i] = variant_end
        for pos, seq in zip(cur_positions, seqs):
            seq.append(ref_seq[pos:interval[1]-interval[0]])
        return ["".join(seq) for seq in seqs]

    def get_variants_from_intervals(self, intervals):
        """ intervals reveresely sorted on start """
        PRUNE = True
        StackElement = namedtuple("StackElement", ["end", "variant_list"])
        intervals = iter(intervals)
        current_intervals = deque([])
        cur_interval = next(intervals)
        is_finished = False
        line_parts = (line.split("\t", 9) for line in self.f
                      if not line.startswith("#"))
        for parts in line_parts:
            pos = int(parts[1])-1
            while pos >= cur_interval[0] and not is_finished:
                current_intervals.append(
                    StackElement(cur_interval[-1],
                                 VariantList(*cur_interval)))
                try:
                    cur_interval = next(intervals)
                except StopIteration:
                    is_finished = True
            if not current_intervals:
                continue

            while current_intervals and current_intervals[0].end < pos:
                tmp = current_intervals.popleft()
                yield tmp.variant_list.finalize()
            if not current_intervals:
                if is_finished:
                    break
                continue
            ref = parts[3].lower()
            alt = parts[4].lower().split(",")
            if PRUNE:
                ref, alt, offset = prune_seqs(ref, alt)
                pos = int(pos) + offset
            precence = VariantPrecence.from_line(parts[-1])
            var = FullVariant(int(pos), ref.lower(), alt, precence)
            variant_end = pos+len(ref)
            for stack_element in current_intervals:
                if stack_element.end > pos:
                    stack_element.variant_list.append(var)

    def get_variants_from(self, start, end):
        variants = []
        cur_end = 0
        intersecting = []
        for line in chain(self.cur_lines, self.f):
            if line.startswith("#"):
                continue
            parts = line.split("\t")
            pos = int(parts[1])-1
            if pos < start:
                continue
            if pos >= cur_end:
                if intersecting:
                    variants.append(self._join_variants(intersecting))
                intersecting = []
            if pos >= end:
                self.cur_lines = [line]
                break

            ref = parts[3]
            intersecting.append(Variant(int(pos-start), ref, parts[4].split(",")))
            cur_end = max(pos+len(ref), cur_end)
            continue
        if intersecting:
            variants.append(self._join_variants(intersecting))

        self.cur_lines = []
        variants.sort(key=lambda x: x.offset)
        return variants

def traverse_variants(alt_seq, ref_seq, variants):
    if not len(variants):
        return None
    Combination = namedtuple("Combination", ["code", "alt_offset", "prev_offset"])
    combinations = [Combination([], 0, 0)]
    for j, variant in enumerate(variants):
        new_combinations = []
        for comb in combinations:
            inter_alt = alt_seq[comb.prev_offset+comb.alt_offset:variant.offset+comb.alt_offset]
            inter_ref = ref_seq[comb.prev_offset:variant.offset]
            if not inter_alt == inter_ref:
                continue
            new_combinations.append(Combination(
                comb.code+[0],
                comb.alt_offset,
                comb.prev_offset))
            if comb.prev_offset > variant.offset:
                continue
            for code, seq in enumerate(variant.alt):
                p_ref, p_alt, p_offset = prune_seqs(variant.ref, [seq])
                seq = p_alt[0]
                start_on_alt = variant.offset+comb.alt_offset+p_offset
                var_alt = alt_seq[start_on_alt:start_on_alt+len(seq)]
                if var_alt == seq.lower():
                    new_combinations.append(Combination(
                        comb.code+[code+1],
                        comb.alt_offset+len(seq)-len(p_ref),
                        variant.offset+len(variant.ref)))

        combinations = new_combinations
        # assert all(len(t[0]) == j+1 for t in combinations), (j, combinations)
    real = []
    for comb in combinations:
        # if not len(alt_seq)-comb.alt_offset == len(ref_seq):
        #     continue
        alt_stub = alt_seq[comb.prev_offset+comb.alt_offset:len(ref_seq)+comb.alt_offset]
        ref_stub = ref_seq[comb.prev_offset:len(ref_seq)]
        l = min(len(alt_stub), len(ref_stub))
        if not alt_stub[:l] == ref_stub[:l]:
            continue
        real.append(comb)

    if not len(real):
        logging.info(alt_seq)
        logging.info(ref_seq)
        logging.info(variants)
        logging.info("---->")
        logging.info("%s / %s", real, combinations)
        print("")
        return ["No Match"]

    haplotypes = []
    for valid in real:
        codes = valid[0]
        haplotypes.extend(find_valid_haplotypes(variants, codes))
        # f = None
        # for code, variant in zip(codes, variants):
        #     f = variant.precence.get_samples(code, f)
        # haplotypes.extend(f)

    return np.sort(np.unique(haplotypes))
    # codes = real[0][0]
    # f = None
    # for code, variant in zip(codes, variants):
    #     f = variant.precence.get_samples(code, f)
    # return f

def find_haplotype(seq, refseq, vcf, start, end):
    refseq = refseq.lower()
    if seq == refseq:
        return ["REF"]
    # logging.info("Finding haplotype for %s->%s, (%s, %s)", refseq, seq, start, end)
    return traverse_variants(seq, refseq, vcf.get_variants_from(start, end))
    cur_offset = 0
    cur_ref_offset = 0
    haplotypes = []

    possible_vars = []
    for i, variant in enumerate(vcf.get_variants_from(start, end)):
        possible_vars.append(variant)
        new_offset = variant.offset
        l = new_offset - cur_ref_offset
        if seq[cur_offset:cur_offset+l] != refseq[cur_ref_offset:new_offset]:
            logging.error("%s != %s", seq[cur_offset:cur_offset+l], refseq[cur_ref_offset:new_offset])
            return ["ERROR"]
        cur_offset += l
        cur_ref_offset = new_offset
        if seq[cur_offset:cur_offset+len(variant.alt)] == variant.alt.lower():
            if (not seq[cur_offset:cur_offset+len(variant.ref)] == variant.ref.lower()) or len(variant.alt) > len(variant.ref):
                haplotypes.append(1)
                cur_offset += len(variant.alt)
                cur_ref_offset += len(variant.ref)
                continue
        # assert seq[cur_offset:cur_offset+len(variant.ref)] == variant.ref
        haplotypes.append(0)
    logging.info(possible_vars)
    logging.info([var for var, h in zip(possible_vars, haplotypes) if h])
    return haplotypes

class Main:

    def __init__(self, indexed_interval, graph, seq_graph, fasta, vcf):
        self.indexed_interval = indexed_interval
        self.graph = graph
        self.seq_graph = seq_graph
        self.fasta = fasta
        self.vcf = vcf
        self.counter = 0
        self.prev_end = -1

    def get_haplo_sequence_around_intervals(self, intervals, haplotype):
        extended = (self.extend_peak_to_reference(interval) for interval in intervals)
        ref_intervals = [
            (self.indexed_interval.get_offset_at_node(peak.region_paths[0]) + peak.start_position.offset,
             self.indexed_interval.get_offset_at_node(peak.region_paths[-1]) + peak.end_position.offset)
            for peak in extended]
        starts, ends = zip(*ref_intervals)
        interval = (min(starts), max(ends))
        ref_seq = self.fasta[int(interval[0]):int(interval[1])].seq
        var_seqs = self.vcf.get_haplotype_sequences(haplotype, interval, ref_seq)
        return var_seqs

    def run_peak_set(self, peaks_dict):
        all_name_tuples = [(key, self.get_sequence_pair(read))
                           for key, reads in peaks_dict.items() for read in reads]
        names, tuples = zip(*sorted(all_name_tuples, key=lambda x: x[1][2][0]))
        alts, refs, intervals = zip(*tuples)
        params = zip(alts, refs, self.vcf.get_variants_from_intervals(intervals))
        haplotypes = (traverse_variants(*param) for param in params)
        haplotypes = (h if h is not None else list(range(1136)) for h in haplotypes)
        result_dict = {name: [] for name in peaks_dict}
        for name, result in zip(names, haplotypes):

            result_dict[name].append(result)
        return result_dict

    def run_peaks(self, peaks):
        alts, refs, intervals = zip(*(self.get_sequence_pair(peak) for peak in peaks))
        haplotypes = [traverse_variants(alt.lower(), ref.lower(), variants)
                      for alt, ref, variants in
                      zip(alts, refs, self.vcf.get_variants_from_intervals(intervals))]
        return haplotypes

    def run_peak(self, peak):
        seq, ref_seq, interval = self.get_sequence_pair(peak)
        if interval[0] < self.prev_end:
            return []
        self.prev_end = interval[-1]
        haplo = find_haplotype(seq, ref_seq, self.vcf, interval[0], interval[1])
        if haplo:
            self.counter += 1
        return haplo

    def extend_peak_to_reference(self, peak):
        if peak.region_paths[0] < 0:
            peak.graph = self.graph
            peak = peak.get_reverse()
        start_node = peak.region_paths[0]
        end_node = peak.region_paths[-1]
        start_offset = peak.start_position.offset
        end_offset = peak.end_position.offset
        rps = peak.region_paths
        if start_node not in self.indexed_interval.nodes_in_interval():
            start_node = min(
                -node for node in self.graph.reverse_adj_list[-start_node]
                if -node in self.indexed_interval.nodes_in_interval())
            start_offset = self.graph.node_size(start_node)-1
            rps = [start_node] + rps
        if end_node not in self.indexed_interval.nodes_in_interval():
            end_node = max(node for node in self.graph.adj_list[end_node]
                           if node in self.indexed_interval.nodes_in_interval())
            end_offset = 1
            rps = rps + [end_node]
        return obg.Interval(start_offset, end_offset, rps)

    def get_sequence_pair(self, peak):
        peak.graph = self.graph
        peak = self.extend_peak_to_reference(peak)
        seq = self.seq_graph.get_interval_sequence(peak)
        assert len(seq) == peak.length(), (len(seq), peak.length())
        interval = (self.indexed_interval.get_offset_at_node(peak.region_paths[0]) + peak.start_position.offset,
                    self.indexed_interval.get_offset_at_node(peak.region_paths[-1]) + peak.end_position.offset)
        ref_seq = self.fasta[int(interval[0]):int(interval[1])].seq
        assert len(ref_seq) == interval[1]-interval[0]
        return seq.lower(), ref_seq.lower(), interval

    @classmethod
    def from_name(cls, name, fasta_file_name, chrom):
        indexed_interval = obg.NumpyIndexedInterval.from_file(
            name+"_linear_pathv2.interval")
        graph = obg.Graph.from_file(name+".nobg")
        seq_graph = obg.SequenceGraph.from_file(name+".nobg.sequences")
        fasta = Fasta(fasta_file_name)[chrom]
        vcf = VCF(name + "_variants.vcf")
        return cls(indexed_interval, graph, seq_graph, fasta, vcf)
    # var_list.create_lookup(name+"_variants.vcf")
    # return var_list

class DummyVCF:

    @staticmethod
    def get_variants_from(start):
        return [Variant(2, "A", "T"), Variant(4, "G", ""), Variant(5, "G", "C")]

def test():
    refseq = "AAAGGGTTT"
    seq = "AATGGCTTT"
    assert find_haplotype(seq, refseq, DummyVCF, 0) == [1, 0, 1]
    print("SUCCESS")

def test2():
    refseq = "AAAGGGTTT"
    seq = "AAAGGCTTT"
    assert find_haplotype(seq, refseq, DummyVCF, 0) == [0, 0, 1]
    print("SUCCESS")

def test3():
    refseq = "AAAGGGTTT"
    seq = "AAAGCTTT"
    assert find_haplotype(seq, refseq, DummyVCF, 0) == [0, 1, 1]
    print("SUCCESS")

def test4():
    vcf = "/home/knut/Documents/phd/data/graphs/1_variants.vcf"
    vcf = VCF(vcf)
    for i, var in enumerate(vcf.get_variants_from(70)):
        print(var)
        if i > 10:
            break

if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG)
    line = ["_", "10", "_",  "aaa", "ccc"] + ["_"]*4
    h = "./.:.:.	1|1:40:8	1|1:40:11	0|0:40:9	./.:.:.	./.:.:.	./.:.:.	0|0:40:8	0|0:40:9	./.:.:.	0|0:40:12	0|0:40:16	0|0:40:40	0|0:40:35	0|0:40:21	0|0:40:13	0|0:40:44	0|0:40:15	0|0:40:86	0|0:40:69	./.:.:.	0|0:40:14	0|0:40:59	0|0:40:51"
    line = "\t".join(line+ [h])
    print(line)
    ref = "ttaaaggcc"
    interval = (7, 7+len(ref))
    print(VCF("tmp.txt").get_haplotype_sequences(1, interval, ref, [line]))
#
#
#     alt = "ccttctttttttg"
#     ref = "ccttctttttttg"
#     variants = [FullVariant(offset=0, ref='ctt', alt=['ttt', 'c', 'ctttt'], precence=None),
#                 FullVariant(offset=2, ref='t', alt=['tc'], precence=None),
#                 FullVariant(offset=5, ref='t', alt=['c'], precence=None),
#                 FullVariant(offset=6, ref='t', alt=['a'], precence=None),
#                 FullVariant(offset=7, ref='t', alt=['c', 'tc'], precence=None)]
#     print(traverse_variants(alt, ref, variants))


import logging

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.ensemble import ExtraTreesClassifier

from .compat import string_, str_cast, unicode_
from .util import get_and_union_features
from .blocks import TagCountNoCSSReadabilityBlockifier

class Extractor(BaseEstimator, ClassifierMixin):
    """
    An sklearn-style classifier that extracts the main content (and/or comments)
    from an HTML document.

    Args:
        blockifier (``Blockifier``)
        features (str or List[str], ``Features`` or List[``Features``], or List[Tuple[str, ``Features``]]):
            One or more features to be used to transform blocks into a matrix of
            numeric values. If more than one, a :class:`FeatureUnion` is
            automatically constructed. See :func:`get_and_union_features`.
        model (:class:`ClassifierMixin`): A scikit-learn classifier that takes
             a numeric matrix of features and outputs a binary prediction of
            1 for content or 0 for not-content. If None, a :class:`ExtraTreesClassifier`
            with default parameters is used.
        to_extract (str or Sequence[str]): Type of information to extract from
            an HTML document: 'content', 'comments', or both via ['content', 'comments'].
        prob_threshold (float): Minimum prediction probability of a block being
            classified as "content" for it actually be taken as such.
        max_block_weight (int): Maximum weight that a single block may be given
            when training the extractor model, where weights are set equal to
            the number of tokens in each block.

    Note:
        If ``prob_threshold`` is not None, then ``model`` must implement the
            ``predict_proba()`` method.
    """

    def __init__(self, blockifier=TagCountNoCSSReadabilityBlockifier,
                 features=('kohlschuetter', 'weninger', 'readability'),
                 model=None,
                 to_extract='content', prob_threshold=0.5, max_block_weight=200):
        self.blockifier = blockifier
        self.features = features
        # initialize model
        if model is None:
            self.model = ExtraTreesClassifier()
        elif isinstance(model, ClassifierMixin):
            self.model = model
        else:
            raise TypeError('invalid `model` type: "{}"'.format(type(model)))
        if isinstance(to_extract, string_):
            self.to_extract = (to_extract,)
        else:
            self.to_extract = tuple(to_extract)
        self.prob_threshold = prob_threshold
        self.max_block_weight = max_block_weight
        self._positive_idx = None

    @property
    def features(self):
        return self._features

    @features.setter
    def features(self, feats):
        self._features = get_and_union_features(feats)

    def fit(self, documents, labels, weights=None):
        """
        Fit :class`Extractor` features and model to a training dataset.

        Args:
            blocks (List[Block])
            labels (``np.ndarray``)
            weights (``np.ndarray``)

        Returns:
            :class`Extractor`
        """
        block_groups = np.array([self.blockifier.blockify(doc) for doc in documents])
        mask = [self._has_enough_blocks(blocks) for blocks in block_groups]
        block_groups = block_groups[mask]
        labels = np.concatenate(np.array(labels)[mask])

        # TODO: This only 'fit's one doc at a time. No feature fitting actually
        # happens for now, but this might be important if the features change
        features_mat = np.concatenate([self.features.fit_transform(blocks)
                                       for blocks in block_groups])
        if weights is None:
            self.model.fit(features_mat, labels)
        else:
            weights = np.concatenate(np.array(weights)[mask])
            self.model.fit(features_mat, labels, sample_weight=weights)
        return self

    def get_html_labels_weights(self, data):
        """
        Gather the html, labels, and weights of many files' data.
        Primarily useful for training/testing an :class`Extractor`.

        Args:
            data: Output of :func:`dragnet.data_processing.prepare_all_data`.

        Returns:
            Tuple[List[Block], np.array(int), np.array(int)]: All blocks, all
                labels, and all weights, respectively.
        """
        all_html = []
        all_labels = []
        all_weights = []
        for html, content, comments in data:
            all_html.append(html)
            labels, weights = self._get_labels_and_weights(
                content, comments)
            all_labels.append(labels)
            all_weights.append(weights)
        return np.array(all_html), np.array(all_labels), np.array(all_weights)

    def _has_enough_blocks(self, blocks):
        if len(blocks) < 3:
            logging.warning(
                'extraction failed: too few blocks (%s)', len(blocks))
            return False
        return True

    def _get_labels_and_weights(self, content, comments):
        """
        Args:
            content (Tuple[np.array[int], np.array[int], List[str]])
            comments (Tuple[np.array[int], np.array[int], List[str]])

        Returns:
            Tuple[np.array[int], np.array[int], List[str]]
        """
        # extract content and comments
        if 'content' in self.to_extract and 'comments' in self.to_extract:
            labels = np.logical_or(content[0], comments[0]).astype(int)
            weights = content[1],
        # extract content only
        elif 'content' in self.to_extract:
            labels = content[0]
            weights = content[1]
        # extract comments only
        else:
            labels = comments[0]
            weights = comments[1]
        if self.max_block_weight is None:
            weights = np.minimum(weights, self.max_block_weight)

        return labels, weights

    def extract(self, html, encoding=None, as_blocks=False):
        """
        Extract the main content and/or comments from an HTML document and
        return it as a string or as a sequence of block objects.

        Args:
            html (str): HTML document as a string.
            encoding (str): Encoding of ``html``. If None (encoding unknown), the
                original encoding will be guessed from the HTML itself.
            as_blocks (bool): If False, return the main content as a combined
                string; if True, return the content-holding blocks as a list of
                block objects.

        Returns:
            str or List[Block]
        """
        preds, blocks = self.predict(html, encoding=encoding, return_blocks=True)
        if as_blocks is False:
            return str_cast(b'\n'.join(blocks[ind].text for ind in np.flatnonzero(preds)))
        else:
            return [blocks[ind] for ind in np.flatnonzero(preds)]

    def predict(self, documents, **kwargs):
        """
        Predict class (content=1 or not-content=0) of the blocks in one or many
        HTML document(s).

        Args:
            documents (str or List[str]): HTML document(s)

        Returns:
            ``np.ndarray`` or List[``np.ndarray``]: array of binary predictions
                for content (1) or not-content (0).
        """
        if isinstance(documents, (str, bytes, unicode_, np.unicode_)):
            return self._predict_one(documents, **kwargs)
        else:
            return np.concatenate([self._predict_one(doc, **kwargs) for doc in documents])

    def _predict_one(self, document, encoding=None, return_blocks=False):
        """
        Predict class (content=1 or not-content=0) of each block in an HTML
        document.

        Args:
            documents (str): HTML document

        Returns:
            ``np.ndarray``: array of binary predictions for content (1) or
            not-content (0).
        """
        # blockify
        blocks = self.blockifier.blockify(document, encoding=encoding)
        # get features
        try:
            features = self.features.transform(blocks)
        except ValueError: # Can't make features, predict no content
            preds = np.zeros((len(blocks)))
        # make predictions
        else:
            if self.prob_threshold is None:
                preds = self.model.predict(features)
            else:
                self._positive_idx = (
                    self._positive_idx or list(self.model.classes_).index(1))
                preds = self.model.predict_proba(features) > self.prob_threshold
                preds = preds[:, self._positive_idx].astype(int)

        if return_blocks:
            return preds, blocks
        else:
            return preds


import numpy as np
import cv2
import random
import math
from scipy import misc
import pandas as pd

class CloudEnvironment():

    def __init__(self,img_size=84,radius=[8,16],speed=1,sequence_stride=1,channels=3,sequence_length=2,ramp_step=0.1,action_type=0,action_nr=2,stochastic_irradiance=False,save_images=False):

        self.img_size=img_size
        self.radius = radius
        self.speed = speed
        self.ramp_step=ramp_step
        self.action_type = action_type

        self.sequence_stride=sequence_stride
        self.sequence_length=sequence_length
        self.channels=channels
        self.stochastic_irradiance= stochastic_irradiance
        self.save_images = save_images

        self.observation_space=self.ObservationSpace((img_size*img_size*sequence_length*channels+sequence_length+1,1))

        #self.observation_space

        self.action_space = self.ActionSpace(action_type,action_nr,ramp_step)

        self.episode_samples = []
        self.episode_ci = []
        self.episode_pointer = 0
        self.current_episode = None
        self.episode_nr = 0

    @property
    def current_train_control_inputs(self):
        return self.episode_ci

    @property
    def current_train_episode(self):
        return self.current_episode

    @property
    def episode_n(self):
        return self.episode_nr

    @property
    def episode_id(self):
        return self.episode_nr

    def reset(self):

        self.episode_samples = []
        self.episode_ci = []
        self.episode_pointer = 0

        episode_data_list = self.create_episode() #[(img,irr)]

        self.episode_samples = [episode_data_list[i:(i + (self.sequence_length * self.sequence_stride)):self.sequence_stride]
                                      for i in
                                      range(len(episode_data_list) - (self.sequence_length - 1) * self.sequence_stride)]

        current_episode_step = self.episode_samples[self.episode_pointer]

        images = [img[0] for img in current_episode_step]

        image_tensor = np.uint8(np.concatenate(images,axis=2)) #image tensor with first 3 channels first image, last threee channels newest image
        current_irr = np.array([img[1] for img in current_episode_step])
        current_ci = current_irr[-1]
        self.episode_ci.append(current_ci)

        #env_obs = [np.uint8(image_tensor.ravel()), current_irr, current_ci]
        env_obs = np.concatenate([image_tensor.ravel(), current_irr, np.reshape(current_irr[-1],(1))], axis=0).astype(np.float16)[:, None]

        return env_obs

    def step(self,action):

        curr_pointer = self.episode_pointer
        self.episode_pointer+=1
        next_pointer = self.episode_pointer

        next_episode_step = self.episode_samples[next_pointer]
        curr_episode_step = self.episode_samples[curr_pointer]

        current_ci = self.episode_ci[-1]

        next_irr = np.array([img[1] for img in next_episode_step])
        curr_irr = np.array([img[1] for img in curr_episode_step])

        next_ci,reward = self.action_space.calculate_step(action,next_irr[-1],curr_irr[-1],current_ci.squeeze())

        self.episode_ci.append(next_ci)

        next_images = [img[0] for img in next_episode_step]

        for i in  range(len(next_images)):
            cv2.imshow('Rotation', next_images[i])
            cv2.waitKey(100)
            print(next_irr[-1],self.episode_ci[-1])

        next_image_tensor = np.uint8(np.concatenate(next_images, axis=2))

        if next_pointer == len(self.episode_samples) - 1:
            done = True
        else:
            done = False

        next_env_obs =  np.concatenate([np.uint8(next_image_tensor.ravel()),next_irr, np.reshape(next_ci,(1))]).astype(np.float16)[:, None]
        """
        next_env_obs = [np.uint8(next_image_tensor.ravel()), next_irr, np.reshape(next_ci, (1))]
        cv2.imshow('next_state_image_32', np.float32(np.reshape(next_env_obs[0:-3], (84, 84, 6))[:, :, 0:3]))
        cv2.waitKey(10)
        """

        return next_env_obs, reward , done, 0

    def create_episode(self):

        self.episode_nr += 1
        episode_id = self.episode_n

        episode_data_list = []
        y1 = 0
        x1 = 31#random.randint(0,self.img_size-1)

        y2 = self.img_size-1
        x2 = 61#random.randint(0,self.img_size-1)
        xdiff = x2-x1
        ydiff = y2-y1

        y_update = 1.0 *self.speed
        x_update = (xdiff/ydiff)*self.speed

        center_pixel = self.img_size//2-1

        x_f=float(x1)
        y_f=float(y1)
        x1 = int(x_f)
        y1 = int(y_f)

        #rotation
        rot = 180#np.random.choice([0,90,180,270])
        rotation_matrix = cv2.getRotationMatrix2D((self.img_size/ 2, self.img_size / 2), rot, 1)

        #radius
        radius = 12#np.random.randint(self.radius[0],self.radius[1])

        assert (y1 < self.img_size and x1 >= 0 and x1 < self.img_size)

        while y1 < self.img_size and x1 >= 0 and x1 < self.img_size:
            episode_log = np.array([x1, y1, x2, y2, rot, radius])
            img = np.zeros((self.img_size, self.img_size, self.channels), np.uint8)
            channel_colors = tuple([255]*self.channels)
            cv2.circle(img, (center_pixel, center_pixel),1, tuple([100]*self.channels), -1)
            cv2.circle(img, (x1, y1), radius,channel_colors , -1)
            img=np.uint8(cv2.warpAffine(img, rotation_matrix, (self.img_size, self.img_size)))
            img = np.reshape(img,(self.img_size,self.img_size,self.channels))

            #cv2.imshow('Rotation', img)
            #cv2.waitKey(10)

            if np.abs(x1-center_pixel)<radius and np.abs(y1-center_pixel) <radius:
                irr = 0.0
            else:
                if self.stochastic_irradiance:
                    irr = np.random.uniform(1.0-self.ramp_step,1.0) #still within the ramp to test a follow_irr 0 action
                else:
                    irr = 1.0

            episode_data_list.append((img,irr,episode_id,episode_log))

            x_f += x_update
            y_f += y_update
            x1 = int(x_f)
            y1 = int(y_f)

        assert len(episode_data_list)>0

        col = ['irr', 'id','img_log']
        self.current_episode = pd.DataFrame(data=np.array(episode_data_list)[:,1:], columns=col)

        return episode_data_list

    class ActionSpace():
        def __init__(self, type=0, action_nr=2,ramp_step=0.1):
            self.n = action_nr
            self.type = type
            self.ramp_step = ramp_step

        def calculate_step(self, action, next_irr, curr_irr, current_ci):
            # return next ci given action

            if self.type == 0:

                if action == 0: #upper target
                    next_ci = np.minimum(1.0,float(current_ci)+self.ramp_step)
                else: #lower target
                    next_ci = np.maximum(0.0,current_ci-self.ramp_step)

            elif self.type==-1:
                #naive policy
                diff = next_irr-current_ci
                ramp = np.abs(diff) > self.ramp_step
                if ramp:
                    next_ci = current_ci+np.sign(diff)*self.ramp_step
                else:
                    next_ci = next_irr

            elif self.type==-2:
                #naive policy
                diff = curr_irr-current_ci
                ramp = np.abs(diff) > self.ramp_step
                if ramp:
                    next_ci = current_ci+np.sign(diff)*self.ramp_step
                else:
                    next_ci = curr_irr

            elif self.type==1:

                if action == 0:
                    diff = curr_irr - current_ci
                    ramp = np.abs(diff) > self.ramp_step
                    if ramp:
                        next_ci = current_ci + np.sign(diff) * self.ramp_step
                    else:
                        next_ci = curr_irr

                elif action == 1: #upper target
                    next_ci = np.minimum(1.0,float(current_ci)+self.ramp_step)
                else: #lower target
                    next_ci = np.maximum(0.0,current_ci-self.ramp_step)

            else:

                raise ValueError('Illegal Action Set')

            reward = -np.abs(next_ci - next_irr).squeeze()

            return next_ci,reward

    class ObservationSpace():
        def __init__(self, shape):
            self.shape = shape

"""
env = CloudEnvironment(84,[9,17],2)

d=0
s = env.reset()
while d==0:
    a = random.randint(0,1)

    sn,r,d,_ = env.step(a)

    #print(sn[1][0])
"""

    #Check reward sizes if still ok
    #Testing facility, print predictions ...
    #TODO try adding follow  action in toy environment check what it learns
    #TODO add real evironment and use originam action space to train ..
    #normalize irr in environment in real environment  so  between 0 and 1 this also normalizes reward .. so we have ame situation as in simple environment
    #more varied environment  changing size of cloud random, rotate images randomly


# This script is used to generate test data for OpenCV deep learning module.
import numpy as np
import tensorflow as tf
import os
import argparse
import struct
import cv2 as cv

from tensorflow.python.tools import optimize_for_inference_lib
from tensorflow.tools.graph_transforms import TransformGraph

np.random.seed(2701)

def gen_data(placeholder):
    shape = placeholder.shape.as_list()
    shape[0] = shape[0] if shape[0] else 1  # batch size = 1 instead None
    return np.random.standard_normal(shape).astype(placeholder.dtype.as_numpy_dtype())

def prepare_for_dnn(sess, graph_def, in_node, out_node, out_graph, dtype, optimize=True, quantize=False):
    # Freeze graph. Replaces variables to constants.
    graph_def = tf.graph_util.convert_variables_to_constants(sess, graph_def, [out_node])
    if optimize:
        # Optimize graph. Removes training-only ops, unused nodes.
        graph_def = optimize_for_inference_lib.optimize_for_inference(graph_def, [in_node], [out_node], dtype.as_datatype_enum)
        # Fuse constant operations.
        transforms = ["fold_constants(ignore_errors=True)"]
        if quantize:
            transforms += ["quantize_weights(minimum_size=0)"]
        transforms += ["sort_by_execution_order"]
        graph_def = TransformGraph(graph_def, [in_node], [out_node], transforms)
    # Serialize
    with tf.gfile.FastGFile(out_graph, 'wb') as f:
        f.write(graph_def.SerializeToString())

tf.reset_default_graph()
tf.Graph().as_default()
tf.set_random_seed(324)
sess = tf.Session()

# Use this variable to switch behavior of layers.
isTraining = tf.placeholder(tf.bool, name='isTraining')

def writeBlob(data, name):
    if data.ndim == 4:
        # NHWC->NCHW
        np.save(name + '.npy', data.transpose(0, 3, 1, 2).astype(np.float32))
    elif data.ndim == 5:
        # NDHWC->NCDHW
        np.save(name + '.npy', data.transpose(0, 4, 1, 2, 3).astype(np.float32))
    else:
        # Save raw data.
        np.save(name + '.npy', data.astype(np.float32))

def runModel(inpName, outName, name):
    with tf.Session(graph=tf.Graph()) as localSession:
        localSession.graph.as_default()

        with tf.gfile.FastGFile(name + '_net.pb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')

        inputData = gen_data(inp)
        outputData = localSession.run(localSession.graph.get_tensor_by_name(outName),
                                      feed_dict={localSession.graph.get_tensor_by_name(inp.name): inputData})
        writeBlob(inputData, name + '_in')
        writeBlob(outputData, name + '_out')

def save(inp, out, name, quantize=False, optimize=True):
    sess.run(tf.global_variables_initializer())

    inputData = gen_data(inp)
    outputData = sess.run(out, feed_dict={inp: inputData, isTraining: False})
    writeBlob(inputData, name + '_in')
    writeBlob(outputData, name + '_out')

    prepare_for_dnn(sess, sess.graph.as_graph_def(), inp.name[:inp.name.rfind(':')],
                    out.name[:out.name.rfind(':')], name + '_net.pb', inp.dtype,
                    optimize, quantize)

    # By default, float16 weights are stored in repeated tensor's field called
    # `half_val`. It has type int32 with leading zeros for unused bytes.
    # This type is encoded by Varint that means only 7 bits are used for value
    # representation but the last one is indicated the end of encoding. This way
    # float16 might takes 1 or 2 or 3 bytes depends on value. To impove compression,
    # we replace all `half_val` values to `tensor_content` using only 2 bytes for everyone.
    with tf.gfile.FastGFile(name + '_net.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        for node in graph_def.node:
            if 'value' in node.attr:
                halfs = node.attr["value"].tensor.half_val
                if not node.attr["value"].tensor.tensor_content and halfs:
                    node.attr["value"].tensor.tensor_content = struct.pack('H' * len(halfs), *halfs)
                    node.attr["value"].tensor.ClearField('half_val')
        tf.train.write_graph(graph_def, "", name + '_net.pb', as_text=False)

# Test cases ###################################################################
# shape: NHWC
for dtype, prefix in zip([tf.float32, tf.float16], ['', 'fp16_']):
    inp = tf.placeholder(dtype, [1, 6, 5, 3], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=3, kernel_size=[1, 1],
                            activation=tf.nn.relu,
                            bias_initializer=tf.random_normal_initializer())
    save(inp, conv, prefix + 'single_conv')
################################################################################
    inp = tf.placeholder(dtype, [3, 7, 5, 4], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=5, kernel_size=[5, 3], padding='SAME',
                            use_bias=False)
    activation_abs = tf.abs(conv)
    save(inp, activation_abs, prefix + 'padding_same')
################################################################################
    inp = tf.placeholder(dtype, [2, 4, 6, 5], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=4, kernel_size=[3, 5], padding='VALID',
                            activation=tf.nn.elu, bias_initializer=tf.random_normal_initializer())
    save(inp, conv, prefix + 'padding_valid')
################################################################################
    inp = tf.placeholder(dtype, [3, 2, 3, 4], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=4, kernel_size=[1, 1], activation=tf.nn.tanh,
                            bias_initializer=tf.random_uniform_initializer(0, 1))
    conv2 = tf.layers.conv2d(inputs=inp, filters=4, kernel_size=[1, 1], activation=tf.nn.sigmoid,
                             bias_initializer=None)
    eltwise_add_mul = (inp * 0.31 + 2 * conv) * conv2
    save(inp, eltwise_add_mul, prefix + 'eltwise_add_mul')
################################################################################
    inp = tf.placeholder(dtype, [1, 4, 5, 1], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=4, kernel_size=[3, 1], padding='VALID')
    padded = tf.pad(conv, [[0, 0], [0, 2], [0, 0], [0, 0]])
    merged = tf.concat([padded, inp], axis=3)
    save(inp, merged, prefix + 'pad_and_concat')
###############################################################################
    inp = tf.placeholder(dtype, [1, 6, 6, 2], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=3, kernel_size=[3, 3], padding='SAME')
    pool = tf.layers.max_pooling2d(inputs=conv, pool_size=2, strides=2)
    save(inp, pool, prefix + 'max_pool_even')
################################################################################
    inp = tf.placeholder(dtype, [1, 7, 7, 2], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=3, kernel_size=[3, 3], padding='SAME')
    pool = tf.layers.max_pooling2d(inputs=conv, pool_size=3, strides=2, padding='VALID')
    save(inp, pool, prefix + 'max_pool_odd_valid')
################################################################################
    inp = tf.placeholder(dtype, [1, 7, 7, 2], 'input')
    conv = tf.layers.conv2d(inputs=inp, filters=3, kernel_size=[3, 3], padding='SAME')
    relu = tf.nn.relu6(conv * 10)
    pool = tf.layers.max_pooling2d(inputs=relu, pool_size=2, strides=2, padding='SAME')
    save(inp, pool, prefix + 'max_pool_odd_same')
################################################################################
    inp = tf.placeholder(dtype, [1, 5, 6, 2], 'input')
    deconv_weights = tf.Variable(tf.random_normal([5, 3, 4, 2], dtype=dtype), name='deconv_weights')
    deconv = tf.nn.conv2d_transpose(value=inp, filter=deconv_weights,
                                    output_shape=[1, 9, 8, 4], strides=[1, 1, 1, 1],
                                    padding='VALID')
    deconv_bias = tf.contrib.layers.bias_add(deconv, activation_fn=tf.nn.relu,
                                             initializer=tf.random_normal_initializer())
    save(inp, deconv_bias, prefix + 'deconvolution')
################################################################################
inp = tf.placeholder(tf.float32, [2, 5, 4, 3], 'input')
bn = tf.contrib.layers.batch_norm(inputs=inp, fused=True, is_training=False,
                                  scale=True, param_initializers={
                                    'beta': tf.random_normal_initializer(),
                                    'gamma': tf.random_normal_initializer(),
                                    'moving_mean': tf.random_uniform_initializer(-2, -1),
                                    'moving_variance': tf.random_uniform_initializer(1, 2)
                                  })
save(inp, bn, 'fused_batch_norm')
################################################################################
inp = tf.placeholder(tf.float32, [2, 5, 6, 3], 'input')
weights = tf.Variable(tf.random_normal([3, 3, 3, 4]), name='weights')
conv = tf.nn.atrous_conv2d(inp, weights, rate=2, padding='VALID')
relu = tf.nn.relu(conv)
save(inp, relu, 'atrous_conv2d_valid')
################################################################################
inp = tf.placeholder(tf.float32, [2, 5, 10, 3], 'input')
weights = tf.Variable(tf.random_normal([3, 5, 3, 4]), name='weights')
conv = tf.nn.atrous_conv2d(inp, weights, rate=2, padding='SAME')
relu = tf.nn.relu(conv)
save(inp, relu, 'atrous_conv2d_same')
################################################################################
inp = tf.placeholder(tf.float32, [2, 5, 4, 3], 'input')
bn = tf.contrib.layers.batch_norm(inputs=inp, is_training=False,
                                  scale=True, param_initializers={
                                    'beta': tf.random_normal_initializer(),
                                    'gamma': tf.random_normal_initializer(),
                                    'moving_mean': tf.random_uniform_initializer(-2, -1),
                                    'moving_variance': tf.random_uniform_initializer(1, 2)
                                  })
save(inp, bn, 'batch_norm')
################################################################################
inp = tf.placeholder(tf.float32, [2, 10, 9, 6], 'input')
weights = tf.Variable(tf.random_normal([5, 3, 6, 4]), name='weights')
conv = tf.nn.depthwise_conv2d(input=inp, filter=weights, strides=[1, 1, 1, 1],
                              padding='SAME')
save(inp, conv, 'depthwise_conv2d')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3], 'input')
biases = tf.Variable(tf.random_normal([4]), name='matmul_biases')
weights = tf.Variable(tf.random_normal([3, 4]), name='matmul_weights')
mm = tf.matmul(inp, weights) + biases
save(inp, mm, 'matmul')
################################################################################
from tensorflow.python.framework import function

@function.Defun(tf.float32, func_name='Dropout')
def my_dropout(x):
    return tf.layers.dropout(x, rate=0.1, training=isTraining)

inp = tf.placeholder(tf.float32, [1, 10, 10, 3], 'input')
conv = tf.layers.conv2d(inp, filters=3, kernel_size=[1, 1])
dropout = my_dropout(conv)
relu = tf.nn.relu(dropout)

save(inp, relu, 'defun_dropout')
################################################################################
# Use not 4 dimensions to save the raw data (see writeBlob function)
inp = tf.placeholder(tf.float32, [2, 3, 4], 'input')
shift = tf.Variable(tf.random_normal([2, 3, 4]), name='shift')
shifted = tf.add(inp, shift, name='shifted')
reshape = tf.reshape(shifted, [4, 3, 2], 'reshaped')
save(inp, reshape, 'shift_reshape_no_reorder')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3], 'input')
reshape = tf.reshape(inp, [3, 1, 2], 'reshaped')
save(inp, reshape, 'reshape_no_reorder')
################################################################################
inp = tf.placeholder(tf.float32, [2, 10, 10, 3], 'input')
pad = tf.pad(inp, [[0, 0], [3, 3], [3, 3], [0, 0]])
conv = tf.layers.conv2d(inp, filters=4, kernel_size=[5, 5], strides=(2, 2),
                        bias_initializer=tf.random_normal_initializer())
save(inp, conv, 'spatial_padding')
################################################################################
inp = tf.placeholder(tf.float32, [1, 10, 10, 3], 'input')
pad = tf.pad(inp, [[0, 0], [3, 3], [3, 3], [0, 0]], mode='REFLECT')
save(inp, pad, 'mirror_pad')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3], 'input')
bn = tf.add(inp, tf.Variable(tf.random_normal(inp.shape)))
reshape = tf.reshape(bn, [-1, 3])
save(inp, reshape, 'reshape_reduce')
################################################################################
times = 4  # Sequence length (number of batches in different time stamps)
batch_size = 2
input_size = 5*6*3  # W*H*C
output_size = 10
# Define LSTM blobk.
inp = tf.placeholder(tf.float32, [times, batch_size, input_size], 'input')
lstm_cell = tf.contrib.rnn.LSTMBlockFusedCell(output_size, forget_bias=0.9,
                                           cell_clip=0.4, use_peephole=True)
outputs, state = lstm_cell(inp, dtype=tf.float32)
# shape(outputs) is a (times, batch_size, output_size)

# Slice the last time iteration:
last_output = tf.slice(outputs, [times - 1, 0, 0], [1, -1, output_size])
# shape(last_output) is a (1, batch_size, output_size)

# Remove time axis.
last_output = tf.reshape(last_output, [-1, 10])
# shape(last_output) is a (batch_size, output_size)

# Fully-connected
weights = tf.Variable(tf.random_normal([10, 2]))
biases = tf.Variable(tf.random_normal([2]))
sigmoid = tf.nn.sigmoid(tf.matmul(last_output, weights) + biases)
save(inp, sigmoid, 'lstm')
################################################################################
bgr = tf.placeholder(tf.float32, [4, 5, 6, 3], 'input')
b, g, r = tf.split(bgr, num_or_size_splits=3, axis=3)
rgb = tf.concat([r, g, b], axis=3)
alpha, beta = tf.split(rgb, num_or_size_splits=2, axis=0)
res = tf.layers.conv2d(alpha, filters=1, kernel_size=[1, 1]) + \
      tf.layers.conv2d(beta, filters=1, kernel_size=[1, 1])
save(bgr, res, 'split_equals')
################################################################################
inp = tf.placeholder(tf.float32, [2, 10, 11, 3], 'input')
conv = tf.layers.conv2d(inp, filters=7, kernel_size=[1, 1])
scaled = tf.image.resize_nearest_neighbor(conv, size=(15, 8))
scaled = tf.image.resize_nearest_neighbor(scaled, size=(9, 12))
save(inp, scaled, 'resize_nearest_neighbor')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
bn = tf.layers.batch_normalization(inp, training=isTraining, fused=False, name='batch_norm',
                                   beta_initializer=tf.random_normal_initializer(),
                                   gamma_initializer=tf.random_normal_initializer(),
                                   moving_mean_initializer=tf.random_uniform_initializer(-2, 1),
                                   moving_variance_initializer=tf.random_uniform_initializer(0.1, 2),)
save(inp, bn, 'batch_norm_text')
################################################################################
inp = tf.placeholder(tf.float32, [2, 4, 5], 'input')
flatten = tf.contrib.layers.flatten(inp)
save(inp, flatten, 'flatten')
################################################################################
inp = tf.placeholder(tf.float32, [2, 2, 2, 2, 3], 'input')
conv = tf.layers.conv3d(inp, filters=2, kernel_size=[1, 1, 1], padding='SAME')
concat = tf.concat([inp, conv], axis=-1)
save(inp, concat, 'concat_3d')
################################################################################
# Generate test data for MobileNet-SSD object detection model from TensorFlow
# model zoo, http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_11_06_2017.tar.gz
# 1. Download and extract an archive
# 2. Place frozen_inference_graph.pb as a ssd_mobilenet_v1_coco.pb nearby this script
with tf.gfile.FastGFile('../ssd_mobilenet_v1_coco.pb') as f:
    # Load the model
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Session(graph=tf.Graph()) as localSession:
    # Restore session
    localSession.graph.as_default()
    tf.import_graph_def(graph_def, name='')

    # Receive output
    inp = cv.imread('../street.png')
    inp = cv.resize(inp, (300, 300))
    inp = inp[:, :, [2, 1, 0]]  # BGR2RGB
    out = localSession.run([localSession.graph.get_tensor_by_name('concat:0'),    # All detections
                            localSession.graph.get_tensor_by_name('concat_1:0'),  # Classification
                            localSession.graph.get_tensor_by_name('num_detections:0'),     # Postprocessed output
                            localSession.graph.get_tensor_by_name('detection_scores:0'),   #
                            localSession.graph.get_tensor_by_name('detection_boxes:0'),    #
                            localSession.graph.get_tensor_by_name('detection_classes:0')], #
                   feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
    np.save('ssd_mobilenet_v1_coco.concat.npy', out[0])
    np.save('ssd_mobilenet_v1_coco.concat_1.npy', out[1])
    # Pack detections in format [id, class_id, confidence, left, top, right, bottom]
    num_detections = int(out[2][0])
    detections = np.zeros([1, 1, num_detections, 7], np.float32)
    detections[0, 0, :, 0] = 0  # bounding boxes ids
    detections[0, 0, :, 1] = out[5][0]
    detections[0, 0, :, 2] = out[3][0]
    detections[0, 0, :, 3:] = out[4][0][:, [1, 0, 3, 2]]
    # Detections are sorted in descending by confidence order. Group them by classes
    # to make OpenCV test more simple.
    detections = sorted(detections[0, 0, :, :], cmp=lambda x, y: -1 if x[1] < y[1] and x[2] < y[2] else 0)
    np.save('ssd_mobilenet_v1_coco.detection_out.npy', detections)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
conv = tf.layers.conv2d(inp, filters=5, kernel_size=[1, 1],
                        activation=tf.nn.relu,
                        bias_initializer=tf.random_normal_initializer())
flattened = tf.reshape(conv, [1, -1], 'reshaped')
biases = tf.Variable(tf.random_normal([10]), name='matmul_biases')
weights = tf.Variable(tf.random_normal([2*3*5, 10]), name='matmul_weights')
mm = tf.matmul(flattened, weights) + biases
save(inp, mm, 'nhwc_reshape_matmul')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
conv = tf.layers.conv2d(inp, filters=5, kernel_size=[1, 1],
                        activation=tf.nn.relu,
                        bias_initializer=tf.random_normal_initializer())
transposed = tf.transpose(conv, [0, 1, 2, 3])
flattened = tf.reshape(transposed, [1, -1], 'reshaped')
biases = tf.Variable(tf.random_normal([10]), name='matmul_biases')
weights = tf.Variable(tf.random_normal([2*3*5, 10]), name='matmul_weights')
mm = tf.matmul(flattened, weights) + biases
save(inp, flattened, 'nhwc_transpose_reshape_matmul')
################################################################################
inp = tf.placeholder(tf.float32, [1, 6, 5, 3], 'input')
conv = tf.layers.conv2d(inputs=inp, filters=3, kernel_size=[1, 1],
                        activation=tf.nn.relu,
                        bias_initializer=tf.random_normal_initializer())
save(inp, conv, 'uint8_single_conv', quantize=True)
runModel(inp, conv.name, 'uint8_single_conv')
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 4, 1], 'input')
conv = tf.layers.conv2d(inp, filters=3, kernel_size=[3, 3], padding='SAME')
pool = tf.layers.average_pooling2d(conv, pool_size=3, strides=1, padding='SAME')
save(inp, pool, 'ave_pool_same')
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 6, 1], 'input')
conv = tf.layers.conv2d(inp, filters=3, kernel_size=[1, 1], padding='SAME')
sliced = tf.slice(conv, [0, 1, 2, 0], [-1, 3, 4, 1])
save(inp, sliced, 'slice_4d')
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 4, 1], 'input')
#                                             ky kx out in
deconv_weights = tf.Variable(tf.random_normal([3, 3, 2, 1], dtype=tf.float32))
deconv = tf.nn.conv2d_transpose(inp, deconv_weights,
                                output_shape=[1, 4, 4, 2], strides=[1, 1, 1, 1],
                                padding='SAME')
leakyRelu = tf.nn.leaky_relu(deconv, alpha=0.2)
save(inp, leakyRelu, 'deconvolution_same')
# ################################################################################
inp = tf.placeholder(tf.float32, [1, 3, 3, 1], 'input')
deconv_weights = tf.Variable(tf.random_normal([3, 3, 2, 1], dtype=tf.float32))
deconv = tf.nn.conv2d_transpose(inp, deconv_weights,
                                output_shape=[1, 5, 5, 2], strides=[1, 2, 2, 1],
                                padding='SAME')
save(inp, deconv, 'deconvolution_stride_2_same')
################################################################################
inp = tf.placeholder(tf.float32, [1, 3, 2, 1], 'input')
deconv_weights = tf.Variable(tf.random_normal([3, 3, 2, 1], dtype=tf.float32))
deconv = tf.nn.conv2d_transpose(inp, deconv_weights,
                                output_shape=[1, 8, 6, 2], strides=[1, 2, 2, 1],
                                padding='VALID')
save(inp, deconv, 'deconvolution_adj_pad_valid')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 2, 1], 'input')
deconv_weights = tf.Variable(np.ones([3, 3, 1, 1]), dtype=tf.float32)
deconv = tf.nn.conv2d_transpose(inp, deconv_weights,
                                output_shape=[1, 4, 4, 1], strides=[1, 2, 2, 1],
                                padding='SAME')
save(inp, deconv, 'deconvolution_adj_pad_same')
################################################################################
inp = tf.placeholder(tf.float32, [1, 3, 4, 5], 'input')
gamma = tf.Variable(tf.random_normal([5], dtype=tf.float32))
beta = tf.Variable(tf.random_normal([5], dtype=tf.float32))
bn = tf.nn.fused_batch_norm(inp, gamma, beta, epsilon=1e-5, is_training=True)[0]
save(inp, bn, 'mvn_batch_norm')
################################################################################
inp = tf.placeholder(tf.float32, [1, 1, 1, 5], 'input')
gamma = tf.Variable(tf.random_normal([5], dtype=tf.float32))
beta = tf.Variable(tf.random_normal([5], dtype=tf.float32))
bn = tf.nn.fused_batch_norm(inp, gamma, beta, epsilon=1e-5, is_training=True)[0]
save(inp, bn, 'mvn_batch_norm_1x1')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
bn = tf.layers.batch_normalization(inp, training=False, fused=False, name='unfused_batch_norm',
                                   beta_initializer=tf.random_normal_initializer(),
                                   gamma_initializer=tf.random_normal_initializer(),
                                   moving_mean_initializer=tf.random_uniform_initializer(-2, 1),
                                   moving_variance_initializer=tf.random_uniform_initializer(0.1, 2),)
save(inp, bn, 'unfused_batch_norm', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
bn = tf.layers.batch_normalization(inp, training=False, fused=True, name='fused_batch_norm_no_gamma',
                                   beta_initializer=tf.random_normal_initializer(),
                                   scale=False,
                                   moving_mean_initializer=tf.random_uniform_initializer(-2, 1),
                                   moving_variance_initializer=tf.random_uniform_initializer(0.1, 2),)
save(inp, bn, 'fused_batch_norm_no_gamma', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
bn = tf.layers.batch_normalization(inp, training=False, fused=False, name='unfused_batch_norm_no_gamma',
                                   beta_initializer=tf.random_normal_initializer(),
                                   scale=False,
                                   moving_mean_initializer=tf.random_uniform_initializer(-2, 1),
                                   moving_variance_initializer=tf.random_uniform_initializer(0.1, 2),)
save(inp, bn, 'unfused_batch_norm_no_gamma', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3], 'input')
flatten = tf.contrib.layers.flatten(inp)
save(inp, flatten, 'unfused_flatten', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [None, 2, 3], 'input')
flatten = tf.contrib.layers.flatten(inp)
save(inp, flatten, 'unfused_flatten_unknown_batch', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
relu = tf.maximum(inp, 0.01 * inp, name='leaky_relu') * 2
save(inp, relu, 'leaky_relu_order1', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
relu = tf.maximum(inp, inp * 0.01, name='leaky_relu') * 2
save(inp, relu, 'leaky_relu_order2', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
relu = tf.maximum(0.01 * inp, inp, name='leaky_relu') * 2
save(inp, relu, 'leaky_relu_order3', optimize=False)
################################################################################
from tensorflow import keras as K
model = K.models.Sequential()
model.add(K.layers.Softmax(name='keras_softmax', input_shape=(2, 3, 4)))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_softmax_input:0'),
     sess.graph.get_tensor_by_name('keras_softmax/truediv:0'), 'keras_softmax', optimize=False)
################################################################################
model = K.models.Sequential()
model.add(K.layers.Conv2D(filters=4, kernel_size=1, data_format='channels_last',
                          name='keras_mobilenet_head_conv', input_shape=(2, 3, 4)))
model.add(K.layers.GlobalAveragePooling2D(name='keras_mobilenet_head_pool'))
model.add(K.layers.Reshape((1, 1, 4), name='keras_mobilenet_head_reshape'))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_mobilenet_head_conv_input:0'),
     sess.graph.get_tensor_by_name('keras_mobilenet_head_reshape/Reshape:0'),
     'keras_mobilenet_head', optimize=False)
################################################################################
def keras_relu6(x):
    return K.activations.relu(x, max_value=6)

inp = K.Input(shape=(2, 3, 4), name='keras_relu6_input')
relu = K.layers.Activation(keras_relu6, name='keras_relu6')(inp)
model = K.Model(inp, relu)
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_relu6_input:0'),
     sess.graph.get_tensor_by_name('keras_relu6/clip_by_value:0'), 'keras_relu6', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [2, 3], 'input')
max_node = tf.clip_by_value(inp, clip_value_min=0, clip_value_max=1)
save(inp, max_node, 'clip_by_value')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
reduced = tf.reduce_mean(inp, axis=[1, 2], keepdims=True)
save(inp, reduced, 'reduce_mean')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
reduced = tf.reduce_sum(inp, axis=[1, 2], keepdims=True)
save(inp, reduced, 'reduce_sum')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
reduced = tf.reduce_sum(inp, axis=[2], keepdims=False)
save(inp, reduced, 'sum_pool_by_axis')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
pool = tf.layers.average_pooling2d(inp, pool_size=1, strides=1, padding='SAME')
l2norm = tf.nn.l2_normalize(pool, axis=-1)
l2norm = tf.nn.l2_normalize(l2norm, axis=[2, 3, 1])
save(inp, l2norm, 'l2_normalize')
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4], 'input')
l2norm = tf.nn.l2_normalize(inp, axis=1)
l2norm = tf.nn.l2_normalize(l2norm, axis=-1)
l2norm = tf.nn.l2_normalize(l2norm, axis=[0, 1])
save(inp, l2norm, 'l2_normalize_3d')
################################################################################
model = K.models.Sequential()
model.add(K.layers.Conv2DTranspose(filters=4, kernel_size=3, strides=(2, 2),
                                   data_format='channels_last', name='keras_deconv_valid',
                                   input_shape=(4, 5, 2)))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_deconv_valid_input:0'),
     sess.graph.get_tensor_by_name('keras_deconv_valid/BiasAdd:0'),
     'keras_deconv_valid', optimize=True)
################################################################################
model = K.models.Sequential()
model.add(K.layers.Conv2DTranspose(filters=4, kernel_size=3, strides=(2, 2),
                                   data_format='channels_last', name='keras_deconv_same',
                                   input_shape=(4, 5, 2), padding='same'))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_deconv_same_input:0'),
     sess.graph.get_tensor_by_name('keras_deconv_same/BiasAdd:0'),
     'keras_deconv_same', optimize=True)
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
resized = tf.image.resize_bilinear(inp, size=[9, 8], name='resize_bilinear')
save(inp, resized, 'resize_bilinear')
################################################################################
inp = tf.placeholder(tf.float32, [None, 3, 4, 5], 'input')
resized = tf.image.resize_bilinear(inp, size=[tf.shape(inp)[1]*2, tf.shape(inp)[2]*3],
                                   name='resize_bilinear_factor')
sub_add = resized - 0.3 + 0.3
save(inp, sub_add, 'resize_bilinear_factor', optimize=False)
################################################################################
model = K.models.Sequential()
model.add(K.layers.SeparableConv2D(filters=4, kernel_size=3, strides=(1, 1),
                                   dilation_rate=(2, 3), name='keras_atrous_conv2d_same',
                                   input_shape=(11, 12, 2), padding='same'))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_atrous_conv2d_same_input:0'),
     sess.graph.get_tensor_by_name('keras_atrous_conv2d_same/BiasAdd:0'),
     'keras_atrous_conv2d_same', optimize=True)
################################################################################
# Generate test data for Faster-RCNN object detection model from TensorFlow
# model zoo, http://download.tensorflow.org/models/object_detection/faster_rcnn_inception_v2_coco_2018_01_28.tar.gz
for name in ['faster_rcnn_inception_v2_coco_2018_01_28', 'faster_rcnn_resnet50_coco_2018_01_28']:
    with tf.gfile.FastGFile(os.path.join('..', name + '.pb')) as f:
        # Load the model
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    with tf.Session(graph=tf.Graph()) as localSession:
        # Restore session
        localSession.graph.as_default()
        tf.import_graph_def(graph_def, name='')

        # Receive output
        inp = cv.imread('../dog416.png')
        inp = cv.resize(inp, (800, 600))
        inp = inp[:, :, [2, 1, 0]]  # BGR2RGB
        out = localSession.run([localSession.graph.get_tensor_by_name('num_detections:0'),   #
                                localSession.graph.get_tensor_by_name('detection_scores:0'),   #
                                localSession.graph.get_tensor_by_name('detection_boxes:0'),    #
                                localSession.graph.get_tensor_by_name('detection_classes:0')], #
                       feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
        # Pack detections in format [id, class_id, confidence, left, top, right, bottom]
        num_detections = int(out[0][0])
        detections = np.zeros([1, 1, num_detections, 7], np.float32)
        detections[0, 0, :, 1] = out[3][0, :num_detections] - 1
        detections[0, 0, :, 2] = out[1][0, :num_detections]
        detections[0, 0, :, 3:] = out[2][:, :num_detections, [1, 0, 3, 2]]
        np.save(name + '.detection_out.npy', detections)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
conv1 = tf.layers.conv2d(inp, filters=4, kernel_size=[1, 1])
conv2 = tf.layers.conv2d(inp, filters=4, kernel_size=[1, 1])
flatten1 = tf.contrib.layers.flatten(conv1)
flatten2 = tf.contrib.layers.flatten(conv2)
concat = tf.concat([flatten1, flatten2], axis=1)
bias = tf.contrib.layers.bias_add(concat)  # Add zeros (it has NHWC data format flag)
save(inp, bias, 'concat_axis_1')
################################################################################
inp = tf.placeholder(tf.float32, [1, 3, 5, 8], 'input')  # NCHW input
conv = tf.layers.conv2d(inp, filters=4, kernel_size=[2, 3], data_format='channels_first')
pool = tf.layers.max_pooling2d(conv, pool_size=2, strides=2, data_format='channels_first')
save(inp, pool, 'conv_pool_nchw')
# Input and output have been transposed (see writeBlob)
for name in ['conv_pool_nchw_in.npy', 'conv_pool_nchw_out.npy']:
    np.save(name, np.load(name).transpose(0, 2, 3, 1))
################################################################################
model = K.models.Sequential()

model.add(K.layers.UpSampling2D(size=(3, 2), data_format='channels_last',
                          name='keras_upsampling2d', input_shape=(2, 3, 4)))
sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_upsampling2d_input:0'),
     sess.graph.get_tensor_by_name('keras_upsampling2d/ResizeNearestNeighbor:0'),
     'keras_upsampling2d')
################################################################################
# Generate test data for MobileNet-SSD object detection model from TensorFlow
# model zoo, http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_ppn_shared_box_predictor_300x300_coco14_sync_2018_07_03.tar.gz
with tf.gfile.FastGFile('../ssd_mobilenet_v1_ppn_coco.pb') as f:
    # Load the model
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Session(graph=tf.Graph()) as localSession:
    # Restore session
    localSession.graph.as_default()
    tf.import_graph_def(graph_def, name='')

    # Receive output
    img = cv.imread('../dog416.png')
    inp = cv.resize(img, (300, 300))
    inp = inp[:, :, [2, 1, 0]]  # BGR2RGB
    out = localSession.run([localSession.graph.get_tensor_by_name('num_detections:0'),   #
                            localSession.graph.get_tensor_by_name('detection_scores:0'),   #
                            localSession.graph.get_tensor_by_name('detection_boxes:0'),    #
                            localSession.graph.get_tensor_by_name('detection_classes:0')], #
                   feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
    # Pack detections in format [id, class_id, confidence, left, top, right, bottom]
    num_detections = int(out[0][0])
    detections = np.zeros([1, 1, num_detections, 7], np.float32)
    detections[0, 0, :, 1] = out[3][0, :num_detections]
    detections[0, 0, :, 2] = out[1][0, :num_detections]
    detections[0, 0, :, 3:] = out[2][:, :num_detections, [1, 0, 3, 2]]
    np.save('ssd_mobilenet_v1_ppn_coco.detection_out.npy', detections)
################################################################################
inp = tf.placeholder(tf.float32, [None, 2, 3], 'input')
flatten = tf.reshape(inp, [-1, 2*3], 'planar')
reshaped = tf.reshape(flatten, tf.shape(inp), 'reshape')
save(inp, reshaped, 'reshape_as_shape', optimize=False)
################################################################################
with tf.gfile.FastGFile('../mask_rcnn_inception_v2_coco_2018_01_28.pb') as f:
    # Load the model
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Session(graph=tf.Graph()) as localSession:
    # Restore session
    localSession.graph.as_default()
    tf.import_graph_def(graph_def, name='')

    # Receive output
    img = cv.imread('../street.png')
    inp = cv.resize(img, (800, 800))
    inp = inp[:, :, [2, 1, 0]]  # BGR2RGB
    out = localSession.run([localSession.graph.get_tensor_by_name('num_detections:0'),
                            localSession.graph.get_tensor_by_name('detection_scores:0'),
                            localSession.graph.get_tensor_by_name('detection_boxes:0'),
                            localSession.graph.get_tensor_by_name('detection_classes:0'),
                            localSession.graph.get_tensor_by_name('detection_masks:0')],
                           feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
    # Pack detections in format [id, class_id, confidence, left, top, right, bottom]
    num_detections = int(out[0][0])
    detections = np.zeros([1, 1, num_detections, 7], np.float32)
    detections[0, 0, :, 1] = out[3][0, :num_detections] - 1
    detections[0, 0, :, 2] = out[1][0, :num_detections]
    detections[0, 0, :, 3:] = out[2][:, :num_detections, [1, 0, 3, 2]]
    np.save('mask_rcnn_inception_v2_coco_2018_01_28.detection_out.npy', detections)
    np.save('mask_rcnn_inception_v2_coco_2018_01_28.detection_masks.npy', out[4])
################################################################################
inp = K.Input(shape=(2, 3, 4), name='keras_pad_concat_input', batch_size=1)
conv = K.layers.Conv2D(filters=4, kernel_size=1, data_format='channels_last',
                       name='keras_pad_concat_conv', input_shape=(2, 3, 4))(inp)

def pad_depth(x, desired_channels):
    y = K.backend.random_uniform_variable(x.shape.as_list()[:-1] + [desired_channels], low=0, high=1)
    return K.layers.concatenate([x, y])

pad = K.layers.Lambda(pad_depth, arguments={'desired_channels': 5}, name='keras_pad_concat')(conv)

sess = K.backend.get_session()
sess.as_default()
save(sess.graph.get_tensor_by_name('keras_pad_concat_input:0'),
     sess.graph.get_tensor_by_name('keras_pad_concat/concatenate/concat:0'),
     'keras_pad_concat', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
conv = tf.layers.conv2d(inp, filters=5, kernel_size=[1, 1],
                        bias_initializer=tf.random_normal_initializer())
sub = conv - inp
save(inp, sub, 'eltwise_sub')
################################################################################
inp = tf.placeholder(tf.float32, [None, 2, 3, 4], 'input')
conv = tf.layers.conv2d(inp, filters=3, kernel_size=[1, 1])
softmax = tf.contrib.slim.softmax(conv)
save(inp, softmax, 'slim_softmax')
################################################################################
# issue https://github.com/opencv/opencv/issues/14224
inp_node = 'img_inputs'
out_node = 'MobileFaceNet/MobileFaceNet/Conv2d_0/add'
with tf.Session(graph=tf.Graph()) as localSession:
    localSession.graph.as_default()

    with tf.gfile.FastGFile('frozen_model.pb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        for node in graph_def.node:
            if node.name == inp_node:
                del node.attr['shape']

    tf.import_graph_def(graph_def, name='')

    inputData = gen_data(tf.placeholder(tf.float32, [1, 4, 5, 3], inp_node))
    outputData = localSession.run(localSession.graph.get_tensor_by_name(out_node + ':0'),
                                  feed_dict={inp_node + ':0': inputData})
    writeBlob(inputData, 'slim_batch_norm_in')
    writeBlob(outputData, 'slim_batch_norm_out')

    graph_def = TransformGraph(graph_def, [inp_node], [out_node], ['fold_constants', 'strip_unused_nodes'])
    with tf.gfile.FastGFile('slim_batch_norm_net.pb', 'wb') as f:
        f.write(graph_def.SerializeToString())

################################################################################
# issue https://github.com/opencv/opencv/issues/13839
inp_node = 'PNet/conv3/add'
out_node = 'PNet/cls_prob'
with tf.Session(graph=tf.Graph()) as localSession:
    localSession.graph.as_default()

    with tf.gfile.FastGFile('PNet_pnet.pb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph_def = TransformGraph(graph_def, [inp_node], [out_node], ['strip_unused_nodes'])

    tf.import_graph_def(graph_def, name='')

    inputData = gen_data(tf.placeholder(tf.float32, [1, 4, 5, 16], inp_node))
    outputData = localSession.run(localSession.graph.get_tensor_by_name(out_node + ':0'),
                                  feed_dict={inp_node + ':0': inputData})
    writeBlob(inputData, 'slim_softmax_v2_in')
    writeBlob(outputData, 'slim_softmax_v2_out')

    with tf.gfile.FastGFile('slim_softmax_v2_net.pb', 'wb') as f:
        f.write(graph_def.SerializeToString())
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 6, 5, 3], 'input') # NDHWC format
conv3d = tf.layers.conv3d(inputs=inp, filters=2, kernel_size=[2, 3, 4], use_bias=True, padding='same')
save(inp, conv3d, 'conv3d')
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 6, 5, 3], 'input') # NDHWC format
maxpool3d = tf.layers.max_pooling3d(inputs=inp, pool_size=(3, 2, 3), strides=(1, 2, 1), padding='same')
save(inp, maxpool3d, 'max_pool3d')
################################################################################
inp = tf.placeholder(tf.float32, [1, 5, 4, 5, 2], 'input') # NDHWC format
avepool3d = tf.layers.average_pooling3d(inputs=inp, pool_size=(3, 3, 2), strides=(2, 1, 1), padding='valid')
save(inp, avepool3d, 'ave_pool3d')
################################################################################
# issue https://github.com/opencv/opencv/issues/13494
inp_node = 'input_image'
out_node = 'SUBPIXEL/SUBPIXEL/subpixel_image/Identity'
with tf.Session(graph=tf.Graph()) as localSession:
    localSession.graph.as_default()

    with tf.gfile.FastGFile('simple_subpixel.optimized.pb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    tf.import_graph_def(graph_def, name='')

    inputData = gen_data(tf.placeholder(tf.float32, [1, 1, 1, 4], inp_node))
    outputData = localSession.run(localSession.graph.get_tensor_by_name(out_node + ':0'),
                                  feed_dict={inp_node + ':0': inputData})
    writeBlob(inputData, 'subpixel_in')
    writeBlob(outputData, 'subpixel_out')

    for node in graph_def.node:
        if node.op == 'Placeholder':
            node.attr["data_format"].s = "NHWC"

    with tf.gfile.FastGFile('subpixel_net.pb', 'wb') as f:
        f.write(graph_def.SerializeToString())
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
conv = tf.layers.conv2d(inp, filters=4, kernel_size=[1, 1])
strided_slice = conv[:, 1:, :2, 2:3]
save(inp, strided_slice, 'strided_slice')
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 5, 3], 'input')
conv_out = tf.keras.layers.Conv2D(filters=3, kernel_size=[3, 3], padding='same')(inp)
crop2d = tf.keras.layers.Cropping2D(((1, 1), (1, 1)))(conv_out)
save(inp, crop2d, 'crop2d')
################################################################################

inp = tf.placeholder(tf.float32, [1, 2, 3, 4, 2], 'input')
bn = tf.layers.batch_normalization(inp, training=False, fused=False, name='batch_norm3d',
                                   beta_initializer=tf.random_normal_initializer(),
                                   gamma_initializer=tf.random_normal_initializer(),
                                   moving_mean_initializer=tf.random_uniform_initializer(-2, 1),
                                   moving_variance_initializer=tf.random_uniform_initializer(0.1, 2),)
save(inp, bn, 'batch_norm3d', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 4, 6, 64], 'activation_8/Elu')
runModel(inp, 'batch_normalization_1/cond/FusedBatchNorm:0', 'switch_identity')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 64], 'Relu_8')
runModel(inp, 'conv2d_transpose_1:0', 'keras_deconv_same_v2')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 4, 3], 'ContentImage')
runModel(inp, 'Relu:0', 'keras_batch_norm_training')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 2, 4], 'Split')
features1 = tf.split(inp, num_or_size_splits=2, axis=3)[0]
features2 = tf.split(inp, num_or_size_splits=2, axis=3)[1]
merged = tf.concat([features1, features2], axis=3)
save(inp, merged, 'split')
################################################################################
from tensorflow.python.ops.nn_grad import _MaxPoolGrad as MaxUnPooling2D

inp = tf.placeholder(tf.float32, [1, 7, 7, 3], 'input')
pool = tf.layers.max_pooling2d(inp, pool_size=(2, 2), strides=(2, 2))
conv = tf.layers.conv2d(inputs=pool, filters=3, kernel_size=[1, 1], padding='VALID')
unpool = MaxUnPooling2D(pool.op, conv)
save(inp, unpool, 'max_pool_grad')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
conv = tf.layers.conv2d(inp, filters=5, kernel_size=[1, 1])
flatten = tf.contrib.layers.flatten(conv)
weights = tf.Variable(tf.random_normal([2*3*5, 4]), name='matmul_weights')
mm = tf.matmul(flatten, weights)
reshape = tf.reshape(mm, [-1, 1, 1, 4], 'reshaped')  # NHWC
save(inp, reshape, 'matmul_layout')
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 2, 4], 'ReduceMean')
out = tf.reduce_mean([inp, inp * 2], axis=0)
save(inp, out, 'global_pool_by_axis')
################################################################################
# issue https://github.com/opencv/opencv/issues/15141
inp_node = 'mobilenetv2_1.00_96_input'
out_node = 'mobilenetv2_1.00_96/Conv1_relu/Relu6'
with tf.Session(graph=tf.Graph()) as localSession:
    localSession.graph.as_default()

    with tf.gfile.FastGFile('normal_and_abnormal_mnet_v2_96_96_Flatten.pb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph_def = optimize_for_inference_lib.optimize_for_inference(graph_def, [inp_node], [out_node], tf.float32.as_datatype_enum)

    tf.import_graph_def(graph_def, name='')

    inputData = gen_data(tf.placeholder(tf.float32, [1, 4, 5, 3], inp_node))
    outputData = localSession.run(localSession.graph.get_tensor_by_name(out_node + ':0'),
                                  feed_dict={inp_node + ':0': inputData})
    writeBlob(inputData, 'keras_learning_phase_in')
    writeBlob(outputData, 'keras_learning_phase_out')

    with tf.gfile.FastGFile('keras_learning_phase_net.pb', 'wb') as f:
        f.write(graph_def.SerializeToString())
################################################################################
inp = tf.placeholder(tf.float32, [None, 2, 3, 4], 'input')
pool = tf.layers.average_pooling2d(inp, pool_size=(2, 3), strides=1, padding='VALID')
out = inp * pool
save(inp, out, 'channel_broadcast', optimize=False)
################################################################################
inp = tf.placeholder(tf.float32, [1, 2, 3, 4], 'input')
resized = tf.image.resize_bilinear(inp, size=[3, 5], name='resize_bilinear', align_corners=True)
conv = tf.layers.conv2d(resized, filters=4, kernel_size=[1, 1])
save(inp, conv, 'fused_resize_conv')
################################################################################
# Uncomment to save model with dynamic shapes
# inp = tf.placeholder(tf.float32, [1, None, None, 2], 'input')
inp = tf.placeholder(tf.float32, [1, 9, 6, 2], 'input')
conv = tf.layers.conv2d(inp, filters=2, kernel_size=[1, 1])
shape_input = tf.shape(inp)
hi = shape_input[1] / 3
wi = shape_input[2] / 2
input_down = tf.image.resize(conv, size=[hi,wi], method=0, name='resize_down')
save(inp, input_down, 'resize_bilinear_down')
################################################################################

# Uncomment to print the final graph.
# with tf.gfile.FastGFile('fused_batch_norm_net.pb', 'rb') as f:
#     graph_def = tf.GraphDef()
#     graph_def.ParseFromString(f.read())
#     print(graph_def)


import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import pickle
import pandas as pd
import numpy as np

class ProGAT(nn.Module):

    def __init__(self, embedding_size, radius, T, p_dropout):
        super(ProGAT, self).__init__()

        '''def get_cm_dict(pickle_dir):

            cm_dict = {}

            for f in os.listdir(pickle_dir):
                f = pickle_dir + f
                data = pickle.load(open(f, 'rb'))
                for index, row in data.iterrows():
                    seq = row['seqs'][:self.max_len]
                    cm = row['contact_map'][0][:self.max_len, :self.max_len]
                    cm_dict[seq] = cm
            return cm_dict

        self.enc_lib = np.eye(input_size)

        self.max_len = 10'''

        # pickle_dir = '/home/jian/protein/dataset/3d_pdbs/'

        # self.contact_map_dict = get_cm_dict(pickle_dir)
        input_size = 26

        self.emb_layer = nn.Linear(input_size, embedding_size)
        self.neighbor_fc = nn.Linear(input_size, embedding_size)
        self.GRUCell = nn.ModuleList([nn.GRUCell(embedding_size, embedding_size) for r in range(radius)])
        self.align = nn.ModuleList([nn.Linear(2 * embedding_size, 1) for r in range(radius)])
        self.attend = nn.ModuleList([nn.Linear(embedding_size, embedding_size) for r in range(radius)])

        self.seq_GRUCell = nn.GRUCell(embedding_size, embedding_size)
        self.seq_align = nn.Linear(2 * embedding_size, 1)
        self.seq_attend = nn.Linear(embedding_size, embedding_size)

        self.dropout = nn.Dropout(p=p_dropout)

        self.radius = radius
        self.T = T

    '''def get_init(self, seq_list):

        mat = []
        for seq in seq_list:
            seq = list(map(lambda ch : ord(ch) - ord('A'), seq[:self.max_len]))

            enc = self.enc_lib[seq]
            if enc.shape[0] < self.max_len:
                enc = np.pad(enc,((0, self.max_len - enc.shape[0]), (0, 0)), 'constant')
            #print(enc.shape)

            mat.append(enc)
        mat = np.stack(mat, 0)

        return torch.tensor(mat, dtype=torch.float32).cuda()

    def get_degree_list(self, seq_list):

        mat = []
        for seq in seq_list:
            seq = seq[:self.max_len]
            if seq in self.contact_map_dict:
                cm = self.contact_map_dict[seq]
            else:
                cm = np.random.rand(self.max_len, self.max_len)

            degree_list = []
            for i in range(len(seq)):
                tmp = np.array(np.where(cm[i] > 0.5)[0])
                tmp = np.pad(tmp, (0, self.max_len - tmp.shape[0]), 'constant', constant_values=(-1,-1))
                degree_list.append(tmp)

            degree_list = np.stack(degree_list, 0)
            degree_list = np.pad(degree_list, ((0, self.max_len - degree_list.shape[0]), (0, 0)), 'constant', constant_values=(-1, -1))
            mat.append(degree_list)
        mat = np.stack(mat, 0)

        return torch.tensor(mat)

    def get_amino_mask(self, seq_list):
        mat = []
        for seq in seq_list:

            mask = np.ones(min(len(seq), self.max_len), dtype=np.int)
            mask = np.pad(mask, (0, self.max_len - len(mask)), 'constant')
            mat.append(mask)
        mat = np.stack(mat, 0)
        #print('mask', mat)
        return torch.tensor(mat, dtype=torch.float32).cuda()'''

    def forward(self, amino_list, amino_degree_list, amino_mask):
        '''
        amino_list is (batch, max_seqlen(512), amino_nums(26)). One mat(512, 26) is a one-hot representation of seq based on atom feature numbers.
        amino_degree_list is (batch, max_seqlen(512), max_neighbors(21?)). One mat(512, 21) is the index of the amino's neighbors, -1 is padding.
        amino_mask is (batch, 512), the existed is 1 while padding is 0.
        '''
        # amino_list = self.get_init(seq_list)
        # amino_degree_list = self.get_degree_list(seq_list)
        # amino_mask = self.get_amino_mask(seq_list)

        amino_mask = amino_mask.unsqueeze(2)

        batch_size, seq_len, num_amino_feature = amino_list.shape

        # print(amino_list.shape)
        # (512, 26) -> (512, 64)  embed amino num_amino_feature
        amino_feature = F.leaky_relu(self.emb_layer(amino_list))

        # (batch, max_len, max_neighbors, amino_nums)
        neighbor = [amino_list[i][amino_degree_list[i]] for i in range(batch_size)]
        neighbor = torch.stack(neighbor, dim=0)

        # (batch, 512, 21, 26) -> (batch, 512, 21, 64)
        neighbor_feature = F.leaky_relu(self.neighbor_fc(neighbor))

        # generate mask to eliminate the influence of blank atoms
        attend_mask = amino_degree_list.clone()
        attend_mask[attend_mask != -1] = 1
        attend_mask[attend_mask == -1] = 0
        attend_mask = attend_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        softmax_mask = amino_degree_list.clone()
        softmax_mask[softmax_mask != -1] = 0
        softmax_mask[softmax_mask == -1] = -9e8  # make the softmax value extremly small
        softmax_mask = softmax_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        batch_size, seq_length, max_neighbor_num, embedding_dim = neighbor_feature.shape
        amino_feature_expand = amino_feature.unsqueeze(-2).expand(batch_size, seq_length, max_neighbor_num,
                                                                  embedding_dim)
        feature_align = torch.cat([amino_feature_expand, neighbor_feature], dim=-1)

        # align[0] relate to radius， align_score -> (batch, max_len, max_neighbor, 1)
        align_score = F.leaky_relu(self.align[0](self.dropout(feature_align)))

        align_score = align_score + softmax_mask
        attention_weight = F.softmax(align_score, -2)
        #             print(attention_weight)
        attention_weight = attention_weight * attend_mask
        #         print(attention_weight)
        neighbor_feature_transform = self.attend[0](self.dropout(neighbor_feature))
        #             print(features_neighbor_transform.shape)
        # 邻节点每维特征*attention权重(21, 64)，然后所有ci相加得到当前节点c(1,64)
        context = torch.sum(torch.mul(attention_weight, neighbor_feature_transform), -2)
        context = F.elu(context)
        # attention_weight -> torch.Size([batch, 512, max_neighbor, 1])
        # context -> torch.Size([batch, 512, 64])

        context_reshape = context.view(batch_size * seq_length, embedding_dim)
        amino_feature_reshape = amino_feature.view(batch_size * seq_length, embedding_dim)
        amino_feature_reshape = self.GRUCell[0](context_reshape, amino_feature_reshape)
        amino_feature = amino_feature_reshape.view(batch_size, seq_length, embedding_dim)

        # do nonlinearity
        activated_features = F.relu(amino_feature)

        for d in range(self.radius - 1):
            # bonds_indexed = [bond_list[i][torch.cuda.LongTensor(bond_degree_list)[i]] for i in range(batch_size)]
            neighbor_feature = [activated_features[i][amino_degree_list[i]] for i in range(batch_size)]

            # neighbor_feature is a list of 3D tensor, so we need to stack them into a 4D tensor first
            neighbor_feature = torch.stack(neighbor_feature, dim=0)
            amino_feature_expand = activated_features.unsqueeze(-2).expand(batch_size, seq_length, max_neighbor_num,
                                                                           embedding_dim)

            feature_align = torch.cat([amino_feature_expand, neighbor_feature], dim=-1)

            align_score = F.leaky_relu(self.align[d + 1](self.dropout(feature_align)))
            #             print(attention_weight)
            align_score = align_score + softmax_mask
            attention_weight = F.softmax(align_score, -2)
            #             print(attention_weight)
            attention_weight = attention_weight * attend_mask
            #             print(attention_weight)
            neighbor_feature_transform = self.attend[d + 1](self.dropout(neighbor_feature))
            #             print(features_neighbor_transform.shape)
            context = torch.sum(torch.mul(attention_weight, neighbor_feature_transform), -2)
            #             print(context.shape)
            context = F.elu(context)
            context_reshape = context.view(batch_size * seq_length, embedding_dim)
            #             atom_feature_reshape = atom_feature.view(batch_size*mol_length, fingerprint_dim)
            amino_feature_reshape = self.GRUCell[d + 1](context_reshape, amino_feature_reshape)
            amino_feature = amino_feature_reshape.view(batch_size, seq_length, embedding_dim)

            # do nonlinearity
            activated_features = F.relu(amino_feature)
            # print('act', activated_features.shape)

        # 将aa特征汇聚成蛋白序列的
        seq_feature = torch.sum(activated_features * amino_mask, dim=-2)  # B * feature, all atom sum = mol

        # do nonlinearity
        activated_features_mol = F.relu(seq_feature)

        seq_softmax_mask = amino_mask.clone()
        seq_softmax_mask[seq_softmax_mask == 0] = -9e8
        seq_softmax_mask[seq_softmax_mask == 1] = 0
        seq_softmax_mask = seq_softmax_mask.type(torch.cuda.FloatTensor)

        for t in range(self.T):
            seq_prediction_expand = activated_features_mol.unsqueeze(-2).expand(batch_size, seq_length, embedding_dim)
            seq_align = torch.cat([seq_prediction_expand, activated_features], dim=-1)
            seq_align_score = F.leaky_relu(self.seq_align(seq_align))
            seq_align_score = seq_align_score + seq_softmax_mask
            seq_attention_weight = F.softmax(seq_align_score, -2)
            seq_attention_weight = seq_attention_weight * amino_mask
            #             print(mol_attention_weight.shape,mol_attention_weight)
            activated_features_transform = self.seq_attend(self.dropout(activated_features))
            #             aggregate embeddings of atoms in a molecule
            seq_context = torch.sum(torch.mul(seq_attention_weight, activated_features_transform), -2)
            #             print(mol_context.shape,mol_context)
            seq_context = F.elu(seq_context)
            seq_feature = self.seq_GRUCell(seq_context, seq_feature)
            #             print(mol_feature.shape,mol_feature)

            # do nonlinearity
            activated_features_seq = F.relu(seq_feature)

            # print(activated_features.shape)

        return activated_features_seq

        # neighbor = [amino_input[i] for i in range(batch_size)]