import torch
import torch.nn as nn
from easydict import EasyDict as edict
from torch.nn import functional as F

from hybrik.models.layers.smplx.body_models import SMPLXLayer

from .builder import SPPE
from .layers.hrnet.hrnet_25d import get_hrnet25d

def flip(x):
    assert (x.dim() == 3 or x.dim() == 4)
    dim = x.dim() - 1

    return x.flip(dims=(dim,))

def norm_heatmap(norm_type, heatmap, tau=5, sample_num=1):
    # Input tensor shape: [N,C,...]
    shape = heatmap.shape
    if norm_type == 'softmax':
        heatmap = heatmap.reshape(*shape[:2], -1)
        # global soft max
        heatmap = F.softmax(heatmap, 2)
        return heatmap.reshape(*shape)
    elif norm_type == 'sampling':
        heatmap = heatmap.reshape(*shape[:2], -1)

        eps = torch.rand_like(heatmap)
        log_eps = torch.log(-torch.log(eps))
        gumbel_heatmap = heatmap - log_eps / tau

        gumbel_heatmap = F.softmax(gumbel_heatmap, 2)
        return gumbel_heatmap.reshape(*shape)
    elif norm_type == 'multiple_sampling':

        heatmap = heatmap.reshape(*shape[:2], 1, -1)

        eps = torch.rand(*heatmap.shape[:2], sample_num, heatmap.shape[3], device=heatmap.device)
        log_eps = torch.log(-torch.log(eps))
        gumbel_heatmap = heatmap - log_eps / tau
        gumbel_heatmap = F.softmax(gumbel_heatmap, 3)
        gumbel_heatmap = gumbel_heatmap.reshape(shape[0], shape[1], sample_num, shape[2])

        # [B, S, K, -1]
        return gumbel_heatmap.transpose(1, 2)
    else:
        raise NotImplementedError

@SPPE.register_module
class HRNetSMPLXCamKid(nn.Module):
    def __init__(self, norm_layer=nn.BatchNorm2d, **kwargs):
        super(HRNetSMPLXCamKid, self).__init__()
        self._norm_layer = norm_layer
        self.num_joints = kwargs['NUM_JOINTS']
        assert self.num_joints == 71
        self.norm_type = kwargs['POST']['NORM_TYPE']
        self.depth_dim = kwargs['EXTRA']['DEPTH_DIM']
        self.height_dim = kwargs['HEATMAP_SIZE'][0]
        self.width_dim = kwargs['HEATMAP_SIZE'][1]
        self.smpl_dtype = torch.float32
        self.use_kid = kwargs['EXTRA']['USE_KID']

        self.preact = get_hrnet25d(
            kwargs['HRNET_TYPE'], num_joints=71,
            depth_dim=self.depth_dim,
            is_train=True, generate_feat=True, generate_hm=True,
            pretrain=kwargs['HR_PRETRAINED'])
        self.pretrain_hrnet = kwargs['HR_PRETRAINED']

        self.smplx_layer = SMPLXLayer(
            # model_path='model_files/smpl_v1.1.0/smplx/SMPLX_NEUTRAL.npz',
            model_path='model_files/smplx/SMPLX_NEUTRAL.npz',
            num_betas=10,
            use_pca=False,
            age='kid',
            kid_template_path='model_files/smplx_kid_template.npy',
        )

        self.root_idx_smpl = 0
        self.body_joint_pairs = [
            (1, 2), (4, 5), (7, 8), (10, 11),
            (13, 14), (16, 17), (18, 19), (20, 21),
            (23, 24)
        ]
        self.hand_joint_pairs = [
            (25 + i, 40 + i) for i in range(40 - 25)
        ]
        self.leaf_joint_pairs = [
            (55, 56), (58, 59)
        ] + [(60 + i, 65 + i) for i in range(5)]
        self.joint_pairs = self.body_joint_pairs + self.hand_joint_pairs + self.leaf_joint_pairs

        # init cam
        init_cam = torch.tensor([0.9])
        self.register_buffer(
            'init_cam',
            torch.Tensor(init_cam).float())

        self.decshape = nn.Linear(2048, 21)

        self.decphi = nn.Linear(2048, 54 * 2)  # [cos(phi), sin(phi)]
        self.deccam = nn.Linear(2048, 1)
        self.decsigma = nn.Linear(2048, 71)

        self.focal_length = kwargs['FOCAL_LENGTH']
        bbox_3d_shape = kwargs['BBOX_3D_SHAPE'] if 'BBOX_3D_SHAPE' in kwargs else (2200, 2200, 2200)
        self.bbox_3d_shape = torch.tensor(bbox_3d_shape).float()
        self.depth_factor = self.bbox_3d_shape[2] * 1e-3
        self.input_size = 256.0

    def _initialize(self):
        self.preact.init_weights(self.pretrain_hrnet)

    def flip_xyz_coord(self, pred_jts, flatten=True):
        if flatten:
            assert pred_jts.dim() == 2
            num_batches = pred_jts.shape[0]
            pred_jts = pred_jts.reshape(num_batches, self.num_joints, 3)
        else:
            assert pred_jts.dim() == 3
            num_batches = pred_jts.shape[0]

        pred_jts[:, :, 0] = - pred_jts[:, :, 0]

        for pair in self.joint_pairs:
            dim0, dim1 = pair
            idx = torch.Tensor((dim0, dim1)).long()
            inv_idx = torch.Tensor((dim1, dim0)).long()
            pred_jts[:, idx] = pred_jts[:, inv_idx]

        if flatten:
            pred_jts = pred_jts.reshape(num_batches, self.num_joints * 3)

        return pred_jts

    def flip_phi(self, pred_phi):
        pred_phi[:, :, 1] = -1 * pred_phi[:, :, 1]

        for pair in self.joint_pairs:
            dim0, dim1 = pair
            if dim0 >= 55 or dim1 >= 55:
                break
            idx = torch.Tensor((dim0 - 1, dim1 - 1)).long()
            inv_idx = torch.Tensor((dim1 - 1, dim0 - 1)).long()
            pred_phi[:, idx] = pred_phi[:, inv_idx]

        return pred_phi

    def flip_sigma(self, pred_sigma):

        for pair in self.joint_pairs:
            dim0, dim1 = pair
            idx = torch.Tensor((dim0, dim1)).long()
            inv_idx = torch.Tensor((dim1, dim0)).long()
            pred_sigma[:, idx] = pred_sigma[:, inv_idx]

        return pred_sigma

    def flip_heatmap(self, heatmaps, shift=True, flip_last_dim=True):
        if flip_last_dim:
            heatmaps = heatmaps.flip(dims=(heatmaps.dim() - 1,))

        for pair in self.joint_pairs:
            dim0, dim1 = pair
            idx = torch.Tensor((dim0, dim1)).long()
            inv_idx = torch.Tensor((dim1, dim0)).long()
            heatmaps[:, idx] = heatmaps[:, inv_idx]

        if shift:
            if heatmaps.dim() == 3:
                heatmaps[:, :, 1:] = heatmaps[:, :, 0:-1]
            elif heatmaps.dim() == 4:
                heatmaps[:, :, :, 1:] = heatmaps[:, :, :, 0:-1]
            else:
                heatmaps[:, :, :, :, 1:] = heatmaps[:, :, :, :, 0:-1]

        return heatmaps

    def update_scale(self, pred_uvd, weight, init_scale, pred_shape_full, pred_phi, **kwargs):
        batch_size = pred_uvd.shape[0]
        cam_depth = self.focal_length / (self.input_size * init_scale + 1e-9)
        pred_phi = pred_phi.reshape(batch_size, -1, 2)

        pred_xyz = torch.zeros_like(pred_uvd)

        weight = weight.clamp_min(0)

        if 'bboxes' in kwargs.keys():
            bboxes = kwargs['bboxes']
            img_center = kwargs['img_center']

            cx = (bboxes[:, 0] + bboxes[:, 2]) * 0.5
            cy = (bboxes[:, 1] + bboxes[:, 3]) * 0.5
            w = (bboxes[:, 2] - bboxes[:, 0])
            h = (bboxes[:, 3] - bboxes[:, 1])

            cx = cx - img_center[:, 0]
            cy = cy - img_center[:, 1]
            cx = cx / w
            cy = cy / h

            bbox_center = torch.stack((cx, cy), dim=1).unsqueeze(dim=1)

            pred_xyz[:, :, 2:] = pred_uvd[:, :, 2:].clone()  # unit: (self.depth_factor m)
            pred_xy = ((pred_uvd[:, :, :2] + bbox_center) * self.input_size / self.focal_length) \
                * (pred_xyz[:, :, 2:] * self.depth_factor + cam_depth)  # unit: m

            pred_xyz[:, :, :2] = pred_xy / self.depth_factor  # unit: (self.depth_factor m)

            camera_root = pred_xyz[:, 0, :] * self.depth_factor
        else:
            # copy z
            pred_xyz[:, :, 2:] = pred_uvd[:, :, 2:].clone()  # unit: (self.depth_factor m)
            # back-project xy
            pred_xy = (pred_uvd[:, :, :2] * self.input_size / self.focal_length) \
                * (pred_xyz[:, :, 2:] * self.depth_factor + cam_depth)  # unit: m

            # unit: (self.depth_factor m)
            pred_xyz[:, :, :2] = pred_xy / self.depth_factor

            # unit: m
            camera_root = pred_xyz[:, 0, :] * self.depth_factor

        pred_xyz = pred_xyz - pred_xyz[:, [0]]

        pred_beta = pred_shape_full[:, :11]
        pred_expression = pred_shape_full[:, 11:]
        output = self.smplx_layer.hybrik(
            betas=pred_beta.type(self.smpl_dtype),
            expression=pred_expression.type(self.smpl_dtype),
            pose_skeleton=pred_xyz.type(self.smpl_dtype) * 2.2,
            phis=pred_phi.type(self.smpl_dtype),
            return_verts=True,
            naive=True
        )

        num_joints = 22
        # unit: m
        pred_xyz55 = output.joints.float()[:, :num_joints]
        pred_xyz55 = pred_xyz55 - pred_xyz55.reshape(batch_size, -1, 3)[:, [0], :]
        pred_xyz55 = pred_xyz55 + camera_root.unsqueeze(dim=1)

        pred_uvd55 = pred_uvd[:, :num_joints, :].clone()
        if 'bboxes' in kwargs.keys():
            pred_uvd55[:, :, :2] = pred_uvd55[:, :, :2] + bbox_center

        bs = pred_uvd.shape[0]
        # [B, K, 1]
        weight_uv55 = weight[:, :num_joints, :].reshape(bs, num_joints, 1)

        Ax = torch.zeros((bs, num_joints, 1), device=pred_uvd.device, dtype=pred_uvd.dtype)
        Ay = torch.zeros((bs, num_joints, 1), device=pred_uvd.device, dtype=pred_uvd.dtype)

        Ax[:, :, 0] = pred_uvd55[:, :, 0]
        Ay[:, :, 0] = pred_uvd55[:, :, 1]

        Ax = Ax * weight_uv55
        Ay = Ay * weight_uv55

        # [B, 2K, 1]
        A = torch.cat((Ax, Ay), dim=1)

        bx = (pred_xyz55[:, :, 0] - self.input_size * pred_uvd55[:, :, 0] / self.focal_length * pred_xyz55[:, :, 2]) * weight_uv55[:, :, 0]
        by = (pred_xyz55[:, :, 1] - self.input_size * pred_uvd55[:, :, 1] / self.focal_length * pred_xyz55[:, :, 2]) * weight_uv55[:, :, 0]

        # [B, 2K, 1]
        b = torch.cat((bx, by), dim=1)[:, :, None]
        res = torch.inverse(A.transpose(1, 2).bmm(A)).bmm(A.transpose(1, 2)).bmm(b)

        scale = 1.0 / res

        assert scale.shape == init_scale.shape

        return scale

    def forward(self, x, flip_test=False, **kwargs):
        batch_size = x.shape[0]

        out_uv, out_z, x0 = self.preact(x)

        out_uv = out_uv.reshape(batch_size, self.num_joints, self.height_dim * self.width_dim)
        out_z = out_z.reshape(batch_size, self.num_joints, self.depth_dim)

        heatmaps_uv = norm_heatmap(self.norm_type, out_uv)
        heatmaps_z = norm_heatmap(self.norm_type, out_z)

        if flip_test:
            flip_x = flip(x)
            flip_out_uv, flip_out_z, flip_x0 = self.preact(flip_x)

            # flip heatmap
            flip_out_uv = flip_out_uv.reshape(batch_size, self.num_joints, self.height_dim, self.width_dim)
            flip_out_z = flip_out_z.reshape(batch_size, self.num_joints, self.depth_dim)
            flip_out_uv = self.flip_heatmap(flip_out_uv)
            flip_out_z = self.flip_heatmap(flip_out_z, shift=False, flip_last_dim=False)

            flip_out_uv = flip_out_uv.reshape(batch_size, self.num_joints, self.height_dim * self.width_dim)
            flip_out_z = flip_out_z.reshape(batch_size, self.num_joints, self.depth_dim)

            flip_heatmaps_uv = norm_heatmap(self.norm_type, flip_out_uv)
            flip_heatmaps_z = norm_heatmap(self.norm_type, flip_out_z)

            heatmaps_uv = (heatmaps_uv + flip_heatmaps_uv) / 2
            heatmaps_z = (heatmaps_z + flip_heatmaps_z) / 2
            # heatmaps_uv = heatmaps_uv
            # heatmaps_z = heatmaps_z

        maxvals_uv, _ = torch.max(heatmaps_uv, dim=2, keepdim=True)
        maxvals_z, _ = torch.max(heatmaps_z, dim=2, keepdim=True)
        maxvals = maxvals_uv * maxvals_z

        heatmaps_uv = heatmaps_uv.reshape((batch_size, self.num_joints, self.height_dim, self.width_dim))

        hm_x0 = heatmaps_uv.sum(axis=2)
        hm_y0 = heatmaps_uv.sum(axis=3)
        hm_z0 = heatmaps_z

        range_tensor = torch.arange(hm_x0.shape[-1], dtype=torch.float32, device=hm_x0.device)
        hm_x = hm_x0 * range_tensor
        hm_y = hm_y0 * range_tensor
        hm_z = hm_z0 * range_tensor

        coord_x = hm_x.sum(dim=2, keepdim=True)
        coord_y = hm_y.sum(dim=2, keepdim=True)
        coord_z = hm_z.sum(dim=2, keepdim=True)

        coord_x = coord_x / float(self.width_dim) - 0.5
        coord_y = coord_y / float(self.height_dim) - 0.5
        coord_z = coord_z / float(self.depth_dim) - 0.5

        #  -0.5 ~ 0.5
        pred_uvd_jts = torch.cat((coord_x, coord_y, coord_z), dim=2)

        x0 = x0.view(x0.size(0), -1)
        init_cam = self.init_cam.expand(batch_size, -1)  # (B, 1,)

        xc = x0

        pred_shape_full = self.decshape(xc)
        pred_beta = pred_shape_full[:, :11]
        pred_expression = pred_shape_full[:, 11:]

        pred_phi = self.decphi(xc)
        pred_camera = self.deccam(xc).reshape(batch_size, -1) + init_cam

        sigma = self.decsigma(xc).reshape(batch_size, self.num_joints, 1).sigmoid()

        pred_phi = pred_phi.reshape(batch_size, -1, 2)

        if flip_test:
            flip_pred_shape = self.decshape(flip_x0)
            pred_shape_full = (flip_pred_shape + pred_shape_full) / 2

            pred_beta = pred_shape_full[:, :11]
            pred_expression = pred_shape_full[:, 11:]

            flip_pred_phi = self.decphi(flip_x0)
            flip_pred_phi = flip_pred_phi.reshape(batch_size, -1, 2)
            flip_pred_phi = self.flip_phi(flip_pred_phi)
            pred_phi = (pred_phi + flip_pred_phi) / 2

            flip_pred_camera = self.deccam(flip_x0).reshape(batch_size, -1) + init_cam
            flip_sigma = self.decsigma(flip_x0).reshape(batch_size, self.num_joints, 1).sigmoid()
            pred_camera = 2 / (1 / flip_pred_camera + 1 / pred_camera)

            flip_sigma = self.flip_sigma(flip_sigma)
            sigma = (sigma + flip_sigma) / 2

        camScale = pred_camera[:, :1].unsqueeze(1)

        # if not self.training:
        ice_step = 3
        for _ in range(ice_step):
            camScale = self.update_scale(
                pred_uvd=pred_uvd_jts,
                weight=1 - sigma * 10,
                init_scale=camScale,
                pred_shape_full=pred_shape_full,
                pred_phi=pred_phi,
                **kwargs)

        camDepth = self.focal_length / (self.input_size * camScale + 1e-9)

        pred_xyz_jts = torch.zeros_like(pred_uvd_jts)
        if 'bboxes' in kwargs.keys():
            bboxes = kwargs['bboxes']
            img_center = kwargs['img_center']

            cx = (bboxes[:, 0] + bboxes[:, 2]) * 0.5
            cy = (bboxes[:, 1] + bboxes[:, 3]) * 0.5
            w = (bboxes[:, 2] - bboxes[:, 0])
            h = (bboxes[:, 3] - bboxes[:, 1])

            cx = cx - img_center[:, 0]
            cy = cy - img_center[:, 1]
            cx = cx / w
            cy = cy / h

            bbox_center = torch.stack((cx, cy), dim=1).unsqueeze(dim=1)

            pred_xyz_jts[:, :, 2:] = pred_uvd_jts[:, :, 2:].clone()  # unit: (self.depth_factor m)
            pred_xy_jts_29_meter = ((pred_uvd_jts[:, :, :2] + bbox_center) * self.input_size / self.focal_length) \
                * (pred_xyz_jts[:, :, 2:] * self.depth_factor + camDepth)  # unit: m

            pred_xyz_jts[:, :, :2] = pred_xy_jts_29_meter / self.depth_factor  # unit: (self.depth_factor m)

            camera_root = pred_xyz_jts[:, 0, :] * self.depth_factor
            camera_root[:, 2] += camDepth[:, 0, 0]
        else:
            pred_xyz_jts[:, :, 2:] = pred_uvd_jts[:, :, 2:].clone()  # unit: (self.depth_factor m)
            pred_xy_jts_29_meter = (pred_uvd_jts[:, :, :2] * self.input_size / self.focal_length) \
                * (pred_xyz_jts[:, :, 2:] * self.depth_factor + camDepth)  # unit: m

            pred_xyz_jts[:, :, :2] = pred_xy_jts_29_meter / self.depth_factor  # unit: (self.depth_factor m)

            camera_root = pred_xyz_jts[:, 0, :] * self.depth_factor
            camera_root[:, 2] += camDepth[:, 0, 0]

        pred_xyz_jts = pred_xyz_jts - pred_xyz_jts[:, [0]]

        pred_phi = pred_phi.reshape(batch_size, -1, 2)

        output = self.smplx_layer.hybrik(
            betas=pred_beta.type(self.smpl_dtype),
            expression=pred_expression.type(self.smpl_dtype),
            pose_skeleton=pred_xyz_jts.type(self.smpl_dtype) * 2.2,
            phis=pred_phi.type(self.smpl_dtype),
            return_verts=True,
            naive=False
        )
        pred_vertices = output.vertices.float()
        #  -0.5 ~ 0.5
        pred_xyz_struct_full = output.joints.float() / 2.2
        #  -0.5 ~ 0.5
        pred_xyz_hybrik_struct = self.smplx_layer.get_extended_joints(output.joints[:, :55].clone(), pred_vertices)
        pred_xyz_hybrik_struct = pred_xyz_hybrik_struct / 2.2

        pred_theta_quat = output.theta_quat.float().reshape(batch_size, -1)
        pred_theta_mat = output.rot_mats.float().reshape(batch_size, -1)

        transl = camera_root - output.joints.float().reshape(batch_size, -1, 3)[:, 0, :]

        gt_output = 0
        if 'gt_labels' in kwargs.keys() and kwargs['gt_labels'] is not None:
            gt_labels = kwargs['gt_labels']
            gt_beta = torch.cat([gt_labels['target_beta'], gt_labels['target_beta_kid']], dim=1)
            with torch.no_grad():
                gt_output = self.smplx_layer.forward_simple(
                    betas=gt_beta,
                    expression=gt_labels['target_expression'],
                    full_pose=gt_labels['target_theta_full'].reshape(-1, 55, 9),
                    return_verts=True,
                    root_align=True
                )

        # project
        pred_xyz_full = (pred_xyz_struct_full * 2.2).reshape(-1, 127, 3)
        pred_uv_full = torch.zeros_like(pred_xyz_full[:, :, :2])
        pred_xyz_full = pred_xyz_full + transl[:, None, :]
        pred_uv_full = (pred_xyz_full[:, :, :2] / pred_xyz_full[:, :, 2:]) * self.focal_length / self.input_size

        if 'bboxes' in kwargs.keys():
            pred_uv_full = pred_uv_full - bbox_center

        output = edict(
            pred_phi=pred_phi,
            pred_shape_full=pred_shape_full,
            pred_beta=pred_beta,
            pred_expression=pred_expression,
            pred_theta_quat=pred_theta_quat,
            pred_theta_mat=pred_theta_mat,
            pred_uvd_jts=pred_uvd_jts.reshape(batch_size, -1),
            pred_xyz_hybrik=pred_xyz_jts.reshape(batch_size, -1),
            pred_xyz_hybrik_struct=pred_xyz_hybrik_struct.reshape(batch_size, -1),
            pred_xyz_full=pred_xyz_struct_full.reshape(batch_size, -1),
            pred_uv_full=pred_uv_full,
            pred_vertices=pred_vertices,
            maxvals=maxvals,
            cam_scale=camScale[:, 0],
            cam_root=camera_root,
            transl=transl,
            img_feat=x0,
            pred_camera=pred_camera,
            sigma=sigma,
            gt_output=gt_output
        )
        return output

    def forward_gt_theta(self, gt_theta, gt_beta):

        output = self.smpl(
            pose_axis_angle=gt_theta,
            betas=gt_beta,
            global_orient=None,
            return_verts=True
        )

        return output


