# core modules
import random
import math
import gym
import numpy as np
from gym import spaces
import tensorflow as tf
from itertools import chain, combinations
# import matplotlib.pyplot as plt
from collections import deque
import numpy as np
import matlab
# import matplotlib.pyplot as plt
import matlab.engine
from mat4py import loadmat, savemat
import logging
import h5py

# import subprocess
import os

LINE_VOLTAGE_MAX_LIMIT = 900
LINE_VOLTAGE_MIN_LIMIT = 650
LINE_CURRENT_MAX_LIMIT = 50
LINE_CURRENT_MIN_LIMIT = 2

LINE_CURRENT_RATED = 10
LINE_VOLTAGE_RATED = 750

SAMPLE_TIME = 0.01
TOTAL_SIM_TIME = 8
TOTAL_TIME_STEPS = int(TOTAL_SIM_TIME / SAMPLE_TIME)
REWARD_AMOUNT_COMMON = 100
"""

Custom implementation of openAI Gym environment API ( https://gym.openai.com/docs )

-- > Multi-agent system with each controller as agent

input state : output of each controller's voltage and current
output state : Increase or decrease del_V to the primary droop controller ( Ops of neural network is equivalent to getting all other V and I values and averaging them, passing them through a PI controller and summing them )
obs dim ( dim of input )        :   num_controllers * 3 ( contimuous values )
action_dim ( dim of output )    :   num_controllers * 2 ( since action space is for the whole env and not for single agent ) ( discrete values )

each step runs the simulation till the current time step.
V and I for all agents are taken
only the local V and I are input to each neural network ( no communication )

xval is created as a timeseries to inpout to the sim in reqd format

input to the NN [ state ]   :   ( timestep, v, i )  //
output [ action ]           :   del_V               //  [0 or 1]  ( inc or dec )]

reward is calculated based on voltage and current levels

"""

class Agent():

    def __init__(self, agent_num):

        self.silent = True
        self.v_level_failed= False
        self.i_level_failed= False
        self.vl = 750
        self.il = 10
        self.del_v = 0
        self.del_i = 0
        self.state = np.asarray([0, LINE_VOLTAGE_RATED / LINE_VOLTAGE_RATED, LINE_CURRENT_RATED / LINE_CURRENT_RATED])
        self.input = np.zeros(TOTAL_TIME_STEPS)
        self.input_timeseries = np.stack([np.around(np.linspace(0, 8, TOTAL_TIME_STEPS), 2), self.input],axis=1)

        # if os.path.exists(self.input_file_name):
        #     os.remove(self.input_file_name)
        # else:
        #     logging.debug("not able to remove file...")

        # with open(self.input_file_name, 'w') as f : pass # .close()
        self.input_file_name = "D:\\maddpg_sec_conrtoller\\custom_microgrid\\input_" + str(agent_num) + ".mat"
        savemat(self.input_file_name, {'input_'+str(agent_num) : np.transpose(self.input_timeseries).tolist()})
        # self.check_input_file = h5py.File(self.input_file_name, "r")
        # self.check_input = self.check_input_file.get('input_' + str(agent_num))
        # self.check_input = np.array(self.check_input)
        self.check_input = loadmat(self.input_file_name)
        self.cost = 0
        self.reward = 0

    def reset(self, agent_num):
        self.state = np.asarray([0, LINE_VOLTAGE_RATED / LINE_VOLTAGE_RATED, LINE_CURRENT_RATED / LINE_CURRENT_RATED])
        self.v_level_failed= False
        self.i_level_failed= False
        self.vl = 750
        self.il = 10
        self.del_v = 0
        self.del_i = 0
        self.input = np.zeros(TOTAL_TIME_STEPS)
        self.input_timeseries = np.stack([np.around(np.linspace(0, 8, TOTAL_TIME_STEPS), 2), self.input],axis=1)
        savemat(self.input_file_name, {'input_'+str(agent_num) : np.transpose(self.input_timeseries).tolist()})
        self.check_input = loadmat(self.input_file_name)
        # open(self.input_file_name, 'w').close()

        self.cost = 0
        self.reward = 0

class Make_World():
    def __init__(self, num_agents = 3):
        self.agents = []
        self.dim_phasors = 3    # may come into use later
        self.dim_c = 0          # may come into use later

        logging.debug('init microgrid .... !!')

        for i in range(num_agents):
            logging.debug("agent_ {0}".format (i))

            self.agents.append(Agent(agent_num = i))

    def get_agents(self):
        return self.agents

