from copy import deepcopy
import logging
from pathlib import Path
from typing import Dict, List, Tuple
import json_tricks

import numpy as np
from torch import Tensor
import torch
from torch.nn import Module

from ........nni.algorithms.compression.v2.pytorch.base import Task, TaskResult
from ........nni.algorithms.compression.v2.pytorch.utils import (
    config_list_canonical,
    compute_sparsity,
    get_model_weights_numel
)
from .base import TaskGenerator

_logger = logging.getLogger(__name__)

class FunctionBasedTaskGenerator(TaskGenerator):
    def __init__(self, total_iteration: int, origin_model: Module, origin_config_list: List[Dict],
                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermediate_result: bool = False):
        """
        Parameters
        ----------
        total_iteration
            The total iteration number.
        origin_model
            The origin unwrapped pytorch model to be pruned.
        origin_config_list
            The origin config list provided by the user. Note that this config_list is directly config the origin model.
            This means the sparsity provided by the origin_masks should also be recorded in the origin_config_list.
        origin_masks
            The pre masks on the origin model. This mask maybe user-defined or maybe generate by previous pruning.
        log_dir
            The log directory use to saving the task generator log.
        keep_intermediate_result
            If keeping the intermediate result, including intermediate model and masks during each iteration.
        """
        self.total_iteration = total_iteration
        super().__init__(origin_model, origin_config_list=origin_config_list, origin_masks=origin_masks,
                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)

    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_iteration = 0
        self.target_sparsity = config_list_canonical(model, config_list)
        super().reset(model, config_list=config_list, masks=masks)

    def init_pending_tasks(self) -> List[Task]:
        origin_model = torch.load(self._origin_model_path)
        origin_masks = torch.load(self._origin_masks_path)

        task_result = TaskResult('origin', origin_model, origin_masks, origin_masks, None)

        return self.generate_tasks(task_result)

    def generate_tasks(self, task_result: TaskResult) -> List[Task]:
        compact_model = task_result.compact_model
        compact_model_masks = task_result.compact_model_masks

        # save intermediate result
        model_path = Path(self._intermediate_result_dir, '{}_compact_model.pth'.format(task_result.task_id))
        masks_path = Path(self._intermediate_result_dir, '{}_compact_model_masks.pth'.format(task_result.task_id))
        torch.save(compact_model, model_path)
        torch.save(compact_model_masks, masks_path)

        # get current2origin_sparsity and compact2origin_sparsity
        origin_model = torch.load(self._origin_model_path)
        current2origin_sparsity, compact2origin_sparsity, _ = compute_sparsity(origin_model, compact_model, compact_model_masks, self.target_sparsity)
        _logger.info('\nTask %s total real sparsity compared with original model is:\n%s', str(task_result.task_id), json_tricks.dumps(current2origin_sparsity, indent=4))
        if task_result.task_id != 'origin':
            self._tasks[task_result.task_id].state['current2origin_sparsity'] = current2origin_sparsity

        # if reach the total_iteration, no more task will be generated
        if self.current_iteration > self.total_iteration:
            return []

        task_id = self._task_id_candidate
        new_config_list = self.generate_config_list(self.target_sparsity, self.current_iteration, compact2origin_sparsity)
        new_config_list = self.allocate_sparsity(new_config_list, compact_model, compact_model_masks)
        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))

        with Path(config_list_path).open('w') as f:
            json_tricks.dump(new_config_list, f, indent=4)
        task = Task(task_id, model_path, masks_path, config_list_path)

        self._tasks[task_id] = task

        self._task_id_candidate += 1
        self.current_iteration += 1

        return [task]

    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        raise NotImplementedError()

    def allocate_sparsity(self, new_config_list: List[Dict], model: Module, masks: Dict[str, Dict[str, Tensor]]):
        return new_config_list

class AGPTaskGenerator(FunctionBasedTaskGenerator):
    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            ori_sparsity = (1 - (1 - iteration / self.total_iteration) ** 3) * target['total_sparsity']
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class LinearTaskGenerator(FunctionBasedTaskGenerator):
    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            ori_sparsity = iteration / self.total_iteration * target['total_sparsity']
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class LotteryTicketTaskGenerator(FunctionBasedTaskGenerator):
    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_iteration = 1
        self.target_sparsity = config_list_canonical(model, config_list)
        super(FunctionBasedTaskGenerator, self).reset(model, config_list=config_list, masks=masks)

    def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
        config_list = []
        for target, mo in zip(target_sparsity, compact2origin_sparsity):
            # NOTE: The ori_sparsity calculation formula in compression v1 is as follow, it is different from the paper.
            # But the formula in paper will cause numerical problems, so keep the formula in compression v1.
            ori_sparsity = 1 - (1 - target['total_sparsity']) ** (iteration / self.total_iteration)
            # The following is the formula in paper.
            # ori_sparsity = (target['total_sparsity'] * 100) ** (iteration / self.total_iteration) / 100
            sparsity = max(0.0, (ori_sparsity - mo['total_sparsity']) / (1 - mo['total_sparsity']))
            assert 0 <= sparsity <= 1, 'sparsity: {}, ori_sparsity: {}, model_sparsity: {}'.format(sparsity, ori_sparsity, mo['total_sparsity'])
            config_list.append(deepcopy(target))
            config_list[-1]['total_sparsity'] = sparsity
        return config_list

class SimulatedAnnealingTaskGenerator(TaskGenerator):
    def __init__(self, origin_model: Module, origin_config_list: List[Dict], origin_masks: Dict[str, Dict[str, Tensor]] = {},
                 start_temperature: float = 100, stop_temperature: float = 20, cool_down_rate: float = 0.9,
                 perturbation_magnitude: float = 0.35, log_dir: str = '.', keep_intermediate_result: bool = False):
        """
        Parameters
        ----------
        origin_model
            The origin unwrapped pytorch model to be pruned.
        origin_config_list
            The origin config list provided by the user. Note that this config_list is directly config the origin model.
            This means the sparsity provided by the origin_masks should also be recorded in the origin_config_list.
        origin_masks
            The pre masks on the origin model. This mask maybe user-defined or maybe generate by previous pruning.
        start_temperature
            Start temperature of the simulated annealing process.
        stop_temperature
            Stop temperature of the simulated annealing process.
        cool_down_rate
            Cool down rate of the temperature.
        perturbation_magnitude
            Initial perturbation magnitude to the sparsities. The magnitude decreases with current temperature.
        log_dir
            The log directory use to saving the task generator log.
        keep_intermediate_result
            If keeping the intermediate result, including intermediate model and masks during each iteration.
        """
        self.start_temperature = start_temperature
        self.stop_temperature = stop_temperature
        self.cool_down_rate = cool_down_rate
        self.perturbation_magnitude = perturbation_magnitude

        super().__init__(origin_model, origin_masks=origin_masks, origin_config_list=origin_config_list,
                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)

    def reset(self, model: Module, config_list: List[Dict] = [], masks: Dict[str, Dict[str, Tensor]] = {}):
        self.current_temperature = self.start_temperature

        self.weights_numel, self.masked_rate = get_model_weights_numel(model, config_list, masks)
        self.target_sparsity_list = config_list_canonical(model, config_list)
        self._adjust_target_sparsity()

        self._temp_config_list = None
        self._current_sparsity_list = None
        self._current_score = None

        super().reset(model, config_list=config_list, masks=masks)

    def _adjust_target_sparsity(self):
        """
        If origin_masks is not empty, then re-scale the target sparsity.
        """
        if len(self.masked_rate) > 0:
            for config in self.target_sparsity_list:
                sparsity, op_names = config['total_sparsity'], config['op_names']
                remaining_weight_numel = 0
                pruned_weight_numel = 0
                for name in op_names:
                    remaining_weight_numel += self.weights_numel[name]
                    if name in self.masked_rate and self.masked_rate[name] != 0:
                        pruned_weight_numel += 1 / (1 / self.masked_rate[name] - 1) * self.weights_numel[name]
                total_mask_rate = pruned_weight_numel / (pruned_weight_numel + remaining_weight_numel)
                config['total_sparsity'] = max(0, (sparsity - total_mask_rate) / (1 - total_mask_rate))

    def _init_temp_config_list(self):
        self._temp_config_list = []
        self._temp_sparsity_list = []
        for config in self.target_sparsity_list:
            sparsity_config, sparsity = self._init_config_sparsity(config)
            self._temp_config_list.extend(sparsity_config)
            self._temp_sparsity_list.append(sparsity)

    def _init_config_sparsity(self, config: Dict) -> Tuple[List[Dict], List]:
        assert 'total_sparsity' in config, 'Sparsity must be set in config: {}'.format(config)
        target_sparsity = config['total_sparsity']
        op_names = config['op_names']

        if target_sparsity == 0:
            return [], []

        while True:
            random_sparsity = sorted(np.random.uniform(0, 1, len(op_names)))
            rescaled_sparsity = self._rescale_sparsity(random_sparsity, target_sparsity, op_names)
            if rescaled_sparsity is not None and rescaled_sparsity[0] >= 0 and rescaled_sparsity[-1] < 1:
                break

        return self._sparsity_to_config_list(rescaled_sparsity, config), rescaled_sparsity

    def _rescale_sparsity(self, random_sparsity: List, target_sparsity: float, op_names: List) -> List:
        assert len(random_sparsity) == len(op_names)

        num_weights = sorted([self.weights_numel[op_name] for op_name in op_names])
        sparsity = sorted(random_sparsity)

        total_weights = 0
        total_weights_pruned = 0

        # calculate the scale
        for idx, num_weight in enumerate(num_weights):
            total_weights += num_weight
            total_weights_pruned += int(num_weight * sparsity[idx])
        if total_weights_pruned == 0:
            return None

        scale = target_sparsity / (total_weights_pruned / total_weights)

        # rescale the sparsity
        sparsity = np.asarray(sparsity) * scale
        return sparsity

    def _sparsity_to_config_list(self, sparsity: List, config: Dict) -> List[Dict]:
        sparsity = sorted(sparsity)
        op_names = [k for k, _ in sorted(self.weights_numel.items(), key=lambda item: item[1]) if k in config['op_names']]
        assert len(sparsity) == len(op_names)
        return [{'total_sparsity': sparsity, 'op_names': [op_name]} for sparsity, op_name in zip(sparsity, op_names)]

    def _update_with_perturbations(self):
        self._temp_config_list = []
        self._temp_sparsity_list = []
        # decrease magnitude with current temperature
        magnitude = self.current_temperature / self.start_temperature * self.perturbation_magnitude
        for config, current_sparsity in zip(self.target_sparsity_list, self._current_sparsity_list):
            if len(current_sparsity) == 0:
                self._temp_sparsity_list.append([])
                continue
            while True:
                perturbation = np.random.uniform(-magnitude, magnitude, len(current_sparsity))
                temp_sparsity = np.clip(0, current_sparsity + perturbation, None)
                temp_sparsity = self._rescale_sparsity(temp_sparsity, config['total_sparsity'], config['op_names'])
                if temp_sparsity is not None and temp_sparsity[0] >= 0 and temp_sparsity[-1] < 1:
                    self._temp_config_list.extend(self._sparsity_to_config_list(temp_sparsity, config))
                    self._temp_sparsity_list.append(temp_sparsity)
                    break

    def _recover_real_sparsity(self, config_list: List[Dict]) -> List[Dict]:
        """
        If the origin masks is not None, then the sparsity in new generated config_list need to be rescaled.
        """
        for config in config_list:
            assert len(config['op_names']) == 1
            op_name = config['op_names'][0]
            if op_name in self.masked_rate:
                config['total_sparsity'] = self.masked_rate[op_name] + config['total_sparsity'] * (1 - self.masked_rate[op_name])
        return config_list

    def init_pending_tasks(self) -> List[Task]:
        origin_model = torch.load(self._origin_model_path)
        origin_masks = torch.load(self._origin_masks_path)

        self.temp_model_path = Path(self._intermediate_result_dir, 'origin_compact_model.pth')
        self.temp_masks_path = Path(self._intermediate_result_dir, 'origin_compact_model_masks.pth')
        torch.save(origin_model, self.temp_model_path)
        torch.save(origin_masks, self.temp_masks_path)

        task_result = TaskResult('origin', origin_model, origin_masks, origin_masks, None)

        return self.generate_tasks(task_result)

    def generate_tasks(self, task_result: TaskResult) -> List[Task]:
        # initial/update temp config list
        if self._temp_config_list is None:
            self._init_temp_config_list()
        else:
            score = self._tasks[task_result.task_id].score
            if self._current_sparsity_list is None:
                self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
                self._current_score = score
            else:
                delta_E = np.abs(score - self._current_score)
                probability = np.exp(-1 * delta_E / self.current_temperature)
                if self._current_score < score or np.random.uniform(0, 1) < probability:
                    self._current_score = score
                    self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
                    self.current_temperature *= self.cool_down_rate
            if self.current_temperature < self.stop_temperature:
                return []
            self._update_with_perturbations()

        task_id = self._task_id_candidate
        new_config_list = self._recover_real_sparsity(deepcopy(self._temp_config_list))
        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))

        with Path(config_list_path).open('w') as f:
            json_tricks.dump(new_config_list, f, indent=4)

        task = Task(task_id, self.temp_model_path, self.temp_masks_path, config_list_path)

        self._tasks[task_id] = task

        self._task_id_candidate += 1

        return [task]


# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
from collections import deque

class Solution(object):
    def levelOrder(self, root):
        """
        :type root: TreeNode
        :rtype: List[List[int]]
        """
        q = deque([root])
        res = []
        while q:
            nd = q.popleft()
            res.append(nd.val)
            if nd.left:
                q.append(nd.left)
            if nd.right:
                q.append(nd.right)
        return res

a = TreeNode(3)
a.left = TreeNode(9)
a.right = TreeNode(20)
a.right.left = TreeNode(15)
a.right.right = TreeNode(7)
x = Solution()
print(x.levelOrder(a))
# 遍历所有node bfs


import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["savefig.dpi"] = 1200
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

water = np.array([77.2 / 100, 22.4 / 100, 0.4 / 100])
plt.pie(water, labels=['冰川、冰盖', '地下水', '可供人类利用的水'], colors=["#d5695d", "#5d8ca8", "#65a479"],
        autopct='%.2f%%', textprops={'fontsize': 7})
plt.title("淡水资源分布扇形图")
plt.show()

淡水水资源分布 扇形统计图.py

淡水水资源分布 扇形统计图.py


import os
import numpy as np
from PIL import Image
from netCDF4 import Dataset
from extract import PathSet

root_path = './Data/ABI-L1b-RadC/'
select_date = '2017253'
save_path = './one-day-picture/'

path_set = PathSet(root_path=root_path)
url = path_set.offer_url(select_date=select_date)

while(True):
    try:
        data_list = next(url)
        if len(data_list) <= 0:
            continue

        dl = data_list[0]

        g16nc = Dataset(dl[1])
        radiance = g16nc.variables['Rad'][:]
        radiance = (radiance + 25.936) / (804.036 + 25.936)

        img = Image.fromarray(radiance)

        save_name = os.path.join(save_path, dl[0]+'.tiff')

        img.save(save_name)

    except StopIteration:
        break


import dill
from sqlalchemy.sql.expression import true
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from dnc import DNC
import dhg
from layers import GraphConvolution
import math
from torch.nn.parameter import Parameter
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as autograd
import torch.optim as optim
from torch.autograd import Variable
from torch.nn.parameter import Parameter
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModel
import numpy as np
import random
import pickle as pkl

longTensor = torch.cuda.LongTensor
floatTensor = torch.cuda.FloatTensor

class MolecularGraphNeuralNetwork(nn.Module):
    def __init__(self, inputs, average_projection, ddi_adj_all, ddi_adj_pn, N_fingerprint, dim, layer_hidden, device):
        super(MolecularGraphNeuralNetwork, self).__init__()
        self.device = device
        self.embed_fingerprint = nn.Embedding(N_fingerprint, dim).to(self.device)
        self.W_fingerprint = nn.ModuleList([nn.Linear(dim, dim).to(self.device)
                                            for _ in range(layer_hidden)])
        self.layer_hidden = layer_hidden
        self.conv = nn.Linear(64, 64)
        self.loss_function = marginLoss()

        self.ddi_adj_all = ddi_adj_all
        self.ddi_adj_pn = ddi_adj_pn

        self.fingerprints, self.adjacencies, self.molecular_sizes = inputs
        self.average_projection = average_projection
        self.adjacencies = self.pad(self.adjacencies, 0)

    def pad(self, matrices, pad_value):
        """Pad the list of matrices
        with a pad_value (e.g., 0) for batch processing.
        For example, given a list of matrices [A, B, C],
        we obtain a new matrix [A00, 0B0, 00C],
        where 0 is the zero (i.e., pad value) matrix.
        """
        shapes = [m.shape for m in matrices]
        M, N = sum([s[0] for s in shapes]), sum([s[1] for s in shapes])
        zeros = torch.FloatTensor(np.zeros((M, N))).to(self.device)
        pad_matrices = pad_value + zeros
        i, j = 0, 0
        for k, matrix in enumerate(matrices):
            m, n = shapes[k]
            pad_matrices[i:i+m, j:j+n] = matrix
            i += m
            j += n
        return pad_matrices

    def update(self, matrix, vectors, layer):
        hidden_vectors = torch.relu(self.W_fingerprint[layer](vectors))
        return hidden_vectors + torch.mm(matrix, hidden_vectors)

    def sum(self, vectors, axis):
        sum_vectors = [torch.sum(v, 0) for v in torch.split(vectors, axis)]
        return torch.stack(sum_vectors)

    def mean(self, vectors, axis):
        mean_vectors = [torch.mean(v, 0) for v in torch.split(vectors, axis)]
        return torch.stack(mean_vectors)

    def triplet_loss(self, MPNN_emb):
        # MPNN_emb triplet loss calculation

        ent_batch = []
        pos_batch = []
        neg_batch = []
        for i in range(128):
            while True:
                ent = random.randint(0, 265)
                idx = np.argwhere(self.ddi_adj_pn[ent] == 1)
                if len(idx) == 0:
                    continue
                pos = idx[random.randint(0, len(idx) - 1)][0]
                idx = np.argwhere(self.ddi_adj_all[ent] == 1)
                if len(idx) == 0:
                    continue
                neg = idx[random.randint(0, len(idx) - 1)][0]
                break

            ent_batch.append(ent)
            pos_batch.append(pos)
            neg_batch.append(neg)

        ent_batch = torch.LongTensor(ent_batch)
        pos_batch = torch.LongTensor(pos_batch)
        neg_batch = torch.LongTensor(neg_batch)

        ent_emb = MPNN_emb[ent_batch, :].unsqueeze(1)
        pos_emb = MPNN_emb[pos_batch, :].unsqueeze(1)
        neg_emb = MPNN_emb[neg_batch, :].unsqueeze(1)

        ent_emb = self.conv(ent_emb).squeeze(1)
        pos_emb = self.conv(pos_emb).squeeze(1)
        neg_emb = self.conv(neg_emb).squeeze(1)

        pos = torch.sum(torch.abs(ent_emb - pos_emb), 1)
        neg = torch.sum(torch.abs(ent_emb - neg_emb), 1)

        margin = torch.FloatTensor([1.0])
        emb_loss = self.loss_function(pos, neg, margin) + normLoss(ent_emb) + normLoss(pos_emb) + normLoss(neg_emb)
        # new_MPNN_emb = self.conv(self.MPNN_emb)

        return emb_loss

    def forward(self):

        """Cat or pad each input data for batch processing."""
        fingerprints = torch.cat(self.fingerprints)
        # self.fingerprints= self.fingerprints[0:2]

        """MPNN layer (update the fingerprint vectors)."""
        fingerprint_vectors = self.embed_fingerprint(fingerprints)

        for l in range(self.layer_hidden):
            hs = self.update(self.adjacencies, fingerprint_vectors, l)
            # fingerprint_vectors = F.normalize(hs, 2, 1)  # normalize.
            fingerprint_vectors = hs

        """Molecular vector by sum or mean of the fingerprint vectors."""
        molecular_vectors = self.sum(fingerprint_vectors, self.molecular_sizes)
        # print(molecular_vectors)
        # molecular_vectors = self.mean(fingerprint_vectors, molecular_sizes)
        MPNN_emb = torch.mm(self.average_projection, molecular_vectors)

        return MPNN_emb

class GCN(nn.Module):
    def __init__(self, voc_size, emb_dim, adj, device=torch.device('cpu:0')):
        super(GCN, self).__init__()
        self.voc_size = voc_size
        self.emb_dim = emb_dim
        self.device = device

        adj = self.normalize(adj + np.eye(adj.shape[0]))

        self.adj = torch.FloatTensor(adj).to(device)
        self.x = torch.eye(voc_size).to(device)

        self.gcn1 = GraphConvolution(voc_size, emb_dim)
        self.dropout = nn.Dropout(p=0.3)
        self.gcn2 = GraphConvolution(emb_dim, emb_dim)

    def forward(self):
        node_embedding = self.gcn1(self.x, self.adj)
        node_embedding = F.relu(node_embedding)
        node_embedding = self.dropout(node_embedding)
        node_embedding = self.gcn2(node_embedding, self.adj)
        return node_embedding

    def normalize(self, mx):
        """Row-normalize sparse matrix"""
        rowsum = np.array(mx.sum(1))
        r_inv = np.power(rowsum, -1).flatten()
        r_inv[np.isinf(r_inv)] = 0.
        r_mat_inv = np.diagflat(r_inv)
        mx = r_mat_inv.dot(mx)
        return mx

class marginLoss(nn.Module):
    def __init__(self):
        super(marginLoss, self).__init__()

    def forward(self, pos, neg, margin):
        zero_tensor = floatTensor(pos.size()).cpu()
        zero_tensor.zero_()
        zero_tensor = autograd.Variable(zero_tensor)
        return torch.sum(torch.max(pos - neg + margin, zero_tensor))

def normLoss(embeddings):
    norm = torch.sum(embeddings)
    return torch.sum(torch.max(norm - autograd.Variable(floatTensor([1.0])), autograd.Variable(floatTensor([0.0])))).cpu()

def weighted_sum(matrix: torch.Tensor, attention: torch.Tensor) -> torch.Tensor:
    if attention.dim() == 2 and matrix.dim() == 3:
        return attention.unsqueeze(1).bmm(matrix).squeeze(1)
    if attention.dim() == 3 and matrix.dim() == 3:
        return attention.bmm(matrix)
    if matrix.dim() - 1 < attention.dim():
        expanded_size = list(matrix.size())
        for i in range(attention.dim() - matrix.dim() + 1):
            matrix = matrix.unsqueeze(1)
            expanded_size.insert(i + 1, attention.size(i + 1))
        matrix = matrix.expand(*expanded_size)
    intermediate = attention.unsqueeze(-1).expand_as(matrix) * matrix
    return intermediate.sum(dim=-2)

def masked_sum(
    vector: torch.Tensor, mask: torch.BoolTensor, dim: int, keepdim: bool = False) -> torch.Tensor:
    replaced_vector = vector.masked_fill(~mask, 0.0)
    value_sum = torch.sum(replaced_vector, dim=dim, keepdim=keepdim)
    return value_sum

