import numpy as np
import pdb
import math
from . import data_generators
import copy
from log import logger

def calc_iou(rois, augmented_annotation, C, class_name_idx_mapping):
    """
    分别计算每一个 roi 和所有 gt_bboxes 的 best iou, 并根据 best_iou 判断 roi 的类型
        如果 best_iou < C.rcnn_min_overlap, 忽略此 roi
        如果 C.rcnn_min_overlap <= best_iou < C.rcnn_max_overlap, 设置该 roi 的 class='bg',regr=[0,0,0,0]
        如果 C.rcnn_max_overlap <= best_iou, 设置 best_iou 对应的 gt_bbox 的 class 为该 roi 的 class, 并计算 regr
    :param rois:
    :param augmented_annotation:
    :param C:
    :param class_name_idx_mapping:
    :return: X,Y21,Y22 X: 表示合格的 rois 组成的数组作为 rcnn 的输入, (1, num_rois,4)
        Y21: 表示对所有 rois 分类的结果, (1, num_rois, num_classes)
        Y22: shape 为 (1,num_rois, (num_classes -1) * 4 * 2)
            前 (num_classes-1)*4 表示 rois 是否是 bg, 不参与回归的 loss 计算
            后 (num_classes-1)*4 表示对所有非 bg 的 rois 的回归梯度
        IoUs: 表示每一个 roi 和所有 gt_bboxes 的 best iou
    """
    bboxes = augmented_annotation['bboxes']
    (width, height) = (augmented_annotation['width'], augmented_annotation['height'])
    # get image dimensions for resizing
    (resized_width, resized_height) = data_generators.get_new_image_size(width, height, C.image_min_size)
    # ground truth bbox 在 feature map 上的坐标
    gt_bboxes = np.zeros((len(bboxes), 4))
    width_ratio = resized_width / float(width)
    height_ratio = resized_height / float(height)
    for bbox_idx, bbox in enumerate(bboxes):
        # get the GT box coordinates, and resize to account for image resizing
        gt_bboxes[bbox_idx, 0] = int(round(bbox['x1'] * width_ratio / C.rpn_stride))
        gt_bboxes[bbox_idx, 1] = int(round(bbox['x2'] * width_ratio / C.rpn_stride))
        gt_bboxes[bbox_idx, 2] = int(round(bbox['y1'] * height_ratio / C.rpn_stride))
        gt_bboxes[bbox_idx, 3] = int(round(bbox['y2'] * height_ratio / C.rpn_stride))

    x_roi = []
    # 分类模型的 y,是 one-hot 形式,(None,len(class_name_idx_mapping))
    y_class = []
    # 回归模型的 y,(None,len(class_name_idx_mapping - 1) * 4)
    y_regr = []
    # 标记 regr 是否参与 loss 计算
    y_regr_valid = []
    IoUs = []  # for debugging only

    for ix in range(rois.shape[0]):
        (x1, y1, x2, y2) = rois[ix, :]
        x1 = int(round(x1))
        y1 = int(round(y1))
        x2 = int(round(x2))
        y2 = int(round(y2))

        best_iou = 0.0
        best_bbox_idx = -1
        # roi 和每一个 bbox 计算 iou, 取 best_iou
        for bbox_idx in range(len(bboxes)):
            current_iou = data_generators.iou([gt_bboxes[bbox_idx, 0], gt_bboxes[bbox_idx, 2], gt_bboxes[bbox_idx, 1], gt_bboxes[bbox_idx, 3]],
                                              [x1, y1, x2, y2])
            if current_iou > best_iou:
                best_iou = current_iou
                best_bbox_idx = bbox_idx

        if best_iou < C.rcnn_min_overlap:
            continue
        else:
            w = x2 - x1
            h = y2 - y1
            x_roi.append([x1, y1, w, h])
            IoUs.append(best_iou)

            if C.rcnn_min_overlap <= best_iou < C.rcnn_max_overlap:
                # hard negative example
                class_name = 'bg'
            elif C.rcnn_max_overlap <= best_iou:
                class_name = bboxes[best_bbox_idx]['class']
                # ground truth bbox 的中心点坐标
                gcx = (gt_bboxes[best_bbox_idx, 0] + gt_bboxes[best_bbox_idx, 1]) / 2.0
                gcy = (gt_bboxes[best_bbox_idx, 2] + gt_bboxes[best_bbox_idx, 3]) / 2.0
                gw = gt_bboxes[best_bbox_idx, 1] - gt_bboxes[best_bbox_idx, 0]
                gh = (gt_bboxes[best_bbox_idx, 3] - gt_bboxes[best_bbox_idx, 2])
                # roi 的中心点坐标
                cx = x1 + w / 2.0
                cy = y1 + h / 2.0
                # 计算梯度
                tx = (gcx - cx) / float(w)
                ty = (gcy - cy) / float(h)
                tw = np.log(gw / float(w))
                th = np.log(gh / float(h))
            else:
                logger.error('roi={}'.format(best_iou))
                raise RuntimeError

        class_idx = class_name_idx_mapping[class_name]
        # 构建分类模型的 y
        class_label = len(class_name_idx_mapping) * [0]
        class_label[class_idx] = 1
        y_class.append(copy.deepcopy(class_label))
        regr_label = [0] * 4 * (len(class_name_idx_mapping) - 1)
        regr_label_valid = [0] * 4 * (len(class_name_idx_mapping) - 1)
        if class_name != 'bg':
            regr_pos = 4 * class_idx
            # UNCLEAR: C.classifier_regr_std 的作用是？
            sx, sy, sw, sh = C.classifier_regr_std
            regr_label[regr_pos:4 + regr_pos] = [sx * tx, sy * ty, sw * tw, sh * th]
            regr_label_valid[regr_pos:4 + regr_pos] = [1, 1, 1, 1]
            y_regr.append(copy.deepcopy(regr_label))
            y_regr_valid.append(copy.deepcopy(regr_label_valid))
        else:
            y_regr.append(copy.deepcopy(regr_label))
            y_regr_valid.append(copy.deepcopy(regr_label_valid))

    if len(x_roi) == 0:
        return None, None, None, None

    X = np.array(x_roi)
    Y_class = np.array(y_class)
    Y_regr = np.concatenate([np.array(y_regr_valid), np.array(y_regr)], axis=-1)
    # expand_dims 是为了 batch 那一维度
    return np.expand_dims(X, axis=0), np.expand_dims(Y_class, axis=0), np.expand_dims(Y_regr, axis=0), IoUs

def apply_regr(x, y, w, h, tx, ty, tw, th):
    """
    对单个矩形框进行修正
    :param x:
    :param y:
    :param w:
    :param h:
    :param tx:
    :param ty:
    :param tw:
    :param th:
    :return:
    """
    try:
        cx = x + w / 2.
        cy = y + h / 2.
        # rcx,rcy 表示修正后的矩形框的中心点 x,y
        # (rcx - cx) / w = tx
        rcx = tx * w + cx
        rcy = ty * h + cy
        # rw,rh 表示修正后的矩形框的宽和高
        # log(rw / w) = tw
        rw = math.exp(tw) * w
        rh = math.exp(th) * h
        # 左上角的点的 x,y
        rx = rcx - rw / 2.
        ry = rcy - rh / 2.
        rx = int(round(rx))
        ry = int(round(ry))
        rw = int(round(rw))
        rh = int(round(rh))
        return rx, ry, rw, rh
    except ValueError:
        return x, y, w, h
    except OverflowError:
        return x, y, w, h
    except Exception as e:
        logger.error(e)
        return x, y, w, h

def apply_regr_np(X, T):
    """
    根据回归梯度对所有的 anchor 进行校正
    :param X: anchor 的坐标, shape 为 (m, n, 4)
    :param T: regr 梯度, shape 为 (m, n, 4)
    :return: 修正后的 anchor 的坐标 array([x,y,w,h]...), shape 为 (m, n, 4)
    """
    try:
        height, width = X.shape[:2]
        # anchor 的坐标
        xa = X[:, :, 0]
        ya = X[:, :, 1]
        wa = X[:, :, 2]
        ha = X[:, :, 3]
        # regr 梯度
        tx = T[:, :, 0]
        ty = T[:, :, 1]
        tw = T[:, :, 2]
        th = T[:, :, 3]
        # anchor 的中心点坐标 x,y
        cxa = xa + wa / 2.
        cya = ya + ha / 2.
        # tx = (cxg - cxa) / wa
        # rectified_anchor 的中心点坐标 x,y
        cxg = tx * wa + cxa
        cyg = ty * ha + cya
        # tw = log(wg / wa)
        wg = np.exp(tw.astype(np.float64)) * wa
        hg = np.exp(th.astype(np.float64)) * ha
        xg = cxg - wg / 2.
        yg = cyg - hg / 2.

        xg = np.round(xg)
        yg = np.round(yg)
        wg = np.round(wg)
        hg = np.round(hg)

        xg = xg.reshape(height, width, 1)
        yg = yg.reshape(height, width, 1)
        wg = wg.reshape(height, width, 1)
        hg = hg.reshape(height, width, 1)
        rectified_anchor_coordinates = np.concatenate((xg, yg, wg, hg), axis=-1)
        return rectified_anchor_coordinates
    except Exception as e:
        logger.exception(e)
        return X