class SecondaryController(gym.Env) :

    def __init__(self, num_agents = 3,  eng_name = None):

        self.num_agents = num_agents
        self.n = num_agents

        self.action_space = []
        self.observation_space = []
        self.discrete_action_space = True
        self.world = Make_World(num_agents=self.num_agents)
        self.agents = self.world.get_agents()
        logging.debug(len(self.agents))
        self.curr_episode = 0

        self.penalty = 10
        self.del_v = 0
        self.del_i = 0
        self.delta_v_change = 10 # (V) for changes in voltage to be multiplied to NN output
        # self.delta_i_change = 2 # (A)

        self.prev_len = 0
        self.full_signal = []

        self.tim = np.linspace(0, 10, num=TOTAL_TIME_STEPS)
        # self.xval = np.ones((1000, 2), dtype=float)

        # ------ action spaces init ------------
        action_dim = 3 # inc, dec and no change
        for agent in self.agents:
            total_action_space = []
            # physical action space
            if self.discrete_action_space:
                # it can either increase or decrease voltage in steps. Dim = 2
                u_action_space = spaces.Discrete(action_dim) # self.world.dim_phasors *  2
                logging.debug("discrete {0}".format( u_action_space))
                # comm action space : for future work
                c_action_space = spaces.Discrete(self.world.dim_c)
            else:
                # can remove since it is not needed
                u_action_space = spaces.Box(low=LINE_VOLTAGE_MIN_LIMIT, high=LINE_VOLTAGE_MAX_LIMIT, shape=(2,2)) # ,dtype=np.float32
                logging.debug("absas {0}".format(u_action_space))
                c_action_space = spaces.Box(low=0.0, high=1.0, shape=(self.world.dim_c,)) # ,dtype=np.float32

            total_action_space.append(u_action_space)
            logging.debug("asas {0}".format( total_action_space))

            if not agent.silent:
                total_action_space.append(c_action_space)

            # Total action space
            # only for comm case
            # not needed
            if len(total_action_space) > 1:
                # all action spaces are discrete, so simplify to MultiDiscrete action space
                if all([isinstance(act_space, spaces.Discrete) for act_space in total_action_space]):
                    act_space = spaces.MultiDiscrete([[0,act_space.n-1] for act_space in total_action_space])
                else:
                    act_space = spaces.Tuple(total_action_space)
                self.action_space.append(act_space)
            else:
                self.action_space.append(total_action_space[0])

            # -------------------------------------- removed print -----------------------------
            # print(np.asarray(self.action_space).shape)
            # ----------------------------------------------------------------

            # Is this Correct?
            # Voltage and current are observed
            obs_dim = 3
            # should spaces be limited to rated values?
            self.observation_space.append(spaces.Box(low=-np.inf, high=+np.inf, shape=(obs_dim,))) # ,dtype=np.float32
            # ------------------------------------------
            # agent.action.c = np.zeros(self.world.dim_c)

        logging.debug("len act sp : {0}".format(len(self.action_space)))
        logging.debug("len obs sp : {0}".format(len(self.observation_space)))

        # ---------- matlab engine init -------------
        if eng_name is None:
            names = matlab.engine.find_matlab()
            logging.debug(names)
            self.eng = matlab.engine.connect_matlab(names[0])

        logging.debug('succesfully connected\n\n')
        logging.debug('starting controls .. \n\n')

    def step(self, action_n):

        obs_n = []
        reward_n = []
        done_n = []
        info_n = {'n': []}

        # check for init state and start simulation:
        if self.curr_step ==0:
            self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','start',nargout=0)

        if self.time_is_over :
            # raise RuntimeError("Episode is done")
            self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','stop',nargout=0)
            logging.debug("Episode done : {0}".format(self.curr_episode))

        # set action for each agent and
        # advance microgrid state
        ## for i, agent in enumerate(self.agents):
        logging.debug("step : {0}".format(self.curr_step))
        self._set_action(action_n)

        # record observation,reward, done and info for each agent
        for k,agent in enumerate(self.agents):
            obs_n.append(self._get_obs(agent))
            reward_n.append(self._get_reward(agent))
            done_n.append(self._get_done())
            info_n['n'].append(self._get_info())

        if not self.time_is_over:
            self.curr_step += 1
        return obs_n, reward_n, done_n, info_n

    def _get_info(self):
        # not implemented
        # for logging
        return {}

    # advance state by one timestep and also give actions to the agents
    # is action_space reqd ?
    def _set_action(self, action_n):

        # logging.debug("len agents : {0}".format(len(self.agents)))
        # logging.debug(len(action_n))
        # generate inputs for each controller chip
        for i, agent in enumerate(self.agents):

            if np.argmax(action_n[i]) == 0 :
                delv = self.delta_v_change
            elif np.argmax(action_n[i]) == 1:
                delv = -self.delta_v_change
            else:
                delv= 0
            logging.debug("delv ----> : {0}".format(agent.del_v))
            if not self.curr_step >= 800: # TOTAL_TIME_STEPS
                # agent.input_timeseries[self.curr_step +1] = [self.tim[self.curr_step +1], agent.del_v]
                agent.input_timeseries[self.curr_step][1] = agent.del_v

            # write to file :
            savemat(agent.input_file_name, {'input' : np.transpose(agent.input_timeseries).tolist()})
            agent.check_input = loadmat(agent.input_file_name)
            logging.debug("checking the last written voltage : {0}".format(agent.check_input))

        # self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','start',nargout=0)
        # and simulate the microgrid
        # while self.eng.get_param('test_v_control_and_i_control_with_droop_copy','SimulationStatus')!=('stopped' or 'terminating'):
        if self.eng.get_param('test_v_control_and_i_control_with_droop_copy','SimulationStatus')=='paused':
            # Trying changing constant value
            for i, agent in enumerate(self.agents):
                if np.argmax(action_n[i]) == 0:
                    curr_ref = int(self.eng.get_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value')) # ,nargout=0)
                    print("0", curr_ref, type(curr_ref))
                    self.eng.set_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value',str(curr_ref+1),nargout=0)
                if np.argmax(action_n[i]) == 1:
                    curr_ref = int(self.eng.get_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value')) # ,nargout=0)
                    print("1", curr_ref, type(curr_ref))
                    self.eng.set_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value',str(curr_ref-1),nargout=0)
                if np.argmax(action_n[i]) == 2:
                    curr_ref = int(self.eng.get_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value')) # ,nargout=0)
                    print("2", curr_ref, type(curr_ref))
                    self.eng.set_param('test_v_control_and_i_control_with_droop_copy/Vref'+str(i),'Value',str(curr_ref),nargout=0)

            # ToDo : operations check
            logging.debug("curr step : {0}".format(self.curr_step))
            self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','update',nargout=0) #if you have updated any simulation parameters

            if self.curr_step >= TOTAL_TIME_STEPS : # * SAMPLE_TIME:
                logging.debug("stopping.. : \t self.curr_step : {1}".format(self.curr_step))
                self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','stop',nargout=0)

            self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','continue',nargout=0)

        else:
            self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','stop',nargout=0)
            logging.debug("stopped *************************** error")

        self.full_signal = self.eng.workspace['V_out_1'].get('signals').get('values')

        if self.curr_step <= 0 :
            # return some safe value and wait till the transient settles
            return 0

        # generate inputs to NN
        for i, agent in enumerate(self.agents):
            agent.vl = np.array(self.eng.workspace['V_out_' + str(i)].get('signals').get('values')[-1], dtype=float)[0]
            logging.debug("Vl : {0}".format(agent.vl))
            agent.il = np.array(self.eng.workspace['I_out_' + str(i)].get('signals').get('values')[-1], dtype=float)[0]
            logging.debug("Il : {0}".format(agent.il))

            # not needed
            # agent.vl = agent.vl + action_n[i,0] * self.delta_v_change
            # agent.il = agent.il + action_n[i,1] * self.delta_i_change
            agent.del_v = float(LINE_VOLTAGE_RATED) - agent.vl
            agent.del_i = float(LINE_CURRENT_RATED) - agent.il

            if  agent.vl  <= LINE_VOLTAGE_MIN_LIMIT and agent.vl >=LINE_VOLTAGE_MAX_LIMIT:
                logging.debug("V failed:")
                agent.v_level_failed = True

            if agent.il  <= LINE_CURRENT_MIN_LIMIT and agent.il >=LINE_CURRENT_MAX_LIMIT:
                logging.debug("I failed:")
                agent.i_level_failed = True
        self.prev_len = len(self.full_signal)

    def reset(self):
        self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','stop',nargout=0)

    def render(self):
        return

    def _get_reward(self, agent):

        if not(agent.v_level_failed) and not(agent.i_level_failed):
            agent.cost = agent.del_v *  self.penalty + agent.del_i * self.penalty
        elif agent.v_level_failed and not(agent.i_level_failed):
            agent.cost = 100 + agent.del_i * self.penalty
        elif agent.i_level_failed and not(agent.v_level_failed):
            agent.cost = agent.del_v * self.penalty + 100
        else:
            agent.cost = 200
        agent.reward = -agent.cost + REWARD_AMOUNT_COMMON

        return agent.reward

    def _get_obs(self, agent):

        agent.state = np.zeros(3)
        if not self.time_is_over:
            agent.state[0] =  self.curr_step+1
        else:
            agent.state[0] =  0
        agent.state[1] =  agent.vl / LINE_VOLTAGE_RATED
        agent.state[2] =  agent.il / LINE_CURRENT_RATED
        # ---------------------------------------- remove print -------------------------------
        logging.debug("state : {0}".format(agent.state))
        # ---------------------------------------------------------
        return agent.state

    def _get_done(self):

        remaining_steps = TOTAL_TIME_STEPS - self.curr_step+1
        self.time_is_over = (remaining_steps <= 0)

        return self.time_is_over

    def _get_initial_state(self):

        self.curr_episode += 1

        self.time_is_over = False
        self.total_rew = 0
        self.curr_step = 0
        self.prev_len = 0
        self.del_v = 0
        self.del_i = 0

        # self.agents = [Agent() for i in range(self.num_agents) ]
        for agent_num, agent in enumerate(self.agents):
            agent.reset(agent_num)

        initial_state_batch = [agent.state for agent in self.agents]

        return initial_state_batch

        # ---------------------------------- old set action ------------------------------------

    #     def _set_action(self, action_n):

    #     print("len agents : ", len(self.agents))
    #     print(len(action_n))
    #     # generate inputs for each controller chip
    #     for i, agent in enumerate(self.agents):

    #         if np.argmax(action_n[i]) == 0 :
    #             delv = self.delta_v_change
    #         elif np.argmax(action_n[i]) == 1:
    #             delv = -self.delta_v_change
    #         else:
    #             delv= 0

    #         # print(agent.input.shape)
    #         # print(np.full((TOTAL_TIME_STEPS - self.prev_len), delv).shape)
    #         agent.input[self.prev_len : ]=np.full((TOTAL_TIME_STEPS - self.prev_len), delv)

    #         agent.input_timeseries = np.stack([self.tim, agent.input], axis=1)
    #         agent.input_timeseries = matlab.double(agent.input_timeseries.tolist())
    #         self.eng.workspace['input_' + str(i)] = agent.input_timeseries

    #     self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','start',nargout=0)

    #     # and simulate the microgrid
    #     j=0
    #     while self.eng.get_param('test_v_control_and_i_control_with_droop_copy','SimulationStatus')!=('stopped' or 'terminating'):
    #         if self.eng.get_param('test_v_control_and_i_control_with_droop_copy','SimulationStatus')=='paused':
    #             # operations check
    # #             print('%f  : %f' ,prev_len, len(self.eng.workspace['ScopeData'].get('signals').get('values')))
    # #             print(self.eng.workspace['ScopeData'].get('signals').get('values')) # [prev_len:]

    #             # print("j : ", j)
    #             # print("curr step : ", self.curr_step)

    #             self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','update',nargout=0) #if you have updated any simulation parameters

    #             if j >= self.curr_step : # * SAMPLE_TIME:
    #                 print("stopping.. : ", j, " : step : ", self.curr_step)
    #                 self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','stop',nargout=0)

    #             self.eng.set_param('test_v_control_and_i_control_with_droop_copy','SimulationCommand','continue',nargout=0)

    #             j+=1

    #     self.full_signal = self.eng.workspace['V_out_1'].get('signals').get('values')

    #     if self.curr_step <= 0 :
    #         # return some safe value and wait till the transient settles
    #         return 0

    #     # generate inputs to NN
    #     for i, agent in enumerate(self.agents):

    #         agent.vl = np.array(self.eng.workspace['V_out_' + str(i)].get('signals').get('values')[-1], dtype=float)[0]
    #         print("Vl : ", agent.vl)
    #         agent.il = np.array(self.eng.workspace['I_out_' + str(i)].get('signals').get('values')[-1], dtype=float)[0]

    #         print("Il : ", agent.il)

    #         # not needed
    #         # agent.vl = agent.vl + action_n[i,0] * self.delta_v_change
    #         # agent.il = agent.il + action_n[i,1] * self.delta_i_change

    #         agent.del_v = float(LINE_VOLTAGE_RATED) - agent.vl
    #         agent.del_i = float(LINE_CURRENT_RATED) - agent.il

    #         if not (agent.vl  <= LINE_VOLTAGE_MAX_LIMIT) and not (agent.vl >=LINE_VOLTAGE_MIN_LIMIT):
    #             print("V failed:")
    #             agent.v_level_failed = True

    #         if not (agent.il  <= LINE_CURRENT_MAX_LIMIT) and not (agent.il >=LINE_CURRENT_MIN_LIMIT):
    #             print("I failed:")
    #             agent.i_level_failed = True

    #     self.prev_len = len(self.full_signal)


from typing import List

from polyaxon._contexts import paths as ctx_paths

def get_artifacts_store_args(artifacts_path: str, clean: bool) -> str:
    get_or_create = 'if [ ! -d "{dir}" ]; then mkdir -m 0777 -p {dir}; fi;'.format(
        dir=artifacts_path
    )
    delete_dir = (
        'if [ -d {path} ] && [ "$(ls -A {path})" ]; '
        "then rm -R {path}/*; fi;".format(path=artifacts_path)
    )
    if clean:
        return "{} {}".format(get_or_create, delete_dir)
    return "{}".format(get_or_create)

def init_artifact_context_args(run_path: str) -> List[str]:
    return [
        'if [ ! -d "{dir}" ]; then mkdir -m 0777 -p {dir}; fi;'.format(
            dir=ctx_paths.CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(run_path)
        ),
        'if [ ! -d "{dir}" ]; then mkdir -m 0777 -p {dir}; fi;'.format(
            dir=ctx_paths.CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(run_path) + "/outputs"
        ),
    ]


#%%
import datetime
import math
import os
import random

import numpy as np
import pylab as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from matplotlib.patches import Rectangle
from torch.autograd import Variable, grad
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms

from config import args
from masked_puzzle import manmade_blocks, mask, traces
from util import (append_file, batch_generator, gumbel_softmax, save_image,
                  write_file)

#%%
class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Reshape(nn.Module):
    def forward(self, input, input_shape):
        return input.view(input.size(0), input_shape)

class GumbelConvAE(nn.Module):
    def __init__(self, parameters):
        super(GumbelConvAE, self).__init__()
        self.params = parameters
        self.create_encoder(self.params['input_shape'])
        self.create_decoder(self.params['input_shape'])
        self.temperature_setting()

        self.path = self.params['model_path']
        if not os.path.isdir(self.path):
            os.mkdir(self.path)

    def forward(self, input):
        q = self.encoder(input)
        q_y = q.view([q.size(0), self.params['N'], self.params['M']])
        z = gumbel_softmax(q_y, self.temperature)
        z_t = z[:,:,0]  # from latplan
        rec = self.decoder(z_t).view(-1, *self.params['input_shape'])
        return rec, F.softmax(q_y, dim=-1), z_t

    def sub_path(self, path):
        return os.path.join(self.path, path)

    def save(self, path):
        torch.save(self.state_dict(), self.sub_path(path))

    def load(self, path):
        self.load_state_dict(torch.load(self.sub_path(path)))

    def create_encoder(self, input_shape):
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels=2,out_channels=self.params['clayer'], kernel_size=(3,3), bias=False),
            nn.Tanh(),
            nn.Dropout2d(p=self.params['dropout']),
            nn.BatchNorm2d(num_features=self.params['clayer'], eps=1e-3),
            nn.MaxPool2d(kernel_size=(2,2)),

            nn.Conv2d(in_channels=self.params['clayer'],out_channels=self.params['clayer'], kernel_size=(3,3), bias=False),
            nn.Tanh(),
            nn.Dropout2d(p=self.params['dropout']),
            nn.BatchNorm2d(num_features=self.params['clayer'], eps=1e-3),
            nn.MaxPool2d(kernel_size=(2,2)),

            Flatten(),

            nn.Linear(self.params['clayer']*(((input_shape[-1]-2)//2-2)//2)**2, self.params['layer']),
            nn.BatchNorm1d(self.params['layer']),
            nn.Linear(self.params['layer'], self.params['N']*self.params['M'])
        )

    def create_decoder(self, input_shape):
        data_dim = int(np.prod(input_shape))
        self.decoder = nn.Sequential(
            Flatten(),
            *([nn.Dropout(p=self.params['dropout'])] if self.params['dropout_z'] else []),
            nn.Linear(self.params['N'], self.params['layer']),
            nn.ReLU(),
            nn.BatchNorm1d(self.params['layer']),
            nn.Dropout(p=self.params['dropout']),

            nn.Linear(self.params['layer'], self.params['layer']),
            nn.ReLU(),
            nn.BatchNorm1d(self.params['layer']),
            nn.Dropout(p=self.params['dropout']),

            nn.Linear(self.params['layer'], data_dim),
            nn.Sigmoid()
        )

    def temperature_setting(self):
        self.min_temperature = self.params['min_temperature']
        self.max_temperature = self.params['max_temperature']
        self.num_epoch = self.params['num_epoch']
        self.anneal_rate = math.log(self.max_temperature/self.min_temperature) / self.num_epoch
        self.temperature = self.max_temperature

    def set_curr_temp(self, temp):
        self.temperature = temp

    def update_temp(self, epoch):
        self.temperature = np.max([self.min_temperature, self.max_temperature * np.exp(- self.anneal_rate * epoch)])
        return self.temperature

class Decoder(nn.Module):
    def __init__(self, input_dim, output_shape, hidden_dim, dropout, dropout_z):
        super(Decoder, self).__init__()
        self.output_shape = output_shape
        data_dim = int(np.prod(output_shape))
        self.decoder = nn.Sequential(
            Flatten(),
            *([nn.Dropout(p=dropout)] if dropout_z else []),
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_dim),
            nn.Dropout(p=dropout),

            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_dim),
            nn.Dropout(p=dropout),

            nn.Linear(hidden_dim, data_dim),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.decoder(x).view([*x.shape[:-1], *self.output_shape])

# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, qy, m=1, cate_dim=2):
    # import ipdb
    # ipdb.set_trace()
    BCE = F.binary_cross_entropy(recon_x*m, x*m, size_average=False)
    MSE = F.mse_loss(recon_x*m, x*m, size_average=False)
    log_qy = torch.log(qy+1e-20)
    g = Variable(torch.log(torch.tensor([1.0/cate_dim])).cuda())
    KLD = torch.sum(qy*(log_qy - g),dim=-1).sum()
    # KLD = torch.sum(qy*(log_qy),dim=-1).mean()
    return BCE, KLD
    # return MSE

#%%
def train(model, epoch, train_data, batch_size, optimizer, log_interval):
    model.train()
    train_loss = 0
    train_bce_loss = 0
    train_kld_loss = 0
    temp = model.update_temp(epoch)
    train_loader = batch_generator(train_data, batch_size, False)
    for batch_idx, data in train_loader:
        data = Variable(data)
        data = data.cuda()
        optimizer.zero_grad()
        recon_batch, qy, z_t = model(data)
        bce_loss, kld_loss = loss_function(recon_batch, data[:,0], qy, data[:,1])
        loss = bce_loss + kld_loss
        loss.backward()
        train_loss += loss.data.item()
        train_bce_loss += bce_loss.data.item()
        train_kld_loss += kld_loss.data.item()
        optimizer.step()

        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_data),
                100. * batch_idx * batch_size / len(train_data),
                loss.data.item() / len(data)))
        if batch_idx == 1:
            n = min(data.size(0), 8)
            comparison = torch.cat([data[:n,0].view(-1, 1, 42, 42),
                                    F.upsample(qy[:n,:,0].view(-1,1,6,6), size=42),
                                    F.upsample(qy[:n,:,0].view(-1,1,6,6)/temp, size=42),
                                    F.upsample(z_t[:n].view(-1,1,6,6), size=42),
                                    recon_batch.view(-1, 1, 42, 42)[:n]])
            save_image(comparison.data.cpu(),
                        model.sub_path('results/train_' + str(epoch) + '.png'), nrow=n)

    train_loss /= len(train_data)
    train_bce_loss /= len(train_data)
    train_kld_loss /= len(train_data)
    print('====> Epoch: {} Average bce loss: {:.4f}, Average kld loss: {:.4f}'.format(
          epoch, train_bce_loss, train_kld_loss))
    return train_bce_loss, train_kld_loss