import numpy as np
import csv
import math
import operator
import matplotlib.pyplot as plt

###################################################################################################################################
#This code has been based on:
#Title:Tutorial To Implement k-Nearest Neighbors in Python From Scratch
#Author:Jason Brownlee
#Date:September 12, 2014
#Link:https://machinelearningmastery.com/tutorial-to-implement-k-nearest-neighbors-in-python-from-scratch/
###################################################################################################################################

with open('training_data.csv', 'r') as csvfile:
        datareader = csv.reader(csvfile, delimiter=',')
        header = next(datareader)
        training_data = []

        for row in datareader:
            row_of_floats = list(map(float, row))
            training_data.append(row_of_floats)

# data is  of type list
trianing_data_as_array = np.array(training_data)

with open('validation_data.csv', 'r') as csvfile1:
        datareader = csv.reader(csvfile1, delimiter=',')
        header = next(datareader)
        validation_data = []

        for row in datareader:
            row_of_floats = list(map(float, row))
            validation_data.append(row_of_floats)

# data is  of type list
validation_data_as_array = np.array(validation_data)

#instance 1 and 2 are two points in n-dimensional space. They represent two separate rows in the table with n  being the number of parameter columns.
def euclideanDistance(instance1, instance2, length):
    distance = 0
    for x in range(length):
        #this will give the sum of the squares of the distance between the 2 points in all the dimensions.
        distance += pow((instance1[x]-instance2[x]), 2)
        # this will give us the Euclidean distance in (length) dimensions.
    return math.sqrt(distance)

def getNeighbors(trainingSet, testInstance, k):
    distances =  []
    #testInstance is a list representing a single point (single row in table).
    #length is the number of dimensions
    length = len(testInstance)-1
    #for each row in training set
    #find the eculidean distance over all dimensions for that point
    #the distances list will contain the distances between the test poiint and all the training points.
    #Actually it will contain the traininSet point (a list of all parameters) and the distance
    #then we sort it from smallest to largest distance
    for x in range(len(trainingSet)):
        dist = euclideanDistance(testInstance, trainingSet[x], length)
        distances.append((trainingSet[x], dist))
    distances.sort(key=operator.itemgetter(1))
    neighbors = []
    for x in range(k):
        # neighbors is a list that will contain the data-points [list of parameters] of the k nearest points
        neighbors.append(distances[x][0])
    return neighbors

