import os
import csv
import copy
import time
import tqdm
import pprint
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
from torch.multiprocessing import Process, Pool
from torch.distributions import Categorical
from torch.utils.tensorboard import SummaryWriter
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from abc import abstractmethod

from virne.solver import Solver
from virne.solver.heuristic.node_rank import *

from .searcher import *
from .buffer import RolloutBuffer
from .shared_adam import SharedAdam, sync_gradients
from ..utils import apply_mask_to_logit, get_observations_sample, RunningMeanStd
from virne.utils import test_running_time

class RLSolver(Solver):
    """General Reinforcement Learning Solve"""
    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(RLSolver, self).__init__(controller, recorder, counter, **kwargs)
        self.rank = 0
        # baseline
        self.use_baseline_solver = kwargs.get('use_baseline_solver', False)
        if self.use_baseline_solver:
            self.baseline_solvers = {}
            baselin_solver_name =  kwargs.get('baselin_solver_name', 'grc')
            if baselin_solver_name == 'grc':
                self.baseline_solver = GRCRankSolver(controller, recorder, counter, **kwargs)
            elif baselin_solver_name == 'self':
                self.baseline_solver = copy.deepcopy(self)
        # training
        self.use_cuda = kwargs.get('use_cuda', True)
        if self.use_cuda and torch.cuda.is_available():
            self.device = torch.device('cuda:0')
            self.device_name = torch.cuda.get_device_name(torch.cuda.current_device())
        else:
            self.device = torch.device('cpu')
            self.device_name = 'CPU'
            self.use_cuda = False
        self.num_workers = kwargs.get('num_workers', 1)
        self.distributed_training = False if self.num_workers == 1 else kwargs.get('distributed_training', False)
        # rl
        self.gamma = kwargs.get('rl_gamma', 0.95)
        self.gae_lambda = kwargs.get('gae_lambda', 0.98)
        self.coef_critic_loss = kwargs.get('coef_critic_loss', 0.5)
        self.coef_entropy_loss = kwargs.get('coef_entropy_loss', 0.01)
        self.coef_mask_loss = kwargs.get('coef_mask_loss', 0.01)
        self.weight_decay = kwargs.get('weight_decay', 0.00001)
        self.lr = kwargs.get('lr', 0.001)
        self.lr_actor = kwargs.get('lr_actor', 0.005)
        self.lr_critic = kwargs.get('lr_critic', 0.001)
        self.lr_scheduler = None
        self.criterion_critic = nn.MSELoss()
        self.compute_advantage_method = kwargs.get('compute_advantage_method', 'gae')
        # nn
        self.embedding_dim = kwargs.get('embedding_dim', 64)
        self.dropout_prob = kwargs.get('dropout_prob', 0.5)
        self.batch_norm = kwargs.get('batch_norm', False)
        # train
        self.batch_size = kwargs.get('batch_size', 128)
        self.use_negative_sample = kwargs.get('use_negative_sample', False)
        self.target_steps = kwargs.get('target_steps', 128)
        self.eval_interval = kwargs.get('eval_interval', 5)
        # eval;
        self.k_searching = kwargs.get('k_searching', 4)
        self.decode_strategy = kwargs.get('decode_strategy', 'sample')
        # tricks
        self.mask_actions = kwargs.get('mask_actions', True)
        self.maskable_policy = kwargs.get('maskable_policy', True)
        self.norm_reward = kwargs.get('norm_reward', False)
        if self.norm_reward:
            self.running_stats = RunningMeanStd(shape=1)
        self.norm_advantage = kwargs.get('norm_advantage', True)
        self.clip_grad = kwargs.get('clip_grad', True)
        self.clip_reward = kwargs.get('clip_reward', False)
        self.clip_reward = kwargs.get('max_reward', 1.0)
        self.max_grad_norm = kwargs.get('max_grad_norm', 1.)
        self.softmax_temp = 1.
        # log
        self.open_tb = kwargs.get('open_tb', True)
        self.log_dir = os.path.join(self.save_dir, 'log')
        self.model_dir = os.path.join(self.save_dir, 'model')
        self.writer = SummaryWriter(self.log_dir) if self.open_tb else None
        self.training_info = []
        self.buffer = RolloutBuffer()
        # save
        self.log_interval = kwargs.get('log_interval', 1)
        self.save_interval = kwargs.get('save_interval', 1)

        for dir in [self.save_dir, self.log_dir, self.model_dir]:
            if not os.path.exists(dir):
                os.makedirs(dir)
        # counter
        self.update_time = 0
        self.time_step = 0
        if self.verbose >= 0:
            self.show_config()
        # optimizer
        self.make_policy = make_policy
        self.policy, self.optimizer = self.make_policy(self)
        self.preprocess_obs = obs_as_tensor

    def show_config(self, ):
        print(f'*' * 50)
        print(f'Key parameters of RL training are as following: ')
        print(f'*' * 50)
        print(f'       device: {self.device_name}')
        print(f'   num_workers: {self.num_workers}')
        print(f'  distributed: {self.distributed_training}')
        print(f'     rl_gamma: {self.gamma}')
        print(f'     lr_actor: {self.lr_actor}')
        print(f'    lr_critic: {self.lr_critic}')
        print(f'   batch_size: {self.batch_size}')
        print(f'embedding_dim: {self.embedding_dim}')
        print(f'coef_ent_loss: {self.coef_entropy_loss}')
        print(f'     norm_adv: {self.norm_advantage}')
        print(f'  norm_reward: {self.norm_advantage}')
        print(f'    clip_grad: {self.clip_grad}')
        print(f'max_grad_norm: {self.max_grad_norm}')
        print(f'save_interval: {self.save_interval}')
        print(f'eval_interval: {self.eval_interval}')
        print(f'*' * 50)
        print()
        print(f'Logging training info at {os.path.join(self.log_dir, "training_info.csv")}')

    @abstractmethod
    def preprocess_obs(self, obs):
        return NotImplementedError

    def solve_with_baseline(self, instance, baseline='grc'):
        if self.baseline_solver is None:
            self.baseline_solver = GRCRankSolver()
        solution = self.baseline_solver.solve(instance)
        solution_info = self.counter.count_solution(instance['v_net'], solution)
        return solution_info

    def log(self, info, update_time, only_tb=False):
        if self.open_tb and hasattr(self, 'writer') and self.writer is not None:
            for key, value in info.items():
                self.writer.add_scalar(key, value, update_time)
        if only_tb:
            return
        log_fpath = os.path.join(self.log_dir, 'training_info.csv')
        write_head = not os.path.exists(log_fpath)
        with open(log_fpath, 'a+', newline='') as f:
            writer = csv.writer(f)
            if write_head:
                writer.writerow(['update_time'] + list(info.keys()))
            writer.writerow([update_time] + list(info.values()))
        if self.verbose >= 1:
            if update_time ==0 or update_time % 1000 == 0:
                info_key_str = ' & '.join([f'{k}' for k, v in info.items() if sum([s in k for s in ['loss', 'prob', 'return', 'penalty']])])
                print(f'             {update_time:06d} | ' + info_key_str)
            info_str = ' & '.join([f'{v:+3.4f}' for k, v in info.items() if sum([s in k for s in ['loss', 'prob', 'return', 'penalty']])])
            print(f'Update time: {update_time:06d} | ' + info_str)

    def get_action_prob_dist(self, observation):
        with torch.no_grad():
            action_logits = self.policy.act(observation)
        if 'action_mask' in observation and self.mask_actions:
            mask = observation['action_mask']
            candicate_action_logits = apply_mask_to_logit(action_logits, mask)
        else:
            candicate_action_logits = action_logits
        action_prob_dist = F.softmax(candicate_action_logits / self.softmax_temp, dim=-1)
        return action_prob_dist, candicate_action_logits

    def select_action(self, observation, sample=True):
        with torch.no_grad():
            action_logits = self.policy.act(observation)
        if 'action_mask' in observation and self.mask_actions:
            mask = observation['action_mask']
            candicate_action_logits = apply_mask_to_logit(action_logits, mask)
        else:
            candicate_action_logits = action_logits
        # if self.mask_actions and self.maskable_policy:
        #     candicate_action_probs = F.softmax(candicate_action_logits / self.softmax_temp, dim=-1)
        #     candicate_action_dist = Categorical(probs=candicate_action_probs)
        # else:
        #     candicate_action_probs = F.softmax(action_logits / self.softmax_temp, dim=-1)
        #     candicate_action_dist = Categorical(probs=candicate_action_probs)
        candicate_action_dist = Categorical(logits=candicate_action_logits / self.softmax_temp)
        raw_action_dist = Categorical(logits=action_logits / self.softmax_temp)

        if self.mask_actions and self.maskable_policy:
            action_dist_for_log_prob = candicate_action_dist
        else:
            action_dist_for_log_prob = raw_action_dist

        if sample:
            action = candicate_action_dist.sample()
        else:
            action = candicate_action_logits.argmax(-1)

        action_logprob = action_dist_for_log_prob.log_prob(action)

        if torch.numel(action) == 1:
            action = action.item()
        else:
            action = action.reshape(-1, ).cpu().detach().numpy()
        # action = action.squeeze(-1).cpu()
        return action, action_logprob.cpu().detach().numpy()

    def evaluate_actions(self, old_observations, old_actions, return_others=False):
        actions_logits = self.policy.act(old_observations)
        actions_probs = F.softmax(actions_logits / self.softmax_temp, dim=-1)
        if 'action_mask' in old_observations:
            masks = old_observations['action_mask']
            candicate_actions_logits = apply_mask_to_logit(actions_logits, masks)
        else:
            masks = None
            candicate_actions_logits = actions_logits

        candicate_actions_probs = F.softmax(candicate_actions_logits, dim=-1)

        dist = Categorical(actions_probs)
        candicate_dist = Categorical(candicate_actions_probs)
        policy_dist = candicate_dist if self.mask_actions and self.maskable_policy else dist

        action_logprobs = policy_dist.log_prob(old_actions)
        dist_entropy = policy_dist.entropy()

        values = self.policy.evaluate(old_observations).squeeze(-1) if hasattr(self.policy, 'evaluate') else None

        if return_others:
            other = {}
            if masks is not None:
                mask_actions_probs = actions_probs * (~masks.bool())
                other['mask_actions_probs'] = mask_actions_probs.sum(-1).mean()
                if hasattr(self.policy, 'predictor'):
                    predicted_masks_logits = self.policy.predict(old_observations)
                    print(predicted_masks_logits)
                    prediction_loss = F.binary_cross_entropy(predicted_masks_logits, masks.float())
                    other['prediction_loss'] = prediction_loss
                    predicted_masks = torch.where(predicted_masks_logits > 0.5, 1., 0.)
                    correct_count = torch.eq(predicted_masks.bool(), masks.bool()).sum(-1).float().mean(0)
                    acc = correct_count / predicted_masks.shape[-1]
                    print(prediction_loss, correct_count, acc)
            return values, action_logprobs, dist_entropy, other

        return values, action_logprobs, dist_entropy

    def estimate_value(self, observation):
        """
        Estimate the value of an observation
        """
        with torch.no_grad():
            estimated_value = self.policy.evaluate(observation).squeeze(-1).detach().cpu().item()
        return estimated_value

    def to_sub_solver(self):
        temp_dict = {}
        unnecessary_attributes = ['policy', 'optimizer', 'lr_scheduler', 'searcher', 'writer']
        for attr_name in unnecessary_attributes:
            if hasattr(self, attr_name):
                temp_dict[attr_name] = getattr(self, attr_name)
                delattr(self, attr_name)
        sub_solver = copy.deepcopy(self)
        sub_solver.policy, _ = sub_solver.make_policy(sub_solver)
        for attr_name in temp_dict.keys():
            setattr(self, attr_name, temp_dict[attr_name])
        sub_solver.policy.load_state_dict(self.policy.state_dict())
        return sub_solver

    def get_worker(self, rank=0):
        self.policy.share_memory()
        shared_policy = self.policy
        self.optimizer = SharedAdam.from_optim(self.optimizer)
        shared_optimizer = self.optimizer

        lr_scheduler = self.lr_scheduler
        writer = self.writer
        unnecessary_attributes = ['policy', 'optimizer', 'lr_scheduler', 'searcher', 'writer']
        for attr_name in unnecessary_attributes:
            if hasattr(self, attr_name):
                delattr(self, attr_name)
        # self.eval()
        worker = copy.deepcopy(self)
        worker.rank = rank
        worker.writer = None
        worker.lr_scheduler = None
        if self.use_cuda:
            num_gpu_devices = torch.cuda.device_count()
            worker_device_id = rank % num_gpu_devices
            worker.device = torch.device(f'cuda:{worker_device_id}')
        worker.policy, worker.optimizer = worker.make_policy(worker)
        worker.policy.load_state_dict(shared_policy.state_dict())
        worker.policy.to(worker.device)
        # worker.optimizer = torch.optim.Adam(worker.policy.parameters(), lr=self.lr)
        worker.shared_policy = shared_policy
        worker.shared_optimizer = shared_optimizer
        self.lr_scheduler = lr_scheduler
        self.writer = writer
        self.policy = shared_policy
        self.optimizer = shared_optimizer
        return worker

    def save_model(self, checkpoint_fname):
        checkpoint_fname = os.path.join(self.model_dir, checkpoint_fname)
        # torch.save(self.policy.state_dict(), checkpoint_fname)
        torch.save({
            'policy': self.policy.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            # 'lr_scheduler_state_dict': self.lr_scheduler.state_dict()
        }, checkpoint_fname)
        print(f'Save model to {checkpoint_fname}\n') if self.verbose >= 0 else None

    def load_model(self, checkpoint_path):
        print('Attempting to load the pretrained model')
        try:
            checkpoint = torch.load(checkpoint_path)
            if 'policy' not in checkpoint:
                self.policy.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))
            else:
                self.policy.load_state_dict(checkpoint['policy'])
                self.optimizer.load_state_dict(checkpoint['optimizer'])
            print(f'Loaded pretrained model from {checkpoint_path}') if self.verbose >= 0 else None
        except Exception as e:
            print(f'Load failed from {checkpoint_path}\nInitilized with random parameters') if self.verbose >= 0 else None

    def train(self):
        """Set the mode to train"""
        self.policy.train()
        if hasattr(self, 'searcher'):
            delattr(self, 'searcher')

    def eval(self, decode_strategy=None, k=None):
        if decode_strategy is None: decode_strategy = self.decode_strategy
        if k is None: k = self.k_searching
        assert k >= 1, f'k should greater than 0. (k={k})'
        self.policy.eval()
        self.searcher = get_searcher(decode_strategy,
                                    policy=self.policy,
                                    preprocess_obs_func=self.preprocess_obs,
                                    make_policy_func=self.make_policy,
                                    k=k, device=self.device,
                                    mask_actions=self.mask_actions,
                                    maskable_policy=self.maskable_policy)

    def update_grad(self, loss):
        # update parameters
        if self.distributed_training:
            with self.lock:
                self.optimizer.zero_grad()
                self.shared_optimizer.zero_grad()
                loss.backward()
                grad_clipped = torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) if self.clip_grad else None
                sync_gradients(self.shared_policy, self.policy)
                self.optimizer.step()
                self.shared_optimizer.step()
        else:
            self.optimizer.zero_grad()
            loss.backward()
            grad_clipped = torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) if self.clip_grad else None
            self.optimizer.step()
        return grad_clipped

    def sync_parameters(self):
        assert self.distributed_training, 'distributed_training should be True'
        with self.lock:
            self.policy.load_state_dict(self.shared_policy.state_dict())

    def learn(self, env, num_epochs=1, **kwargs):
        self.start_time = time.time()
        if self.distributed_training:
            self.learn_distributedly(env, num_epochs)
        else:
            print(f'Start to learn singly')
            self.learn_singly(env, num_epochs)
        print(f'Start to validate')
        self.save_model(f'model.pkl')
        # self.validate(env)
        self.end_time = time.time()
        print(f'\nTotal training time: {(self.end_time - self.start_time) / 3600:4.6f} h')

    def learn_distributedly(self, env, num_epochs, **kwargs):
        assert self.distributed_training, 'distributed_training should be True'
        assert num_epochs % self.num_workers == 0, 'num_epochs should be divisible by num_workers'
        job_list = []
        mp.set_start_method('spawn')
        lock = mp.Lock()
        worker_num_epochs = int(num_epochs // self.num_workers)
        worker_save_interval = int(np.ceil(self.save_interval / self.num_workers))
        worker_eval_interval = int(np.ceil(self.eval_interval / self.num_workers))
        print(f'Distributed training with {self.num_workers} workers')
        print(f'Worker num epochs:    {   worker_num_epochs:3d} epochs')
        print(f'Worker save interval: {worker_save_interval:3d} epochs')
        print(f'Worker eval interval: {worker_eval_interval:3d} epochs')
        print()
        for worker_rank in range(self.num_workers):
            env = copy.deepcopy(env)
            worker = self.get_worker(worker_rank)
            worker.lock = lock
            worker.save_interval = worker_save_interval
            worker.eval_interval = worker_eval_interval
            if worker_rank != 0:
                worker.verbose = 0
                env.verbose = 0
            job = mp.Process(target=worker.learn_singly, args=(env, worker_num_epochs))
            job_list.append(job)
            job.start()
        for job in job_list:
            job.join()

class PGSolver(RLSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(PGSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)

    def update(self, ):
        observations = self.preprocess_obs(self.buffer.observations, self.device)
        actions = torch.LongTensor(np.array(self.buffer.actions)).to(self.device)
        returns = torch.FloatTensor(np.array(self.buffer.returns)).to(self.device)
        _, action_logprobs, _, _ = self.evaluate_actions(observations, actions, return_others=True)

        loss = - (action_logprobs * returns).mean()

        grad_clipped = self.update_grad(loss)

        info = {
            'lr': self.optimizer.defaults['lr'],
            'loss/loss': loss.detach().cpu(),
            'value/logprob': action_logprobs.detach().mean().cpu(),
            'value/return': returns.detach().mean().cpu(),
        }
        self.log(info, self.update_time, only_tb=False)

        self.buffer.clear()
        self.lr_scheduler.step() if self.lr_scheduler is not None else None
        self.update_time += 1

        if self.distributed_training: self.sync_parameters()
        return loss.detach()

class A2CSolver(RLSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(A2CSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)
        self.repeat_times = 1

    def update(self, ):
        observations = self.preprocess_obs(self.buffer.observations, self.device)
        actions = torch.LongTensor(np.array(self.buffer.actions)).to(self.device)
        returns = torch.FloatTensor(self.buffer.returns).to(self.device)
        values, action_logprobs, dist_entropy, other = self.evaluate_actions(observations, actions, return_others=True)
        advantages = returns - values.detach()
        if self.norm_advantage:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        actor_loss = - (action_logprobs * advantages).mean()
        critic_loss = F.mse_loss(returns, values)
        entropy_loss = dist_entropy.mean()
        loss = actor_loss + self.coef_critic_loss * critic_loss + self.coef_entropy_loss * entropy_loss

        grad_clipped = self.update_grad(loss)

        info = {
            'lr': self.optimizer.defaults['lr'],
            'loss/loss': loss.detach().cpu().numpy(),
            'loss/actor_loss': actor_loss.detach().cpu().numpy(),
            'loss/critic_loss': critic_loss.detach().cpu().numpy(),
            'loss/entropy_loss': entropy_loss.detach().cpu().numpy(),
            'value/logprob': action_logprobs.detach().mean().cpu().numpy(),
            'value/return': returns.detach().mean().cpu().numpy()
        }
        self.log(info, self.update_time, only_tb=False)

        self.buffer.clear()
        self.lr_scheduler.step() if self.lr_scheduler is not None else None
        self.update_time += 1

        if self.distributed_training: self.sync_parameters()
        return loss.detach()

# class A3CSolver(A2CSolver):

#     def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
#         super(A3CSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)
#         self.repeat_times = 1

class PPOSolver(RLSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(PPOSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)
        self.repeat_times = kwargs.get('repeat_times', 10)
        self.gae_lambda = kwargs.get('gae_lambda', 0.98)
        self.eps_clip = kwargs.get('eps_clip', 0.2)

    def update(self, ):
        # assert self.buffer.size() >= self.batch_size
        device = torch.device('cpu')

        batch_observations = self.preprocess_obs(self.buffer.observations, device)
        # # batch_actions = torch.LongTensor(np.concatenate(self.buffer.actions, axis=0)).to(self.device)
        batch_actions = torch.LongTensor(np.array(self.buffer.actions)).to(self.device)
        batch_old_action_logprobs = torch.FloatTensor(np.concatenate(self.buffer.logprobs, axis=0))
        batch_rewards = torch.FloatTensor(self.buffer.rewards)
        batch_returns = torch.FloatTensor(self.buffer.returns)
        if self.norm_reward:
            batch_returns = (batch_returns - batch_returns.mean()) / (batch_returns.std() + 1e-9)
        sample_times = 1 + int(self.buffer.size() * self.repeat_times / self.batch_size)
        for i in range(sample_times):
            sample_indices = torch.randint(0, self.buffer.size(), size=(self.batch_size,)).long()
            # observations  = get_observations_sample(batch_observations, sample_indices, self.device)
            sample_obersevations = [self.buffer.observations[i] for i in sample_indices]
            observations = self.preprocess_obs(sample_obersevations, self.device)
            actions = batch_actions[sample_indices].to(self.device)
            returns = batch_returns[sample_indices].to(self.device)
            old_action_logprobs = batch_old_action_logprobs[sample_indices].to(self.device)
            # evaluate actions and observations
            values, action_logprobs, dist_entropy, other = self.evaluate_actions(observations, actions, return_others=True)

            # calculate advantage
            advantages = returns - values.detach()
            if self.norm_advantage and values.numel() != 0:
                advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-9)

            ratio = torch.exp(action_logprobs - old_action_logprobs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1. - self.eps_clip, 1. + self.eps_clip) * advantages
            actor_loss = - torch.min(surr1, surr2).mean()
            critic_loss = self.criterion_critic(returns, values)
            entropy_loss = dist_entropy.mean()

            mask_loss = other.get('mask_actions_probs', 0)
            prediction_loss = other.get('prediction_loss', 0)

            loss = actor_loss + self.coef_critic_loss * critic_loss - self.coef_entropy_loss * entropy_loss + self.coef_mask_loss * mask_loss + prediction_loss
            # update parameters
            grad_clipped = self.update_grad(loss)

            if self.open_tb and self.update_time % self.log_interval == 0:
                info = {
                    'lr': self.optimizer.defaults['lr'],
                    'loss/loss': loss.detach().cpu().numpy(),
                    'loss/actor_loss': actor_loss.detach().cpu().numpy(),
                    'loss/critic_loss': critic_loss.detach().cpu().numpy(),
                    'loss/entropy_loss': entropy_loss.detach().cpu().numpy(),
                    'value/logprob': action_logprobs.detach().mean().cpu().numpy(),
                    'value/old_action_logprob': old_action_logprobs.mean().cpu().numpy(),
                    'value/value': values.detach().mean().cpu().numpy(),
                    'value/return': returns.mean().cpu().numpy(),
                    'value/advantage': advantages.detach().mean().cpu().numpy(),
                    'value/reward': batch_rewards.mean().cpu().numpy(),
                    'grad/grad_clipped': grad_clipped.detach().cpu().numpy()
                }
                only_tb = not (i == sample_times-1)
                self.log(info, self.update_time, only_tb=only_tb)

            self.update_time += 1

        # print(f'loss: {loss.detach():+2.4f} = {actor_loss.detach():+2.4f} & {critic_loss:+2.4f} & {entropy_loss:+2.4f} & {mask_loss:+2.4f}, ' +
        #         f'action log_prob: {action_logprobs.mean():+2.4f} (old: {batch_old_action_logprobs.detach().mean():+2.4f}), ' +
        #         f'mean reward: {returns.detach().mean():2.4f}', file=self.fwriter) if self.verbose >= 0 else None
        self.lr_scheduler.step() if self.lr_scheduler is not None else None

        self.buffer.clear()

        if self.distributed_training: self.sync_parameters()
        return loss.detach()

class A3CSolver(PPOSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(A3CSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)

class ARPPOSolver(RLSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(ARPPOSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)
        self.repeat_times = kwargs.get('repeat_times', 10)
        self.gae_lambda = kwargs.get('gae_lambda', 0.98)
        self.eps_clip = kwargs.get('eps_clip', 0.2)

    def update(self, ):
        assert self.buffer.size() >= self.batch_size
        device = torch.device('cpu')
        batch_observations = self.preprocess_obs(self.buffer.observations, device)
        batch_actions = torch.LongTensor(np.array(self.buffer.actions)).to(self.device)
        batch_old_action_logprobs = torch.FloatTensor(np.concatenate(self.buffer.logprobs, axis=0))
        batch_rewards = torch.FloatTensor(self.buffer.rewards)
        mean_batch_rewards = batch_rewards.mean()

        batch_returns = torch.FloatTensor(self.buffer.returns)
        sample_times = 1 + int(self.buffer.size() * self.repeat_times / self.batch_size)
        for i in range(sample_times):
            sample_indices = torch.randint(0, self.buffer.size(), size=(self.batch_size,)).long()
            observations = get_observations_sample(batch_observations, sample_indices, device=self.device)
            actions = batch_actions[sample_indices].to(self.device)
            returns = batch_returns[sample_indices].to(self.device)
            old_action_logprobs = batch_old_action_logprobs[sample_indices].to(self.device)
            # masks = batch_masks[sample_indices].to(self.device) if batch_masks is not None else None
            # evaluate actions and observations
            values, action_logprobs, dist_entropy, other = self.evaluate_actions(observations, actions, return_others=True)

            # calculate advantage
            advantages = returns - values.detach()
            if self.norm_advantage and values.numel() != 0:
                advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-9)
            ratio = torch.exp(action_logprobs - old_action_logprobs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1. - self.eps_clip, 1. + self.eps_clip) * advantages
            actor_loss = - torch.min(surr1, surr2).mean()
            critic_loss = self.criterion_critic(returns, values)
            entropy_loss = dist_entropy.mean()

            mask_loss = other.get('mask_actions_probs', 0)
            prediction_loss = other.get('prediction_loss', 0)

            loss = actor_loss + self.coef_critic_loss * critic_loss - self.coef_entropy_loss * entropy_loss + self.coef_mask_loss * mask_loss + prediction_loss
            # update parameters
            grad_clipped = self.update_grad(loss)

            if self.open_tb and self.update_time % self.log_interval == 0:
                info = {
                    'lr': self.optimizer.defaults['lr'],
                    'loss/loss': loss.detach().cpu().numpy(),
                    'loss/actor_loss': actor_loss.detach().cpu().numpy(),
                    'loss/critic_loss': critic_loss.detach().cpu().numpy(),
                    'loss/entropy_loss': entropy_loss.detach().cpu().numpy(),
                    'value/logprob': action_logprobs.detach().mean().cpu().numpy(),
                    'value/old_action_logprob': old_action_logprobs.mean().cpu().numpy(),
                    'value/value': values.detach().mean().cpu().numpy(),
                    'value/return': returns.mean().cpu().numpy(),
                    'value/advantages': advantages.mean().cpu().numpy(),
                    'value/reward': batch_rewards.mean().cpu().numpy(),
                    'grad/grad_clipped': grad_clipped.detach().cpu().numpy()
                }
                only_tb = not (i == sample_times-1)
                self.log(info, self.update_time, only_tb=only_tb)

            self.update_time += 1

        # print(f'loss: {loss.detach():+2.4f} = {actor_loss.detach():+2.4f} & {critic_loss:+2.4f} & {entropy_loss:+2.4f} & {mask_loss:+2.4f}, ' +
        #         f'action log_prob: {action_logprobs.mean():+2.4f} (old: {batch_old_action_logprobs.detach().mean():+2.4f}), ' +
        #         f'mean reward: {returns.detach().mean():2.4f}', file=self.fwriter) if self.verbose >= 0 else None
        self.lr_scheduler.step() if self.lr_scheduler is not None else None
        self.buffer.clear()

        if self.distributed_training: self.sync_parameters()
        return loss.detach()

class DDPGSolver(RLSolver):

    def __init__(self, controller, recorder, counter, make_policy, obs_as_tensor, **kwargs):
        super(DDPGSolver, self).__init__(controller, recorder, counter, make_policy, obs_as_tensor, **kwargs)
        self.repeat_times = kwargs.get('repeat_times', 10)
        self.gae_lambda = kwargs.get('gae_lambda', 0.98)
        self.eps_clip = kwargs.get('eps_clip', 0.2)

    def update(self, ):
        # assert self.buffer.size() >= self.batch_size
        device = torch.device('cpu')

        batch_observations = self.preprocess_obs(self.buffer.observations, device)
        # # batch_actions = torch.LongTensor(np.concatenate(self.buffer.actions, axis=0)).to(self.device)
        batch_actions = torch.LongTensor(np.array(self.buffer.actions)).to(self.device)
        batch_old_action_logprobs = torch.FloatTensor(np.concatenate(self.buffer.logprobs, axis=0))
        batch_rewards = torch.FloatTensor(self.buffer.rewards)
        batch_returns = torch.FloatTensor(self.buffer.returns)
        if self.norm_reward:
            batch_returns = (batch_returns - batch_returns.mean()) / (batch_returns.std() + 1e-9)
        sample_times = 1 + int(self.buffer.size() * self.repeat_times / self.batch_size)
        for i in range(sample_times):
            sample_indices = torch.randint(0, self.buffer.size(), size=(self.batch_size,)).long()
            # observations  = get_observations_sample(batch_observations, sample_indices, self.device)
            sample_obersevations = [self.buffer.observations[i] for i in sample_indices]
            observations = self.preprocess_obs(sample_obersevations, self.device)
            actions = batch_actions[sample_indices].to(self.device)
            returns = batch_returns[sample_indices].to(self.device)
            old_action_logprobs = batch_old_action_logprobs[sample_indices].to(self.device)
            # evaluate actions and observations
            values, action_logprobs, dist_entropy, other = self.evaluate_actions(observations, actions, return_others=True)

            # calculate advantage
            advantages = returns - values.detach()
            if self.norm_advantage and values.numel() != 0:
                advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-9)

            ratio = torch.exp(action_logprobs - old_action_logprobs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1. - self.eps_clip, 1. + self.eps_clip) * advantages
            actor_loss = - torch.min(surr1, surr2).mean()
            critic_loss = self.criterion_critic(returns, values)
            entropy_loss = dist_entropy.mean()

            mask_loss = other.get('mask_actions_probs', 0)
            prediction_loss = other.get('prediction_loss', 0)

            loss = actor_loss + self.coef_critic_loss * critic_loss - self.coef_entropy_loss * entropy_loss + self.coef_mask_loss * mask_loss + prediction_loss
            # update parameters
            grad_clipped = self.update_grad(loss)

            if self.open_tb and self.update_time % self.log_interval == 0:
                info = {
                    'lr': self.optimizer.defaults['lr'],
                    'loss/loss': loss.detach().cpu().numpy(),
                    'loss/actor_loss': actor_loss.detach().cpu().numpy(),
                    'loss/critic_loss': critic_loss.detach().cpu().numpy(),
                    'loss/entropy_loss': entropy_loss.detach().cpu().numpy(),
                    'value/logprob': action_logprobs.detach().mean().cpu().numpy(),
                    'value/old_action_logprob': old_action_logprobs.mean().cpu().numpy(),
                    'value/value': values.detach().mean().cpu().numpy(),
                    'value/return': returns.mean().cpu().numpy(),
                    'value/advantage': advantages.detach().mean().cpu().numpy(),
                    'value/reward': batch_rewards.mean().cpu().numpy(),
                    'grad/grad_clipped': grad_clipped.detach().cpu().numpy()
                }
                only_tb = not (i == sample_times-1)
                self.log(info, self.update_time, only_tb=only_tb)

            self.update_time += 1

        # print(f'loss: {loss.detach():+2.4f} = {actor_loss.detach():+2.4f} & {critic_loss:+2.4f} & {entropy_loss:+2.4f} & {mask_loss:+2.4f}, ' +
        #         f'action log_prob: {action_logprobs.mean():+2.4f} (old: {batch_old_action_logprobs.detach().mean():+2.4f}), ' +
        #         f'mean reward: {returns.detach().mean():2.4f}', file=self.fwriter) if self.verbose >= 0 else None
        self.lr_scheduler.step() if self.lr_scheduler is not None else None

        self.buffer.clear()

        if self.distributed_training: self.sync_parameters()
        return loss.detach()


import pandas as pd
import os
import sys
import json
import re

with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'id_city.json')) as f:
    id_city = json.load(f)
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'city_region.json')) as f:
    city_region = json.load(f)