def test(model, epoch, test_data, batch_size):
    model.eval()
    test_loss = 0
    test_bce_loss = 0
    test_kld_loss = 0
    temp = model.update_temp(epoch)
    test_loader = batch_generator(test_data, batch_size, False)
    for batch_idx, data in test_loader:
        data = data.cuda()
        data = Variable(data, volatile=True)
        recon_batch, qy, z_t = model(data)
        bce_loss, kld_loss = loss_function(recon_batch, data[:,0], qy, data[:,1])
        test_bce_loss += bce_loss.data.item()
        test_kld_loss += kld_loss.data.item()
        test_loss += test_bce_loss + test_kld_loss
        if batch_idx == 1:
            n = min(data.size(0), 8)
            comparison = torch.cat([data[:n,0].view(-1, 1, 42, 42),
                                  F.upsample(qy[:n,:,0].view(-1,1,6,6), size=42),
                                  F.upsample(qy[:n,:,0].view(-1,1,6,6)/temp, size=42),
                                  F.upsample(z_t[:n].view(-1,1,6,6), size=42),
                                  recon_batch.view(-1, 1, 42, 42)[:n]])
            save_image(comparison.data.cpu(),
                     model.sub_path('results/reconstruction_' + str(epoch) + '.png'), nrow=n)

    test_loss /= len(test_data)
    test_bce_loss /= len(test_data)
    test_kld_loss /= len(test_data)
    print('====> Test set bce loss: {:.4f}, kld loss: {:.4f}'.format(test_bce_loss, test_kld_loss))
    return test_bce_loss, test_kld_loss

#%%
def run():
    parameters = {
        'layer'             :1000,# [400,4000],
        'clayer'            :16,# [400,4000],
        'dropout'           :0.4, #[0.1,0.4],
        'noise'             :0.4,
        'dropout_z'         :False,
        'activation'        :'tanh',
        'epoch'             :300,#150,
        'batch_size'        :4000,#4000,
        'lr'                :0.001,
        'N'                 :36,
        'M'                 :2,
        'input_shape'       :(42,42),
        'num_epoch'         :300,#150,
        'min_temperature'   :0.7,
        'max_temperature'   :5.0,
    }
    parameters['model_path'] = os.path.join(os.getcwd(), 'models', 'GumbelConvAE')

    log_interval = 3

    model = GumbelConvAE(parameters).cuda()
    optimizer = optim.Adam(model.parameters(), lr=parameters['lr'])

    log_file = model.sub_path('record_{}.txt'.format(datetime.datetime.now().strftime('%m%d%H%M')))
    import json
    write_file(json.dumps(parameters), path=log_file)

    trace_num = 3750
    trace_len = 8
    imgs, _ = traces(3, 3, trace_num=trace_num, trace_len=trace_len)
    n_train = int(trace_num*trace_len * 0.85)

    mask_mat = torch.tensor([manmade_blocks(mask_pos=random.sample(range(0,9), 2)) for l in range(trace_len*trace_num)]).view(-1, 1, 42, 42)
    mask_mat[n_train:] = 1
    data = (torch.from_numpy(imgs).view(-1, 1, 42, 42) * mask_mat)
    data_mask = torch.cat([data, mask_mat], dim=1)

    train_data = data_mask[:n_train]
    test_data = data_mask[n_train:]
    #%%
    best_test_loss = None
    for epoch in range(1, parameters['num_epoch'] + 1):
        batch_size = parameters['batch_size']
        train_data=train_data[torch.randperm(train_data.size()[0])]
        test_data=test_data[torch.randperm(test_data.size()[0])]
        train_bce_loss, train_kld_loss = train(model, epoch, train_data, batch_size, optimizer, log_interval)
        test_bce_loss, test_kld_loss = test(model, epoch, test_data, batch_size)
        if best_test_loss is None or best_test_loss > test_loss:
            model.save('model.pkl')
            # model.save('ep_%03d' % epoch)

        loss_info = map(lambda x: '%.6f'%x, [train_bce_loss, train_kld_loss, test_bce_loss, test_kld_loss])
        append_file(loss_info, log_file, num=epoch)

        sample_z = Variable(torch.randn(8, 36).round()).cuda()
        sample_m = mask_mat[:8].view(-1,1,42,42).cuda()
        sample_x = model.decoder(sample_z).view(-1,1,42,42)
        sample_rec, qy, z_t = model(torch.cat([sample_x, sample_m], dim=1))
        sample_rec2, qy2, z_t2 = model(torch.cat([sample_rec.view(-1, 1, 42, 42), sample_m], dim=1))
        sample_rec3, qy3, z_t3 = model(torch.cat([sample_rec2.view(-1, 1, 42, 42), sample_m], dim=1))
        sample = torch.cat([F.upsample(sample_z.view(-1,1,6,6), size=42),
                            sample_x.view(-1, 1, 42, 42),
                            sample_m.view(-1, 1, 42, 42),
                            F.upsample(qy[:,:,0].view(-1,1,6,6), size=42),
                            F.upsample(z_t.view(-1,1,6,6), size=42),
                            sample_rec.view(-1, 1, 42, 42),
                            F.upsample(qy2[:,:,0].view(-1,1,6,6), size=42),
                            F.upsample(z_t2.view(-1,1,6,6), size=42),
                            sample_rec2.view(-1, 1, 42, 42),
                            F.upsample(qy3[:,:,0].view(-1,1,6,6), size=42),
                            F.upsample(z_t3.view(-1,1,6,6), size=42),
                            sample_rec3.view(-1, 1, 42, 42)], dim=0).cpu()
        save_image(sample.data.view(-1, 1, 42, 42),
                    model.sub_path('results/sample_' + str(epoch) + '.png'))

def test_sae(path):

    parameters = {
        'layer'             :1000,# [400,4000],
        'clayer'            :16,# [400,4000],
        'dropout'           :0.4, #[0.1,0.4],
        'noise'             :0.4,
        'dropout_z'         :False,
        'activation'        :'tanh',
        'epoch'             :150,
        'batch_size'        :4000,
        'lr'                :0.001,
        'N'                 :36,
        'M'                 :2,
        'input_shape'       :(42,42),
        'num_epoch'         :150,
        'min_temperature'   :0.7,
        'max_temperature'   :5.0,
    }
    parameters['model_path'] = os.path.join(os.getcwd(), 'models', 'GumbelConvAE')

    ae = GumbelConvAE(parameters)
    ae.load(path)
    sample_z = Variable(torch.randn(8, 36).round())
    sample_z = sample_z.cuda()
    sample_x = ae.decoder(sample_z).view(-1,1,42,42)
    sample_rec = ae(sample_x)[0].view(-1,1,42,42)
    sample_rec2 = ae(sample_rec)[0].view(-1,1,42,42)
    sample_rec3 = ae(sample_rec2)[0].view(-1,1,42,42)
    sample = torch.stack((sample_x, sample_rec, sample_rec2, sample_rec3), dim=0).cpu()
    save_image(sample.data.view(-1, 1, 42, 42),'test.png')

#%%
if __name__ == '__main__':
    run()

# %%


import pygame, sys
from settings import *
from level import Level
from player import Player

# pygame setup ======================================
pygame.init()
screen = pygame.display.set_mode(RES, vsync=1)
clock = pygame.time.Clock()

# Game setup ========================================
background = pygame.image.load('bg_minecraft.jpg').convert()

level = Level()

# Game Loop =========================================
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        # events ===================================

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                level.player.jump()

    # Draw - Update =================================
    screen.blit(background, (0,0))

    level.run()

# pygame setup ======================================
    pygame.display.update()
    clock.tick(60)
    pygame.display.set_caption(f'fps: {clock.get_fps():.0f}          Platformer')


#!/usr/bin/env python3
"""
This module sets up a basic Flask app with a single route.
"""
from typing import Any
from flask import Flask, render_template

app = Flask(__name__)

@app.route('/')
def index() -> Any:
    """
    Render the index.html template.
    """
    return render_template('0-index.html')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)


import os
import numpy as np
import torch
import torchvision
from torch import nn
from torchinfo import summary
import matplotlib.pyplot as plt
from glob import glob

''' SEED Everything '''
from src.utils import seed_everything
seed_everything(seed=42)
print("Using Torch", torch.__version__, 'and Torchvision:', torchvision.__version__)

