# Copyright 2016-2021 The Van Valen Lab at the California Institute of
# Technology (Caltech), with support from the Paul Allen Family Foundation,
# Google, & National Institutes of Health (NIH) under Grant U24CA224309-01.
# All rights reserved.
#
# Licensed under a modified Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.github.com/vanvalenlab/deepcell-tf/LICENSE
#
# The Work provided may be used for non-commercial academic purposes only.
# For any other use of the Work, including commercial use, please contact:
# vanvalenlab@gmail.com
#
# Neither the name of Caltech nor the names of its contributors may be used
# to endorse or promote products derived from this software without specific
# prior written permission.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Custom metrics for pixel-based and object-based classification accuracy.

The schema for this analysis was adopted from the description of object-based
statistics in Caicedo et al. (2018) Evaluation of Deep Learning Strategies for
Nucleus Segmentation in Fluorescence Images. BioRxiv 335216.

The SEG metric was adapted from Maska et al. (2014). A benchmark for comparison
of cell tracking algorithms. Bioinformatics 30, 1609-1617.

The linear classification schema used to match objects in truth and prediction
frames was adapted from Jaqaman et al. (2008). Robust single-particle tracking
in live-cell time-lapse sequences. Nature Methods 5, 695-702.
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import datetime
import json
import logging
import operator
import os
import warnings

import numpy as np
import pandas as pd
import networkx as nx

from scipy.optimize import linear_sum_assignment
from scipy.stats import hmean
from skimage.measure import regionprops
from skimage.segmentation import relabel_sequential
from sklearn.metrics import confusion_matrix
from tqdm import tqdm

from deepcell_toolbox import erode_edges
from deepcell_toolbox.compute_overlap import compute_overlap  # pylint: disable=E0401
from deepcell_toolbox.compute_overlap import compute_overlap_3D

def _cast_to_tuple(x):
    try:
        tup_x = tuple(x)
    except TypeError:
        tup_x = () if x is None else (x,)
    return tup_x

class Detection(object):  # pylint: disable=useless-object-inheritance
    """Object to hold relevant information about a given detection."""

    def __init__(self, true_index=None, pred_index=None):
        # cast the indices as tuples if possible to make them immutable
        try:
            self.true_index = tuple(true_index)
        except TypeError:
            self.true_index = true_index
        try:
            self.pred_index = tuple(pred_index)
        except TypeError:
            self.pred_index = pred_index

    def __eq__(self, other):
        """Custom comparator. Detections with the same indices are the same."""
        try:
            is_true_same = self.true_index == other.true_index
            is_pred_same = self.pred_index == other.pred_index
            return is_true_same and is_pred_same
        except AttributeError:
            return False

    def __hash__(self):
        """Custom hasher, allow Detections to be hashable."""
        return tuple((self.true_index, self.pred_index)).__hash__()

    def __repr__(self):
        return 'Detection({}, {})'.format(self.true_index, self.pred_index)

    @property
    def is_correct(self):
        is_linked = self.true_index is not None and self.pred_index is not None
        return is_linked and not self.is_split and not self.is_merge

    @property
    def is_gained(self):
        return self.true_index is None and self.pred_index is not None

    @property
    def is_missed(self):
        return self.true_index is not None and self.pred_index is None

    @property
    def is_split(self):
        if self.is_gained or self.is_missed:
            return False

        try:
            is_many_pred = len(self.pred_index) > 1
        except TypeError:
            is_many_pred = False

        try:
            is_single_true = len(tuple(self.true_index)) == 1
        except TypeError:
            is_single_true = isinstance(self.true_index, int)

        return is_single_true and is_many_pred

    @property
    def is_merge(self):
        if self.is_gained or self.is_missed:
            return False

        try:
            is_many_true = len(self.true_index) > 1
        except TypeError:
            is_many_true = False

        try:
            is_single_pred = len(tuple(self.pred_index)) == 1
        except TypeError:
            is_single_pred = isinstance(self.pred_index, int)

        return is_single_pred and is_many_true

    @property
    def is_catastrophe(self):
        if self.is_gained or self.is_missed:
            return False

        try:
            is_many_true = len(self.true_index) > 1
        except TypeError:
            is_many_true = False

        try:
            is_many_pred = len(self.pred_index) > 1
        except TypeError:
            is_many_pred = False

        return is_many_true and is_many_pred

class BaseMetrics(object):  # pylint: disable=useless-object-inheritance

    """Base class for Metrics classes."""

    def __init__(self, y_true, y_pred):
        if y_pred.shape != y_true.shape:
            raise ValueError('Input shapes must match. Shape of prediction '
                             'is: {}.  Shape of y_true is: {}'.format(
                                 y_pred.shape, y_true.shape))

        if not np.issubdtype(y_true.dtype, np.integer):
            warnings.warn('Casting y_true from {} to int'.format(y_true.dtype))
            y_true = y_true.astype('int32')

        if not np.issubdtype(y_pred.dtype, np.integer):
            warnings.warn('Casting y_pred from {} to int'.format(y_pred.dtype))
            y_pred = y_pred.astype('int32')

        self.y_true = y_true
        self.y_pred = y_pred

class PixelMetrics(BaseMetrics):
    """Calculates pixel-based statistics.
    (Dice, Jaccard, Precision, Recall, F-measure)

    Takes in raw prediction and truth data in order to calculate accuracy
    metrics for pixel based classfication. Statistics were chosen according
    to the guidelines presented in Caicedo et al. (2018) Evaluation of Deep
    Learning Strategies for Nucleus Segmentation in Fluorescence Images.
    BioRxiv 335216.

    Args:
        y_true (numpy.array): Binary ground truth annotations for a single
            feature, (batch,x,y)
        y_pred (numpy.array): Binary predictions for a single feature,
            (batch,x,y)

    Raises:
        ValueError: Shapes of y_true and y_pred do not match.

    Warning:
        Comparing labeled to unlabeled data will produce low accuracy scores.
        Make sure to input the same type of data for y_true and y_pred
    """

    def __init__(self, y_true, y_pred):
        super(PixelMetrics, self).__init__(
            y_true=(y_true != 0).astype('int'),
            y_pred=(y_pred != 0).astype('int'))

        self._y_true_sum = np.count_nonzero(self.y_true)
        self._y_pred_sum = np.count_nonzero(self.y_pred)

        # Calculations for IOU
        self._intersection = np.count_nonzero(np.logical_and(self.y_true, self.y_pred))
        self._union = np.count_nonzero(np.logical_or(self.y_true, self.y_pred))

    @classmethod
    def get_confusion_matrix(cls, y_true, y_pred, axis=-1):
        """Calculate confusion matrix for pixel classification data.

        Args:
            y_true (numpy.array): Ground truth annotations after any
                necessary transformations
            y_pred (numpy.array): Prediction array
            axis (int): The channel axis of the input arrays.

        Returns:
            numpy.array: nxn confusion matrix determined by number of features.
        """
        # Argmax collapses on feature dimension to assign class to each pixel
        # Flatten is required for confusion matrix
        y_true = y_true.argmax(axis=axis).flatten()
        y_pred = y_pred.argmax(axis=axis).flatten()
        return confusion_matrix(y_true, y_pred)

    @property
    def recall(self):
        try:
            _recall = self._intersection / self._y_true_sum
        except ZeroDivisionError:
            _recall = np.nan
        return _recall

    @property
    def precision(self):
        try:
            _precision = self._intersection / self._y_pred_sum
        except ZeroDivisionError:
            _precision = 0
        return _precision

    @property
    def f1(self):
        _recall = self.recall
        _precision = self.precision

        # f1 is nan if recall is nan and no false negatives
        if np.isnan(_recall) and _precision == 0:
            return np.nan

        f_measure = hmean([_recall, _precision])
        # f_measure = (2 * _precision * _recall) / (_precision + _recall)
        return f_measure

    @property
    def dice(self):
        y_sum = self._y_true_sum + self._y_pred_sum
        if y_sum == 0:
            warnings.warn('DICE score is technically 1.0, '
                          'but prediction and truth arrays are empty.')
            return 1.0

        return 2.0 * self._intersection / y_sum

    @property
    def jaccard(self):
        try:
            _jaccard = self._intersection / self._union
        except ZeroDivisionError:
            _jaccard = np.nan
        return _jaccard

    def to_dict(self):
        return {
            'jaccard': self.jaccard,
            'recall': self.recall,
            'precision': self.precision,
            'f1': self.f1,
            'dice': self.dice,
        }

def get_box_labels(arr):
    """Get the bounding box and label for all objects in the image.

    Args:
        arr (np.array): integer label array of objects.

    Returns:
        tuple(list(np.array), list(int)): A tuple of bounding boxes and
            the corresponding integer labels.
    """
    props = regionprops(np.squeeze(arr.astype('int')), cache=False)
    boxes, labels = [], []
    for prop in props:
        boxes.append(np.array(prop.bbox))
        labels.append(int(prop.label))
    boxes = np.array(boxes).astype('double')
    return boxes, labels