def non_max_suppression_fast(boxes, probs, overlap_thresh=0.9, max_boxes=300):
    """

    :param boxes: proposal 的 box 的坐标 np.array[[x1,y1,x2,y2]...]
    :param probs: box 包含物体的概率 np.array([p1,p2...])
    :param overlap_thresh: 覆盖阈值
    :param max_boxes: 最大返回的 box 个数
    :return: filtered_boxes and related probs
    """
    # code used from here: http://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/
    # https://www.pyimagesearch.com/2014/11/17/non-maximum-suppression-object-detection-python
    # if there are no boxes, return an empty list
    if len(boxes) == 0:
        return np.array([]), np.array([])

    # grab the coordinates of the bounding boxes
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    # https://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.testing.assert_array_less.html
    # 如果两数组的 shape 不同,会抛出异常
    # 如果 shape 相同但是有前一个数组的元素大于后一个数组的元素,也会抛出异常
    np.testing.assert_array_less(x1, x2)
    np.testing.assert_array_less(y1, y2)

    # if the bounding boxes integers, convert them to floats --
    # this is important since we'll be doing a bunch of divisions
    if boxes.dtype.kind == "i":
        boxes = boxes.astype("float")

    # initialize the list of picked indexes
    pick = []

    # calculate the areas
    # 数组的每个元素分别相乘 [area1,area2...]
    area = (x2 - x1) * (y2 - y1)

    # sort the bounding boxes, 返回的结果是下标数组,每一个元素的值为原数组的元素的下标
    # 如最后一个元素的值表示的是原数组最大值的下标,第一个元素的值表示原数组最小值的下标
    # https://docs.scipy.org/doc/numpy/reference/generated/numpy.argsort.html
    idxs = np.argsort(probs)
    # keep looping while some indexes still remain in the indexes list
    while len(idxs) > 0:
        # grab the last index in the indexes list and add the
        # index value to the list of picked indexes
        last = len(idxs) - 1
        # 最大 prob 的 box 在原 boxes 中的下标
        i = idxs[last]
        pick.append(i)

        # 计算最大 prob 的 box 和其他所有 box 的交集
        # x1 的 shape (num_boxes,), x1[i] 的 shape 为 (1,)
        # x1[idxs[:last]] 的 shape 为 (last, )
        # x1_intersection 的 shape 为 (last, )
        #  _____________________
        # |  (max_x1,max_y1)    |
        # |                 ____|________________
        # |                |    |                |
        # |                |    |                |
        # |                |    |                |
        # |________________|____|(min_x2,min_y2) |
        #                  |                     |
        #                  |_____________________|
        # x1[idxs[:last]] 其实这 last 个 x1 就已经按照 probs 排好序了
        x1_intersection = np.maximum(x1[i], x1[idxs[:last]])
        y1_intersection = np.maximum(y1[i], y1[idxs[:last]])
        x2_intersection = np.minimum(x2[i], x2[idxs[:last]])
        y2_intersection = np.minimum(y2[i], y2[idxs[:last]])
        w_intersection = np.maximum(0, x2_intersection - x1_intersection)
        h_intersection = np.maximum(0, y2_intersection - y1_intersection)
        # shape 为 (last,)
        area_intersection = w_intersection * h_intersection
        # 计算最大 prob 的 box 和其他所有 box 的并集
        area_union = area[i] + area[idxs[:last]] - area_intersection
        # compute the ratio of overlap, 就是 IOU
        # shape 为 (last,)
        overlap = area_intersection / (area_union + 1e-6)

        # delete all indexes from the index list that have
        # 删除当前最大 prob 的 box 和当前最大 prob 的 box 重叠超过阈值的其他 box
        # 这里因为 overlap 是按照 idxs 获得的, 所以 np.where 其实返回的 idxs 的下标
        idxs = np.delete(idxs, np.concatenate(([last],
                                               np.where(overlap > overlap_thresh)[0])))

        if len(pick) >= max_boxes:
            break

    # return only the bounding boxes that were picked using the integer data type
    boxes = boxes[pick].astype("int")
    probs = probs[pick]
    return boxes, probs

def rpn_to_roi(rpn_class, rpn_regr, C, max_rois=300, overlap_thresh=0.9):
    """
    对所有的 anchor, 根据 rpn 的分类和回归结果进行筛选和修正, 生成 roi
    :param rpn_class: rpn 分类的结果 (1,m,n,9) m,n 表示 feature map 的高和宽
    :param rpn_regr: rpn 回归的结果 (1,m,n,36)
    :param C: config 对象
    :param max_rois: 最大 roi 个数
    :param overlap_thresh:
    :return: 返回 rois 在 feature map 上的坐标 (x1,y1,x2,y2)
    """
    # 在生成 rpn target 时乘过 C.std_scaling
    rpn_regr = rpn_regr / C.std_scaling
    anchor_scales = C.anchor_scales
    anchor_ratios = C.anchor_ratios

    assert rpn_class.shape[0] == 1
    assert rpn_regr.shape[0] == 1
    # feature map 的高,宽和深, d=9
    (m, n, d) = rpn_class.shape[1:]

    anchor_idx = 0
    # (m,n,9,4)
    A = np.zeros((m, n, d, 4))

    for anchor_scale in anchor_scales:
        for anchor_ratio in anchor_ratios:
            # anchor 在特征图上的宽
            anchor_width = (anchor_scale * anchor_ratio[0]) / C.rpn_stride
            # anchor 在特征图上的高
            anchor_height = (anchor_scale * anchor_ratio[1]) / C.rpn_stride
            # 该 anchor 的回归梯度
            regr = rpn_regr[0, :, :, 4 * anchor_idx:4 * anchor_idx + 4]
            # X: (m, n) [[0,1...,n-1],[0,1...,n-1]...]
            # Y: (m, n) [[0,0...,0],[1,1...,1],...[m-1,m-1...,m-1]]
            # 其实生成的是一个坐标,特征图上的每个点都有了相应的坐标
            # X 表示的是所有 anchor 的中心点的 x 坐标
            # Y 表示的是所有 anchor 的中心点的 y 坐标
            X, Y = np.meshgrid(np.arange(n), np.arange(m))
            # A 的最后一维的四个值表示 anchor 的 (x,y,w,h), 其中 x,y 表示左上方点的坐标
            A[:, :, anchor_idx, 0] = X - anchor_width / 2
            A[:, :, anchor_idx, 1] = Y - anchor_height / 2
            A[:, :, anchor_idx, 2] = anchor_width
            A[:, :, anchor_idx, 3] = anchor_height
            # 根据回归梯度修正
            A[:, :, anchor_idx, :] = apply_regr_np(A[:, :, anchor_idx, :], regr)
            # 宽度最小为 1
            A[:, :, anchor_idx, 2] = np.maximum(1, A[:, :, anchor_idx, 2])
            # 高度最小为 1
            A[:, :, anchor_idx, 3] = np.maximum(1, A[:, :, anchor_idx, 3])
            # A[:,:,anchor_idx, 2] 变成了 x2, 即右下角 x 坐标
            A[:, :, anchor_idx, 2] += A[:, :, anchor_idx, 0]
            # A[:,:,anchor_idx, 3] 变成了 y2, 即右下角 y 坐标
            A[:, :, anchor_idx, 3] += A[:, :, anchor_idx, 1]
            # x1 最小为 0
            A[:, :, anchor_idx, 0] = np.maximum(0, A[:, :, anchor_idx, 0])
            # y1 最小为 0
            A[:, :, anchor_idx, 1] = np.maximum(0, A[:, :, anchor_idx, 1])
            # x2 最大为 n - 1
            A[:, :, anchor_idx, 2] = np.minimum(n - 1, A[:, :, anchor_idx, 2])
            # y2 最大为 m - 1
            A[:, :, anchor_idx, 3] = np.minimum(m - 1, A[:, :, anchor_idx, 3])

            anchor_idx += 1

    # (m*n*9,4)
    all_rois = np.reshape(A, (-1, 4))
    # (m*n*9,)
    all_probs = rpn_class.reshape((-1))

    x1 = all_rois[:, 0]
    y1 = all_rois[:, 1]
    x2 = all_rois[:, 2]
    y2 = all_rois[:, 3]

    # 删除不合理的矩形存在, 左上角的坐标不能大于右下角的坐标
    idxs = np.where((x1 - x2 >= 0) | (y1 - y2 >= 0))
    all_rois = np.delete(all_rois, idxs, axis=0)
    all_probs = np.delete(all_probs, idxs, axis=0)

    rois = non_max_suppression_fast(all_rois, all_probs, overlap_thresh=overlap_thresh, max_boxes=max_rois)[0]

    return rois


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

# In[4]:

# required libraries
import numpy as np
import matplotlib.pyplot as plt
import panel as pn
import pandas as pd
pn.extension("katex")

#from IPython.display import Image, Video

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

# ## Lecture 1 - Course Introduction/Water Cycle ##
#
# _(The contents presented in this section were re-developed principally by Dr. P. K. Yadav. The original contents were developed by Prof. Rudolf Liedl)_

# ### Contents of “Groundwater Module ” ###
#
# The entire contents of this interactive book can be listed with the following points:
#
# + general overview
# + types of aquifers, properties of aquifers
# + forces and pressure in the subsurface
# + laws of groundwater flow and some applications (e.g. groundwater wells)
# + quantification of aquifer parameter values via pumping tests
# + transport of chemicals (solutes) in groundwater
# + retardation and degradation of chemicals in groundwater
# + groundwater modelling
#

# ### Suggested Literature: ###
#
#
# + Brassington R. (1988): Field hydrogeology, Wiley & Sons.
#
# + Domenico P. A., Schwartz F. W. (1990): Physical and chemical hydrogeology, Wiley & Sons.
#
# + Fetter C. W. (2001): Applied hydrogeology, Prentice Hall.
#
# + Freeze R. A., Cherry J. A. (1979): Groundwater, Prentice Hall.
#
# + Heath R. C. (1987): Basic groundwater hydrology, USGS Water Supply Paper 2220.
#
# + Price M. (1996): Introducing groundwater, Chapman and Hall.
#
# Additional literature details are provided in the text when used.

# ### What is Hydrogeology? ###
#
# _**Hydrogeology**_ is the study of the laws governing the movement of subterranean water, the mechanical, chemical, and thermal interaction of this water with the porous solid, and the transport of energy and chemical constituents by the flow.
# (Domenico and Schwartz, 1990)
#
#
# The dominant reliance of groundwater for the drinking globally has made hydrogeology a very important academic course. Also, it is a very important research field. Therefore, several **techniques** and **methods** are now available to explore and understand **Hydrogeological Process**. The methods and techniques can be broadly categorized to:
#
# 1. Field works
# 2. Laboratory experiments
# 3. Computer modeling
#
#
# _Computer modelling_ is often the most economical method but its usefullness rely of data obtained from _Field works_ and _Laboratory experiments._ Thus, the sequence of techniques/methods to be adopted depends on the available site information.
#

# In[2]:

im1 = pn.pane.PNG("images/L01_f_1c.png", width=250)
im2 = pn.pane.PNG("images/L01_f_1b.png", width=275)
im3 = pn.pane.PNG("images/L01_f_1a.png", width=280)

pn.Row(im1, im2, im3)

# ### Example: Groundwater Extraction Well
#
# Groundwater is extracted using a groundwater well applying _hydrogeological_ methods and techniques. The procedure followed can be summarized in the following steps:
#
# 1. The appropriate extraction location is identified
# 2. Drilling machine are used to obtain sub-surface structure, i.e. or well logs are obtained. The process is also called well logging.
# 3. Well logs are studied in detail to identify the characteristics of the subsurface- e.g., how thick is the aquifer or identify environmental consequence of water extraction.
# 4. The construction of well begins
#
# Groundwater extraction using well is a challenge when aquifers are located very deep from the surface, e.g., in deserts.
#

# In[3]:

video1 = pn.pane.Video("images/L01_f_2.mp4", width=600, height=400, loop=False)
video1

# In[4]:

#gif_pane = pn.pane.GIF('images/L01_f_2.gif', width=500)
#gif_pane
video2 = pn.pane.Video("images/L01_f_3.mp4", width=600, height=400, loop=False)
#Video("images/L01_f_3.mp4", width=600, embed=True)
video2