''' Load training data '''
root_dir        = 'path/to/dataset/dir/'
train_img_dir   = root_dir + '/train_data/images'       # train images inside this dir
train_gt_dir    = root_dir + '/train_data/ground-truth' # train GTs inside this dir
test_img_dir    = root_dir + '/test_data/images'        # test images inside this dir
test_gt_dir     = root_dir + '/test_data/ground-truth'  # test GTs inside this dir
train_img_paths = glob(train_img_dir + '/*.jpg')
test_img_paths  = glob(test_img_dir + '/*.jpg')
train_gt_paths  = glob(train_gt_dir + '/*.mat')
test_gt_paths   = glob(test_gt_dir + '/*.mat')
train_img_paths.sort(key=lambda f: int('_'.join(filter(str.isdigit, f))))
test_img_paths.sort(key=lambda f: int('_'.join(filter(str.isdigit, f))))
train_gt_paths.sort(key=lambda f: int('_'.join(filter(str.isdigit, f))))
test_gt_paths.sort(key=lambda f: int('_'.join(filter(str.isdigit, f))))

''' Create dataset '''
from src.Dataset import CrowdDataset2 as CrowdDataset
img_size   = (800, 800, 3) # size of the input image
gt_size    = (200, 200)  # size of the density map
sigma      = 10 # Gaussian kernal
batch_size = 2
num_workers= 1
train_ds   = CrowdDataset(train_img_paths, train_gt_paths, img_size, gt_size, sigma=sigma, augmentation=True)
test_ds    = CrowdDataset(test_img_paths, test_gt_paths, img_size, gt_size, sigma=sigma, augmentation=False)
train_dl   = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=True)
test_dl    = torch.utils.data.DataLoader(test_ds,  batch_size=1, num_workers=num_workers, pin_memory=True, shuffle=False)

###############  Taking a subset for test (Comment the below three lines to use full dataset)
from torch.utils.data import Subset
indices    = torch.randint(low=0, high=20, size=(32,))
train_ds   = Subset(train_ds, indices)
test_ds    = Subset(test_ds, indices)

from src.models.CSRNet import CSRNet
# torch.backends.cudnn.enabled   = True
# torch.backends.cudnn.benchmark = True

''' Define crowd counting model '''
device          = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
model           = CSRNet().to(device)
model_name      = model.__class__.__name__
optimizer       = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
n_epochs        = 100
weights         = None # Define path to model weights if want to load for finetuning
# weights         = f'{os.getcwd()}/checkpoints/{model_name}_{ds_name}.pth'
checkpoint      = f'{os.getcwd()}/checkpoints/{model_name}.pth'
val_metric_best = np.Inf
params          = {'model': model,
                   'train_dl':train_dl,
                   'val_dl': test_dl,
                   'optimizer': optimizer,
                   'n_epochs': n_epochs,
                   'weights': weights,
                   'checkpoint': checkpoint,
                   'val_metric_best': val_metric_best,
                   'device': device
                   }
print(f'Models will be saved to >>>> {checkpoint}')

''' Start training '''
from src.trainer import Trainer
trainer = Trainer(params)
history = trainer.fit()


# import urllib3
# http = urllib3.PoolManager()
# r = http.request("GET", "https://www.google.com")

# from lxml import html

# data_string = r.data.decode('utf-8', errors = 'ignore')
# tree = html.fromstring(data_string)

# links = tree.xpath("//a")
# for link in links:
#     print(link.get('href'))

import requests
from bs4 import BeautifulSoup

r = requests.get("https://news.ycombinator.com")
soup = BeautifulSoup(r.text, 'html.parser')
links = soup.findAll('tr', class_='athing')

formatted_links = []

for link in links:
    data = {
        'id': link['id'],
        'title': link.find_all('td')[2].a.text,
		"url": link.find_all('td')[2].a['href'],
		"rank": int(links[0].td.span.text.replace('.', ''))
    }
    formatted_links.append(data)

print(formatted_links[0])


# -*- coding: utf-8 -*-

from sys import exit as sys_exit
from typing import Callable, List, Optional

from avplayer.apps.default import default_main
from avplayer.arguments import get_default_arguments
from avplayer.logging.logging import (
    SEVERITY_NAME_DEBUG,
    logger,
    set_colored_formatter_logging_config,
    set_root_level,
    set_simple_logging_config,
)
from avplayer.variables import PRINTER_NAMESPACE_ATTR_KEY

def main(
    cmdline: Optional[List[str]] = None,
    printer: Callable[..., None] = print,
) -> int:
    args = get_default_arguments(cmdline)

    colored_logging = args.colored_logging
    simple_logging = args.simple_logging
    severity = args.severity
    debug = args.debug
    verbose = args.verbose

    assert isinstance(colored_logging, bool)
    assert isinstance(simple_logging, bool)
    assert isinstance(severity, str)
    assert isinstance(debug, bool)
    assert isinstance(verbose, int)

    assert not hasattr(args, PRINTER_NAMESPACE_ATTR_KEY)
    setattr(args, PRINTER_NAMESPACE_ATTR_KEY, printer)
    assert hasattr(args, PRINTER_NAMESPACE_ATTR_KEY)

    if colored_logging:
        set_colored_formatter_logging_config()
    elif simple_logging:
        set_simple_logging_config()

    if debug:
        set_root_level(SEVERITY_NAME_DEBUG)
    else:
        set_root_level(severity)

    logger.debug(f"Parsed arguments: {args}")
    return default_main(args)

if __name__ == "__main__":
    sys_exit(main())


from secondary_functions.data_sheets import data_sheets
from collections import Counter
import time

update_position = []
insert_position = []

def update_document(post_cur, data, USD, new_data):
    global update_position
    try:
        if new_data != None:
            if len(new_data) == len(data):
                for i in new_data:
                    if i in data:
                        pass
                    else:
                        update_position.append(i)  # Список даннах на обновления

            for i in update_position:  # обновление данных
                if len(new_data) == len(data):  # Если списки равны по элементом то это точно update

                    flat = [y for x in new_data for y in x]  # Поиск дубликатов
                    c = Counter(flat)
                    print(c[i[0]])
                    if int(c[i[0]]) > 1:
                        print("Есть дубликат Обновлять не могу ")
                        update_position.remove(i)
                        break

                    item_tuple_update = (
                        i[1], i[2], USD * int(i[2]), str(i[3]), i[0])  # VALUES (%s, %s, %s, %s ,%s)
                    post_cur.update_table(item_tuple_update)
                    update_position = []
        else:
            new_data = []
    except Exception as err:
        print(err)

def insert_and_del_document(post_cur, data, USD, new_data):
    global insert_position
    try:
        if new_data != None:
            if len(new_data) != len(data):
                if len(new_data) > len(data):  # Добавление данных в таблицу
                    for i in new_data:
                        if i in data:
                            pass
                        else:
                            insert_position.append(i)

                    if insert_position == []:  # Проверка на дубликат если это повтор то пишем оповешение и не добовляем его в бд
                        print("Вы написали дубликат!")

                    for i in insert_position:  # Ишем новые значения для добавления
                        item_tuple_insert = (i[0], i[1], i[2], USD * int(i[2]), str(i[3]))
                        post_cur.insert_table(item_tuple_insert)
                        insert_position = []

                elif len(new_data) < len(data):  # Удаление данных
                    for i in data:
                        if i in new_data:
                            pass
                        else:
                            flat = [y for x in new_data for y in x]
                            c = Counter(flat)
                            if int(c[i[
                                0]]) > 0:  # В случаи если есть дубликат то блакируем удалаение пока есть дубликаты
                                print("Вы удалили дубликат")
                                break
                            post_cur.delete_table(i[0])
        else:
            post_cur.delete_table(insert_query="DELETE FROM item;",
                                  item_query=())  # если в таблицк 0 значений то полностью чистим бд
    except Exception as err:
        print(err)

def run_full_def_app(post_cur, data, USD):

    while True:
        new_data = data_sheets()
        update_document(post_cur, data, USD, new_data)
        insert_and_del_document(post_cur, data, USD, new_data)
        data = new_data
        time.sleep(3)


from os import listdir
from os.path import isfile, join
import numpy as np
import cv2
import os
import sigpy as sp
#import cupy as cp
import sigpy.plot as plt
import sigpy.mri as spmri
from PIL import Image
from matplotlib import cm

# pylint: disable=unnecessary-semicolon

def cartisian2D( traj_shape, densitySAxis = (1.,1.), densityScale = 1., dtype=np.float):
    '''
    img_shape : tuple of two ints
        shape of the resultant trajctory
    densitySAxis :  tuple of two ints
        density in of sample in y and x axis respectivly per densityScale
    densityScale : float
        controlles the dinsity in each aixs
    dtype:   numpy dtype
        type of the output
    '''
    imgDimY, imgDimX = traj_shape;
    imgDimY = imgDimY/2;
    imgDimX = imgDimX/2;
    denistyY, densityX = densitySAxis;
    if(densityScale>0):
        factor1 =  denistyY / densityScale;
        factor2 =  densityX / densityScale;
    else:
        #{
        print("WARNING! density scale should be bigger than zero.");
        return None;
        #}
    #m1 , m2 = np.meshgrid(np.arange(-imgDimX, imgDimX + factor2, factor2), np.arange(-imgDimY, imgDimY + factor1, factor1));
    m1 , m2 = np.meshgrid(np.arange(-imgDimX, imgDimX , factor2), np.arange(-imgDimY, imgDimY , factor1));

    y, x = m2.shape;
    mesh = np.zeros(( y,x,2));

    mesh[:,:,0] = m1;
    mesh[:,:,1] = m2;
    return mesh.astype(np.float);

name  = 'img.jpg';
image = Image.open(name).convert('L');

arr = np.array(image) + 1j;
traj  = cartisian2D(arr.shape,[1,1],1);
plt.ScatterPlot(traj, title='Trajectory');
image.close()
arr=arr/np.max(arr[...]);

print(traj.shape)

kspaceNUFFT = sp.nufft(arr, traj);

plt.ImagePlot(np.log(kspaceNUFFT), title='k-space data from NUFFT');

kspaceFFT  = sp.fft(arr);

plt.ImagePlot(np.log(kspaceFFT), title='k-space data from FFT')
print(kspaceFFT.shape);
print(kspaceNUFFT.shape);
sumNUFFT = np.sum(kspaceNUFFT);
sumFFT  = np.sum(kspaceFFT);
if(np.allclose(kspaceNUFFT,kspaceFFT,rtol= 10,atol = 10) and np.isclose(sumNUFFT,sumFFT,rtol= 50,atol = 50)):
    print('Outputs are  similar!')
else:
    print('Outputs are NOT similar!')


# -*- coding: utf-8 -*-
import os
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow_addons as tfa
from PIL import Image