class HGNN_conv(nn.Module):
    def __init__(self, in_ft, out_ft, bias=True):
        super(HGNN_conv, self).__init__()

        self.weight = Parameter(torch.Tensor(in_ft, out_ft))
        if bias:
            self.bias = Parameter(torch.Tensor(out_ft))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def reset_parameters_xavier(self):
        nn.init.xavier_uniform_(self.weight)
        if self.bias is not None:
            nn.init.xavier_uniform_(self.bias)

    def forward(self, x: torch.Tensor, G: torch.Tensor):
        x = x.matmul(self.weight)
        if self.bias is not None:
            x = x + self.bias
        x = G.matmul(x)
        return x

class HGNN_fc(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(HGNN_fc, self).__init__()
        self.fc = nn.Linear(in_ch, out_ch)

    def forward(self, x):
        return self.fc(x)

class HGNN_sg_attn_simple(nn.Module):

    def __init__(self, edim, ddim, atype='additive'):
        super(HGNN_sg_attn, self).__init__()
        if atype == 'additive':
            self.attention = attention.AdditiveAttention(edim, ddim)
        elif atype == 'multiplicative':
            self.attention = attention.MultiplicativeAttention(edim, ddim)
        else:
            exit(f'unrecognized attention type {atype}')

    def forward(self, x, sgs):
        xsize = list(x.size())
        bsize = sgs.shape[0]
        b_attn = torch.matmul(sgs, x)
        y = self.attention(b_attn, x)

        return y

class HGNN_sg_attn_Allen(nn.Module):

    def __init__(self, vdim, mdim, atype='additive'):
        super(HGNN_sg_attn, self).__init__()
        if atype == 'additive':
            self.attention = attention.AdditiveAttention(vdim, mdim)
        elif atype == 'dotprod':
            self.attention = attention.DotProductAttention()
        else:
            exit(f'unrecognized attention type {atype}')

    def forward(self, x, sgs):
        xsize = list(x.size())
        bsize = sgs.shape[0]
        b_attn = torch.matmul(sgs, x)
        attn_wts = self.attention(b_attn, x.unsqueeze(0).expand(bsize, *xsize))
        x = torch.matmul(attn_wts, x)
        return x

class HGNN_sg_attn(nn.Module):

    def __init__(self, vdim, mdim, atype='additive'):
        super(HGNN_sg_attn, self).__init__()
        self.attn_vector = torch.nn.Parameter(torch.zeros((vdim,1), dtype=torch.float), requires_grad=True)

        stdv = 1. / math.sqrt(vdim)
        self.attn_vector.data.uniform_(-stdv, stdv)

    def forward(self, x, sgs):
        xsize = list(x.size())
        bsize = sgs.shape[0]
        attn_wts = torch.matmul(x, self.attn_vector)
        attn_wts = attn_wts.squeeze().unsqueeze(0).expand(bsize, xsize[0])
        x = torch.matmul(sgs*attn_wts, x)
        return x

class HGNN_sg_attn_multiplicative(nn.Module):

    def __init__(self, vdim, mdim, atype='additive'):
        super(HGNN_sg_attn, self).__init__()
        self.W = nn.Parameter(torch.FloatTensor(mdim, vdim), requires_grad=True)
        nn.init.xavier_uniform_(self.W)

    def forward(self, x, sgs):
        x = sgs.matmul(x).matmul(self.W)
        return x

class HGNN_embedding(nn.Module):
    def __init__(self, in_ch, n_hid, dropout=0.5):
        super(HGNN_embedding, self).__init__()
        self.dropout = dropout
        self.hgc1 = HGNN_conv(in_ch, n_hid)
        self.hgc2 = HGNN_conv(n_hid, n_hid)

    def forward(self, x, G):
        x = F.relu(self.hgc1(x, G))
        x = F.dropout(x, self.dropout)
        x = F.relu(self.hgc2(x, G))
        return x

class HGNN_classifier(nn.Module):
    def __init__(self, n_hid, n_class):
        super(HGNN_classifier, self).__init__()
        self.fc1 = nn.Linear(n_hid, n_class)

    def forward(self, x):
        x = self.fc1(x)
        return x

class MEGACare(nn.Module):
    def __init__(self, vocab_size, ddi_adj_pn, ddi_adj_all, MPNNSet, N_fingerprints, average_projection, emb_dim=256, device=torch.device('cpu:0')):
        super(MEGACare, self).__init__()

        self.device = device

        self.average_projection = average_projection.to(device=self.device)

        # pre-embedding
        self.embeddings = nn.ModuleList(
            [nn.Embedding(vocab_size[i], emb_dim) for i in range(2)])
        self.dropout = nn.Dropout(p=0.5)
        self.encoders = nn.ModuleList([nn.GRU(emb_dim, emb_dim, batch_first=True) for _ in range(2)])
        self.query = nn.Sequential(
                nn.ReLU(),
                nn.Linear(2 * emb_dim, emb_dim)
        )

        # MPNN global embedding
        MPNN_molecule_Set = list(zip(*MPNNSet))

        self.MPNN_emb = MolecularGraphNeuralNetwork(N_fingerprints, emb_dim, layer_hidden=2, device=device).forward(self.MPNN_molecule_Set)
        self.MPNN_emb = torch.mm(average_projection.to(device=self.device), self.MPNN_emb.to(device=self.device))

        # self.MPNN_emb.to(device=self.device)
        self.MPNN_emb = torch.tensor(self.MPNN_emb, requires_grad=True)

        self.MPNN_net = MolecularGraphNeuralNetwork(MPNN_molecule_Set, self.average_projection, ddi_adj_all, ddi_adj_pn, \
                                                     N_fingerprints, emb_dim, layer_hidden=2, device=self.device)
        self.MPNN_emb = torch.load("../MPNN_net_embeddings").detach().to(device)
        # self.MPNN_net.load_state_dict(stat_dic)
        self.MPNN_output = nn.Linear(vocab_size[2], vocab_size[2])
        self.MPNN_layernorm = nn.LayerNorm(vocab_size[2])

        self.loss_function = marginLoss()

        # graphs, bipartite matrix
        # Create POS and NEG ddi metrix
        self.tensor_ddi_adj_all = torch.FloatTensor(ddi_adj_all).to(device)
        self.tensor_ddi_adj_pn = torch.FloatTensor(ddi_adj_pn).to(device)
        self.ddi_adj_all = ddi_adj_all
        self.ddi_adj_pn = ddi_adj_pn

        # self.tensor_ddi_mask_H = torch.FloatTensor(ddi_mask_H).to(device)
        self.init_weights()

        # self.add_Linear = add_Linear

        # if add_Linear == True:
        self.mlp = nn.Linear(64, 64)
        # self.new_MPNN_emb = self.conv()

    def triplet_loss_v2(self):
        # MPNN_emb triplet loss calculation
        ent_batch = []
        pos_batch = []
        neg_batch = []
        for batch in range(128):
            while True:
                ent = random.randint(0, 265)
                idx = np.argwhere(self.ddi_adj_pn[ent] == 1)
                if len(idx) == 0:
                    continue
                pos = idx[random.randint(0, len(idx) - 1)][0]
                idx = np.argwhere(self.ddi_adj_all[ent] == 1)
                if len(idx) == 0:
                    continue
                neg = idx[random.randint(0, len(idx) - 1)][0]
                break

            ent_batch.append(ent)
            pos_batch.append(pos)
            neg_batch.append(neg)

        ent_batch = autograd.Variable(longTensor(ent_batch))
        pos_batch = autograd.Variable(longTensor(pos_batch))
        neg_batch = autograd.Variable(longTensor(neg_batch))

        ent_emb = self.MPNN_emb[ent_batch, :]
        pos_emb = self.MPNN_emb[pos_batch, :]
        neg_emb = self.MPNN_emb[neg_batch, :]

        pos = torch.sum(torch.abs(ent_emb - pos_emb), 1)
        neg = torch.sum(torch.abs(ent_emb - neg_emb), 1)

        margin = autograd.Variable(floatTensor([1.0]))
        emb_loss = self.loss_function(pos, neg, margin) + normLoss(ent_emb) + normLoss(pos_emb) + normLoss(neg_emb)
        # new_MPNN_emb = self.conv(self.MPNN_emb)
        return emb_loss

    def forward(self, input):

        loss_1, loss_2 = 0, 0
        all_patient_loss_1 = torch.zeros(size=[0, self.vocab_size[2]]).to(self.device)
        all_patient_loss_2 = torch.zeros(size=[0, self.vocab_size[2]], dtype=torch.long).to(self.device)

        self.MPNN_emb = self.MPNN_net().detach()
        self.MPNN_emb = torch.tensor(self.MPNN_emb, requires_grad=True)
        loss_tri = self.triplet_loss()

	    # patient health representation
        i1_seq = []
        i2_seq = []
        def sum_embedding(embedding):
            return embedding.sum(dim=1).unsqueeze(dim=0)  # (1,1,dim)
        for adm in input:
            i1 = sum_embedding(self.dropout(self.embeddings[0](torch.LongTensor(adm[0]).unsqueeze(dim=0).to(self.device)))) # (1,1,dim)
            i2 = sum_embedding(self.dropout(self.embeddings[1](torch.LongTensor(adm[1]).unsqueeze(dim=0).to(self.device))))
            i1_seq.append(i1)
            i2_seq.append(i2)
        i1_seq = torch.cat(i1_seq, dim=1) #(1,seq,dim)
        i2_seq = torch.cat(i2_seq, dim=1) #(1,seq,dim)

        o1, h1 = self.encoders[0](
            i1_seq
        )
        o2, h2 = self.encoders[1](
            i2_seq
        )
        patient_representations = torch.cat([o1, o2], dim=-1).squeeze(dim=0) # (seq, dim*2)
        query = self.query(patient_representations)[-1:, :] # (seq, dim)

	    # MPNN embedding

        # MPNN with a MLP layer
        # new_MPNN_emb = self.mlp(self.MPNN_emb)
        # MPNN_match = F.sigmoid(torch.mm(query, new_MPNN_emb.t()))

        # MPNN without a MLP layer
        MPNN_match = F.sigmoid(torch.mm(query, self.MPNN_emb.t()))
        MPNN_att = self.MPNN_layernorm(MPNN_match + self.MPNN_output(MPNN_match))

        result = MPNN_att

        # result = torch.mul(bipartite_emb, MPNN_att)
        # result = torch.cat(bipartite_emb, MPNN_att)

        neg_pred_prob = F.sigmoid(result)
        neg_pred_prob = neg_pred_prob.t() * neg_pred_prob  # (voc_size, voc_size)

        # tensor

        batch_neg = 0.0005 *  neg_pred_prob.mul(self.tensor_ddi_adj_all).sum()
        batch_pos = 0.0005 *  neg_pred_prob.mul(self.tensor_ddi_adj_pn).sum()

        # LOSS
        loss = (- torch.log(batch_pos / (batch_pos + batch_neg) ))

        return all_patient_loss_1, all_patient_loss_2, prediction, loss_1, loss_2, batch_neg, encoder_visit

    def init_weights(self):
        """Initialize weights."""
        initrange = 0.1
        for item in self.embeddings:
            item.weight.data.uniform_(-initrange, initrange)

class MaskLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(MaskLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input, mask):
        weight = torch.mul(self.weight, mask)
        output = torch.mm(input, weight)

        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -> ' \
               + str(self.out_features) + ')'