pn1 = pn.pane.Markdown("""
**Wells** are placed on the layer or that aquifer part which allows feasible extraction of groundwater.
The extraction leads to drop of groundwater level. To ensure that there is sustainable extraction,
the drops in the level has to be monitored. Quite often this is done through _computer modelling_. There
already exists several computer models that can use the well logs data (also called Borehole) and provide
good estimations of the effects due to extraction. The _computer models_ are also able to predict the effects
at larger scales, e.g., regional scales. _Computer models_ are oftenly used these days be agencies to determine
quantities such as **travel time**, **capture zones** or obtain **isochrones**, which are used for deciding on
groundwater extraction programmes.
""")

pn.Row(pn1, video2)

# ### Groundwater and Global Water Cycle ###
#
# Water bodies that exist on earth is connected, and they function as a cycle, called **Global Water Cycle**. It is estimated that over 57, 700 Km$^3$ of water actively participates in the cycle each year. **Precipitation** and **evaporation** are the two main components of the cycle in which **temperature** plays the critical role. In the cycle, **Groundwater** receives water from _precipitation,_ It then contributes to _evaporation_ through subsurface flow or through mostly human intervention (e.g., use for drinking water).
#
# The water cycle provides an approach to judge the sustainability of groundwater extraction. The sustainability of extraction can be obtained if extraction rate approximately equals the replenishing rate. Often the replenishing rate of groundwater is much slower and this has led to groundwater stress in many parts of the world.

# In[2]:

#gif_pane = pn.pane.GIF('images/L01_f_2.gif', width=500)
#gif_pane
video3 = pn.pane.Video("images/L01_f_4.mp4", width=600, height=400, loop=False)
#Video("images/L01_f_3.mp4", width=600, embed=True)
video3

# In[3]:

#gif_pane = pn.pane.GIF('images/L01_f_2.gif', width=500)
#gif_pane
fig5 = pn.pane.PNG("images/L01_f_5.png", width=600)
#Video("images/L01_f_3.mp4", width=600, embed=True)

pn1 = pn.pane.Markdown("""

### Water balance by continents

Groundwater receives water from the _infiltration_ of **runoff** water.

""")

pn.Row(pn1, fig5)

# ### The Hydrological Balance ###
#
# Since _groundwater_ is part of the global water cycle, the balance of the cycle becomes an important topic. In general:
#
# + The _hydrological balance_ provides a relationship between various flow rates for a certain area. It is based on the conservation of water volume.
# + expressed in words:  _inflow_ equals _outflow_ plus _change in storage_
# + expressed by a formula:
#
# $$
# P = ET + R + \Delta S
# $$
#
# ```{margin} Where,
# where, <br>
# $P$ = _Precipitation, <br> $ET$ = Evapotranspiration, <br> $R$ = Runoff,_  and <br> $\Delta S$ = _Change in Storage_
#
# ```
#
# The _change in storage_ can be interpreted in the following way:
#
# + change in storage $\Delta S > 0$ : Water volume is increasing with time in the investigation area.
#
# + change in storage $\Delta S < 0$:Water volume is decreasing with time in the investigation area.
#
# + change in storage $\Delta S = 0$: Water volume does not change with time in the investigation area (steady-state or stationary situation, i.e. inflow equals outflow).

# ### Water Volume
#
#
# **So how much water do we have?**

# It is estimated* that the total volume of water on Earth amounts to ca. 1 358 710 150 km$^3$ ($\approx$ 1018 m$^3$).
#
#
#
# The total volume of fresh water on Earth amounts to ca. $38\times 10^6$ km$^3$ ($\approx$ 1016 m$^3$).

#
# _*Gleick P. (1996): Water re-sources, in: Schneider S. H. (ed.), Encyclopedia of climate and weather 2, Oxford Univ. Press._
#

# ### Volume of Available Fresh Water ###
#
# **Fresh water** are water with low concentrations of dissolved salts and other total dissolved solids, i.e., sea/ocean water or brackish water are not fresh water. Human activities (drinking water) are directly dependent on fresh activities.
#
# **So how much _fresh water_ do we have?**
#
#
# It is estimated* that the total volume of available fresh water (liquid) on Earth amounts to ca. 8 831 600 km$^3$ ($\approx$ 1016 m$^3$).
#
#
#
#
# _*Gleick P. (1996): Water re-sources, in: Schneider S. H. (ed.), Encyclopedia of climate and weather 2, Oxford Univ. Press._

# ### Continental distribution of fresh water components  ###
#
#
#

# ### Volume and Mass Budget
#
#
# Very basics of volume and mass budget - let us start with _budget._
#
# **Budget** = quantitative comparison of _growth_ (or _production_) and _loss_ in a system
#
# Budgets can be put together for various quantities:
# + energy
# + mass $\leftarrow$ needed to quantify transport of solutes in groundwater
# + volume $\leftarrow$  needed to quantify groundwater flow
# + momentum
# + electric charge
# + number of inhabitants
# + animal population
# + money (bank account!)
# + and many others
#
# In this course we focus on _Mass Budget_ and _Volume Budget._

# ### Volume Budget ###
#
# As discussed in the last topic a _budget_ represents the change (e.g., growth and loss). Thus, it is more suitable to quantify the **volume budget** in terms of a change, representing two different states (e.g., time $(t)$). More formally, the **volume budget** ($\Delta V$) can be obtained from:
#
# $$
# \Delta V = Q_{in} \cdot \Delta t - Q{out} \cdot \Delta t
# $$
#

# with,

#
# $\Delta t$ = time interval [T]

# $\Delta V$ = change of volume in the system [L$^3$]

# $Q_{in}$ = volumetric rate of flow into the system [L$^3$/T]

# $Q_{out} =$ volumetric rate of flow out of the system [L$^3$/T]

#
# The following points have to be considered when using the above equation:
#
# + Inflow and outflow may each consist of several individual components.

#
# + $\Delta V = 0$ (no change in  volume) is tantamount to steady-state or stationary (= time-independent) conditions.

#
# + For steady-state conditions we have: $Q_{in} = Q_{out}$
#

# #### Water Budget for a Catchment ####
#
# The equation provided for _volume budget_ looks simple but in practice it is very complicated as several _inflow_ and _outflow_ components must be considered. Quantifying these components can be a challenging task.
#
# For quantifying water budget of a catchment, one has to consider the following components:
#
# **To be considered:**
# + precipitation
# + evapotranspiration
# + surface runoff
# + subsurface runoff
#
# Among the above components, quantification of evapotranspiration and subsurface runoff have very high level of uncertainties.
#
# ```{image} images/L01_f_9.png
# :height: 500px
# :align: center
# :name: Water-Budget
# ```

# ### Example: Estimation of Subsurface Runoff ###
#
# Most numbers used in the example do not refer to the catchment shown before!
#
# To calculate the following four-steps are to be followed:
#
# + Step 1: determine rate of inflow in m³/a
# + step 2: determine rate of outflow due to evapotranspiration (ET.A) in m³/a
# + Step 3: express rate of outflow due to surface runoff in m³/a
# + step 4: determine rate of outflow due to subsurface runoff
#
# An Example:

# For given data, determine the rate of outflow Qout,sub due to subsurface runoff for steady-state conditions

# In[4]:

A = 4500 # km², catchment area
P = 550 # mm/a, precipitation
ET = 200 # mm/a, evapotranspiration
Qout_surf = 40 # m³/s, surface runoff
Delta_V = 0 # m³, change in volume = 0 Steady-state conditions

#Volume budget in this example: P·A = ET·A + Qout,surf + Qout,sub

#Step 1
Qin = P*A*10**3  #m³/a, 10^3 for unit conversion

#step 2:
ET_A = ET*A*10**3 #m³/a, 10^3 for unit conversion

#Step 3:
Qout_surf = Qout_surf *365*24*3600 #  m³/a

# step 4:
Qout_sub = Qin - ET_A - Qout_surf # m³/a

print("The rate of inflow, Qin is {0:1.1E}".format(Qin),"m\u00b3/a \n"); print("The outflow rate due to Evapotranspiration is {0:1.1E}".format(ET_A),"m\u00b3/a \n")
print("The surface outflow rate, Q_out_surf in m\u00b3/a is {0:1.1E}".format(Qout_surf),"m\u00b3/a \n");print("The subsurface outflow rate, Qout_surf in m\u00b3/a is {0:1.1E}".format(Qout_sub),"m\u00b3/a \n")

# ### Mass Budget ###
#
# The **mass budget** is quantified similar to the _volume budget._ Mathematically, the _mass budget_ is:
#
# $$\Delta M = J_{in}\cdot \Delta t - J_{out} \cdot \Delta t$$
#
# with

# $\Delta t$ = time interval [T]

# 	$\Delta M$ = change of mass in the system [M]

# 	$J_{in}$ = rate of mass flow into the system [M/T]

# 	$J_{out}$ = rate of mass flow out of the system [M/T]
#
# Similar to _volume budget,_ the following points have to be considered in quantifying mass budget:
#
# + Inflow and outflow may each consist of several individual components.

# + $\Delta M$ = 0 (no change in mass) is tantamount to steady-state or stationary
(= time-independent) conditions.

# + For steady-state conditions we have: $J_{in}$= $J_{out}$

# ### Example of Mass Budget: Radioactive Decay ###
#
# Consider a decay chain comprising of the three chemicals: **A**, **B** and **C**
#
# + decay chain: A $\rightarrow$ B $\rightarrow$ C

# + 30% of $\text{A}$ and 20% of $\text{B}$  decay each year.

#
# + decay rate of $\text{A}$   = production rate of $\text{B}$   = $0.3 \cdot a^{-1}\cdot M_A$

#
# + decay rate of $\text{B}$ = production rate of $\text{C}$ = $0.2\cdot a^{-1}\cdot M_B$

#
# + mass budgets for $\text{A}$, $\text{B}$ and $\text{C}$:

#
#
#
# \begin{equation*}
# \begin{split}
# \Delta M_A &= 0.3 \text{ a $^{-1}$ } \cdot M_A  \cdot \Delta t  \\
# \Delta M_B  &= 0.3 \text{a$^{-1}$} \cdot M_A  \cdot \Delta t - 0.2 \text{ a$^{-1}$} \cdot M_B  \cdot \Delta t \\
# \Delta M_C &= 0.2 \text{a$^{-1}$} \cdot M_B  \cdot \Delta t
# \end{split}
# \end{equation*}
#
#
# + Similar equations hold for quantitative descriptions of some chemical reactions which correspond to the type A $\rightarrow$ B $\rightarrow$ C

# In[2]:

def mass_bal(n_simulation, MA, MB, MC, R_A, R_B):

    A = np.zeros(n_simulation)
    B = np.zeros(n_simulation)
    C = np.zeros(n_simulation)
    time  = np.arange(n_simulation)

    for i in range(0,n_simulation-1):
        A[0] = MA
        B[0] = MB
        C[0] = MC
        A[i+1] = A[i]-R_A*A[i]
        B[i+1] = B[i]+R_A*A[i]-R_B*B[i]
        C[i+1] = C[i]+R_B*B[i]
        summ = A[i]+B[i]+C[i]

    d = {"Mass_A": A, "Mass_B": B, "Mass_C": C, "Total Mass": summ}
    df = pd.DataFrame(d) # Generating result table
    label = ["Mass A (g)", "Mass B (g)", "Mass C (g)"]
    fig = plt.figure(figsize=(6,4))
    plt.plot(time, A, time, B, time, C, linewidth=3);  # plotting the results
    plt.xlabel("Time [Time Unit]"); plt.ylabel("Mass [g]") # placing axis labels
    plt.legend(label, loc=0);plt.grid(); plt.xlim([0,20]); plt.ylim(bottom=0) # legends, grids, x,y limits
    plt.show() # display plot

    df_pane = pn.pane.DataFrame(df)
    return print(df.round(2))

N = widgets.BoundedIntText(value=20,min=0,max=100,step=1,description= 'Δ t (day)',disabled=False)

A = widgets.BoundedFloatText(value=100,min=0,max=1000.0,step=1,description='MA (kg)',disabled=False)

B = widgets.BoundedFloatText(value=5,min=0,max=1000.0,step=1,description='MB (kg)',disabled=False)

C = widgets.BoundedFloatText(value=10,min=0,max=1000,step=0.1,description='MC (kg)',disabled=False)

RA = widgets.BoundedFloatText(value=0.2,min=0,max=100,step=0.1,description='RA (day-1 )',disabled=False)

RB = widgets.BoundedFloatText(value=0.2,min=0,max=100,step=0.1,description='RB (day-1 )',disabled=False)

interactive_plot = widgets.interactive(mass_bal, n_simulation = N, MA=A, MB=B, MC=C, R_A=RA, R_B=RB,)
output = interactive_plot.children[-1]
#output.layout.height = '350px'
interactive_plot

# ### Comparison of Mass and Volume Budgets ###
#
# **mass budget**:	$\Delta M = J_{in} \cdot \Delta t - J_{out} \cdot \Delta t$
#
# **volume budget**:	$\Delta V = Q_{in} \cdot \Delta t - Q_{out} \cdot \Delta t $
#
#
# + Mass and volume budgets are equivalent if there is no change of density $\rho$ [M/L$^3$] with time. In this case the well known relationship $\Delta M$ = $\rho \cdot \Delta V$ holds and each equation given above can be directly transformed into the other one.
#
#
# + If density changes have to be considered (e.g. for gas flow), the mass budget equation remains valid but the volume budget equation must be modified because $\Delta M = \rho \cdot \Delta V + \Delta \rho \cdot V$ with $\Delta \rho$= change in density.
#
#
# + Cases with changing density have proven to be more easily tractable if the mass budget equation is used.


# MIT License
#
# Copyright (C) IBM Corporation 2019
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
# Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
The Wishart mechanism in differential privacy, for producing positive semi-definite perturbed second-moment matrices
"""
from numbers import Real
import warnings

import numpy as np

from diffprivlib.mechanisms.base import DPMechanism
from diffprivlib.utils import copy_docstring

class Wishart(DPMechanism):
    r"""
    The Wishart mechanism in differential privacy.

    Used to achieve differential privacy on 2nd moment matrices.

    Paper link: https://ieeexplore.ieee.org/abstract/document/7472095/

    .. deprecated:: 0.4
        `Wishart` is deprecated and will be removed in version 0.5. The Wishart mechanism has been shown not to satisfy
        differential privacy, and its continued use is not recommended.

    Parameters
    ----------
    epsilon : float
        The value of epsilon for achieving :math:`(\epsilon,\delta)`-differential privacy with the mechanism.  Must be
        > 0.

    sensitivity : float
        The maximum l2-norm of the data.  Must be >= 0.

    """
    def __init__(self, epsilon, sensitivity):
        warnings.warn("The Wishart mechanism has been shown not to satisfy differential privacy as originally "
                      "proposed.  As a result, the Wishart mechanism is deprecated as of version 0.4, and will be "
                      "removed in version 0.5.  To get a differentially private estimate of a covariance matrix, it is "
                      "recommended to use `models.utils.covariance_eig` instead.", DeprecationWarning)

        super().__init__(epsilon=epsilon, delta=0.0)
        self.sensitivity = self._check_sensitivity(sensitivity)

        self._rng = np.random.default_rng()

    @classmethod
    def _check_epsilon_delta(cls, epsilon, delta):
        if not delta == 0:
            raise ValueError("Delta must be zero")

        return super()._check_epsilon_delta(epsilon, delta)

    @classmethod
    def _check_sensitivity(cls, sensitivity):
        if not isinstance(sensitivity, Real):
            raise TypeError("Sensitivity must be numeric")

        if sensitivity < 0:
            raise ValueError("Sensitivity must be non-negative")

        return float(sensitivity)

    def _check_all(self, value):
        super()._check_all(value)
        self._check_sensitivity(self.sensitivity)

        if not isinstance(value, np.ndarray):
            raise TypeError("Value to be randomised must be a numpy array, got %s" % type(value))
        if value.ndim != 2:
            raise ValueError("Array must be 2-dimensional, got %d dimensions" % value.ndim)
        if value.shape[0] != value.shape[1]:
            raise ValueError("Array must be square, got %d x %d" % (value.shape[0], value.shape[1]))

        return True

    @copy_docstring(DPMechanism.bias)
    def bias(self, value):
        raise NotImplementedError

    @copy_docstring(DPMechanism.variance)
    def variance(self, value):
        raise NotImplementedError

    def randomise(self, value):
        """Randomise `value` with the mechanism.

        Parameters
        ----------
        value : numpy array
            The data to be randomised.

        Returns
        -------
        numpy array
            The randomised array.

        """
        self._check_all(value)

        scale = 1 / 2 / self.epsilon
        n_features = value.shape[0]

        noise_array = self._rng.standard_normal((n_features, n_features + 1)) * scale * self.sensitivity
        noise_array = np.dot(noise_array, noise_array.T)

        return value + noise_array


import math
import torch.nn.functional as F
import collections
import torch
import numpy
from collections import defaultdict
from dataclasses import dataclass
from fairseq import metrics
#
# from fairseq.sequence_generator  import SequenceGenerator
# from sequence_generator import
from fairseq import utils #, bleu
from fairseq.dataclass import FairseqDataclass
from fairseq.criterions import FairseqCriterion, register_criterion
from omegaconf import II
from fairseq.models import FairseqMultiModel
from fairseq.sequence_generator import SequenceGenerator

@register_criterion('v2')
class V2Criterion(FairseqCriterion):

    def __init__(self, task, sentence_avg): #(self, args, task):
        super().__init__(task)
        self.cnt=0

    # def forward(self, model, sample, reduce=True):
    #     """Compute the loss for the given sample.
    #
    #     Returns a tuple with three elements:
    #     1) the loss
    #     2) the sample size, which is used as the denominator for the gradient
    #     3) logging outputs to display while training
    #     """
    #     net_output = model(**sample['net_input'])
    #     lprobs = model.get_normalized_probs(net_output, log_probs=True)
    #     lprobs = lprobs.view(-1, lprobs.size(-1))
    #     target = model.get_targets(sample, net_output).view(-1)
    #     loss = F.nll_loss(lprobs, target, size_average=False,
    #                       ignore_index=self.padding_idx,
    #                       reduce=reduce)
    #     sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']
    #     logging_output = {
    #         'loss': utils.item(loss.data) if reduce else loss.data,
    #         'ntokens': sample['ntokens'],
    #         'sample_size': sample_size,
    #     }
    #     return loss, sample_size, logging_output

    def forward(self, model, sample, reduce=True):
        # sample mode
        #print('!!!RL loss.')
        model.eval()
        # src_dict = self.task.source_dictionary
        tgt_dict = self.task.target_dictionary
        eos_idx = self.task.target_dictionary.eos()
        sample_beam = 4 #self.args.sample_beam

        translator = SequenceGenerator([model], tgt_dict=tgt_dict,beam_size=sample_beam)

        translator.cuda()
        ct = 0
        translations = []

        s = utils.move_to_cuda(sample)
        input = s['net_input']
        max_len = 200
        with torch.no_grad():
            hypos = translator.generate([model],s)
        for i, id in enumerate(s['id'].data):
            src = input['src_tokens'].data[i, :]
            # remove padding from ref
            ref = utils.strip_pad(s['target'].data[i, :], tgt_dict.pad()) if s['target'] is not None else None
            translations.append((id, src, ref, hypos[i]))

            ct += 1
        # print("sample batch size:", ct)

        model.train()

        # MLE loss
        mle_net_output = model(**sample['net_input'])
        mle_lprobs = model.get_normalized_probs(mle_net_output, log_probs=True)
        mle_lprobs = mle_lprobs.view(-1, mle_lprobs.size(-1))
        mle_target = model.get_targets(sample, mle_net_output).view(-1)
        mle_loss = F.nll_loss(mle_lprobs, mle_target, size_average=False,
                              ignore_index=self.padding_idx, reduce=reduce)
        mle_tokens = sample['ntokens']
        avg_mle_loss = mle_loss / mle_tokens
        self.cnt  += 1
        if self.cnt % 4 !=0 :
            total_loss = avg_mle_loss
            total_tokens =  sample["ntokens"]
            logging_output = {
                'loss': utils.item(total_loss.data),
                'ntokens': total_tokens,
                'sample_size': total_tokens,
            }
    #         print('total: ',total_loss)
            return total_loss, total_tokens, logging_output

        print('avg_mle_loss:', avg_mle_loss)
        # RL loss
        batch_rl_loss = 0
        batch_tokens = 0
        sample_ind = 0
        for sample_id, src_tokens, tgt_tokens, hypos in translations:
            # calculate bleu
            sample_ind += 1
            rewards = torch.Tensor(sample_beam).float().cuda()
            logprobs = torch.Tensor(sample_beam).float().cuda()
            for i in range(sample_beam):
                hypo = hypos[i]
                trans_tokens = hypo['tokens']
                rewards[i] = self.compute_gleu(tgt_tokens.cpu(), trans_tokens.cpu(), max_order=4, gram=0).cuda()
                # one_sample loss calculation
                tgt_input_tokens = trans_tokens.new(trans_tokens.shape).fill_(0)
                assert trans_tokens[-1] == eos_idx
                tgt_input_tokens[0] = eos_idx
                tgt_input_tokens[1:] = trans_tokens[:-1]
                train_sample = {
                    'net_input': {
                        'src_tokens': src_tokens.view(1, -1),
                        'src_lengths': torch.LongTensor(src_tokens.numel()).view(1, -1),
                        'prev_output_tokens': tgt_input_tokens.view(1, -1),
                    },
                    'target': trans_tokens.view(1, -1)
                }
                train_sample = utils.move_to_cuda(train_sample)
                net_output = model(**train_sample['net_input'])
                lprobs = model.get_normalized_probs(net_output, log_probs=True)
                lprobs = lprobs.view(-1, lprobs.size(-1))
                target = model.get_targets(train_sample, net_output).view(-1, 1)
                non_pad_mask = target.ne(tgt_dict.pad())
                lprob = -lprobs.gather(dim=-1, index=target)[non_pad_mask]
                logprobs[i] = torch.sum(lprob)
                ntokens = len(train_sample['target'])
                batch_tokens += ntokens
            rl_loss = torch.sum(logprobs * (rewards - rewards.mean()))  # one sample loss
            batch_rl_loss += rl_loss

        avg_rl_loss = batch_rl_loss / batch_tokens
        print('avg_rl_loss:', avg_rl_loss)
        if True:
            total_loss = avg_mle_loss +  avg_rl_loss
            total_tokens = batch_tokens + mle_tokens
        else:
            total_loss = avg_rl_loss
            total_tokens = batch_tokens
        logging_output = {
            'loss': utils.item(total_loss.data),
            'ntokens': total_tokens,
            'sample_size': total_tokens,
        }
        print('total: ',total_loss)
        return total_loss, total_tokens, logging_output

    def _get_ngrams(self, segment, max_order):
        ngram_counts = collections.Counter()
        for order in range(1, max_order + 1):
            for i in range(0, len(segment) - order + 1):
                ngram = tuple(segment[i:i + order])
                ngram_counts[ngram] += 1
        return ngram_counts

    def compute_gleu(self, reference_corpus, translation_corpus, max_order=4, gram=0, smooth=False):
        scores = torch.zeros(max_order)
        reference_array = numpy.array(reference_corpus)
        translation_array = numpy.array(translation_corpus)
        matches_by_order = [0] * max_order
        possible_matches_by_order_ref = [0] * max_order
        possible_matches_by_order_trans = [0] * max_order
        reference_length = 0
        translation_length = 0
        reference_length += reference_array.shape[0]
        translation_length += translation_array.shape[0]
        merged_ref_ngram_counts = collections.Counter()
        merged_ref_ngram_counts |= self._get_ngrams(reference_array, max_order)
        translation_ngram_counts = self._get_ngrams(translation_array, max_order)
        overlap = translation_ngram_counts & merged_ref_ngram_counts
        for ngram in overlap:
            matches_by_order[len(ngram)-1] += overlap[ngram]
        for order in range(1, max_order+1):
            possible_matches_trans = translation_length - order + 1
            if possible_matches_trans > 0:
                possible_matches_by_order_trans[order-1] += possible_matches_trans
            possible_matches_ref = reference_length - order + 1
            if possible_matches_ref > 0:
                possible_matches_by_order_ref[order-1] += possible_matches_ref
        precisions = [0] * max_order
        recalls = [0] * max_order

        for i in range(0, max_order):
            if smooth:
                precisions[i] = ((matches_by_order[i] + 1.) / (possible_matches_by_order_trans[i] + 1.))
                recalls[i] = ((matches_by_order[i] + 1.) / (possible_matches_by_order_ref[i] + 1.))
            else:
                if possible_matches_by_order_trans[i] > 0:
                    precisions[i] = (float(matches_by_order[i]) / possible_matches_by_order_trans[i])
                else:
                    precisions[i] = 0.0

                if possible_matches_by_order_ref[i] > 0:
                    recalls[i] = (float(matches_by_order[i]) / possible_matches_by_order_ref[i])
                else:
                    recalls[i] = 0.0
        for i in range(max_order):
            scores[i] = min(precisions[i],recalls[i])

        if False :#self.args.modgleu:
            if reference_length < max_order and translation_length < max_order:
                order = max(reference_length, translation_length)
                scores = scores[0:order]
            else:
                order = max_order
        else:
            order = max_order
        if gram == 0:
            if min(scores) > 0:
                log_scores = torch.log(scores)
                p_log_sum = torch.sum((1. / order) * log_scores)
                geo_mean = torch.exp(p_log_sum)
                return geo_mean
            else:
                return torch.tensor(0.0)
        else:
            if scores[gram] > 0:
                return scores[gram]
            else:
                return torch.tensor(0.0)

    @staticmethod
    def aggregate_logging_outputs(logging_outputs):
        """Aggregate logging outputs from data parallel training."""
        loss_sum = sum(log.get('loss', 0) for log in logging_outputs)
        ntokens = sum(log.get('ntokens', 0) for log in logging_outputs)
        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)
        agg_output = {
            'loss': loss_sum / sample_size / math.log(2),
            'sample_size': sample_size,
        }
        if sample_size != ntokens:
            agg_output['nll_loss'] = loss_sum / ntokens / math.log(2)
        return agg_output

    @staticmethod
    def reduce_metrics(logging_outputs) -> None:
        """Aggregate logging outputs from data parallel training."""
        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)

        # we divide by log(2) to convert the loss from base e to base 2
        metrics.log_scalar(
            "loss", loss_sum / sample_size / math.log(2), sample_size, round=3
        )
        if sample_size != ntokens:
            metrics.log_scalar(
                "nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3
            )
            metrics.log_derived(
                "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
            )
        else:
            metrics.log_derived(
                "ppl", lambda meters: utils.get_perplexity(meters["loss"].avg)
            )

    @staticmethod
    def logging_outputs_can_be_summed() -> bool:
        """
        Whether the logging outputs returned by `forward` can be summed
        across workers prior to calling `reduce_metrics`. Setting this
        to True will improves distributed training speed.
        """
        return True