class ObjectMetrics(BaseMetrics):
    """Classifies object prediction errors as TP, FP, FN, merge or split

    The schema for this analysis was adopted from the description of
    object-based statistics in Caicedo et al. (2018) Evaluation of Deep
    Learning Strategies for Nucleus Segmentation in Fluorescence Images.
    BioRxiv 335216.
    The SEG metric was adapted from Maska et al. (2014). A benchmark for
    comparison of cell tracking algorithms.
    Bioinformatics 30, 1609-1617.
    The linear classification schema used to match objects in truth and
    prediction frames was adapted from Jaqaman et al. (2008).
    Robust single-particle tracking in live-cell time-lapse sequences.
    Nature Methods 5, 695-702.

    Args:
        y_true (numpy.array): Labeled ground truth annotation
        y_pred (numpy.array): Labled object prediction, same size as y_true
        cutoff1 (:obj:`float`, optional): Threshold for overlap in cost matrix,
            smaller values are more conservative, default 0.4
        cutoff2 (:obj:`float`, optional): Threshold for overlap in unassigned
            cells, smaller values are better, default 0.1
        seg (:obj:`bool`, optional): Calculates SEG score for cell tracking
            competition
        force_event_links(:obj:'bool, optional): Flag that determines whether to modify IOU
            calculation so that merge or split events with cells of very different sizes are
            never misclassified as misses/gains.
        is_3d(:obj:'bool', optional): Flag that determines whether or not the input data
            should be treated as 3-dimensional.

    Raises:
        ValueError: If y_true and y_pred are not the same shape
        ValueError: If data_type is 2D, if input shape does not have ndim 2 or 3
        ValueError: If data_type is 3D, if input shape does not have ndim 3
    """
    def __init__(self,
                 y_true,
                 y_pred,
                 cutoff1=0.4,
                 cutoff2=0.1,
                 force_event_links=False,
                 is_3d=False):

        # If 2D, dimensions can be 3 or 4 (with or without channel dimension)
        if not is_3d and y_true.ndim not in {2, 3}:
            raise ValueError('Expected dimensions for y_true (2D data) are 2 '
                             '(x, y) and 3 (x, y, chan). '
                             'Got ndim: {}'.format(y_true.ndim))

        elif is_3d and y_true.ndim != 3:
            raise ValueError('Expected dimensions for y_true (3D data) is 3.'
                             'Requires format is: (z, x, y)'
                             'Got ndim: {}'.format(y_true.ndim))

        super(ObjectMetrics, self).__init__(y_true=y_true, y_pred=y_pred)

        self.cutoff1 = cutoff1
        self.cutoff2 = cutoff2
        self.is_3d = is_3d

        self.compute_overlap = compute_overlap_3D if is_3d else compute_overlap

        self.n_true = len(np.unique(self.y_true[np.nonzero(self.y_true)]))
        self.n_pred = len(np.unique(self.y_pred[np.nonzero(self.y_pred)]))

        # keep track of every pair of objects through the detections dict
        # using tuple(true_index, pred_index): Detection as a key/vaue pair
        self._detections = set()

        # store the keys of relevant Detections in a set for easy fetching
        # types of detections
        self._splits = set()
        self._gained = set()
        self._missed = set()
        # types of errors
        self._merges = set()
        self._catastrophes = set()
        self._correct = set()

        # IoU: used to determine relative overlap of y_pred and y_true
        self.iou = np.zeros((self.n_true, self.n_pred))

        # used to determine seg score
        self.seg_thresh = np.zeros((self.n_true, self.n_pred))

        # Check if either frame is empty before proceeding
        if self.n_true == 0:
            logging.info('Ground truth frame is empty')

        if self.n_pred == 0:
            logging.info('Prediction frame is empty')

        self._calc_iou()  # set self.iou and update self.seg_thresh

        self.iou_modified = self._get_modified_iou(force_event_links)

        matrix = self._linear_assignment()

        # Identify direct matches as true positives
        correct_index = np.nonzero(matrix[:self.n_true, :self.n_pred])

        for i, j in zip(correct_index[0], correct_index[1]):
            self._add_detection(true_index=int(i), pred_index=int(j))

        # Calc seg score for true positives if requested
        iou_mask = np.where(self.seg_thresh == 0, self.iou, np.nan)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore', category=RuntimeWarning)
            # correct_index may be empty, suppress mean of empty slice warning
            self.seg_score = np.nanmean(iou_mask[correct_index])

        # Classify other errors using a graph
        G = self._array_to_graph(matrix)
        self._classify_graph(G)

        # Calculate pixel-level stats
        self.pixel_stats = PixelMetrics(y_true, y_pred)

    def _add_detection(self, true_index=None, pred_index=None):
        detection = Detection(true_index=true_index, pred_index=pred_index)

        self._detections.add(detection)

        # keep track of all error types
        # TODO: better way to do this?
        if detection.is_correct:
            self._correct.add(detection)
        if detection.is_gained:
            self._gained.add(detection)
        if detection.is_missed:
            self._missed.add(detection)
        if detection.is_split:
            self._splits.add(detection)
        if detection.is_merge:
            self._merges.add(detection)
        if detection.is_catastrophe:
            self._catastrophes.add(detection)

    def _calc_iou(self):
        """Calculates IoU matrix for each pairwise comparison between true and
        predicted. Additionally, if seg is True, records a 1 for each pair of
        objects where $|Tbigcap P| > 0.5 * |T|$
        """
        # Use bounding boxes to find masks that are likely to overlap
        y_true_boxes, y_true_labels = get_box_labels(self.y_true)
        y_pred_boxes, y_pred_labels = get_box_labels(self.y_pred)

        if not y_true_boxes.shape[0] or not y_pred_boxes.shape[0]:
            return  # cannot compute overlaps of nothing

        # has the form [gt_bbox, res_bbox]
        overlaps = self.compute_overlap(y_true_boxes, y_pred_boxes)

        # Find the bboxes that have any overlap
        # (ind_ corresponds to box number - starting at 0)
        ind_true, ind_pred = np.nonzero(overlaps)

        # TODO: this accounts for ~50+% of the time spent on calc_iou
        for index in range(ind_true.shape[0]):
            iou_y_true_idx = y_true_labels[ind_true[index]]
            iou_y_pred_idx = y_pred_labels[ind_pred[index]]

            is_true = self.y_true == iou_y_true_idx
            is_pred = self.y_pred == iou_y_pred_idx

            intersection = np.count_nonzero(np.logical_and(is_true, is_pred))
            union = np.count_nonzero(np.logical_or(is_true, is_pred))

            iou = intersection / union

            # Subtract 1 from index to account for skipping 0
            self.iou[iou_y_true_idx - 1, iou_y_pred_idx - 1] = iou

            if intersection > 0.5 * np.count_nonzero(self.y_true == index):
                self.seg_thresh[iou_y_true_idx - 1, iou_y_pred_idx - 1] = 1

    def _get_modified_iou(self, force_event_links):
        """Modifies the IoU matrix to boost the value for small cells.

        Args:
            force_event_links (:obj:`bool'): Whether to modify IOU values of
                large objects if they have been split or merged by
                a small object.

        Returns:
            np.array: The modified IoU matrix.
        """
        # identify cells that have matches in IOU but may be too small
        true_labels, pred_labels = np.nonzero(
            np.logical_and(self.iou > 0, self.iou < 1 - self.cutoff1)
        )

        iou_modified = self.iou.copy()

        for idx in range(len(true_labels)):
            # add 1 to get back to original label id
            true_idx, pred_idx = true_labels[idx], pred_labels[idx]
            true_label, pred_label = true_idx + 1, pred_idx + 1
            true_mask = self.y_true == true_label
            pred_mask = self.y_pred == pred_label

            # fraction of true cell that is contained within pred cell, vice versa
            true_in_pred = np.count_nonzero(
                self.y_true[pred_mask] == true_label) / np.sum(true_mask)
            pred_in_true = np.count_nonzero(
                self.y_pred[true_mask] == pred_label) / np.sum(pred_mask)

            iou_val = self.iou[true_idx, pred_idx]
            max_val = np.max([true_in_pred, pred_in_true])

            # if this cell has a small IOU due to its small size,
            # but is at least half contained within the big cell,
            # we bump its IOU value up so it doesn't get dropped from the graph
            if iou_val <= self.cutoff1 and max_val > 0.5:
                iou_modified[true_idx, pred_idx] = self.cutoff2

                # optionally, we can also decrease the IOU value of the cell
                # that swallowed up the small cell so that it doesn't directly
                # match a different cell
                if force_event_links and true_in_pred > 0.5:
                    fix_idx = np.nonzero(self.iou[:, pred_idx] >= 1 - self.cutoff1)
                    iou_modified[fix_idx, pred_idx] = 1 - self.cutoff1 - 0.01

                if force_event_links and pred_in_true > 0.5:
                    fix_idx = np.nonzero(self.iou[true_idx, :] >= 1 - self.cutoff1)
                    iou_modified[true_idx, fix_idx] = 1 - self.cutoff1 - 0.01

        return iou_modified

    def _get_cost_matrix(self):
        """Assembles cost matrix using the iou matrix and cutoff1

        The previously calculated iou matrix is cast into the top left and
        transposed for the bottom right corner. The diagonals of the two
        remaining corners are populated according to cutoff1. The lower the
        value of cutoff1 the more likely it is for the linear sum assignment
        to pick unmatched assignments for objects.
        """
        n_obj = self.n_true + self.n_pred
        matrix = np.ones((n_obj, n_obj))

        # Assign 1 - iou to top left and bottom right
        cost = 1 - self.iou_modified
        matrix[:self.n_true, :self.n_pred] = cost
        matrix[n_obj - self.n_pred:, n_obj - self.n_true:] = cost.T

        # Calculate diagonal corners
        bl = (self.cutoff1 * np.eye(self.n_pred)
              + np.ones((self.n_pred, self.n_pred))
              - np.eye(self.n_pred))
        tr = (self.cutoff1 * np.eye(self.n_true)
              + np.ones((self.n_true, self.n_true))
              - np.eye(self.n_true))

        # Assign diagonals to cm
        matrix[n_obj - self.n_pred:, :self.n_pred] = bl
        matrix[:self.n_true, n_obj - self.n_true:] = tr
        return matrix

    def _linear_assignment(self):
        """Runs linear sun assignment on cost matrix, identifies true
        positives and unassigned true and predicted cells.

        True positives correspond to assignments in the top left or bottom
        right corner. There are two possible unassigned positions: true cell
        unassigned in bottom left or predicted cell unassigned in top right.
        """
        cost_matrix = self._get_cost_matrix()

        results = linear_sum_assignment(cost_matrix)

        # Map results onto cost matrix
        assignment_matrix = np.zeros_like(cost_matrix)
        assignment_matrix[results] = 1
        return assignment_matrix

    def _array_to_graph(self, matrix):
        """Transform matrix for unassigned cells into a graph object

        In order to cast the iou matrix into a graph form, we treat each
        unassigned cell as a node. The iou values for each pair of cells is
        treated as an edge between nodes/cells. Any iou values equal to 0 are
        dropped because they indicate no overlap between cells.

        Args:
            matrix (np.array): Assignment matrix.
        """
        # Collect unassigned objects
        x, y = matrix.shape
        gained, _ = np.nonzero(matrix[x - self.n_pred:, :self.n_pred])
        missed, _ = np.nonzero(matrix[:self.n_true, y - self.n_true:])

        # Use meshgrid to get true and predicted object index for each val
        tt, pp = np.meshgrid(missed, gained, indexing='ij')

        true_nodes = tt.flatten()
        pred_nodes = pp.flatten()

        # construct list of edges for networkx
        G = nx.Graph()

        for t, p in zip(true_nodes, pred_nodes):
            # edges between overlapping objects only
            if self.iou_modified[t, p] >= self.cutoff2:
                G.add_edge('true_{}'.format(t), 'pred_{}'.format(p))

        # Add nodes to ensure all cells are included
        G.add_nodes_from(('true_{}'.format(n) for n in missed))
        G.add_nodes_from(('pred_{}'.format(n) for n in gained))

        return G

    def _classify_graph(self, G):
        """Assign each node in graph to an error type

        Nodes with a degree (connectivity) of 0 correspond to either false
        positives or false negatives depending on the origin of the node from
        either the predicted objects (false positive) or true objects
        (false negative). Any nodes with a connectivity of 1 are considered to
        be true positives that were missed during linear assignment.
        Finally any nodes with degree >= 2 are indicative of a merge or split
        error. If the top level node is a predicted cell, this indicates a merge
        event. If the top level node is a true cell, this indicates a split event.
        """
        # Find subgraphs, e.g. merge/split
        for g in (G.subgraph(c) for c in nx.connected_components(G)):
            # Get the highest degree node
            _, max_d = max(dict(g.degree).items(), key=operator.itemgetter(1))

            true_indices, pred_indices = [], []

            for node in g.nodes:
                node_type, index = node.split('_')
                index = int(index) + 1

                if node_type == 'true':
                    if max_d > 1:
                        true_indices.append(index)
                    else:
                        self._add_detection(true_index=index)

                if node_type == 'pred':
                    if max_d > 1:
                        pred_indices.append(index)
                    else:
                        self._add_detection(pred_index=index)

            self._add_detection(
                true_index=tuple(true_indices) if true_indices else None,
                pred_index=tuple(pred_indices) if pred_indices else None,
            )

    def _get_props(self, detection_type):
        prediction_types = {
            'gained',
        }
        is_pred_type = detection_type in prediction_types
        arr = self.y_pred if is_pred_type else self.y_true
        label_image = np.zeros_like(arr)
        attrname = '_{}'.format(detection_type)

        try:
            detections = getattr(self, attrname)
        except AttributeError:
            raise ValueError('Invalid detection_type: {}'.format(
                detection_type))

        for det in detections:
            idx = det.pred_index if is_pred_type else det.true_index
            idx = idx if isinstance(idx, tuple) else (idx,)
            for i in idx:
                label_image[arr == i] = i

        return regionprops(label_image)

    def __repr__(self):
        """Format the calculated statistics as a ``pd.DataFrame``."""
        return json.dumps(self.to_dict())

    def to_dict(self):
        """Return a dictionary representation of the calclulated metrics."""
        return {
            'n_pred': self.n_pred,
            'n_true': self.n_true,
            'correct_detections': self.correct_detections,
            'missed_detections': self.missed_detections,
            'gained_detections': self.gained_detections,
            'missed_det_from_merge': self.missed_det_from_merge,
            'gained_det_from_split': self.gained_det_from_split,
            'true_det_in_catastrophe': self.true_det_in_catastrophe,
            'pred_det_in_catastrophe': self.pred_det_in_catastrophe,
            'merge': self.merges,
            'split': self.splits,
            'catastrophe': self.catastrophes,
            'precision': self.precision,
            'recall': self.recall,
            'f1': self.f1,
            'seg': self.seg_score,
            'jaccard': self.jaccard,
            'dice': self.dice,
        }

    @property
    def correct_detections(self):
        return len(self._correct)

    @property
    def missed_detections(self):
        return len(self._missed)

    @property
    def gained_detections(self):
        return len(self._gained)

    @property
    def splits(self):
        return len(self._splits)

    @property
    def merges(self):
        return len(self._merges)

    @property
    def catastrophes(self):
        return len(self._catastrophes)

    @property
    def gained_det_from_split(self):
        gained_dets = 0
        for det in self._splits:
            true_idx = _cast_to_tuple(det.true_index)
            pred_idx = _cast_to_tuple(det.pred_index)
            gained_dets += len(true_idx) + len(pred_idx) - 2
        return gained_dets

    @property
    def missed_det_from_merge(self):
        missed_dets = 0
        for det in self._merges:
            true_idx = _cast_to_tuple(det.true_index)
            pred_idx = _cast_to_tuple(det.pred_index)
            missed_dets += len(true_idx) + len(pred_idx) - 2
        return missed_dets

    @property
    def true_det_in_catastrophe(self):
        return sum([len(d.true_index) for d in self._catastrophes])

    @property
    def pred_det_in_catastrophe(self):
        return sum([len(d.pred_index) for d in self._catastrophes])

    @property
    def split_props(self):
        return self._get_props('splits')

    @property
    def merge_props(self):
        return self._get_props('merges')

    @property
    def missed_props(self):
        return self._get_props('missed')

    @property
    def gained_props(self):
        return self._get_props('gained')

    @property
    def recall(self):
        try:
            recall = self.correct_detections / self.n_true
        except ZeroDivisionError:
            recall = 0
        return recall

    @property
    def precision(self):
        try:
            precision = self.correct_detections / self.n_pred
        except ZeroDivisionError:
            precision = 0
        return precision

    @property
    def f1(self):
        return hmean([self.recall, self.precision])

    @property
    def jaccard(self):
        return self.pixel_stats.jaccard

    @property
    def dice(self):
        return self.pixel_stats.jaccard

    def plot_errors(self):
        """Plots the errors identified from linear assignment code.

        This must be run with sequentially relabeled data.

        TODO: this is not working!
        """

        import matplotlib as mpl
        import matplotlib.pyplot as plt

        # erode edges for easier visualization of adjacent cells
        y_true = erode_edges(self.y_true.copy(), 1)
        y_pred = erode_edges(self.y_pred.copy(), 1)

        # semantic labels for each error
        categories = ['Background', 'missed', 'splits', 'merges',
                      'gained', 'catastrophes', 'correct']

        # Background is set to zero
        plotting_tif = np.zeros_like(y_true)

        # missed detections are tracked with true labels
        misses = [d.true_index for d in self._missed]
        plotting_tif[np.isin(y_true, misses)] = 1

        # skip background and misses, already done
        for i, category in enumerate(categories[2:]):
            # the rest are all on y_pred
            labels = list(getattr(self, '_{}'.format(category)))
            plotting_tif[np.isin(y_pred, labels)] = i + 2

        plotting_colors = ['Black', 'Pink', 'Blue', 'Green',
                           'tan', 'Red', 'Grey']

        cmap = mpl.colors.ListedColormap(plotting_colors)

        fig, ax = plt.subplots(nrows=1, ncols=1)
        mat = ax.imshow(plotting_tif, cmap=cmap,
                        vmin=np.min(plotting_tif) - .5,
                        vmax=np.max(plotting_tif) + .5)

        # tell the colorbar to tick at integers
        ticks = np.arange(len(categories))
        cbar = fig.colorbar(mat, ticks=ticks)
        cbar.ax.set_yticklabels(categories)
        fig.tight_layout()