# DMNC
class DMNC(nn.Module):
    def __init__(self, vocab_size, emb_dim=64, device=torch.device('cpu:0')):
        super(DMNC, self).__init__()
        K = len(vocab_size)
        self.K = K
        self.vocab_size = vocab_size
        self.device = device

        self.token_start = vocab_size[2]
        self.token_end = vocab_size[2] + 1

        self.embeddings = nn.ModuleList(
            [nn.Embedding(vocab_size[i] if i != 2 else vocab_size[2] + 2, emb_dim) for i in range(K)])
        self.dropout = nn.Dropout(p=0.5)

        self.encoders = nn.ModuleList([DNC(
            input_size=emb_dim,
            hidden_size=emb_dim,
            rnn_type='gru',
            num_layers=1,
            num_hidden_layers=1,
            nr_cells=16,
            cell_size=emb_dim,
            read_heads=1,
            batch_first=True,
            gpu_id=0,
            independent_linears=False
        ) for _ in range(K - 1)])

        self.decoder = nn.GRU(emb_dim + emb_dim * 2, emb_dim * 2,
                              batch_first=True)  # input: (y, r1, r2,) hidden: (hidden1, hidden2)
        self.interface_weighting = nn.Linear(emb_dim * 2, 2 * (emb_dim + 1 + 3))  # 2 read head (key, str, mode)
        self.decoder_r2o = nn.Linear(2 * emb_dim, emb_dim * 2)

        self.output = nn.Linear(emb_dim * 2, vocab_size[2] + 2)

    def forward(self, input, i1_state=None, i2_state=None, h_n=None, max_len=20):
        # input (3, code)
        i1_input_tensor = self.embeddings[0](
            torch.LongTensor(input[0]).unsqueeze(dim=0).to(self.device))  # (1, seq, codes)
        i2_input_tensor = self.embeddings[1](
            torch.LongTensor(input[1]).unsqueeze(dim=0).to(self.device))  # (1, seq, codes)

        o1, (ch1, m1, r1) = \
            self.encoders[0](i1_input_tensor, (None, None, None) if i1_state is None else i1_state)
        o2, (ch2, m2, r2) = \
            self.encoders[1](i2_input_tensor, (None, None, None) if i2_state is None else i2_state)

        # save memory state
        i1_state = (ch1, m1, r1)
        i2_state = (ch2, m2, r2)

        predict_sequence = [self.token_start] + input[2]
        if h_n is None:
            h_n = torch.cat([ch1[0], ch2[0]], dim=-1)

        output_logits = []
        r1 = r1.unsqueeze(dim=0)
        r2 = r2.unsqueeze(dim=0)

        if self.training:
            for item in predict_sequence:
                # teacher force predict drug
                item_tensor = self.embeddings[2](
                    torch.LongTensor([item]).unsqueeze(dim=0).to(self.device))  # (1, seq, codes)

                o3, h_n = self.decoder(torch.cat([item_tensor, r1, r2], dim=-1), h_n)
                read_keys, read_strengths, read_modes = self.decode_read_variable(h_n.squeeze(0))

                # read from i1_mem, i2_mem and i3_mem
                r1, _ = self.read_from_memory(self.encoders[0],
                                              read_keys[:, 0, :].unsqueeze(dim=1),
                                              read_strengths[:, 0].unsqueeze(dim=1),
                                              read_modes[:, 0, :].unsqueeze(dim=1), i1_state[1])

                r2, _ = self.read_from_memory(self.encoders[1],
                                              read_keys[:, 1, :].unsqueeze(dim=1),
                                              read_strengths[:, 1].unsqueeze(dim=1),
                                              read_modes[:, 1, :].unsqueeze(dim=1), i2_state[1])

                output = self.decoder_r2o(torch.cat([r1, r2], dim=-1))
                output = self.output(output + o3).squeeze(dim=0)
                output_logits.append(output)
        else:
            item_tensor = self.embeddings[2](
                torch.LongTensor([self.token_start]).unsqueeze(dim=0).to(self.device))  # (1, seq, codes)
            for idx in range(max_len):
                # predict
                # teacher force predict drug
                o3, h_n = self.decoder(torch.cat([item_tensor, r1, r2], dim=-1), h_n)
                read_keys, read_strengths, read_modes = self.decode_read_variable(h_n.squeeze(0))

                # read from i1_mem, i2_mem and i3_mem
                r1, _ = self.read_from_memory(self.encoders[0],
                                              read_keys[:, 0, :].unsqueeze(dim=1),
                                              read_strengths[:, 0].unsqueeze(dim=1),
                                              read_modes[:, 0, :].unsqueeze(dim=1), i1_state[1])

                r2, _ = self.read_from_memory(self.encoders[1],
                                              read_keys[:, 1, :].unsqueeze(dim=1),
                                              read_strengths[:, 1].unsqueeze(dim=1),
                                              read_modes[:, 1, :].unsqueeze(dim=1), i2_state[1])

                output = self.decoder_r2o(torch.cat([r1, r2], dim=-1))
                output = self.output(output + o3).squeeze(dim=0)
                output = F.softmax(output, dim=-1)
                output_logits.append(output)

                input_token = torch.argmax(output, dim=-1)
                input_token = input_token.item()
                item_tensor = self.embeddings[2](
                    torch.LongTensor([input_token]).unsqueeze(dim=0).to(self.device))  # (1, seq, codes)

        return torch.cat(output_logits, dim=0), i1_state, i2_state, h_n

    def read_from_memory(self, dnc, read_key, read_str, read_mode, m_hidden):
        read_vectors, hidden = dnc.memories[0].read(read_key, read_str, read_mode, m_hidden)
        return read_vectors, hidden

    def decode_read_variable(self, input):
        w = 64
        r = 2
        b = input.size(0)

        input = self.interface_weighting(input)
        # r read keys (b * w * r)
        read_keys = F.tanh(input[:, :r * w].contiguous().view(b, r, w))
        # r read strengths (b * r)
        read_strengths = F.softplus(input[:, r * w:r * w + r].contiguous().view(b, r))
        # read modes (b * 3*r)
        read_modes = F.softmax(input[:, (r * w + r):].contiguous().view(b, r, 3), -1)
        return read_keys, read_strengths, read_modes

# GAMENet
class GCN(nn.Module):
    def __init__(self, voc_size, emb_dim, adj, device=torch.device('cpu:0')):
        super(GCN, self).__init__()
        self.voc_size = voc_size
        self.emb_dim = emb_dim
        self.device = device

        adj = self.normalize(adj + np.eye(adj.shape[0]))

        self.adj = torch.FloatTensor(adj).to(device)
        self.x = torch.eye(voc_size).to(device)

        self.gcn1 = GraphConvolution(voc_size, emb_dim)
        self.dropout = nn.Dropout(p=0.3)
        self.gcn2 = GraphConvolution(emb_dim, emb_dim)

    def forward(self):
        node_embedding = self.gcn1(self.x, self.adj)
        node_embedding = F.relu(node_embedding)
        node_embedding = self.dropout(node_embedding)
        node_embedding = self.gcn2(node_embedding, self.adj)
        return node_embedding

    def normalize(self, mx):
        """Row-normalize sparse matrix"""
        rowsum = np.array(mx.sum(1))
        r_inv = np.power(rowsum, -1).flatten()
        r_inv[np.isinf(r_inv)] = 0.
        r_mat_inv = np.diagflat(r_inv)
        mx = r_mat_inv.dot(mx)
        return mx

class GAMENet(nn.Module):
    def __init__(self, vocab_size, ehr_adj, ddi_adj, emb_dim=64, device=torch.device('cpu:0'), ddi_in_memory=True):
        super(GAMENet, self).__init__()
        K = len(vocab_size)
        self.K = K
        self.vocab_size = vocab_size
        self.device = device
        self.tensor_ddi_adj = torch.FloatTensor(ddi_adj).to(device)
        self.ddi_in_memory = ddi_in_memory
        self.embeddings = nn.ModuleList(
            [nn.Embedding(vocab_size[i], emb_dim) for i in range(K-1)])
        self.dropout = nn.Dropout(p=0.5)

        self.encoders = nn.ModuleList([nn.GRU(emb_dim, emb_dim * 2, batch_first=True) for _ in range(K-1)])

        self.query = nn.Sequential(
            nn.ReLU(),
            nn.Linear(emb_dim * 4, emb_dim),
        )

        self.ehr_gcn = GCN(voc_size=vocab_size[2], emb_dim=emb_dim, adj=ehr_adj, device=device)
        self.ddi_gcn = GCN(voc_size=vocab_size[2], emb_dim=emb_dim, adj=ddi_adj, device=device)
        self.inter = nn.Parameter(torch.FloatTensor(1))

        self.output = nn.Sequential(
            nn.ReLU(),
            nn.Linear(emb_dim * 3, emb_dim * 2),
            nn.ReLU(),
            nn.Linear(emb_dim * 2, vocab_size[2])
        )

        self.init_weights()

    def forward(self, input):
        # input (adm, 3, codes)

        # generate medical embeddings and queries
        i1_seq = []
        i2_seq = []
        def mean_embedding(embedding):
            return embedding.mean(dim=1).unsqueeze(dim=0)  # (1,1,dim)
        for adm in input:
            i1 = mean_embedding(self.dropout(self.embeddings[0](torch.LongTensor(adm[0]).unsqueeze(dim=0).to(self.device)))) # (1,1,dim)
            i2 = mean_embedding(self.dropout(self.embeddings[1](torch.LongTensor(adm[1]).unsqueeze(dim=0).to(self.device))))
            i1_seq.append(i1)
            i2_seq.append(i2)
        i1_seq = torch.cat(i1_seq, dim=1) #(1,seq,dim)
        i2_seq = torch.cat(i2_seq, dim=1) #(1,seq,dim)

        o1, h1 = self.encoders[0](
            i1_seq
        ) # o1:(1, seq, dim*2) hi:(1,1,dim*2)
        o2, h2 = self.encoders[1](
            i2_seq
        )
        patient_representations = torch.cat([o1, o2], dim=-1).squeeze(dim=0) # (seq, dim*4)
        queries = self.query(patient_representations) # (seq, dim)

        # graph memory module
        '''I:generate current input'''
        query = queries[-1:] # (1,dim)

        '''G:generate graph memory bank and insert history information'''
        if self.ddi_in_memory:
            drug_memory = self.ehr_gcn() - self.ddi_gcn() * self.inter  # (size, dim)
        else:
            drug_memory = self.ehr_gcn()

        if len(input) > 1:
            history_keys = queries[:(queries.size(0)-1)] # (seq-1, dim)

            history_values = np.zeros((len(input)-1, self.vocab_size[2]))
            for idx, adm in enumerate(input):
                if idx == len(input)-1:
                    break
                history_values[idx, adm[2]] = 1
            history_values = torch.FloatTensor(history_values).to(self.device) # (seq-1, size)

        '''O:read from global memory bank and dynamic memory bank'''
        key_weights1 = F.softmax(torch.mm(query, drug_memory.t()), dim=-1)  # (1, size)
        fact1 = torch.mm(key_weights1, drug_memory)  # (1, dim)

        if len(input) > 1:
            visit_weight = F.softmax(torch.mm(query, history_keys.t())) # (1, seq-1)
            weighted_values = visit_weight.mm(history_values) # (1, size)
            fact2 = torch.mm(weighted_values, drug_memory) # (1, dim)
        else:
            fact2 = fact1
        '''R:convert O and predict'''
        output = self.output(torch.cat([query, fact1, fact2], dim=-1)) # (1, dim)

        if self.training:
            neg_pred_prob = F.sigmoid(output)
            neg_pred_prob = neg_pred_prob.t() * neg_pred_prob  # (voc_size, voc_size)
            batch_neg = neg_pred_prob.mul(self.tensor_ddi_adj).mean()

            return output, batch_neg
        else:
            return output

    def init_weights(self):
        """Initialize weights."""
        initrange = 0.1
        for item in self.embeddings:
            item.weight.data.uniform_(-initrange, initrange)

        self.inter.data.uniform_(-initrange, initrange)