# Introduction

Classical mechanics is a topic which has been taught intensively over
several centuries. It is, with its many variants and ways of
presenting the educational material, normally the first **real** physics
course many of us meet and it lays the foundation for further physics
studies. Many of the equations and ways of reasoning about the
underlying laws of motion and pertinent forces, shape our approaches and understanding
of the scientific method and discourse, as well as the way we develop our insights
and deeper understanding about physical systems.

There is a wealth of
well-tested (from both a physics point of view and a pedagogical
standpoint) exercises and problems which can be solved
analytically. However, many of these problems represent idealized and
less realistic situations.  The large majority of these problems are
solved by paper and pencil and are traditionally aimed
at what we normally refer to as continuous models from which we may find an analytical solution.  As a consequence,
when teaching mechanics, it implies that we can seldomly venture beyond an idealized case
in order to develop our understandings and insights about the
underlying forces and laws of motion.

On the other hand, numerical algorithms call for approximate discrete
models and much of the development of methods for continuous models
are nowadays being replaced by methods for discrete models in science and
industry, simply because **much larger classes of problems can be addressed** with discrete models, often  by simpler and more
generic methodologies.

As we will see below, when properly scaling the equations at hand,
discrete models open up for more advanced abstractions and the possibility to
study  real life systems, with the added bonus that we can explore and
deepen our basic understanding of various physical systems

Analytical solutions are as important as before. In addition, such
solutions provide us with invaluable benchmarks and tests for our
discrete models. Such benchmarks, as we will see below, allow us
to discuss possible sources of errors and their behaviors.  And
finally, since most of our models are based on various algorithms from
numerical mathematics, we have a unique oppotunity to gain a deeper
understanding of the mathematical approaches we are using.

With computing and data science as important elements in essentially
all aspects of a modern society, we could  then try to define Computing as
**solving scientific problems using all possible tools, including
symbolic computing, computers and numerical algorithms, and analytical
paper and pencil solutions**.
Computing provides us with the tools to develope our own understanding of the scientific method by enhancing algorithmic thinking.

The way we will teach this course reflects
this definition of computing. The course contains both classical paper
and pencil exercises as well as  computational projects and exercises. The
hope is that this will allow you to explore the physics of systems
governed by the degrees of freedom of classical mechanics at a deeper
level, and that these insights about the scientific method will help
you to develop a better understanding of how the underlying forces and
equations of motion and how they impact a given system. Furthermore, by introducing various numerical methods
via computational projects and exercises, we aim at developing your competences and skills about these topics.

These competences will enable you to

* understand how algorithms are used to solve mathematical problems,

* derive, verify, and implement algorithms,

* understand what can go wrong with algorithms,

* use these algorithms to construct reproducible scientific outcomes and to engage in science in ethical ways, and

* think algorithmically for the purposes of gaining deeper insights about scientific problems.

All these elements are central for maturing and gaining a better understanding of the modern scientific process *per se*.

The power of the scientific method lies in identifying a given problem
as a special case of an abstract class of problems, identifying
general solution methods for this class of problems, and applying a
general method to the specific problem (applying means, in the case of
computing, calculations by pen and paper, symbolic computing, or
numerical computing by ready-made and/or self-written software). This
generic view on problems and methods is particularly important for
understanding how to apply available, generic software to solve a
particular problem.

*However, verification of algorithms and understanding their limitations requires much of the classical knowledge about continuous models.*

## A well-known examples to illustrate many of the above concepts

Before we venture into a reminder on Python and mechanics relevant applications, let us briefly outline some of the
abovementioned topics using an example many of you may have seen before in for example CMSE201.
A simple algorithm for integration is the Trapezoidal rule.
Integration of a function $f(x)$ by the Trapezoidal Rule is given by following algorithm for an interval $x \in [a,b]$