class Metrics(object):
    """Class to calculate and save various segmentation metrics.

    Args:
        model_name (str): Name of the model which determines output file names
        outdir (:obj:`str`, optional): Directory to save json file, default ''
        cutoff1 (:obj:`float`, optional): Threshold for overlap in cost matrix,
            smaller values are more conservative, default 0.4
        cutoff2 (:obj:`float`, optional): Threshold for overlap in unassigned
            cells, smaller values are better, default 0.1
        pixel_threshold (:obj:`float`, optional): Threshold for converting
            predictions to binary
        ndigits (:obj:`int`, optional): Sets number of digits for rounding,
            default 4
        feature_key (:obj:`list`, optional): List of strings, feature names
        json_notes (:obj:`str`, optional): Str providing any additional
            information about the model
        force_event_links(:obj:`bool`, optional): Flag that determines whether to modify IOU
            calculation so that merge or split events with cells of very different sizes are
            never misclassified as misses/gains.
        is_3d(:obj:`bool`, optional): Flag that determines whether or not the input data
            should be treated as 3-dimensional.

    Examples:
        >>> from deepcell import metrics
        >>> m = metrics.Metrics('model_name')
        >>> all_metrics = m.run_all(y_true, y_pred)
        >>> m.save_to_json(all_metrics)
    """
    def __init__(self, model_name,
                 outdir='',
                 cutoff1=0.4,
                 cutoff2=0.1,
                 pixel_threshold=0.5,
                 ndigits=4,
                 crop_size=None,
                 feature_key=[],
                 json_notes='',
                 force_event_links=False,
                 is_3d=False,
                 **kwargs):
        self.model_name = model_name
        self.outdir = outdir
        self.cutoff1 = cutoff1
        self.cutoff2 = cutoff2
        self.pixel_threshold = pixel_threshold
        self.ndigits = ndigits
        self.crop_size = crop_size
        self.feature_key = feature_key
        self.json_notes = json_notes
        self.force_event_links = force_event_links
        self.is_3d = is_3d

        if 'seg' in kwargs:
            warnings.warn('seg is deprecated and will be removed '
                          'in a future release', DeprecationWarning)

        # Initialize output list to collect stats
        self.object_metrics = []
        self.pixel_metrics = []

    def df_to_dict(self, df, stat_type='pixel'):
        """Output pandas df as a list of dictionary objects

        Args:
            df (pandas.DataFrame): Dataframe of statistics for each channel
            stat_type (str): Category of statistic.

        Returns:
            list: List of dictionaries
        """

        # Initialize output dictionary
        L = []

        # Write out average statistics
        for k, v in df.mean().iteritems():
            L.append(dict(
                name=k,
                value=v,
                feature='average',
                stat_type=stat_type,
            ))

        # Save individual stats to list
        for i, row in df.iterrows():
            for k, v in row.iteritems():
                L.append(dict(
                    name=k,
                    value=v,
                    feature=i,
                    stat_type=stat_type,
                ))

        return L

    def calc_pixel_stats(self, y_true, y_pred, axis=-1):
        """Calculate pixel statistics for each feature.

        ``y_true`` should have the appropriate transform applied to match
        ``y_pred``. Each channel is converted to binary using the threshold
        ``pixel_threshold`` prior to calculation of accuracy metrics.

        Args:
            y_true (numpy.array): Ground truth annotations after transform
            y_pred (numpy.array): Model predictions without labeling

        Returns:
            list: list of dictionaries with each stat being a key.

        Raises:
            ValueError: If y_true and y_pred are not the same shape
        """
        n_features = y_pred.shape[axis]

        pixel_metrics = []

        slc = [slice(None)] * y_pred.ndim
        for i in range(n_features):
            slc[axis] = slice(i, i + 1)
            yt = y_true[slc] > self.pixel_threshold
            yp = y_pred[slc] > self.pixel_threshold
            pm = PixelMetrics(yt, yp)
            pixel_metrics.append(pm.to_dict())

        pixel_df = pd.DataFrame.from_records(pixel_metrics)

        # Calculate confusion matrix
        cm = PixelMetrics.get_confusion_matrix(y_true, y_pred, axis=axis)

        print('\n____________Pixel-based statistics____________\n')
        print(pixel_df)
        print('\nConfusion Matrix')
        print(cm)

        output = self.df_to_dict(pixel_df)

        output.append(dict(
            name='confusion_matrix',
            value=cm.tolist(),
            feature='all',
            stat_type='pixel'
        ))
        return output

    def calc_pixel_confusion_matrix(self, y_true, y_pred, axis=-1):
        """DEPRECATED: Use ``PixelMetrics.get_confusion_matrix``.

        Calculate confusion matrix for pixel classification data.

        Args:
            y_true (numpy.array): Ground truth annotations after any
                necessary transformations
            y_pred (numpy.array): Prediction array
            axis (int): The channel axis of the input arrays.

        Returns:
            numpy.array: nxn confusion matrix determined by number of features.
        """
        return PixelMetrics.get_confusion_matrix(y_true, y_pred, axis=axis)

    def calc_object_stats(self, y_true, y_pred, progbar=True):
        """Calculate object statistics and save to output

        Loops over each frame in the zeroth dimension, which should pass in
        a series of 2D arrays for analysis. 'metrics.split_stack' can be
        used to appropriately reshape the input array if necessary

        Args:
            y_true (numpy.array): Labeled ground truth annotations
            y_pred (numpy.array): Labeled prediction mask
            progbar (bool): Whether to show the progress tqdm progress bar

        Returns:
            list: list of dictionaries with each stat being a key.

        Raises:
            ValueError: If y_true and y_pred are not the same shape
            ValueError: If data_type is 2D, if input shape does not have ndim 3 or 4
            ValueError: If data_type is 3D, if input shape does not have ndim 4
        """
        if y_pred.shape != y_true.shape:
            raise ValueError('Input shapes need to match. Shape of prediction '
                             'is: {}.  Shape of y_true is: {}'.format(
                                 y_pred.shape, y_true.shape))

        # If 2D, dimensions can be 3 or 4 (with or without channel dimension)
        if not self.is_3d:
            if y_true.ndim not in {3, 4}:
                raise ValueError('Expected dimensions for y_true (2D data) are 3 or 4.'
                                 'Accepts: (batch, x, y), or (batch, x, y, chan)'
                                 'Got ndim: {}'.format(y_true.ndim))

        # If 3D, inputs must have 4 dimensions (batch, z, x, y) - cannot have channel dimension or
        # _classify_graph breaks, as it expects input to be 2D or 3D
        # TODO - add compatibility for multi-channel 3D-data
        else:
            if y_true.ndim != 4:
                raise ValueError('Expected dimensions for y_true (3D data) is 4. '
                                 'Required format is: (batch, z, x, y) '
                                 'Got ndim: {}'.format(y_true.ndim))

        all_object_metrics = []  # store all calculated metrics
        is_batch_relabeled = False  # used to warn if batches were relabeled

        for i in tqdm(range(y_true.shape[0]), disable=not progbar):
            # check if labels aren't sequential, raise warning on first occurence if so
            true_batch, pred_batch = y_true[i], y_pred[i]
            true_batch_relabel, _, _ = relabel_sequential(true_batch)
            pred_batch_relabel, _, _ = relabel_sequential(pred_batch)

            # check if segmentations were relabeled
            if not is_batch_relabeled:  # only one True is required
                is_batch_relabeled = not (
                    np.array_equal(true_batch, true_batch_relabel)
                    and np.array_equal(pred_batch, pred_batch_relabel)
                )

            o = ObjectMetrics(
                true_batch_relabel,
                pred_batch_relabel,
                cutoff1=self.cutoff1,
                cutoff2=self.cutoff2,
                force_event_links=self.force_event_links,
                is_3d=self.is_3d)

            all_object_metrics.append(o)

        if is_batch_relabeled:
            warnings.warn(
                'Provided data is being relabeled. Cell ids from metrics will not match '
                'cell ids in original data. Relabel your data prior to running the '
                'metrics package if you wish to maintain cell ids. ')

        # print the object report
        object_metrics = pd.DataFrame.from_records([
            o.to_dict() for o in all_object_metrics
        ])
        self.print_object_report(object_metrics)
        return object_metrics

    def summarize_object_metrics_df(self, df):
        correct_detections = int(df['correct_detections'].sum())
        n_true = int(df['n_true'].sum())
        n_pred = int(df['n_pred'].sum())

        _round = lambda x: round(x, self.ndigits)

        seg = df['seg'].mean()
        jaccard = df['jaccard'].mean()

        try:
            recall = correct_detections / n_true
        except ZeroDivisionError:
            recall = np.nan
        try:
            precision = correct_detections / n_pred
        except ZeroDivisionError:
            precision = 0

        errors = [
            'gained_detections',
            'missed_detections',
            'split',
            'merge',
            'catastrophe',
        ]

        bad_detections = [
            'gained_det_from_split',
            'missed_det_from_merge',
            'true_det_in_catastrophe',
            'pred_det_in_catastrophe',
        ]

        summary = {
            'correct_detections': correct_detections,
            'n_true': n_true,
            'n_pred': n_pred,
            'recall': _round(recall),
            'precision': _round(precision * 100),
            'seg': _round(seg * 100),
            'jaccard': _round(jaccard),
            'total_errors': 0,
        }
        # update bad detections
        for k in bad_detections:
            summary[k] = int(df[k].sum())
        # update error counts
        for k in errors:
            count = int(df[k].sum())
            summary[k] = count
            summary['total_errors'] += count
        return summary

    def print_object_report(self, object_metrics):
        """Print neat report of object based statistics

        Args:
            object_metrics (pd.DataFrame): DataFrame of all calculated metrics
        """
        summary = self.summarize_object_metrics_df(object_metrics)
        errors = [
            'gained_detections',
            'missed_detections',
            'split',
            'merge',
            'catastrophe'
        ]

        bad_detections = [
            'gained_det_from_split',
            'missed_det_from_merge',
            'true_det_in_catastrophe',
            'pred_det_in_catastrophe',
        ]

        print('\n____________Object-based statistics____________\n')
        print('Number of true cells:\t\t', summary['n_true'])
        print('Number of predicted cells:\t', summary['n_pred'])

        print('\nCorrect detections:  {}\tRecall: {}%'.format(
            summary['correct_detections'], summary['recall']))

        print('Incorrect detections: {}\tPrecision: {}%'.format(
            summary['n_pred'] - summary['correct_detections'],
            summary['precision']))

        print('\n')
        for k in errors:
            v = summary[k]
            name = k.replace('_', ' ').capitalize()
            if not name.endswith('s'):
                name += 's'

            try:
                err_fraction = v / summary['total_errors']
            except ZeroDivisionError:
                err_fraction = 0

            print('{name}: {val}{tab}Perc Error {percent}%'.format(
                name=name, val=v,
                percent=round(100 * err_fraction, self.ndigits),
                tab='\t' * (1 if ' ' in name else 2)))

        for k in bad_detections:
            name = k.replace('_', ' ').capitalize().replace(' det ', ' detections')
            print('{name}: {val}'.format(name=name, val=summary[k]))

        print('SEG:', round(summary['seg'], self.ndigits), '\n')

        print('Average Pixel IOU (Jaccard Index):',
              round(summary['jaccard'], self.ndigits), '\n')

    def run_all(self, y_true, y_pred, axis=-1):
        object_metrics = self.calc_object_stats(y_true, y_pred)
        pixel_metrics = self.calc_pixel_stats(y_true, y_pred, axis=axis)

        object_list = self.df_to_dict(object_metrics, stat_type='object')
        all_output = object_list + pixel_metrics
        self.save_to_json(all_output)

    def save_to_json(self, L):
        """Save list of dictionaries to json file with file metadata

        Args:
            L (list): List of metric dictionaries
        """
        todays_date = datetime.datetime.now().strftime('%Y-%m-%d')
        outname = os.path.join(
            self.outdir, '{}_{}.json'.format(self.model_name, todays_date))

        # Configure final output
        D = {}

        # Record metadata
        D['metadata'] = dict(
            model_name=self.model_name,
            date=todays_date,
            notes=self.json_notes
        )

        # Record metrics
        D['metrics'] = L

        with open(outname, 'w') as outfile:
            json.dump(D, outfile)

        logging.info('Saved to {}'.format(outname))