def getResponse(neighbors,k):
    #a dictionary containing the different claasses and the number of neighbors that belong to that class
    total_sum = 0
    for x in range(len(neighbors)):
        response = neighbors[x][-1]
        total_sum += response
    return total_sum/k
# In case of a  draw this selects one response, ideally we would select an ubiased random response.

#  Simple function that evaluates accuracy of predictions
def rmsError(testSet, predictions):
    errors = []
    for x in range(len(testSet)):
        error = (testSet[x][-1]-predictions[x])
        errors.append(error**2)
    rms_error =  math.sqrt(np.mean(errors))
    return rms_error

def mean_absolute_percentage_error(testSet, predictions):
    errors = []
    for x in range(len(testSet)):
        individual__percentage_error= ((abs((testSet[x][-1]-predictions[x])))/testSet[x][-1])*100
        errors.append(individual__percentage_error)
    return np.mean(errors)

max_neighbours = 200
accuracies = []

# this code is for plotting rms errors
'''
for y in range(1,max_neighbours):

    predictions = []
    for x in range(len(validation_data)):
        neighbors = getNeighbors(training_data, validation_data[x], y)
        result = getResponse(neighbors,y)
        predictions.append(result)
        #print('> predicted=' + repr(result) + ', actual=' + repr(validation_data[x][-1]))

    accuracy = rmsError(validation_data, predictions)
    #print('Accuracy: ' + repr(accuracy) + '%')

    accuracies.append([y,accuracy])

npAcurracies = np.array(accuracies)

'''

# this code is for plotting mean absolute percentage errors

for y in range(1,max_neighbours):

    predictions = []
    for x in range(len(validation_data)):
        neighbors = getNeighbors(training_data, validation_data[x], y)
        result = getResponse(neighbors,y)
        predictions.append(result)
        #print('> predicted=' + repr(result) + ', actual=' + repr(validation_data[x][-1]))

    accuracy = mean_absolute_percentage_error(validation_data, predictions)
    #print('Accuracy: ' + repr(accuracy) + '%')

    accuracies.append([y,accuracy])

npAcurracies = np.array(accuracies)

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(npAcurracies[:,0:1],npAcurracies[:,1:2])
ax.set_xlabel("Number of nearest neighbors chosen")
ax.set_ylabel("Mean Absolute Percentage Error of predictions")
ax.set_title("Mean Absolute Percentage Error of KNN REGRESSION method versus number of nearest neighbors chosen",fontsize = 8)
fig.savefig("KNN regression accuracy Percentage.pdf", fmt="pdf")
plt.show()


import argparse, datetime, os, re, sys
import Wrangler

USAGE = """

Substitute transit lines for an operator.

Writes out new line file as transitLines_[date_time].lin

"""
USERNAME     = os.environ["USERNAME"]
TM2_INPUTS   = os.path.join(r"C:\\Users", USERNAME, "Box\\Modeling and Surveys\\Development\\Travel Model Two Development\\Model Inputs")
TRN_NETFILE  = os.path.join(TM2_INPUTS,"2015_revised_mazs","trn")
LOG_FILENAME = "substitute_TransitLines_{}.log"

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter,)
    parser.add_argument("del_regex", help="Regex for line names to delete")
    parser.add_argument("new_line_file", type=str, nargs="+", help="New line file(s) to incorporate")
    args = parser.parse_args()

    now_str = datetime.datetime.now().strftime("%Y%m%d_%H%M")

    log_filename = LOG_FILENAME.format(now_str)
    Wrangler.setupLogging(None, log_filename)

    # read the transit line file
    trn_net = Wrangler.TransitNetwork(modelType=Wrangler.Network.MODEL_TYPE_TM2, modelVersion=1.0,
                                      basenetworkpath=TRN_NETFILE, isTiered=True, networkName="transitLines")
    # remove the ones we want to remove
    del_re = re.compile(args.del_regex)
    keep_lines = []
    for trn_line in trn_net.lines:
        # keep comments
        if isinstance(trn_line, str):
            keep_lines.append(trn_line)
            continue

        if del_re.match(trn_line.name):
            Wrangler.WranglerLogger.info("  Removing line {}".format(trn_line))
        else:
            keep_lines.append(trn_line)
    trn_net.lines = keep_lines

    # add the new ones
    for new_line_file in args.new_line_file:
        new_line_realpath = os.path.realpath(new_line_file)
        Wrangler.WranglerLogger.info("Reading line file {}".format(new_line_realpath))

        path_file = os.path.split(new_line_realpath)
        file_ext  = os.path.splitext(path_file[1])

        # create empty files fares.far, fareMatrix.txt
        for empty_files in ["fares.far", "fareMatrix.txt"]:
            f = open(os.path.join(path_file[0], empty_files),"w")
            f.write("")
            f.close()

        new_net = Wrangler.TransitNetwork(modelType=Wrangler.Network.MODEL_TYPE_TM2, modelVersion=1.0,
                                          basenetworkpath=path_file[0], isTiered=True, networkName=file_ext[0])
        for new_net_line in new_net.lines:
            if isinstance(new_net_line, str): continue

            Wrangler.WranglerLogger.info("  Adding {}".format(new_net_line))
            trn_net.lines.append(new_net_line)

    # write it
    out_file = "transitLines_{}".format(now_str)
    trn_net.write(name=out_file, writeEmptyFiles=False, suppressQuery=False, suppressValidation=True)
    Wrangler.WranglerLogger.info("Wrote {}.lin".format(out_file))


import time

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import DGMnets

# Set random seeds
np.random.seed(42)
tf.set_random_seed(42)

# PDE parameters
r = 0.05            # Interest rate
sigma = 0.25        # Volatility
mu = 0.2            # Mean
lambd = (mu-r)/sigma
gamma = 1           # Utility decay

# Time limits
T0 = 0.0 + 1e-10    # Initial time
T  = 1.0            # Terminal time

# Space limits
S1 = 0.0 + 1e-10    # Low boundary
S2 = 1              # High boundary

# Merton's analytical known solution
def analytical_solution(t, x):
    return -np.exp(-x*gamma*np.exp(r*(T-t)) - (T-t)*0.5*lambd**2)

def analytical_control_solution(t, x):
    return np.full_like(x, lambd/(gamma*sigma)*np.exp(-r*(T - t)))

def terminal_condition(x):
    return -tf.exp(-gamma*x)

# Loss relative to the value function
def loss_value_function(value_function, optimal_control, t1, x1, t3, x3):
    # Loss term 1: PDE
    u1 = value_function(t1,x1)
    u1_t = tf.gradients(u1,t1)[0]
    u1_x = tf.gradients(u1,x1)[0]
    u1_xx = tf.gradients(u1_x,x1)[0]
    pi = optimal_control(t1,x1)
    f = u1_t + pi*(mu-r)*u1_x + r*x1*u1_x + 0.5*sigma**2*pi**2*u1_xx
    L1 = tf.reduce_mean(tf.square(f))

    # Loss term 2: Terminal condition
    u3 = value_function(t3,x3)
    tc = terminal_condition(x3)
    L3 = tf.reduce_mean(tf.square(u3 - tc))

    return L1,L3

# Loss relative to the optimal control function
def loss_optimal_control(value_function, optimal_control, t1, x1):
    u1 = value_function(t1,x1)
    u1_x = tf.gradients(u1,x1)[0]
    u1_xx = tf.gradients(u1_x,x1)[0]
    pi = optimal_control(t1,x1)
    f = pi*(mu-r)*u1_x + r*x1*u1_x + 0.5*sigma**2*pi**2*u1_xx
    L = tf.reduce_mean(f)
    return L

# Sampling
def sampler(N1, N2, N3):
    # Sampler #1: PDE domain
    t1 = np.random.uniform(low=T0 - 0.5*(T - T0),
                           high=T,
                           size=[N1,1])
    s1 = np.random.uniform(low=S1 - (S2 - S1)*0.5,
                           high=S2 + (S2 - S1)*0.5,
                           size=[N1,1])

    # Sampler #2: boundary condition
    t2 = np.zeros(shape=(1, 1))
    s2 = np.zeros(shape=(1, 1))

    # Sampler #3: initial/terminal condition
    t3 = T * np.ones((N3,1)) #Terminal condition
    s3 = np.random.uniform(low=S1 - (S2 - S1)*0.5,
                           high=S2 + (S2 - S1)*0.5,
                           size=[N3,1])

    return (t1, s1, t2, s2, t3, s3)

# Neural Networks definition
num_layers = 3
nodes_per_layer = 50

# Value function net
with tf.variable_scope("value_function", reuse=True):
    value_function = DGMnets.DGMNet(num_layers, nodes_per_layer)
# Optimal control function net
with tf.variable_scope("optimal_control", reuse=True):
    optimal_control = DGMnets.DGMNet(num_layers, nodes_per_layer)

# Nets' input placeholders
t1_t = tf.placeholder(tf.float32, [None,1])
x1_t = tf.placeholder(tf.float32, [None,1])
t3_t = tf.placeholder(tf.float32, [None,1])
x3_t = tf.placeholder(tf.float32, [None,1])

t4_t = tf.placeholder(tf.float32, [None,1])
x4_t = tf.placeholder(tf.float32, [None,1])

# Loss terms
L1_t, L3_t = loss_value_function(value_function, optimal_control,
                                 t1_t, x1_t, t3_t, x3_t)
loss_value_function_t = L1_t + L3_t

loss_optimal_control_t = loss_optimal_control(value_function, optimal_control,
                                              t4_t, x4_t)

# Optimizers
value_function_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="value_function")
value_function_optimizer = tf.train.AdamOptimizer(learning_rate=0.001). \
                           minimize(loss_value_function_t, var_list=value_function_vars)

optimal_control_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="optimal_control")
optimal_control_optimizer = tf.train.AdamOptimizer(learning_rate=0.001). \
                            minimize(-loss_optimal_control_t, var_list=optimal_control_vars)

# Training parameters
steps_per_sample = 10
sampling_stages = 800

# Number of samples
NS_1 = 1000
NS_2 = 0
NS_3 = 100

# Plot tensors
tplot_t = tf.placeholder(tf.float32, [None,1], name="tplot_t") # We name to recover it later
xplot_t = tf.placeholder(tf.float32, [None,1], name="xplot_t")
vplot_t = tf.identity(value_function(tplot_t, xplot_t), name="vplot_t") # Trick for naming the trained model
cplot_t = tf.identity(optimal_control(tplot_t, xplot_t), name="cplot_t") # Trick for naming the trained model

# Training data holders
sampling_stages_list = []
elapsed_time_list = []
LVF_list = []
LOC_list = []

# Train network!!
init_op = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init_op)

for i in range(sampling_stages):
    t1, x1, _, _, t3, x3 = sampler(NS_1, NS_2, NS_3)

    start_time = time.clock()
    for _ in range(steps_per_sample):
        loss_value_function, _ = sess.run([loss_value_function_t, value_function_optimizer],
                                          feed_dict = {t1_t:t1, x1_t:x1, t3_t:t3, x3_t:x3})
    for _ in range(steps_per_sample):
        loss_optimal_control, _ = sess.run([loss_optimal_control_t, optimal_control_optimizer],
                                           feed_dict = {t4_t:t1, x4_t:x1})
    end_time = time.clock()
    elapsed_time = end_time - start_time

    sampling_stages_list.append(i)
    elapsed_time_list.append(elapsed_time)
    LVF_list.append(loss_value_function)
    LOC_list.append(loss_optimal_control)

    text = "Stage: {:04d}, LVF: {:e}, LOC: {:e}, {:f} seconds".format(i, loss_value_function, loss_optimal_control, elapsed_time)
    print(text)

# Plot results
N = 41      # Points on plot grid

times_to_plot = [0*T, 0.33*T, 0.66*T, T]
tplot = np.linspace(T0, T, N)
xplot = np.linspace(S1, S2, N)