class Dataset:
    """
    Class handling tf_dataset creation for a segmentation dataset from two
    directories, containing 2D images and corresponding labels as png.

    All images in image_dir have to have a label with the same filename in
    label_dir. Images which do not have a corresponding label file in
    label_dir are assumed to be test images.
    This class can split these images in train/val folds and optionally
    normalise them. The way categories are calculated from label images can
    be changed in the function _load_label_images.

    Usage: An example on how to use this class can be found in experiment.py
    """
    def __init__(self, image_dir, label_dir, val_ratio=0.2,
                 cross_validation=False, n_folds=None, normalise=False,
                 ignore=None):
        """
        :param image_dir: string, directory where images can be found
        :param label_dir: string, directory where labels can be found -
        labels are expected to have the same name as corresponding image.
        :param val_ratio: float, ratio of samples with existing groundtruth to
        be used as validation set
        :param cross_validation: bool, whether to calculate crossvalidation
        folds
        :param n_folds: int, number of folds for crossvalidation
        :param normalise: whether to normalise images when loading
        :param ignore: list, image names to be ignored (the corresponding labels
        will automatically be ignored)
        """
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.val_ratio = val_ratio
        self.cross_validation = cross_validation
        self.n_folds = n_folds

        self.img_shape = None
        self.ignore = ignore or []
        self.normalise = normalise
        self.split_df = self._calculate_data_split(self.val_ratio,
                                                   self.cross_validation,
                                                   self.n_folds)

    def _calculate_data_split(self, val_ratio=0.2, cross_validation=False,
                              n_folds=None):
        """
        Split dataset into purpose (train/val/test) or crossvalidation folds
        and test.

        :param val_ratio:  ratio of samples with existing groundtruth to be used
        as validation set, ignored if cross_validation=True
        :param cross_validation: whether to calculate crossvalidation folds
        :param n_folds: number of folds for crossvalidation
        :return: dataframe with entries for each sample and its corresponding
        purpose / fold.
        """
        if not cross_validation:
            if not (0 <= val_ratio <= 1):
                raise ValueError('Please choose a val_ratio between 0 and 1.')
        else:
            if type(n_folds) is not int:
                raise TypeError('Only integers are allowed for n_folds.')

        # fix seed to ensure reproducible results
        np.random.seed(1234)

        img_list = list(os.walk(self.image_dir))[0][2]
        label_list = list(os.walk(self.label_dir))[0][2]

        # remove samples to be ignored
        for name in self.ignore:
            img_list.remove(name)
            label_list.remove(name)

        # test images are missing corresponding labels with same name as image
        is_test = np.array([img not in label_list for img in img_list])

        split_df = pd.DataFrame(data={'img_name': img_list})
        test_idx = np.where(is_test)[0]
        split_df.loc[test_idx, 'purpose'] = 'test'

        if not cross_validation:
            self.n_val = int(val_ratio * len(label_list))
            self.n_train = len(label_list) - self.n_val
            # indices available for training and validation
            train_val_idx = np.where(~is_test)[0]
            # choose training indices randomly
            train_idx = np.random.choice(train_val_idx, self.n_train,
                                         replace=False)
            # take remaining indices for validation
            val_idx = np.setdiff1d(train_val_idx, train_idx)
            split_df.loc[train_idx, 'purpose'] = 'train'
            split_df.loc[val_idx, 'purpose'] = 'val'
        else:
            # choose indices for folds of crossvalidation
            n_train_samples = len(split_df) - len(test_idx)
            self.n_sample_per_fold = int(n_train_samples / n_folds)
            split_df.loc[:, 'fold'] = None
            for i in range(n_folds-1):
                unassigned_idx = np.where(split_df.fold.isnull())[0]
                available_idx = np.setdiff1d(unassigned_idx, test_idx)
                # choose fold indices randomly from remaining indices
                fold_idx = np.random.choice(available_idx,
                                            self.n_sample_per_fold,
                                            replace=False)
                split_df.loc[fold_idx, 'fold'] = i

            # take remaining samples for last fold
            unassigned_idx = np.where(split_df.fold.isnull())[0]
            remaining_idx = np.setdiff1d(unassigned_idx, test_idx)
            split_df.loc[remaining_idx, 'fold'] = i + 1

        return split_df

    def create_tf_dataset(self, purpose=None, cross_validation=False, fold=None,
                          return_stats=False):
        """
        Create a tensorflow dataset object, containing the images specified in
        purpose/fold.

        :param purpose: string, 'train'/'val'/'test'
        :param cross_validation: bool, whether to get crossvalidation folds
        :param fold: int, fold for crossvalidation
        :param return_stats: bool, whether to return class_ratios as list
        :return: tensorflow dataset (, class_ratio)
        """
        np_images = self._load_images(purpose, cross_validation, fold) / 255
        np_labels, class_ratio = self._load_labels(purpose,
                                                   cross_validation,
                                                   fold,
                                                   return_stats=True)

        tf_images = tf.convert_to_tensor(np_images, dtype=tf.float32)
        tf_labels = tf.convert_to_tensor(np_labels, dtype=tf.float32)

        tf_set = tf.data.Dataset.from_tensor_slices({'image': tf_images,
                                                     'label': tf_labels
                                                     })

        if return_stats:
            return tf_set, class_ratio
        else:
            return tf_set

    def _load_images(self, purpose=None, cross_validation=False, fold=None):
        """
        Load all images with specified properties as one numpy array.

        For documentation on properties see self._filter_df.
        :param purpose: string, 'train'/'val'/'test'
        :param cross_validation: bool, whether to get crossvalidation folds
        :param fold: int, fold for crossvalidation
        :return: numpy array, shape (n_sample, x, y, n_channel)
        """
        filtered_df = self._filter_df(purpose, cross_validation, fold)
        images = []
        for path in filtered_df.img_name:
            img = Image.open(os.path.join(self.image_dir, path)).convert('RGB')
            if self.normalise:
                img = tf.keras.utils.normalize(img)
            images.append(img)

        stack = np.stack(images)
        self.img_shape = stack.shape[1:]

        return stack

    def _load_labels(self, purpose=None, cross_validation=False, fold=None,
                     return_stats=False):
        """
        Load all labels with specified properties as one numpy array.

        For documentation on properties see self._filter_df.
        :param purpose: string,  'train'/'val'/'test'
        :param cross_validation: bool, whether to get crossvalidation folds
        :param fold: int, fold for crossvalidation
        :param return_stats: bool, whether to return ratio of different labels
        over all loaded data
        :return: numpy array, shape (n_samples, x, y, 1)
        or tuple of numpy array, class_ratios if return_stats is True
        """
        if purpose == 'test':
            raise AttributeError('Test images have no corresponding label.')

        filtered_df = self._filter_df(purpose, cross_validation, fold)
        labels = []
        for path in filtered_df.img_name:
            img_path = os.path.join(self.label_dir, path)
            label_img = Image.open(img_path).convert(mode='L')
            label = np.array(label_img)
            # assume three classes: pixel == 0, pixel == 2, and everything else
            label_mask = np.zeros_like(label)
            label_mask[(0 < label) & (label < 255)] = 1
            label_mask[label == 255] = 2
            label_mask = np.expand_dims(label_mask.astype('int32'), axis=-1)
            labels.append(label_mask)

        label_stack = np.stack(labels)

        if not return_stats:
            return label_stack

        class_count = np.bincount(label_stack.astype('int').flatten())
        class_ratio = class_count / np.size(label_stack)
        return label_stack, class_ratio

    def _filter_df(self, purpose=None, cross_validation=False, fold=None):
        """
        Filter dataframe with all samples for specified properties.

        Purpose is ignored, when cross_validation=True
        :param purpose: one of 'train'/'val'/'test'
        :param cross_validation: whether to use cross validation folds
        :param fold: crossvalidation fold to filter for
        :return: dataframe containing only filtered entries
        """
        if cross_validation:
            assert type(fold) is int
            filtered_df = self.split_df[self.split_df.fold == fold]
        else:
            assert type(purpose) is str
            filtered_df = self.split_df[self.split_df.purpose == purpose]
        return filtered_df

class DataAugmenter:
    """
    Class handling data augmentation for a tf_dataset.

    Usage: An example on how to use this class can be found in experiment.py
    """
    def __init__(self, flip_left_right=0.5, flip_up_down=0.5,
                 brightness=True, gaussian=0.2, gaussian_std=0.03,
                 rotate_prob=0.2, rotate_range=90):
        """
        :param flip_left_right: float, probability for doing a left-right flip
        for data augmentation
        :param flip_up_down: float, probability for doing a up_down flip for
        data augmentation
        :param brightness: bool, whether to do brightness augmentation on image
        :param gaussian: float, probability for adding gaussian noise to image
        as data augmentation
        :param gaussian_std: float, standard deviation for gaussian noise,
        only considered if gaussian > 0
        :param rotate_prob: float, probability for rotating data as augmentation
        :param rotate_range: float, if rotate_prob > 0, data will be rotated by
        a random angle in [-rotate_range, rotate_range] in degree.
        """
        self.flip_left_right = flip_left_right
        self.flip_up_down = flip_up_down
        self.brightness = brightness
        self.gaussian = gaussian
        self.gaussian_std = gaussian_std
        self.rotate_prob = rotate_prob
        self.rotate_range = rotate_range

    @tf.function
    def augment_train_images(self, datapoint):
        """
        Augment a datapoint consisting of an image and corresponding label.

        Augmentation functions are specified for class.
        :param datapoint: dictionary for one datapoint. Needs keys 'image' and
        'label'. Images and labels can be numpy arrays or tf.Tensor objects.
        :return: tuple augmented image, augmented label
        """
        # take arguments as class functions
        img = datapoint['image']
        label = datapoint['label']

        if self.flip_left_right:
            do_flip = tf.random.uniform([]) < self.flip_left_right
            img = tf.cond(do_flip,
                          lambda: tf.image.flip_left_right(img),
                          lambda: img)
            label = tf.cond(do_flip,
                            lambda: tf.image.flip_left_right(label),
                            lambda: label)

        if self.flip_up_down:
            do_flip = tf.random.uniform([]) < self.flip_up_down
            img = tf.cond(do_flip,
                          lambda: tf.image.flip_up_down(img),
                          lambda: img)
            label = tf.cond(do_flip,
                            lambda: tf.image.flip_up_down(label),
                            lambda: label)

        if self.brightness:
            do_bright = tf.random.uniform([]) < self.brightness
            img = tf.cond(do_bright,
                          lambda: tf.image.random_brightness(img,
                                                             max_delta=0.1),
                          lambda: img)

        if self.gaussian:
            do_noise = tf.random.uniform([]) < self.gaussian
            std = tf.random.normal(shape=[1], stddev=self.gaussian_std)
            noise = tf.random.normal(shape=tf.shape(img), stddev=std)
            img = tf.cond(do_noise, lambda: tf.add(img, noise), lambda: img)
            img = tf.clip_by_value(img, 0.0, 1.0)

        if self.rotate_prob > 0:
            angle = tf.random.uniform([],
                                      minval=-self.rotate_range,
                                      maxval=self.rotate_range)
            do_rotate = tf.random.uniform([]) < self.rotate_prob
            img = tf.cond(do_rotate,
                          lambda: tfa.image.rotate(img, angle,
                                                   interpolation='BILINEAR'),
                          lambda: img)
            label = tf.cond(do_rotate,
                            lambda: tfa.image.rotate(label, angle,
                                                     interpolation='NEAREST'),
                            lambda: label)

        return img, label

    @tf.function
    def load_validation_images(self, datapoint):
        """
        Load a datapoint without any augmentations.

        :param datapoint: dictionary for one datapoint. Needs keys 'image' and
        'label'. Images and labels can be numpy arrays or tf.Tensor objects.
        :return: tuple image, label
        """
        return datapoint['image'], datapoint['label']


#coding:utf-8
import re
import time
import queue
import threading

import chardet

from model.book.book import Book

