# --------------
# import packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import seaborn as sns
import re
from nltk.corpus import stopwords
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.multiclass import OneVsRestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score ,confusion_matrix

# Code starts here

# load data
news = pd.read_csv(path)

# subset data
news = news[['TITLE','CATEGORY']]
# distribution of classes
dist = news['CATEGORY'].value_counts()

# display class distribution
print(dist.head())

# display data
print(news.head())

# Code ends here

# --------------
# Code starts here

# stopwords
stop = (set(stopwords.words('english')))
# retain only alphabets
news['TITLE'] = news['TITLE'].apply(lambda x : re.sub("[^a-zA-Z]", " ",x) )

# convert to lowercase and tokenize
news['TITLE'] = news['TITLE'].apply(lambda x : x.lower().split())

# remove stopwords
news['TITLE'] = news['TITLE'].apply(lambda x : [i for i in x if i not in stop])

# join list elements
print(news['TITLE'].head(2))
news['TITLE'] = news['TITLE'].apply(lambda x : ' '.join(x))
print(news['TITLE'].head(2))
# split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(news['TITLE'], news['CATEGORY'], test_size=0.2, random_state=3)

# Code ends here

# --------------
# Code starts here

# initialize count vectorizer
count_vectorizer = CountVectorizer()

# initialize tfidf vectorizer
tfidf_vectorizer = TfidfVectorizer(ngram_range=(1,3))

# fit and transform with count vectorizer
X_train_count= count_vectorizer.fit_transform(X_train)
X_test_count = count_vectorizer.transform(X_test)

# fit and transform with tfidf vectorizer
X_train_tfidf= tfidf_vectorizer.fit_transform(X_train)
X_test_tfidf = tfidf_vectorizer.transform(X_test)
# Code ends here

# --------------
# Code starts here

# initialize multinomial naive bayes
nb_1 = MultinomialNB()
nb_2 = MultinomialNB()

# fit on count vectorizer training data
nb_1.fit(X_train_count, y_train)
# fit on tfidf vectorizer training data
nb_2.fit(X_train_tfidf, y_train)

# accuracy with count vectorizer
acc_count_nb = accuracy_score(nb_1.predict(X_test_count), y_test)

# accuracy with tfidf vectorizer
acc_tfidf_nb = accuracy_score(nb_2.predict(X_test_tfidf), y_test)

# display accuracies
print('Count Vectorizer accuracy is', acc_count_nb)
print('TFIDF accuracy is', acc_tfidf_nb)

# Code ends here

# --------------
import warnings
warnings.filterwarnings('ignore')

# initialize logistic regression
logreg_1 = OneVsRestClassifier(LogisticRegression(random_state=10))
logreg_2 = OneVsRestClassifier(LogisticRegression(random_state=10))

# fit on count vectorizer training data
logreg_1.fit(X_train_count, y_train)
logreg_2.fit(X_train_tfidf, y_train)
# fit on tfidf vectorizer training data
acc_count_logreg = accuracy_score(logreg_1.predict(X_test_count), y_test)
acc_tfidf_logreg = accuracy_score(logreg_2.predict(X_test_tfidf), y_test)
# accuracy with count vectorizer
print('Count vectorizer accurancy is', acc_count_logreg)
# accuracy with tfidf vectorizer
print('TFIDF accuracy is', acc_tfidf_logreg)

# display accuracies

# Code ends here


from bisect import bisect_right
from itertools import accumulate
from math import inf, sqrt
from numbers import Number

class ApproximateHistogram:
    """
    Streaming, approximate histogram

    Based on http://jmlr.org/papers/volume11/ben-haim10a/ben-haim10a.pdf

    Performance of adding a point is about 5x faster than
    https://github.com/carsonfarmer/streamhist (unmaintained).

    The output of quantile() will match numpy.quantile() exactly until
    the number of points reaches max_bins, and then gracefully transition
    to an approximation.
    """

    def __init__(self, max_bins):
        self._max_bins = max_bins
        self._bins = []  # (point, count)
        self._costs = []  # item i is _bins[i+1].point - _bins[i].point
        self._count = 0
        # TODO: maintain min/max as bin entries with infinite merge cost
        self._min = inf
        self._max = -inf

    @staticmethod
    def _update_costs(costs, l, i, val):
        """update costs array to reflect l.insert(i, val)"""
        if i > 0:
            new_cost = val[0] - l[i - 1][0]
            costs.insert(i - 1, new_cost)
            if i < len(costs):
                costs[i] = l[i + 1][0] - val[0]
        elif len(l) > 1:
            costs.insert(0, l[1][0] - val[0])
        # assert costs == approx([b - a for (a, _), (b, _) in zip(l, l[1:])], rel=1e-4)

    @staticmethod
    def _update_costs_for_merge(costs, l, i, val):
        """update costs array to reflect l[i:i+2] = (val, )"""
        # TODO: combine with update_costs()
        if 0 < i < len(costs) - 1:
            costs[i - 1:i + 2] = val[0] - l[i - 1][0], l[i + 1][0] - val[0]
        elif i > 0:
            costs[i - 1:i + 1] = (val[0] - l[i - 1][0], )
        else:
            costs[i:i + 2] = (l[i + 1][0] - val[0], )
        # assert costs == approx([b - a for (a, _), (b, _) in zip(l, l[1:])], rel=1e-4)

    @classmethod
    def _insert_with_cost(cls, costs, l, val):
        i = bisect_right(l, val)
        l.insert(i, val)
        cls._update_costs(costs, l, i, val)

    def add(self, point):
        """Add point to histogram"""
        # optimization:  maintain cost array
        self._count += 1
        self._min = min(self._min, point)
        self._max = max(self._max, point)
        bins = self._bins
        costs = self._costs
        self._insert_with_cost(costs, bins, (point, 1))
        if len(bins) > self._max_bins:
            i = costs.index(min(costs))
            (q0, k0), (q1, k1) = bins[i:i+2]
            _count = k0 + k1
            median = (q0 * k0 + q1 * k1) / _count
            bins[i:i+2] = ((median, _count), )
            self._update_costs_for_merge(costs, bins, i, (median, _count))

    @property
    def count(self):
        """Return number of points represented by this histogram."""
        return self._count

    @property
    def min(self):
        """Return minimum point represented by this histogram"""
        return self._min

    @property
    def max(self):
        """Return maximum point represented by this histogram"""
        return self._max

    def mean(self):
        """Return mean;  O(max_bins) complexity."""
        return sum(p * count for p, count in self._bins) / self._count

    def std(self):
        """Return standard deviation;  O(max_bins) complexity."""
        mean = self.mean()
        sum_squares = sum((p - mean) ** 2 * count for p, count in self._bins)
        return sqrt(sum_squares / self._count)

    def _quantile(self, sums, q):
        if q <= 0:
            return self._min
        if q >= 1:
            return self._max
        bins = self._bins
        target_sum = q * (self._count - 1) + 1
        i = bisect_right(sums, target_sum) - 1
        left = bins[i] if i >= 0 else (self._min, 0)
        right = bins[i+1] if i+1 < len(bins) else (self._max, 0)
        l0, r0 = left[0], right[0]
        l1, r1 = left[1], right[1]
        s = target_sum - (sums[i] if i >= 0 else 1)
        if l1 <= 1 and r1 <= 1:
            # We have exact info at this quantile.  Match linear interpolation
            # strategy of numpy.quantile().
            b = l0 + (r0 - l0) * s / r1 if r1 > 0 else l0
        else:
            if r1 == 1:
                # For exact bin on RHS, compensate for trapezoid interpolation using
                # only half of count.
                r1 = 2
            if l1 == r1:
                bp_ratio = s / l1
            else:
                bp_ratio = (l1 - (l1 ** 2 - 2 * s * (l1 - r1)) ** .5) / (l1 - r1)
                assert bp_ratio.imag == 0
            b = bp_ratio * (r0 - l0) + l0
        return b

    def sum(self):
        """Return sum of points;  O(max_bins) complexity."""
        return sum(x * count for x, count in self._bins)

    def quantile(self, q):
        """Return list of values at given quantile fraction(s);  O(max_bins) complexity."""
        # Deviation from Ben-Haim sum strategy:
        #     * treat count 1 bins as "exact" rather than dividing the count at the point
        #     * for neighboring exact bins, use simple linear interpolation matching
        #       numpy.quantile()
        if isinstance(q, Number):
            q = (q, )
        bins = self._bins
        sums = [x - (y/2 if y > 1 else 0) for x, (_, y) in \
                zip(accumulate(bin[1] for bin in bins), bins)]
        return list(self._quantile(sums, q_item) for q_item in q)


import numpy as np

def scroll(clip, h=None, w=None, x_speed=0, y_speed=0,
           x_start=0, y_start=0, apply_to="mask"):
    """ Scrolls horizontally or vertically a clip, e.g. to make fin
        credits """
    if h is None: h = clip.h
    if w is None: w = clip.w

    xmax = clip.w-w-1
    ymax = clip.h-h-1

    def f(gf,t):
        x = max(0, min(xmax, x_start+ np.round(x_speed*t)))
        y = max(0, min(ymax, y_start+ np.round(y_speed*t)))
        return gf(t)[y:y+h, x:x+w]

    return clip.fl(f, apply_to = apply_to)


"""Non-Maximum Suppression module."""
import numpy as np
import torch

def nms(detections, threshold):
    """Apply Non-Maximum Suppression over the detections.
    The detections must be a tensor with two dimensions: (number of detections, 5).
    Why 5? Because a detection has x1, y1, x2, y2 and score.

    Heavily inspired by Adrian Rosebrock at:
    https://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/

    Why not the version of GPU? Because I couldn't make it work in my GPU.

    Args:
        detections (torch.Tensor): A tensor with all the detections. The shape must be
            (number of detections, 5) with the score as the last value of the second
            dimension.
        threshold (float): The threshold for the IoU (intersection over union) to take
            two detections as detecting the same object.

    Returns:
        torch.Tensor: A tensor with the indexes of the detections to keep.
    """
    # If there aren't detections return empty
    if detections.shape[0] == 0:
        return torch.zeros((0))

    # Get the numpy version
    was_cuda = detections.is_cuda
    detections = detections.cpu().numpy()

    # Start the picked indexes list empty
    picked = []

    # Get the coordinates
    x1 = detections[:, 0]
    y1 = detections[:, 1]
    x2 = detections[:, 2]
    y2 = detections[:, 3]
    scores = detections[:, 4]

    # Compute the area of the bounding boxes
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    # Get the indexes of the detections sorted by score
    indexes = np.argsort(scores)

    while len(indexes) > 0:
        # Take the last index (highest score) and add it to the picked
        last = len(indexes) - 1
        actual = indexes[last]
        picked.append(actual)

        # We need to find the overlap of the bounding boxes with the actual picked bounding box

        # Find the largest (more to the bottom-right) (x,y) coordinates for the start
        # (top-left) of the bounding box
        xx1 = np.maximum(x1[actual], x1[indexes[:last]])
        yy1 = np.maximum(y1[actual], y1[indexes[:last]])
        # Find the smallest (more to the top-left) (x,y) coordinates for the end (bottom-right)
        # of the bounding box
        xx2 = np.minimum(x2[actual], x2[indexes[:last]])
        yy2 = np.minimum(y2[actual], y2[indexes[:last]])

        # Compute width and height to compute the intersection over union
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)
        intersection = (w * h)
        union = areas[actual] + areas[indexes[:last]] - intersection
        iou = intersection / union

        # Delete the last index and all that overlap is bigger than threshold
        indexes = np.delete(indexes, np.concatenate(([last], np.where(iou > threshold)[0])))

    # Return the picked indexes
    picked = torch.Tensor(picked).long()
    if was_cuda:
        picked = picked.cuda()

    return picked


import os
from functools import partial
import numpy as np
import pandas as pd
import tables
import matplotlib
import warnings

from PyQt5.QtCore import Qt, QPointF
from PyQt5.QtGui import QPixmap, QPainter, QFont, QPen, QPolygonF, QColor, QKeySequence, QBrush
from PyQt5.QtWidgets import QApplication, QMessageBox