plt.figure(figsize=(15,7))
i = 1
for t in times_to_plot:
    v_solution_plot = analytical_solution(t, xplot)
    c_solution_plot = analytical_control_solution(t, xplot)

    tt = t*np.ones_like(xplot.reshape(-1,1))
    nnv_plot, nnc_plot, = sess.run([vplot_t, cplot_t],
                        feed_dict={tplot_t:tt, xplot_t:xplot.reshape(-1,1)})

    plt.subplot(2, 4, i)
    plt.plot(xplot, v_solution_plot, 'b')
    plt.plot(xplot, nnv_plot, 'r')

    plt.ylim(-1.1, -0.2)
    plt.xlabel("S")
    plt.ylabel("V")
    plt.title("t = %.2f"%t, loc="left")

    plt.subplot(2, 4, i+len(times_to_plot))
    plt.plot(xplot, c_solution_plot, 'b')
    plt.plot(xplot, nnc_plot, 'r')

    plt.ylim(2, 3)
    plt.xlabel("S")
    plt.ylabel(r"$\pi$")
    plt.title("t = %.2f"%t, loc="left")

    i = i+1

plt.subplots_adjust(wspace=0.4, hspace=0.3)
plt.show()

# Save the trained model
saver = tf.train.Saver()
saver.save(sess, './SavedNets/merton_with_sup')

# Save the time tracking
np.save('./TimeTracks/merton_with_sup',
        (sampling_stages_list, elapsed_time_list, LVF_list, LOC_list))

# Plot losses X stages
plt.figure(figsize=(12,5))

plt.subplot(1,2,1)
plt.semilogy(sampling_stages_list,LVF_list)
plt.title("Value Function Loss", loc="left")
plt.xlabel("Stage")
plt.ylim(1e-8, 1)

plt.subplot(1,2,2)
plt.plot(sampling_stages_list,LOC_list)
plt.title("Optimal Control Loss", loc="left")
plt.xlabel("Stage")
plt.ylim(0.1, 0.2)

plt.show()


import cvlib as cv
from cvlib.object_detection import draw_bbox
import cv2
import json
import time
import os

class Stopwatch:
    def __init__(self):
        self.creationTime = time.time()

    def elapsedTime(self):
        return time.time() - self.creationTime

    def setZero(self):
        self.creationTime = time.time()

class Pole:
    def __init__(self):
        self.g = 0
        self.time = Stopwatch()

def ForStart(carsNorth, carsWest):
    if (carsNorth > carsWest):
        poleNorth.g = 1
        poleNorth.time.setZero()
        return [(1, 0), (0, 1)]

    poleWest.g = 1
    poleWest.time.setZero()
    return [(0, 1), (1, 0)]

def lights(north, west):
    if north[0].g == 0:
        toWork = west
        noWork = north
    else:
        toWork = north
        noWork = west

    if (toWork[0].time.elapsedTime() <= 10):
        return None

    elif (toWork[0].time.elapsedTime() >= 60):
        ##change Light
        toWork[0].g = 0
        noWork[0].g = 1
        noWork[0].time.setZero()

    else:
        loadFactor = toWork[1] / (toWork[1] + noWork[1])
        if (loadFactor < threshHold):
            ##change Light
            toWork[0].g = 0
            noWork[0].g = 1
            noWork[0].time.setZero()

    return [(north[0].g, (north[0].g + 1) % 2), (west[0].g, (west[0].g + 1) % 2)]

# open webcam
webcam = cv2.VideoCapture(0)

if not webcam.isOpened():
    # print("Could not open webcam")
    exit()

person = {}
person['people'] = []

# loop through frames
while webcam.isOpened():

    # read frame from webcam
    status, frame = webcam.read()

    if not status:
        # print("Could not read frame")
        exit()

    # apply object detection
    bbox, label, conf = cv.detect_common_objects(frame)
    person['people'].append(label)
    # print(bbox, label, conf)
    count = 0
    size = len(label)
    for i in range(0, size):
        if (label[i] == 'car'):
            count += 1
    # print (count)

    # draw bounding box over detected objects
    out = draw_bbox(frame, bbox, label, conf)

    # display output
    cv2.imshow("Real-time object detection", out)
    time.sleep(3)
    # press "Q" to stop
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# release resources
webcam.release()
cv2.destroyAllWindows()

min = 10  # sec
max = 60  # sec

threshHold = 5 / 12

carsNorth = count
carsWest = 5

poleNorth = Pole()
poleWest = Pole()

lightStatus = ForStart(carsNorth, carsWest)
print(lightStatus)

while (1):
    lightStatus = lights((poleNorth, carsNorth), (poleWest, carsWest))
    print(lightStatus)
    time.sleep(5)
    carsWest = 5


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn import svm
from sklearn.preprocessing import MinMaxScaler

# Carregar o conjunto de dados
input_file = '0-Datasets/diabetesClear.data'
names = ['Número Gestações', 'Glucose', 'pressao Arterial', 'Expessura da Pele', 'Insulina', 'IMC', 'Função Pedigree Diabete', 'Idade', 'Resultado']
features = ['Número Gestações', 'Glucose', 'pressao Arterial', 'Expessura da Pele', 'Insulina', 'IMC', 'Função Pedigree Diabete', 'Idade']
target = 'Resultado'
df = pd.read_csv(input_file, names=names)

# Dados de treinamento
X = df[features].values
y = df[target].values

# Normalização dos dados
scaler = MinMaxScaler()
X_normalized = scaler.fit_transform(X)

# Criação do classificador SVM
clf = svm.SVC(kernel='linear', C=1.0)
clf.fit(X_normalized, y)

# Função para plotar o hiperplano de separação em 3D
def plot_hyperplane(clf, ax):
    # Extrai os coeficientes e intercepto do hiperplano
    w = clf.coef_[0]
    a = -w[0] / w[1]

    # Determina as coordenadas do hiperplano
    xx = np.linspace(np.min(X_normalized[:, 0]), np.max(X_normalized[:, 0]), 10)
    yy = np.linspace(np.min(X_normalized[:, 1]), np.max(X_normalized[:, 1]), 10)
    XX, YY = np.meshgrid(xx, yy)
    ZZ = (-w[0] * XX - w[1] * YY - clf.intercept_[0]) / w[2]

    # Plota o hiperplano em 3D
    ax.plot_surface(XX, YY, ZZ, alpha=0.5, color='gray')

# Plotagem dos pontos de dados e hiperplano em 3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Rótulos das classes
labels = np.unique(y)
for label in labels:
    indices = np.where(y == label)
    ax.scatter(X_normalized[indices, 0], X_normalized[indices, 1], X_normalized[indices, 2], label=f'Classe {label}')

plot_hyperplane(clf, ax)

# Configurações do gráfico
ax.set_xlabel('Número de Gestacoes')
ax.set_ylabel('Glucose')
ax.set_zlabel('pressao Arterial')
ax.set_title('Classificação SVM em 3D')
ax.legend()
plt.show()


from triangle import gen_triangle

triangle = gen_triangle(30)
for i, _ in enumerate(triangle):
    ok = True

    for num in triangle[i]:
        if num % 5 == 0:
            ok = False

    if ok:
        print(f"row {i + 1} is OK")


""":class:`domain_intel.AlexaApi` unit test cases.

"""
import os
import calendar
from datetime import datetime
import pytest
import mock

import domain_intel.awis.actions
import domain_intel.utils

def test_awis_urlinfo_init():
    """Initialise a domain_intel.awis.actions.UrlInfo object.
    """
    # When I initialise an Alexa Web Information object
    awis = domain_intel.awis.actions.UrlInfo()

    # I should get a domain_intel.awis.actions.UrlInfo instance
    msg = 'Object is not a domain_intel.awis.actions.UrlInfo instance'
    assert isinstance(awis, domain_intel.awis.actions.UrlInfo), msg

@pytest.mark.usefixtures('docker_compose', 'kafka_ready')
def test_add_domains(add_domains):
    """Load of GTR domains into Kafka topic.
    """
    msg = 'Load of GTR domains into Kafka error'
    assert add_domains == 1000, msg

@pytest.mark.usefixtures('docker_compose', 'kafka_ready', 'add_domains')
def test_read_domains(read_domains):
    """Read GTR domains from Kafka topic.
    """
    msg = 'Read GTR domains from Kafka error'
    assert read_domains == 1000, msg

def test_flatten_batched_domain_xml():
    """Flatten a batched-domain XML response from Alexa AWIS.
    """
    # Given an Alexa AWIS batched XML response
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'multi-domain-result.xml')) as _fh:
        source_xml = _fh.read().rstrip()

    # when I flatten the XML
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.flatten_batched_xml(source_xml)

    # then I should receive a list of JSON-encode domains
    msg = 'Flattened domains should produce list of JSON-encoded domains'
    assert len(received) == 2, msg

def test_flatten_single_domain_xml():
    """Flatten a single-domain XML response from Alexa AWIS.
    """
    # Given an Alexa AWIS batched XML response
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'single-domain-result.xml')) as _fh:
        source_xml = _fh.read().rstrip()

    # when I flatten the XML
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.flatten_batched_xml(source_xml)

    # then I should receive a list of JSON-encode domains
    msg = 'Flattened domain should produce JSON-encoded domain list'
    assert len(received) == 1, msg

@pytest.mark.usefixtures('docker_compose', 'kafka_ready')
def test_parse_raw_urlinfo(parse_raw_urlinfo):
    """Parse raw Alexa XML.
    """
    msg = 'Real load of raw Alexa load into Kafka error'
    assert parse_raw_urlinfo == (9385, 16, 16), msg

@pytest.mark.usefixtures('docker_compose',
                         'kafka_ready',
                         'parse_raw_urlinfo')
def test_flatten_domains(flatten_domains):
    """Flatten batched Alexa XML and publish.
    """
    msg = 'Flattening batched Alexa XML back into Kafka error'
    assert flatten_domains == 16, msg

@pytest.mark.usefixtures('docker_compose',
                         'kafka_ready',
                         'parse_raw_urlinfo',
                         'flatten_domains',
                         'init_arango')
def test_persist(persist_domains):
    """Persist domain to the database.
    """
    # ... then I should receive a record count if domains stored
    msg = 'Count against the domain collection incorrect'
    assert persist_domains == (80, 80), msg

@pytest.mark.usefixtures('docker_compose',
                         'kafka_ready',
                         'parse_raw_urlinfo',
                         'flatten_domains')
def test_alexa_csv_dump():
    """Alexa CSV dump.
    """
    # Given a unique group ID that ensure isolated consumption of messages
    group_id = domain_intel.utils.id_generator()

    # ehen I dump the flattened Alexa JSON to CSV
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.alexa_csv_dump(max_read_count=7,
                                   group_id=group_id,
                                   dry=True)

    # then I should receive a record count of 7
    msg = 'Count against the flattened Alexa JSON read is incorrect'
    assert received == 7, msg

@mock.patch('domain_intel.awis.actions.urlinfo.time')
def test_alexa_flattened_extract(mock_dt):
    """Rules to extract fields from flattened Alexa.
    """
    # Given a source of flattened Alexa JSON
    alexa_json_file = os.path.join('domain_intel',
                                   'test',
                                   'files',
                                   'samples',
                                   'flattened_alexa.json')
    # when I parse the JSON
    mock_epoch = datetime(2017, 5, 16).utctimetuple()
    mock_dt.time.return_value = calendar.timegm(mock_epoch)
    received = []
    with open(alexa_json_file) as _fh:
        for line in _fh:
            stats = domain_intel.awis.actions.UrlInfo.alexa_flattened_extract(line)
            received.append(stats)

    # then I should receive domain name ranks
    msg = 'CSV dump extraction error: domains processed'
    expected = [
        'partymegapop.info',
        'new-rutor.org.pl',
        'watchseriesonline.io',
        'kibergrad.com',
        'darktorrent.pl',
        'kickassproxy.club',
        'inevil.com',
        'partymegapop.info',
    ]
    assert [x[0][0] for x in received] == expected, msg

    msg = 'CSV dump extraction error: exported row'
    expected = ['watchseriesonline.io', 1494892800.0, 588080]
    assert received[2][0] == expected, msg