# LEAP
class Leap(nn.Module):
    def __init__(self, voc_size, emb_dim=64, device=torch.device('cpu:0')):
        super(Leap, self).__init__()
        self.voc_size = voc_size
        self.device = device
        self.SOS_TOKEN = voc_size[2]
        self.END_TOKEN = voc_size[2]+1

        self.enc_embedding = nn.Sequential(
            nn.Embedding(voc_size[0], emb_dim, ),
            nn.Dropout(0.3)
        )
        self.dec_embedding = nn.Sequential(
            nn.Embedding(voc_size[2] + 2, emb_dim, ),
            nn.Dropout(0.3)
        )

        self.dec_gru = nn.GRU(emb_dim*2, emb_dim, batch_first=True)

        self.attn = nn.Linear(emb_dim*2, 1)

        self.output = nn.Linear(emb_dim, voc_size[2]+2)

    def forward(self, input, max_len=20):
        device = self.device
        # input (3, codes)
        input_tensor = torch.LongTensor(input[0]).to(device)
        # (len, dim)
        input_embedding = self.enc_embedding(input_tensor.unsqueeze(dim=0)).squeeze(dim=0)

        output_logits = []
        hidden_state = None
        if self.training:
            for med_code in [self.SOS_TOKEN] + input[2]:
                dec_input = torch.LongTensor([med_code]).unsqueeze(dim=0).to(device)
                dec_input = self.dec_embedding(dec_input).squeeze(dim=0) # (1,dim)

                if hidden_state is None:
                    hidden_state = dec_input
                hidden_state_repeat = hidden_state.repeat(input_embedding.size(0), 1) # (len, dim)

                combined_input = torch.cat([hidden_state_repeat, input_embedding], dim=-1) # (len, dim*2)
                attn_weight = F.softmax(self.attn(combined_input).t(), dim=-1) # (1, len)
                input_embedding = attn_weight.mm(input_embedding) # (1, dim)

                _, hidden_state = self.dec_gru(torch.cat([input_embedding, dec_input], dim=-1).unsqueeze(dim=0), hidden_state.unsqueeze(dim=0))
                hidden_state = hidden_state.squeeze(dim=0) # (1,dim)

                output_logits.append(self.output(F.relu(hidden_state)))
            return torch.cat(output_logits, dim=0)

        else:
            for di in range(max_len):
                if di == 0:
                    dec_input = torch.LongTensor([[self.SOS_TOKEN]]).to(device)
                dec_input = self.dec_embedding(dec_input).squeeze(dim=0) # (1,dim)
                if hidden_state is None:
                    hidden_state = dec_input
                hidden_state_repeat = hidden_state.repeat(input_embedding.size(0), 1)  # (len, dim)
                combined_input = torch.cat([hidden_state_repeat, input_embedding], dim=-1)  # (len, dim*2)
                attn_weight = F.softmax(self.attn(combined_input).t(), dim=-1)  # (1, len)
                input_embedding = attn_weight.mm(input_embedding)  # (1, dim)
                _, hidden_state = self.dec_gru(torch.cat([input_embedding, dec_input], dim=-1).unsqueeze(dim=0),
                                               hidden_state.unsqueeze(dim=0))
                hidden_state = hidden_state.squeeze(dim=0)  # (1,dim)
                output = self.output(F.relu(hidden_state))
                topv, topi = output.data.topk(1)
                output_logits.append(F.softmax(output, dim=-1))
                dec_input = topi.detach()
            return torch.cat(output_logits, dim=0)

# RETAIN
class Retain(nn.Module):
    def __init__(self, voc_size, emb_size=64, device=torch.device('cpu:0')):
        super(Retain, self).__init__()
        self.device = device
        self.voc_size = voc_size
        self.emb_size = emb_size
        self.input_len = voc_size[0] + voc_size[1] + voc_size[2]
        self.output_len = voc_size[2]

        self.embedding = nn.Sequential(
            nn.Embedding(self.input_len + 1, self.emb_size, padding_idx=self.input_len),
            nn.Dropout(0.5)
        )

        self.alpha_gru = nn.GRU(emb_size, emb_size, batch_first=True)
        self.beta_gru = nn.GRU(emb_size, emb_size, batch_first=True)

        self.alpha_li = nn.Linear(emb_size, 1)
        self.beta_li = nn.Linear(emb_size, emb_size)

        self.output = nn.Linear(emb_size, self.output_len)

    def forward(self, input):
        device = self.device
        # input: (visit, 3, codes )
        max_len = max([(len(v[0]) + len(v[1]) + len(v[2])) for v in input])
        input_np = []
        for visit in input:
            input_tmp = []
            input_tmp.extend(visit[0])
            input_tmp.extend(list(np.array(visit[1]) + self.voc_size[0]))
            input_tmp.extend(list(np.array(visit[2]) + self.voc_size[0] + self.voc_size[1]))
            if len(input_tmp) < max_len:
                input_tmp.extend( [self.input_len]*(max_len - len(input_tmp)) )

            input_np.append(input_tmp)

        visit_emb = self.embedding(torch.LongTensor(input_np).to(device)) # (visit, max_len, emb)
        visit_emb = torch.sum(visit_emb, dim=1) # (visit, emb)

        g, _ = self.alpha_gru(visit_emb.unsqueeze(dim=0)) # g: (1, visit, emb)
        h, _ = self.beta_gru(visit_emb.unsqueeze(dim=0)) # h: (1, visit, emb)

        g = g.squeeze(dim=0) # (visit, emb)
        h = h.squeeze(dim=0) # (visit, emb)
        attn_g = F.softmax(self.alpha_li(g), dim=-1) # (visit, 1)
        attn_h = F.tanh(self.beta_li(h)) # (visit, emb)

        c = attn_g * attn_h * visit_emb # (visit, emb)
        c = torch.sum(c, dim=0).unsqueeze(dim=0) # (1, emb)

        return self.output(c)


import datetime
import pytz
import pandas
import robin_stocks as robin

def movers(source='robinhood', top=20):

    if source == 'robinhood':
        return movers_robinhood_top20()
    else:
        raise ValueError('Provided source is not supported')

def movers_robinhood_top20():
    r = robin.markets.get_top_movers()

    for sec_data in r:
        sec_data['ask_price'] = float(sec_data['ask_price'])
        sec_data['ask_size'] = int(sec_data['ask_size'])
        sec_data['bid_price'] = float(sec_data['bid_price'])
        sec_data['bid_size'] = int(sec_data['bid_size'])
        sec_data['last_trade_price'] = float(sec_data['last_trade_price'])
        sec_data['previous_close'] = float(sec_data['previous_close'])

        sec_data['adjusted_previous_close'] = (
            float(sec_data['adjusted_previous_close'])
        )

        # This value doesn't exist during trading hours, returns None
        # You'll need to use 'last_trade_price'
        if sec_data['last_extended_hours_trade_price'] is not None:
            sec_data['last_extended_hours_trade_price'] = (
                float(sec_data['last_extended_hours_trade_price'])
            )

            sec_data['growth'] = (sec_data['last_extended_hours_trade_price'] /
                                  sec_data['adjusted_previous_close'] - 1)
        else:
            sec_data['growth'] = (sec_data['last_trade_price'] /
                                  sec_data['adjusted_previous_close'] - 1)

        sec_data['previous_close_date'] = (
            datetime.date.fromisoformat(sec_data['previous_close_date'])
        )

        sec_data['trading_halted'] = bool(sec_data['trading_halted'])
        sec_data['has_traded'] = bool(sec_data['has_traded'])

        sec_data['updated_at'] = sec_data['updated_at'].replace('Z', '')
        sec_data['updated_at'] = (
            datetime.datetime.fromisoformat(sec_data['updated_at'])
        )

        sec_data['updated_at'] = pytz.UTC.localize(sec_data['updated_at'],
                                                   is_dst=False)

    result = pandas.DataFrame.from_records(r)
    result = result.set_index('symbol').drop('instrument', axis=1)

    return result


#!/usr/bin/env python2

"""
Andriy Rudyk
Tim Sizemore

Camera operation module

"""

import picamera
#from time import sleep
import time
import os
import io
import cv2

PICTURE_NAME = 'capture'
WIDTH = 1024
HEIGHT = 768

def take_picture(filename, width, height):
    """
    Takes a picture width x height and saves us filename.
    """
    with picamera.PiCamera() as cam:
        picamera.resolution = (width, height)
        # According to API camera needs 0.2 seconds warm up at start.
        time.sleep(0.2)
        cam.capture(filename)

def ocv_object():
    """
    Takes a picture then converts that data to an array that Open CV can use.
    """
    stream = io.BytesIO()
    with picamera.PiCamera() as cam:
    	cam.start_preview()
    time.sleep(2)
    camera.capture(stream, format='jpeg')

    data = np.fromstring(stream.getvalue(), dtype=np.uint8)
    image = cv2.imdecode(data, 1)
    image = image[:, :, ::-1]

def remove_picture(filename):
    """
    Removes filename
    """
    os.remove(filename)

def timelapse(filename):
    """
    Takes a timelapse sequence of pictures.
    """
    with picamera.PiCamera() as cam:
        cam.start_preview()
        time.sleep(2)
        for filename in cam.capture_continuous('img{counter:03d}.jpg'):
            print('Captured %s' % filename)
            time.sleep(60)

def video(height, width):
    """
    Captures a video of specified height and width
    """
    with picamera.PiCamera() as cam:
        cam.resolution = (height, width)
        cam.start_recording('video.h264')
        cam.wait_recording(60)
        cam.stop_recording()

def detect(path):
    img = cv2.imread(path)
    cascade = cv2.CascadeClassifier("haarcascade_frontalface_alt.xml")
    rects = cascade.detectMultiScale(img, 1.3, 4, cv2.cv.CV_HAAR_SCALE_IMAGE, (20, 20))

    if len(rects) == 0:
        return [], img
    rects[:, 2:] += rects[:, :2]
    return rects, img

def box(rects, img):
    for x1, y1, x2, y2 in rects:
        cv2.rectangle(img, (x1, y1), (x2, y2), (127, 255, 0), 2)
    cv2.imwrite(PICTURE_NAME + '_DETECTED.jpg', img)

def main():
    take_picture(PICTURE_NAME + '.jpg', WIDTH, HEIGHT)
    rects, img = detect(PICTURE_NAME + '.jpg')
    box(rects, img)
    os.system('display capture_DETECTED.jpg')

if __name__ == '__main__':
    main()


import numpy as np
import pickle
import sys
import os
import itertools

import matplotlib
import pylab
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

from train_doc2vec import load_data

import pandas as pd
import seaborn as sns

from sklearn.preprocessing import OneHotEncoder, LabelEncoder
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score

def PDF_SAVECHOP(f, PNG=False):

    if PNG == False:
        pylab.savefig(f + '.pdf', pad_inches=0, transparent=False)
        os.system('pdfcrop %s.pdf' % f)
        os.system('mv %s-crop.pdf %s.pdf' % (f, f))
    if PNG:
        pylab.savefig(f + '.jpg', pad_inches=0)
        os.system('convert -trim %s.jpg %s.jpg' % (f, f))

def fig_window(_id=0, scale_x=1, scale_y=1):
    plt.rc(
        'font', **{
            'family': 'serif',
            'serif': ['Computer Modern Roman'],
            'monospace': ['Computer Modern Typewriter']
        })
    fig_width_pt = 252  # Get this from LaTeX using \showthe\columnwidth
    inches_per_pt = 1.0 / 72.27  # Convert pt to inch
    golden_mean = (np.sqrt(5) - 1.0) / 2.0  # Aesthetic ratio
    fig_width = 7  # fig_width_pt*inches_per_pt  # width in inches
    fig_height = fig_width * golden_mean  # height in inches
    fig_size = [fig_width, fig_height]
    params = {
        'backend': 'ps',
        'axes.labelsize': 12,
        'font.size': 12,
        'legend.fontsize': 10,
        'legend.handlelength': 1,
        'xtick.labelsize': 12,
        'ytick.labelsize': 12,
        # 'axes.linewidth': 0.1,
        'text.usetex': True,
        'text.latex.preamble':
        [r"\usepackage{amstext}", r"\usepackage{mathpazo}"],
        # 'xtick.major.pad': 10,
        # 'ytick.major.pad': 10
    }
    fig = plt.figure(
        _id,
        figsize=(scale_x * fig_width, scale_y * fig_height),
        facecolor='w',
        edgecolor='k')
    pylab.rcParams.update(params)
    return fig