class Book_parser(object):
    def __init__(self):
        # 读取文件, user_id=None, book_path=None
        # 创建队列
        self.Q = queue.Queue()
        def loop():
            while True:
                if self.Q.empty():
                    time.sleep(5)
                else:
                    self.parser(book_info=self.Q.get())

        threading.Thread(target=loop).start()

    def parser(self,book_info=None):
        # 检查文本中的章卷节等结构
        user_id = book_info["user_id"]
        book_id = book_info["book_id"]
        book_path = book_info["book_path"]

        # 打开文件
        with open(book_path,"rb")as f:
            # 读取少量字节做编码测试
            b = f.read(1024)
            encoding = chardet.detect(b)["encoding"]
            if not encoding:
                # 不支持的编码类型
                Book().update_book_status(user_id=user_id,book_id=book_id,book_status=-2)

            else:
                b = b + f.read()
                lines = b.decode(encoding,'ignore')
                c = re.compile(r'\S{1,10}?[章节卷回篇册集话]\s+.{1,31}?\n')
                m = c.findall(lines)
                l = c.split(lines)
                if m and l:
                    del l[0]
                    for i in range(len(m)):
                        title = m[i].replace("\n","")
                        content = l[i]
                        time.sleep(0.05)
                        Book().new_chapter(book_id=book_id,chapter_index=i,chapter_title=title,chapter_content=content)
                    # 解析完成
                    Book().update_book_status(user_id=user_id,book_id=book_id,book_status=1)
                else:
                    # 正则匹配失败
                    Book().update_book_status(user_id=user_id,book_id=book_id,book_status=-1)


from django.urls import path
from .views import *

urlpatterns=[
    path('' ,ListTaskView.as_view(),name = 'tasklist'),
    path('item/create/',CreateTaskView.as_view(),name='taskcreate'),
    path('item/<int:pk>/update', UpdateTaskView.as_view(), name='taskupdate'),
    path('item/<int:pk>/delete',DeleteTaskView.as_view(), name = 'taskdelete')
]


#!/usr/bin/env python
import argparse
import os

import dask.diagnostics
import fsspec
import numpy as np
import vcm
import xarray as xr

from pathlib import Path

def open_zarr(path):
    mapper = fsspec.get_mapper(path)
    return xr.open_zarr(mapper)

def open_nudged(root):
    path = os.path.join(root, "reference_state.zarr")
    return open_zarr(path)

FIELDS = {
    "air_temperature_reference": "air_temperature",
    "specific_humidity_reference": "specific_humidity",
    "eastward_wind_reference": "eastward_wind",
    "northward_wind_reference": "northward_wind"
}

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("SIMULATION_ROOT")
    parser.add_argument("destination", help="The path to output store.")

    args, _ = parser.parse_known_args()

    nudged_root = Path(args.SIMULATION_ROOT) / "nudged"
    destination = args.destination

    ds = open_nudged(nudged_root)
    delp = ds.pressure_thickness_of_atmospheric_layer_reference
    interpolated = vcm.interpolate_to_pressure_levels(ds[FIELDS.keys()], delp, dim="z")
    interpolated = interpolated.rename(FIELDS)

    # Resample to daily frequency to match the diagnostics output from the
    # baseline and ML-corrected runs; label with the center of the interval.
    interpolated = interpolated.resample(time="D", loffset="12H").mean()

    destination_mapper = fsspec.get_mapper(destination)
    with dask.diagnostics.ProgressBar():
        interpolated.to_zarr(destination_mapper, consolidated=True)


from _md5 import md5
from functools import lru_cache

from read import read

@lru_cache
def hash_value(value):
    hash = md5()
    hash.update(value.encode())
    return hash.hexdigest()

def solve(value):
    i = 0
    ans = ''
    while len(ans) != 8:
        hash = hash_value(f'{value}{i}')
        if hash[:5] == '00000':
            ans += hash[5]
            print(ans)
        i += 1

def main():
    data = read()
    solve(data)

if __name__ == '__main__':
    main()


#coding:utf-8

"""
ID:          issue-3902
ISSUE:       3902
TITLE:       Aliases for the RETURNING clause
DESCRIPTION:
JIRA:        CORE-3546
FBTEST:      bugs.core_3546
"""

import pytest
from firebird.qa import *

db = db_factory()

test_script = """
    recreate table t(id int, x int default 123, y int default 456);
    commit;
    set list on;
    insert into t(id) values(1) returning x+y as i234567890123456789012345678901;
    insert into t(id) values(2) returning x-y    "/** That's result of (x-y) **/";
"""

act = isql_act('db', test_script)

expected_stdout = """
    I234567890123456789012345678901 579
    /** That's result of (x-y) **/  -333
"""

@pytest.mark.version('>=3.0')
def test_1(act: Action):
    act.expected_stdout = expected_stdout
    act.execute()
    assert act.clean_stdout == act.clean_expected_stdout


from pydantic import BaseSettings

class Messages(BaseSettings):
    # Base messages
    NO_API_KEY = "No API key provided."
    AUTH_REQ = "Authentication required."
    HTTP_500_DETAIL = "Internal server error."

    # Templates
    NO_VALID_PAYLOAD = "{} is not a valid payload."

messages = Messages()


from logging.config import dictConfig

from config.default import *

SQLALCHEMY_DATABASE_URI = f'sqlite:///{os.path.join(BASE_DIR, "hbb.db")}'
SQLALCHEMY_TRACK_MODIFICATIONS = False

SECRET_KEY = b"\xd3DLYd\x99\x88\xb6\xf6\xcd\xbe\x12j\xcc\x9d\xdb"

SIMPLEMDE_JS_IIFE = True
SIMPLEMDE_USE_CDN = True

MAIL_SERVER = "smtp.gmail.com"
MAIL_PORT = 465
MAIL_USERNAME = ""
MAIL_PASSWORD = ""
MAIL_USE_TLS = False
MAIL_USE_SSL = True

dictConfig(
    {
        "version": 1,
        "formatters": {
            "default": {
                "format": "[%(asctime)s] %(levelname)s in %(module)s: %(message)s"
            }
        },
        "handlers": {
            "file": {
                "level": "INFO",
                "class": "logging.handlers.RotatingFileHandler",
                "filename": os.path.join(BASE_DIR, "logs/hbb.log"),
                "maxBytes": 1024 * 1024 * 5,
                "backupCount": 5,
                "formatter": "default",
            }
        },
        "root": {"level": "INFO", "handlers": ["file"]},
    }
)


import os
import logging
import subprocess
import numpy as np
import pandas as pd
import nibabel as nib

from pathlib import Path
from typing import Tuple, Union, List, Dict
from fastprogress import progress_bar

__all__ = ['compute_mp_fn', 'ComputeSubjectROIStats', 'PairedTTest']

# ============================== FUNCTIONS ==============================
def fill_missing_roi_values(df: pd.DataFrame, n_rois: int) -> pd.DataFrame:
    """Fill missing values with mean of that column.
    This function will check `roi_num` column from 1 to `n_rois` values,
    if there is any missing number then it will be filled using the mean
    value of the column.

    Args:
        df: pandas dataframe containing data
        n_rois: number of ROI's in the dataframe.
    """
    if df.roi_num.unique().shape[0] != n_rois:
        missing_rois = list(set(range(1, n_rois + 1)
                                ).difference(set(df.roi_num.unique())))
        total_missing = len(missing_rois)
        if total_missing != 0:
            # Too many missing values
            if total_missing > (df.shape[0] * .20):  # no more than 20%
                raise ValueError("Too much missing values.")
            for val in missing_rois:
                row = {'roi_num': val}
                columns = df.columns.to_list()
                columns.remove('roi_num')
                for colname in columns:
                    row[colname] = df[colname].mean()
                df = df.append(row, ignore_index=True)
                df = df.sort_values(by='roi_num').reset_index(drop=True)
            print(f"Interpolated {total_missing} missing value(s).")

    return df

def compute_mp_fn(kwargs):
    """Compute wrapper function for compute-stats-multi multiprocessing command
    """

    input_path = kwargs['input_path']
    template_path = kwargs['template_path']
    output_path = kwargs['output_path']
    logging.info(f"Computing ROI stats for {input_path}...")
    sstats = ComputeSubjectROIStats(input_path=input_path,
                                    subject_space_pcl_path=template_path,
                                    output_path=output_path,
                                    show_pb=False)
    ret_code = sstats.run()
    if ret_code == 0:
        logging.info("done!")
    else:
        logging.error(f"Error in computing stats for {input_path}")
    return 0

# ============================== CLASSES ==============================
class ComputeSubjectROIStats:
    def __init__(self,
                 input_path: Union[str, Path],
                 subject_space_pcl_path: Union[str, Path],
                 output_path: Union[str, Path],
                 stats_filename: str = 'dti_stats.csv',
                 show_pb: bool = True):
        """Compute ROI stats for one subject.

            Stats include mean and standard deviation of
            fractional anisotropy (FA), axial diffusivity (AD),
            and radial diffusivity (RD).

        Args:
            input_path: subject's DTI volume nifti file path.
            subject_space_pcl_path: Transformed to subject space ROI
                atlas nifti file path.
            output_path: where the output ROI files will be saved.
            stats_filename: name of the csv file containing computed stats.
                This file will be saved in `output_path`.
            show_pb: Enable/disable progress bar visualization.
                Useful to disable in multiprocessing.
        """
        self.input_path = Path(input_path)
        self.subject_space_pcl_path = Path(subject_space_pcl_path)
        self.output_path = Path(output_path)
        self.output_path.mkdir(parents=True, exist_ok=True)
        self.show_pb = show_pb

        self.subject_name = self.input_path.parent.stem
        # Set save path for csv file
        self.filepath = self.output_path / \
            f"{self.subject_name}__{stats_filename}"
        # Set a temporary file path for computation operations
        self.tmp_path = self.output_path/f"{self.subject_name}_tmp_roi.nii.gz"

    def run(self) -> int:
        # Get number of ROIs in parcellation / roi file.
        pcl_data = nib.load(self.subject_space_pcl_path)
        self.rois_values = [
            int(i)
            for i in np.unique(pcl_data.get_fdata())
            if i != 0  # Zero is for background
        ]
        self.n_rois = len(self.rois_values)
        print(f"Found {self.n_rois} unique ROI values.")

        # Compute ROI stats
        self.stats_df = self.make_subject_rois()
        self.stats_df.to_csv(self.filepath, index=False)

        # Remove temporary files
        if os.path.exists(self.tmp_path):
            os.remove(self.tmp_path)

        return 0

    def make_subject_rois(self) -> pd.DataFrame:
        """Split each ROI, binary threshold, and save as
            nifti files for further computation
        """
        from fsl.wrappers.fslmaths import fslmaths

        input_path = str(self.input_path)
        subject_stats = {
            'roi_num': [],
            'fa_mean': [], 'fa_std': [],
            'ad_mean': [], 'ad_std': [],
            'rd_mean': [], 'rd_std': [],
            'md_mean': [], 'md_std': [],
        }

        for roi_num in progress_bar(self.rois_values, display=self.show_pb):
            print("ROI_NUM =", roi_num)
            # Create binarized single ROI image
            (fslmaths(self.subject_space_pcl_path)
             .thr(roi_num)  # Threshold
             .uthr(roi_num)  # Upper threshold
             .bin()  # Binary thresholding
             .run(self.tmp_path))
            # Apply the ROI on subject
            fslmaths(input_path).mul(self.tmp_path).run(self.tmp_path)

            # Compute ROI stats
            roi_stats_dict = self._compute(self.tmp_path)
            subject_stats['roi_num'].append(roi_num)
            for k, v in roi_stats_dict.items():
                subject_stats[k].append(v)

        # Remove temporary file
        os.remove(self.tmp_path)

        return pd.DataFrame(subject_stats)

    def _compute(self, roi_path) -> dict:
        """Compute stats using given metrics and operations"""
        roi_stats = {}

        # Fractional Anisotropy (FA)
        ret_code = subprocess.run([
            'TVtool', '-in', roi_path, '-fa', '-out', 'tmp_roi.nii.gz'
        ]).returncode
        if ret_code == 0:
            fa_img = nib.load('tmp_roi.nii.gz').get_fdata()
            roi_stats['fa_mean'] = fa_img.mean()
            roi_stats['fa_std'] = fa_img.std()

        # Axial Diffusivity (AD)
        ret_code = subprocess.run([
            'TVtool', '-in', roi_path, '-ad', '-out', 'tmp_roi.nii.gz'
        ]).returncode
        if ret_code == 0:
            ad_img = nib.load('tmp_roi.nii.gz').get_fdata()
            roi_stats['ad_mean'] = ad_img.mean()
            roi_stats['ad_std'] = ad_img.std()

        # Radial Diffusivity (RD)
        ret_code = subprocess.run([
            'TVtool', '-in', roi_path, '-rd', '-out', 'tmp_roi.nii.gz'
        ]).returncode
        if ret_code == 0:
            rd_img = nib.load('tmp_roi.nii.gz').get_fdata()
            roi_stats['rd_mean'] = rd_img.mean()
            roi_stats['rd_std'] = rd_img.std()

        # Mean Diffusivity (MD)
        # To compute the map of trace (TR) (equal to 3 times the mean diffusivity)
        ret_code = subprocess.run([
            'TVtool', '-in', roi_path, '-tr', '-out', 'tmp_roi.nii.gz'
        ]).returncode
        if ret_code == 0:
            md_img = nib.load('tmp_roi.nii.gz').get_fdata() / 3
            roi_stats['md_mean'] = md_img.mean()
            roi_stats['md_std'] = md_img.std()
        return roi_stats