def clean(text):
    if type(text) != str:
        return ''

    text = text.lower()
    text = text.replace('ё', 'е')
    return text

def get_hashtags(text):
    if type(text) != str:
        return ''

    hashtag_regex = r'#[А-Яа-я]+'
    return list(map(lambda x: clean(x[1:]), re.findall(hashtag_regex, text)))

def concat_json1(path):
    json_files = list(filter(lambda x: x[-5:]  == '.json', os.listdir(path)))

    media = []
    for json_file in json_files:
        with open(path + '/' + json_file, 'r') as file:
            media += map(lambda x: [x, json_file[:-5].lower()], json.load(file))

    media = pd.DataFrame(media)
    media.rename({0: 'news', 1: 'city'}, axis="columns", inplace=True)

    media['city'] = media.city.apply(clean)
    media.news = media.news.apply(lambda x: x.replace('\n', ' '))
    media['date'] = ''
    media['tags'] = '[]'
    media['clarification'] = ''
    media['region'] = media.city.apply(lambda x: city_region[x])
    media['link'] = ''

    media.to_csv(path + '/collected.csv')

def concat_json2(path):
    json_files = list(filter(lambda x: x[-5:]  == '.json', os.listdir(path)))

    media = []
    for json_file in json_files:
        with open(os.path.join(path, json_file), 'r') as file:
            media += map(lambda x: [x[0], x[1], json_file[:-5].lower()], json.load(file))

    media = pd.DataFrame(media)
    media.rename({0: 'news', 1: 'link', 2: 'city'}, axis="columns", inplace=True)
    media['city'] = media.city.apply(clean)
    media['date'] = ''
    media['tags'] = '[]'
    media['clarification'] = ''
    media['region'] = media.city.apply(lambda x: city_region[x])

    media.to_csv(path + '/collected.csv')