def class_imbalance(df, label_index):

    label = ['publisher', 'fact', 'bias'][label_index]

    fig = fig_window(0, 1, 1)

    le = LabelEncoder()
    df[label] = le.fit_transform(df[label])

    if label_index == 0:
        ax = sns.countplot(
            data=df, x=label, order=df[label].value_counts().index, log=True)
    else:
        ax = sns.countplot(
            data=df, x=label, order=df[label].value_counts().index)
    ax.set_title("Articles per " + label + " class")

    if label_index == 0:
        ticks = range(0, max(df[label]))
        xticks = []
        for x in ticks:
            if (x % 100) == 0:
                xticks.append(x)
            else:
                xticks.append(" ")
        ax.set_xticklabels(xticks)

    if label_index == 1:
        ax.set_xticklabels(
            le.inverse_transform(range(3)), rotation=15, ha="right")
    if label_index == 2:
        ax.set_xticklabels(
            le.inverse_transform(range(7)), rotation=15, ha="right")

    PDF_SAVECHOP("results/class_histogram_{}".format(label))

if __name__ == '__main__':

    WEIGHTS_FPATH = "models/classifier/model.h5.pkl"

    with open(WEIGHTS_FPATH, "rb") as fid:
        classifier = pickle.load(fid)

    with open("data/plot_data.pkl", "rb") as fid:
        history, N_classes, X_train, y_input, X_test, y_input_test = pickle.load(
            fid)

    with open("data/label_encoders.pkl", "rb") as fid:
        encoders = pickle.load(fid)

    # with open("data/plot_data.pkl", "rb") as fid:
    # 	classifier, history, N_classes, X_train, y_input, X_test, y_input_test = pickle.load(fid)
    #
    # with open("data/label_encoders.pkl", "rb") as fid:
    # 	encoders = pickle.load(fid)

    # Load data into pandas dataframe
    doc2vec_training_csv = 'data/Doc2vecTrainingData.csv'
    doc2vec_preprocessed = 'data/Doc2vecTrainingDataProcessed.pkl'
    df = load_data(
        doc2vec_training_csv=doc2vec_training_csv,
        doc2vec_preprocessed=doc2vec_preprocessed)
    # if os.path.isfile("data/doc2vec_train.pkl"):
    # 	with open("data/doc2vec_train.pkl", "rb") as fid:
    # 		df = pickle.load(fid)
    # else:
    # 	doc2vec_train = 'data/doc2vec_train.csv'
    # 	df, company_to_token, token_to_company = load_data(doc2vec_train = doc2vec_train)
    # 	with open("data/doc2vec_train.pkl", "wb") as fid:
    # 		pickle.dump(df,fid)
    df.rename(
        columns={
            'fact_score': 'fact',
            'bias_score': 'bias'
        }, inplace=True)

    for label_index in [0, 1, 2]:
        class_imbalance(df, label_index)


import json
import os
import pandas as pd
import random
import scipy.stats as stats
from statistics import mean, stdev

def get_questionnaire(questionnaire_name):
    try:
        with open('questionnaires.json') as dataset:
            data = json.load(dataset)
    except FileNotFoundError:
        raise FileNotFoundError("'questionnaires.json' file not exist.")

    # Matching by questionnaire_name in dataset
    questionnaire = None
    for item in data:
        if item["name"] == questionnaire_name:
            questionnaire = item

    if questionnaire is None:
        raise ValueError("Questionnaire not found.")

    return questionnaire

def get_scenarios():
    try:
        with open('situations.json') as dataset:
            scenarios = json.load(dataset)
    except FileNotFoundError:
        raise FileNotFoundError("'situations.json' file not exist.")

    return scenarios

def generate_scenarios(scenario_file, select_time, target_emotions):
    scenarios = get_scenarios()
    headers_list = []
    scenarios_list = []

    for emotion in scenarios["emotions"]:
        emotion_name = emotion["name"]

        if ('ALL' not in target_emotions) and (emotion_name not in target_emotions):
            continue

        for index, factor in enumerate(emotion["factors"]):
            headers_list.append(f'{emotion_name}-{index}')
            if select_time > len(factor["scenarios"]):
                selected_scenarios = factor["scenarios"]
            else:
                selected_scenarios = random.sample(factor["scenarios"], select_time)
            scenarios_list.append([factor["name"]] + selected_scenarios)

    # Convert each list to a DataFrame
    df_list = [pd.DataFrame(data, columns=[header]) for header, data in zip(headers_list, scenarios_list)]
    output_df = pd.concat(df_list, axis=1)
    output_df.to_csv(scenario_file, index=False)

    return output_df

def generate_testfile(questionnaire, args):
    scenarios_csv = args.scenarios_file
    output_csv = args.testing_file

    test_times = args.test_count
    default_shuffle_times = args.default_shuffle_count
    emotion_shuffle_times = args.emotion_shuffle_count
    shuffle_times = max(default_shuffle_times, emotion_shuffle_times)

    # Extract the scenarios file
    emotions_selection = args.emotion
    select_times = args.select_count
    if emotions_selection != 'Customize':
        scenarios_df = generate_scenarios(scenarios_csv, select_times, emotions_selection.split(','))
    else:
        try:
            scenarios_df = pd.read_csv(scenarios_csv)
        except:
            raise FileNotFoundError("'situations.csv' file not exist.")

    output_df = pd.DataFrame()
    questions_list = questionnaire["questions"]
    question_num = len(list(questions_list.keys()))
    headers_list = []
    data_list = []

    for shuffle_count in range(shuffle_times + 1):
        question_indices = list(questions_list.keys())

        # Shuffle the question indices
        if shuffle_count != 0:
            random.shuffle(question_indices)
        questions = [f'{index}. {questions_list[question]}' for index, question in enumerate(question_indices, 1)]

        # Append questions column
        headers_list.append(f"question-{shuffle_count}")
        data_list.append([f'Prompt: {questionnaire["prompt"]}'] + questions)

        # Append order column
        headers_list.append(f"order-{shuffle_count}")
        data_list.append([questionnaire["prompt"]] + question_indices)

    output_df.to_csv(output_csv, index=False)

    # For each shuffled order,
    for shuffle_count in range(default_shuffle_times + 1):
        for test_count in range(test_times):
            headers_list.append(f'General_test-{test_count}_order-{shuffle_count}')
            data_list.append([' '] + [''] * question_num)

    # For each scenario, create a new column in output_df for each question order
    headers = scenarios_df.columns.tolist()
    for factor in headers:
        for s_index, scenario in enumerate(scenarios_df[factor].iloc[1:].dropna().astype(str)):
            for shuffle_count in range(emotion_shuffle_times + 1):
                for test_count in range(test_times):
                    # Append scenario column
                    headers_list.append(f'{factor}_scenario-{s_index}_test-{test_count}_order-{shuffle_count}')
                    data_list.append([f'Imagine you are the protagonist in the scenario: "{scenario}"'] + [''] * question_num)

    # Create a DataFrame from the data_list and new_header_list, and save it to a CSV file
    data_dict = {header: data for header, data in zip(headers_list, data_list)}
    output_df = pd.DataFrame(data_dict)
    output_df.to_csv(output_csv, index=False)

def convert_data(questionnaire, scenarios_csv, testing_csv):
    try:
        # Read scenarios_csv to extract all headers
        scenarios_df = pd.read_csv(scenarios_csv)
        scenarios_headers = scenarios_df.columns.tolist()

        # Extract unique emotions from the column headers
        emotions_list = []
        for h in scenarios_headers:
            if h.split("_")[0].split("-")[0] not in emotions_list:
                emotions_list.append(h.split("_")[0].split("-")[0])

        # Create a dictionary of factors with the headers as keys and names as values (first row of the DataFrame)
        factors_list = {h : scenarios_df[h].iloc[0] for h in scenarios_headers}

    except Exception as e:
        raise ValueError(f'{scenarios_csv} has problems: {str(e)}')

    try:
        # Read testing_csv file to extract all headers
        testing_df = pd.read_csv(testing_csv)
        testing_headers = testing_df.columns.tolist()

        # Extract all questions orders
        orders_list = {h: testing_df[h].iloc[1:].astype(int).tolist() for h in testing_headers if h.startswith("order")}

    except Exception as e:
        raise ValueError(f'{testing_csv} has problems: {str(e)}')

    # Store the tested data into a dictionary with proper mapping
    tested_data = {}
    for h in testing_headers:
        if not h.startswith(("question", "order")):
            try:
                tested_data[h] = {
                    # Map the data to its corresponding order
                    i : questionnaire["scale"] - int(val) if i in questionnaire["reverse"] else int(val)
                    for i,val in zip(orders_list[h.split("_")[-1]], testing_df[h].iloc[1:].tolist())
                }
            except:
                raise ValueError(f'Error in {testing_csv}: Some cells in column "{h}" cannot be converted to integers.')

    # Create a dictionary to store the mapped data according to the corresponding emotion and factor
    organized_data = {
        emotion : {
            factor : {
                "factor_name" : factors_list[factor],
                "data" : []
            } for factor in factors_list if factor.startswith(emotion)
        } for emotion in emotions_list
    }
    organized_data['General'] = []

    # Organize the data into the mapped_data dictionary based on emotions and factors
    for key in tested_data:
        emotion = key.split("_")[0].split("-")[0]
        if emotion == 'General':
            # If the emotion is 'General', add the data directly to the 'General' category
            organized_data['General'].append(tested_data[key])
        else:
            # Add other data to the corresponding emotion and factor category based on their keys
            factor = key.split("_")[0]
            organized_data[emotion][factor]["data"].append(tested_data[key])

    return organized_data

def compute_statistics(questionnaire, data_list):
    cat_list = []
    results = []

    for cat in questionnaire["categories"]:
        scores_list = []

        for data in data_list:
            scores = []
            for key in data:
                if key in cat["cat_questions"]:
                    scores.append(data[key])

            # Getting the computation mode (SUM or AVG)
            if questionnaire["compute_mode"] == "SUM":
                scores_list.append(sum(scores))
            elif questionnaire["compute_mode"] == "SUM*2":
                scores_list.append(sum(scores)*2)
            else:
                scores_list.append(mean(scores))

        results.append([mean(scores_list), stdev(scores_list), len(scores_list)])
        cat_list.append(cat["cat_name"])

    return results, cat_list    # ([mean, std, size], cat_list)