def split_stack(arr, batch, n_split1, axis1, n_split2, axis2):
    """Crops an array in the width and height dimensions to produce
    a stack of smaller arrays

    Args:
        arr (numpy.array): Array to be split with at least 2 dimensions
        batch (bool): True if the zeroth dimension of arr is a batch or
            frame dimension
        n_split1 (int): Number of sections to produce from the first split axis
            Must be able to divide arr.shape[axis1] evenly by n_split1
        axis1 (int): Axis on which to perform first split
        n_split2 (int): Number of sections to produce from the second split axis
            Must be able to divide arr.shape[axis2] evenly by n_split2
        axis2 (int): Axis on which to perform first split

    Returns:
        numpy.array: Array after dual splitting with frames in the zeroth dimension

    Raises:
        ValueError: arr.shape[axis] must be evenly divisible by n_split
            for both the first and second split

    Examples:
        >>> from deepcell import metrics
        >>> from numpy import np
        >>> arr = np.ones((10, 100, 100, 1))
        >>> out = metrics.split_stack(arr, True, 10, 1, 10, 2)
        >>> out.shape
        (1000, 10, 10, 1)
        >>> arr = np.ones((100, 100, 1))
        >>> out = metrics.split_stack(arr, False, 10, 1, 10, 2)
        >>> out.shape
        (100, 10, 10, 1)
    """
    # Check that n_split will divide equally
    if ((arr.shape[axis1] % n_split1) != 0) | ((arr.shape[axis2] % n_split2) != 0):
        raise ValueError(
            'arr.shape[axis] must be evenly divisible by n_split'
            'for both the first and second split')

    split1 = np.split(arr, n_split1, axis=axis1)

    # If batch dimension doesn't exist, create and adjust axis2
    if batch is False:
        split1con = np.stack(split1)
        axis2 += 1
    else:
        split1con = np.concatenate(split1, axis=0)

    split2 = np.split(split1con, n_split2, axis=axis2)
    split2con = np.concatenate(split2, axis=0)

    return split2con

def match_nodes(y_true, y_pred):
    """Loads all data that matches each pattern and compares the graphs.

    Args:
        y_true (numpy.array): ground truth array with all cells labeled uniquely.
        y_pred (numpy.array): data array to match to unique.

    Returns:
        numpy.array: IoU of ground truth cells and predicted cells.
    """
    num_frames = y_true.shape[0]
    # TODO: does max make the shape bigger than necessary?
    iou = np.zeros((num_frames, np.max(y_true) + 1, np.max(y_pred) + 1))

    # Compute IOUs only when neccesary
    # If bboxs for true and pred do not overlap with each other, the assignment
    # is immediate. Otherwise use pixelwise IOU to determine which cell is which

    # Regionprops expects one frame at a time
    for frame in range(num_frames):
        gt_frame = y_true[frame]
        res_frame = y_pred[frame]

        gt_props = regionprops(np.squeeze(gt_frame.astype('int')))
        gt_boxes = [np.array(gt_prop.bbox) for gt_prop in gt_props]
        gt_boxes = np.array(gt_boxes).astype('double')
        gt_box_labels = [int(gt_prop.label) for gt_prop in gt_props]

        res_props = regionprops(np.squeeze(res_frame.astype('int')))
        res_boxes = [np.array(res_prop.bbox) for res_prop in res_props]
        res_boxes = np.array(res_boxes).astype('double')
        res_box_labels = [int(res_prop.label) for res_prop in res_props]

        # has the form [gt_bbox, res_bbox]
        overlaps = compute_overlap(gt_boxes, res_boxes)

        # Find the bboxes that have overlap at all
        # (ind_ corresponds to box number - starting at 0)
        ind_gt, ind_res = np.nonzero(overlaps)

        # frame_ious = np.zeros(overlaps.shape)
        for index in range(ind_gt.shape[0]):
            iou_gt_idx = gt_box_labels[ind_gt[index]]
            iou_res_idx = res_box_labels[ind_res[index]]
            intersection = np.logical_and(
                gt_frame == iou_gt_idx, res_frame == iou_res_idx)
            union = np.logical_or(
                gt_frame == iou_gt_idx, res_frame == iou_res_idx)
            iou[frame, iou_gt_idx, iou_res_idx] = intersection.sum() / union.sum()

    return iou