@pytest.mark.usefixtures('docker_compose', 'init_arango')
def test_write_to_store():
    """Persistent store abstraction.
    """
    # Given an Alexa UrlInfo record
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'kafka_consumer_record.json')) as _fh:
        record = _fh.read().rstrip().encode('utf')

    # when I persist the record
    awis = domain_intel.awis.actions.UrlInfo()
    awis.write_to_store(record, dry=False)

    # then I should receive a record count of 1 domain
    msg = 'Count against the domain collection incorrect'
    assert awis.store.get_collection_count() == 1, msg

    # and I should receive a record count of 1 url_info records
    msg = 'Count against the url_info collection incorrect'
    assert awis.store.get_collection_count(collection_name='url-info') == 1, msg

@pytest.mark.usefixtures('docker_compose', 'init_arango')
def test_persist_country_rank():
    """Persistent country rank abstraction.
    """
    # Given an Alexa UrlInfo record
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'kafka_consumer_record.json')) as _fh:
        record = _fh.read().rstrip().encode('utf')
    url_info = domain_intel.awisapi.parser.UrlInfo(record)

    # when I persist the country ranks
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.persist_country_rank(url_info, dry=False)

    # then I should receive a record count of 16 countries
    msg = 'Country rank return count error'
    assert received == 16, msg

@pytest.mark.usefixtures('docker_compose', 'init_arango')
def test_persist_related_links():
    """Persistent related links abstraction.
    """
    # Given an Alexa UrlInfo record
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'kafka_consumer_record.json')) as _fh:
        record = _fh.read().rstrip().encode('utf')
    url_info = domain_intel.awisapi.parser.UrlInfo(record)

    # when I persist the country ranks
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.persist_related_links(url_info, dry=False)

    # then I should receive a record count of 10 links
    msg = 'Related links return count error'
    assert received == 10, msg

    msg = 'Count against the country collection incorrect'
    assert awis.store.get_collection_count(collection_name='link') == 10, msg

@pytest.mark.usefixtures('docker_compose', 'init_arango')
def test_contributing_subdomains():
    """Persistent contributing subdomains abstraction.
    """
    # Given an Alexa UrlInfo record
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'kafka_consumer_record.json')) as _fh:
        record = _fh.read().rstrip().encode('utf')
    url_info = domain_intel.awisapi.parser.UrlInfo(record)

    # when I persist the country ranks
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.persist_contributing_subdomains(url_info,
                                                    dry=False)

    # then I should receive a record count of 7 contributing subdomains
    msg = 'Contributing subdomains links return count error'
    assert received == 7, msg

    msg = 'Count against the subdomain collection incorrect'
    assert awis.store.get_collection_count(collection_name='subdomain') == 7, msg

@pytest.mark.usefixtures('docker_compose', 'init_arango')
def test_persist_url_info():
    """Persistent URL info JSON records.
    """
    # Given an Alexa UrlInfo record
    with open(os.path.join('domain_intel',
                           'test',
                           'files',
                           'samples',
                           'kafka_consumer_record.json')) as _fh:
        record = _fh.read().rstrip().encode('utf')
    url_info = domain_intel.awisapi.parser.UrlInfo(record)

    # when I persist the country ranks
    awis = domain_intel.awis.actions.UrlInfo()
    awis.persist_url_info(url_info, dry=False)
    received = awis.store.get_collection_count(collection_name='url-info')

    # then I should receive a record count of 1 URL info JSON
    msg = 'Count against the url-info'
    assert received == 1, msg

@pytest.mark.usefixtures('docker_compose',
                         'kafka_ready',
                         'parse_raw_urlinfo',
                         'flatten_domains',
                         'init_arango',
                         'persist_domains',
                         'add_domain_labels')
def test_traverse_relationships(traverse_relationships):
    """Traverse relationships.
    """
    # ... then I should get a count of relationship JSON structures
    msg = 'Graph relationship count incorrect'
    assert traverse_relationships == 80, msg

@pytest.mark.usefixtures('docker_compose',
                         'kafka_ready',
                         'parse_raw_urlinfo',
                         'flatten_domains',
                         'init_arango',
                         'persist_domains',
                         'add_domain_labels',
                         'traverse_relationships')
def test_wide_column_dump_worker():
    """Dump the wide-column CSV format.
    """
    # When I dump the wide-column CSV
    awis = domain_intel.awis.actions.UrlInfo()
    received = awis.wide_column_dump()

    # then I should get a count of CSV lines
    msg = 'Wide-column CSV count incorrect'
    assert received == (80, 242), msg


from datetime import tzinfo
import json

from django.contrib.auth.decorators import login_required
from django.http import HttpResponse, JsonResponse
from django.views.decorators.csrf import csrf_protect

from comments.models import Comment
from task.models import Task
from user_account.models import UserProfile, Achievement, AchievementsSettings

def set_achievements_at_commenting(user):
    achivements_names = {5: 'Commentator1', 10: 'Commentator2', 20: 'Commentator3', 30: 'Commentator4'}
    count = Comment.objects.count_comments_for_user(user)
    if count in achivements_names.keys():
        profile = UserProfile.objects.get_or_create_profile(user)
        achievement = Achievement.objects.get(name=achivements_names[count])
        if not AchievementsSettings.objects.filter(userProfile=profile, achievement=achievement).exists():
            achSetting = AchievementsSettings(userProfile=profile, achievement=achievement, count=1)
            achSetting.save()

@login_required
@csrf_protect
def add(request):
    comment = None
    func = lambda x: x if x is not None else '/static/user_account/pictures/unknown.png'
    try:
        if request.method == 'POST':
            pk = request.POST['pk']
            text = request.POST['text']
            task = Task.objects.get(pk=pk)
            comment = Comment(user=request.user, task=task, text=text)
            comment.save()
            set_achievements_at_commenting(request.user)
            comment = {"username": request.user.username, 'pk': request.user.pk,
                       "url": func(UserProfile.objects.get_or_create_profile(request.user).pictureUrl),
                       'time': comment.creation_time}
    except Exception as e:
        print(e)
        return HttpResponse(status=500)
    return JsonResponse(comment, status=200)


import logging
from enum import Enum
from collections import defaultdict

from src.common import *
from src.isoform_assignment import *

logger = logging.getLogger('CSA')

class MultimapResolvingStrategy(Enum):
    ignore_multimapper = 0
    merge = 1
    take_best = 2

class MultimapResolver:

    def __init__(self, strategy):
        self.strategy = strategy

    # list of ShortReadAssignments
    def resolve(self, assignment_list):
        if assignment_list is None or len(assignment_list) <= 1:
            return assignment_list

        if self.strategy == MultimapResolvingStrategy.ignore_multimapper:
            for a in assignment_list:
                a.assignment_type = ReadAssignmentType.noninformative
            return assignment_list

        elif self.strategy == MultimapResolvingStrategy.merge:
            informative_assignments = set()
            for i in range(len(assignment_list)):
                if assignment_list[i].assignment_type != ReadAssignmentType.noninformative:
                    informative_assignments.add(i)
            if len(informative_assignments) == 1:
                return self.suspend_assignments(assignment_list, informative_assignments)
            return self.suspend_assignments(assignment_list, informative_assignments, ReadAssignmentType.ambiguous)

        elif self.strategy == MultimapResolvingStrategy.take_best:
            return self.select_best_assignment(assignment_list)
        else:
            raise ValueError("Unsupported multimap strategy")

    def select_best_assignment(self, assignment_list):
        logger.debug("Resolving read %s" % assignment_list[0].read_id)
        logger.debug("Read assignment types %s" % " ".join(a.assignment_type.name for a in assignment_list))
        primary_unique = set()
        consistent_assignments = set()
        inconsistent_assignments = set()
        primary_inconsistent = set()
        for i, a in enumerate(assignment_list):
            if a.assignment_type in [ReadAssignmentType.unique, ReadAssignmentType.unique_minor_difference] and \
                    not a.multimapper:
                primary_unique.add(i)
            elif a.assignment_type == ReadAssignmentType.inconsistent and not a.multimapper:
                primary_inconsistent.add(i)
            if a.assignment_type in [ReadAssignmentType.unique,
                                     ReadAssignmentType.unique_minor_difference,
                                     ReadAssignmentType.ambiguous]:
                consistent_assignments.add(i)
            elif a.assignment_type == ReadAssignmentType.inconsistent:
                inconsistent_assignments.add(i)

        if primary_unique:
            if len(primary_unique) > 1:
                # TODO silence warn
                logger.debug("Multiple primary unique %s: %s " % (assignment_list[0].read_id, ",".join([assignment_list[i].gene_id for i in primary_unique])))
                return self.suspend_assignments(assignment_list, primary_unique, ReadAssignmentType.ambiguous)
            # primary unique is found, rest is noninformative
            logger.debug("Primary unique assignment selected: %s" %
                         " ".join([assignment_list[i].gene_id for i in primary_unique]))
            return self.suspend_assignments(assignment_list, primary_unique)

        if consistent_assignments:
            logger.debug("Merging %d consistent assignments " % len(consistent_assignments))
            return self.suspend_assignments(assignment_list, consistent_assignments, ReadAssignmentType.ambiguous)

        if primary_inconsistent:
            return self.select_best_inconsistent(assignment_list, primary_inconsistent)

        logger.debug("Merging inconsistent from %d assignments" % len(inconsistent_assignments))
        return self.select_best_inconsistent(assignment_list, inconsistent_assignments)

    def select_best_inconsistent(self, assignment_list, inconsistent_assignments):
        if len(inconsistent_assignments) > 1:
            logger.debug("= Multiple primary inconsistent " +
                         ",".join([assignment_list[i].gene_id for i in inconsistent_assignments]))
            assignment_scores = []
            for i in inconsistent_assignments:
                assignment_scores.append((assignment_list[i].score, i))
            best_score = min(assignment_scores, key=lambda x: x[0])[0]
            logger.debug("= Best score " + str(best_score))
            best_isoforms = [x[1] for x in filter(lambda x: x[0] == best_score, assignment_scores)]
            return self.suspend_assignments(assignment_list, best_isoforms,
                                            ReadAssignmentType.inconsistent if len(best_isoforms) > 1 else None)

        logger.debug("Primary inconsistent assignment selected")
        return self.suspend_assignments(assignment_list, inconsistent_assignments)

    def suspend_assignments(self, assignment_list, assignments_to_keep, new_type=None):
        for i in range(len(assignment_list)):
            if i in assignments_to_keep:
                if new_type is not None:
                    assignment_list[i].assignment_type = new_type
                    assignment_list[i].multimapper = True
                continue
            assignment_list[i].assignment_type = ReadAssignmentType.noninformative
        return assignment_list


import requests
import datetime
import re
from bs4 import BeautifulSoup as bs