def hypothesis_testing(result1, result2, cat_list, significance_level, title):
    output_list = f'| {title} |'
    output_text = f'### {title}\n'

    for i, cat_name in enumerate(cat_list):
        output_text += f'\n##### {cat_name}'

        # Extract the mean, std and size for both data sets
        mean1, std1, n1 = result1[i]
        mean2, std2, n2 = result2[i]
        # Add an epsilon to prevent the zero standard deviarion
        epsilon = 1e-8
        std1 += epsilon
        std2 += epsilon

        output_text += '\n- **Statistic**:\n'
        output_text += f'Corresponding Factor:\tmean1 = {mean1:.1f},\tstd1 = {std1:.1f},\tn1 = {n1}\n'
        output_text += f'Default:\tmean2 = {mean2:.1f},\tstd2 = {std2:.1f},\tn2 = {n2}\n'

        # Perform F-test
        output_text += '\n- **F-Test:**\n\n'

        if std1 > std2:
            f_value = std1 ** 2 / std2 ** 2
            df1, df2 = n1 - 1, n2 - 1
        else:
            f_value = std2 ** 2 / std1 ** 2
            df1, df2 = n2 - 1, n1 - 1

        p_value = (1 - stats.f.cdf(f_value, df1, df2)) * 2
        equal_var = True if p_value > significance_level else False

        output_text += f'\tf-value = {f_value:.4f}\t($df_1$ = {df1}, $df_2$ = {df2})\n\n'
        output_text += f'\tp-value = {p_value:.4f}\t(two-tailed test)\n\n'
        output_text += '\tNull hypothesis $H_0$ ($s_1^2$ = $s_2^2$): '

        if p_value > significance_level:
            output_text += f'\tSince p-value ({p_value:.4f}) > α ({significance_level}), $H_0$ cannot be rejected.\n\n'
            output_text += f'\t**Conclusion ($s_1^2$ = $s_2^2$):** The variance of LLM\'s average responses in this factor is statistically equal to the variance of general.\n\n'
        else:
            output_text += f'\tSince p-value ({p_value:.4f}) < α ({significance_level}), $H_0$ is rejected.\n\n'
            output_text += f'\t**Conclusion ($s_1^2$ ≠ $s_2^2$):** The variance of LLM\'s average responses in this factor is statistically unequal to the variance of general.\n\n'

        # Performing T-test
        output_text += '- **Two Sample T-Test (Equal Variance):**\n\n' if equal_var else '- **Two Sample T-test (Welch\'s T-Test):**\n\n'

        df = n1 + n2 - 2 if equal_var else ((std1**2 / n1 + std2**2 / n2)**2) / ((std1**2 / n1)**2 / (n1 - 1) + (std2**2 / n2)**2 / (n2 - 1))
        t_value, p_value = stats.ttest_ind_from_stats(mean1, std1, n1, mean2, std2, n2, equal_var=equal_var)

        output_text += f'\tt-value = {t_value:.4f}\t($df$ = {df:.1f})\n\n'
        output_text += f'\tp-value = {p_value:.4f}\t(two-tailed test)\n\n'

        output_text += '\tNull hypothesis $H_0$ ($µ_1$ = $µ_2$): '
        if p_value > significance_level:
            output_text += f'\tSince p-value ({p_value:.4f}) > α ({significance_level}), $H_0$ cannot be rejected.\n\n'
            output_text += f'\t**Conclusion ($µ_1$ = $µ_2$):** The average of LLM\'s responses in this factor is assumed to be equal to the average of general.\n\n'

            # output_list += f' $-$ ({'+' if (mean1 - mean2) > 0 else ''}{(mean1 - mean2):.1f}) |'
            # output_list += f' $-$ ({'+' if (mean1 - mean2) > 0 else ''}{(mean1 - mean2):.1f}) |'
            output_list += f''' $-$ ({'+' if (mean1 - mean2) > 0 else ''}{(mean1 - mean2):.1f}) |'''

        else:
            output_text += f'Since p-value ({p_value:.4f}) < α ({significance_level}), $H_0$ is rejected.\n\n'
            if t_value > 0:
                output_text += '\tAlternative hypothesis $H_1$ ($µ_1$ > $µ_2$): '
                output_text += f'\tSince p-value ({(1-p_value/2):.1f}) > α ({significance_level}), $H_1$ cannot be rejected.\n\n'
                output_text += f'\t**Conclusion ($µ_1$ > $µ_2$):** The average of LLM\'s responses in this factor is assumed to be larger than the average of general.\n\n'

                output_list += f' $\\uparrow$ (+{(mean1 - mean2):.1f}) |'
            else:
                output_text += '\tAlternative hypothesis $H_1$ ($µ_1$ < $µ_2$): '
                output_text += f'\tSince p-value ({(1-p_value/2):.1f}) > α ({significance_level}), $H_1$ cannot be rejected.\n\n'
                output_text += f'\t**Conclusion ($µ_1$ < $µ_2$):** The average of LLM\'s responses in this factor is assumed to be smaller than the average of general emotion.\n\n'

                output_list += f' $\\downarrow$ ({(mean1 - mean2):.1f}) |'

    output_list += f' {result1[0][2]} |\n'
    return (output_text, output_list)

def analysis_results(questionnaire, args):
    scenarios_csv = args.scenarios_file
    testing_csv = args.testing_file
    result_file = args.results_file
    significance_level = args.significance_level

    overall_list = '# PANAS Results Analysis\n'
    markdown_output = ''    # overall markdown output text
    overall_data = []

    data = convert_data(questionnaire, scenarios_csv, testing_csv)

    general_results, cat_list = compute_statistics(questionnaire, data['General'])

    overall_list += '| Emotions | ' + ' | '.join(cat_list) + ' | N |\n'
    overall_list += '| :---: |' + ' | '.join([':---:' for i in cat_list]) + ' | :---: |\n'
    overall_list += '| Default |' + ' | '.join([f'{r[0]:.1f} $\\pm$ {r[1]:.1f}' for r in general_results]) + f' | {general_results[0][2]} |\n'

    # Analyze the results for each emotion
    for emotion in data:
        if emotion == 'General':
            continue

        emotion_output = ''     # markdown output text for the current emotion

        emotion_list = f'## {emotion}\n'
        emotion_list += '| Factors | ' + ' | '.join(cat_list) + ' | N |\n'
        emotion_list += '| :---: |' + ' | '.join([':---:' for i in cat_list]) + ' | :---: |\n'
        emotion_list += '| Default |' + ' | '.join([f'{r[0]:.1f} $\\pm$ {r[1]:.1f}' for r in general_results]) + f' | {general_results[0][2]} |\n'

        emotion_data = []   # the data that belongs to the current emotion

        # Analyze the results for each factor
        for factor in data[emotion]:
            emotion_data += data[emotion][factor]["data"]
            overall_data += data[emotion][factor]["data"]
            results, _ = compute_statistics(questionnaire, data[emotion][factor]["data"])
            text_msg, list_msg = hypothesis_testing(results, general_results, cat_list, significance_level, data[emotion][factor]["factor_name"])
            emotion_output += text_msg
            emotion_list += list_msg

        results, _ = compute_statistics(questionnaire, emotion_data)
        text_msg, list_msg = hypothesis_testing(results, general_results, cat_list, significance_level, 'Overall')
        emotion_output += text_msg
        emotion_list += list_msg
        overall_list += list_msg.replace('Overall', emotion)
        markdown_output += emotion_list + '\n'+ emotion_output

    markdown_output += f'## Emotions Overall\n'
    results, _ = compute_statistics(questionnaire, overall_data)
    text_msg, list_msg = hypothesis_testing(results, general_results, cat_list, significance_level, 'Overall')
    markdown_output += text_msg
    overall_list += list_msg

    with open(f'{result_file}.md', "w") as f:
        f.write(overall_list + '\n\n' + markdown_output)

def run_emotionbench(args, generator):
    # Get questionnaire
    questionnaire = get_questionnaire(args.questionnaire)
    questionnaire_name = questionnaire["name"]
    args.scenarios_file = f'situations/{args.name_exp}-{questionnaire_name}-situations.csv' if args.name_exp is not None else f'situations/{args.model}-{questionnaire_name}-situations.csv'
    args.testing_file = f'results/{args.name_exp}-{questionnaire_name}-testing.csv' if args.name_exp is not None else f'results/{args.model}-{questionnaire_name}-testing.csv'
    args.results_file = f'results/{args.name_exp}-{questionnaire_name}-results' if args.name_exp is not None else f'results/{args.model}-{questionnaire_name}-results'

    os.makedirs("results", exist_ok=True)
    os.makedirs("situations", exist_ok=True)

    # Generation
    if args.mode in ['generation', 'auto']:
        generate_testfile(questionnaire, args)

    # Testing
    if args.mode in ['testing', 'auto']:
        generator(questionnaire, args)

    # Analysis
    if args.mode in ['analysis', 'auto']:
        analysis_results(questionnaire, args)


import paho.mqtt.client as mqtt
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import threading
import time

#this is to check wifi reconncetion
#if machine is not connceted to any wifi, it will show timeout after 3 second
import socket
socket.setdefaulttimeout(3)

temp="0"
y=0
fig=plt.figure()
plt.ylim(0,10)

#standard mqtt setup definition
def setup_mqtt():
	client=mqtt.Client("movement_level2")
	client.on_connect=OnConnect
	client.on_message=OnMessage
	client.connect("192.168.4.1",1883,120)
	return client

#standard mqtt onConnect definition
def OnConnect(client,userdata,flags,rc):
	client.subscribe("sensors/accelerometer/magnitude")

#standard mqtt onMessage definition
def OnMessage(client,userdata,msg):
    global temp
    if msg.topic =="sensors/accelerometer/magnitude":
		#for every new message we receive, we store the sensor reading in temp variable
        temp=msg.payload.decode()

def assign_value():
    global y
    global temp
    while True:
		#if sensor reading is less than 0.43, this is the sensor reading (most of the time) when bear is not moving
		#then we assume there is no movement, make y=0
        if float(temp)<0.43:
            y=0
        else:
		#else we assign this lastest sensor reading to y
            y=float(temp)
        time.sleep(1)

def animate(i):
    global value
    global label
    global y
	#clear the graph
    plt.clf()
	#y axies is 0 to 10, on average if u are moving a bear, we get a reading between 1-5
	#its mostly just 1.ish, unless you are throwing it across the room, u hardlt get it over 10
    plt.ylim(0,10)
    print(y)
    plt.bar(['Magnitude'],y)

#this while loop is to check network and mqtt connections
while True:
	#if anything fails, it indicates you are not connceted to the required wifi
    try:
        client = setup_mqtt()
        client.loop_start()
        print("mqtt connected")
		#loop is executed until both passes and break out the loop
        break
    except:
		#if not conncetion, print message to indicate
        print("Still waiting for mqtt connection")
		#sleep for a second to allow reconncetion, should be fine without it but while loop runs too fast, it may cause error
    time.sleep(1)

#if mqtt is connected then we start the assigning value thread
assign_ = threading.Thread(target=assign_value,daemon=True)
assign_.start()

#start the animation graph with interval of 500ms
ani= animation.FuncAnimation(fig, animate,interval=500)
plt.show()


#!/usr/bin/env python

import rospy
import rosbag
from geometry_msgs.msg import TransformStamped
from sensor_msgs.msg import LaserScan
import tf2_ros
import tf_conversions
import tf2_msgs

rospy.init_node('coding')

file = open('/home/vboxuser/catkin_ws/src/my_gmapping/src/examp3.txt', 'r')

bag = rosbag.Bag('data_gmapping.bag', 'w')

tf_broadcaster = tf2_ros.TransformBroadcaster()
tf_buffer = tf2_ros.Buffer()

count = 1

for line in file:

    odom_data = line.split(';')[0]
    lidar_data = line.split(';')[1]
    timestamp = count
    count += 1
    x = float(odom_data.split(',')[0])
    y = float(odom_data.split(',')[1])
    yaw = float(odom_data.split(',')[2])

    transform_stamped = TransformStamped()
    transform_stamped.header.stamp = rospy.Time.from_sec(timestamp)
    transform_stamped.header.frame_id = 'odom'
    transform_stamped.child_frame_id = 'base_link'
    transform_stamped.transform.translation.x = x
    transform_stamped.transform.translation.y = y
    transform_stamped.transform.translation.z = 0.0
    quaternion = tf_conversions.transformations.quaternion_from_euler(0, 0, yaw)
    transform_stamped.transform.rotation.x = quaternion[0]
    transform_stamped.transform.rotation.y = quaternion[1]
    transform_stamped.transform.rotation.z = quaternion[2]
    transform_stamped.transform.rotation.w = quaternion[3]

    tf_broadcaster.sendTransform(transform_stamped)
    tf_buffer.set_transform(transform_stamped, 'default_authority')
    tf_msg = tf2_msgs.msg.TFMessage()
    tf_msg.transforms.append(transform_stamped)

    scan = LaserScan()
    scan.header.stamp = rospy.Time.from_sec(timestamp)
    scan.header.frame_id = 'base_link'
    scan.angle_min = -2.094
    scan.angle_max = 2.094
    scan.angle_increment = 0.3524
    scan.time_increment = 0.1
    scan.scan_time = 0.1
    scan.range_min = 0.0
    scan.range_max = 5.6
    scan.ranges = [float(x) for x in lidar_data.split(',')]

    rospy.loginfo(scan)
    bag.write("/tf", tf_msg, rospy.Time.from_sec(timestamp))
    bag.write('/base_scan', scan, rospy.Time.from_sec(timestamp))