import numpy,from numpy,from scipy,import networkx


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import numpy as np
import re

from .linalg import dot, matmul, transpose
from .manipulation import squeeze, unsqueeze, reshape
from .math import multiply
from .math import sum as paddle_sum
from ..fluid.framework import _in_legacy_dygraph
from paddle import _C_ops
from ..fluid.data_feeder import check_variable_and_dtype, check_type, check_dtype
from ..fluid.layer_helper import LayerHelper
from ..fluid.framework import _non_static_mode, in_dygraph_mode, _in_legacy_dygraph
import collections
import string
import opt_einsum

from paddle.common_ops_import import dygraph_only

__all__ = []

def parse_op_labels(labelstr, operand):
    '''
    Parse labels for an input operand.

    Parameters
    ----------
    labelstr:
        the input label string
    operand:
        the input operand

    Returns
    -------
    the input operand's full label string in which all anonymous dimensions are
    labeled in dots.
    '''
    # Sanity checks
    for c in labelstr.replace('.', ''):
        assert c.isalpha(), (
            f"Invalid equation: {c} is not a valid label, which should be letters."
        )

    assert labelstr.replace('...', '', 1).find('.') == -1, (
        f"Invalid equation: `.` is found outside of an ellipsis.")

    # Check shape. Note, in Paddle a tensor rank is always nonzero
    ndims = len(operand.shape)
    assert ndims > 0

    full_labelstr = labelstr.replace('...', '.' * (ndims - len(labelstr) + 3))

    assert len(full_labelstr) == ndims, (
        f"Invalid equation: the label string '{labelstr}' misses dimensions.")

    return full_labelstr

def parse_labels(labelstr, operands):
    '''
    Parse label strings for all input operands.

    Parameters
    ----------
    labelstr:
        The equation's label string
    operands:
        The input operands

    Returns
    -------
    list of full label strings for all input operands
    '''

    nop_labels = labelstr.split(',')
    assert len(nop_labels) == len(operands), (
        f"Invalid equation: the number of operands is {len(operands)}, "
        f"but found {len(nop_labels)} segments in the label equation.")

    return list(map(parse_op_labels, nop_labels, operands))

def validate_rhs(rhs, input_labels, n_bcast_dims):
    '''
    Check whether the equation's right hand side is valid
    '''
    # Sanity check.
    if n_bcast_dims > 0:
        assert '...' in rhs, (
            f"Invalid equation: missing ellipsis in output labels.")

    rhs = rhs.replace('...', '')
    rhs_set = set(rhs)

    # Hidden assumption: availble labels don't include '.'
    assert '.' not in input_labels

    # Verify that output labels all come from the set of input labels
    non_input_labels = rhs_set.difference(input_labels)
    assert not non_input_labels, (
        f"Invalid equation: "
        f"output label {sorted(non_input_labels)} not used by any input.")
    # Verify that output labels are not duplicate
    assert len(rhs) == len(rhs_set), (
        f"Invalid equation: duplicate output labels are found.")

def build_view(in_labels, out_labels):
    '''
    Build an inverse map of dimension indices. Three conditions must hold for
    the result to be meaningful.
    First, no duplicate letter labels in each label string.
    Second, the number of dots in dimout_labels >= that in in_labels.
    Third, dots are contiguous in each label string.

    Parameters
    ----------
    in_labels:
        The dimension labels to map to
    out_labels:
        The dimension labels to map from

    Returns
    -------
    The inverse map from out_labels to in_labels. The length of the inverse map equals that of
    out_labels. -1 is filled if there's no matching intput dimension for a specific label.

    Examples
    --------
    in_labels = 'ij..', out_labels = '..ji'
    inv_map = [2, 3, 1, 0]
    in_labels = 'ij..', out_labels = '..kji'
    inv_map = [2, 3, -1, 1, 0]
    '''

    inv_map = [-1] * len(out_labels)

    # First build the broadcast dimension mapping
    # Find the broadcast index range in out_labels
    r = re.search(r'\.+', out_labels)
    if r:
        start, end = r.start(), r.end()
        s = re.search(r'\.+', in_labels)
        # fill the broadcast dimension indices from right to left.
        if s:
            for ax, dim in zip(
                    range(start, end)[::-1], range(s.start(), s.end())[::-1]):
                inv_map[ax] = dim

    # Now work on non-broadcast dimensions
    if r:
        it = itertools.chain(range(start), range(end, len(out_labels)))
    else:
        it = iter(range(len(out_labels)))

    for i in it:
        inv_map[i] = in_labels.find(out_labels[i])

    return inv_map

def build_global_view(nop_labels, rhs, n_bcast_dims):
    '''
    Build the global view, which is a layout of all dimension labels
    plus an index table that maps from the layout to the dimensions
    in each operand. In the global view, the dimensions are arranged
    such that output ones are put on the left and contraction ones
    are put on the right.

    Parameters
    ----------
    nop_labels:
        The input full label strings of all input operands
    rhs:
        The equation right hand side
    n_bcast_dims:
        The maxium number of broadcast dimensions

    Returns
    -------
    A tuple of g_labels, g_view, g_nout, g_count
    g_labels:
        the layout of all labels in a string
    g_view:
        the index table
    g_nout:
        the number of output dimensions
    g_count:
        the counter array for dimension contractions
    '''
    # Put all labels in alphabetical order
    concat = sorted(''.join(nop_labels).replace('.', ''))
    labels, count = [], []
    for a, b in zip(['.'] + concat, concat):
        if a != b:
            labels.append(b)
            count.append(1)
        else:
            count[-1] += 1

    if rhs != None:
        validate_rhs(rhs, labels, n_bcast_dims)
        g_labels_out = rhs.replace('...', '.' * n_bcast_dims)
    else:
        g_labels_out = '.' * n_bcast_dims + ''.join(
            l for l, c in zip(labels, count) if c == 1)

    for i in range(len(count))[::-1]:
        if labels[i] in g_labels_out:
            labels.pop(i)
            count.pop(i)

    g_labels_sum = ''.join(labels)
    g_labels = g_labels_out + g_labels_sum
    g_view = list(map(lambda i: build_view(i, g_labels), nop_labels))
    g_nout = len(g_labels_out)
    g_count = count

    return g_labels, g_view, g_nout, g_count

def build_global_shape(g_view, g_labels, op_shapes):
    '''
    The global shape is the shape of all dimensions rearranged and broadcasting
    to the global view. It's a reference data structure for einsum planning.

    Parameters
    ----------
    g_view:
        the global view
    op_shapes:
        the shapes of the all operands

    Returns
    -------
    g_shape:
        the global shape vector
    g_masks:
        list of shape masks for each operand. A dimension's shape mask is a boolean
        indicating whether its size > 1, in other words, it's not squeezable
    '''
    view_shapes = []
    g_masks = []

    for view, op_shape in zip(g_view, op_shapes):
        view_shapes.append([op_shape[dim] if dim > -1 else 1 for dim in view])

    g_shape = [set(sizes_per_ax) - {1} for sizes_per_ax in zip(*view_shapes)]

    non_bcastable = [ax for ax, sizes in enumerate(g_shape) if len(sizes) > 1]

    assert not non_bcastable, (
        f"Invalid operands: label {g_labels[non_bcastable[0]]} "
        f"corresponds to non-broadcastable dimensions.")

    g_shape = [sizes.pop() if len(sizes) > 0 else 1 for sizes in g_shape]

    g_masks = [[s > 1 or s == -1 for s in view_shape]
               for view_shape in view_shapes]

    return g_shape, g_masks

def has_duplicated_labels(labels):
    '''
    Returns True if there is any duplicate label.
    '''
    labels = labels.replace('.', '')
    return len(labels) > len(set(labels))

def diagonalize(labels, operand):
    '''
    Merges dimensions with duplicate labels.

    For those dimensions with duplicate labels, merge them into one dimension
    which represents the diagonal elements. This requires the dimensions with
    duplicate labels are equal sized.

    Examples
    --------
    'ijj...i' would be merged into 'ij...'
    '''
    assert not has_duplicated_labels(labels), (
        f'Duplicate labels are not supported.')

    return labels, operand

def plan_reduce(plan, op, reduce_dims, keepdim):
    '''
    Add reduce to the plan
    '''
    varname = f'op{op}'

    f = lambda var, dims: paddle_sum(var, dims, keepdim=keepdim)
    step = f, [varname], varname, reduce_dims
    plan.add_step(step)

def plan_scalar_prod(plan, op1, op2):
    varnames = [f'op{op1}', f'op{op2}']
    f = lambda var1, var2: paddle_sum(var1) * var2
    # f = lambda var1, var2: var1 * var2
    step = f, varnames, varnames[1]
    plan.add_step(step)