url='https://us-cert.cisa.gov/ics/advisories?items_per_page=All'
outfile=open('CISA_Advisory_List_'+datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')+'.csv', 'a')
error=open('CISA_ERROR_FILE_'+datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')+'.txt', 'a')
page = requests.get(url)
soup = bs(page.content, 'html.parser')

for elem in soup.find_all('a', href=re.compile('/ics/advisories/')):
    url='https://us-cert.cisa.gov' + elem['href']
    advisory=url.replace('https://us-cert.cisa.gov/ics/advisories/','')
    try:
        print('https://us-cert.cisa.gov' + elem['href'] + ',' + advisory.replace(",",";") + ',' + elem.text.replace(',',';'), file=outfile)
    except Exception:
        print('CISA Alert: '+advisory+' cannot be parsed correctly', file=error)
outfile.close
error.close


from abc import ABC

import torch
import torch.nn as nn
import torch.nn.functional as tf

from vidar.arch.losses.BaseLoss import BaseLoss
from vidar.geometry.camera import Camera
from vidar.utils.depth import inv2depth
from vidar.utils.tensor import match_scales

def view_synthesis(ref_image, depth, ref_cam, cam,
                   mode='bilinear', padding_mode='zeros', align_corners=True):
    """Grid sample for view synthesis"""
    assert depth.shape[1] == 1, 'Depth map should have C=1'
    # Reconstruct world points from target_camera
    world_points = cam.reconstruct(depth, frame='w')
    # Project world points onto reference camera
    ref_coords = ref_cam.project(world_points, frame='w')
    # View-synthesis given the projected reference points
    return tf.grid_sample(ref_image, ref_coords, mode=mode,
                          padding_mode=padding_mode, align_corners=align_corners)

def gradient_x(image):
    """ Calculates image gradient along x-axis"""
    return image[:, :, :, :-1] - image[:, :, :, 1:]

def gradient_y(image):
    """ Calculates image gradient along y-axis"""
    return image[:, :, :-1, :] - image[:, :, 1:, :]

def inv_depths_normalize(inv_depths):
    """ Normalizes inverse depth maps by their mean"""
    mean_inv_depths = [inv_depth.mean(2, True).mean(3, True) for inv_depth in inv_depths]
    return [inv_depth / mean_inv_depth.clamp(min=1e-6)
            for inv_depth, mean_inv_depth in zip(inv_depths, mean_inv_depths)]

def calc_smoothness(inv_depths, images, num_scales):
    """ Calculates smoothness loss for inverse depth maps"""
    inv_depths_norm = inv_depths_normalize(inv_depths)
    inv_depth_gradients_x = [gradient_x(d) for d in inv_depths_norm]
    inv_depth_gradients_y = [gradient_y(d) for d in inv_depths_norm]

    image_gradients_x = [gradient_x(image) for image in images]
    image_gradients_y = [gradient_y(image) for image in images]

    weights_x = [torch.exp(-torch.mean(torch.abs(g), 1, keepdim=True)) for g in image_gradients_x]
    weights_y = [torch.exp(-torch.mean(torch.abs(g), 1, keepdim=True)) for g in image_gradients_y]

    # Note: Fix gradient addition
    smoothness_x = [inv_depth_gradients_x[i] * weights_x[i] for i in range(num_scales)]
    smoothness_y = [inv_depth_gradients_y[i] * weights_y[i] for i in range(num_scales)]
    return smoothness_x, smoothness_y

def SSIM(x, y, C1=1e-4, C2=9e-4, kernel_size=3, stride=1):
    """
    Structural Similarity (SSIM) distance between two images.

    Parameters
    ----------
    x,y : torch.Tensor
        Input images [B,3,H,W]
    C1,C2 : float
        SSIM parameters
    kernel_size,stride : int
        Convolutional parameters

    Returns
    -------
    ssim : torch.Tensor
        SSIM distance [1]
    """
    pool2d = nn.AvgPool2d(kernel_size, stride=stride)
    refl = nn.ReflectionPad2d(1)

    x, y = refl(x), refl(y)
    mu_x = pool2d(x)
    mu_y = pool2d(y)

    mu_x_mu_y = mu_x * mu_y
    mu_x_sq = mu_x.pow(2)
    mu_y_sq = mu_y.pow(2)

    sigma_x = pool2d(x.pow(2)) - mu_x_sq
    sigma_y = pool2d(y.pow(2)) - mu_y_sq
    sigma_xy = pool2d(x * y) - mu_x_mu_y
    v1 = 2 * sigma_xy + C2
    v2 = sigma_x + sigma_y + C2

    ssim_n = (2 * mu_x_mu_y + C1) * v1
    ssim_d = (mu_x_sq + mu_y_sq + C1) * v2
    ssim = ssim_n / ssim_d

    return ssim

class MultiViewPhotometricLoss(BaseLoss, ABC):
    """
    Self-Supervised multiview photometric loss.
    It takes two images, a depth map and a pose transformation to produce a
    reconstruction of one image from the perspective of the other, and calculates
    the difference between them

    Parameters
    ----------
    num_scales : int
        Number of inverse depth map scales to consider
    ssim_loss_weight : float
        Weight for the SSIM loss
    occ_reg_weight : float
        Weight for the occlusion regularization loss
    smooth_loss_weight : float
        Weight for the smoothness loss
    C1,C2 : float
        SSIM parameters
    photometric_reduce_op : str
        Method to reduce the photometric loss
    disp_norm : bool
        True if inverse depth is normalized for
    clip_loss : float
        Threshold for photometric loss clipping
    progressive_scaling : float
        Training percentage for progressive scaling (0.0 to disable)
    padding_mode : str
        Padding mode for view synthesis
    automask_loss : bool
        True if automasking is enabled for the photometric loss
    kwargs : dict
        Extra parameters
    """
    def __init__(self, num_scales=4, ssim_loss_weight=0.85, occ_reg_weight=0.1, smooth_loss_weight=0.1,
                 C1=1e-4, C2=9e-4, photometric_reduce_op='mean', disp_norm=True, clip_loss=0.5,
                 progressive_scaling=0.0, padding_mode='zeros', automask_loss=False, **kwargs):
        super().__init__()
        self.n = num_scales
        self.ssim_loss_weight = ssim_loss_weight
        self.occ_reg_weight = occ_reg_weight
        self.smooth_loss_weight = smooth_loss_weight
        self.C1 = C1
        self.C2 = C2
        self.photometric_reduce_op = photometric_reduce_op
        self.disp_norm = disp_norm
        self.clip_loss = clip_loss
        self.padding_mode = padding_mode
        self.automask_loss = automask_loss

        # Asserts
        if self.automask_loss:
            assert self.photometric_reduce_op == 'min', \
                'For automasking only the min photometric_reduce_op is supported.'

    @property
    def logs(self):
        """Returns class logs."""
        return {
            'num_scales': self.n,
        }

    def warp_ref_image(self, inv_depths, ref_image, K, ref_K, pose):
        """
        Warps a reference image to produce a reconstruction of the original one.

        Parameters
        ----------
        inv_depths : list[torch.Tensor]
            Inverse depth map of the original image [B,1,H,W]
        ref_image : torch.Tensor
            Reference RGB image [B,3,H,W]
        K : torch.Tensor
            Original camera intrinsics [B,3,3]
        ref_K : torch.Tensor
            Reference camera intrinsics [B,3,3]
        pose : Pose
            Original -> Reference camera transformation

        Returns
        -------
        ref_warped : torch.Tensor
            Warped reference image (reconstructing the original one) [B,3,H,W]
        """
        B, _, H, W = ref_image.shape
        device = ref_image.device
        # Generate cameras for all scales
        cams, ref_cams = [], []
        for i in range(self.n):
            _, _, DH, DW = inv_depths[i].shape
            scale_factor = DW / float(W)
            cams.append(Camera(K=K.float()).scaled(scale_factor).to(device))
            ref_cams.append(Camera(K=ref_K.float(), Tcw=pose).scaled(scale_factor).to(device))
        # View synthesis
        depths = [inv2depth(inv_depths[i]) for i in range(self.n)]
        ref_images = match_scales(ref_image, inv_depths, self.n)
        ref_warped = [view_synthesis(
            ref_images[i], depths[i], ref_cams[i], cams[i],
            padding_mode=self.padding_mode) for i in range(self.n)]
        # Return warped reference image
        return ref_warped

    def SSIM(self, x, y, kernel_size=3):
        """
        Calculates the SSIM (Structural Similarity) loss

        Parameters
        ----------
        x,y : torch.Tensor
            Input images [B,3,H,W]
        kernel_size : int
            Convolutional parameter

        Returns
        -------
        ssim : torch.Tensor
            SSIM loss [1]
        """
        ssim_value = SSIM(x, y, C1=self.C1, C2=self.C2, kernel_size=kernel_size)
        return torch.clamp((1. - ssim_value) / 2., 0., 1.)

    def calc_photometric_loss(self, t_est, images):
        """
        Calculates the photometric loss (L1 + SSIM)
        Parameters
        ----------
        t_est : list[torch.Tensor]
            List of warped reference images in multiple scales [B,3,H,W]
        images : list[torch.Tensor]
            List of original images in multiple scales [B,3,H,W]

        Returns
        -------
        photometric_loss : list[torch.Tensor]
            Photometric loss [B,1,H,W]
        """
        # L1 loss
        l1_loss = [torch.abs(t_est[i] - images[i])
                   for i in range(self.n)]
        # SSIM loss
        if self.ssim_loss_weight > 0.0:
            ssim_loss = [self.SSIM(t_est[i], images[i], kernel_size=3)
                         for i in range(self.n)]
            # Weighted Sum: alpha * ssim + (1 - alpha) * l1
            photometric_loss = [self.ssim_loss_weight * ssim_loss[i].mean(1, True) +
                                (1 - self.ssim_loss_weight) * l1_loss[i].mean(1, True)
                                for i in range(self.n)]
        else:
            photometric_loss = l1_loss
        # Clip loss
        if self.clip_loss > 0.0:
            for i in range(self.n):
                mean, std = photometric_loss[i].mean(), photometric_loss[i].std()
                photometric_loss[i] = torch.clamp(
                    photometric_loss[i], max=float(mean + self.clip_loss * std))
        # Return total photometric loss
        return photometric_loss

    def reduce_photometric_loss(self, photometric_losses):
        """
        Combine the photometric loss from all context images

        Parameters
        ----------
        photometric_losses : list[list[torch.Tensor]]
            Pixel-wise photometric losses from the entire context [B,3,H,W]

        Returns
        -------
        photometric_loss : torch.Tensor
            Reduced photometric loss [1]
        """
        # Reduce function
        def reduce_function(losses):
            if self.photometric_reduce_op == 'mean':
                return sum([l.mean() for l in losses]) / len(losses)
            elif self.photometric_reduce_op == 'min':
                return torch.cat(losses, 1).min(1, True)[0].mean()
            else:
                raise NotImplementedError(
                    'Unknown photometric_reduce_op: {}'.format(self.photometric_reduce_op))
        # Reduce photometric loss
        photometric_loss = sum([reduce_function(photometric_losses[i])
                                for i in range(self.n)]) / self.n
        # Store and return reduced photometric loss
        return photometric_loss

    def calc_smoothness_loss(self, inv_depths, images):
        """
        Calculates the smoothness loss for inverse depth maps.

        Parameters
        ----------
        inv_depths : list[torch.Tensor]
            Predicted inverse depth maps for all scales [B,1,H,W]
        images : list[torch.Tensor]
            Original images for all scales [B,3,H,W]

        Returns
        -------
        smoothness_loss : torch.Tensor
            Smoothness loss [1]
        """
        # Calculate smoothness gradients
        smoothness_x, smoothness_y = calc_smoothness(inv_depths, images, self.n)
        # Calculate smoothness loss
        smoothness_loss = sum([(smoothness_x[i].abs().mean() +
                                smoothness_y[i].abs().mean()) / 2 ** i
                               for i in range(self.n)]) / self.n
        # Apply smoothness loss weight
        smoothness_loss = self.smooth_loss_weight * smoothness_loss
        # Store and return smoothness loss
        return smoothness_loss


import logging
from .utils import (
    answer,
    get_args,
    get_args_raw,
    get_args_html,
)

logging = logging.getLogger(__name__)

class MessageMagic:
    def __init__(self, message) -> None:
        self.message = message
        self._init_msg()

    def _init_msg(self):
        """
        Initializes the `message` object with various methods and attributes.

        Parameters:
            None

        Returns:
            None

        Raises:
            Exception: If there is an error initializing the `MessageMagic` module.

        Side Effects:
            Sets the `answer` attribute of the `message` object.
            Sets the `get_args` attribute of the `message` object.
            Sets the `get_args_raw` attribute of the `message` object.
            Sets the `get_args_html` attribute of the `message` object.
            Logs a success message if initialization is successful.
            Logs an error message if there is an error initializing `MessageMagic`.
        """
        try:
            self.message.answer = answer
            self.message.get_args = get_args
            self.message.get_args_raw = get_args_raw
            self.message.get_args_html = get_args_html

        except Exception as e:
            logging.error(f"Error initializing MessageMagic: {str(e)}")


#=================================================================================================================
# AutoVEM2
# Author: Xi Binbin
# Email: 201766841276@mail.scut.edu.cn
#=================================================================================================================

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime
from math import ceil
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.patches as mpatches
import matplotlib.colors as mcolors

def plot(file, directory, days=7):
    '''
    Visualization

    :param file: data_plot.tsv
    :param directory: ouput directory
    :param days: interval days
    '''

    pdfFile = os.path.join(directory, 'hap_date.pdf')
    if os.path.exists(pdfFile):
        os.system(f"rm -rf {pdfFile}")
    df = pd.read_csv(file, sep='\t')
    df.index = pd.to_datetime(df["Date"])
    df = df.sort_index(ascending=True)

    n_genomes = df.shape[0]
    haplotypes = df["Name"].unique().tolist()
    countries = df["Country"].unique().tolist()
    n_haplotypes=len(haplotypes)
    n_countries = len(countries)

    n=days
    if n_genomes <= 1:
        print("Too little genomes.")
        sys.exit()
    start_day = df.index[0]
    end_day = df.index[-1]
    delta = end_day - start_day
    m = ceil(delta.days/n)
    n_periods = list(range(m))
    date_point = list()
    date_point.append(start_day.strftime('%Y/%m/%d'))
    for c in n_periods:
        b = pd.to_datetime(start_day+datetime.timedelta(days = int((c+1)*n)))
        date_point.append(b.strftime('%Y/%m/%d'))

    def name(n):
            haps = list()
            for i in list(range(n-1)):
                s = 'H' + str(i+1)
                haps.append(s)
            haps.append('other')
            return haps

    def size(n):
        m=0
        if(n>=1000):
            m=1
        elif(n>=100):
            m=0.85
        elif(n>=10):
            m=0.7
        else:
            m=0.5
        return m

    sort_hap = name(n_haplotypes)
    colors_all=list(mcolors.CSS4_COLORS.keys())
    colors=colors_all[10:(n_haplotypes+10)]
    colors_dict = dict(zip(sort_hap, colors))

    x_size = 2.5  + 0.396 + m*0.396 + 0.396 + 0.1 + 0.396*7
    y_size = 1.2 + 0.396 + n_countries*0.396 + 0.396 + 0.396*2
    a = x_origin = 2.5/x_size
    b = y_origin = 1.2/y_size
    c = x_len = (m+2)*0.396/x_size
    d = y_len = (n_countries+2)*0.396/y_size
    lx = 0.396/x_size
    ly = 0.396/y_size

    # pdf = PdfPages(pdfFile)
    plt.figure(figsize=(x_size, y_size))
    axes_main = plt.axes([a,b,c,d])
    axes_main.set_xlim(0, m+2)
    axes_main.set_ylim(0, n_countries+2)
    xticks = list(range(1, m+2))
    yticks = list(np.arange(1.5, n_countries+1, 1))
    xticklabels = date_point[:]
    yticklabels = countries[::-1]
    axes_main.set_xticks(xticks)
    axes_main.set_yticks(yticks)
    axes_main.set_xticklabels(xticklabels)
    axes_main.xaxis.set_tick_params(rotation=30, labelsize=12)
    axes_main.set_yticklabels(yticklabels, size=12)

    for c in n_periods:
        a = pd.to_datetime(start_day+datetime.timedelta(days = int(c*n)))
        b = pd.to_datetime(start_day+datetime.timedelta(days = int((c+1)*n)))
        if b == end_day:
            df2 = df[(df.index >=a) & (df.index <= b)]
        else :
            df2 = df[(df.index >=a) & (df.index < b)]
        for r, j in enumerate(yticklabels):
            df3 = df2[df2["Country"] == j]
            haps_count = list()
            haps_color = list()
            for i, k in enumerate(sort_hap):
                df4 = df3[df3["Name"] == k]
                hap_num = len(df4.index)
                haps_count.append(hap_num)
                haps_color.append(colors_dict[k])
            num_all = sum(haps_count)
            if num_all == 0:
                continue
            p1 = x_origin + lx + lx*c
            p2 = y_origin + ly + ly*r
            p3 = lx
            p4 = ly
            ax = plt.axes([p1, p2, p3, p4])
            ax.axis("off")
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            radius = size(num_all)
            ax.pie(haps_count, colors=haps_color, radius=radius, normalize=True)

    handle = list()
    for hap in sort_hap:
        hap_patch = mpatches.Patch(color=colors_dict[hap], label=str(hap))
        handle.append(hap_patch)
    p1 = (x_size-7*0.396)/x_size
    p2 = 1.2/y_size
    p3 = 4*0.396/x_size
    p4 = (y_size-1.2-0.396*2)/y_size
    axes_main = plt.axes([p1,p2,p3,p4])
    axes_main.axis("off")
    axes_main.set_xticks([])
    axes_main.set_yticks([])
    axes_main.legend(handles=handle, loc="upper center", ncol=1, frameon=False, fontsize=12)

    p2 = (x_size-4*0.396)/x_size
    en = (y_size-0.396*3)/y_size
    ax4=plt.axes([p2, en, lx, ly])
    ax4.pie([1],labels=[""],colors=['white'],rotatelabels=True, wedgeprops={'linewidth': 1, 'edgecolor': "black"}, radius=1.0)
    ax4.annotate("n>=1000",xy=(p1, en), fontsize=12)
    ax5=plt.axes([p2, en-ly, lx, ly])
    ax5.pie([1],labels=[""],colors=['white'],rotatelabels=True, wedgeprops={'linewidth': 1, 'edgecolor': "black"}, radius=0.85)
    ax5.annotate("n>=100",xy=(p1, en-ly), fontsize=12)
    ax6=plt.axes([p2, en-2*ly, lx, ly])
    ax6.pie([1],labels=[""],colors=['white'],rotatelabels=True, wedgeprops={'linewidth': 1, 'edgecolor': "black"}, radius=0.7)
    ax6.annotate("n>=10",xy=(p1, en-2*ly), fontsize=12)
    ax7=plt.axes([p2, en-3*ly, lx, ly])
    ax7.pie([1],labels=[""],colors=['white'],rotatelabels=True, wedgeprops={'linewidth': 1, 'edgecolor': "black"}, radius=0.5)
    ax7.annotate("n>=1",xy=(p1, en-3*ly), fontsize=12)
    plt.savefig(pdfFile)
    plt.close()

def module3(dataPlot, outDirectory, days):
    file = dataPlot
    if not os.path.exists(os.path.abspath(file)):
        print("Error: can't find the %s file"%(os.path.abspath(file)))
        sys.exit()

    print('Plotting...')
    plot(dataPlot, outDirectory, days)
    print('Done!')


"""Dashboard for the ninerchat server"""
from flask import Blueprint, render_template, request, flash, url_for, redirect
from flask_login import current_user
from flask import current_app as app
from application.auth import admin_only
from application.forms import SelectRoomForm
from application.models import db, User, Chatroom, MemberList, Messages
from flask_login import login_required

# Blueprints
dashboard_bp = Blueprint('dashboard_bp', __name__,
    template_folder='templates',
    static_folder='static',
    url_prefix='/admin')

@dashboard_bp.route('/')
@admin_only
def main():
    """ Main dashboard """
    return render_template('main_dashboard.html',
        title='NinerChat | Admin Dashboard',
        template='admin-dashboard',
        current_user=current_user,
        body="Admin Dashboard")

@dashboard_bp.route('/users')
@admin_only
def user():
    """ User dashboard """
    users = User.query.all()
    sort_by = request.args.get('sort')
    user_list = [
        {
            'id':u.id,
            'name':u.username,
            'email':u.email,
            'college':u.college,
            'major':u.major
        } for u in users
    ]
    if sort_by:
        user_list.sort(key=lambda x: x[sort_by])

    return render_template('user_dashboard.html',
        title='NinerChat | User Dashboard',
        template='user-dashboard',
        current_user=current_user,
        users=user_list,
        body="User Dashboard")

@dashboard_bp.route('/users/<id>', methods=['GET','POST'])
@admin_only
def user_profile(id):
    """ User profile for admins """

    room_select = SelectRoomForm(request.form)

    if request.method == 'POST' and room_select.validate():
        room_id = request.form.get('room_id')
        user_id = id
        return redirect(url_for(
            'room_bp.update_members',
            room_id=room_id,
            user_id=user_id
        ))
    elif request.method == 'GET':
        user = User.query.filter_by(id=int(id)).first()
        memberlist = MemberList.query.filter_by(user_id=int(id)).all()
        rooms = Chatroom.query.all()
        room_list = [
            (room.id,room.name) for room in rooms
        ]

        return render_template('user_profile_dashboard.html',
            title='NinerChat | User Profile',
            template='user-profile-dashboard',
            current_user=current_user,
            user=user,
            memberlist=memberlist,
            form=SelectRoomForm(),
            body="User Profile Dashboard")

@dashboard_bp.route('/users/<id>/messages', methods=['GET','POST'])
@admin_only
def user_messages(id):
    """ User messages for admins """
    user = User.query.filter_by(id=id).first()
    messages = Messages.query.filter_by(user_id=int(id)).all()
    messages_list = [
        {
            'id':int(m.id),
            'ts':str(m.ts),
            'chatroom_id':m.chatroom_id,
            'chatroom_name':m.chatroom.name,
            'text':m.text
        } for m in messages
    ]
    sort_by = request.args.get('sort')
    if sort_by:
        messages_list.sort(key=lambda x: x[sort_by])
    return render_template('user_message_dashboard.html',
            title='NinerChat | User Messages',
            template='user-message-dashboard',
            current_user=current_user,
            user=user,
            messages=messages_list,
            body="{} Messages Dashboard".format(user.username))

@dashboard_bp.route('/users/<id>/change')
@admin_only
def change_admin(id):
    """ Toggle user admin True/False """
    user = User.query.filter_by(id=int(id)).first()
    if user.email != 'ninerchat@uncc.edu':
        user.admin = not user.admin
        db.session.commit()
    else:
        flash('Ninerchat admin can not be changed')
    return redirect(url_for('dashboard_bp.user_profile', id = user.id))

@dashboard_bp.route('/rooms')
@admin_only
def room():
    """ Room dashboard """
    rooms = Chatroom.query.all()
    sort_by = request.args.get('sort')
    room_list = [
        {
            'id':room.id,
            'name': room.name,
            'public': room.public
        } for room in rooms
    ]
    if sort_by:
        room_list.sort(key=lambda x: x[sort_by])

    return render_template('room_dashboard.html',
        title='NinerChat | Room Dashboard',
        template='room-dashboard',
        rooms=room_list,
        current_user=current_user,
        body="Room Dashboard")

@dashboard_bp.route('/rooms/<id>')
@admin_only
def room_info(id):
    """ Room Info Dashboard """
    room = Chatroom.query.filter_by(id=id).first()
    members = MemberList.query.filter_by(chatroom_id=id).all()
    sort_by = request.args.get('sort')
    member_list = [
        {
            'id': member.user_id,
            'name': member.user.username,
            'email': member.user.email,
            'admin': member.user.admin
        } for member in members
    ]
    if sort_by:
        member_list.sort(key=lambda x: x[sort_by])

    return render_template('room_info_dashboard.html',
        title='NinerChat | Room Dashboard',
        template='room-dashboard',
        room=room,
        members=member_list,
        test=members,
        current_user=current_user,
        body="Room Info Dashboard")

@dashboard_bp.route('/reports')
@admin_only
def report():
    """ Report dashboard """
    return render_template('report_dashboard.html',
        title='NinerChat | Report Dashboard',
        template='report-dashboard',
        current_user=current_user,
        body="Report Dashboard")


#################################################
#################### IMPORTS ####################
#################################################

import json
import math
import os
import shutil
import ssl
import sys

import torch
import torchvision
import torchvision.transforms as transforms
import yaml
from roboflow import Roboflow
from tqdm import tqdm

import wandb

###########################################################
#################### GENERAL UTILITIES ####################
###########################################################

def initialize_wandb(cfg_fname):
    """
    Function to initialize wandb experiment, configured by the specified config file

    Arguments
    - cfg_fname (string): specifies the file location of model_cfg.json

    Returns:
    - model_cfg (dict): dictionary storing data from model_cfg.json
    - config (wandb.Config): config object for wandb experiment
    """
    with open(cfg_fname, "r") as f:
        model_cfg = json.load(f)
    model_cfg["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_cfg["project"] = f"oakd-research-{model_cfg['task'].replace('_', '-')}"
    model_cfg["lr_steps"] = [int(0.4*model_cfg["epochs"]), int(0.8*model_cfg["epochs"])]

    # WandB – Initialize a new run
    wandb.login()
    wandb.init(
        project=model_cfg["project"],
        group=model_cfg["model_arch"],
        name=model_cfg["name"],
    )
    # WandB – Config is a variable that holds and saves hyperparameters and inputs
    config = wandb.config
    for k, v in model_cfg.items():
        setattr(config, k, v)
    return model_cfg, config

########################################################
#################### DATA UTILITIES ####################
########################################################

def prepareTorchDataset(model_cfg):
    """
    Prepare a torch dataset, parameterized by the model config

    Arguments:
    - model_cfg (wandb.Config): the wandb Config object

    Returns:
    - trainloader (torch.utils.data.DataLoader): training torch DataLoader
    - testloader (torch.utils.data.DataLoader): testing torch DataLoader
    """
    # Hacky fix for turning off SSL verification to handle URLError
    ssl._create_default_https_context = ssl._create_unverified_context

    # Define transforms
    transform_train_list = []
    transform_test_list = []

    # AutoAugmentation Policy for all tasks except generative ones
    if "gan" not in model_cfg["name"]:
        transform_train.append(transforms.AutoAugment(
                policy=transforms.autoaugment.AutoAugmentPolicy.CIFAR10,
                interpolation=transforms.functional.InterpolationMode.BILINEAR
            )
        )

    # Ensure Tensors
    transform_train_list.append(transforms.ToTensor())
    transform_test_list.append(transforms.ToTensor())

    # Resize Tensors
    if not isinstance(model_cfg["imgsz"], list):
        model_cfg["imgsz"] = [model_cfg["imgsz"]]
    transform_train_list.append(transforms.Resize(model_cfg["imgsz"][1:], antialias=True))
    transform_test_list.append(transforms.Resize(model_cfg["imgsz"][1:], antialias=True))

    # Normalize Tensors
    # transform_train_list.append(
    #     transforms.Normalize(
    #         (0.5, 0.5, 0.5),
    #         (0.5, 0.5, 0.5)
    #     )
    # )
    # transform_test_list.append(
    #     transforms.Normalize(
    #         (0.5, 0.5, 0.5),
    #         (0.5, 0.5, 0.5)
    #     )
    # )

    ############ Temp for MNIST ############
    transform_train_list.append(
        transforms.Normalize((0.1307,), (0.3081,))
    )
    transform_test_list.append(
        transforms.Normalize((0.1307,), (0.3081,))
    )
    transform_train = transforms.Compose(transform_train_list)
    transform_test = transforms.Compose(transform_test_list)

    # Read the datasets for training and testing
    try:
        dataset_name_split = model_cfg["dataset_name"].split("/")
        dataset_source = str(dataset_name_split[0])
        dataset_name = str(dataset_name_split[1])
        data_root = f"./datasets/{model_cfg['task']}/{dataset_source}"
        trainset = getattr(torchvision.datasets, dataset_name)(
            root=data_root,
            train=True,
            download=True,
            transform=transform_train
        )
        testset = getattr(torchvision.datasets, dataset_name)(
            root=data_root,
            train=False,
            download=True,
            transform=transform_test
        )
    except Exception as e1:
        try:
            trainset = getattr(torchvision.datasets, dataset_name)(
                root=data_root,
                split="train",
                download=True,
                transform=transform_train
            )
            testset = getattr(torchvision.datasets, dataset_name)(
                root=data_root,
                split="valid",
                download=True,
                transform=transform_test
            )
        except Exception as e2:
            print("Error Loading Torch Dataset, dataset does not have argument train or split")
            print(e1)
            print(e2)
            sys.exit(1)

    # Create the corresponding dataloaders
    trainloader = torch.utils.data.DataLoader(
        trainset,
        batch_size=model_cfg.batch_size,
        shuffle=True,
        pin_memory=True,
        num_workers=8
    )
    testloader = torch.utils.data.DataLoader(
        testset,
        batch_size=model_cfg.test_batch_size,
        shuffle=False,
        pin_memory=True,
        num_workers=2
    )

    return trainloader, testloader

def prepareRoboFlowDataset(model_cfg):
    """
    Prepare a RoboFlow Dataset in the yolov8 annotation format

    Arguments:
    - model_cfg (wandb.Config): the wandb Config object
    """
    rf_key = os.getenv("ROBOFLOW_API_KEY")
    if rf_key is None:
        rf_key = input("Roboflow API Key: ")
        os.environ["ROBOFLOW_API_KEY"] = rf_key
    rf = Roboflow(api_key=rf_key)

    try:
        ws, pr, vs = model_cfg["dataset_name"].split("/")
        project = rf.workspace(ws).project(pr)
        dataset = project.version(vs).download(model_format="yolov8", location=f"./datasets/{model_cfg['task']}/{model_cfg['dataset_name']}/", overwrite=False)
    except:
        print("For roboflow datasets, you must specify the workspace, project, and version in the model_cfg.json field=dataset_name")
        sys.exit(1)

def prepareDetectionDataset(model_cfg):
    """
    Simple Function for ensuring that dataset exists and is saved locally

    Arguments:
    - model_cfg (wandb.Config): the wandb Config object
    """
    dataset_path = f"./datasets/{model_cfg['task']}/{model_cfg['dataset_name']}/"
    if not os.path.exists(dataset_path):
        prepareRoboFlowDataset(model_cfg)

def coco2yolo(model_cfg):
    """
    Function to convert MS-COCO Labels to YOLOv8 format
    """
    data_root = f"./datasets/{model_cfg['task']}/{model_cfg['dataset_name']}/"
    for root, folders, files in os.walk(data_root):
        if "coco" in root:
            if "train.json" in files:
                nc_list = [] # Collect list of classes
                yolo_dir = root.replace("coco", "yolo")
                os.makedirs(os.path.join(yolo_dir, "train", "images"), exist_ok=True)
                os.makedirs(os.path.join(yolo_dir, "train", "labels"), exist_ok=True)
                os.makedirs(os.path.join(yolo_dir, "valid", "images"), exist_ok=True)
                os.makedirs(os.path.join(yolo_dir, "valid", "labels"), exist_ok=True)
                os.makedirs(os.path.join(yolo_dir, "test", "images"), exist_ok=True)
                os.makedirs(os.path.join(yolo_dir, "test", "labels"), exist_ok=True)

                ##### Go through train.json, val.json, test.json
                for dataset_json in ["train", "val", "test"]:
                    dataset_json_fpath = os.path.join(root, dataset_json)
                    with open(dataset_json_fpath + ".json", "r") as f:
                        dataset_json_dict = json.load(f)
                        id_sorted_ims = sorted(dataset_json_dict["images"], key=lambda im: im["id"])
                    if dataset_json == "val":
                        dataset_json = "valid"
                    ##### Go though the "images" and "annotations" keys
                    with tqdm(iterable=id_sorted_ims) as tq:
                        for im in id_sorted_ims:
                            coco_image_fpath = os.path.join(root, "images", im["file_name"])
                            yolo_image_fpath = os.path.join(yolo_dir, dataset_json, "images", im["file_name"])
                            coco_label_fpath = os.path.join(root, "annotations", im["file_name"].replace("jpg", "json"))
                            yolo_label_fpath = os.path.join(yolo_dir, dataset_json, "labels", im["file_name"].replace("jpg", "txt"))

                            ##### Copy image to corresponding images folder
                            try:
                                shutil.copy(coco_image_fpath, yolo_image_fpath)
                            except Exception as e:
                                print(e)
                                continue

                            ##### Create txt file of same name in a labels folder
                            with open(coco_label_fpath, "r") as coco_f:
                                coco_label_dict = json.load(coco_f)

                            img_wh = [coco_label_dict["imageWidth"], coco_label_dict["imageHeight"]] # Use this to normalize labels
                            yolo_label_list = []
                            for shape in coco_label_dict["shapes"]: # for each instance
                                if shape["label"] not in nc_list:
                                    nc_list.append(shape["label"]) # add the class to our list if we haven't seen it before
                                label_ind = nc_list.index(shape["label"])
                                points = [str(label_ind)]
                                for pt in shape["points"]:
                                    points.extend([str(pt[0] / img_wh[0]), str(pt[1] / img_wh[1])])
                                yolo_label_list.append(" ".join(points))
                            with open(yolo_label_fpath, "w") as yolo_f:
                                yolo_f.write("\n".join(yolo_label_list))
                            tq.update()
                # Create data.yaml file
                dataset_yaml = {
                    "nc": len(nc_list),
                    "names": nc_list,
                    "path": yolo_dir.replace("./datasets/", "./"),
                    "train": "train/images",
                    "val": "valid/images",
                    "test": "test/images"
                }
                with open(os.path.join(yolo_dir, "data.yaml"), "w") as yolo_f:
                    yaml.dump(dataset_yaml, yolo_f)

###############################################################
#################### ALGORITHMIC UTILITIES ####################
###############################################################

def lr_finder_algo(net, optimizer, scheduler, criterion, train_loader, model_cfg):
	"""
    Function to perform a search for the optimal initial learning rate for the Cosine-Annealing with Warmup scheduler.
    Algorithm is implemented according to https://arxiv.org/abs/1506.01186

    Arguments:
    - model (torch.nn.Module): the torch model to be trained
    - optimizer (torch.optim): the torch optimizer to use
    - scheduler (torch.optim.lr_scheduler)
    - criterion (torch.nn.Loss): the torch loss function
    - train_loader (torch.utils.data.DataLoader): the torch Dataloader
    - model_cfg (wandb.Config): the wandb Config object

    Returns:
    - eta_max (float): the maximized initial learning rate
	"""
	device, epochs = model_cfg["device"], max(150, model_cfg["epochs"])
	model = net.to(device)
	losses = []
	lr_list = []
	for epoch, (images, labels) in enumerate(train_loader):
		if epoch == epochs:
			break
		# Move tensors to configured device
		images = images.to(device)
		labels = labels.to(device)
		# Forward Pass
		outputs = model(images)
		loss = criterion(outputs, labels)
		# Backpropogation and SGD
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()
		curr_lr = optimizer.param_groups[0]['lr']
		wandb.log(
			{
				"Training Loss (LR Finder)": loss.item(),
				"LR": curr_lr,
				"Epoch": epoch
			}
		)
		losses.append(loss.item())
		lr_list.append(curr_lr)
		scheduler.step()
	eta_max = lr_list[min(enumerate(losses), key=lambda x: x[1])[0]] / 10.0
	return eta_max

def custom_cos_annealing_warmup(num_batches, T, eta_max, epoch, batch):
    """
    Custom Implementation of a cosine annealing learning rate scheduler with warmup

    Arguments:
    - num_batches (int): number of batches in the dataset
    - T (int): num_batches*epochs
    - eta_max (float): the maximized initial learning rate
    - epoch (int): the current epoch
    - batch (int): the current batch number/id

    Returns:
    - next_lr (float): the next scheduled learning rate
	"""
	# fixed hyperparameters used for Cosine Annealing w/ Warmup Schedule
    T_0 = int(T / 5)
    t = epoch*num_batches + batch
    next_lr = 1e-5
    if t <= T_0:
        next_lr += (t / T_0)*eta_max
    else:
        next_lr += eta_max*math.cos((math.pi/2.0)*((t - T_0)/(T - T_0)))
    return next_lr


import constants
from shop.storage import memoryStorage

from flask import jsonify, request

def get_products():
    options = dict()
    if "offset" in request.args:
        options["offset"] = request.args.get("offset", type=int)
    if "count" in request.args:
        options["count"] = request.args.get("count", type=int)

    products, status = memoryStorage.get_products(**options)
    http_status = constants.responses[status]
    total_count = memoryStorage.get_products_count()
    if len(options) != 0:
        options["count"] = min(max(0, total_count - options.get("offset", 0)), options.get("count", 0))
    return jsonify(products=products, total_count=total_count, status=status, **options), http_status

def get_product(prod_id=None):
    prod_id = prod_id or request.args.get("id")
    product, status = memoryStorage.get_product(prod_id)
    http_status = constants.responses[status]

    if status == constants.statuses["product"]["returned"]:
        body = jsonify(product=product, status=status)
    else:
        body = f"no such product id: {prod_id}"
    return body, http_status