$$
\int_a^b(f(x) dx = \frac{1}{2}\left [f(a)+2f(a+h)+\dots+2f(b-h)+f(b)\right] +O(h^2),
$$

where $h$ is the so-called stepsize defined by the number of integration points $N$ as $h=(b-a)/(n)$.
Python offers an  extremely versatile programming  environment, allowing for
the inclusion of analytical studies in a numerical program. Here we show an
example code with the **trapezoidal rule**. We use also **SymPy** to evaluate the exact value of the integral and compute the absolute error
with respect to the numerically evaluated one of the integral
$\int_0^1 dx x^2 = 1/3$.
The following code for  the trapezoidal rule allows you  to plot the relative error by comparing with the exact result. By increasing to $10^8$ points one arrives at a region where numerical errors start to accumulate.

%matplotlib inline

from math import log10
import numpy as np
from sympy import Symbol, integrate
import matplotlib.pyplot as plt
# function for the trapezoidal rule
def Trapez(a,b,f,n):
   h = (b-a)/float(n)
   s = 0
   x = a
   for i in range(1,n,1):
       x = x+h
       s = s+ f(x)
   s = 0.5*(f(a)+f(b)) +s
   return h*s
#  function to compute pi
def function(x):
    return x*x
# define integration limits
a = 0.0;  b = 1.0;
# find result from sympy
# define x as a symbol to be used by sympy
x = Symbol('x')
exact = integrate(function(x), (x, a, b))
# set up the arrays for plotting the relative error
n = np.zeros(9); y = np.zeros(9);
# find the relative error as function of integration points
for i in range(1, 8, 1):
    npts = 10**i
    result = Trapez(a,b,function,npts)
    RelativeError = abs((exact-result)/exact)
    n[i] = log10(npts); y[i] = log10(RelativeError);
plt.plot(n,y, 'ro')
plt.xlabel('n')
plt.ylabel('Relative error')
plt.show()

This example shows the potential of combining numerical algorithms with symbolic calculations, allowing us to

* Validate and verify  their  algorithms.

* Including concepts like unit testing, one has the possibility to test and test several or all parts of the code.

* Validation and verification are then included *naturally* and one can develop a better attitude to what is meant with an ethically sound scientific approach.

* The above example allows the student to also test the mathematical error of the algorithm for the trapezoidal rule by changing the number of integration points. The students get **trained from day one to think error analysis**.

* With a Jupyter notebook you can keep exploring similar examples and turn them in as your own notebooks.

In this process we can easily bake in
1. How to structure a code in terms of functions

2. How to make a module

3. How to read input data flexibly from the command line

4. How to create graphical/web user interfaces

5. How to write unit tests (test functions or doctests)

6. How to refactor code in terms of classes (instead of functions only)

7. How to conduct and automate large-scale numerical experiments

8. How to write scientific reports in various formats (LaTeX, HTML)

The conventions and techniques outlined here will save you a lot of time when you incrementally extend software over time from simpler to more complicated problems. In particular, you will benefit from many good habits:
1. New code is added in a modular fashion to a library (modules)

2. Programs are run through convenient user interfaces

3. It takes one quick command to let all your code undergo heavy testing

4. Tedious manual work with running programs is automated,

5. Your scientific investigations are reproducible, scientific reports with top quality typesetting are produced both for paper and electronic devices.


import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd.function import Function
from torch.autograd import Variable

class CenterTripletLoss(nn.Module):
    """ Hetero-center-triplet-loss-for-VT-Re-ID
   "Parameters Sharing Exploration and Hetero-Center Triplet Loss for Visible-Thermal Person Re-Identification"
   [(arxiv)](https://arxiv.org/abs/2008.06223).

    Args:
    - margin (float): margin for triplet.
    """

    def __init__(self, batch_size, margin=0.3):
        super(CenterTripletLoss, self).__init__()
        self.margin = margin
        self.ranking_loss = nn.MarginRankingLoss(margin=margin)

    def forward(self, feats, labels):
        """
        Args:
        - inputs: feature matrix with shape (batch_size, feat_dim)
        - targets: ground truth labels with shape (num_classes)
        """
        label_uni = labels.unique()
        targets = torch.cat([label_uni,label_uni])
        label_num = len(label_uni)
        feat = feats.chunk(label_num*2, 0)
        center = []
        for i in range(label_num*2):
            center.append(torch.mean(feat[i], dim=0, keepdim=True))
        inputs = torch.cat(center)

        n = inputs.size(0)

        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
        dist = dist + dist.t()
        dist.addmm_(1, -2, inputs, inputs.t())
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability

        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
        dist_ap, dist_an = [], []
        for i in range(n):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)

        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)

        # compute accuracy
        correct = torch.ge(dist_an, dist_ap).sum().item()
        return loss, correct

class CrossEntropyLabelSmooth(nn.Module):
    """Cross entropy loss with label smoothing regularizer.
    Reference:
    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
    Equation: y = (1 - epsilon) * y + epsilon / K.
    Args:
        num_classes (int): number of classes.
        epsilon (float): weight.
    """
    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):
        super(CrossEntropyLabelSmooth, self).__init__()
        self.num_classes = num_classes
        self.epsilon = epsilon
        self.use_gpu = use_gpu
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, inputs, targets):
        """
        Args:
            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
            targets: ground truth labels with shape (num_classes)
        """
        log_probs = self.logsoftmax(inputs)
        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)
        if self.use_gpu: targets = targets.cuda()
        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
        loss = (- targets * log_probs).mean(0).sum()
        return loss

class OriTripletLoss(nn.Module):
    """Triplet loss with hard positive/negative mining.

    Reference:
    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.

    Args:
    - margin (float): margin for triplet.
    """

    def __init__(self, batch_size, margin=0.3):
        super(OriTripletLoss, self).__init__()
        self.margin = margin
        self.ranking_loss = nn.MarginRankingLoss(margin=margin)

    def forward(self, inputs, targets):
        """
        Args:
        - inputs: feature matrix with shape (batch_size, feat_dim)
        - targets: ground truth labels with shape (num_classes)
        """
        n = inputs.size(0)

        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
        dist = dist + dist.t()
        dist.addmm_(1, -2, inputs, inputs.t())
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability

        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
        dist_ap, dist_an = [], []
        for i in range(n):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)

        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)

        # compute accuracy
        correct = torch.ge(dist_an, dist_ap).sum().item()
        return loss, correct

# Adaptive weights
def softmax_weights(dist, mask):
    max_v = torch.max(dist * mask, dim=1, keepdim=True)[0]
    diff = dist - max_v
    Z = torch.sum(torch.exp(diff) * mask, dim=1, keepdim=True) + 1e-6 # avoid division by zero
    W = torch.exp(diff) * mask / Z
    return W

def normalize(x, axis=-1):
    """Normalizing to unit length along the specified dimension.
    Args:
      x: pytorch Variable
    Returns:
      x: pytorch Variable, same shape as input
    """
    x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
    return x

class TripletLoss_WRT(nn.Module):
    """Weighted Regularized Triplet'."""

    def __init__(self):
        super(TripletLoss_WRT, self).__init__()
        self.ranking_loss = nn.SoftMarginLoss()

    def forward(self, inputs, targets, normalize_feature=False):
        if normalize_feature:
            inputs = normalize(inputs, axis=-1)
        dist_mat = pdist_torch(inputs, inputs)

        N = dist_mat.size(0)
        # shape [N, N]
        is_pos = targets.expand(N, N).eq(targets.expand(N, N).t()).float()
        is_neg = targets.expand(N, N).ne(targets.expand(N, N).t()).float()

        # `dist_ap` means distance(anchor, positive)
        # both `dist_ap` and `relative_p_inds` with shape [N, 1]
        dist_ap = dist_mat * is_pos
        dist_an = dist_mat * is_neg

        weights_ap = softmax_weights(dist_ap, is_pos)
        weights_an = softmax_weights(-dist_an, is_neg)
        furthest_positive = torch.sum(dist_ap * weights_ap, dim=1)
        closest_negative = torch.sum(dist_an * weights_an, dim=1)

        y = furthest_positive.new().resize_as_(furthest_positive).fill_(1)
        loss = self.ranking_loss(closest_negative - furthest_positive, y)

        # compute accuracy
        correct = torch.ge(closest_negative, furthest_positive).sum().item()
        return loss, correct

def pdist_torch(emb1, emb2):
    '''
    compute the eucilidean distance matrix between embeddings1 and embeddings2
    using gpu
    '''
    m, n = emb1.shape[0], emb2.shape[0]
    emb1_pow = torch.pow(emb1, 2).sum(dim = 1, keepdim = True).expand(m, n)
    emb2_pow = torch.pow(emb2, 2).sum(dim = 1, keepdim = True).expand(n, m).t()
    dist_mtx = emb1_pow + emb2_pow
    dist_mtx = dist_mtx.addmm_(1, -2, emb1, emb2.t())
    # dist_mtx = dist_mtx.clamp(min = 1e-12)
    dist_mtx = dist_mtx.clamp(min = 1e-12).sqrt()
    return dist_mtx

def pdist_np(emb1, emb2):
    '''
    compute the eucilidean distance matrix between embeddings1 and embeddings2
    using cpu
    '''
    m, n = emb1.shape[0], emb2.shape[0]
    emb1_pow = np.square(emb1).sum(axis = 1)[..., np.newaxis]
    emb2_pow = np.square(emb2).sum(axis = 1)[np.newaxis, ...]
    dist_mtx = -2 * np.matmul(emb1, emb2.T) + emb1_pow + emb2_pow
    # dist_mtx = np.sqrt(dist_mtx.clip(min = 1e-12))
    return dist_mtx


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

# # Naive Bayes
# Welcome to week two of this specialization. You will learn about Naive Bayes. Concretely, you will be using Naive Bayes for sentiment analysis on tweets. Given a tweet, you will decide if it has a positive sentiment or a negative one. Specifically you will:
#
# * Train a naive bayes model on a sentiment analysis task
# * Test using your model
# * Compute ratios of positive words to negative words
# * Do some error analysis
# * Predict on your own tweet
#
# You may already be familiar with Naive Bayes and its justification in terms of conditional probabilities and independence.
# * In this week's lectures and assignments we used the ratio of probabilities between positive and negative sentiments.
# * This approach gives us simpler formulas for these 2-way classification tasks.
#
# Load the cell below to import some packages.
# You  may want to browse the documentation of unfamiliar libraries and functions.

# In[2]:

from utils import process_tweet, lookup
import pdb
from nltk.corpus import stopwords, twitter_samples
import numpy as np
import pandas as pd
import nltk
import string
from nltk.tokenize import TweetTokenizer
from os import getcwd

# If you are running this notebook in your local computer,
# don't forget to download the twitter samples and stopwords from nltk.
#
# ```
# nltk.download('stopwords')
# nltk.download('twitter_samples')
# ```

# In[3]:

# add folder, tmp2, from our local workspace containing pre-downloaded corpora files to nltk's data path
filePath = f"{getcwd()}/../tmp2/"
nltk.data.path.append(filePath)

# In[4]:

# get the sets of positive and negative tweets
all_positive_tweets = twitter_samples.strings('positive_tweets.json')
all_negative_tweets = twitter_samples.strings('negative_tweets.json')

# split the data into two pieces, one for training and one for testing (validation set)
test_pos = all_positive_tweets[4000:]
train_pos = all_positive_tweets[:4000]
test_neg = all_negative_tweets[4000:]
train_neg = all_negative_tweets[:4000]

train_x = train_pos + train_neg
test_x = test_pos + test_neg

# avoid assumptions about the length of all_positive_tweets
train_y = np.append(np.ones(len(train_pos)), np.zeros(len(train_neg)))
test_y = np.append(np.ones(len(test_pos)), np.zeros(len(test_neg)))

# # Part 1: Process the Data
#
# For any machine learning project, once you've gathered the data, the first step is to process it to make useful inputs to your model.
# - **Remove noise**: You will first want to remove noise from your data -- that is, remove words that don't tell you much about the content. These include all common words like 'I, you, are, is, etc...' that would not give us enough information on the sentiment.
# - We'll also remove stock market tickers, retweet symbols, hyperlinks, and hashtags because they can not tell you a lot of information on the sentiment.
# - You also want to remove all the punctuation from a tweet. The reason for doing this is because we want to treat words with or without the punctuation as the same word, instead of treating "happy", "happy?", "happy!", "happy," and "happy." as different words.
# - Finally you want to use stemming to only keep track of one variation of each word. In other words, we'll treat "motivation", "motivated", and "motivate" similarly by grouping them within the same stem of "motiv-".
#
# We have given you the function `process_tweet()` that does this for you.

# In[5]:

custom_tweet = "RT @Twitter @chapagain Hello There! Have a great day. :) #good #morning http://chapagain.com.np"

# print cleaned tweet
print(process_tweet(custom_tweet))

# ## Part 1.1 Implementing your helper functions
#
# To help train your naive bayes model, you will need to build a dictionary where the keys are a (word, label) tuple and the values are the corresponding frequency.  Note that the labels we'll use here are 1 for positive and 0 for negative.
#
# You will also implement a `lookup()` helper function that takes in the `freqs` dictionary, a word, and a label (1 or 0) and returns the number of times that word and label tuple appears in the collection of tweets.
#
# For example: given a list of tweets `["i am rather excited", "you are rather happy"]` and the label 1, the function will return a dictionary that contains the following key-value pairs:
#
# {
#     ("rather", 1): 2
#     ("happi", 1) : 1
#     ("excit", 1) : 1
# }
#
# - Notice how for each word in the given string, the same label 1 is assigned to each word.
# - Notice how the words "i" and "am" are not saved, since it was removed by process_tweet because it is a stopword.
# - Notice how the word "rather" appears twice in the list of tweets, and so its count value is 2.
#
# #### Instructions
# Create a function `count_tweets()` that takes a list of tweets as input, cleans all of them, and returns a dictionary.
# - The key in the dictionary is a tuple containing the stemmed word and its class label, e.g. ("happi",1).
# - The value the number of times this word appears in the given collection of tweets (an integer).

#
#
#     Hints
#
#

#

#
Please use the `process_tweet` function that was imported above, and then store the words in their respective dictionaries and sets.

#
You may find it useful to use the `zip` function to match each element in `tweets` with each element in `ys`.

#
Remember to check if the key in the dictionary exists before adding that key to the dictionary, or incrementing its value.

#
Assume that the `result` dictionary that is input will contain clean key-value pairs (you can assume that the values will be integers that can be incremented).  It is good practice to check the datatype before incrementing the value, but it's not required here.

#

#

# In[6]:

# UNQ_C1 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def count_tweets(result, tweets, ys):
    '''
    Input:
        result: a dictionary that will be used to map each pair to its frequency
        tweets: a list of tweets
        ys: a list corresponding to the sentiment of each tweet (either 0 or 1)
    Output:
        result: a dictionary mapping each pair to its frequency
    '''

    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###
    for y, tweet in zip(ys, tweets):
        for word in process_tweet(tweet):
            # define the key, which is the word and label tuple
            pair = (word, y)

            # if the key exists in the dictionary, increment the count
            if pair in result:
                result[pair] += 1

            # else, if the key is new, add it to the dictionary and set the count to 1
            else:
                result[pair] = 1
    ### END CODE HERE ###

    return result

# In[7]:

# Testing your function

result = {}
tweets = ['i am happy', 'i am tricked', 'i am sad', 'i am tired', 'i am tired']
ys = [1, 0, 0, 0, 0]
count_tweets(result, tweets, ys)

# **Expected Output**: {('happi', 1): 1, ('trick', 0): 1, ('sad', 0): 1, ('tire', 0): 2}

# # Part 2: Train your model using Naive Bayes
#
# Naive bayes is an algorithm that could be used for sentiment analysis. It takes a short time to train and also has a short prediction time.
#
# #### So how do you train a Naive Bayes classifier?
# - The first part of training a naive bayes classifier is to identify the number of classes that you have.
# - You will create a probability for each class.
# $P(D_{pos})$ is the probability that the document is positive.
# $P(D_{neg})$ is the probability that the document is negative.
# Use the formulas as follows and store the values in a dictionary:
#
# $$P(D_{pos}) = \frac{D_{pos}}{D}\tag{1}$$
#
# $$P(D_{neg}) = \frac{D_{neg}}{D}\tag{2}$$
#
# Where $D$ is the total number of documents, or tweets in this case, $D_{pos}$ is the total number of positive tweets and $D_{neg}$ is the total number of negative tweets.

# #### Prior and Logprior
#
# The prior probability represents the underlying probability in the target population that a tweet is positive versus negative.  In other words, if we had no specific information and blindly picked a tweet out of the population set, what is the probability that it will be positive versus that it will be negative? That is the "prior".
#
# The prior is the ratio of the probabilities $\frac{P(D_{pos})}{P(D_{neg})}$.
# We can take the log of the prior to rescale it, and we'll call this the logprior
#
# $$\text{logprior} = log \left( \frac{P(D_{pos})}{P(D_{neg})} \right) = log \left( \frac{D_{pos}}{D_{neg}} \right)$$.
#
# Note that $log(\frac{A}{B})$ is the same as $log(A) - log(B)$.  So the logprior can also be calculated as the difference between two logs:
#
# $$\text{logprior} = \log (P(D_{pos})) - \log (P(D_{neg})) = \log (D_{pos}) - \log (D_{neg})\tag{3}$$

# #### Positive and Negative Probability of a Word
# To compute the positive probability and the negative probability for a specific word in the vocabulary, we'll use the following inputs:
#
# - $freq_{pos}$ and $freq_{neg}$ are the frequencies of that specific word in the positive or negative class. In other words, the positive frequency of a word is the number of times the word is counted with the label of 1.
# - $N_{pos}$ and $N_{neg}$ are the total number of positive and negative words for all documents (for all tweets), respectively.
# - $V$ is the number of unique words in the entire set of documents, for all classes, whether positive or negative.
#
# We'll use these to compute the positive and negative probability for a specific word using this formula:
#
# $$ P(W_{pos}) = \frac{freq_{pos} + 1}{N_{pos} + V}\tag{4} $$
# $$ P(W_{neg}) = \frac{freq_{neg} + 1}{N_{neg} + V}\tag{5} $$
#
# Notice that we add the "+1" in the numerator for additive smoothing.  This [wiki article](https://en.wikipedia.org/wiki/Additive_smoothing) explains more about additive smoothing.

# #### Log likelihood
# To compute the loglikelihood of that very same word, we can implement the following equations:
#
# $$\text{loglikelihood} = \log \left(\frac{P(W_{pos})}{P(W_{neg})} \right)\tag{6}$$

# ##### Create `freqs` dictionary
# - Given your `count_tweets()` function, you can compute a dictionary called `freqs` that contains all the frequencies.
# - In this `freqs` dictionary, the key is the tuple (word, label)
# - The value is the number of times it has appeared.
#
# We will use this dictionary in several parts of this assignment.

# In[8]:

# Build the freqs dictionary for later uses

freqs = count_tweets({}, train_x, train_y)

# #### Instructions
# Given a freqs dictionary, `train_x` (a list of tweets) and a `train_y` (a list of labels for each tweet), implement a naive bayes classifier.
#
# ##### Calculate $V$
# - You can then compute the number of unique words that appear in the `freqs` dictionary to get your $V$ (you can use the `set` function).
#
# ##### Calculate $freq_{pos}$ and $freq_{neg}$
# - Using your `freqs` dictionary, you can compute the positive and negative frequency of each word $freq_{pos}$ and $freq_{neg}$.
#
# ##### Calculate $N_{pos}$ and $N_{neg}$
# - Using `freqs` dictionary, you can also compute the total number of positive words and total number of negative words $N_{pos}$ and $N_{neg}$.
#
# ##### Calculate $D$, $D_{pos}$, $D_{neg}$
# - Using the `train_y` input list of labels, calculate the number of documents (tweets) $D$, as well as the number of positive documents (tweets) $D_{pos}$ and number of negative documents (tweets) $D_{neg}$.
# - Calculate the probability that a document (tweet) is positive $P(D_{pos})$, and the probability that a document (tweet) is negative $P(D_{neg})$
#
# ##### Calculate the logprior
# - the logprior is $log(D_{pos}) - log(D_{neg})$
#
# ##### Calculate log likelihood
# - Finally, you can iterate over each word in the vocabulary, use your `lookup` function to get the positive frequencies, $freq_{pos}$, and the negative frequencies, $freq_{neg}$, for that specific word.
# - Compute the positive probability of each word $P(W_{pos})$, negative probability of each word $P(W_{neg})$ using equations 4 & 5.
#
# $$ P(W_{pos}) = \frac{freq_{pos} + 1}{N_{pos} + V}\tag{4} $$
# $$ P(W_{neg}) = \frac{freq_{neg} + 1}{N_{neg} + V}\tag{5} $$
#
# **Note:** We'll use a dictionary to store the log likelihoods for each word.  The key is the word, the value is the log likelihood of that word).
#
# - You can then compute the loglikelihood: $log \left( \frac{P(W_{pos})}{P(W_{neg})} \right)\tag{6}$.

# In[24]:

# UNQ_C2 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def train_naive_bayes(freqs, train_x, train_y):
    '''
    Input:
        freqs: dictionary from (word, label) to how often the word appears
        train_x: a list of tweets
        train_y: a list of labels correponding to the tweets (0,1)
    Output:
        logprior: the log prior. (equation 3 above)
        loglikelihood: the log likelihood of you Naive bayes equation. (equation 6 above)
    '''
    loglikelihood = {}
    logprior = 0

    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###

    # calculate V, the number of unique words in the vocabulary
    vocab = set([pair[0] for pair in freqs.keys()])
    V = len(vocab)

    # calculate N_pos and N_neg
    N_pos = N_neg = 0

    for pair in freqs.keys():
        # if the label is positive (greater than zero)
        if pair[1] > 0:

            # Increment the number of positive words by the count for this (word, label) pair
            N_pos += freqs[pair]

        # else, the label is negative
        else:

            # increment the number of negative words by the count for this (word,label) pair
            N_neg += freqs[pair]

    # Calculate D, the number of documents
    D = len(train_y)

    # Calculate D_pos, the number of positive documents (*hint: use sum(<np_array>))
    D_pos = (train_y == 1).sum()

    # Calculate D_neg, the number of negative documents (*hint: compute using D and D_pos)
    D_neg = (train_y == 0).sum()

    # Calculate logprior
    logprior = np.log(D_pos) - np.log(D_neg)

    # For each word in the vocabulary...
    for word in vocab:
        # get the positive and negative frequency of the word
        if (word, 1) in freqs:
            freq_pos = freqs[(word, 1)]
        else:
            freq_pos = 0

        if (word, 0) in freqs:
            freq_neg = freqs[(word, 0)]
        else:
            freq_neg = 0

        # calculate the probability that each word is positive, and negative
        p_w_pos = (freq_pos + 1)/ (V + N_pos)
        p_w_neg = (freq_neg + 1)/ (V + N_neg)

        # calculate the log likelihood of the word
        loglikelihood[word] = np.log(p_w_pos) - np.log(p_w_neg)

    ### END CODE HERE ###

    return logprior, loglikelihood

# In[25]:

# UNQ_C3 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# You do not have to input any code in this cell, but it is relevant to grading, so please do not change anything
logprior, loglikelihood = train_naive_bayes(freqs, train_x, train_y)
print(logprior)
print(len(loglikelihood))