file.close()
bag.close()


import sys
input = sys.stdin.readline
from collections import deque
def bfs(x0, y0, x1, y1):
    q = deque()
    q.append([x0, y0])
    T[x0][y0] = 1 #지나간 곳은 체크 다시 지나갈 필요가 없다. 반복되기 때문!
    while q:
        a, b = q.popleft()
        if a == x1 and b == y1:
            print(T[x1][y1]-1)    #종착점까지 걸린 단계 출력!
            return
        for i in range(8):
            x = a + idx[i]
            y = b + idy[i]
            if 0 <= x < l and 0 <= y < l and T[x][y] == 0:
                q.append([x, y])
                T[x][y] = T[a][b] + 1
for _ in range(int(input())):
    l = int(input())
    x0, y0 = map(int, input().split()) # 초기 위치 설정
    x1, y1 = map(int, input().split()) # 종착점
    idx = [-2, -2, -1, -1, 1, 1, 2, 2]
    idy = [-1, 1, -2, 2, -2, 2, -1, 1]
    T = [[0 for _ in range(l)] for _ in range(l)]
    bfs(x0, y0, x1, y1)


import csv
import enum
import threading
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from functools import partial
from multiprocessing import Manager, Process
from multiprocessing import Queue as mpQueue
from multiprocessing.dummy import Pool
from queue import Queue

from api.api_client import YandexWeatherAPI
from api.common import get_logger
from api.handlers.base import QueueMixin
from utils import Args

logging = get_logger()

class DataFetchingTask(YandexWeatherAPI, QueueMixin):
    def fetch_cities(self, cities: list[str] | str):
        logging.info("DataFetchingTask::fetch_cities")
        with ThreadPoolExecutor() as pool:
            for city in cities:
                self.queue = pool.submit(self.get_forecasting, city)
        return self._queue

class DataCalculationTask:
    FINDING_HOURS = (9, 20)
    NORMAL_CONDITIONS = ("clear", "partly-cloudy", "overcast", "cloudy")

    def __init__(self):
        self.processes: list[Process] = []
        self.queue = Manager().Queue()

    @staticmethod
    def __count_common_overage(dates_info: list[dict]) -> tuple:
        temp_list = [
            d["overage_temperature"]
            for d in dates_info
            if isinstance(d["overage_temperature"], float)
        ]
        cond_list = [d["condition"] for d in dates_info if isinstance(d["condition"], int)]
        return (
            round(sum(temp_list) / len(temp_list), 1),
            round(sum(cond_list) / len(cond_list), 1),
        )

    @staticmethod
    def __count_overage_temperature(hours: list[dict]) -> float | str:
        temps = []
        for hour in hours:
            if int(hour["hour"]) in range(*DataCalculationTask.FINDING_HOURS):
                temps.append(hour["temp"])
        try:
            return round(sum(temps) / len(temps), 1)
        except ZeroDivisionError:
            return "Недостаточно данных"

    @staticmethod
    def __count_conditions(hours: list[dict]) -> int | str:
        cond = 0
        for hour in hours:
            if (
                int(hour["hour"]) in range(*DataCalculationTask.FINDING_HOURS)
                and hour["condition"] in DataCalculationTask.NORMAL_CONDITIONS
            ):
                cond += 1
        return "Недостаточно данных" if len(hours) != 24 and cond == 0 else cond

    @staticmethod
    def __sorted_and_get_rating(result: list[dict]) -> list[dict]:
        result = sorted(result, key=lambda k: (-k["temp_avg"], -k["cond_avg"]))
        for i, value in enumerate(result):
            value["rating"] = i + 1
        return result

    @classmethod
    def format_response_info(cls, city_response: dict, queue: mpQueue) -> None:
        city_info = {
            "name": city_response["geo_object"]["locality"]["name"],
            "dates_info": [],
            "rating": None,
        }

        for date in city_response["forecasts"]:
            city_info["dates_info"].append(
                {
                    "date": datetime.strptime(date["date"], "%Y-%m-%d").strftime("%d-%m"),
                    "overage_temperature": cls.__count_overage_temperature(date["hours"]),
                    "condition": cls.__count_conditions(date["hours"]),
                }
            )
            logging.debug(
                "DataCalculationTask::format_response_info::"
                "hours - %(hours)s ||| temp - %(temp)s ||| cond - %(cond)s",
                Args(
                    hours=[i["hour"] for i in date["hours"]],
                    temp=partial(cls.__count_overage_temperature, date["hours"]),
                    cond=partial(cls.__count_conditions, date["hours"]),
                ),
            )
        temp_avg, cond_avg = DataCalculationTask.__count_common_overage(city_info["dates_info"])
        city_info["temp_avg"] = temp_avg
        city_info["cond_avg"] = cond_avg
        queue.put(city_info)

        logging.debug("DataCalculationTask::format_response_info - {}".format(city_info))

    def __run_processes(self) -> None:
        [proc.start() for proc in self.processes]
        logging.debug(
            "DataCalculationTask::__run_processes: pid - name \n {}".format(
                [f"{i.pid}-{i.name}" for i in self.processes]
            )
        )

    def __join_processes(self) -> None:
        [proc.join() for proc in self.processes]
        logging.debug(
            "DataCalculationTask::__join_processes: pid - name \n {}".format(
                [f"{i.pid}-{i.name}" for i in self.processes]
            )
        )

    def calculate_cities(self, queue: Queue) -> list[dict]:
        for _ in range(queue.qsize()):
            self.processes.append(
                Process(
                    target=DataCalculationTask.format_response_info,
                    args=(queue.get().result(), self.queue),
                )
            )
        self.__run_processes()
        self.__join_processes()

        result = []
        while not self.queue.empty():
            result.append(self.queue.get())

        result = DataCalculationTask.__sorted_and_get_rating(result)
        return result

class DataAggregationTask(YandexWeatherAPI):
    class CSVFields(str, enum.Enum):
        CITY = "Город/День"
        AVERAGE = "Среднее"
        RATING = "Рейтинг"
        EMPTY = ""
        TEMP_AVERAGE = "Температура, среднее"
        CONDITION_AVERAGE = "Без осадков, среднее"

    def __init__(self) -> None:
        self.csv_locking = threading.Lock()
        self.writer: None | csv.DictWriter = None

    def __create_base_fields(self, city_result: dict) -> list:
        return [
            self.CSVFields.CITY.value,
            self.CSVFields.EMPTY.value,
            *[i["date"] for i in city_result["dates_info"]],
            self.CSVFields.AVERAGE.value,
            self.CSVFields.RATING.value,
        ]

    def __save_csv(self, city) -> None:
        row_temp = {
            self.CSVFields.CITY.value: city["name"],
            self.CSVFields.EMPTY.value: self.CSVFields.TEMP_AVERAGE.value,
            self.CSVFields.AVERAGE.value: city["temp_avg"],
            self.CSVFields.RATING.value: city["rating"],
        }

        row_cond = {
            self.CSVFields.CITY.value: self.CSVFields.EMPTY.value,
            self.CSVFields.EMPTY.value: self.CSVFields.CONDITION_AVERAGE.value,
            self.CSVFields.AVERAGE.value: city["cond_avg"],
            self.CSVFields.RATING.value: self.CSVFields.EMPTY.value,
        }

        for date in city["dates_info"]:
            row_temp[date["date"]] = date["overage_temperature"]
            row_cond[date["date"]] = date["condition"]
        self.writer.writerows([row_temp, row_cond])
        logging.info("DataAggregationTask::__save_csv::written - {}".format(city["name"]))

    def aggreagate_result(self, result: list[dict]) -> None:
        with open("cities.csv", "w", newline="") as file:
            self.writer = csv.DictWriter(file, fieldnames=self.__create_base_fields(result[0]))
            self.writer.writeheader()
            logging.info("DataAggregationTask::aggreagate_result::writeheader")

            with ThreadPoolExecutor() as pool:
                pool.map(self.__save_csv, result)

class DataAnalyzingTask(YandexWeatherAPI):
    max_temp: int = 0
    max_cond: int = 0
    max_city_name: str = ""

    def __init__(self, result: list[dict]) -> None:
        self.result = result
        self.__set_max_temp_cond()

    def __set_max_temp_cond(self) -> None:
        max_city = self.result[0]
        self.__class__.max_temp = max_city["temp_avg"]
        self.__class__.max_cond = max_city["cond_avg"]
        self.__class__.max_city_name = max_city["name"]

    @classmethod
    def checking_max_cities(cls, result: dict) -> str | None:
        if result["temp_avg"] == cls.max_temp and result["cond_avg"] == cls.max_cond:
            return result["name"]

    def analyze_result(self):
        with Pool(processes=4) as pool:
            result = pool.map(DataAnalyzingTask.checking_max_cities, self.result)
        return [city for city in result if city is not None]

class WeatherCityAnalyzer:
    def __init__(self, cities: list[str] | str):
        self.cities = cities

    def get_good_cities(self):
        queue = DataFetchingTask().fetch_cities(self.cities)
        result = DataCalculationTask().calculate_cities(queue)
        DataAggregationTask().aggreagate_result(result)
        return DataAnalyzingTask(result).analyze_result()


import copy
from einops import rearrange

import torch.nn as nn
from torch import einsum

#################### * MSA * ####################

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        inner_dim = dim_head * heads
        self.heads = heads
        self.scale = dim_head ** -0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def forward(self, x):
        N, T, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, "n t (h d) -> n h t d", h=h), qkv)

        dots = einsum("n h i d, n h j d -> n h i j", q, k) * self.scale
        attn = dots.softmax(dim=-1)

        out = einsum("n h i j, n h j d -> n h i d", attn, v)
        out = rearrange(out, "n h t d -> n t (h d)")
        out = self.to_out(out)
        return out

class ATULayer(nn.Module):
    def __init__(self, T):
        super(ATULayer, self).__init__()
        # self.attn = Residual(
        #     PreNorm(1600, Attention(1600, heads=8, dim_head=128, dropout=0.0))
        # )
        self.attn = Residual(
            PreNorm(2400, Attention(2400, heads=8, dim_head=128, dropout=0.0))
        )
        self.linear = nn.Linear(T, T)
        self.Tanh = nn.Tanh()

    def forward(self, x):
        N, C, V, T = x.shape
        x = x.permute(0, 3, 1, 2).reshape(N, T, C * V)
        # x = x.reshape(-1, 2400)
        x = self.attn(x)
        x = x.reshape(N, T, C, V).permute(0, 2, 3, 1)

        x = self.linear(x)  # N,C,V,T
        x = self.Tanh(x)
        return x

class ATNet(nn.Module):
    def __init__(self, layers, T):
        super(ATNet, self).__init__()
        self.layer = layers
        ATUlayer = ATULayer(T)
        self.ATUNet = nn.ModuleList(
            [copy.deepcopy(ATUlayer) for _ in range(self.layer)]
        )

    def forward(self, x):
        N, C, T, V = x.shape
        x = x.permute(0, 1, 3, 2)
        for ATU in self.ATUNet:
            x = ATU(x)
        x = x.permute(0, 1, 3, 2)
        return x