from tierpsy.analysis.ske_create.helperIterROI import getWormROI
from tierpsy.analysis.split_fov.FOVMultiWellsSplitter import FOVMultiWellsSplitter

from tierpsy.gui.MWTrackerViewer_ui import Ui_MWTrackerViewer
from tierpsy.gui.TrackerViewerAux import TrackerViewerAuxGUI
from tierpsy.gui.PlotFeatures import PlotFeatures

from tierpsy.helper.misc import WLAB, save_modified_table
from tierpsy.analysis.split_fov.helper import get_well_color

class WellsDrawer(TrackerViewerAuxGUI):
    '''
    Dummy class with the wells division drawer functions
    '''
    def __init__(self, ui):
        super().__init__(ui)
        # colour
        self.fovsplitter_mask = None
        self.fovsplitter_feat = None
        self.fovsplitter = None
        self.is_fov_tosplit = None

    def updateVideoFile(self, vfilename):
        super().updateVideoFile(vfilename)
        # check if /fov_wells exists in masked video
        if self.fid is not None:
            if '/fov_wells' not in self.fid:
                self.is_fov_tosplit = False
            else:
                self.is_fov_tosplit = True
            # if it exists, read it
            if self.is_fov_tosplit:
#                self.wells_in_mask = pd.DataFrame(
#                        self.fid.get_node('/fov_wells').read())
                self.fovsplitter_mask = FOVMultiWellsSplitter(self.vfilename)

    def updateSkelFile(self, skeletons_file):
        super().updateSkelFile(skeletons_file)
        # if no skeletons, skip
        if not self.skeletons_file:
            return
        # check if /fov_wells exists in features video
        with tables.File(self.skeletons_file, 'r') as fid:
            if '/fov_wells' not in fid:
                self.is_fov_tosplit = False
#                print("didn't find fov wells though")
            else:
                self.is_fov_tosplit = True
#                print("found fov wells in featuresN")
            # if it exists, read it
            if self.is_fov_tosplit:
#                print('reading fov_wells from featuresN')
#                print('pre-reading:')
#                print(self.wells)
#                self.wells_in_feat = pd.DataFrame(
#                        fid.get_node('/fov_wells').read())
                self.fovsplitter_feat = FOVMultiWellsSplitter(self.skeletons_file)

    def draw_wells(self, image):
        '''
        Draw wells.
        '''
        if self.is_fov_tosplit:
            if self.fovsplitter_feat is not None:
                self.fovsplitter = self.fovsplitter_feat
            else: # fall back to mask ones
                print('falling back')
                self.fovsplitter = self.fovsplitter_mask

            # prepare constants for drawing
            self.fontsize = max(1, max(image.height(), image.width()) // 60)
            penwidth = max(1, max(image.height(), image.width()) // 400)
            self.penwidth = penwidth if penwidth % 2 == 1 else penwidth + 1
#            self.wellsC = QColor(250, 140, 0)
            if 'is_good_well' in self.fovsplitter.wells.columns:
                is_color_by_well = True
            else:
                is_color_by_well = False
            # Qt drawing code
            painter = QPainter()
            painter.begin(image)
            pen = QPen()
            pen.setWidth(self.penwidth)
            painter.setFont(QFont('Decorative', self.fontsize))
            # loop on wells
            for _, well in self.fovsplitter.wells.iterrows():
                # update color every time
                if is_color_by_well:
                    wellC = get_well_color(well['is_good_well'], forCV=True)
                    wellC = QColor(*wellC)
                else:
                    wellC = QColor(250, 140, 0)
                pen.setColor(wellC)
                painter.setPen(pen)
                # draw well name
                painter.drawText(well['x_min'] + self.fontsize*0.4,
                                 well['y_min'] + self.fontsize*1.2,
                                 well['well_name'])
                # draw rectangle
                painter.drawRect(well['x_min'],
                                 well['y_min'],
                                 well['x_max'] - well['x_min'],
                                 well['y_max'] - well['y_min'])
                if well['is_good_well'] == False:
                    painter.drawLine(well['x_min'],
                                     well['y_min'],
                                     well['x_max'],
                                     well['y_max'])
            painter.end()

#        super().keyPressEvent(event)

class ContourDrawer(TrackerViewerAuxGUI):
    '''
    Dummy class with the contour functions
    '''
    def __init__(self, ui):
        super().__init__(ui)

        self.food_coordinates = None
        self.wlabC = {
            WLAB['U']: Qt.white,
            WLAB['WORM']: Qt.green,
            WLAB['WORMS']: Qt.blue,
            WLAB['BAD']: Qt.darkRed,
            WLAB['GOOD_SKE']: Qt.darkCyan
            }
        self.ui.checkBox_showFood.stateChanged.connect(self.updateImage)
        self.ui.checkBox_showFood.setEnabled(False)
        self.ui.checkBox_showFood.setChecked(True)

    def updateSkelFile(self, skeletons_file):

        super().updateSkelFile(skeletons_file)
        if not self.skeletons_file or self.trajectories_data is None:
            self.food_coordinates = None
            return

        with tables.File(self.skeletons_file, 'r') as fid:
            if not '/food_cnt_coord' in fid:
                self.food_coordinates = None
                self.ui.checkBox_showFood.setEnabled(False)
            else:
                #change from microns to pixels
                self.food_coordinates = fid.get_node('/food_cnt_coord')[:]
                self.food_coordinates /= self.microns_per_pixel

                self.ui.checkBox_showFood.setEnabled(True)

    def draw_food_contour(self, image):
        if self.food_coordinates is None or not self.ui.checkBox_showFood.isChecked():
            return

        painter = QPainter()
        painter.begin(image)

        penwidth = max(1, max(image.height(), image.width()) // 800)
        col = Qt.darkMagenta
        p = QPolygonF()
        for x,y in self.food_coordinates:
            p.append(QPointF(x,y))

        pen = QPen()
        pen.setWidth(penwidth)
        pen.setColor(col)
        painter.setPen(pen)

        painter.drawPolyline(p)
        painter.end()

class IntensityLabeler(TrackerViewerAuxGUI):
    def __init__(self, ui):
        super().__init__(ui)

        self.mean_intensity = None
        self.ui.intensity_label.setStyleSheet('') #avoid displaying color at the start of the programı

    def updateVideoFile(self, vfilename):
        super().updateVideoFile(vfilename)
        if self.fid is not None:
            #get mean intensity information.
            #Useful for the optogenetic experiments.
            try:
                mean_int = self.fid.get_node('/mean_intensity')[:]

                #calculate the intensity range and normalize the data.
                #I am ignoring any value less than 1. The viewer only works with uint8 data.

                dd = mean_int[mean_int>=1]
                if dd.size == 0:
                    raise ValueError

                bot = np.min(dd)
                top = np.max(dd)
                rr = top-bot

                # if the mean value change is less than 1 (likely continous image do nothing)
                if rr <= 1:
                    raise ValueError

                self.mean_intensity = (mean_int-bot)/(rr)

            except (tables.exceptions.NoSuchNodeError, ValueError):
                self.mean_intensity = None
                self.ui.intensity_label.setStyleSheet('')

    def display_intensity(self):
        if self.mean_intensity is not None and self.frame_number < self.mean_intensity.size:
            d = int(self.mean_intensity[self.frame_number]*255)
            self.ui.intensity_label.setStyleSheet('QLabel {background-color: rgb(%i, %i, %i);}' % (0, 0, d))

class BlobLabeler(TrackerViewerAuxGUI):
    def __init__(self, ui):
        super().__init__(ui)
        self.wlab = WLAB
        self.label_type = 'worm_label'

        self.ui.pushButton_U.clicked.connect(
            partial(self._h_tag_worm, self.wlab['U']))
        self.ui.pushButton_W.clicked.connect(
            partial(self._h_tag_worm, self.wlab['WORM']))
        self.ui.pushButton_WS.clicked.connect(
            partial(self._h_tag_worm, self.wlab['WORMS']))
        self.ui.pushButton_B.clicked.connect(
            partial(self._h_tag_worm, self.wlab['BAD']))

        self.ui.pushButton_W.setShortcut(QKeySequence(Qt.Key_W))
        self.ui.pushButton_U.setShortcut(QKeySequence(Qt.Key_U))
        self.ui.pushButton_WS.setShortcut(QKeySequence(Qt.Key_C))
        self.ui.pushButton_B.setShortcut(QKeySequence(Qt.Key_B))

    def enable_label_buttons(self, value):
        self.ui.pushButton_U.setEnabled(value)
        self.ui.pushButton_W.setEnabled(value)
        self.ui.pushButton_WS.setEnabled(value)
        self.ui.pushButton_B.setEnabled(value)

    def _h_tag_worm(self, label_ind):
        if not self.worm_index_type == 'worm_index_manual':
            return

        worm_ind = self.current_worm_index

        if self.frame_data is None:
            return

        if not worm_ind in self.frame_data['worm_index_manual'].values:
            QMessageBox.critical(
                self,
                'The selected worm is not in this frame.',
                'Select a worm in the current frame to label.',
                QMessageBox.Ok)
            return

        good = self.trajectories_data['worm_index_manual'] == worm_ind
        self.trajectories_data.loc[good, 'worm_label'] = label_ind
        self.updateImage()

class ROIWorm():
    def __init__(self, wormCanvas, comboBox_ROI, checkBox_ROI):
        self.worm_index = None
        self.wormCanvas = wormCanvas
        self.comboBox_ROI = comboBox_ROI
        self.checkBox_ROI = checkBox_ROI

        self.comboBox_ROI.activated.connect(self.selectROI)

    def selectROI(self, index):
        try:
            self.worm_index = int(self.comboBox_ROI.itemText(index))
        except ValueError:
            self.worm_index = None

    @property
    def isDrawSkel(self):
        return self.checkBox_ROI.isChecked()

class ROIManager(TrackerViewerAuxGUI):
    def __init__(self, ui):

        super().__init__(ui)
        self.rois = [
            ROIWorm(
                self.ui.wormCanvas1,
                self.ui.comboBox_ROI1,
                self.ui.checkBox_ROI1
                ),
            ROIWorm(
                self.ui.wormCanvas2,
                self.ui.comboBox_ROI2,
                self.ui.checkBox_ROI2
                )
        ]

        self.ui.radioButton_ROI1.setShortcut(QKeySequence(Qt.Key_Up))
        self.ui.radioButton_ROI2.setShortcut(QKeySequence(Qt.Key_Down))

        self.ui.checkBox_ROI1.stateChanged.connect(partial(self._updateROI, self.rois[0]))
        self.ui.checkBox_ROI2.stateChanged.connect(partial(self._updateROI, self.rois[1]))

        self.ui.comboBox_ROI1.activated.connect(partial(self._updateROI, self.rois[0]))
        self.ui.comboBox_ROI2.activated.connect(partial(self._updateROI, self.rois[1]))

        # flags for RW and FF
        self.RW, self.FF = 1, 2
        self.ui.pushButton_ROI1_RW.clicked.connect(partial(self.roiRWFF, self.RW, self.rois[0]))
        self.ui.pushButton_ROI1_FF.clicked.connect(partial(self.roiRWFF, self.FF, self.rois[0]))
        self.ui.pushButton_ROI2_RW.clicked.connect(partial(self.roiRWFF, self.RW, self.rois[1]))
        self.ui.pushButton_ROI2_FF.clicked.connect(partial(self.roiRWFF, self.FF, self.rois[1]))

    @property
    def current_roi(self):
        if self.ui.radioButton_ROI1.isChecked():
            return self.rois[0]
        elif self.ui.radioButton_ROI2.isChecked():
            return self.rois[1]
        else:
            raise ValueError("I shouldn't be here")

    @property
    def current_worm_index(self):
        return self.current_roi.worm_index

    def updateSkelFile(self, skeletons_file):
        for roi in self.rois:
            roi.worm_index = None
        super().updateSkelFile(skeletons_file)

    def keyPressEvent(self, event):
        #MORE SHORTCUTS
        # go the the start of end of a trajectory
        if event.key() == Qt.Key_BracketLeft:

            self.roiRWFF(self.RW, self.current_roi)

        elif event.key() == Qt.Key_BracketRight:

            self.roiRWFF(self.FF, self.current_roi)

        super().keyPressEvent(event)

    def updateROIcomboBox(self, roi):
        # update valid index for the comboBox
        roi.comboBox_ROI.clear()

        if roi.worm_index is not None:
            roi.comboBox_ROI.addItem(str(int(roi.worm_index)))

        for ind in self.frame_data[self.worm_index_type]:
            roi.comboBox_ROI.addItem(str(int(ind)))

        if roi.worm_index is None:
            w_ind = float(roi.comboBox_ROI.itemText(0))
            roi.worm_index = int(w_ind)

    # function that generalized the updating of the ROI
    def _updateROI(self, roi):

        if self.frame_data is None or not self.worm_index_type:
            # no trajectories data presented, nothing to do here
            roi.wormCanvas.clear()
            return

        self.updateROIcomboBox(roi)

        # extract individual worm ROI
        good = self.frame_data[self.worm_index_type] == roi.worm_index
        row_data = self.frame_data.loc[good].squeeze()

        if row_data.size == 0 or \
        np.isnan(row_data['coord_x']) or \
        np.isnan(row_data['coord_y']):
            # invalid data nothing to do here
            roi.wormCanvas.clear()
            return

        worm_img, roi_corner = getWormROI(self.frame_img,
                                        row_data['coord_x'],
                                        row_data['coord_y'],
                                        row_data['roi_size']
                                        )

        roi_ori_size = worm_img.shape
        worm_img = np.ascontiguousarray(worm_img)
        worm_qimg = self._convert2Qimg(worm_img)

        canvas_size = min(roi.wormCanvas.height(), roi.wormCanvas.width())
        worm_qimg = worm_qimg.scaled(
            canvas_size, canvas_size, Qt.KeepAspectRatio)

        worm_qimg = self.drawSkelResult(worm_img, worm_qimg, row_data, roi.isDrawSkel, roi_corner, read_center=False)

        pixmap = QPixmap.fromImage(worm_qimg)
        roi.wormCanvas.setPixmap(pixmap)

    def updateROIs(self):
        for roi in self.rois:
            self._updateROI(roi)

    def clearROIs(self):
        for roi in self.rois:
            roi.wormCanvas.clear()

    # move to the first or the last frames of a trajectory
    def roiRWFF(self, rwff, roi):

        if self.frame_data is None:
            return

        # use 1 for rewind RW or 2 of fast forward
        good = self.trajectories_data[self.worm_index_type] == roi.worm_index
        frames = self.trajectories_data.loc[good, 'frame_number']

        if frames.size == 0:
            return

        if rwff == self.RW:
            self.frame_number = frames.min()
        elif rwff == self.FF:
            self.frame_number = frames.max()
        else:
            raise ValueError('Invalid rwff value : {} '.format(rwff))

        self.ui.spinBox_frame.setValue(self.frame_number)

class TrajectoryEditor(ROIManager):
    def __init__(self, ui):
        super().__init__(ui)
        self.ui.pushButton_join.clicked.connect(self.joinTraj)
        self.ui.pushButton_split.clicked.connect(self.splitTraj)

        #SHORTCUTS
        self.ui.pushButton_join.setShortcut(QKeySequence(Qt.Key_J))
        self.ui.pushButton_split.setShortcut(QKeySequence(Qt.Key_S))

    def enable_trajectories_buttons(self, value):
        self.ui.pushButton_join.setEnabled(value)
        self.ui.pushButton_split.setEnabled(value)

    def joinTraj(self):
        if self.worm_index_type != 'worm_index_manual' \
        or self.frame_data is None:
            return

        worm_ind1 = self.rois[0].worm_index
        worm_ind2 = self.rois[1].worm_index

        if worm_ind1 == worm_ind2:
            QMessageBox.critical(
                self,
                'Cannot join the same trajectory with itself',
                'Cannot join the same trajectory with itself.',
                QMessageBox.Ok)
            return

        index1 = (self.trajectories_data[
                  'worm_index_manual'] == worm_ind1).values
        index2 = (self.trajectories_data[
                  'worm_index_manual'] == worm_ind2).values

        # if the trajectories do not overlap they shouldn't have frame_number
        # indexes in commun
        frame_number = self.trajectories_data.loc[
            index1 | index2, 'frame_number']

        if frame_number.size != np.unique(frame_number).size:
            QMessageBox.critical(
                self,
                'Cannot join overlaping trajectories',
                'Cannot join overlaping trajectories.',
                QMessageBox.Ok)
            return

        if not (worm_ind1 in self.frame_data[
                'worm_index_manual'].values or worm_ind2 in self.frame_data['worm_index_manual'].values):
            reply = QMessageBox.question(
                self,
                'Message',
                "The none of the selected worms to join is not in this frame. Are you sure to continue?",
                QMessageBox.Yes | QMessageBox.No,
                QMessageBox.No)

            if reply == QMessageBox.No:
                return

        # get the first row for each segment to extract some data
        first_row1 = self.trajectories_data.loc[index1, :].iloc[0]
        first_row2 = self.trajectories_data.loc[index2, :].iloc[0]

        # join trajectories
        self.trajectories_data.loc[
            index2, 'worm_label'] = first_row1['worm_label']
        self.trajectories_data.loc[index2, 'worm_index_manual'] = worm_ind1

        self.rois[0].worm_index = worm_ind1
        self.rois[1].worm_index = worm_ind1

        #this might be too slow. I might need to change it
        self.traj_worm_index_grouped = self.trajectories_data.groupby(self.worm_index_type)

        self.updateImage()

    def splitTraj(self):
        if self.worm_index_type != 'worm_index_manual' \
        or self.frame_data is None:
            return

        worm_ind = self.current_worm_index

        if not worm_ind in self.frame_data['worm_index_manual'].data:
            QMessageBox.critical(
                self,
                'Worm index is not in the current frame.',
                'Worm index is not in the current frame. Select a valid index.',
                QMessageBox.Ok)
            return

        last_index = self.trajectories_data['worm_index_manual'].max()

        new_ind1 = last_index + 1
        new_ind2 = last_index + 2

        good = self.trajectories_data['worm_index_manual'] == worm_ind
        frames = self.trajectories_data.loc[good, 'frame_number']
        frames = frames.sort_values(inplace=False)

        good = frames < self.frame_number
        index1 = frames[good].index
        index2 = frames[~good].index
        self.trajectories_data.ix[index1, 'worm_index_manual'] = new_ind1
        self.trajectories_data.ix[index2, 'worm_index_manual'] = new_ind2

        self.rois[0].index = new_ind1
        self.rois[1].index = new_ind2

        #this might be too slow. I might need to change it
        self.traj_worm_index_grouped = self.trajectories_data.groupby(self.worm_index_type)

        self.updateImage()

class FeatureReaderBase(TrackerViewerAuxGUI):
    index_cols = ['worm_index', 'timestamp', 'motion_modes', 'skeleton_id', 'well_name']
    valid_fields = ['/timeseries_data', '/features_timeseries']

    def __init__(self, ui):
        self.timeseries_data = None
        self.feat_column = ''

        super().__init__(ui)

    def updateSkelFile(self, skeletons_file):
        super().updateSkelFile(skeletons_file)
        try:
            self.traj_colors = {}
            with pd.HDFStore(self.skeletons_file, 'r') as ske_file_id:
                for field in self.valid_fields:
                    if field in ske_file_id:
                        self.timeseries_data = ske_file_id[field]

                        if field == '/timeseries_data':
                            blob_features = ske_file_id['/blob_features']
                            blob_features.columns = ['blob_' + x for x in blob_features.columns]
                            self.timeseries_data = pd.concat((self.timeseries_data, blob_features), axis=1)
                        break
                else:
                    raise KeyError

            if not len(self.timeseries_data) != len(self.trajectories_data):
                ValueError('timeseries_data and trajectories_data does not match. You might be using an old version of featuresN.hdf5')

            self.valid_features = [x for x in self.timeseries_data.columns if x not in self.index_cols]

        except (TypeError, AttributeError, IOError, KeyError, tables.exceptions.HDF5ExtError):
            self.valid_features = None
            self.timeseries_data = None

class MarkersDrawer(FeatureReaderBase):
    def __init__(self, ui):
        super().__init__(ui)

        self.traj_colors = {}
        self.n_points_traj = 250
        self.n_colors = 256
        cmap = matplotlib.cm.get_cmap("bwr")
        palette = [cmap(x) for x in np.linspace(0, 1, self.n_colors)]
        #palette = sns.color_palette("RdBu_r", self.n_colors)
        palette = np.round(np.array(palette)*255).astype(np.int)
        self.palette = [QColor(*x) for x in palette]

        self.drawT = {x: self.ui.comboBox_drawType.findText(x , flags=Qt.MatchContains)
                                for x in ['boxes', 'traj']}

        self.showT = {x: self.ui.comboBox_showLabels.findText(x , flags=Qt.MatchContains)
                                for x in ['hide', 'all', 'filter']}

        self.ui.comboBox_showLabels.setCurrentIndex(self.showT['all'])

        self.ui.comboBox_showLabels.currentIndexChanged.connect(self.updateImage)
        self.ui.comboBox_drawType.currentIndexChanged.connect(self.updateImage)

        self.ui.feature_column.currentIndexChanged.connect(self.change_feature)

        self.ui.feat_max_value.valueChanged.connect(self.updateImage)
        self.ui.feat_min_value.valueChanged.connect(self.updateImage)
        self.ui.is_color_features.stateChanged.connect(self.updateImage)

        self.enable_color_feats(False)

        self.ui.spinBox_step.valueChanged.connect(self.updateImage)

    def updateSkelFile(self, skeletons_file):
        self.ui.is_color_features.setChecked(False)

        super().updateSkelFile(skeletons_file)

        self.ui.feature_column.clear()
        if self.timeseries_data is None:
            #no feature data
            self.enable_color_feats(False)
        else:
            self.enable_color_feats(True)
            self.ui.feature_column.addItems(self.valid_features)
            self._h_find_feat_limits()

    def change_feature(self):
        self._h_find_feat_limits()
        self.updateImage()

    def _h_find_feat_limits(self):
        self.feat_column = str(self.ui.feature_column.currentText())
        print(self.feat_column)

        if self.feat_column and self.timeseries_data is not None:
            f_max = self.timeseries_data[self.feat_column].max()
            f_min = self.timeseries_data[self.feat_column].min()
            q1, q2 = self.timeseries_data[self.feat_column].quantile([0.02, 0.98])

        else:
            f_min, f_max, q1, q2  = 0,0,0,0

        self.ui.feat_max_value.setRange(f_min, f_max)
        self.ui.feat_min_value.setRange(f_min, f_max)
        self.ui.feat_min_value.setValue(q1)
        self.ui.feat_max_value.setValue(q2)

    def enable_color_feats(self, value):
        self.ui.feature_column.setEnabled(value)
        self.ui.feat_min_value.setEnabled(value)
        self.ui.feat_max_value.setEnabled(value)
        self.ui.is_color_features.setEnabled(value)

    def _h_assign_feat_color(self, irow):

        feat_val = self.timeseries_data.loc[irow, self.feat_column]

        if (feat_val != feat_val):
            return Qt.black

        #this function can and should be optimized
        f_min = self.ui.feat_min_value.value()
        f_max = self.ui.feat_max_value.value()

        if f_min == f_max: #dummy range in case all the values are the same
            f_min, f_max = -1, 1
        elif f_min > f_max:
            return Qt.black

        nn = np.clip((feat_val - f_min)/(f_max - f_min), 0, 1)
        ind = int(np.round(nn*(self.n_colors-1)))

        col = self.palette[ind]
        return col

    def draw_worm_markers(self, image):
        '''
        Draw traj worm trajectory.
        '''
        if not self.worm_index_type in self.frame_data or \
        self.ui.comboBox_showLabels.currentIndex() == self.showT['hide']:
            return

        if hasattr(self, 'current_worm_index'):
            current_index = self.current_worm_index
        else:
            current_index = -1

        painter = QPainter()
        painter.begin(image)

        self.fontsize = max(1, max(image.height(), image.width()) // 120)

        penwidth = max(1, max(image.height(), image.width()) // 800)
        self.penwidth = penwidth if penwidth % 2 == 1 else penwidth + 1

        if not self.label_type in self.frame_data:
            self.frame_data[self.label_type] = self.wlab['U']

        for row_id, row_data in self.frame_data.iterrows():
            # check if the coordinates are nan
            if np.isnan(row_data['coord_x']) or np.isnan(row_data['coord_y']):
                continue

            #if select between showing filtered index or not
            if self.ui.comboBox_showLabels.currentIndex() == self.showT['filter']:
                continue

            is_current_index =  current_index == int(row_data[self.worm_index_type])

            cb_ind = self.ui.comboBox_drawType.currentIndex()
            if cb_ind == self.drawT['boxes']:
                self.draw_boxes(painter, row_id, row_data, is_current_index)
            elif cb_ind == self.drawT['traj']:
                self.draw_trajectories(painter, row_data, is_current_index)

        painter.end()

    def _h_get_trajectory(self, worm_index, current_frame):
        worm_data = self.traj_worm_index_grouped.get_group(worm_index)
        valid_index = worm_data.index[worm_data['frame_number']<= current_frame]

        ini = max(0, valid_index.size - self.frame_step*self.n_points_traj)
        traj_ind = valid_index.values[ini::self.frame_step]
        traj_data = worm_data.loc[traj_ind]
        return traj_data

    def draw_trajectories(self, painter, row_data, is_current_index):
        if self.traj_worm_index_grouped is None:
            return
        worm_index = int(row_data[self.worm_index_type])
        current_frame = row_data['frame_number']
        traj_data = self._h_get_trajectory(worm_index, current_frame)
        traj_data = traj_data.dropna(subset=['coord_x', 'coord_y'])

        x_v = traj_data['coord_x'].round()
        y_v = traj_data['coord_y'].round()
        points = [QPointF(*map(int, c)) for c in zip(x_v, y_v)]

        if self.ui.is_color_features.isChecked():

            vec_color = [self._h_assign_feat_color(x) for x in traj_data.index]

            pen = QPen()
            pen.setWidth(self.penwidth)
            for p1, p2, c in zip(points[1:], points[:-1], vec_color):
                pen.setColor(c)
                painter.setPen(pen)
                painter.drawLine(p1, p2)
        else:
            pol = QPolygonF()
            for p in points:
                pol.append(p)

            if not worm_index in self.traj_colors:
                self.traj_colors[worm_index] = QColor(*np.random.randint(50, 230, 3))
            col = self.traj_colors[worm_index]

            pen = QPen()
            pen.setWidth(self.penwidth)
            pen.setColor(col)
            painter.setPen(pen)
            painter.drawPolyline(pol)

    def draw_boxes(self, painter, row_id, row_data,  is_current_index):
        '''
        Draw traj worm trajectory.
        '''
        worm_index = int(row_data[self.worm_index_type])
        x = int(round(row_data['coord_x']))
        y = int(round(row_data['coord_y']))

        label_color = self.wlabC[int(row_data[self.label_type])]
        if not self.ui.is_color_features.isChecked():
            label_color = self.wlabC[int(row_data[self.label_type])]
        else:
            label_color = self._h_assign_feat_color(row_id)

        pen = QPen()
        pen.setColor(label_color)
        pen.setWidth(self.penwidth)
        painter.setPen(pen)
        painter.setFont(QFont('Decorative', self.fontsize))

        painter.drawText(x, y, str(worm_index))

        bb = row_data['roi_size']
        painter.drawRect(x - bb / 2, y - bb / 2, bb, bb)

        if is_current_index:

            b_size = bb//5
            offset = bb/2 - b_size
            painter.fillRect(x + offset, y + offset, b_size, b_size, QBrush(label_color))

class PlotCommunicator(FeatureReaderBase, ROIManager):
    def __init__(self, ui=''):
        super().__init__(ui)
        self.ui.pushButton_plot.setEnabled(False)
        self.ui.pushButton_plot.clicked.connect(self.show_plot)
        self.plotter = None

    def closePrev(self):
        if self.plotter is not None:
            self.plotter.close()
            self.plotter = None

    def updateSkelFile(self, skeletons_file):
        super().updateSkelFile(skeletons_file)
        self.closePrev()
        if self.timeseries_data is None:
            self.ui.pushButton_plot.setEnabled(False)
        else:
            self.ui.pushButton_plot.setEnabled(True)

    def show_plot(self):
        self.closePrev()

        self.plotter = PlotFeatures(self.skeletons_file,
                                   self.timeseries_data,
                                   self.traj_worm_index_grouped,
                                   self.time_units,
                                   self.xy_units,
                                   self.fps,
                                   parent = self)

        self.plotter.setWindowFlags(self.plotter.windowFlags() | Qt.WindowStaysOnTopHint)

        self.plotter.show()
        self.update_plot()

    def update_plot(self):
        if self.plotter:
            self.plotter.plot(self.current_worm_index, self.feat_column)

class MWTrackerViewer_GUI( MarkersDrawer, PlotCommunicator,
    ContourDrawer, BlobLabeler, IntensityLabeler, TrajectoryEditor, WellsDrawer):

    def __init__(self, ui='', argv=''):
        if not ui:
            super().__init__(Ui_MWTrackerViewer())
        else:
            super().__init__(ui)

        self.setWindowTitle("Multi-Worm Viewer")

        self.vfilename = '' if len(argv) <= 1 else argv[1]
        self.videos_dir = r"/Volumes/behavgenom$/GeckoVideo/MaskedVideos/"
        self.results_dir = ''
        self.skeletons_file = ''
        self.worm_index_type = 'worm_index_manual'
        self.frame_data = None

        self.ui.comboBox_labelType.currentIndexChanged.connect(self.selectWormIndexType)

        self.ui.pushButton_save.clicked.connect(self.saveData)

        # select worm ROI when doubleclick a worm
        self.mainImage._canvas.mouseDoubleClickEvent = self.selectWorm
        self.mainImage._canvas.mouseRightClickEvent = self.toggleWellStatus

        self.ui.comboBox_ROI1.activated.connect(self.update_plot)
        self.ui.comboBox_ROI2.activated.connect(self.update_plot)

    def saveData(self):
        '''save data from manual labelling. pytables saving format is more convenient than pandas'''

        if os.name == 'nt':
            # I Windows the paths return by QFileDialog use / as the file
            # separation character. We need to correct it.
            for field_name in ['vfilename', 'skeletons_file']:
                setattr(
                    self, field_name, getattr(
                        self, field_name).replace(
                        '/', os.sep))
        has_skeletons_file = ((self.skeletons_file is not None)
                              and (self.skeletons_file != ''))
        if has_skeletons_file:
            save_modified_table(self.skeletons_file,
                                self.trajectories_data,
                                'trajectories_data')

        if self.is_fov_tosplit:
            if has_skeletons_file:
                self.fovsplitter.write_fov_wells_to_file(self.skeletons_file)
            else:
                warnings.warn('No skeletons file. Saving wells info in masked video')
                self.fid.close()
                self.fovsplitter.write_fov_wells_to_file(self.vfilename)
                # self.fid = tables.File(self.vfilename, 'r')
                self.updateVideoFile(self.vfilename)

        if has_skeletons_file:
            self.updateSkelFile(self.skeletons_file)

    def updateVideoFile(self, vfilename):
        super().updateVideoFile(vfilename)
        self.updateImage()

    def updateSkelFile(self, skeletons_file):
        super().updateSkelFile(skeletons_file)

        if self.trajectories_data is None:
            #empty file nothing to do here
            self.updateImage()
            return

        #correct the `worm_index_N` to the actual name `worm_index_manual`
        if 'worm_index_N' in self.trajectories_data:
            self.trajectories_data = self.trajectories_data.rename(
                columns={'worm_index_N': 'worm_index_manual'})

        #if this is really a trajectories_data not (_features.hdf5) add `worm_index_manual` if it does not exists
        if not 'worm_index_manual' in self.trajectories_data and not self.is_estimated_trajectories_data:
            self.trajectories_data['worm_label'] = self.wlab['U']
            self.trajectories_data['worm_index_manual'] = self.trajectories_data['worm_index_joined']

        #deactiate the save option if we are dealing with estimated data...
        self.ui.pushButton_save.setEnabled(not self.is_estimated_trajectories_data)

        #add this column if it does not exist
        if not 'has_skeleton' in self.trajectories_data:
            self.trajectories_data['has_skeleton'] = self.trajectories_data['skeleton_id'] >= 0

        self.updateWormIndexTypeMenu()
        self.updateImage()

    def updateWormIndexTypeMenu(self):
        possible_indexes = [x.replace('worm_index_', '') for x in self.trajectories_data.columns if x.startswith('worm_index_')]
        assert len(set(possible_indexes)) == len(possible_indexes) #all indexes ending must be different

        menu_names = sorted([x + ' index' for x in possible_indexes])
        self.ui.comboBox_labelType.clear()
        self.ui.comboBox_labelType.addItems(menu_names)
        if 'manual' in possible_indexes:
            dd = self.ui.comboBox_labelType.findText('manual index')
            self.ui.comboBox_labelType.setCurrentIndex(dd);

        self.selectWormIndexType()

    def selectWormIndexType(self):
        index_option = self.ui.comboBox_labelType.currentText()

        if not index_option:
            return
        assert index_option.endswith(' index')
        self.worm_index_type = 'worm_index_' + index_option.replace(' index', '')

        # select between automatic and manual worm indexing and label
        if self.worm_index_type == 'worm_index_manual':
            self.label_type = 'worm_label'
            self.enable_trajectories_buttons(True)
            self.enable_label_buttons(True)
        else:
            self.label_type = 'auto_label'
            self.enable_trajectories_buttons(False)
            self.enable_label_buttons(False)

        #recalculate the grouped indexes
        self.traj_worm_index_grouped = self.trajectories_data.groupby(self.worm_index_type)

        self.updateImage()

    # update image
    def updateImage(self):
        if (self.image_group is None) and (self.isimgstore is False):
            return

        super(TrackerViewerAuxGUI, self).readCurrentFrame()

        # read the data of the particles that exists in the frame
        self.frame_data = self.getFrameData(self.frame_number)

        #draw extra info only if the worm_index_type is valid
        if self.frame_data is not None and \
        self.worm_index_type in self.frame_data:
            #filter any -1 index
            self.frame_data = self.frame_data[self.frame_data[self.worm_index_type]>=0]
            if self.frame_data.size > 0:
                self.draw_worm_markers(self.frame_qimg)
                self.draw_food_contour(self.frame_qimg)

                self.updateROIs()

        else:
            self.clearROIs()
        # plot wells
        self.draw_wells(self.frame_qimg)

        # create the pixmap for the label
        self.mainImage.setPixmap(self.frame_qimg)

        self.display_intensity()

    def selectWorm(self, event):

        x = event.pos().x()
        y = event.pos().y()
        print(x,y)

        if self.frame_data is None or self.frame_data.size == 0:
            return

        R = (x - self.frame_data['coord_x'])**2 + \
            (y - self.frame_data['coord_y'])**2

        ind = R.idxmin()

        good_row = self.frame_data.loc[ind]
        if np.sqrt(R.loc[ind]) < good_row['roi_size']:
            self.current_roi.worm_index = int(good_row[self.worm_index_type])
            self.update_plot()

        self.updateImage()

    def toggleWellStatus(self, event):
        # abort if not multifov
        if self.is_fov_tosplit != True:
            return
        # event is for sure a right click or this does not get called
        x = event.pos().x()
        y = event.pos().y()
        # this will always return something. n/a if clicking outside a well
        well_name = self.fovsplitter.find_well_of_xy(x, y)[0].decode('utf-8')
        idx = self.fovsplitter.wells['well_name'] == str(well_name)
        self.fovsplitter.wells.loc[idx, 'is_good_well'] = \
            np.mod(self.fovsplitter.wells.loc[idx, 'is_good_well']+1, 2)
#        print(self.fovsplitter.wells)
        self.updateImage()

    def joinTraj(self):
        super().joinTraj()
        self.update_plot()

    def splitTraj(self):
        super().splitTraj()
        self.update_plot()

    def change_feature(self):
        super().change_feature()
        self.update_plot()

if __name__ == '__main__':
    import sys

    app = QApplication(sys.argv)
    main = MWTrackerViewer_GUI(argv=sys.argv)

    #mask_file = '/Users/avelinojaver/OneDrive - Imperial College London/tierpsy_examples/mutliworm_example/BRC20067_worms10_food1-10_Set2_Pos5_Ch2_02062017_121709.hdf5'
    #mask_file = '/Volumes/rescomp1/data/WormData/screenings/Pratheeban/First_Set/MaskedVideos/Old_Adult/16_07_22/W3_ELA_1.0_Ch1_22072016_131149.hdf5'
    #mask_file = '/Users/avelinojaver/Documents/GitHub/tierpsy-tracker/tests/data/AVI_VIDEOS/MaskedVideos/AVI_VIDEOS_1.hdf5'
#    mask_file = '/Users/avelinojaver/Documents/GitHub/tierpsy-tracker/tests/data/WT2/MaskedVideos/WT2.hdf5'
    mask_file = '/Users/lferiani/Hackathon/multiwell_tierpsy/12_FEAT_TIERPSY_forGUI/MaskedVideos/20191205/syngenta_screen_run1_bluelight_20191205_151104.22956805/metadata.hdf5'
    main.updateVideoFile(mask_file)

    main.show()
    sys.exit(app.exec_())


from typing import List

import numpy as np

class DNNLayer:
    def __init__(self, out_shape, depends_on: List["DNNLayer"] = tuple(), param_count=0):
        assert out_shape is not None  # get around varargs restriction
        self.extra_repr_params = {}
        self.unique_idx = "{}{:02d}".format(self.__class__.__name__, id(self) % 100)
        self.out_shape = out_shape
        self.depends_on = depends_on
        self.param_count = param_count

    def __repr__(self):
        args = self.extra_repr_params
        args["out_shape"] = self.out_shape
        args["param_count"] = self.param_count
        args["depends_on"] = "[{}]".format(", ".join([x.unique_idx for x in self.depends_on]))
        return "{}({})".format(self.unique_idx, ",".join(["{}={}".format(k, v) for k, v in args.items()]))

class QueryKeyValueMatrix(DNNLayer):
	# Fusing Query, Key, And Value into 1
	def __init__(self, SEQ_LEN, HIDDEN_DIM, I, ATTN_HEADS, input):
		super().__init__(
			out_shape=(3 * SEQ_LEN,I,ATTN_HEADS), # [seq_lean X intermediate_vector_dim] for 12 heads
			depends_on=[input] if input is not None else [],
			param_count=3 * HIDDEN_DIM*I*ATTN_HEADS)
		self.flop = 3 * SEQ_LEN*HIDDEN_DIM*I*ATTN_HEADS

class QKTMatrix(DNNLayer):
	# Fusing Masking and Dropout
	def __init__(self, SEQ_LEN, HIDDEN_DIM, I, ATTN_HEADS, input):
		super().__init__(
			out_shape=(SEQ_LEN,I,ATTN_HEADS),
			depends_on=[input] if input is not None else [], # Different to accept a list
			param_count=0)
		self.flop = SEQ_LEN*HIDDEN_DIM*I*ATTN_HEADS + np.prod(self.out_shape) + np.prod(self.out_shape) # QKT + mask + dropout

class Mask(DNNLayer):
    def __init__(self, input: DNNLayer):
        super().__init__(
        	out_shape=input.out_shape,
        	depends_on=[input] if input is not None else [],
        	param_count=0)
        self.flop = np.prod(self.out_shape)

class QKTVMatrix(DNNLayer):
	# QKTV + Concat
	def __init__(self, SEQ_LEN, HIDDEN_DIM, I, ATTN_HEADS, input):
		super().__init__(
			out_shape=(SEQ_LEN,I * ATTN_HEADS),
			depends_on=[input] if input is not None else [],
			param_count=0)
		self.flop = SEQ_LEN*HIDDEN_DIM*I*ATTN_HEADS + SEQ_LEN*HIDDEN_DIM*I*ATTN_HEADS # QKTVMatrix + Concat

class Concat(DNNLayer):
	def __init__(self, SEQ_LEN, HIDDEN_DIM, I, ATTN_HEADS, input):
		super().__init__(
			out_shape=(SEQ_LEN,I * ATTN_HEADS),
			depends_on=[input] if input is not None else [],
			param_count=HIDDEN_DIM*I*ATTN_HEADS)
		# self.flop = SEQ_LEN*HIDDEN_DIM*I*ATTN_HEADS
		self.flop = 0

class LinearLayerReLU(DNNLayer):
    def __init__(self, in_features: int, out_features: int, input: DNNLayer):
        super().__init__(
            self.find_outshape(in_features, out_features, input),
            [input] if input is not None else [],
            param_count=((in_features + 1) * out_features),
        )
        self.extra_repr_params["in_features"] = in_features
        self.extra_repr_params["out_features"] = out_features
        self.in_features = in_features
        self.out_features = out_features
        self.flop = 2 * self.param_count + self.out_features + np.prod(self.out_shape) # (Linear) + ReLU

    def find_outshape(self, in_features, out_features, input):
        assert len(input.out_shape) == 2 and input.out_shape[1] == in_features, f"{input.out_shape}, {in_features}"
        return (input.out_shape[0], out_features)

def selfattn_flop(B, H, K, Tc, Tg, cache_length=0):
	assert cache_length >= 0, "cache_length should be non-negative"

	x = DNNLayer(out_shape=(B, Tc, H))
	qkt = QKTMatrix(SEQ_LEN=Tc, HIDDEN_DIM=H, I=H//K, ATTN_HEADS=K, input=x)
	mask = Mask(input=x)
	flops = qkt.flop + mask.flop
	for i in range(1, Tg):
		x = DNNLayer(out_shape=(B, Tc + i, H))
		if i <= cache_length:
			qkt = QKTMatrix(SEQ_LEN=1, HIDDEN_DIM=H, I=H//K, ATTN_HEADS=K, input=x)
		else:
			qkt = QKTMatrix(SEQ_LEN=Tc + i, HIDDEN_DIM=H, I=H//K, ATTN_HEADS=K, input=x)
		flops += qkt.flop

	print(f"selfattn_flop: {flops}")
	return flops

if __name__ == "__main__":

	hparams = {"117M": (12, 768), "345M": (24, 1024), "762M": (36, 1280), "1542M": (48, 1600)}
	K = 4
	B, H = hparams["117M"]
	Tc = 128
	Tg = 128

	selfattn_flop(B=B, H=H, K=K, Tc=Tc, Tg=Tg, cache_length=0)
	selfattn_flop(B=B, H=H, K=K, Tc=Tc, Tg=Tg, cache_length=64)
	selfattn_flop(B=B, H=H, K=K, Tc=Tc, Tg=Tg, cache_length=128)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json

from alipay.aop.api.constant.ParamConstants import *

class TradeFundBillDetail(object):

    def __init__(self):
        self._amount = None
        self._asset_type_code = None
        self._asset_user_id = None
        self._biz_pay_type = None
        self._create_time = None
        self._payment_no = None

    @property
    def amount(self):
        return self._amount

    @amount.setter
    def amount(self, value):
        self._amount = value
    @property
    def asset_type_code(self):
        return self._asset_type_code

    @asset_type_code.setter
    def asset_type_code(self, value):
        self._asset_type_code = value
    @property
    def asset_user_id(self):
        return self._asset_user_id

    @asset_user_id.setter
    def asset_user_id(self, value):
        self._asset_user_id = value
    @property
    def biz_pay_type(self):
        return self._biz_pay_type

    @biz_pay_type.setter
    def biz_pay_type(self, value):
        self._biz_pay_type = value
    @property
    def create_time(self):
        return self._create_time

    @create_time.setter
    def create_time(self, value):
        self._create_time = value
    @property
    def payment_no(self):
        return self._payment_no

    @payment_no.setter
    def payment_no(self, value):
        self._payment_no = value

    def to_alipay_dict(self):
        params = dict()
        if self.amount:
            if hasattr(self.amount, 'to_alipay_dict'):
                params['amount'] = self.amount.to_alipay_dict()
            else:
                params['amount'] = self.amount
        if self.asset_type_code:
            if hasattr(self.asset_type_code, 'to_alipay_dict'):
                params['asset_type_code'] = self.asset_type_code.to_alipay_dict()
            else:
                params['asset_type_code'] = self.asset_type_code
        if self.asset_user_id:
            if hasattr(self.asset_user_id, 'to_alipay_dict'):
                params['asset_user_id'] = self.asset_user_id.to_alipay_dict()
            else:
                params['asset_user_id'] = self.asset_user_id
        if self.biz_pay_type:
            if hasattr(self.biz_pay_type, 'to_alipay_dict'):
                params['biz_pay_type'] = self.biz_pay_type.to_alipay_dict()
            else:
                params['biz_pay_type'] = self.biz_pay_type
        if self.create_time:
            if hasattr(self.create_time, 'to_alipay_dict'):
                params['create_time'] = self.create_time.to_alipay_dict()
            else:
                params['create_time'] = self.create_time
        if self.payment_no:
            if hasattr(self.payment_no, 'to_alipay_dict'):
                params['payment_no'] = self.payment_no.to_alipay_dict()
            else:
                params['payment_no'] = self.payment_no
        return params

    @staticmethod
    def from_alipay_dict(d):
        if not d:
            return None
        o = TradeFundBillDetail()
        if 'amount' in d:
            o.amount = d['amount']
        if 'asset_type_code' in d:
            o.asset_type_code = d['asset_type_code']
        if 'asset_user_id' in d:
            o.asset_user_id = d['asset_user_id']
        if 'biz_pay_type' in d:
            o.biz_pay_type = d['biz_pay_type']
        if 'create_time' in d:
            o.create_time = d['create_time']
        if 'payment_no' in d:
            o.payment_no = d['payment_no']
        return o


from typing import TYPE_CHECKING

# typing doesnt understand aenum so im pretending its stdlib enum while type checking

if TYPE_CHECKING:
    import enum
else:
    import aenum as enum

__all__ = (
    "ChannelType",
    "PresenceType",
    "RelationshipType",
    "AssetType",
    "SortType",
)

class ChannelType(enum.Enum):
    saved_message = "SavedMessage"
    direct_message = "DirectMessage"
    group = "Group"
    text_channel = "TextChannel"
    voice_channel = "VoiceChannel"

class PresenceType(enum.Enum):
    busy = "Busy"
    idle = "Idle"
    invisible = "Invisible"
    online = "Online"

class RelationshipType(enum.Enum):
    blocked = "Blocked"
    blocked_other = "BlockedOther"
    friend = "Friend"
    incoming_friend_request = "Incoming"
    none = "None"
    outgoing_friend_request = "Outgoing"
    user = "User"

class AssetType(enum.Enum):
    image = "Image"
    video = "Video"
    text = "Text"
    audio = "Audio"
    file = "File"

class SortType(enum.Enum):
    latest = "Latest"
    oldest = "Oldest"
    relevance = "Relevance"


'''
There are a total of numCourses courses you have to take, labeled from 0 to numCourses-1.

Some courses may have prerequisites, for example to take course 0 you have to first take course 1, which is expressed as a pair: [0,1]

Given the total number of courses and a list of prerequisite pairs, is it possible for you to finish all courses?

**Example 1**

`Input: numCourses = 2, prerequisites = [[1,0]]`
`Output: true`

Explanation: There are a total of 2 courses to take.
To take course 1 you should have finished course 0. So it is possible.

**Example 2**
`Input: numCourses = 2, prerequisites = [[1,0],[0,1]]`
`Output: false`

Explanation: There are a total of 2 courses to take.
To take course 1 you should have finished course 0, and to take course 0 you should
also have finished course 1. So it is impossible.

**Note**
You may assume that there are no duplicate edges in the input prerequisites.
'''

from collections import defaultdict

class Solution(object):
    def __init__(self):
        self.eligibleCourses = []
        self.visited = []

    def seedEligibleCourses(self, g):
        for index, node in g.items():
            if len(node) == 0 and index not in self.visited:
                self.eligibleCourses.append(index)

    def dfs(self, node, g):
        if node in self.visited:
            return

        self.visited.append(node)
        for _, n in g.items():
            if node in n:
                n.remove(node)

        for successor in g[node]:
            if successor not in self.visited:
                self.eligibleCourses.append(successor)
                self.dfs(node, g)

    def canFinish(self, numCourses, prerequisites):
        if not prerequisites:
            return True

        graph = defaultdict(list)
        for relation in prerequisites:
            currentCourse, prerequisite = relation[0], relation[1]
            graph[prerequisite].append(currentCourse)  # post order!!
            if currentCourse not in graph:
                graph[currentCourse] = []

        self.seedEligibleCourses(graph)
        while self.eligibleCourses:
            current = self.eligibleCourses.pop(0)
            self.dfs(current, graph)
            self.seedEligibleCourses(graph)

        for _, n in graph.items():
            if len(n) > 0:
                return False
        return True


#%% load the background
from __future__ import print_function, division
import torch
from torchvision import datasets, transforms
import os
import matplotlib.pyplot as plt
import seaborn as sns; sns.set()
import pandas as pd
import numpy as np
import torch.nn as nn

#%% define the datasets
list_datasets = ['/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/original',
                '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_HE',
                '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_tumorLymphnode_165',
                '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_onlyH',
                '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/tumorLymphnode/patches/size_165/original',
                 '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/tumorLymphnode/patches/size_165/normalized_to_HE_165',
                 '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/tumorLymphnode/patches/size_165/normalized_to_camelyon_165',
                '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/tumorLymphnode/patches/size_165/normalized_to_onlyH_165'
                 ]
list_dataset_names = ['camelyon_ori', 'camelyon_to_HE', 'camelyon_to_tL', 'camelyon_to_H',
                      'tumorLymphnode_ori', 'tumorLymphnode_to_HE', 'tumorLymphnode_to_ca', 'tumorLymphnode_to_H']

list_models = ['/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/original/model_ResNet152.pt',
            '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_HE/model_ResNet152.pt',
            '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_tumorLymphnode_165/model_ResNet152.pt',
            '/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/patchCamelyon/patches/normalized_to_onlyH/model_ResNet152.pt' ]

list_model_names = ['ResNet_original', "ResNet_normalized_to_HE", "ResNet_normalized_to_tumorLymphnode", "ResNet_normalized_to_H"]

#%% iterate over all datasets (and later over all models)
list_model = []
list_dataset = []
list_kappa = []
list_accuracy = []
list_loss = []

for idataset, tdataset in enumerate(list_datasets):
    #print(idataset)

    #%% define the folder
    if tdataset.find("patches") > 0:
        dataset2use = "val"
    else:
        dataset2use = 'test'

    # %%define the function to get the data
    def get_datatransform(inputSize, data_dir):

        data_transforms = {
            dataset2use: transforms.Compose([
                transforms.Resize([inputSize, inputSize]),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
        }

        image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                                  data_transforms[x])
                          for x in [dataset2use]}
        dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                                      shuffle=False, num_workers=4)
                       for x in [dataset2use]}

        return(data_transforms, image_datasets, dataloaders)

    #%% prepare the transformations and the dataset
    data_transforms , image_datasets, dataloaders= get_datatransform(259, tdataset)

    class_names = dataloaders[dataset2use].dataset.classes
    nb_classes = len(class_names)
    confusion_matrix = torch.zeros(nb_classes, nb_classes)

    #%% visualize the input data (to look if evey class is evenly)
    class_names =  ['normal', 'tumor']

    df = pd.DataFrame(dataloaders[dataset2use].dataset.samples)
    df.columns = ['file', 'class_nr']

    df.class_nr = np.array(df.class_nr)

    class_labels = ['NaN' for x in range(df.shape[0])]
    for i in range(0,df.shape[0]):
        class_labels[i] = class_names[df.class_nr[int(i)]]
    df = df.assign(class_labels = class_labels)
    sns.set_palette("Set1", n_colors = 12)
    sns.countplot(df.class_labels)
    plt.xlabel('Pattern')
    plt.ylabel('Count [n]')
    plt.savefig('DataBase_' + dataset2use + '.jpg')
    plt.show()
    plt.close()

    n_normal = sum(map(lambda x : x == "normal", class_labels))
    n_tumor = sum(map(lambda x: x == "tumor", class_labels))
    print("n = " + str(n_normal) + " tiles without and n = " + str(n_tumor) + " tiles with tumor.")

    #%% iterate over the models
    from sklearn.metrics import cohen_kappa_score
    from sklearn.metrics import accuracy_score
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    n = 0
    df_values = pd.DataFrame(list(range(0,len(dataloaders[dataset2use].sampler.data_source.imgs))))

    for imodel, tmodel in enumerate(list_models):
        print(imodel)

        #%% prepare the dataset
        inputSize = 224
        data_transforms, image_datasets, dataloaders = get_datatransform(inputSize, tdataset)

        #%% apply model on test data set (and get a confusion matrix)
        model_ft = torch.load(tmodel)
        model_ft.eval()
        vector_prd = []
        vector_exp = []

        with torch.no_grad():
            for i, (inputs, classes) in enumerate(dataloaders[dataset2use]):
                inputs = inputs.to(device)
                classes = classes.to(device)
                outputs = model_ft(inputs)
                _, preds = torch.max(outputs, 1)

                if i == 0:
                    outputs_matrix = outputs
                else:
                    outputs_matrix = torch.cat((outputs_matrix, outputs), 0)

                vector_prd = vector_prd + preds.view(-1).cpu().tolist()
                vector_exp = vector_exp + classes.view(-1).cpu().tolist()

        confusion_matrix = torch.zeros(nb_classes, nb_classes)
        for x, y in zip(vector_exp, vector_prd):
            confusion_matrix[y, x] += 1

        loss_function = nn.CrossEntropyLoss()
        loss_value = loss_function(outputs_matrix.to('cpu'), torch.tensor(vector_exp))
        print(confusion_matrix)

        #%% calcualte the comparison values
        list_model.append(list_model_names[imodel])
        list_dataset.append(list_dataset_names[idataset])
        list_kappa.append(cohen_kappa_score(vector_prd, vector_exp))
        list_accuracy.append(accuracy_score(vector_prd, vector_exp))
        list_loss.append(loss_value.tolist())
        print('Kappa-value: ' + str(list_kappa[-1]))
        print('Accurary-value: ' + str(list_accuracy[-1]))

        #%% plot a confusion matrix
        matrix2plot = confusion_matrix.numpy()
        matrix2plot = matrix2plot.astype(int)

        ax = sns.heatmap(matrix2plot,
                         annot = True, linewidths=5, annot_kws={"size": 10},
                         xticklabels=class_names, yticklabels=class_names,
                         cmap = "Blues")
        plt.xlabel('Ground Truth')
        plt.ylabel('Model ' + list_model[-1] + " on " + list_dataset[-1])
        plt.savefig('ConfMat_' +'Model ' + list_model[-1] + " on " + list_dataset[-1] + '.jpg')
        plt.show()
        plt.close()

#%% make a dataframe
df = pd.DataFrame(list(zip(list_model, list_dataset, list_kappa)), columns=['model', 'data', 'kappa'])
df = df.pivot_table(index = ["model"], columns = ["data"], values = "kappa")
df.to_csv('/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/table.csv')
df.to_excel('/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/table.xlsx')
with open('/home/cw9/sds_hd/sd18a006/marlen/datasets/stainNormalization/table.tex', 'w') as tf:
    tf.write(df.to_latex())


def is_element(element):

    pass


"""
Limits
======

Implemented according to the PhD thesis
http://www.cybertester.com/data/gruntz.pdf, which contains very thorough
descriptions of the algorithm including many examples.  We summarize here
the gist of it.

All functions are sorted according to how rapidly varying they are at
infinity using the following rules. Any two functions f and g can be
compared using the properties of L:

L=lim  log|f(x)| / log|g(x)|           (for x -> oo)

We define >, < ~ according to::

    1. f > g .... L=+-oo

        we say that:
        - f is greater than any power of g
        - f is more rapidly varying than g
        - f goes to infinity/zero faster than g

    2. f < g .... L=0

        we say that:
        - f is lower than any power of g

    3. f ~ g .... L!=0, +-oo

        we say that:
        - both f and g are bounded from above and below by suitable integral
          powers of the other

Examples
========
::
    2 < x < exp(x) < exp(x**2) < exp(exp(x))
    2 ~ 3 ~ -5
    x ~ x**2 ~ x**3 ~ 1/x ~ x**m ~ -x
    exp(x) ~ exp(-x) ~ exp(2x) ~ exp(x)**2 ~ exp(x+exp(-x))
    f ~ 1/f

So we can divide all the functions into comparability classes (x and x^2
belong to one class, exp(x) and exp(-x) belong to some other class). In
principle, we could compare any two functions, but in our algorithm, we
do not compare anything below the class 2~3~-5 (for example log(x) is
below this), so we set 2~3~-5 as the lowest comparability class.

Given the function f, we find the list of most rapidly varying (mrv set)
subexpressions of it. This list belongs to the same comparability class.
Let's say it is {exp(x), exp(2x)}. Using the rule f ~ 1/f we find an
element "w" (either from the list or a new one) from the same
comparability class which goes to zero at infinity. In our example we
set w=exp(-x) (but we could also set w=exp(-2x) or w=exp(-3x) ...). We
rewrite the mrv set using w, in our case {1/w, 1/w^2}, and substitute it
into f. Then we expand f into a series in w::

    f = c0*w^e0 + c1*w^e1 + ... + O(w^en),       where e0<e1<...<en, c0!=0

but for x->oo, lim f = lim c0*w^e0, because all the other terms go to zero,
because w goes to zero faster than the ci and ei. So::

    for e0>0, lim f = 0
    for e0<0, lim f = +-oo   (the sign depends on the sign of c0)
    for e0=0, lim f = lim c0

We need to recursively compute limits at several places of the algorithm, but
as is shown in the PhD thesis, it always finishes.

Important functions from the implementation:

compare(a, b, x) compares "a" and "b" by computing the limit L.
mrv(e, x) returns list of most rapidly varying (mrv) subexpressions of "e"
rewrite(e, Omega, x, wsym) rewrites "e" in terms of w
leadterm(f, x) returns the lowest power term in the series of f
mrv_leadterm(e, x) returns the lead term (c0, e0) for e
limitinf(e, x) computes lim e  (for x->oo)
limit(e, z, z0) computes any limit by converting it to the case x->oo

All the functions are really simple and straightforward except
rewrite(), which is the most difficult/complex part of the algorithm.
When the algorithm fails, the bugs are usually in the series expansion
(i.e. in SymPy) or in rewrite.

This code is almost exact rewrite of the Maple code inside the Gruntz
thesis.

Debugging
---------

Because the gruntz algorithm is highly recursive, it's difficult to
figure out what went wrong inside a debugger. Instead, turn on nice
debug prints by defining the environment variable SYMPY_DEBUG. For
example:

[user@localhost]: SYMPY_DEBUG=True ./bin/isympy

In [1]: limit(sin(x)/x, x, 0)
limitinf(_x*sin(1/_x), _x) = 1
+-mrv_leadterm(_x*sin(1/_x), _x) = (1, 0)
| +-mrv(_x*sin(1/_x), _x) = set([_x])
| | +-mrv(_x, _x) = set([_x])
| | +-mrv(sin(1/_x), _x) = set([_x])
| |   +-mrv(1/_x, _x) = set([_x])
| |     +-mrv(_x, _x) = set([_x])
| +-mrv_leadterm(exp(_x)*sin(exp(-_x)), _x, set([exp(_x)])) = (1, 0)
|   +-rewrite(exp(_x)*sin(exp(-_x)), set([exp(_x)]), _x, _w) = (1/_w*sin(_w), -_x)
|     +-sign(_x, _x) = 1
|     +-mrv_leadterm(1, _x) = (1, 0)
+-sign(0, _x) = 0
+-limitinf(1, _x) = 1

And check manually which line is wrong. Then go to the source code and
debug this function to figure out the exact problem.

"""
from functools import reduce

from sympy.core import Basic, S, Mul, PoleError
from sympy.core.cache import cacheit
from sympy.core.numbers import ilcm, I, oo
from sympy.core.symbol import Dummy, Wild
from sympy.core.traversal import bottom_up

from sympy.functions import log, exp, sign as _sign
from sympy.series.order import Order
from sympy.simplify import logcombine
from sympy.simplify.powsimp import powsimp, powdenest

from sympy.utilities.misc import debug_decorator as debug
from sympy.utilities.timeutils import timethis

timeit = timethis('gruntz')

def compare(a, b, x):
    """Returns "<" if a<b, "=" for a == b, ">" for a>b"""
    # log(exp(...)) must always be simplified here for termination
    la, lb = log(a), log(b)
    if isinstance(a, Basic) and (isinstance(a, exp) or (a.is_Pow and a.base == S.Exp1)):
        la = a.exp
    if isinstance(b, Basic) and (isinstance(b, exp) or (b.is_Pow and b.base == S.Exp1)):
        lb = b.exp

    c = limitinf(la/lb, x)
    if c == 0:
        return "<"
    elif c.is_infinite:
        return ">"
    else:
        return "="

class SubsSet(dict):
    """
    Stores (expr, dummy) pairs, and how to rewrite expr-s.

    Explanation
    ===========

    The gruntz algorithm needs to rewrite certain expressions in term of a new
    variable w. We cannot use subs, because it is just too smart for us. For
    example::

        > Omega=[exp(exp(_p - exp(-_p))/(1 - 1/_p)), exp(exp(_p))]
        > O2=[exp(-exp(_p) + exp(-exp(-_p))*exp(_p)/(1 - 1/_p))/_w, 1/_w]
        > e = exp(exp(_p - exp(-_p))/(1 - 1/_p)) - exp(exp(_p))
        > e.subs(Omega[0],O2[0]).subs(Omega[1],O2[1])
        -1/w + exp(exp(p)*exp(-exp(-p))/(1 - 1/p))

    is really not what we want!

    So we do it the hard way and keep track of all the things we potentially
    want to substitute by dummy variables. Consider the expression::

        exp(x - exp(-x)) + exp(x) + x.

    The mrv set is {exp(x), exp(-x), exp(x - exp(-x))}.
    We introduce corresponding dummy variables d1, d2, d3 and rewrite::

        d3 + d1 + x.

    This class first of all keeps track of the mapping expr->variable, i.e.
    will at this stage be a dictionary::

        {exp(x): d1, exp(-x): d2, exp(x - exp(-x)): d3}.

    [It turns out to be more convenient this way round.]
    But sometimes expressions in the mrv set have other expressions from the
    mrv set as subexpressions, and we need to keep track of that as well. In
    this case, d3 is really exp(x - d2), so rewrites at this stage is::

        {d3: exp(x-d2)}.

    The function rewrite uses all this information to correctly rewrite our
    expression in terms of w. In this case w can be chosen to be exp(-x),
    i.e. d2. The correct rewriting then is::

        exp(-w)/w + 1/w + x.
    """
    def __init__(self):
        self.rewrites = {}

    def __repr__(self):
        return super().__repr__() + ', ' + self.rewrites.__repr__()

    def __getitem__(self, key):
        if not key in self:
            self[key] = Dummy()
        return dict.__getitem__(self, key)

    def do_subs(self, e):
        """Substitute the variables with expressions"""
        for expr, var in self.items():
            e = e.xreplace({var: expr})
        return e

    def meets(self, s2):
        """Tell whether or not self and s2 have non-empty intersection"""
        return set(self.keys()).intersection(list(s2.keys())) != set()

    def union(self, s2, exps=None):
        """Compute the union of self and s2, adjusting exps"""
        res = self.copy()
        tr = {}
        for expr, var in s2.items():
            if expr in self:
                if exps:
                    exps = exps.xreplace({var: res[expr]})
                tr[var] = res[expr]
            else:
                res[expr] = var
        for var, rewr in s2.rewrites.items():
            res.rewrites[var] = rewr.xreplace(tr)
        return res, exps

    def copy(self):
        """Create a shallow copy of SubsSet"""
        r = SubsSet()
        r.rewrites = self.rewrites.copy()
        for expr, var in self.items():
            r[expr] = var
        return r

@debug
def mrv(e, x):
    """Returns a SubsSet of most rapidly varying (mrv) subexpressions of 'e',
       and e rewritten in terms of these"""
    e = powsimp(e, deep=True, combine='exp')
    if not isinstance(e, Basic):
        raise TypeError("e should be an instance of Basic")
    if not e.has(x):
        return SubsSet(), e
    elif e == x:
        s = SubsSet()
        return s, s[x]
    elif e.is_Mul or e.is_Add:
        i, d = e.as_independent(x)  # throw away x-independent terms
        if d.func != e.func:
            s, expr = mrv(d, x)
            return s, e.func(i, expr)
        a, b = d.as_two_terms()
        s1, e1 = mrv(a, x)
        s2, e2 = mrv(b, x)
        return mrv_max1(s1, s2, e.func(i, e1, e2), x)
    elif e.is_Pow and e.base != S.Exp1:
        e1 = S.One
        while e.is_Pow:
            b1 = e.base
            e1 *= e.exp
            e = b1
        if b1 == 1:
            return SubsSet(), b1
        if e1.has(x):
            base_lim = limitinf(b1, x)
            if base_lim is S.One:
                return mrv(exp(e1 * (b1 - 1)), x)
            return mrv(exp(e1 * log(b1)), x)
        else:
            s, expr = mrv(b1, x)
            return s, expr**e1
    elif isinstance(e, log):
        s, expr = mrv(e.args[0], x)
        return s, log(expr)
    elif isinstance(e, exp) or (e.is_Pow and e.base == S.Exp1):
        # We know from the theory of this algorithm that exp(log(...)) may always
        # be simplified here, and doing so is vital for termination.
        if isinstance(e.exp, log):
            return mrv(e.exp.args[0], x)
        # if a product has an infinite factor the result will be
        # infinite if there is no zero, otherwise NaN; here, we
        # consider the result infinite if any factor is infinite
        li = limitinf(e.exp, x)
        if any(_.is_infinite for _ in Mul.make_args(li)):
            s1 = SubsSet()
            e1 = s1[e]
            s2, e2 = mrv(e.exp, x)
            su = s1.union(s2)[0]
            su.rewrites[e1] = exp(e2)
            return mrv_max3(s1, e1, s2, exp(e2), su, e1, x)
        else:
            s, expr = mrv(e.exp, x)
            return s, exp(expr)
    elif e.is_Function:
        l = [mrv(a, x) for a in e.args]
        l2 = [s for (s, _) in l if s != SubsSet()]
        if len(l2) != 1:
            # e.g. something like BesselJ(x, x)
            raise NotImplementedError("MRV set computation for functions in"
                                      " several variables not implemented.")
        s, ss = l2[0], SubsSet()
        args = [ss.do_subs(x[1]) for x in l]
        return s, e.func(*args)
    elif e.is_Derivative:
        raise NotImplementedError("MRV set computation for derviatives"
                                  " not implemented yet.")
    raise NotImplementedError(
        "Don't know how to calculate the mrv of '%s'" % e)

def mrv_max3(f, expsf, g, expsg, union, expsboth, x):
    """
    Computes the maximum of two sets of expressions f and g, which
    are in the same comparability class, i.e. max() compares (two elements of)
    f and g and returns either (f, expsf) [if f is larger], (g, expsg)
    [if g is larger] or (union, expsboth) [if f, g are of the same class].
    """
    if not isinstance(f, SubsSet):
        raise TypeError("f should be an instance of SubsSet")
    if not isinstance(g, SubsSet):
        raise TypeError("g should be an instance of SubsSet")
    if f == SubsSet():
        return g, expsg
    elif g == SubsSet():
        return f, expsf
    elif f.meets(g):
        return union, expsboth

    c = compare(list(f.keys())[0], list(g.keys())[0], x)
    if c == ">":
        return f, expsf
    elif c == "<":
        return g, expsg
    else:
        if c != "=":
            raise ValueError("c should be =")
        return union, expsboth

def mrv_max1(f, g, exps, x):
    """Computes the maximum of two sets of expressions f and g, which
    are in the same comparability class, i.e. mrv_max1() compares (two elements of)
    f and g and returns the set, which is in the higher comparability class
    of the union of both, if they have the same order of variation.
    Also returns exps, with the appropriate substitutions made.
    """
    u, b = f.union(g, exps)
    return mrv_max3(f, g.do_subs(exps), g, f.do_subs(exps),
                    u, b, x)

@debug
@cacheit
@timeit
def sign(e, x):
    """
    Returns a sign of an expression e(x) for x->oo.

    ::

        e >  0 for x sufficiently large ...  1
        e == 0 for x sufficiently large ...  0
        e <  0 for x sufficiently large ... -1

    The result of this function is currently undefined if e changes sign
    arbitrarily often for arbitrarily large x (e.g. sin(x)).

    Note that this returns zero only if e is *constantly* zero
    for x sufficiently large. [If e is constant, of course, this is just
    the same thing as the sign of e.]
    """
    if not isinstance(e, Basic):
        raise TypeError("e should be an instance of Basic")

    if e.is_positive:
        return 1
    elif e.is_negative:
        return -1
    elif e.is_zero:
        return 0

    elif not e.has(x):
        e = logcombine(e)
        return _sign(e)
    elif e == x:
        return 1
    elif e.is_Mul:
        a, b = e.as_two_terms()
        sa = sign(a, x)
        if not sa:
            return 0
        return sa * sign(b, x)
    elif isinstance(e, exp):
        return 1
    elif e.is_Pow:
        if e.base == S.Exp1:
            return 1
        s = sign(e.base, x)
        if s == 1:
            return 1
        if e.exp.is_Integer:
            return s**e.exp
    elif isinstance(e, log):
        return sign(e.args[0] - 1, x)

    # if all else fails, do it the hard way
    c0, e0 = mrv_leadterm(e, x)
    return sign(c0, x)

@debug
@timeit
@cacheit
def limitinf(e, x, leadsimp=False):
    """Limit e(x) for x-> oo.

    Explanation
    ===========

    If ``leadsimp`` is True, an attempt is made to simplify the leading
    term of the series expansion of ``e``. That may succeed even if
    ``e`` cannot be simplified.
    """
    # rewrite e in terms of tractable functions only

    if not e.has(x):
        return e  # e is a constant
    if e.has(Order):
        e = e.expand().removeO()
    if not x.is_positive or x.is_integer:
        # We make sure that x.is_positive is True and x.is_integer is None
        # so we get all the correct mathematical behavior from the expression.
        # We need a fresh variable.
        p = Dummy('p', positive=True)
        e = e.subs(x, p)
        x = p
    e = e.rewrite('tractable', deep=True, limitvar=x)
    e = powdenest(e)
    c0, e0 = mrv_leadterm(e, x)
    sig = sign(e0, x)
    if sig == 1:
        return S.Zero  # e0>0: lim f = 0
    elif sig == -1:  # e0<0: lim f = +-oo (the sign depends on the sign of c0)
        if c0.match(I*Wild("a", exclude=[I])):
            return c0*oo
        s = sign(c0, x)
        # the leading term shouldn't be 0:
        if s == 0:
            raise ValueError("Leading term should not be 0")
        return s*oo
    elif sig == 0:
        if leadsimp:
            c0 = c0.simplify()
        return limitinf(c0, x, leadsimp)  # e0=0: lim f = lim c0
    else:
        raise ValueError("{} could not be evaluated".format(sig))

def moveup2(s, x):
    r = SubsSet()
    for expr, var in s.items():
        r[expr.xreplace({x: exp(x)})] = var
    for var, expr in s.rewrites.items():
        r.rewrites[var] = s.rewrites[var].xreplace({x: exp(x)})
    return r

def moveup(l, x):
    return [e.xreplace({x: exp(x)}) for e in l]

@debug
@timeit
def calculate_series(e, x, logx=None):
    """ Calculates at least one term of the series of ``e`` in ``x``.

    This is a place that fails most often, so it is in its own function.
    """
    from sympy.polys import cancel

    for t in e.lseries(x, logx=logx):
        # bottom_up function is required for a specific case - when e is
        # -exp(p/(p + 1)) + exp(-p**2/(p + 1) + p). No current simplification
        # methods reduce this to 0 while not expanding polynomials.
        t = bottom_up(t, lambda w: getattr(w, 'normal', lambda: w)())
        t = cancel(t, expand=False).factor()

        if t.has(exp) and t.has(log):
            t = powdenest(t)

        if not t.is_zero:
            break

    return t

@debug
@timeit
@cacheit
def mrv_leadterm(e, x):
    """Returns (c0, e0) for e."""
    Omega = SubsSet()
    if not e.has(x):
        return (e, S.Zero)
    if Omega == SubsSet():
        Omega, exps = mrv(e, x)
    if not Omega:
        # e really does not depend on x after simplification
        return exps, S.Zero
    if x in Omega:
        # move the whole omega up (exponentiate each term):
        Omega_up = moveup2(Omega, x)
        exps_up = moveup([exps], x)[0]
        # NOTE: there is no need to move this down!
        Omega = Omega_up
        exps = exps_up
    #
    # The positive dummy, w, is used here so log(w*2) etc. will expand;
    # a unique dummy is needed in this algorithm
    #
    # For limits of complex functions, the algorithm would have to be
    # improved, or just find limits of Re and Im components separately.
    #
    w = Dummy("w", real=True, positive=True)
    f, logw = rewrite(exps, Omega, x, w)
    series = calculate_series(f, w, logx=logw)
    try:
        lt = series.leadterm(w, logx=logw)
    except (ValueError, PoleError):
        lt = f.as_coeff_exponent(w)
        # as_coeff_exponent won't always split in required form. It may simply
        # return (f, 0) when a better form may be obtained. Example (-x)**(-pi)
        # can be written as (-1**(-pi), -pi) which as_coeff_exponent does not return
        if lt[0].has(w):
            base = f.as_base_exp()[0].as_coeff_exponent(w)
            ex = f.as_base_exp()[1]
            lt = (base[0]**ex, base[1]*ex)
    return (lt[0].subs(log(w), logw), lt[1])

def build_expression_tree(Omega, rewrites):
    r""" Helper function for rewrite.

    We need to sort Omega (mrv set) so that we replace an expression before
    we replace any expression in terms of which it has to be rewritten::

        e1 ---> e2 ---> e3
                 \
                  -> e4

    Here we can do e1, e2, e3, e4 or e1, e2, e4, e3.
    To do this we assemble the nodes into a tree, and sort them by height.

    This function builds the tree, rewrites then sorts the nodes.
    """
    class Node:
        def __init__(self):
            self.before = []
            self.expr = None
            self.var = None
        def ht(self):
            return reduce(lambda x, y: x + y,
                          [x.ht() for x in self.before], 1)
    nodes = {}
    for expr, v in Omega:
        n = Node()
        n.var = v
        n.expr = expr
        nodes[v] = n
    for _, v in Omega:
        if v in rewrites:
            n = nodes[v]
            r = rewrites[v]
            for _, v2 in Omega:
                if r.has(v2):
                    n.before.append(nodes[v2])

    return nodes

@debug
@timeit
def rewrite(e, Omega, x, wsym):
    """e(x) ... the function
    Omega ... the mrv set
    wsym ... the symbol which is going to be used for w

    Returns the rewritten e in terms of w and log(w). See test_rewrite1()
    for examples and correct results.
    """
    if not isinstance(Omega, SubsSet):
        raise TypeError("Omega should be an instance of SubsSet")
    if len(Omega) == 0:
        raise ValueError("Length cannot be 0")
    # all items in Omega must be exponentials
    for t in Omega.keys():
        if not isinstance(t, exp):
            raise ValueError("Value should be exp")
    rewrites = Omega.rewrites
    Omega = list(Omega.items())

    nodes = build_expression_tree(Omega, rewrites)
    Omega.sort(key=lambda x: nodes[x[1]].ht(), reverse=True)

    # make sure we know the sign of each exp() term; after the loop,
    # g is going to be the "w" - the simplest one in the mrv set
    for g, _ in Omega:
        sig = sign(g.exp, x)
        if sig != 1 and sig != -1:
            raise NotImplementedError('Result depends on the sign of %s' % sig)
    if sig == 1:
        wsym = 1/wsym  # if g goes to oo, substitute 1/w
    # O2 is a list, which results by rewriting each item in Omega using "w"
    O2 = []
    denominators = []
    for f, var in Omega:
        c = limitinf(f.exp/g.exp, x)
        if c.is_Rational:
            denominators.append(c.q)
        arg = f.exp
        if var in rewrites:
            if not isinstance(rewrites[var], exp):
                raise ValueError("Value should be exp")
            arg = rewrites[var].args[0]
        O2.append((var, exp((arg - c*g.exp).expand())*wsym**c))

    # Remember that Omega contains subexpressions of "e". So now we find
    # them in "e" and substitute them for our rewriting, stored in O2

    # the following powsimp is necessary to automatically combine exponentials,
    # so that the .xreplace() below succeeds:
    # TODO this should not be necessary
    f = powsimp(e, deep=True, combine='exp')
    for a, b in O2:
        f = f.xreplace({a: b})

    for _, var in Omega:
        assert not f.has(var)

    # finally compute the logarithm of w (logw).
    logw = g.exp
    if sig == 1:
        logw = -logw  # log(w)->log(1/w)=-log(w)

    # Some parts of SymPy have difficulty computing series expansions with
    # non-integral exponents. The following heuristic improves the situation:
    exponent = reduce(ilcm, denominators, 1)
    f = f.subs({wsym: wsym**exponent})
    logw /= exponent

    return f, logw

def gruntz(e, z, z0, dir="+"):
    """
    Compute the limit of e(z) at the point z0 using the Gruntz algorithm.

    Explanation
    ===========

    ``z0`` can be any expression, including oo and -oo.

    For ``dir="+"`` (default) it calculates the limit from the right
    (z->z0+) and for ``dir="-"`` the limit from the left (z->z0-). For infinite z0
    (oo or -oo), the dir argument doesn't matter.

    This algorithm is fully described in the module docstring in the gruntz.py
    file. It relies heavily on the series expansion. Most frequently, gruntz()
    is only used if the faster limit() function (which uses heuristics) fails.
    """
    if not z.is_symbol:
        raise NotImplementedError("Second argument must be a Symbol")

    # convert all limits to the limit z->oo; sign of z is handled in limitinf
    r = None
    if z0 == oo:
        e0 = e
    elif z0 == -oo:
        e0 = e.subs(z, -z)
    else:
        if str(dir) == "-":
            e0 = e.subs(z, z0 - 1/z)
        elif str(dir) == "+":
            e0 = e.subs(z, z0 + 1/z)
        else:
            raise NotImplementedError("dir must be '+' or '-'")

    try:
        r = limitinf(e0, z)
    except ValueError:
        r = limitinf(e0, z, leadsimp=True)

    # This is a bit of a heuristic for nice results... we always rewrite
    # tractable functions in terms of familiar intractable ones.
    # It might be nicer to rewrite the exactly to what they were initially,
    # but that would take some work to implement.
    return r.rewrite('intractable', deep=True)


from .control_sequence import ControlSequence
from .formatter import format_table

class DocumentClass(ControlSequence):
    def __init__(self, name, descr=''):
        ControlSequence.__init__(self, name, descr)
        # A document class may have options.
        self.has_opts = True
        return


from unicodedata import name
from django.contrib.auth import get_user_model
from django.urls import reverse
from django.test import TestCase

from rest_framework import status
from rest_framework.test import APIClient

from core.models import Ingredient,Recipe
from recipe.serializers import IngredientSerializer

INGREDIENTS_URL = reverse('recipe:ingredient-list')

class PublicIngredientsApiTests(TestCase):
    #Test tha publicly available ingredients API

    def setUp(self):
        self.client = APIClient()

    def test_login_required(self):
        #Test tha login is required to access the endpoint
        res = self.client.get(INGREDIENTS_URL)

        self.assertEqual(res.status_code, status.HTTP_401_UNAUTHORIZED)

class PrivateIngredientsApiTests(TestCase):
    # Test that private ingredient api

    def setUp(self) -> None:
        self.client = APIClient()
        self.user = get_user_model().objects.create_user(
            'test@gmail.com',
            'pss123'
        )
        self.client.force_authenticate(self.user)

    def test_retrive_ingredient_list(self):
        #Test that retrivinng the list of ingredients
        Ingredient.objects.create(user = self.user,name='Kale')
        Ingredient.objects.create(user = self.user,name='Salt')

        res = self.client.get(INGREDIENTS_URL)

        ingredients = Ingredient.objects.all().order_by('-name')
        serializer = IngredientSerializer(ingredients, many=True)
        self.assertEqual(res.status_code,status.HTTP_200_OK)
        self.assertEqual(res.data,serializer.data)

    def test_ingredients_limits_to_user(self):
        #Test that ingredients for authenticated user and return it
        user2 = get_user_model().objects.create_user(
            'other@gamil.com'
            'pass321'
        )

        Ingredient.objects.create(user=user2,name='Vinegar')
        ingredient = Ingredient.objects.create(user = self.user,name = 'Tumeric')

        res = self.client.get(INGREDIENTS_URL)

        self.assertEqual(res.status_code,status.HTTP_200_OK)
        self.assertEqual(len(res.data),1)
        self.assertEqual(res.data[0]['name'], ingredient.name)

    def test_create_ingredient_successful(self):
        #Test creating a new ingredient
        payload = {'name':'Cabbage'}
        self.client.post(INGREDIENTS_URL,payload)
        exists = Ingredient.objects.filter(
            user=self.user,
            name=payload['name']
        ).exists()
        self.assertTrue(exists)

    def test_create_ingredient_invalid(self):
        #Test creating a new ingredient with invalid payload
        payload = {'name':''}
        res = self.client.post(INGREDIENTS_URL,payload)

        self.assertEqual(res.status_code,status.HTTP_400_BAD_REQUEST)

    def test_retrieve_ingredients_assigned_to_recipes(self):
        #Test filtering ingredients by those assigned to recipes
        ingredient1 = Ingredient.objects.create(user=self.user,name='Apples')
        ingredient2 = Ingredient.objects.create(user=self.user,name='Turkey')

        recipe = Recipe.objects.create(
            title ='Apple crumble',
            time_minutes=5,
            price=10,
            user = self.user
        )
        recipe.ingredients.add(ingredient1)

        res = self.client.get(INGREDIENTS_URL, {'assigned_only':1})

        serializer1 = IngredientSerializer(ingredient1)
        serializer2 = IngredientSerializer(ingredient2)
        self.assertIn(serializer1.data, res.data)
        self.assertNotIn(serializer2.data, res.data)

    def test_retreive_ingredients_assigned_unique(self):
        #Test filtering ingredients by assigned return unique items
        ingredient = Ingredient.objects.create(user=self.user,name='Eggs')
        Ingredient.objects.create(user=self.user,name='Cheese')

        recipe1 = Recipe.objects.create(
            title = 'Eggs benedict',
            time_minutes = 30,
            price = 12.00,
            user=self.user
        )
        recipe1.ingredients.add(ingredient)

        recipe2 = Recipe.objects.create(
            title = 'Coriander eggs',
            time_minutes=20,
            price=5.00,
            user=self.user
        )
        recipe2.ingredients.add(ingredient)

        res = self.client.get(INGREDIENTS_URL,{'assigned_only':1})

        self.assertEqual(len(res.data),1)