class DataLoader:
    from scipy import stats

    def __init__(self,
                 pre: List[Union[str, Path]],
                 post: List[Union[str, Path]],
                 n_rois: int,
                 fill_missing_roi: bool = True) -> None:
        """Data loader class for analysis using different tests

        Args:
            pre: list of paths of pre or group 1 csv data files.
            post: list of paths of post or group 2 csv data files.
            n_rois: Total number of ROIs present in the data.
                Make sure the ROI's numbering starts from 1 and consecutive.
            fill_missing_roi: if True, fill missing values with mean of the
                column.
        """
        # load data and set `roi_num` as index
        self.pre_dfs = self._load_data_safely(pre, fill_missing_roi, n_rois)
        self.post_dfs = self._load_data_safely(post, fill_missing_roi, n_rois)

        # set `roi_num` as index
        self.pre_dfs = {
            name: df.set_index('roi_num')
            for name, df in self.pre_dfs.items()
        }
        self.post_dfs = {
            name: df.set_index('roi_num')
            for name, df in self.post_dfs.items()
        }

        # Check if all dataframes are of same size
        size = [df.shape[0] == n_rois for name, df in self.pre_dfs.items()]
        _errmsg = f"All files must be of same size. Found {size}"
        assert all(size) == True, _errmsg

        self.n_rois = n_rois

    def _load_data_safely(self,
                          paths_list:  List[Union[str, Path]],
                          fill_missing_roi: bool,
                          n_rois: int):
        """Check and fix data before testing"""
        dfs = {}
        for p in paths_list:
            df = pd.read_csv(p)
            if fill_missing_roi:
                df = fill_missing_roi_values(df, n_rois)
            if df.shape[0] != n_rois:
                _errmsg = f"# ROIs in {p} != n_rois ({df.shape[0]} != {n_rois})."
                _errmsg += " set `fill_missing_roi = True` to fill missing values"
                raise ValueError(_errmsg)
            dfs[Path(p).stem] = df
        return dfs

    def paired_t_test(self,
                      var_name: str,
                      roi_num: Union[None, str, int, list],
                      roi_groups: Union[None, Dict[int, list]],
                      alpha: float = 0.05) -> dict:
        """Perform paired T-test on the given variable

        Args:
            var_name: column name of the variable to perform T-test.
            roi_num: ROI number i.e. specific brain region.
                [default is -1, means perform t-test on all regions one-by-one]
            roi_groups: Apply paried T-test on the groups of ROIs.
                There will be one p-value per group.
            alpha: Significance value
        """
        if (roi_num is not None) and (roi_groups is not None):
            raise AttributeError("Either pass `roi_num` or `roigrps`")

        elif roi_num is not None:
            if isinstance(roi_num, str) and (roi_num == 'all'):
                ret_dict = {}
                for i in range(1, self.n_rois + 1):
                    t_stat, p_value = self._one_roi_paired_t_test(
                        var_name, roi_num=i
                    )
                    ret_dict[i] = {
                        't_stat': t_stat,
                        'two_sided_p_value': p_value,
                        'p_value': p_value / 2,
                    }
                return ret_dict
            elif isinstance(roi_num, list):
                ret_dict = {}
                for i in roi_num:
                    t_stat, p_value = self._one_roi_paired_t_test(
                        var_name, roi_num=i
                    )
                    ret_dict[i] = {
                        't_stat': t_stat,
                        'two_sided_p_value': p_value,
                        'p_value': p_value / 2,
                    }
                return ret_dict
            else:
                try:
                    roi_num = int(roi_num)
                    if roi_num == 0:
                        raise ValueError("ROI number must be > 0.")
                except ValueError:
                    _errmsg = "`roi_num` must be an integer > 0 or `all`. "
                    _errmsg += f"Found roi_num = {roi_num}."
                    raise ValueError(_errmsg)
                t_stat, p_value = self._one_roi_paired_t_test(
                    var_name, roi_num=roi_num)
                return {roi_num: {
                        't_stat': t_stat,
                        'two_sided_p_value': p_value,
                        'p_value': p_value / 2,
                        }}

        elif roi_groups is not None:
            ret_dict = {}
            for group_num, group in roi_groups.items():
                ret = self._group_roi_paired_t_test(
                    var_name, group
                )
                group_t_stat, group_two_sided_p_value = ret
                group_p_value = group_two_sided_p_value / 2
                p_value = group_p_value / len(group)
                ret_dict[group_num] = {
                    'group_t_stat': group_t_stat,
                    'group_two_sided_p_value': group_two_sided_p_value,
                    'group_p_value': group_p_value,
                    'p_value': p_value
                }
            return ret_dict

    def _group_roi_paired_t_test(self, var_name: str, roi_group: list
                                 ) -> Tuple[float, float]:
        """Perform T-test on group of ROIs
        Args:
            var_name: Name of the variable (column of dataframe) to use.
            roi_group: list of roi numbers as a group.

        Returns:
            t-statistic, p-value
        """
        pre_values = [
            df.loc[roi_num, var_name]
            for df in self.pre_dfs.values() for roi_num in roi_group
        ]
        post_values = [
            df.loc[roi_num, var_name]
            for df in self.post_dfs.values() for roi_num in roi_group
        ]
        t_stat, two_sided_p_value = self.stats.ttest_rel(
            a=pre_values,
            b=post_values
        )
        return t_stat, two_sided_p_value

    def _one_roi_paired_t_test(self, var_name, roi_num
                               ) -> Tuple[float, float]:
        """Perform T-test on single ROI.
        Args:
            var_name: Name of the variable (column of dataframe) to use.
            roi_group: list of roi numbers as a group.

        Returns:
            t-statistic, p-value
        """
        pre_values = [
            df.loc[roi_num, var_name] for df in self.pre_dfs.values()
        ]
        post_values = [
            df.loc[roi_num, var_name] for df in self.post_dfs.values()
        ]
        t_stat, two_sided_p_value = self.stats.ttest_rel(
            a=pre_values,
            b=post_values
        )
        return t_stat, two_sided_p_value