def concat_csv(path, delim='№%:№%:'):
    csv_files = list(filter(lambda x: x[:2] == 'pu', os.listdir(path)))
    vk_publics = []
    for csv_file in csv_files:
        vk_publics.append(pd.read_csv(path + '/' + csv_file, sep=delim, engine='python', header=None))
    vk_publics = pd.concat(vk_publics, sort=True, ignore_index=True)

    vk_publics.rename({0: 'news', 1: 'date', 2: 'link'}, axis="columns", inplace=True)

    vk_publics['tags'] = vk_publics.news.apply(get_hashtags)
    # vk_publics['news'] = vk_publics['news'].apply(clean)
    vk_publics['clarification'] = ''
    vk_publics['city'] = vk_publics.link.apply(lambda x: clean(id_city[str(x)]))
    vk_publics['region'] = vk_publics.city.apply(lambda x: city_region[x])

    links = vk_publics.link.apply(lambda x: 'vk.com/public' + str(x))
    vk_publics.drop('link', inplace=True, axis=1)
    vk_publics['link'] = links

    vk_publics.to_csv(path + '/collected.csv')

def main():
    file_type = sys.argv[1]
    path = sys.argv[2]

    if file_type == 'json1':
        concat_json1(path)
    elif file_type == 'json2':
        concat_json2(path)
    elif file_type == 'csv':
        if len(sys.argv) >= 4:
            concat_csv(path, sys.argv[3])
        else:
            concat_csv(path)