def plan_matmul(plan, g_view, op1, op2, g_supports, g_shape, I, J1, J2, K):
    '''
    plan matmul
    '''
    # Transpose and re-shape op1 and op2 in I, J1, K and I, J2, K
    # Then apply matmul(x, y, transpose_x=False, tranpose_y=True)
    var1, var2 = f'op{op1}', f'op{op2}'

    op1_view, op2_view = [g_view[op] for op in (op1, op2)]

    I1 = [idx for idx in I if op1_view[idx] >= 0]
    I2 = [idx for idx in I if op2_view[idx] >= 0]
    op1_view = np.array(op1_view)
    op1_dims = op1_view[I1 + J1 + K]

    op2_view = np.array(op2_view)
    op2_dims = op2_view[I2 + J2 + K]

    op1_mask, op2_mask = [g_supports[op] for op in (op1, op2)]
    op1_vshape = np.array([s if m else 1 for s, m in zip(g_shape, op1_mask)])
    op2_vshape = np.array([s if m else 1 for s, m in zip(g_shape, op2_mask)])
    vshape = np.maximum(op1_vshape, op2_vshape)

    i1, i2, j1, j2, k = map(len, (I1, I2, J1, J2, K))

    if any(op1_dims != np.arange(len(op1_dims))):
        # print(f'perm1: {perm1}')
        step = transpose, [var1], var1, list(op1_dims)
        plan.add_step(step)

    if any(op2_dims != np.arange(len(op2_dims))):
        # print(f'perm2: {perm2}')
        step = transpose, [var2], var2, list(op2_dims)
        plan.add_step(step)

    # Check if conditions hold for turnning the operation into a matmul
    if j1 + j2 > 0 and k > 0 and -1 not in np.concatenate(
        (op1_vshape, op2_vshape)):
        op1_shape = list(op1_vshape[I]) + [np.prod(op1_vshape[J1])
                                           ] + [np.prod(op1_vshape[K])]
        op2_shape = list(op2_vshape[I]) + [np.prod(op2_vshape[J2])
                                           ] + [np.prod(op2_vshape[K])]

        # Merge J dims and K dims by reshaping
        step = reshape, [var1], var1, op1_shape
        plan.add_step(step)
        step = reshape, [var2], var2, op2_shape
        plan.add_step(step)

        # Matmul
        step = matmul, [var1, var2], var2, False, True
        plan.add_step(step)

        # Reshape back
        shape = list(vshape[I + J1 + J2])
        step = reshape, [var2], var2, shape
        plan.add_step(step)

    elif j1 == j2 == k == 1:
        # Can still do matmul even unknown shapes are present
        step = matmul, [var1, var2], var2, False, True
        plan.add_step(step)

    # In the rest cases we opt for ops other than matmul
    else:
        # unsqueeze operands include J1...J2... dimensions
        if j2:
            fill = list(range(i1 + j1, i1 + j1 + j2))
            step = unsqueeze, [var1], var1, fill
            plan.add_step(step)
        if j1:
            fill = list(range(i2, i2 + j1))
            step = unsqueeze, [var2], var2, fill
            plan.add_step(step)
        # In case of no dimensions to contract, do an elementwise multiply
        if k == 0:
            # make broadcast
            step = multiply, [var1, var2], var2
            plan.add_step(step)
        # Contract and no join, turn into a dot
        elif j1 + j2 == 0 and k == 1:
            step = unsqueeze, [var1], var1, [-2]
            plan.add_step(step)
            step = unsqueeze, [var2], var2, [-1]
            plan.add_step(step)
            step = matmul, [var1, var2], var2
            plan.add_step(step)
            step = squeeze, [var2], var2, [-1, -2]
            plan.add_step(step)
        elif j1 + j2 == 0 and not-1 in np.concatenate(
            (op1_vshape[K], op2_vshape[K])):
            assert all(op1_vshape[K] == op2_vshape[K])
            step = reshape, [var1], var1, list(op1_vshape[
                I]) + [1] + [np.prod(op1_vshape[K])]
            plan.add_step(step)
            step = reshape, [var2], var2, list(op2_vshape[
                I]) + [1] + [np.prod(op2_vshape[K])]
            plan.add_step(step)
            step = matmul, [var1, var2], var2, False, True
            plan.add_step(step)
            step = squeeze, [var2], var2, [-1, -2]
            plan.add_step(step)
        else:
            step = multiply, [var1, var2], var2
            plan.add_step(step)
            reduce_dims = list(range(-k, 0))
            plan_reduce(plan, op2, reduce_dims, keepdim=False)

    # Wrap up, updating auxiliary data
    # Updating g_mask for I and J axes
    for ax in I + J1 + J2:
        op2_mask[ax] = vshape[ax] > 1 or vshape[ax] == -1

    for ax in K:
        op2_mask[ax] = False

    for ax in range(len(op2_view)):
        op2_view[ax] = -1
    dim = 0
    for ax in I + J1 + J2:
        op2_view[ax], dim = dim, dim + 1

    g_view[op2] = list(op2_view)

def plan_summation(plan, g_view, op1, op2, g_supports, g_shape, g_count,
                   n_bcast):
    '''
    Plan various kinds of summation
    '''
    op1_view, op2_view = g_view[op1], g_view[op2]
    op1_mask, op2_mask = g_supports[op1], g_supports[op2]

    ndim = len(op1_view)
    nout = ndim - len(g_count)

    count = [0] * nout + g_count

    I, K, J1, J2 = list(range(n_bcast)), [], [], []

    for ax, dim1, dim2 in zip(
            range(n_bcast, ndim), op1_view[n_bcast:], op2_view[n_bcast:]):

        if (dim1 != -1) != (dim2 != -1):
            if dim1 != -1:
                J1.append(ax)
            else:
                J2.append(ax)
        elif dim1 != -1:
            fold = int(op1_mask[ax]) + int(op2_mask[ax])
            if ax >= nout and fold == count[ax]:
                # Ready to fold the dimensions
                K.append(ax)
                count[ax] -= fold
            else:
                I.append(ax)
                count[ax] -= max(fold - 1, 0)

    # Update g_count
    g_count[:] = count[nout:]

    # Now it's OK to merge the K dims as the same shape holds
    # print(f'I: {I}   J1: {J1}    J2: {J2}   K: {K}')
    plan_matmul(plan, g_view, op1, op2, g_supports, g_shape, I, J1, J2, K)

def rearrange(axes):
    perm, fill = [], []
    for ax, dim in enumerate(axes):
        if dim < 0:
            fill.append(ax)
        else:
            perm.append(dim)
    # Trivial permutation returns []
    if all(i == dim for i, dim in enumerate(perm)):
        perm = []

    return perm, fill

def plan_broadcast(plan, operands, nop_axes):
    '''
    Plan broadcast across
    '''
    nop = len(operands)
    varnames = [f'op{i}' for i in range(nop)]

    for i, op_axes in zip(range(nop), nop_axes):
        # Re-arrange the dimesions according to the global layout
        perm, fill = rearrange(op_axes)
        var = varnames[i]
        if perm:
            step = transpose, [var], var, perm
            plan.add_step(step)
        if fill:
            step = unsqueeze, [var], var, fill
            plan.add_step(step)

    def f(*args):
        expr = ' * '.join(varnames)
        return eval(expr, dict(zip(varnames, args)))

    step = f, varnames, None
    plan.add_step(step)

class Plan:
    def __init__(self):
        self.env = {}
        self.steps = []

    def add_step(self, step):
        self.steps.append(step)

    def get_var(self, varname):
        return self.env[varname] if varname in self.env else None

    def set_var(self, varname, var):
        self.env[varname] = var

    def show(self):
        res = None
        for f, in_varnames, out_varname, *args in self.steps:
            print(repr((out_varname, f, *in_varnames, *args)))
        return res

    def execute(self):
        res = None
        for f, in_varnames, out_varname, *args in self.steps:
            res = f(*map(self.get_var, in_varnames), *args)
            if out_varname:
                self.set_var(out_varname, res)
        return res

def plan_einsum(operands, g_view, g_shape, g_supports, g_count, n_bcast):
    '''
    Plans the actual execution steps.
    Results
    -------
    the execution plan
    '''
    nop = len(operands)
    ndim = len(g_view[0])
    nout = ndim - len(g_count)

    # Initialize a plan with an environment
    plan = Plan()
    op_names = [f'op{i}' for i in range(nop)]
    list(map(plan.set_var, op_names, operands))

    # In case no dimensions to combine, do broadcast straight across
    if not g_count:
        plan_broadcast(plan, operands, g_view)
        return plan

    # Down count degenerate contraction dimensions.
    for view, support in zip(g_view, g_supports):
        # To collect the down count number, we use a type casting trick
        down_count = [
            int((d + 1) and (not s))
            for d, s in zip(view[nout:], support[nout:])
        ]
        for i, count in enumerate(down_count):
            g_count[i] -= count

    # Reduce any dimension for which g_support is set and g_count == 1
    for i, view, mask in zip(range(nop), g_view, g_supports):
        to_reduce = []
        for dim, masked, count in zip(view[nout:], mask[nout:], g_count):
            to_reduce.append(dim if (masked and count == 1) else -1)

        reduce_dims = list(filter(lambda x: x > -1, to_reduce))
        if reduce_dims:
            plan_reduce(plan, i, reduce_dims, keepdim=True)

        # Unset mask and decrease g_count for the reduced dimensions
        for i, d in enumerate(to_reduce):
            ax = i + nout
            mask[ax] = mask[ax] and (d == -1)
            g_count[i] -= 0 if d == -1 else 1

    # Plan the summations over the operand sequence
    for i in range(nop):
        # plan a single step

        if i == 0:
            continue

        # We'd like to arrange the dimensions in the following way:
        # [I...  J... K...]
        # [I...  J... K...]
        # where
        #       I... are aligned and not to be combined immediately
        #       J... are not aligned and not to be combined immediately
        #       K... are aligned and should be immediately combined
        # At this point the non-trivial broadcast dimensinos in K are already reduced
        # and removed. That means all K dimensions are aligned and their sizes are not 1.
        # We then inspect the layout of I,J,K plus the above observation to make
        # specializatoin decisions.  The current strategy is set as follows:
        #  (1) if I... J... K... are all empty, it's multiplying a scalar
        #  (2) if K... are empty, better use a broadcast
        #  (3) if I... J... empty and K... not empty, a vector-vector multiply (or a dot)
        #  (4) Elsewise, either I... or J... not empty, and K... not empty, use a general matmul

        # Resolve the summation kind: dot, matmul or *
        if not any(g_supports[i - 1]):
            # op1 is a one element tensor.
            plan_scalar_prod(plan, i - 1, i)
        else:
            plan_summation(plan, g_view, i - 1, i, g_supports, g_shape, g_count,
                           n_bcast)

    # for ax, dim in enumerate(g_view[nop-1][:nout]):
    #     assert dim == ax
    assert all(not masked for masked in g_supports[nop - 1][nout:])

    view = g_view[-1]
    if any(ax != dim for ax, dim in enumerate(view[:nout])):
        perm = [dim for dim in view if dim >= 0]
        if sorted(perm) != perm:
            varname = f'op{nop-1}'
            step = transpose, [varname], varname, perm
            plan.add_step(step)
        dim = 0
        unsqueeze_dims = []
        for ax, d in enumerate(view):
            if d != -1:
                view[ax], dim = dim, dim + 1
        for ax, d in enumerate(view[:nout]):
            if d == -1:
                unsqueeze_dims.append(ax)
        if unsqueeze_dims:
            varname = f'op{nop-1}'
            step = unsqueeze, [varname], varname, unsqueeze_dims
            plan.add_step(step)

    squeeze_dims = [dim for dim in view[nout:] if dim != -1]
    if squeeze_dims:
        # plan_reduce(plan, nop-1, reduce_dims, keepdim=False)
        varname = f'op{nop-1}'
        step = squeeze, [varname], varname, squeeze_dims
        plan.add_step(step)

    return plan