class PairedTTest:
    from scipy import stats

    def __init__(self,
                 pre: List[Union[str, Path]],
                 post: List[Union[str, Path]],
                 n_rois: int,
                 fill_missing_roi: bool = True) -> None:
        """Perform Paired T-test
        Args:
            pre: list of paths of pre or group 1 csv data files.
            post: list of paths of post or group 2 csv data files.
            n_rois: Total number of ROIs present in the data.
                Make sure the ROI's numbering starts from 1 and consecutive.
            fill_missing_roi: if True, fill missing values with mean of the
                column.
        """
        # load data and set `roi_num` as index
        self.pre_dfs = self._load_data_safely(pre, fill_missing_roi, n_rois)
        self.post_dfs = self._load_data_safely(post, fill_missing_roi, n_rois)

        # set `roi_num` as index
        self.pre_dfs = {
            name: df.set_index('roi_num')
            for name, df in self.pre_dfs.items()
        }
        self.post_dfs = {
            name: df.set_index('roi_num')
            for name, df in self.post_dfs.items()
        }

        # Check if all dataframes are of same size
        size = [df.shape[0] == n_rois for name, df in self.pre_dfs.items()]
        _errmsg = f"All files must be of same size. Found {size}"
        assert all(size) == True, _errmsg

        # stack all subjects in an array

        self.n_rois = n_rois

    def run_roi(self,
                var_names: list,
                roi_num: Union[None, str, int, list],
                combine: bool = False
                ) -> dict:
        """Perform paired T-test on the given variable(s) collectively

        Args:
            var_names: list of column name(s) of the variable to
                perform T-test.
            roi_num: ROI number i.e. specific brain region.
                Pass an integer, 'all', or list of rois.
                Where `all` means perform t-test on all regions one-by-one.
            combine: If True, perform t-test on all variables and roi
                combined. If False, perform t-test all variables and for
                each roi separately
        Returns:
            dict containing t-statistic and p-value
        """
        # Pre and post data shapes must be (n_subjects, n_rois, n_variables)
        if isinstance(roi_num, str):
            if roi_num == 'all':
                pre_data, post_data = self._get_selected_roi_num_data(
                    roi_num,
                    var_names
                )
            else:
                try:
                    roi_num = int(roi_num)
                    if roi_num == 0:
                        raise ValueError("ROI number must be > 0.")
                    pre_data, post_data = self._get_selected_roi_num_data(
                        roi_num,
                        var_names
                    )
                except ValueError:
                    _errmsg = "`roi_num` must be an integer > 0 or `all`. "
                    _errmsg += f"Found roi_num = {roi_num}."
                    raise ValueError(_errmsg)
        elif isinstance(roi_num, (list, int)):
            if isinstance(roi_num, int) and (roi_num == 0):
                raise ValueError("ROI number must be > 0.")
            pre_data, post_data = self._get_selected_roi_num_data(
                roi_num,
                var_names)

        # number of subjects
        assert len(self.pre_dfs) == pre_data.shape[0]
        assert len(self.post_dfs) == post_data.shape[0]
        assert pre_data.shape == post_data.shape

        # Perform paired T-test
        if combine:
            t_stat, two_sided_p_value = self.stats.ttest_rel(
                a=pre_data.flatten(),
                b=post_data.flatten()
            )
            ret_dict = {
                't_stat': t_stat,
                'two_sided_p_value': two_sided_p_value,
                'p_value': two_sided_p_value / 2,
            }
        else:
            if isinstance(roi_num, list):
                ret_dict = {}
                for i, _roi in enumerate(roi_num):
                    t_stat, two_sided_p_value = self.stats.ttest_rel(
                        a=pre_data[:, i, :].flatten(),  # per ROI
                        b=post_data[:, i, :].flatten()
                    )
                    ret_dict[_roi] = {
                        't_stat': t_stat,
                        'two_sided_p_value': two_sided_p_value,
                        'p_value': two_sided_p_value / 2,
                    }
            elif roi_num == 'all':
                ret_dict = {}
                for i, _roi in enumerate(range(1, self.n_rois+1)):
                    t_stat, two_sided_p_value = self.stats.ttest_rel(
                        a=pre_data[:, i, :].flatten(),  # per ROI
                        b=post_data[:, i, :].flatten()
                    )
                    ret_dict[_roi] = {
                        't_stat': t_stat,
                        'two_sided_p_value': two_sided_p_value,
                        'p_value': two_sided_p_value / 2,
                    }
            else:
                # assert len(pre_data.shape) == 2
                t_stat, two_sided_p_value = self.stats.ttest_rel(
                    a=pre_data.flatten(),  # per ROI
                    b=post_data.flatten()
                )
                ret_dict = {
                    't_stat': t_stat,
                    'two_sided_p_value': two_sided_p_value,
                    'p_value': two_sided_p_value / 2,
                }
        return ret_dict

    def run_roi_groups(self,
                       var_names: list,
                       roi_groups: Union[None, Dict[int, list]]) -> dict:
        """Perform paired T-test on the given variable

        Args:
            var_names: list of column name(s) of the variable to
                perform T-test.
            roi_groups: Apply paried T-test on the groups of ROIs.
                There will be one p-value per group.
        """
        ret_dict = {}
        for grp_num, roi_list in roi_groups.items():
            results = self.run_roi(
                var_names=var_names,
                roi_num=roi_list,
                combine=True
            )
            results['group_p_value'] = results['p_value']
            # Normalize p_value
            results['p_value'] = results['group_p_value'] / len(roi_list)
            ret_dict[grp_num] = results
        return ret_dict

    def _load_data_safely(self,
                          paths_list:  List[Union[str, Path]],
                          fill_missing_roi: bool,
                          n_rois: int):
        """Check and fix data before testing"""
        dfs = {}
        for p in paths_list:
            df = pd.read_csv(p)
            try:
                if fill_missing_roi:
                    df = fill_missing_roi_values(df, n_rois)
            except ValueError as e:
                print(e)
                raise ValueError(f"Error @ `{p}`")

            if df.shape[0] != n_rois:
                _errmsg = f"# ROIs in {p} != n_rois ({df.shape[0]} != {n_rois})."
                _errmsg += " set `fill_missing_roi = True` to fill missing values"
                raise ValueError(_errmsg)
            dfs[Path(p).stem] = df
        return dfs

    def _get_selected_roi_num_data(
        self,
        roi_num: Union[str, int, list],
        var_names: List[str]
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Get selected ROIs data from subject dataframes in np.array format
        """
        if roi_num == 'all':
            pre_data = np.array([
                s[var_names].values
                for s in self.pre_dfs.values()
            ])
            post_data = np.array([
                s[var_names].values
                for s in self.post_dfs.values()
            ])
        else:
            pre_data = np.array([
                s[var_names].loc[roi_num].values
                for s in self.pre_dfs.values()
            ])
            post_data = np.array([
                s[var_names].loc[roi_num].values
                for s in self.post_dfs.values()
            ])
        return pre_data, post_data


from collections import defaultdict
from typing import Dict, Optional

import discord
from discord.utils import SnowflakeList
from redbot.core import Config, checks, commands
from redbot.core.bot import Red

class ModOnlyMode(commands.Cog):
    """
    Cog to limit the bot to mods and higher.
    """

    __version__ = "2021.03"

    async def red_delete_data_for_user(self, **kwargs):
        """ Nothing to delete """
        return

    def format_help_for_context(self, ctx):
        pre_processed = super().format_help_for_context(ctx)
        return f"{pre_processed}\nCog Version: {self.__version__}"

    def __init__(
        self,
        bot: Red,
        config: Config,
        cache: SnowflakeList,
        exclusion_cache: Dict[str, SnowflakeList],
    ):
        self.bot: Red = bot
        self.cache: SnowflakeList = cache
        self.config: Config = config
        self.exclusion_cache: Dict[str, SnowflakeList] = exclusion_cache

    @classmethod
    async def setup(cls, bot: Red):
        config = Config.get_conf(None, 78631113035100160, cog_name="ModOnlyMode")
        config.register_guild(active=False, exclusions=[])

        all_guild_data = await config.all_guilds()

        cache = SnowflakeList(())
        exclusion_cache: Dict[str, SnowflakeList] = defaultdict(
            lambda: SnowflakeList(())
        )

        for guild_id, guild_data in all_guild_data.items():
            if guild_data.get("active", False):
                cache.add(guild_id)
            for exclusion in guild_data.get("exclusions", []):
                exclusion_cache[exclusion].add(guild_id)

        cog = cls(bot, config, cache, exclusion_cache)
        bot.add_cog(cog)

    async def __permissions_hook(self, ctx: commands.Context) -> Optional[bool]:
        if ctx.guild and self.cache.has(ctx.guild.id):
            assert isinstance(ctx.author, discord.Member), "mypy"  # nosec
            if not await self.bot.is_mod(ctx.author):
                if ctx.cog and self.exclusion_cache[ctx.cog.qualified_name].has(
                    ctx.guild.id
                ):
                    return None
                return False

        return None

    @checks.admin()
    @commands.guild_only()
    @commands.group(name="modonlymodeset", aliases=["momset"])
    async def momset(self, ctx: commands.Context):
        """
        Settings for ModOnlyMode
        """

    @momset.command()
    async def enable(self, ctx: commands.GuildContext):
        """
        Makes the bot's commands only work for mods and above in the guild.

        Note: The licenseinfo command will still be available.
          Disabling this command is not allowed under red's license,
          please do not look for alternative means of doing so.
        """

        if self.cache.has(ctx.guild.id):
            return await ctx.send("Mod only mode is already enabled here")

        await self.config.guild(ctx.guild).active.set(True)
        self.cache.add(ctx.guild.id)
        await ctx.tick()

    @momset.command()
    async def disable(self, ctx: commands.GuildContext):
        """
        Disable mod only mode in this guild.
        """

        if not self.cache.has(ctx.guild.id):
            return await ctx.send("Mod only mode is not enabled here")

        await self.config.guild(ctx.guild).active.clear()
        self.cache.remove(ctx.guild.id)
        await ctx.tick()

    @momset.command()
    async def excludecog(self, ctx: commands.GuildContext, *, cog_name: str):
        """
        Exclude a cog from mod only mode.
        """

        if self.bot.get_cog(cog_name) is None:
            return await ctx.send_help()

        if self.exclusion_cache[cog_name].has(ctx.guild.id):
            return await ctx.send("That cog was already excluded.")

        async with self.config.guild(ctx.guild).exclusions() as cog_list:
            cog_list.append(cog_name)
            self.exclusion_cache[cog_name].add(ctx.guild.id)
            await ctx.tick()

    @momset.command()
    async def reincludecog(self, ctx: commands.GuildContext, *, cog_name: str):
        """
        Reinclude a cog in mod only mode.
        """

        if self.bot.get_cog(cog_name) is None:
            return await ctx.send_help()

        if not self.exclusion_cache[cog_name].has(ctx.guild.id):
            return await ctx.send("That cog was not excluded.")

        async with self.config.guild(ctx.guild).exclusions() as cog_list:
            cog_list.remove(cog_name)
            self.exclusion_cache[cog_name].remove(ctx.guild.id)
            await ctx.tick()


import torch
import torch.nn as nn
import torch.optim

import itertools

from tqdm import tqdm

import data
import models
import numpy as np

import pathlib

import utils

import click

def basic_train(model, dataloader, parameter_set, optimizer, loss_function, device):
    model.train()
    model.to(device)

    max_epochs = parameter_set['epochs']
    l2_coefficient = parameter_set['l2']

    iterator = tqdm(range(max_epochs), desc='Training')

    for _ in iterator:
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)

            y_pred = model(images)
            loss = loss_function(y_pred, labels)

            l2_loss = torch.sum(torch.stack([(parameter ** 2).sum() for parameter in model.parameters()]))
            loss += l2_loss * l2_coefficient
            del l2_loss

            optimizer.zero_grad()
            loss.backward()

            optimizer.step()

def create_basic_set(**parameters):
    parameter_pairs = []
    for key, value_list in parameters.items():
        parameter_pairs.append([(key, val) for val in value_list])

    parameter_sets = []

    for product in itertools.product(*parameter_pairs):
        parameter_set = {}
        for key, value in product:
            parameter_set[key] = value

        parameter_sets.append(parameter_set)

    return parameter_sets

def get_optimizer(model_parameters, parameter_set):
    if parameter_set['optimizer_type'] == 'adam':
        return torch.optim.Adam(model_parameters, lr=parameter_set['lr'])
    elif parameter_set['optimizer_type'] == 'sgd':
        return torch.optim.SGD(model_parameters, lr=parameter_set['lr'])
    else:
        raise NotImplementedError

def accuracy(model, loader, device):
    correct_count = 0
    total_count = 0

    model.to(device)

    for images, true_labels in tqdm(loader, desc='Accuracy Test'):
        total_count += len(images)
        images = images.to(device)
        true_labels = true_labels.to(device)

        predicted_labels = utils.get_labels(model, images).detach()

        correct = torch.eq(predicted_labels, true_labels)
        correct_count += len(torch.nonzero(correct))

    return correct_count / total_count

@click.command()
@click.argument('coefficient', type=str)
@click.argument('run', type=int)
def main(coefficient, run):
    device = 'cuda'
    batch_size = 128
    loss_function = nn.CrossEntropyLoss()

    dataset = data.parse_dataset('cifar10', 'std:train')
    dataloader = torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)

    base_folder = pathlib.Path('trained-models/standard-2') / coefficient / str(run)
    base_folder.mkdir(parents=True, exist_ok=True)

    test_dataset = data.parse_dataset('cifar10', 'std:test')
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size, shuffle=False)

    # parameter_sets = create_basic_set(optimizer_type=['adam'], lr=[1e-3, 1e-4, 1e-5], epochs=[100, 200, 300])

    lrs = [1e-4, 1e-5]
    epochs = [200, 300]
    l2s = [1e-3, 1e-4]

    parameter_sets = create_basic_set(optimizer_type=['adam'], lr=lrs, epochs=epochs, l2=l2s)

    for i, parameter_set in enumerate(parameter_sets):
        lr_string = '1e' + str(int(np.log10(parameter_set['lr'])))
        l2_string = '1e' + str(int(np.log10(parameter_set['l2'])))

        instance_name = f'{lr_string}_{parameter_set["epochs"]}_{l2_string}'

        print('Running ', instance_name)

        instance_path = base_folder / f'{instance_name}.pth'
        if not instance_path.exists():
            model = models.mini_model(coeff=float(coefficient))

            optimizer = get_optimizer(model.parameters(), parameter_set)
            basic_train(model, dataloader, parameter_set, optimizer, loss_function, device)

            torch.save(model.state_dict(), instance_path)
            acc = accuracy(model, test_dataloader, device)
            print(parameter_set)
            print(f'{instance_name} Accuracy: ', acc * 100.0, '%')

if __name__ == '__main__':
    main()