if __name__ == '__main__':
    main()


class Solution():                                  # 字典方法
    def singleNumber(self, nums):
        count ={}
        for element in nums:
            if element in count:count[element]+=1
            else: count[element]=1
        res = []
        for num, c in count.items():
            if c == 1:
                res.append(num)
        return res


# zignoruj te linijke
from lib import *

wysokosc = 500
szerokosc = 500
okno = stworz_okno_gry(szerokosc, wysokosc)
szerokosc_paletki = 10
wysokosc_paletki = 60

pilka = stworz_pilke()
pilka.ustaw_rozmiary(10, 10)
pilka.przesun_do(szerokosc / 2, wysokosc / 2)
pilka.ustaw_predkosc(2.5)
pilka.zawsze_odbijaj_od_scian()
pilka.zawsze_odbijaj_od_paletek()

lewa_paletka = stworz_paletke()
lewa_paletka.ustaw_rozmiary(szerokosc_paletki, wysokosc_paletki)
lewa_paletka.przesun_do(20, (wysokosc - wysokosc_paletki) / 2)
lewa_paletka.zawsze_odbijaj_od_scian()

prawa_paletka = stworz_paletke()
prawa_paletka.ustaw_rozmiary(szerokosc_paletki, wysokosc_paletki)
prawa_paletka.przesun_do(szerokosc - 20 - szerokosc_paletki, (wysokosc - wysokosc_paletki) / 2)
prawa_paletka.zawsze_odbijaj_od_scian()

jesli(nacisniety_klawisz, gora).to(prawa_paletka.przesun_do_gory)
jesli(nacisniety_klawisz, dol).to(prawa_paletka.przesun_w_dol)
jesli(nacisniety_klawisz, w).to(lewa_paletka.przesun_do_gory)
jesli(nacisniety_klawisz, s).to(lewa_paletka.przesun_w_dol)

jesli(pilka.dotyka_prawej_sciany).to(lewa_paletka.zwieksz_punkty)
jesli(pilka.dotyka_prawej_sciany).to(rozpocznij_nowa_runde)

jesli(pilka.dotyka_lewej_sciany).to(prawa_paletka.zwieksz_punkty)
jesli(pilka.dotyka_lewej_sciany).to(rozpocznij_nowa_runde)

def w_kazdej_klatce():
	pilka.aktualizuj_stan()
	pilka.rysuj(okno)

	lewa_paletka.rysuj(okno)
	lewa_paletka.aktualizuj_stan()

	prawa_paletka.rysuj(okno)
	prawa_paletka.aktualizuj_stan()

	wyswietl_tekst(okno, lewa_paletka.punkty, 30, 30)
	wyswietl_tekst(okno, prawa_paletka.punkty, szerokosc - 50, 30)


import sys
import zipfile
from zipfile import ZipFile
import os
from os.path import basename
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job

import boto3
from botocore.exceptions import ClientError

## @params: [JOB_NAME]
args = getResolvedOptions(sys.argv, ['JOB_NAME'])

sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session
logger = glueContext.get_logger()
job = Job(glueContext)
job.init(args['JOB_NAME'], args)

# Parameters
TARGET_BUCKET='your-bucket-name-with-rw-acces'

# Send file to S3
def upload_file(file_name, bucket, object_name=None):
    """Upload a file to an S3 bucket

    :param file_name: File to upload
    :param bucket: Bucket to upload to
    :param object_name: S3 object name. If not specified then file_name is used
    :return: True if file was uploaded, else False
    """

    # If S3 object_name was not specified, use file_name
    if object_name is None:
        object_name = os.path.basename(file_name)

    # Upload the file
    s3_client = boto3.client('s3')
    try:
        response = s3_client.upload_file(file_name, bucket, object_name)
    except ClientError as e:
        logger.error(e)
        return False
    return True

def zipfolder(zipFileName, target_dir):
    zipobj = zipfile.ZipFile(zipFileName, 'w', compression=zipfile.ZIP_DEFLATED)
    rootlen = len(target_dir) + 1
    for base, dirs, files in os.walk(target_dir):
        for file in files:
            fn = os.path.join(base, file)
            zipobj.write(fn, fn[rootlen:])

def env():
    logger.info('Env:')
    e = ""
    for k, v in sorted(os.environ.items()):
        e = e + f'\n{k} : {v}'
    logger.info(e)

def content():
    logger.info('Content:')
    out = ""
    for root, dirs, files in os.walk("/opt/amazon"):
        path = root.split(os.sep)
        out = out + f'\n{(len(path) - 1) * "---"} {os.path.basename(root)}'
        for file in files:
            out = out + f'\n{len(path) * "---"} {file}'
    logger.info(out)

content()
env()
zipfolder('/tmp/glue3.zip', '/opt/amazon')
upload_file('/tmp/glue3.zip', TARGET_BUCKET)

job.commit()


#!/usr/bin/python
# mk-abstract-book.py
#
# input: AKN preregistration data : TSV format

import pandas as pd
import copy
import re
import math
import sys

if (len(sys.argv) != 2):
  print("usage: ./mk-abstract-book.py <csv-filename>")
  exit()

#--------------------
# CONFIG
#--------------------
year = 2022
filename = sys.argv[1]
lastname = "Last name (in English)"    # column title for lastname
firstname = "First name (in English)"  # column title for firsname
#--------------------