def preprocess(equation, *operands):
    """
    check equation / raise error, default right labels generation
    """
    equation = equation.replace(" ", "")
    nop = len(operands)
    assert nop > 0, "Required at least one operand in Einsum API, but received %s " % nop

    # Part the equation to left hand side and right hand side
    lhs, *rhs = equation.lower().split('->')
    assert len(rhs) < 2, "Invalid equation: multiple `->` were found."

    labels = parse_labels(lhs, operands)
    # Note, we distinguish between 'ij->' and 'ij' by setting rhs to '' and None
    rhs = rhs[0] if rhs else None
    if rhs is None:
        rhs = rhs_inference(lhs)

    assert len(lhs.split(',')) == len(operands), (
        f"Invalid equation: the number of operands is {len(operands)}, "
        f"but found {len(lhs.split(','))} segments in the label equation.")

    assert not ('...' in lhs and '...' not in rhs
                ), f'Invalid equation: missing ellipsis in output labels.'

    assert not (len(list(filter(has_duplicated_labels, lhs.split(',')))) > 0
                ), f'Duplicate labels are not supported.'

    assert not has_duplicated_labels(
        rhs), f'Invalid equation: duplicate output labels are found.'

    return lhs, rhs, labels

def parse_fake_shape(equation, operands, labels):
    """
    this shape is just used for operands planning. may differ with the original shape.
    for example:
    ... is replaced by 1
    -1  is replaced by 1
    Results
    -------
    list of shape
    """
    shaped = collections.namedtuple('shaped', ['shape'])

    def fake_shape(label, op):
        assert len(op.shape) == len(
            label
        ), "length of shape and length of label must be the same, but received %d != %d" % (
            len(op.shape), len(label))
        fakes = [s for i, (l, s) in enumerate(zip(label, op.shape)) if l != '.']
        fakes = list(map(abs, fakes))  # make -1 -> 1
        if '.' in label:
            fakes.insert(label.index('.'), 1)
        return shaped(fakes)

    out = list(map(fake_shape, labels, operands))
    return out

def rhs_inference(lhs):
    def is_free(key):
        return cnt.get(key) == 1 and key not in ['.', ',']

    cnt = collections.Counter(lhs)
    rhs = "..." if '...' in lhs else ""
    rhs = rhs + "".join(filter(is_free, sorted(cnt.elements())))
    return rhs

def gen_equation_for_opteinsum(lhs, rhs):
    """
    1. gen rhs if rhs is None
    2. '...' -> 'A'
    """

    def get_used_label(counter):
        used = set(counter.elements())
        for c in string.ascii_lowercase:
            if c not in used: return c
        raise ValueError(
            "You have used all `a` - `z`, there can't find a unused for einsum optimization"
        )

    cnt = collections.Counter(lhs)
    broadcast_label = get_used_label(cnt)
    if rhs is None:
        rhs = rhs_inference(lhs)
    lhs = lhs.replace("...", broadcast_label)
    rhs = rhs.replace("...", broadcast_label)
    return lhs + "->" + rhs, broadcast_label

def einsum_v2(equation, *operands):
    """
    einsum v2 implementation.
    1. Implement C++ EinsumOp.
    2. V2 create the EinsumOp to calculate, so just a little verifty work in python.
    3. V2 use opt_einsum.contract_path to optimize the multivariable einsum.
    """
    n_op = len(operands)
    lhs, rhs, labels = preprocess(equation, *operands)

    if n_op <= 2:
        return gen_einsum_op(lhs + '->' + rhs, *operands)

    shapes = parse_fake_shape(lhs, operands, labels)
    opt_equation, broadcast_label = gen_equation_for_opteinsum(lhs, rhs)
    _, cons = opt_einsum.contract_path(opt_equation, *shapes, einsum_call=True)
    var_list = list(operands)
    for path in cons:
        (a, b), _, eq, *__ = path
        assert a > b, "Assume the first var_idx is smaller than the second_idx. opt_einsum can guarantee it."
        var_s = [var_list.pop(a), var_list.pop(b)]
        eq = eq.replace(broadcast_label, "...")
        var_list.append(gen_einsum_op(eq, *var_s))
    assert len(
        var_list
    ) == 1, "There must be one elements in list, but received %d." % len(
        var_list)
    return var_list[0]

def gen_einsum_op(equation, *operands):
    """
    EinsumOp Python Interface:
    """
    assert len(operands) <= 2, "Only support two operands in EinsumOp."
    if in_dygraph_mode():
        return _C_ops.final_state_einsum(operands, equation)[0]

    if _in_legacy_dygraph():
        # dygraph
        return _C_ops.einsum(operands, len(operands), 'equation', equation)[0]

    # static graph
    for inp in operands:
        check_variable_and_dtype(inp, 'dtype', ['float32', 'float64'], 'einsum')
    check_type(equation, 'equation', str, 'einsum')
    helper = LayerHelper('einsum', **locals())
    out = helper.create_variable_for_type_inference(dtype=operands[0].dtype)
    attrs = dict()
    attrs['equation'] = equation
    caches = [
        helper.create_variable_for_type_inference(dtype=operands[0].dtype)
        for i in range(len(operands))
    ]
    helper.append_op(
        type='einsum',
        inputs={'Operands': operands},
        outputs={'Out': out,
                 "InnerCache": caches},
        attrs=attrs)
    return out

def einsum(equation, *operands):
    r"""
    einsum(equation, *operands)

    The current version of this API should be used in dygraph only mode.

    Einsum offers a tensor operation API which allows using the Einstein summation
    convention or Einstain notation. It takes as input one or multiple tensors and
    produces as output one tensor.

    Einsum is able to perform a variety of tensor operations. Following lists a few:

        - for single operand
            - trace
            - diagonal
            - transpose
            - sum
        - for double operands
            - dot
            - outer
            - broadcasting and elementwise multiply
            - matrix multiply
            - batched matrix multiply
        - for many operads
            - broadcasting multiply
            - chained matrix multiply

    **The summation notation**

        - The tensor dimensions are labeled using uncased English letters. E.g., `ijk`
        relates to a three dimensional tensor whose dimensions are labeled i, j, and k.
        - The equation is `,` separated into terms, each being a distinct input's
        dimension label string.
        - Ellipsis `...` enables broadcasting by automatically converting the unlabeled
        dimensions into broadcasting dimensions.
        - Singular labels are called free labels, duplicate are dummy labels. Dummy labeled
        dimensions will be reduced and removed in the output.
        - Output labels can be explicitly specified on the right hand side of `->` or omitted.
        In the latter case, the output labels will be inferred from the input labels.
            - Inference of output labels
                - Broadcasting label `...`, if present, is put on the leftmost position.
                - Free labels are reordered alphabetically and put after `...`.
            - On explicit output labels
                - If broadcasting is enabled, then `...` must be present.
                - The output labels can be an empty, an indication to output as a scalar
                the sum over the original output.
                - Non-input labels are invalid.
                - Duplicate labels are invalid.
                - For any dummmy label which is present for the output, it's promoted to
                a free label.
                - For any free label which is not present for the output, it's lowered to
                a dummy label.
        - Examples
            - '...ij, ...jk', where i and k are free labels, j is dummy. The output label
            string is '...ik'
            - 'ij -> i', where i is a free label and j is a dummy label.
            - '...ij, ...jk -> ...ijk', where i, j and k are all free labels.
            - '...ij, ...jk -> ij', an invalid equation since `...` is not present for
            the output.

    **The summation rule**

    The summation procedure can be outlined as follows, although the actual steps taken
    may vary significantly due to implementation specific optimization.

        - Step 1: preparation for broadcasting, that is, transposing and unsqueezing
        the input operands to have each resulting dimension identically labeled across
        all the input operands.
        - Step 2: broadcasting multiply all the resulting operands from step 1.
        - Step 3: reducing dummy labeled dimensions.
        - Step 4: transposing the result tensor to match the output labels.

    **On trace and diagonal**

    The trace and diagonal are planned yet unimplemented features.

    Args:
        equation (`str`):
            The summation terms using the Einstein summation notation.
        operands (`list|Tensor`):
            The input tensors over which to compute the Einstein summation. The number of
            operands should equal the number of input terms in the equation.

    Returns:
        result (`Tensor`): the result tensor.

    Examples:
        .. code-block:: python

        import paddle
        paddle.seed(102)
        x = paddle.rand([4])
        y = paddle.rand([5])

        # sum
        print(paddle.einsum('i->', x))
        # Tensor(shape=[], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   1.95791852)

        # dot
        print(paddle.einsum('i,i->', x, x))
        # Tensor(shape=[1], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [1.45936954])

        # outer
        print(paddle.einsum("i,j->ij", x, y))
        # Tensor(shape=[4, 5], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [[0.00079869, 0.00120950, 0.00136844, 0.00187187, 0.00192194],
        #    [0.23455200, 0.35519385, 0.40186870, 0.54970956, 0.56441545],
        #    [0.11773264, 0.17828843, 0.20171674, 0.27592498, 0.28330654],
        #    [0.32897076, 0.49817693, 0.56364071, 0.77099484, 0.79162055]])

        A = paddle.rand([2, 3, 2])
        B = paddle.rand([2, 2, 3])

        # transpose
        print(paddle.einsum('ijk->kji', A))
        #  Tensor(shape=[2, 3, 2], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [[[0.95649719, 0.49684682],
        #     [0.80071914, 0.46258664],
        #     [0.49814570, 0.33383518]],
        #
        #    [[0.07637714, 0.29374704],
        #     [0.51470858, 0.51907635],
        #     [0.99066722, 0.55802226]]])

        # batch matrix multiplication
        print(paddle.einsum('ijk, ikl->ijl', A,B))
        # Tensor(shape=[2, 3, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [[[0.32172769, 0.50617385, 0.41394392],
        #     [0.51736701, 0.49921003, 0.38730967],
        #     [0.69078457, 0.42282537, 0.30161136]],
        #
        #    [[0.32043904, 0.18164253, 0.27810261],
        #     [0.50226176, 0.24512935, 0.39881429],
        #     [0.51476848, 0.23367381, 0.39229113]]])

        # Ellipsis transpose
        print(paddle.einsum('...jk->...kj', A))
        # Tensor(shape=[2, 2, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [[[0.95649719, 0.80071914, 0.49814570],
        #     [0.07637714, 0.51470858, 0.99066722]],
        #
        #    [[0.49684682, 0.46258664, 0.33383518],
        #     [0.29374704, 0.51907635, 0.55802226]]])

        # Ellipsis batch matrix multiplication
        print(paddle.einsum('...jk, ...kl->...jl', A,B))
        # Tensor(shape=[2, 3, 3], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
        #   [[[0.32172769, 0.50617385, 0.41394392],
        #     [0.51736701, 0.49921003, 0.38730967],
        #     [0.69078457, 0.42282537, 0.30161136]],
        #
        #    [[0.32043904, 0.18164253, 0.27810261],
        #     [0.50226176, 0.24512935, 0.39881429],
        #     [0.51476848, 0.23367381, 0.39229113]]])
    """
    import os
    if int(os.environ.get('FLAGS_new_einsum', "1")):
        return einsum_v2(equation, *operands)

    nop = len(operands)
    assert nop > 0, "At least one operand is expected."

    # Part the equation to left hand side and right hand side
    lhs, *rhs = equation.lower().replace(' ', '').split('->')
    assert len(rhs) < 2, "Invalid equation: multiple `->` were found."

    # Note, we distinguish between 'ij->' and 'ij' by setting rhs to '' and None
    rhs = rhs[0] if rhs else None

    # Parse labels for each operand and count the number of occurrences for each alphabet label
    nop_labels = parse_labels(lhs, operands)

    # Diagonalize the operands which have duplicate labels
    nop_labels, operands = list(zip(*map(diagonalize, nop_labels, operands)))

    # To handle broadcasting, we should first know how many dimensions are there
    # We need to use that number to generate output labels
    # e.g. 1 for ['ij', 'i.', '.k']
    n_bcast_dims = max(map(lambda s: s.count('.'), nop_labels))

    # Build the data structures for planning. It's helpful to think of all the operands
    # broadcasting together from a global view. In this view, dimensions from multiple
    # operands are mapped to the same position if they are labeled uniquely. Broadcasting
    # dimensions are mapped to adjacent positions with the right bound fixed. Subject to
    # each operand, the map is injective but for all operands the map is on-to.
    # g_labels:
    #   The labels of the global view
    # g_view:
    #   Includes a list of maps from each operand's dimensions to the global view's dimensions
    #   which we refer to as ax or axes in the code to distinguish from operand's dims
    # g_shape:
    #   The shape of the global view. The size of each dimension is what the aligned dimensions
    #   should broadcast to
    # g_nout:
    #   Number of output axes
    # g_supports
    #   Booleans indicating each operand's non-trivial dimensions
    # g_count
    #   Counting how many non-trivial dimensions remain for each ax

    g_labels, g_view, g_nout, g_count = build_global_view(nop_labels, rhs,
                                                          n_bcast_dims)
    g_shape, g_supports = build_global_shape(g_view, g_labels,
                                             [op.shape for op in operands])

    # Now we're ready to build up an execution plan
    args = operands, g_view, g_shape, g_supports, g_count, n_bcast_dims
    plan = plan_einsum(*args)
    result = plan.execute()

    return result