# **Expected Output**:
#
# 0.0
#
# 9089

# # Part 3: Test your naive bayes
#
# Now that we have the `logprior` and `loglikelihood`, we can test the naive bayes function by making predicting on some tweets!
#
# #### Implement `naive_bayes_predict`
# **Instructions**:
# Implement the `naive_bayes_predict` function to make predictions on tweets.
# * The function takes in the `tweet`, `logprior`, `loglikelihood`.
# * It returns the probability that the tweet belongs to the positive or negative class.
# * For each tweet, sum up loglikelihoods of each word in the tweet.
# * Also add the logprior to this sum to get the predicted sentiment of that tweet.
#
# $$ p = logprior + \sum_i^N (loglikelihood_i)$$
#
# #### Note
# Note we calculate the prior from the training data, and that the training data is evenly split between positive and negative labels (4000 positive and 4000 negative tweets).  This means that the ratio of positive to negative 1, and the logprior is 0.
#
# The value of 0.0 means that when we add the logprior to the log likelihood, we're just adding zero to the log likelihood.  However, please remember to include the logprior, because whenever the data is not perfectly balanced, the logprior will be a non-zero value.

# In[19]:

# UNQ_C4 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def naive_bayes_predict(tweet, logprior, loglikelihood):
    '''
    Input:
        tweet: a string
        logprior: a number
        loglikelihood: a dictionary of words mapping to numbers
    Output:
        p: the sum of all the logliklihoods of each word in the tweet (if found in the dictionary) + logprior (a number)

    '''
    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###
    # process the tweet to get a list of words
    word_l = process_tweet(tweet)

    # initialize probability to zero
    p = 0

    # add the logprior
    p += logprior

    for word in word_l:

        # check if the word exists in the loglikelihood dictionary
        if word in loglikelihood:
            # add the log likelihood of that word to the probability
            p += loglikelihood[word]

    ### END CODE HERE ###

    return p

# In[26]:

# UNQ_C5 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# You do not have to input any code in this cell, but it is relevant to grading, so please do not change anything

# Experiment with your own tweet.
my_tweet = 'She smiled.'
p = naive_bayes_predict(my_tweet, logprior, loglikelihood)
print('The expected output is', p)

# **Expected Output**:
# - The expected output is around 1.57
# - The sentiment is positive.

# #### Implement test_naive_bayes
# **Instructions**:
# * Implement `test_naive_bayes` to check the accuracy of your predictions.
# * The function takes in your `test_x`, `test_y`, log_prior, and loglikelihood
# * It returns the accuracy of your model.
# * First, use `naive_bayes_predict` function to make predictions for each tweet in text_x.

# In[27]:

# UNQ_C6 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def test_naive_bayes(test_x, test_y, logprior, loglikelihood):
    """
    Input:
        test_x: A list of tweets
        test_y: the corresponding labels for the list of tweets
        logprior: the logprior
        loglikelihood: a dictionary with the loglikelihoods for each word
    Output:
        accuracy: (# of tweets classified correctly)/(total # of tweets)
    """
    accuracy = 0  # return this properly

    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###
    y_hats = []
    for tweet in test_x:
        # if the prediction is > 0
        if naive_bayes_predict(tweet, logprior, loglikelihood) > 0:
            # the predicted class is 1
            y_hat_i = 1
        else:
            # otherwise the predicted class is 0
            y_hat_i = 0

        # append the predicted class to the list y_hats
        y_hats.append(y_hat_i)

    # error is the average of the absolute values of the differences between y_hats and test_y
    error = (1/len(test_y))*(test_y != y_hats).sum()

    # Accuracy is 1 minus the error
    accuracy = 1- error

    ### END CODE HERE ###

    return accuracy

# In[28]:

print("Naive Bayes accuracy = %0.4f" %
      (test_naive_bayes(test_x, test_y, logprior, loglikelihood)))

# **Expected Accuracy**:
#
# 0.9940

# In[29]:

# UNQ_C7 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# You do not have to input any code in this cell, but it is relevant to grading, so please do not change anything

# Run this cell to test your function
for tweet in ['I am happy', 'I am bad', 'this movie should have been great.', 'great', 'great great', 'great great great', 'great great great great']:
    # print( '%s -> %f' % (tweet, naive_bayes_predict(tweet, logprior, loglikelihood)))
    p = naive_bayes_predict(tweet, logprior, loglikelihood)
#     print(f'{tweet} -> {p:.2f} ({p_category})')
    print(f'{tweet} -> {p:.2f}')

# **Expected Output**:
# - I am happy -> 2.15
# - I am bad -> -1.29
# - this movie should have been great. -> 2.14
# - great -> 2.14
# - great great -> 4.28
# - great great great -> 6.41
# - great great great great -> 8.55

# In[30]:

# Feel free to check the sentiment of your own tweet below
my_tweet = 'you are bad :('
naive_bayes_predict(my_tweet, logprior, loglikelihood)

# # Part 4: Filter words by Ratio of positive to negative counts
#
# - Some words have more positive counts than others, and can be considered "more positive".  Likewise, some words can be considered more negative than others.
# - One way for us to define the level of positiveness or negativeness, without calculating the log likelihood, is to compare the positive to negative frequency of the word.
#     - Note that we can also use the log likelihood calculations to compare relative positivity or negativity of words.
# - We can calculate the ratio of positive to negative frequencies of a word.
# - Once we're able to calculate these ratios, we can also filter a subset of words that have a minimum ratio of positivity / negativity or higher.
# - Similarly, we can also filter a subset of words that have a maximum ratio of positivity / negativity or lower (words that are at least as negative, or even more negative than a given threshold).
#
# #### Implement `get_ratio()`
# - Given the `freqs` dictionary of words and a particular word, use `lookup(freqs,word,1)` to get the positive count of the word.
# - Similarly, use the `lookup()` function to get the negative count of that word.
# - Calculate the ratio of positive divided by negative counts
#
# $$ ratio = \frac{\text{pos_words} + 1}{\text{neg_words} + 1} $$
#
# Where pos_words and neg_words correspond to the frequency of the words in their respective classes.
#

#

#
#             Words
#
#
#         Positive word count
#
#
#         Negative Word Count
#
#

#

#
#         glad
#
#
#         41
#
#
#         2
#
#

#

#
#         arriv
#
#
#         57
#
#
#         4
#
#

#

#
#         :(
#
#
#         1
#
#
#         3663
#
#

#

#
#         :-(
#
#
#         0
#
#
#         378
#
#

#

# In[31]:

# UNQ_C8 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def get_ratio(freqs, word):
    '''
    Input:
        freqs: dictionary containing the words
        word: string to lookup

    Output: a dictionary with keys 'positive', 'negative', and 'ratio'.
        Example: {'positive': 10, 'negative': 20, 'ratio': 0.5}
    '''
    pos_neg_ratio = {'positive': 0, 'negative': 0, 'ratio': 0.0}
    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###
    # use lookup() to find positive counts for the word (denoted by the integer 1)
    pos_neg_ratio['positive'] = lookup(freqs, word, 1)

    # use lookup() to find negative counts for the word (denoted by integer 0)
    pos_neg_ratio['negative'] = lookup(freqs, word, 0)

    # calculate the ratio of positive to negative counts for the word
    pos_neg_ratio['ratio'] = (pos_neg_ratio['positive'] + 1)/( pos_neg_ratio['negative'] + 1)
    ### END CODE HERE ###
    return pos_neg_ratio

# In[32]:

get_ratio(freqs, 'happi')

# #### Implement `get_words_by_threshold(freqs,label,threshold)`
#
# * If we set the label to 1, then we'll look for all words whose threshold of positive/negative is at least as high as that threshold, or higher.
# * If we set the label to 0, then we'll look for all words whose threshold of positive/negative is at most as low as the given threshold, or lower.
# * Use the `get_ratio()` function to get a dictionary containing the positive count, negative count, and the ratio of positive to negative counts.
# * Append a dictionary to a list, where the key is the word, and the dictionary is the dictionary `pos_neg_ratio` that is returned by the `get_ratio()` function.
# An example key-value pair would have this structure:
# ```
# {'happi':
#     {'positive': 10, 'negative': 20, 'ratio': 0.5}
# }
# ```

# In[35]:

# UNQ_C9 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
def get_words_by_threshold(freqs, label, threshold):
    '''
    Input:
        freqs: dictionary of words
        label: 1 for positive, 0 for negative
        threshold: ratio that will be used as the cutoff for including a word in the returned dictionary
    Output:
        word_set: dictionary containing the word and information on its positive count, negative count, and ratio of positive to negative counts.
        example of a key value pair:
        {'happi':
            {'positive': 10, 'negative': 20, 'ratio': 0.5}
        }
    '''
    word_list = {}

    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###
    for key in freqs.keys():
        word, _ = key

        # get the positive/negative ratio for a word
        pos_neg_ratio = get_ratio(freqs, word)['ratio']

        # if the label is 1 and the ratio is greater than or equal to the threshold...
        if label == 1 and pos_neg_ratio >= threshold:

            # Add the pos_neg_ratio to the dictionary
            word_list[word] = pos_neg_ratio

        # If the label is 0 and the pos_neg_ratio is less than or equal to the threshold...
        elif label == 0 and pos_neg_ratio <= threshold:

            # Add the pos_neg_ratio to the dictionary
            word_list[word] = pos_neg_ratio

        # otherwise, do not include this word in the list (do nothing)

    ### END CODE HERE ###
    return word_list

# In[36]:

# Test your function: find negative words at or below a threshold
get_words_by_threshold(freqs, label=0, threshold=0.05)

# In[37]:

# Test your function; find positive words at or above a threshold
get_words_by_threshold(freqs, label=1, threshold=10)

# Notice the difference between the positive and negative ratios. Emojis like :( and words like 'me' tend to have a negative connotation. Other words like 'glad', 'community', and 'arrives' tend to be found in the positive tweets.

# # Part 5: Error Analysis
#
# In this part you will see some tweets that your model missclassified. Why do you think the misclassifications happened? Were there any assumptions made by the naive bayes model?

# In[38]:

# Some error analysis done for you
print('Truth Predicted Tweet')
for x, y in zip(test_x, test_y):
    y_hat = naive_bayes_predict(x, logprior, loglikelihood)
    if y != (np.sign(y_hat) > 0):
        print('%d\t%0.2f\t%s' % (y, np.sign(y_hat) > 0, ' '.join(
            process_tweet(x)).encode('ascii', 'ignore')))

# # Part 6: Predict with your own tweet
#
# In this part you can predict the sentiment of your own tweet.

# In[39]:

# Test with your own tweet - feel free to modify `my_tweet`
my_tweet = 'I am happy because I am learning :)'

p = naive_bayes_predict(my_tweet, logprior, loglikelihood)
print(p)

# Congratulations on completing this assignment. See you next week!