'''

  functions

'''
def print_header (year):

  print("
AKN Abstracts "+str(year)+"
")
  print("
")
  print("
	Name	&nbsp	 Affiliation	 Abst	 Abst	 Abst	 Abst	 Abst	  Theme
")

def print_footer ():

  print("

")

def td_print (txt):

  print(" "+str(txt)+" 	")

def td_print_url (tag,url):

  print(' '+tag+' 	')

def capitalize (name):
  '''
  quick hack to properly capitalize names
  '''

  # 0. Only if all caps
  if len(re.findall("[a-z]",name))>0:
    return name

  dash_loc = name.find('-')

  # 1. two words, and "-" in the middle
  words = re.split('\s|-',name)

  ret = ""
  for w in words:
    ret = ret+w.capitalize()+" "
  ret = ret[:-1]

  # put back dash
  if dash_loc>0:
    ret = ret[:dash_loc] + "-" + ret[dash_loc+1:]

  # 2. parenthesized nicknames
  words = re.split('\(',ret)

  if (len(words)>1):
    ret = ""
    n = 0
    for w in words:
      if n==0:
        ret = ret+w+"("
      else:
        ret = ret+w.capitalize()
      n += 1

  return ret

'''

  1. load CSV file

'''
df = pd.read_csv(filename)
df = df.sort_values(by=['Last name (in English)','First name (in English)'],ascending=True)

'''

  2. get header for abstract columns

'''

cols = df.columns
abst_cols = copy.deepcopy(cols)
abst_cols = abst_cols.tolist()

for c in cols:
  if (not re.search("bstract",c)):
    abst_cols.remove(c)

print_header(year)

'''

  3. print name, affiliation, and abstracts

'''
n = 1
for idx in df.index:

  print('
')

  td_print(str(n)+' ')
  n+=1

  # 1. Name
  td_print( capitalize(df['Last name (in English)'][idx]) )
  td_print( capitalize(df['First name (in English)'][idx]) )

  # 2. Affiliation
  td_print(df['Affiliation'][idx])

  # 3. Abstracts
  for i in range(0,5):
    url = df[abst_cols[i]][idx]
    if (not isinstance(url,str) and math.isnan(url)):
      td_print('   ')
    else:
      if (re.search("abstractsonline",url) and re.search("presentation", url)):
        td_print_url('Abst'+str(i+1),url)
      else:
        td_print_url('invalid url',url)

  # 4. Theme
  td_print(df['Research area'][idx])

  # endline
  print('
')

print_footer()


#!/usr/bin/env python3
__author__ = "Matthew Ng"

import serial
import numpy as np
import time
import cv2
import os
from matplotlib import pyplot as plt
from imutils.object_detection import non_max_suppression
import imutils
import csv

def intersect(p1_si, p2_si):
   if len(p1_si) != 2 and len(p2_si) != 2:
      print("Incorrect slope intercept form")
      return
   b = p2_si[1] - p1_si[1]
   s = p1_si[0] - p2_si[0]
   x = b/s
   y = p1_si[0]*x + p1_si[1]
   return (x, y)

def writeListToCSV(data):
   with open("eric.csv", 'w') as myfile:
      wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)
      wr.writerow([int(x) for x in data])

csvlist = []

cap = cv2.VideoCapture(0)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('output.avi', fourcc, 20.0, (320, 240))
#hog oink
hog = cv2.HOGDescriptor()
hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

while (cap.isOpened()):
   #image testing
   #frame = cv2.imread('hall4.jpg')
   ret, frame = cap.read()

   (rects, weights) = hog.detectMultiScale(frame, winStride=(8,8),padding=(16,16),scale=1.05)
   rects = np.array([[x/2,y/2,(x+w)/2,(y+h)/2] for (x,y,w,h) in rects])
   pick = non_max_suppression(rects, probs=None, overlapThresh=0.65)
   pick = []
   frame = cv2.resize(frame, (0,0), fx=0.5, fy=0.5)
   height, width, channels = frame.shape
   frame = cv2.GaussianBlur(frame, (3,3),0)
   slopeIntercept = []
   newIntersects = []
   lineIntersects = []

   #our operations on the frame come here
   gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
   edges = cv2.Canny(gray, 50, 150)
   for(x1,y1,x2,y2) in pick:
      cv2.rectangle(edges, (x1,y1), (x2,y2), 0,thickness=-1)
   lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi/180, threshold=(50 - len(pick)*3), minLineLength=50, maxLineGap=30)

   for xa,ya,xb,yb in pick:
      lines = [[(x1,y1,x2,y2) for x1,y1,x2,y2 in line if not ((xa < x1 and x1 < xb and xa < x2 and x2 < xb)and(ya<y1 and y1<yb and ya<y2 and y2<yb)) ] for line in lines]

   for line in lines:
      for x1,y1,x2,y2 in line:
         cv2.line(frame, (x1,y1), (x2,y2), (0,255,0),2)
         if x2-x1 == 0:
            continue
         slope = (y2-y1)/(x2-x1)
         intercept = y1 - x1*slope
         slopeIntercept.append([slope,intercept])

   #align the horizon line
   #slopeIntercept.append([0, height/2])
   cv2.line(frame, (0, round(height/2)), (width, round(height/2)), (0, 0, 255), 2)

   slopeIntercept = [x for x in slopeIntercept if (abs(x[0]) > 0.1 and abs(x[0]) < 9)] #eliminate lines close to zero
   [cv2.line(frame, (0, int(x[1])), (1000, int(x[0]*1000 + x[1])), (0,0,255), 2) for x in slopeIntercept]

   #calculate intersects
   for i, si1 in enumerate(slopeIntercept):
      for si2 in slopeIntercept[i+1:]:
         if abs(si1[0] - si2[0]) > 0.5:
            lineIntersects.append(intersect(si1, si2))

   #gets rid of intersects in the top 1/3 or bottom 1/3 of the image
   lineIntersects = [li for li in lineIntersects if (li[1] > height/3 and li[1] < height*2/3)]

   [cv2.circle(frame, (int(intersect[0]), int(intersect[1])), 10, (255, 0, 0), -1) for intersect in lineIntersects]

   numIntersects = len(lineIntersects)
   if numIntersects != 0:
      sum_x = sum([x[0] for x in lineIntersects])
      sum_y = sum([y[1] for y in lineIntersects])
      avg_x = int(sum_x/numIntersects)
      avg_y = int(sum_y/numIntersects)

      color_x = avg_x if avg_x < width else width
      color_x = color_x if color_x > 0 else 0

      #print(int(255/width*color_x))
      speed = chr(int(255/width*color_x))

      color_x = 2*abs(width/2 - color_x)
      #this is where the avg circle is drawn
      cv2.circle(frame, (avg_x, avg_y), 10, (0, 255, 0), -1)
      csvlist.append(avg_x)

   for(x1,y1,x2,y2) in pick:
      cv2.rectangle(frame, (x1,y1), (x2,y2), (25, 25, 25),thickness=2)

   if cv2.waitKey(1) & 0xFF == ord('q'):
      break
   #display the resulting frame
   cv2.imshow('frame', frame)
   out.write(frame)

writeListToCSV(csvlist)
cap.release()
out.release()
cv2.destroyAllWindows()


from transformers import Trainer, TrainingArguments
from tqdm import tqdm
import pandas as pd
import numpy as np
import torch, os

def train(model, train_dataset, val_dataset, save_dir):
    training_args = TrainingArguments(
        output_dir=save_dir,
        num_train_epochs=3,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        warmup_steps=500,
        weight_decay=0.01,
        save_steps = 10000,
        save_total_limit=1,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset
    )

    print('[Fine-tuning] Training..')
    trainer.train()

    trainer.save_model(save_dir)
    print('[Fine-tuning] Saved trained model at {}'.format(save_dir))

def inference(model, input_ids, attention_mask):
    model.eval()
    with torch.no_grad():
        try: outputs = model(input_ids, attention_mask=attention_mask, return_dict=True)
        except: return None
    logits = outputs['logits'].detach().cpu().numpy()
    predicted_label = np.argmax(logits)
    return predicted_label


from vtool.maya_lib import rigs
from vtool.maya_lib import space
from vtool.maya_lib import attr
from vtool.maya_lib import geo

def main():

    size = put.size
    put.control_arm_fk = {}

    for side in 'LR':

        joints = put.joint_arm[side]
        rig = rigs.FkRig('fk_arm',side)
        rig.set_joints(joints)

        rig.set_control_offset_axis('Y')
        rig.set_control_size(size*10)
        rig.set_buffer(True)
        rig.set_control_shape('cylinder')
        rig.set_control_set([side, 'arm_%s' % side])

        if side == 'L':
            rig.set_control_color_hue(.65)
            rig.set_control_color_increment_hue(-0.02)
        if side == 'R':
            rig.set_control_color_hue(1)
            rig.set_control_color_increment_hue(-0.04)

        rig.set_control_parent(put.control_spine[-1])
        rig.set_setup_parent(put.group_setup)
        rig.create()

        cmds.parentConstraint(put.joint_spine[-1], rig.control_group, mo = True)

        put.control_arm_fk[side] = rig.controls

        attr.hide_attributes(rig.controls[1],['rotateY','rotateZ'])
        attr.hide_translate(rig.controls[1])


from cloudscheduler.lib.db_config import Config

rowid_dict = {}
attr_list_dict = {}

def map_attribute_names(src, dest, attribute_names):
    """
    Example:
        from cloudscheduler.lib.attribute_mapper import map_attribute_names
        column_map, bad_columns = map_attribute_names('ec2_regions', 'csv2', tables[config.ec2_region_and_endpoint_table]['heads'])
        if len(bad_columns) > 1:
            logging.error(<some_message>)
            exit(1)

        for attribute_name in tables[config.ec2_region_and_endpoint_table]['heads']:
            if attribute_name in column_map:
                print("Attribute: %s (ec2_regions) = %s (csv2)" % (attribute_name, column_map[attribute_name]))

    """

    global rowid_dict
    global attr_list_dict

    if isinstance(attribute_names, list):
        attribute_name_list = attribute_names
    else:
        attribute_name_list = [str(attribute_names)]

    if len(rowid_dict) == 0 or len(attr_list_dict) == 0:
        build_mapping_dictionaries(config)

    mapped_dict = {}
    unmapped_list = []
    for attribute_name in attribute_name_list:
        if attribute_name in rowid_dict[src]:
            mapped_dict[attribute_name] = attr_list_dict[dest][rowid_dict[src][attribute_name]]
        else:
            unmapped_list.append(attribute_name)

    return mapped_dict, unmapped_list

def map_attributes(src, dest, attr_dict, config):
    global rowid_dict
    global attr_list_dict
    if len(rowid_dict) == 0 or len(attr_list_dict) == 0:
        build_mapping_dictionaries(config)

    unmapped_keys = []
    mapped_dict = {}
    for key, value in attr_dict.items():
        if key not in rowid_dict[src]:
            unmapped_keys.append((key, value))
        else:
            mapped_dict[attr_list_dict[dest][rowid_dict[src][key]]] = str(value)

    return mapped_dict, unmapped_keys

def build_mapping_dictionaries(config):
    global rowid_dict
    global attr_list_dict

    config.db_open()

    Mappings = "csv2_attribute_mapping"

    rc, msg, mapping_rows = config.db_query(Mappings)
    mapping_names = mapping_rows[0].keys()

    # Generate rowid dict
    for language in mapping_names:
        row_id = 0
        row_dict = {}
        for row in mapping_rows:
            # These ifs will need to be updated every time a new translation language is added
            # since there is no way to use the contents of a variable as a variable name reliably
            # (possible with exec())
            row_dict[row[language]] = row_id
            '''
            if language == "csv2":
                row_dict[row["csv2"]] = row_id
            elif language == "os_limits":
                row_dict[row["os_limits"]] = row_id
            elif language == "os_flavors":
                row_dict[row["os_flavors"]] = row_id
            elif language == "os_images":
                row_dict[row["os_images"]] = row_id
            elif language == "os_networks":
                row_dict[row["os_networks"]] = row_id
            elif language == "os_vms":
                row_dict[row["os_vms"]] = row_id
            elif language == "os_sec_grps":
                row_dict[row["os_sec_grps"]] = row_id
            elif language == "condor":
                row_dict[row["condor"]] = row_id
            elif language == "ec2_flavors":
                row_dict[row["ec2_flavors"]] = row_id
            elif language == "ec2_limits":
                row_dict[row["ec2_limits"]] = row_id
            elif language == "ec2_networks":
                row_dict[row["ec2_limits"]] = row_id
            elif language == "ec2_regions":
                row_dict[row["ec2_regions"]] = row_id
            else:
                print("Found column not implemented in code, breaking")
                break
            '''
            row_id = row_id + 1
        rowid_dict[language] = row_dict

    # Generate attribute list dict  (mapping_names are the columns)
    for language in mapping_names:
        attr_list = []
        for row in mapping_rows:
            # These ifs will need to be updated every time a new translation language is added
            # since there is no way to use the contents of a variable as a variable name reliably
            # (possible with exec())
            if row[language] is not None:
                attr_list.append(row[language])
            '''
            if language == "csv2":
                attr_list.append(row["csv2"])
            elif language == "os_limits":
                attr_list.append(row["os_limits"])
            elif language == "os_flavors":
                attr_list.append(row["os_flavors"])
            elif language == "os_images":
                attr_list.append(row["os_images"])
            elif language == "os_networks":
                attr_list.append(row["os_networks"])
            elif language == "os_vms":
                attr_list.append(row["os_vms"])
            elif language == "os_sec_grps":
                attr_list.append(row["os_sec_grps"])
            elif language == "condor":
                attr_list.append(row["condor"])
            elif language == "ec2_flavors":
                attr_list.append(row["ec2_flavors"])
            elif language == "ec2_limits":
                attr_list.append(row["ec2_limits"])
            elif language == "ec2_networks":
                attr_list.append(row["ec2_networks"])
            elif language == "ec2_regions":
                attr_list.append(row["ec2_regions"])
            else:
                print("Found column not implemented in code, breaking")
                break
            '''
        attr_list_dict[language] = attr_list

    return True

def dump_dicts():
    global attr_list_dict
    global rowid_dict

    print(attr_list_dict)
    print(rowid_dict)

if __name__ == "__main__":
    #build_mapping_dictionaries()
    #dump_dicts()
    src = "os_limits"
    dest = "csv2"
    a_dict = {
        "maxTotalCores": 4
    }
    trans_dict = map_attributes(src=src, dest=dest, attr_dict=a_dict)

    print("Input dict:")
    print(a_dict)
    print("Translated dict:")
    print(trans_dict)
    print()

    a_dict = {
        "maxServerMeta": 500,
        "maxTotalRAMSize": 12000
    }
    trans_dict = map_attributes(src=src, dest=dest, attr_dict=a_dict)

    print("Input dict:")
    print(a_dict)
    print("Translated dict:")
    print(trans_dict)
    print()


from distutils.core import setup
from distutils.command.install import INSTALL_SCHEMES
import os

root = os.path.dirname(os.path.abspath(__file__))
os.chdir(root)

VERSION = '1.1'

# Make data go to the right place.
# http://groups.google.com/group/comp.lang.python/browse_thread/thread/35ec7b2fed36eaec/2105ee4d9e8042cb
for scheme in INSTALL_SCHEMES.values():
    scheme['data'] = scheme['purelib']

setup(
    name='django-location-field',
    version=VERSION,
    description="Location field for Django",
    long_description="This module provides a location field for Django applications.",
    author="Codasus Technologies",
    author_email="contact@codasus.com",
    url="http://github.com/codasus/django-location-field",
    license="MIT License",
    platforms=["any"],
    packages=['location_field'],
    #data_files=[(template_dir, templates)],
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Environment :: Web Environment",
        "Intended Audience :: Developers",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
        "Programming Language :: Python",
        "Framework :: Django",
    ],
    include_package_data=True,
)


# This Python file uses the following encoding: ascii
# This module does all dirty work for the program.
import json
import re
from datetime import datetime

import requests
import telethon
import schedule
import telebot
from bs4 import BeautifulSoup
from telethon.sync import TelegramClient
from telethon.tl.functions.messages import GetHistoryRequest

import src.config as c
from modules import functions as f
from classes.database import Database
# from bot_main import main as bot_func

tg_client = TelegramClient(c.username, c.api_id, c.api_hash)
# tg_client.start()

async def dump_comments(channel, post_id):
    # Error processing if channel comments turned off
    try:
        all_messages = []
        async for message in tg_client.iter_messages(channel, reply_to=post_id, reverse=True):
            if isinstance(message.sender, telethon.tl.types.User):
                mess_date = message.date.timestamp()
                comm_length = c.comments_length - 5
                message_text = message.text[:comm_length] + '...' if len(message.text) > comm_length else message.text
                new_rec = [mess_date, message.sender.username, message.sender.first_name, message_text]
            else:
                new_rec = None
            all_messages.append(new_rec)

        for i, mess in enumerate(all_messages):
            if all_messages[i] is None:
                del all_messages[i]

    except telethon.errors.rpcerrorlist.MsgIdInvalidError:
        all_messages = None

    return all_messages

async def dump_post_info(channel, post_id):
    all_messages = []

    class DateTimeEncoder(json.JSONEncoder):
        def default(self, o):
            if isinstance(o, datetime):
                return o.isoformat()
            if isinstance(o, bytes):
                return list(o)
            return json.JSONEncoder.default(self, o)

    history = await tg_client(GetHistoryRequest(
        peer=channel,
        offset_id=(post_id + 1),
        offset_date=None,
        add_offset=0,
        limit=1, max_id=0, min_id=0,
        hash=0))

    messages = history.messages

    for message in messages:
        all_messages.append(message.to_dict())

    # Check if the ID received match the ID asked
    if all_messages[0]['id'] != post_id:
        all_data = False
    else:
        publication_date = all_messages[0]['date']
        views = all_messages[0]['views']
        forwards = all_messages[0]['forwards']

        publication_date = publication_date.timestamp()

        all_data = {'publication_date': publication_date,
                    'views': views,
                    'forwards': forwards}

    return all_data

def dump_subscribers_count(channel_name: str):
    """Takes channel name and returns the number of its subscribers."""
    url = f'https://tgstat.ru/channel/@{channel_name}'
    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.9; rv:45.0) Gecko/20100101 Firefox/45.0'
    }
    r = requests.get(url, headers=headers)
    r = r.text

    try:
        # Extract raw subscribers count
        soup = BeautifulSoup(r, 'html.parser').find('h2').get_text()

        # Convert subscribers count to int
        subscribers_count = re.findall(r'\d+', soup)
        subscribers_count = int(''.join(subscribers_count))
    except AttributeError:
        subscribers_count = 0

    return subscribers_count

async def collecting_data(post: list):
    """Takes post_id and flag if it is first time."""

    # Create channel link and channel object
    channel_link = f'https://t.me/{post[0][1]}'
    channel = await tg_client.get_entity(channel_link)
    channel_post_id = post[0][2]

    # Dump comments
    all_comments = await dump_comments(channel, channel_post_id)

    # Dump post info
    all_data = await dump_post_info(channel, channel_post_id)
    if not all_data:
        f.track_status(post[0], to_archive=True)

    if post[1]:
        all_data['subscribers_count'] = dump_subscribers_count(post[0][1])

    return all_data, all_comments

def put_data_to_posts(post_id, data):
    query = 'UPDATE posts SET publication_date = %s, subscribers_count = %s WHERE post_id = %s'
    params = (data['publication_date'], data['subscribers_count'], post_id)

    with Database() as db:
        db.update_query(query, params)

def put_data_to_stats(post_id, data):
    # Take last added data from base
    query = 'SELECT views, forwards FROM stats WHERE post_id = %s'
    params = (post_id,)

    with Database() as db:
        old_data = db.select_query(query, params)

    # Put data into stats
    if len(old_data) > 0:
        if old_data[-1][0] == data['views'] and old_data[-1][1] == data['forwards']:
            pass
    else:
        date_now = datetime.now().timestamp()
        query = 'INSERT INTO stats (post_id, date, views, forwards) VALUES (%s, %s, %s, %s)'
        params = (post_id, date_now, data['views'], data['forwards'])

        with Database() as db:
            db.update_query(query, params)

def put_data_to_comments(post_id, comments, is_first_collecting: bool):
    comments_to_load = []
    if not is_first_collecting:

        # Load all previous comments from the base
        query = 'SELECT comment_date FROM comments WHERE post_id = %s'
        params = (post_id,)

        with Database() as db:
            old_comments = db.select_query(query, params)

        # Compare old and new comment sets
        old_comments_1 = set()
        [old_comments_1.add(comm[0]) for comm in old_comments]

        new_comments = set()
        [new_comments.add(comm[0]) for comm in comments]
        diff_set = new_comments - old_comments_1

        # Update list of comments to load
        for comm in comments:
            if comm[0] in diff_set:
                comm.insert(0, post_id)
                comm.append('1')
                comm = tuple(comm)
                comments_to_load.append(comm)

    else:
        comments_to_load = []
        for comm in comments:
            comm.insert(0, post_id)
            comm.append('1')
            comm = tuple(comm)
            comments_to_load.append(comm)

    # Put comments into comments
    query = 'INSERT INTO comments (post_id, comment_date, author_username, author, comment_text, is_new) ' \
            'VALUES (%s, %s, %s, %s, %s, %s)'

    with Database() as db:
        db.update_many_query(query, comments_to_load)

def put_data_tobase(post_id, data, comments, is_first_collecting: bool):
    if is_first_collecting:
        put_data_to_posts(post_id, data)

    put_data_to_stats(post_id, data)

    if comments is not None:
        put_data_to_comments(post_id, comments, is_first_collecting)

def get_active_post_list():
    # Get post attributes from the base
    query = 'SELECT post_id, channel_name, channel_post_id, publication_date FROM posts WHERE is_archive = 0'
    with Database() as db:
        post_info = db.select_query(query)

    full_data = []
    for post in post_info:
        full_post = []
        is_first_collecting = False if post[3] is not None else True
        full_post.append(post)
        full_post.append(is_first_collecting)
        full_data.append(full_post)

    return full_data

def send_comments():
    # Create bot connection
    bot = telebot.TeleBot(c.bot_token)
    bot.delete_webhook()

    # Main activity
    result = f.prepare_comments()
    if result.flag:
        for user in result.list_of_users:
            bot.send_message(user, result.notification_message, reply_markup=result.list_of_buttons)

def main_main_main():

    with tg_client:
        tg_client.loop.run_until_complete(main())
    send_comments()
    print(f'Done at {datetime.now().strftime("%H:%M")}')

def with_tg_client(main):

    def another_main():
        def main_main_main():

            with tg_client:
                tg_client.loop.run_until_complete(main())
            send_comments()
            print(f'Done at {datetime.now().strftime("%H:%M")}')

        # schedule.every(15).minutes.do(main_main_main)
        schedule.every(15).minutes.do(main_main_main)
        while True:
            schedule.run_pending()

    return another_main

@with_tg_client
async def main():
    print('Start backend working (internal message)')

    # Get all active posts from the base
    full_data = get_active_post_list()

    for post in full_data:
        data, comments = await collecting_data(post)
        if not data:
            continue
        put_data_tobase(post[0][0], data, comments, post[1])

if __name__ == '__main__':
    print('Start backend working!')
    main()


import ixmp
import numpy as np
import numpy.testing as npt
import pandas as pd
import pandas.testing as pdt
import pytest

from message_ix import Scenario
from message_ix.testing import SCENARIO, make_dantzig

@pytest.fixture
def dantzig_message_scenario(message_test_mp):
    yield Scenario(message_test_mp, **SCENARIO["dantzig"])

def test_year_int(test_mp):
    scen = make_dantzig(test_mp, solve=True, multi_year=True)

    # Dimensions indexed by 'year' are returned as integers for all item types
    assert scen.set("cat_year").dtypes["year"] == "int"
    assert scen.par("demand").dtypes["year"] == "int"
    assert scen.par("bound_activity_up").dtypes["year_act"] == "int"
    assert scen.var("ACT").dtypes["year_vtg"] == "int"
    assert scen.equ("COMMODITY_BALANCE_GT").dtypes["year"] == "int"

def test_add_spatial_single(test_mp):
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    data = {"country": "Austria"}
    scen.add_spatial_sets(data)

    exp = ["World", "Austria"]
    obs = scen.set("node")
    npt.assert_array_equal(obs, exp)

    exp = ["World", "global", "country"]
    obs = scen.set("lvl_spatial")
    npt.assert_array_equal(obs, exp)

    exp = [["country", "Austria", "World"]]
    obs = scen.set("map_spatial_hierarchy")
    npt.assert_array_equal(obs, exp)

def test_add_spatial_multiple(test_mp):
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    data = {"country": ["Austria", "Germany"]}
    scen.add_spatial_sets(data)

    exp = ["World", "Austria", "Germany"]
    obs = scen.set("node")
    npt.assert_array_equal(obs, exp)

    exp = ["World", "global", "country"]
    obs = scen.set("lvl_spatial")
    npt.assert_array_equal(obs, exp)

    exp = [["country", "Austria", "World"], ["country", "Germany", "World"]]
    obs = scen.set("map_spatial_hierarchy")
    npt.assert_array_equal(obs, exp)

def test_add_spatial_hierarchy(test_mp):
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    data = {"country": {"Austria": {"state": ["Vienna", "Lower Austria"]}}}
    scen.add_spatial_sets(data)

    exp = ["World", "Vienna", "Lower Austria", "Austria"]
    obs = scen.set("node")
    npt.assert_array_equal(obs, exp)

    exp = ["World", "global", "state", "country"]
    obs = scen.set("lvl_spatial")
    npt.assert_array_equal(obs, exp)

    exp = [
        ["state", "Vienna", "Austria"],
        ["state", "Lower Austria", "Austria"],
        ["country", "Austria", "World"],
    ]
    obs = scen.set("map_spatial_hierarchy")
    npt.assert_array_equal(obs, exp)

@pytest.mark.parametrize(
    "args, kwargs, exp",
    [
        # Two periods of duration 20, 1 each of duration 10 and 15
        (
            ([2020, 2030, 2050, 2070, 2085],),
            dict(),
            {
                "year": [2020, 2030, 2050, 2070, 2085],
                "fmy": [2020],
                # 20 is chosen for the first period
                "dp": [20, 10, 20, 20, 15],
            },
        ),
        # Mix of positional and keyword arguments; firstmodelyear given
        (
            ([2020, 2030, 2040, 2060],),
            dict(firstmodelyear=2030),
            {
                "year": [2020, 2030, 2040, 2060],
                # firstmodelyear as selected
                "fmy": [2030],
                # 10 is chosen for the first period
                "dp": [10, 10, 10, 20],
            },
        ),
        # Years out of order
        (
            ([2030, 2010, 2020],),
            dict(),
            {
                "year": [2010, 2020, 2030],
                "fmy": [2010],
                "dp": [10, 10, 10],
            },
        ),
        # Deprecated usage with a dict as the first positional argument
        (
            (dict(year=[2010, 2020]),),
            dict(),
            {"year": [2010, 2020], "fmy": [2010], "dp": [10, 10]},
        ),
        (
            (dict(year=[2010, 2020], firstmodelyear=2020),),
            dict(),
            {"year": [2010, 2020], "fmy": [2020], "dp": [10, 10]},
        ),
        # Deprecated usage with user errors
        pytest.param(
            (dict(firstmodelyear=2010),),
            dict(),
            None,
            marks=pytest.mark.xfail(raises=ValueError),
        ),
        pytest.param(
            (dict(year=[2010, 2020], firstmodelyear=2010),),
            dict(firstmodelyear=2020),
            None,
            marks=pytest.mark.xfail(raises=ValueError),
        ),
        pytest.param(
            (dict(year=[2010, 2020], foo="bar"),),
            dict(),
            None,
            marks=pytest.mark.xfail(raises=ValueError),
        ),
        pytest.param(
            (dict(year=[2010, 2020]),),
            dict(data=dict(year=[2010, 2020])),
            None,
            marks=pytest.mark.xfail(raises=ValueError),
        ),
    ],
)
def test_add_horizon(test_mp, args, kwargs, exp):
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")

    # Call completes successfully
    if isinstance(args[0], dict):
        with pytest.warns(
            DeprecationWarning,
            match=(
                r"dict\(\) argument to add_horizon\(\); use year= and "
                "firstmodelyear="
            ),
        ):
            scen.add_horizon(*args, **kwargs)
    else:
        scen.add_horizon(*args, **kwargs)

    # Sets and parameters have the expected contents
    npt.assert_array_equal(exp["year"], scen.set("year"))
    npt.assert_array_equal(exp["fmy"], scen.cat("year", "firstmodelyear"))
    npt.assert_array_equal(exp["dp"], scen.par("duration_period")["value"])

def test_add_horizon_repeat(test_mp, caplog):
    """add_horizon() does not handle scenarios with existing years."""
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")

    # Create a base scenario
    scen.add_horizon([2010, 2020, 2030])
    npt.assert_array_equal([10, 10, 10], scen.par("duration_period")["value"])

    with pytest.raises(
        ValueError,
        match=r"Scenario has year=\[2010, 2020, 2030\] and related values",
    ):
        scen.add_horizon([2015, 2020, 2025], firstmodelyear=2010)

def test_cat_all(dantzig_message_scenario):
    scen = dantzig_message_scenario
    df = scen.cat("technology", "all")
    npt.assert_array_equal(
        df, ["canning_plant", "transport_from_seattle", "transport_from_san-diego"]
    )

def test_cat_list(test_mp):
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")

    # cat_list() returns default 'year' categories in a new message_ix.Scenario
    exp = ["firstmodelyear", "lastmodelyear", "initializeyear_macro"]
    assert exp == scen.cat_list("year")

def test_add_cat(dantzig_message_scenario):
    scen = dantzig_message_scenario
    scen2 = scen.clone(keep_solution=False)
    scen2.check_out()
    scen2.add_cat(
        "technology", "trade", ["transport_from_san-diego", "transport_from_seattle"]
    )
    df = scen2.cat("technology", "trade")
    npt.assert_array_equal(df, ["transport_from_san-diego", "transport_from_seattle"])
    scen2.discard_changes()

def test_add_cat_unique(message_test_mp):
    scen = Scenario(message_test_mp, **SCENARIO["dantzig multi-year"])
    scen2 = scen.clone(keep_solution=False)
    scen2.check_out()
    scen2.add_cat("year", "firstmodelyear", 1963, True)
    assert [1963] == scen2.cat("year", "firstmodelyear")

def test_years_active(test_mp):
    test_mp.add_unit("year")
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    scen.add_set("node", "foo")
    scen.add_set("technology", "bar")

    # Periods of uneven length
    years = [1990, 1995, 2000, 2005, 2010, 2020, 2030]

    # First period length is immaterial
    duration = [1900, 5, 5, 5, 5, 10, 10]
    scen.add_horizon(year=years, firstmodelyear=years[-1])
    scen.add_par(
        "duration_period", pd.DataFrame(zip(years, duration), columns=["year", "value"])
    )

    # 'bar' built in period '1995' with 25-year lifetime:
    # - is constructed in 1991-01-01.
    # - by 1995-12-31, has operated 5 years.
    # - operates until 2015-12-31. This is within the period '2020'.
    scen.add_par(
        "technical_lifetime",
        pd.DataFrame(
            dict(
                node_loc="foo",
                technology="bar",
                unit="year",
                value=25,
                year_vtg=years[1],
            ),
            index=[0],
        ),
    )

    result = scen.years_active("foo", "bar", years[1])

    # Correct return type
    assert isinstance(result, list)
    assert isinstance(result[0], int)

    # Years 1995 through 2020
    npt.assert_array_equal(result, years[1:-1])

def test_years_active_extend(message_test_mp):
    scen = Scenario(message_test_mp, **SCENARIO["dantzig multi-year"])

    # Existing time horizon
    years = [1963, 1964, 1965]
    result = scen.years_active("seattle", "canning_plant", years[1])
    npt.assert_array_equal(result, years[1:])

    # Add years to the scenario
    years.extend([1993, 1995])
    scen.check_out()
    scen.add_set("year", years[-2:])
    scen.add_par("duration_period", "1993", 28, "y")
    scen.add_par("duration_period", "1995", 2, "y")

    # technical_lifetime of seattle/canning_plant/1964 is 30 years.
    # - constructed in 1964-01-01.
    # - by 1964-12-31, has operated 1 year.
    # - by 1965-12-31, has operated 2 years.
    # - operates until 1993-12-31.
    # - is NOT active within the period '1995' (1994-01-01 to 1995-12-31)
    result = scen.years_active("seattle", "canning_plant", 1964)
    npt.assert_array_equal(result, years[1:-1])

def test_years_active_extended2(test_mp):
    test_mp.add_unit("year")
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    scen.add_set("node", "foo")
    scen.add_set("technology", "bar")

    # Periods of uneven length
    years = [1990, 1995, 2000, 2005, 2010, 2020, 2030]

    # First period length is immaterial
    duration = [1900, 5, 5, 5, 5, 10, 10]
    scen.add_horizon(year=years, firstmodelyear=years[-1])
    scen.add_par(
        "duration_period", pd.DataFrame(zip(years, duration), columns=["year", "value"])
    )

    # 'bar' built in period '2020' with 10-year lifetime:
    # - is constructed in 2011-01-01.
    # - by 2020-12-31, has operated 10 years.
    # - operates until 2020-12-31. This is within the period '2020'.
    # The test ensures that the correct lifetime value is retrieved,
    # i.e. the lifetime for the vintage 2020.
    scen.add_par(
        "technical_lifetime",
        pd.DataFrame(
            dict(
                node_loc="foo",
                technology="bar",
                unit="year",
                value=[20, 20, 20, 20, 20, 10, 10],
                year_vtg=years,
            ),
        ),
    )

    result = scen.years_active("foo", "bar", years[-2])

    # Correct return type
    assert isinstance(result, list)
    assert isinstance(result[0], int)

    # Years 2020
    npt.assert_array_equal(result, years[-2])

def test_years_active_extend3(test_mp):
    test_mp.add_unit("year")
    scen = Scenario(test_mp, **SCENARIO["dantzig"], version="new")
    scen.add_set("node", "foo")
    scen.add_set("technology", "bar")

    # Periods of uneven length
    years = [1990, 1995, 2000, 2005, 2010, 2020, 2030]

    scen.add_horizon(year=years, firstmodelyear=2010)

    scen.add_set("year", [1992])
    scen.add_par("duration_period", "1992", 2, "y")
    scen.add_par("duration_period", "1995", 3, "y")

    scen.add_par(
        "technical_lifetime",
        pd.DataFrame(
            dict(
                node_loc="foo",
                technology="bar",
                unit="year",
                value=[20],
                year_vtg=1990,
            ),
        ),
    )

    obs = scen.years_active("foo", "bar", 1990)

    assert obs == [1990, 1992, 1995, 2000, 2005]

def test_new_timeseries_long_name64(message_test_mp):
    scen = Scenario(message_test_mp, **SCENARIO["dantzig multi-year"])
    scen = scen.clone(keep_solution=False)
    scen.check_out(timeseries_only=True)
    df = pd.DataFrame(
        {
            "region": [
                "India",
            ],
            "variable": [
                ("Emissions|CO2|Energy|Demand|Transportation|Aviation|" "Domestic|Fre"),
            ],
            "unit": [
                "Mt CO2/yr",
            ],
            "2012": [
                0.257009,
            ],
        }
    )
    scen.add_timeseries(df)
    scen.commit("importing a testing timeseries")

def test_new_timeseries_long_name64plus(message_test_mp):
    scen = Scenario(message_test_mp, **SCENARIO["dantzig multi-year"])
    scen = scen.clone(keep_solution=False)
    scen.check_out(timeseries_only=True)
    df = pd.DataFrame(
        {
            "region": [
                "India",
            ],
            "variable": [
                (
                    "Emissions|CO2|Energy|Demand|Transportation|Aviation|"
                    "Domestic|Freight|Oil"
                ),
            ],
            "unit": [
                "Mt CO2/yr",
            ],
            "2012": [
                0.257009,
            ],
        }
    )
    scen.add_timeseries(df)
    scen.commit("importing a testing timeseries")

def test_rename_technology(dantzig_message_scenario):
    scen = dantzig_message_scenario
    assert scen.par("output")["technology"].isin(["canning_plant"]).any()

    clone = scen.clone("foo", "bar")
    clone.rename("technology", {"canning_plant": "foo_bar"})
    assert not clone.par("output")["technology"].isin(["canning_plant"]).any()
    assert clone.par("output")["technology"].isin(["foo_bar"]).any()
    clone.solve(quiet=True)
    assert np.isclose(clone.var("OBJ")["lvl"], 153.675)

def test_rename_technology_no_rm(dantzig_message_scenario):
    scen = dantzig_message_scenario
    assert scen.par("output")["technology"].isin(["canning_plant"]).any()

    clone = scen.clone("foo", "bar")
    # also test if already checked out
    clone.check_out()

    clone.rename("technology", {"canning_plant": "foo_bar"}, keep=True)
    assert clone.par("output")["technology"].isin(["canning_plant"]).any()
    assert clone.par("output")["technology"].isin(["foo_bar"]).any()

def test_excel_read_write(message_test_mp, tmp_path):
    # Path to temporary file
    tmp_path /= "excel_read_write.xlsx"
    # Convert to string to ensure this can be handled
    fname = str(tmp_path)

    scen1 = Scenario(message_test_mp, **SCENARIO["dantzig"])
    scen1 = scen1.clone(keep_solution=False)
    scen1.check_out()
    scen1.init_set("new_set")
    scen1.add_set("new_set", "member")
    scen1.init_par("new_par", idx_sets=["new_set"])
    scen1.add_par("new_par", "member", 2, "-")
    scen1.commit("new set and parameter added.")

    # Writing to Excel without solving
    scen1.to_excel(fname)

    # Writing to Excel when scenario has a solution
    scen1.solve(quiet=True)
    scen1.to_excel(fname)

    scen2 = Scenario(message_test_mp, model="foo", scenario="bar", version="new")

    # Fails without init_items=True
    with pytest.raises(ValueError, match="no set 'new_set'"):
        scen2.read_excel(fname)

    # Succeeds with init_items=True
    scen2.read_excel(fname, init_items=True, commit_steps=True)

    exp = scen1.par("input")
    obs = scen2.par("input")
    pdt.assert_frame_equal(exp, obs)

    assert scen2.has_par("new_par")
    assert 2 == scen2.par("new_par").at[0, "value"]

    scen2.solve(quiet=True)
    assert np.isclose(scen2.var("OBJ")["lvl"], scen1.var("OBJ")["lvl"])

def test_clone(tmpdir):
    # Two local platforms
    mp1 = ixmp.Platform(driver="hsqldb", path=tmpdir / "mp1")
    mp2 = ixmp.Platform(driver="hsqldb", path=tmpdir / "mp2")

    # A minimal scenario
    scen1 = Scenario(mp1, model="model", scenario="scenario", version="new")
    scen1.add_spatial_sets({"country": "Austria"})
    scen1.add_set("technology", "bar")
    scen1.add_horizon(year=[2010, 2020])
    scen1.commit("add minimal sets for testing")

    assert len(mp1.scenario_list(default=False)) == 1

    # Clone
    scen2 = scen1.clone(platform=mp2)

    # Return type of ixmp.Scenario.clone is message_ix.Scenario
    assert isinstance(scen2, Scenario)

    # Close and re-open both databases
    mp1.close_db()  # TODO this should be done automatically on del
    mp2.close_db()  # TODO this should be done automatically on del
    del mp1, mp2
    mp1 = ixmp.Platform(driver="hsqldb", path=tmpdir / "mp1")
    mp2 = ixmp.Platform(driver="hsqldb", path=tmpdir / "mp2")

    # Same scenarios present in each database
    assert all(mp1.scenario_list(default=False) == mp2.scenario_list(default=False))

    # Load both scenarios
    scen1 = Scenario(mp1, "model", "scenario")
    scen2 = Scenario(mp2, "model", "scenario")

    # Contents are identical
    assert all(scen1.set("node") == scen2.set("node"))
    assert all(scen1.set("year") == scen2.set("year"))


"""Worker impl module."""
from io import BytesIO
from typing import List

import cv2
import numpy as np
from imutils.video import FileVideoStream
from loguru import logger
from PIL import Image

from server.base.Message import BaseFailed
from server.base.Worker import BaseWorker
from server.engine.detection.Message import (
    Box,
    DetectionInputType,
    DetectionResult,
    DetectionTask,
    FrameData,
    PatchData,
)
from server.engine.detection.Model import DetectionModel
from server.util.types import EngineDevice
from server.util.util import create_timestamp

class DetectionWorker(BaseWorker):
    __model: DetectionModel

    __batch_size: int

    def __init__(
        self, device: EngineDevice = None, batch_size: int = 32, queue_size: int = 100
    ):
        self.__batch_size = batch_size
        if device is None:
            device = EngineDevice()

        super().__init__(queue_size=queue_size, device=device)

    def init(self):
        self.__model = DetectionModel(device=self.device)
        logger.info(f"[{self.__class__.__name__}] worker ready")

    def __transform(self, frame):
        if frame is None:
            return None
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        return frame

    def __process_video(self, video_path: str) -> List[FrameData]:

        frame_indices: List[int] = []
        frames: List[np.ndarray] = []

        stream = FileVideoStream(path=video_path, transform=self.__transform).start()

        index = 0
        while stream.running():
            frame = stream.read()  # HWC
            if frame is None:
                continue
            frames.append(frame)
            frame_indices.append(index)
            index += 1

        results: List[FrameData] = [
            FrameData(frame_index=i, patches=[]) for i in range(index)
        ]

        target_length = len(frame_indices)
        for start in range(0, target_length, self.__batch_size):
            end = min(start + self.__batch_size, target_length)
            batch = frames[start:end]

            ts_start = create_timestamp()
            indices, boxes, class_ids, scores, class_names = self.__model(frames=batch)
            ts_end = create_timestamp()

            logger.info(
                f"[{__class__.__name__}] Batch size {len(batch)}: {ts_end - ts_start} ms"
            )

            for index, box, class_id, score, class_name in zip(
                indices, boxes, class_ids, scores, class_names
            ):
                patch = PatchData(
                    box=Box(x1=box[0], y1=box[1], x2=box[2], y2=box[3]),
                    class_id=class_id,
                    score=score,
                    class_name=class_name,
                )
                results[start + index].patches.append(patch)

        return results

    def __process_image(self, bimages: List[bytes]) -> List[FrameData]:
        batch = []
        for bimage in bimages:
            image = np.asarray(Image.open(BytesIO(bimage)).convert("RGB"))  # HWC
            batch.append(image)

        num_images = len(batch)

        ts_start = create_timestamp()
        indices, boxes, class_ids, scores, class_names = self.__model(frames=batch)
        ts_end = create_timestamp()

        logger.info(
            f"[{__class__.__name__}] Batch size {len(batch)}: {ts_end - ts_start} ms"
        )

        results: List[FrameData] = [
            FrameData(frame_index=i, patches=[]) for i in range(num_images)
        ]

        for index, box, class_id, score, class_name in zip(
            indices, boxes, class_ids, scores, class_names
        ):
            patch = PatchData(
                box=Box(x1=box[0], y1=box[1], x2=box[2], y2=box[3]),
                class_id=class_id,
                score=score,
                class_name=class_name,
            )
            results[index].patches.append(patch)
        return results

    def process(self, task: DetectionTask):
        if task.input_type == DetectionInputType.IMAGE:
            ret = self.__process_image(task.image_data)
            self.sink.put(DetectionResult(data=ret, task_id=task.id))
        elif task.input_type == DetectionInputType.VIDEO:
            ret = self.__process_video(task.video_path)
            self.sink.put(DetectionResult(data=ret, task_id=task.id))
        else:
            raise RuntimeError(f"Invalid input type: {task.input_type}")

    def handle(self, msg: DetectionTask) -> DetectionResult:
        try:
            self.process(msg)
        except Exception as e:
            self.sink.put(BaseFailed(task_id=msg.id))
            raise RuntimeError(f"Failed to process task `{msg.id}`: {e}")