#!/usr/bin/env python
# Copyright 2018-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Pseudo-spectral methods (COSX, PS, SN-K)
'''

import copy
import numpy
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import mcscf
from pyscf.scf import _vhf
from pyscf.lib import logger
from pyscf.sgx import sgx_jk
from pyscf.df import df_jk
from pyscf import __config__

def sgx_fit(mf, auxbasis=None, with_df=None):
    '''For the given SCF object, update the J, K matrix constructor with
    corresponding SGX or density fitting integrals.

    Args:
        mf : an SCF object

    Kwargs:
        auxbasis : str or basis dict
            Same format to the input attribute mol.basis.  If auxbasis is
            None, optimal auxiliary basis based on AO basis (if possible) or
            even-tempered Gaussian basis will be used.

    Returns:
        An SCF object with a modified J, K matrix constructor which uses density
        fitting integrals to compute J and K

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = sgx_fit(scf.RHF(mol))
    >>> mf.scf()
    -100.00978770917165

    >>> mol.symmetry = 1
    >>> mol.build(0, 0)
    >>> mf = sgx_fit(scf.UHF(mol))
    >>> mf.scf()
    -100.00978770951018
    '''
    assert(isinstance(mf, scf.hf.SCF))

    if with_df is None:
        with_df = SGX(mf.mol)
        with_df.max_memory = mf.max_memory
        with_df.stdout = mf.stdout
        with_df.verbose = mf.verbose
        with_df.auxbasis = auxbasis

    mf_class = mf.__class__

    if isinstance(mf, _SGXHF):
        if mf.with_df is None:
            mf = mf_class(mf, with_df, auxbasis)
        elif mf.with_df.auxbasis != auxbasis:
            #logger.warn(mf, 'DF might have been initialized twice.')
            mf = copy.copy(mf)
            mf.with_df = with_df
        return mf

    class SGXHF(_SGXHF, mf_class):
        def __init__(self, mf, df, auxbasis):
            self.__dict__.update(mf.__dict__)
            self._eri = None
            self.auxbasis = auxbasis
            self.with_df = df

            # Grids/Integral quality varies during SCF. VHF cannot be
            # constructed incrementally.
            self.direct_scf = False

            self._last_dm = 0
            self._in_scf = False
            self._keys = self._keys.union(['auxbasis', 'with_df'])

        def build(self, mol=None, **kwargs):
            if self.direct_scf:
                self.with_df.build(level=self.with_df.grids_level_f)
            else:
                self.with_df.build(level=self.with_df.grids_level_i)
            return mf_class.build(self, mol, **kwargs)

        def reset(self, mol=None):
            self.with_df.reset(mol)
            return mf_class.reset(self, mol)

        def pre_kernel(self, envs):
            self._in_scf = True

        def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
                  omega=None):
            if dm is None: dm = self.make_rdm1()
            with_df = self.with_df
            if not with_df:
                return mf_class.get_jk(self, mol, dm, hermi, with_j, with_k, omega)

            if self._in_scf and not self.direct_scf:
                if numpy.linalg.norm(dm - self._last_dm) < with_df.grids_switch_thrd:
                    logger.debug(self, 'Switching SGX grids')
                    with_df.build(level=with_df.grids_level_f)
                    self._in_scf = False
                    self._last_dm = 0
                else:
                    self._last_dm = numpy.asarray(dm)

            return with_df.get_jk(dm, hermi, with_j, with_k,
                                  self.direct_scf_tol, omega)

        def post_kernel(self, envs):
            self._in_scf = False
            self._last_dm = 0

        def nuc_grad_method(self):
            raise NotImplementedError

    return SGXHF(mf, with_df, auxbasis)

# A tag to label the derived SCF class
class _SGXHF(object):
    def method_not_implemented(self, *args, **kwargs):
        raise NotImplementedError
    nuc_grad_method = Gradients = method_not_implemented
    Hessian = method_not_implemented
    NMR = method_not_implemented
    NSR = method_not_implemented
    Polarizability = method_not_implemented
    RotationalGTensor = method_not_implemented
    MP2 = method_not_implemented
    CISD = method_not_implemented
    CCSD = method_not_implemented
    CASCI = method_not_implemented
    CASSCF = method_not_implemented

scf.hf.SCF.COSX = sgx_fit
mcscf.casci.CASCI.COSX = sgx_fit

def _make_opt(mol):
    '''Optimizer to genrate 3-center 2-electron integrals'''
    intor = mol._add_suffix('int3c2e')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor)
    # intor 'int1e_ovlp' is used by the prescreen method
    # 'SGXnr_ovlp_prescreen' only. Not used again in other places.
    # It can be released early
    vhfopt = _vhf.VHFOpt(mol, 'int1e_ovlp', 'SGXnr_ovlp_prescreen',
                         'SGXsetnr_direct_scf')
    vhfopt._intor = intor
    vhfopt._cintopt = cintopt
    return vhfopt

class SGX(lib.StreamObject):
    def __init__(self, mol, auxbasis=None):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.max_memory = mol.max_memory
        self.grids_thrd = 1e-10
        self.grids_level_i = 0  # initial grids level
        self.grids_level_f = 1  # final grids level
        self.grids_switch_thrd = 0.03
        # compute J matrix using DF and K matrix using SGX. It's identical to
        # the RIJCOSX method in ORCA
        self.dfj = False
        self._auxbasis = auxbasis

        # debug=True generates a dense tensor of the Coulomb integrals at each
        # grids. debug=False utilizes the sparsity of the integral tensor and
        # contracts the sparse tensor and density matrices on the fly.
        self.debug = False

        self.grids = None
        self.blockdim = 1200
        self.auxmol = None
        self._vjopt = None
        self._opt = None
        self._last_dm = 0
        self._rsh_df = {}  # Range separated Coulomb DF objects
        self._keys = set(self.__dict__.keys())

    @property
    def auxbasis(self):
        return self._auxbasis
    @auxbasis.setter
    def auxbasis(self, x):
        if self._auxbasis != x:
            self._auxbasis = x
            self.auxmol = None

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('******** %s ********', self.__class__)
        log.info('max_memory = %s', self.max_memory)
        log.info('grids_level_i = %s', self.grids_level_i)
        log.info('grids_level_f = %s', self.grids_level_f)
        log.info('grids_thrd = %s', self.grids_thrd)
        log.info('grids_switch_thrd = %s', self.grids_switch_thrd)
        log.info('df_j = %s', self.df_j)
        log.info('auxbasis = %s', self.auxbasis)
        return self

    # To mimic DF object, so that SGX can be used as in DF-SCF method by setting
    # mf.with_df = SGX(mol)
    @property
    def _cderi(self):
        return self.grids

    def build(self, level=None):
        if level is None:
            level = self.grids_level_f
        self.grids = sgx_jk.get_gridss(self.mol, level, self.grids_thrd)
        self._opt = _make_opt(self.mol)

        # In the RSH-integral temporary treatment, recursively rebuild SGX
        # objects in _rsh_df.
        if self._rsh_df:
            for k, v in self._rsh_df.items():
                v.build(level)
        return self

    def kernel(self, *args, **kwargs):
        return self.build(*args, **kwargs)

    def reset(self, mol=None):
        '''Reset mol and clean up relevant attributes for scanner mode'''
        if mol is not None:
            self.mol = mol
        self.grids = None
        self.auxmol = None
        self._vjopt = None
        self._opt = None
        self._last_dm = 0
        self._rsh_df = {}
        return self

    def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
               direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
               omega=None):
        if omega is not None:
            # A temporary treatment for RSH integrals
            key = '%.6f' % omega
            if key in self._rsh_df:
                rsh_df = self._rsh_df[key]
            else:
                rsh_df = copy.copy(self)
                rsh_df._rsh_df = None  # to avoid circular reference
                # Not all attributes need to be reset. Resetting _vjopt
                # because it is used by get_j method of regular DF object.
                rsh_df._vjopt = None
                self._rsh_df[key] = rsh_df
                logger.info(self, 'Create RSH-SGX object %s for omega=%s', rsh_df, omega)

            with rsh_df.mol.with_range_coulomb(omega):
                return rsh_df.get_jk(dm, hermi, with_j, with_k,
                                     direct_scf_tol)

        if with_j and self.dfj:
            vj = df_jk.get_j(self, dm, hermi, direct_scf_tol)
            if with_k:
                vk = sgx_jk.get_jk(self, dm, hermi, False, with_k, direct_scf_tol)[1]
            else:
                vk = None
        else:
            vj, vk = sgx_jk.get_jk(self, dm, hermi, with_j, with_k, direct_scf_tol)
        return vj, vk

if __name__ == '__main__':
    from pyscf import scf
    mol = gto.Mole()
    mol.build(
        atom = [["O" , (0. , 0.     , 0.)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)] ],
        basis = 'ccpvdz',
    )
    method = sgx_fit(scf.RHF(mol), 'weigend')
    energy = method.scf()
    print(energy - -76.02673747045691)

    method.with_df.dfj = True
    energy = method.scf()
    print(energy - -76.02686422219752)