try:
    import cPickle # Python 2
except:
    import pickle as cPickle # Python 3

import os
import numpy as np
import json
from prev_ob_models.Birgiolas2020.isolated_cells import *
from blenderneuron.nrn.neuronnode import NeuronNode
from olfactorybulb.database import Odor, OdorGlom, CellModel, database
from math import pow
from LFPsimpy import LfpElectrode
import sys
from heapq import *
from matplotlib import pyplot as plt
from hashlib import sha1
from random import random, seed

from olfactorybulb.paramsets.base import *
from olfactorybulb.paramsets.case_studies import *
from olfactorybulb.paramsets.sensitivity import *

class OlfactoryBulb:
    """
    The main class used to build and simulate the olfactory bulb network model.
    """

    def __init__(self, params="ParameterSetBase", autorun=True):
        """
        :param params: The name of the class defined in olfactorybulb.paramsets that defines the network parameters
        :param autorun: When true, after the network model is built, starts the simulation
        """

        if type(params) == str:
            params = eval(params)()

        self.params = params

        self.rnd_seed = params.rnd_seed

        self.slice_dir = os.path.abspath(os.path.join(params.slice_dir, params.slice_name))
        self.cells = {}
        self.inputs = []

        self.gj_source_gids = set()
        self.gjs = []

        # Just use the BlenderNEURON package functions (e.g. no server/client)
        self.bn_server = NeuronNode(server_end='Package')

        from neuron import h, load_mechanisms
        self.h = h
        self.pc = h.ParallelContext()
        self.mpimap = {}
        self.nranks = int(self.pc.nhost())
        self.mpirank = self.pc.id()

        # Keep track of rank complexities with a min-heap
        self.rank_complexities = [(0, r) for r in range(self.nranks)]

        self.t_vec = h.Vector()
        self.t_vec.record(h._ref_t, params.recording_period)
        self.v_vectors = {}
        self.input_vectors = []

        for cell_type in ['MC', 'GC', 'TC']:
            self.load_cells(cell_type)

        if self.mpirank == 0:
            complexities = np.array([c[0] for c in self.rank_complexities])
            min = np.min(complexities)
            max = np.max(complexities)
            mean = np.mean(complexities)

            print('Rank Complexity min: %s, mean: %s, max: %s' % (min, mean, max))

        for synapse_set in ['GCs__MCs', 'GCs__TCs']:
            self.load_synapse_set(synapse_set)

        # Load glom->cell links
        self.load_glom_cells()

        # Create gap junctions between MC and TC tufts
        for cell_type, g_gap in params.gap_juction_gmax.items():
            self.add_gap_junctions(cell_type, g_gap)

        # Set synapse parameters
        for syn_mech, syn_values in params.synapse_properties.items():
            if hasattr(h, syn_mech):
                for syn_attrib, attrib_value in syn_values.items():
                    [setattr(s, syn_attrib, attrib_value) for s in getattr(h, syn_mech)]

        # Add glomerular inputs
        for time, odor_info in params.input_odors.items():
            self.add_inputs(odor=odor_info["name"], t=time, rel_conc=odor_info["rel_conc"])

        # LFP
        self.electrode = self.create_lfp_electrode(*params.lfp_electrode_location,
                                                   sampling_period=params.recording_period)

        self.setup_status_reporter()

        for cell_type in params.record_from_somas:
            self.record_from_somas(cell_type)

        if self.mpirank == 0 and self.nranks == 1:
            from neuron import gui
            # h.load_file('1x1x1-testbed.ses')

            h.newPlotI()
            [g for g in h.Graph][-1].addvar('LfpElectrode[0].value')

        if autorun:
            self.run(params.tstop)

            if self.mpirank == 0:
                self.results_dir = os.path.join('results', params.name)
                if not os.path.exists(self.results_dir):
                    os.makedirs(self.results_dir)

            self.save_recorded_vectors()

            if self.mpirank == 0:
                t, lfp = self.get_lfp()

            # Cleanup on MPI
            if self.nranks > 1:
                database.close()
                self.h.quit()

    def stim_glom_segments(self, time, input_segs, intensity):
        """
        Adds input synapses onto glomerular tufts at specified start time and intensity

        The inhalation part of a sniff cycle is modeled as a gaussian probability that is centered at
        the midpoint of the inhalation onset and end. The probability is translated into spikes. The spikes
        then trigger the excitatory synapses placed at the mitral/tufted cell tufts.

        Intensity regulates how many spikes to pick from the gaussian.

        :param time: the inhalation onset time in ms

        :param input_segs: a list containing tuples of:
            a) The name of the segment to stimulate as it appears on the current MPI rank
            b) segment gid
            c) segment name as it appears when there is only one rank. If not using MPI a) and c) are same.
        :param intensity: 0-1 representing odor intensity

        :return: None
        """

        h = self.h

        inhale_duration = self.params.inhale_duration

        # ORN firing rate
        max_firing_rate = self.params.max_firing_rate

        # Translate intensity to number of spikes per inhalation
        spike_count = int(round(max_firing_rate * intensity * (inhale_duration / 1000.0)))

        for seg_name, seg_gid, single_rank_seg_name in input_segs:
            # Randomize spikes to each tufted segment
            seed_source = "%s|%s|%s|%s" % (self.rnd_seed, time, single_rank_seg_name, intensity)
            np.random.seed(self.stable_hash(seed_source))

            # Odor is modeled as a gaussian spike train representing OSN spikes during inhalation
            # exhalation is assumed to not generate OSN spikes
            spike_times = self.get_gaussian_spike_train(spike_count, time, inhale_duration)

            # Create synapse point process
            seg = eval(seg_name.replace('(1)', '(.999)'))
            syn = h.Exp2Syn(seg)
            syn.tau1 = self.params.input_syn_tau1
            syn.tau2 = self.params.input_syn_tau2

            if "MC" in seg_name:  # MCs
                delay = self.params.mc_input_delay
                weight = self.params.mc_input_weight

            else:  # "TC"
                delay = self.params.tc_input_delay
                weight = self.params.tc_input_weight

            # VecStim will deliver events to synapse at vector times
            ns = h.VecStim()
            ns.play(h.Vector(spike_times + delay))

            # Netcon to trigger the synapse
            netcon = h.NetCon(
                ns,
                syn,
                0,  # thresh
                0,  # delay
                weight  # weight uS
            )

            # Record odor input events
            input_vec = h.Vector()
            netcon.record(input_vec)
            self.input_vectors.append((single_rank_seg_name, input_vec))

            self.inputs.append((syn, ns, netcon))

    def stable_hash(self, source, digits=9):
        """
        Creates a hash code of digits long that is stable across different machines.

        :param source: The string to hash, in this case a section name
        :param digits: The number of digits to keep of the hash
        :return: The hash code as an integer
        """

        return int(sha1(source.encode()).hexdigest(), 16) % (10 ** digits)

    def run(self, tstop):
        """
        Runs the NEURON simulation until the specified stop time

        :param tstop: Simulation stop time
        """

        if self.mpirank == 0:
            print('Starting simulation...')

        h = self.h
        h.dt = self.params.sim_dt
        h.tstop = tstop

        if self.nranks == 1:
            h.cvode_active(0)
            h.cvode.cache_efficient(1)
            h.run()

        else:
            self.pc.setup_transfer()
            self.pc.timeout(1)
            # h.cvode.cache_efficient(0) # This line causes gap junction Seg Faults
            h.cvode_active(0)
            self.pc.set_maxstep(1)
            h.stdinit()
            self.pc.psolve(h.tstop)

        # Clear status updater line
        if self.mpirank == 0:
            print('')

    def print_status(self):
        """
        Prints the current simulation time on the same line (no new line)
        """

        sys.stdout.write("\rTime: %s ms" % self.h.t)
        sys.stdout.flush()

    def setup_status_reporter(self):
        """
        Sets up the NEURON simulation to report the simulation time
        """

        if self.mpirank == 0:
            h = self.h

            collector_stim = h.NetStim(0.5)
            collector_stim.start = 0
            collector_stim.interval = 1
            collector_stim.number = 1e9
            collector_stim.noise = 0

            collector_con = h.NetCon(collector_stim, None)
            collector_con.record(self.print_status)

            self.collector_stim = collector_stim
            self.collector_con = collector_con

    def create_lfp_electrode(self, x, y, z, sampling_period, method='Line'):
        """
        Uses the LFPsimpy package to add an LFP electrode at the specified x,y,z location

        See `LFPsimpy package <https://github.com/justasb/LFPsimpy>`_.

        :param x: y, z coordinates in um
        :param sampling_period: How often to compute the LFP signal in ms
        :param method: One of 'Line', 'Point', or 'RC'.
        :return: an LFPsimpy LfpElectrode object
        """

        return LfpElectrode(x, y, z, sampling_period, method)

    def get_lfp(self):
        """
        Returns the LFP signal in nV

        :return: a tuple of LFP times, and voltages (nV)
        """

        if self.electrode is None or not any(self.electrode.times):
            raise Exception('Run simulation first to get the LFP')

        t = self.electrode.times
        lfp = self.electrode.values

        with open(os.path.join(self.results_dir, 'lfp.pkl'), 'wb') as f:
            cPickle.dump((t, lfp), f)

        return t, lfp

    def get_model_inputsegs(self):
        """
        Queries the model database to get the 'root' segments of the tufted dendrites
        of the mitral and tufted cells

        :return: A dict that maps the cell model's class name to the name of the root tufted dendrite section
        """

        # Get all the different cell models used in the slice
        input_models = set()
        for cells in self.glom_cells.values():
            for cell in cells:
                input_models.add(cell[:cell.find('[')])

        # Get each model's input segments (in the tuft)
        model_inputsegs = {m.class_name: m.tufted_dend_root
                           for m in CellModel \
                               .select(CellModel.class_name, CellModel.tufted_dend_root) \
                               .where(CellModel.class_name.in_(list(input_models)))}

        return model_inputsegs

    def add_gap_junctions(self, in_name, g_gap):
        """
        Adds gap junctions between tufted dendrites of specified cells

        :param in_name: A part of a cell class name (e.g. 'Mitral') used to select a cell to which the GJ is added
        :param g_gap: The conductance of the gap junctions
        """

        model_inputsegs = self.get_model_inputsegs()

        for glom_id, cells in self.glom_cells.items():

            input_segs = []
            for cell in cells:
                if in_name not in cell:
                    continue

                model_class = cell[:cell.find('[')]
                input_seg = model_inputsegs[model_class]

                single_rank_address = 'h.' + cell + '.' + input_seg
                single_rank_gid = self.stable_hash(single_rank_address)

                rank_cell = self.bn_server.rank_section_name(cell)

                if rank_cell is not None:
                    seg_address = 'h.' + rank_cell + '.' + input_seg
                else:
                    seg_address = None

                input_segs.append((seg_address, single_rank_gid))

            if len(input_segs) > 0:
                self.create_gap_junctions_between(input_segs, g_gap)

        self.pc.setup_transfer()

    def create_gap_junctions_between(self, input_segs, g_gap):
        """
        Creates gap junctions between a list of specified segments. GJs are connected in a chain
        (e.g. Seg1 <-GJ1-> Seg2 <-GC2-> Seg3)

        :param input_segs: List of segments to connect by gap junctions
        :param g_gap: Gap junction conductance
        """

        count = len(input_segs)

        if count < 2:
            return

        h = self.h

        first_seg = input_segs[0]
        last_seg = input_segs[-1]

        if count > 2:
            for i, seg in enumerate(input_segs[:-1]):
                next_seg = input_segs[i + 1]

                self.create_gap_junction(seg, next_seg, g_gap)

        self.create_gap_junction(first_seg, last_seg, g_gap)

    def create_gap_junction(self, seg_1_info, seg_2_info, g_gap):
        """
        Creates a gap junction between two segments

        :param seg_1_info: Tuple of the name and gid of the first segment
        :param seg_2_info: Tuple of the name and gid of the second segment
        :param g_gap: Gap junction conductance
        """

        h = self.h

        seg_1_name, seg_1_gid = seg_1_info
        seg_2_name, seg_2_gid = seg_2_info

        if seg_1_name is not None:
            seg1 = eval(seg_1_name.replace('(1)', '(.999)'))

            if seg_1_gid not in self.gj_source_gids:
                self.pc.source_var(seg1._ref_v, seg_1_gid, sec=seg1.sec)
                self.gj_source_gids.add(seg_1_gid)

            gap1 = h.GapJunction(seg1.x, sec=seg1.sec)
            gap1.g = g_gap
            self.pc.target_var(gap1._ref_v_other, seg_2_gid)
            self.gjs.append(gap1)

        if seg_2_name is not None:
            seg2 = eval(seg_2_name.replace('(1)', '(.999)'))

            if seg_2_gid not in self.gj_source_gids:
                self.pc.source_var(seg2._ref_v, seg_2_gid, sec=seg2.sec)
                self.gj_source_gids.add(seg_2_gid)

            gap2 = h.GapJunction(seg2)
            gap2.g = g_gap
            self.pc.target_var(gap2._ref_v_other, seg_1_gid)
            self.gjs.append(gap2)

    def add_inputs(self, odor, t, rel_conc):
        """
        Add odor stimulation to the tufts of the principal cells

        :param odor: The name of the odor
        :param t: Onset time
        :param rel_conc: Relative concentration 0-1
        """

        model_inputsegs = self.get_model_inputsegs()

        # Get input odor glomeruli
        glom_intensities = {g.glom_id: g.intensity
                            for g in OdorGlom
                                .select(OdorGlom.glom_id, OdorGlom.intensity)
                                .join(Odor)
                                .where(Odor.name == odor)}

        for glom_id, cells in self.glom_cells.items():
            glom_id = int(glom_id)

            input_segs = []
            for cell in cells:
                rank_cell = self.bn_server.rank_section_name(cell)

                # Add inputs only to cells that are on this rank
                if rank_cell is None:
                    continue

                model_class = rank_cell[:rank_cell.find('[')]
                input_seg = model_inputsegs[model_class]
                seg_address = 'h.' + rank_cell + '.' + input_seg

                single_rank_address = 'h.' + cell + '.' + input_seg
                single_rank_gid = int(sha1(single_rank_address.encode()).hexdigest(), 16) % (10 ** 9)

                input_segs.append((seg_address, single_rank_gid, single_rank_address))

            if len(input_segs) > 0:
                glom_intensity = glom_intensities[glom_id] * rel_conc
                self.stim_glom_segments(t, input_segs, glom_intensity)

    def load_glom_cells(self):
        """
        Loads a dict that maps glomeruli ids to cells that are attached to each glomerulus
        """

        with open(os.path.join(self.slice_dir, 'glom_cells.json')) as f:
            self.glom_cells = json.load(f)

    def get_gaussian_spike_train(self, spikes=50, start_time=100, duration=10):
        """
        Gets a spike train from a gaussian probability distribution whose 99% range starts
        at the specified time and lasts for the specified duration.

        :param spikes: The number of spikes to generate
        :param start_time: The onset time of the gaussian
        :param duration: The duration of the gaussian
        :return: A numpy array of spike times in chronological order
        """

        # Create a gaussian whose 99% range starts at start_time
        # and ends at start_time + duration
        normal_stdev = duration / (2.576 * 2)

        times = np.random.normal(start_time + (duration / 2.0), normal_stdev, spikes)

        # Remove any spikes outside this range
        times = times[np.where((times > start_time) & (times < start_time + duration))]
        times.sort()

        return times

    def load_cells(self, cell_type):
        """
        Load the cells of the specified type onto least busy MPI ranks.

        'Busyness' of a rank is the sum of all cell complexities on that rank, as measured by the number
        of segments of each cell.

        :param cell_type: One of 'MC', 'GC', 'TC'
        """

        # Load the cell json file
        path = os.path.join(self.slice_dir, cell_type + 's.json')

        with open(path, 'r') as f:
            group_dict = json.load(f)

        # Count how many of each cell model will be on each rank
        rank_cell_counts = {r: {} for r in range(self.nranks)}

        for ri, root in enumerate(group_dict['roots']):
            # Get the least loaded rank
            min_complexity, min_complexity_rank = heappop(self.rank_complexities)

            # Cell nseg count is used as a proxy for complexity
            nsegs = self.get_nseg_count(root)

            # Add to rank complexity and push back onto the heap
            heappush(self.rank_complexities, (min_complexity + nsegs, min_complexity_rank))

            # Assign cell to least busy rank
            cell_rank = min_complexity_rank

            name = root['name']
            name = name[0:name.find('[')]

            count = rank_cell_counts[cell_rank].get(name, 0)

            self.mpimap[root['name'][:root['name'].find(']') + 1]] = {
                'name': name + '[' + str(count * 2) + ']',
                'rank': cell_rank
            }

            count += 1
            rank_cell_counts[cell_rank][name] = count

        # Load that many base instances of each model
        self.cells[cell_type] = []
        for cell_model_name, count in rank_cell_counts[self.mpirank].items():
            cell_models = [eval(cell_model_name + '()') for _ in range(count)]
            self.cells[cell_type].extend(cell_models)

        # Update section index with the new cells
        self.bn_server.update_section_index()

        # Apply the cell json file onto the base instances
        self.bn_server.init_mpi(self.pc, self.mpimap)
        self.bn_server.update_groups([group_dict])

    def record_from_somas(self, cell_type):
        """
        Adds NEURON vector recorders to the somas of the specified cell types

        :param cell_type: One of 'MC', 'GC', 'TC'
        """

        h = self.h

        for cell_model in self.cells[cell_type]:
            v_vec = h.Vector()
            v_vec.record(cell_model.soma(0.5)._ref_v, self.params.recording_period)
            self.v_vectors[str(cell_model.soma)] = v_vec

    def save_recorded_vectors(self):
        """
        Saves soma voltage traces and odor input spike times to Pickle files for later processing

        Saves to the results directory as 'soma_vs.pkl' and 'input_times.pkl'
        """

        # Gather cell voltage vectors
        all_v_vecs = self.pc.py_gather(self.v_vectors, 0)

        if all_v_vecs is not None:
            t = self.t_vec.to_python()
            result = []
            for rank_v_vecs in all_v_vecs:
                for cell, v_vec in rank_v_vecs.items():
                    result.append((cell, t, v_vec.to_python()))

            with open(os.path.join(self.results_dir, 'soma_vs.pkl'), 'wb') as f:
                cPickle.dump(result, f)

        # Gather input event time vectors
        all_input_vecs = self.pc.py_gather(self.input_vectors, 0)

        if all_input_vecs is not None:
            result = []
            for rank_input_vecs in all_input_vecs:
                for seg_name, t_vec in rank_input_vecs:
                    result.append((seg_name, t_vec.to_python()))

            with open(os.path.join(self.results_dir, 'input_times.pkl'), 'wb') as f:
                cPickle.dump(result, f)

    def get_nseg_count(self, root_dict):
        """
        Recursively counts the number of segments of a cell provided its BlenderNEURON root segment dict

        :param root_dict: The root segment dict of a cell as saved by BlenderNEURON
        :return: The total number of segments of the cell
        """

        count = root_dict["nseg"]

        for child_dict in root_dict['children']:
            count += self.get_nseg_count(child_dict)

        return count

    def load_synapse_set(self, synapse_set):
        """
        Uses BlenderNEURON to load a previously saved set of synapses between a population of cells

        :param synapse_set: One of 'GCs__MCs' or 'GCs__TCs' as seen in the olfactorybulb.slices.DorsalColumnSlice folder.
        """

        path = os.path.join(self.slice_dir, synapse_set + '.json')

        with open(path, 'r') as f:
            synapse_set_dict = json.load(f)

        self.bn_server.create_synapses(synapse_set_dict)


import numpy as np
import tensorflow as tf
import tqdm
from agents.maddpg import MADDPGAgent
from agents.util import ReplayBuffer
import time
import copy

# A class which stores the MADDPG agents and acts as an interface between them and the environment
class MADDPGRunner():
    def __init__(self,
            env,
            batch_size=64
            ):
        self.env = env
        self.batch_size = batch_size
        # Number of unique observations=number of agents
        self.num_agents = len(env.observation_space)
        self.agents = []
        self.obs_space_size = 0
        self.act_space_size = 0
        for i in range(self.num_agents):
            # TODO: Add an interface to modify individual hyperparameters
            agent = MADDPGAgent(env, i)
            self.agents.append(agent)
            self.obs_space_size += env.observation_space[i].shape[0]
            self.act_space_size += env.action_space[i].shape[0]

        self.buffer = ReplayBuffer(self.obs_space_size, self.act_space_size, num_agents=self.num_agents)

    # Get all of the explorative actions for the different agents
    def get_acts(self, states):
        acts = []
        for i in range(self.num_agents):
            # Reshape our states and actions to work with the NN
            state = tf.expand_dims(tf.convert_to_tensor(states[i]), 0)
            state = tf.reshape(state, (1, -1))
            acts.extend(self.agents[i].policy(state))

        return tf.convert_to_tensor(acts)

    # Get the non-exploring actions for demonstrations
    def get_non_exploring_acts(self, states, policy_param):
        acts = []
        for i in range(self.num_agents):
            state = states[i]
            acts.extend(self.agents[i].non_exploring_policy(state, policy_param))

        return tf.convert_to_tensor(acts)

    # Reshape and push experiences to the experience replay buffer
    def distribute_to_buffers(self, prev_states, actions, rewards, states, dones):
        prev_state = np.squeeze(np.reshape(prev_states, (-1, 1)))
        state = np.squeeze(np.reshape(states, (-1, 1)))
        self.buffer.add((prev_state, actions, rewards, state, dones))

    # Splits the observation blob into each agents observation
    def split_obs(self, observations):
        pointer = 0
        pointer_end = 0
        split_obs = []
        for i in range(self.num_agents):
            pointer_end = pointer + self.env.observation_space[i].shape[0]
            split_obs.append(observations[:, pointer:pointer_end])
            pointer = pointer_end

        return split_obs

    # Loop through agents and perform an update
    def update_agents(self, update_batch, next_acts):
        for i in range(self.num_agents):
            self.agents[i].perform_update_step(update_batch, next_acts)

    # Check if the episode is terminal
    def is_done(self, dones):
        for i in range(len(dones)):
            if dones[i]:
                return True
        return False

    # Train the agents for the given number of episodes
    def train(self, num_episodes=100, num_steps=100):
        ep_reward_list =[]
        avg_reward_list = []
        info_buffer = []
        best_reward = None
        best_reward_avg = None

        with tqdm.trange(num_episodes) as t:
            for i in t:
                episodic_reward, info = self.train_episode(num_steps)
                ep_reward_list.append(episodic_reward)
                # Mean of last 10 episodes
                avg_reward = np.mean(ep_reward_list[-10:])
                t.set_description(f'Episode {i}')
                episodic_reward_str = "⠀{:6.0f}".format(episodic_reward) #Extra spacing was ignored. So I added an invisible character right before
                avg_reward_str = "⠀{:6.0f}".format(avg_reward)
                t.set_postfix(episodic_reward=episodic_reward_str, average_reward=avg_reward_str)
                avg_reward_list.append(avg_reward)
                info_buffer.append(info)

                # We save the agents based on best single and episodic rewards for the whole system
                if(best_reward == None or episodic_reward > best_reward):
                    best_reward = episodic_reward
                    for i, agent in enumerate(self.agents):
                        agent.cache_best_single()

                if(i >= 10 and (best_reward_avg == None or avg_reward > best_reward_avg)):
                    best_reward_avg = avg_reward
                    for i, agent in enumerate(self.agents):
                        agent.cache_best_average()

        return ep_reward_list, avg_reward_list, info_buffer

    # Run the training for a single episode
    def train_episode(self, num_steps, render=False):
        prev_states = self.env.reset()
        episodic_reward = 0
        episode_info = []

        for i in range(num_steps):
            if render:
                self.env.render()

            actions = self.get_acts(prev_states)
            # Recieve state and reward from environment.
            states, rewards, dones, info = self.env.step(actions)
            rewards = tf.cast(rewards, dtype=tf.float32)
            actions = np.squeeze(np.reshape(actions, (-1,1)))
            # Push experience to the replay buffers
            self.distribute_to_buffers(prev_states, actions, rewards, states, dones)
            update_batch = self.buffer.sample_batch(self.batch_size)
            # Split the observations from the replay buffer for each agent
            next_obs = self.split_obs(update_batch[-2])
            # Get non-exploring actions
            next_actions = self.get_non_exploring_acts(next_obs, "last")
            next_actions = np.squeeze(np.reshape(next_actions, (self.batch_size, -1, 1)))
            # Perform agent NN updates
            self.update_agents(update_batch, next_actions)
            episodic_reward += sum(rewards)
            episode_info.append(info)

            prev_states = states

            # End this episode when `done` is True
            if self.is_done(dones):
                break

        return episodic_reward, episode_info

    # Run a noiseless episode to observe agent behaviour
    def run_episode(self, num_steps=100, render=True, waitTime=0.1, policy_param='last', mode='human'):
        prev_states = self.env.reset()
        episodic_reward = 0
        episode_info = []
        states = []
        images = []

        for i in range(num_steps):
            #print("Step {}".format(i))
            if render:
                img = self.env.render(mode)
                images.append(img[0])
            prev_states = self.split_obs(tf.reshape(prev_states, (1, -1)))
            actions = self.get_non_exploring_acts(prev_states, policy_param)

            states, rewards, dones, info = self.env.step(actions)
            rewards = tf.cast(rewards, dtype=tf.float32)
            episodic_reward += sum(rewards)

            # End this episode when `done` is True
            if self.is_done(dones):
                break
            prev_states = states
            time.sleep(waitTime)

            episode_info.append(copy.deepcopy(info))

        return states, episodic_reward, episode_info, images

    # Save all of the agents neural nets
    def save_agents(self, suffix="", policy_param="best_average"):
        for i, agent in enumerate(self.agents):
            agent.save_models(suffix=(suffix + str(i)))

    # Load the given agents neural nets
    def load_agents(self, suffix=""):
        for i, agent in enumerate(self.agents):
            agent.load_models(suffix=(suffix + str(i)))


import lump
import line
import vertex
import sqlite3
import contextlib

def create_line_schema(conn):
    conn.execute("DROP TABLE IF EXISTS Line")
    conn.execute("""
    CREATE TABLE IF NOT EXISTS Line(
      line_id INTEGER PRIMARY KEY,
      v1 INTEGER,
      v2 INTEGER,
      FOREIGN KEY(v1) REFERENCES Vertex (vertex_id),
      FOREIGN KEY(v2) REFERENCES Vertex (vertex_id)
    )""")

def create_vertex_schema(conn):
    conn.execute("DROP TABLE IF EXISTS Vertex")
    conn.execute("""
    CREATE TABLE IF NOT EXISTS Vertex(
       vertex_id INTEGER PRIMARY KEY,
       x, y INTEGER
    )""")

def _build_vertex_table(conn):
    conn.executemany("""
    INSERT INTO Vertex (vertex_id, x, y)
    VALUES (:vertex_id, :x, :y)
    """,vertex.do_read())

def _build_line_table(conn):
    fix_line = lambda line_entry: dict(line_id=line_entry.line_id,
                                       v1=line_entry.v1,
                                       v2=line_entry.v2)
    conn.executemany("""
    INSERT INTO Line (line_id, v1, v2)
    VALUES (:line_id, :v1, :v2)
    """, map(fix_line, line.do_read()))

@contextlib.contextmanager
def awesome_manager(filename, foreign_keys=True):
    with sqlite3.connect(filename) as conn:
        if foreign_keys:
            conn.execute("PRAGMA foreign_keys=ON")
        yield conn

def build_level_data():
    lump.build_lump_database()

    with awesome_manager("leveldata.db", foreign_keys=False) as conn:
        create_vertex_schema(conn)
        create_line_schema(conn)

    with awesome_manager("leveldata.db") as conn:
        _build_vertex_table(conn)
        _build_line_table(conn)

def main():
    build_level_data()

if __name__ == "__main__":
    main()


import lmfit
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import streamlit as st

# # Functions ##

# Define the model
def langmuir_model(x, Bmax, hill, kd):
    return Bmax * (x**hill) / (kd**hill + x**hill)

def langmuir_fitting(x, y):
    # Create a parameters object for the model
    params = lmfit.Parameters()
    params.add("Bmax", value=1.0, min=0)
    params.add("hill", value=1.0, min=0, max=8)
    params.add("kd", value=1.0, min=0)

    # Define the objective function to minimize (residual sum of squares)
    def objective(params, x, y):
        return (
            y - langmuir_model(x, params["Bmax"], params["hill"], params["kd"])
        ) ** 2

    # Perform the fit using the lmfit module
    result = lmfit.minimize(objective, params, args=(x, y))

    # Extract the values from the fitted parameters
    Bmax = result.params["Bmax"].value
    hill = result.params["hill"].value
    kd = result.params["kd"].value

    # Calculate the R-squared value of the fit
    y_pred = langmuir_model(x, Bmax, hill, kd)
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.average(y)) ** 2)
    r_squared = 1 - ss_res / ss_tot

    # print(f"{Bmax=}")
    # print(f"{hill=}")
    # print(f"{kd=}")
    # print(f"{r_squared=}")
    return Bmax, hill, kd, r_squared

def two_site_langmuir_model(x, Bmax1, hill1, kd1, Bmax2, hill2, kd2):
    return Bmax1 * (x**hill1) / (kd1**hill1 + x**hill1) + Bmax2 * (x**hill2) / (
        kd2**hill2 + x**hill2
    )

def two_site_langmuir_fitting(x, y):
    # Create a parameters object for the model
    params = lmfit.Parameters()
    params.add("Bmax1", value=1.0, min=0)
    params.add("hill1", value=1.0, min=0, max=8)
    params.add("kd1", value=1.0, min=0)
    params.add("Bmax2", value=1.0, min=0)
    params.add("hill2", value=1.0, min=0, max=8)
    params.add("kd2", value=1.0, min=0)

    # Define the objective function to minimize (residual sum of squares)
    def objective(params, x, y):
        return (
            y
            - two_site_langmuir_model(
                x,
                params["Bmax1"],
                params["hill1"],
                params["kd1"],
                params["Bmax2"],
                params["hill2"],
                params["kd2"],
            )
        ) ** 2

    # Perform the fit using the lmfit module
    result = lmfit.minimize(objective, params, args=(x, y))

    # Extract the values from the fitted parameters
    Bmax1 = result.params["Bmax1"].value
    hill1 = result.params["hill1"].value
    kd1 = result.params["kd1"].value
    Bmax2 = result.params["Bmax2"].value
    hill2 = result.params["hill2"].value
    kd2 = result.params["kd2"].value

    # Calculate the R-squared value of the fit
    y_pred = two_site_langmuir_model(x, Bmax1, hill1, kd1, Bmax2, hill2, kd2)
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.average(y)) ** 2)
    r_squared = 1 - ss_res / ss_tot

    # print(f"{Bmax1=}")
    # print(f"{hill1=}")
    # print(f"{kd1=}")
    # print(f"{Bmax2=}")
    # print(f"{hill2=}")
    # print(f"{kd2=}")
    # print(f"{r_squared=}")
    return Bmax1, hill1, kd1, Bmax2, hill2, kd2, r_squared

st.title("Langmuir Fitting")

st.markdown(
    """
Please upload an excel (.xlsx) file with Target [] in column A and Current response in column B, with a header row 1."""
)

uploaded_file = st.file_uploader("Choose a file")

if uploaded_file is not None:
    df = pd.read_excel(uploaded_file)
    x = df.iloc[1:, 0].values
    y = df.iloc[1:, 1].values

    Bmax, hill, kd, r_squared = langmuir_fitting(x, y)

    fake_x = np.linspace(0, max(x), 30000)
    fit_y = langmuir_model(fake_x, Bmax, hill, kd)

    with plt.style.context("seaborn-deep"):
        fig = plt.figure()
        fig.set_figheight(5)
        fig.set_figwidth(10)
        plt.rcParams.update({"font.size": 16})
        plt.plot(x, y, "o")
        plt.plot(fake_x, fit_y)
        plt.annotate(f"Kd = {kd:.2e} M", xy=(0.2, 0.5), xycoords="axes fraction")
        plt.annotate(f"hill = {hill:.2g}", xy=(0.2, 0.4), xycoords="axes fraction")
        plt.title(f"{uploaded_file.name}")
        plt.xlabel("[Target], M")
        plt.ylabel("Current, nA")
        plt.xscale("log")
        plt.figure(facecolor="white")

        st.pyplot(fig)


""" Unit test for rascil_rcal app.

"""
import os
import logging
import unittest
import shutil
import glob
import tempfile
import sys

import numpy
from astropy import units as u
from astropy.coordinates import SkyCoord

from rascil.apps.rascil_rcal import (
    cli_parser,
    rcal_simulator,
    get_gain_data,
    gt_single_plot,
    read_skycomponent_from_txt_with_external_frequency,
    _rfi_flagger,
    apply_beam_correction,
    realtime_single_bvis_solver,
)
from rascil.data_models import (
    Skycomponent,
    import_gaintable_from_hdf5,
    export_skycomponent_to_hdf5,
    rascil_path,
)
from rascil.data_models.polarisation import PolarisationFrame
from rascil.processing_components import (
    export_blockvisibility_to_ms,
    dft_skycomponent_visibility,
    create_gaintable_from_blockvisibility,
    simulate_gaintable,
    apply_gaintable,
    qa_visibility,
    qa_gaintable,
)
from rascil.processing_components.simulation import create_named_configuration
from rascil.processing_components.simulation import ingest_unittest_visibility

log = logging.getLogger("rascil-logger")
log.setLevel(logging.WARNING)
log.addHandler(logging.StreamHandler(sys.stdout))

class TestRASCILRcal(unittest.TestCase):
    def pre_setup(self, dopol=False):
        """Create and fill values into the MeassurementSet

        :param dopol: Use polarisation?
        """

        self.persist = os.getenv("RASCIL_PERSIST", False)

        self.low = create_named_configuration("LOW-AA0.5")
        # For the flagger to work, the number of channels and times should be larger than 32
        # Because the largest sequence number is 32
        self.freqwin = 40
        self.ntimes = 48
        self.times = numpy.linspace(-2.0, +2.0, self.ntimes) * numpy.pi / 12.0
        self.frequency = numpy.linspace(0.8e8, 1.2e8, self.freqwin)

        if self.freqwin > 1:
            self.channelwidth = numpy.array(
                self.freqwin * [self.frequency[1] - self.frequency[0]]
            )
        else:
            self.channelwidth = numpy.array([1e6])

        if dopol:
            self.vis_pol = PolarisationFrame("linear")
            self.image_pol = PolarisationFrame("stokesIQUV")
            f = [100.0, 20.0, 0.0, 0.0]
        else:
            self.vis_pol = PolarisationFrame("stokesI")
            self.image_pol = PolarisationFrame("stokesI")
            f = [100.0]

        self.flux = numpy.array(self.freqwin * [f])

        self.phasecentre = SkyCoord(
            ra=+180.0 * u.deg, dec=-60.0 * u.deg, frame="icrs", equinox="J2000"
        )
        self.bvis_original = ingest_unittest_visibility(
            self.low,
            self.frequency,
            self.channelwidth,
            self.times,
            self.vis_pol,
            self.phasecentre,
        )

    def makeMS(self, flux):

        comp = self.create_dft_components(flux)

        export_skycomponent_to_hdf5(
            [comp],
            self.tempdir + "/test_rascil_rcal_components.hdf",
        )

        self.bvis_error = self.create_apply_gains()

        export_blockvisibility_to_ms(
            self.tempdir + "/test_rascil_rcal.ms", [self.bvis_error]
        )

    def create_dft_components(self, flux):
        """Create the components, save to file, dft into visibility

        :param flux:

        :return pointsource: Point source skycomponent
        """
        pointsource = Skycomponent(
            direction=self.phasecentre,
            polarisation_frame=self.image_pol,
            flux=flux,
            frequency=self.frequency,
        )
        self.bvis_original = dft_skycomponent_visibility(
            self.bvis_original, pointsource
        )

        return pointsource

    def write_to_txt(self, comp):

        self.txtfile = self.tempdir + "/test_rascil_rcal_components.txt"

        coord_ra = comp.direction.ra.degree
        coord_dec = comp.direction.dec.degree
        f = open(self.txtfile, "w")
        f.write(
            "%.6f, %.6f, %10.6e, %10.6e, %10.6e, %10.6e\n"
            % (
                coord_ra,
                coord_dec,
                comp.flux[0][0],
                0.0,
                0.0,
                0.0,
            )
        )
        f.close()

    def create_apply_gains(self):
        """Create the gaintable, apply to the visibility, write as MeasurementSet

        :return: bvis_error: BlockVisibility
        """
        self.gt = create_gaintable_from_blockvisibility(
            self.bvis_original, jones_type="B"
        )
        self.gt = simulate_gaintable(self.gt, phase_error=0.1)
        qa_gt = qa_gaintable(self.gt)
        assert qa_gt.data["rms-amp"] < 1e-12, str(qa_gt)
        assert qa_gt.data["rms-phase"] > 0.0, str(qa_gt)
        bvis_error = apply_gaintable(self.bvis_original, self.gt)
        assert numpy.std(numpy.angle(bvis_error["vis"].data)) > 0.0

        return bvis_error

    def persist_data_files(self):
        """Persist the temporary data files"""
        try:
            shutil.copyfile(
                self.tempdir + "/test_rascil_rcal.ms", rascil_path("test_results")
            )
        except FileNotFoundError:
            pass

        to_copy = self.tempdir + "/test_rascil_rcal*"
        for f in glob.glob(to_copy):
            shutil.copy(f, rascil_path("test_results"))

    def setUp(self) -> None:

        self.tempdir_root = tempfile.TemporaryDirectory(dir=rascil_path("test_results"))
        self.tempdir = self.tempdir_root.name

        parser = cli_parser()
        self.args = parser.parse_args([])
        self.args.ingest_msname = self.tempdir + "/test_rascil_rcal.ms"
        self.args.ingest_components_file = (
            self.tempdir + "/test_rascil_rcal_components.hdf"
        )
        self.args.do_plotting = "False"
        self.args.plot_dir = self.tempdir + "/"
        self.args.use_previous_gaintable = "True"

    # Regression test
    def test_rcal(self):
        self.pre_setup()
        self.makeMS(self.flux)

        gtfile = rcal_simulator(self.bvis_original, self.args)

        # Check that the gaintable exists and is correct by applying it to
        # the corrupted visibility
        assert os.path.exists(gtfile)
        gain_table = import_gaintable_from_hdf5(gtfile)
        assert (
            gain_table["weight"].data != 0
        ).all()  # un-flagged data, all weights are non-zero
        log.info(f"\nFinal gaintable: {gain_table}")

        qa_gt = qa_gaintable(gain_table)
        log.info(qa_gt)
        assert qa_gt.data["rms-phase"] > 0.0, str(qa_gt)

        bvis_difference = apply_gaintable(self.bvis_error, gain_table, inverse=True)
        bvis_difference["vis"] -= self.bvis_original["vis"]
        qa = qa_visibility(bvis_difference)
        assert qa.data["maxabs"] < 1e-12, str(qa)
        assert qa.data["minabs"] < 1e-12, str(qa)

        # Test the plot does not exist
        self.plotfile = self.tempdir + "/test_rascil_rcal_plot.png"
        assert os.path.exists(self.plotfile) is False

        # Test that when we flag RFI, the results are different from
        # when we flag after gains were calculated
        os.remove(gtfile)
        self.args.flag_rfi = "True"  # flag before gains are calculated
        gtfile = rcal_simulator(self.bvis_original, self.args)
        gain_table_w_flag = import_gaintable_from_hdf5(gtfile)

        assert (gain_table_w_flag["weight"].data != gain_table["weight"].data).any()

        if self.persist is True:
            self.persist_data_files()

    def test_rcal_plot(self):
        self.pre_setup()
        self.create_dft_components(self.flux)
        self.bvis_error = self.create_apply_gains()

        self.plotfile = self.tempdir + "/test_rascil_rcal_plot.png"
        plot_name = self.plotfile.replace(".png", "")
        gt_single_plot(self.gt, plot_name=plot_name)

        assert os.path.exists(self.plotfile)

        if self.persist is True:
            self.persist_data_files()

    # Unit tests for additional functions
    def test_read_txtfile(self):
        """Test for read_skycomponent_from_txt_with_external_frequency"""
        self.pre_setup()
        comp = self.create_dft_components(self.flux)
        self.write_to_txt(comp)

        components_read = read_skycomponent_from_txt_with_external_frequency(
            self.txtfile, self.frequency, self.vis_pol
        )
        assert components_read.direction == self.phasecentre
        assert components_read.flux[:, 0].all() == self.flux.all()

        if self.persist is True:
            self.persist_data_files()

    def test_apply_beam_correction(self):
        """Test for apply_beam_correction
        Currently only test for LOW"""

        self.pre_setup()
        new_bvis = self.bvis_original.copy(deep=True)
        comp = self.create_dft_components(self.flux)
        new_comp = apply_beam_correction(new_bvis, [comp], None, telescope_name="LOW")

        assert len(new_comp) == 1
        assert new_comp[0].direction == self.phasecentre
        assert numpy.any(numpy.not_equal(new_comp[0].flux, self.flux))

        if self.persist is True:
            self.persist_data_files()

    def test_get_gain_data(self):
        self.pre_setup()
        self.create_dft_components(self.flux)
        self.bvis_error = self.create_apply_gains()

        gain_data = get_gain_data(self.gt)
        assert len(gain_data[0]) == 1  # time dimension
        assert len(gain_data[1]) == 6  # gain dimension (number of antennas)
        assert len(gain_data[2]) == 6  # phase dimension
        assert len(gain_data[3]) == 1  # residual dimension
        assert len(gain_data[4]) == 6  # weight dimension

        if self.persist is True:
            self.persist_data_files()

    def test_rfi_flagger(self):
        self.pre_setup()
        new_bvis = self.bvis_original.copy(deep=True)
        # update new_bvis to have a value that will be flagged
        new_bvis["vis"].data[0, 0, 0, 0] = complex(100, 0)

        _rfi_flagger(new_bvis)

        assert new_bvis != self.bvis_original
        # the flagger flagged that single data point
        assert new_bvis["flags"].data[0, 0, 0, 0] == 1
        # the flags array is all 0s in the original bvis
        assert (self.bvis_original["flags"] == 0).all()

        if self.persist is True:
            self.persist_data_files()

    def test_realtime_single_bvis_solver(self):
        """
        use_previous=True surfaced a bug in the code;
        this test makes sure that at least the code doesn't
        break when that argument is set.
        """
        self.pre_setup()
        bvis = self.bvis_original.copy(deep=True)
        model_components = None
        previous_solution = None
        use_previous = True

        result_gain_table, previous_solution = realtime_single_bvis_solver(
            bvis, model_components, previous_solution, use_previous=use_previous
        )

        assert previous_solution is not None

        # shape: (time, antenna, frequency, receptor1, receptor2
        # time and frequency are same as in bvis
        assert result_gain_table["gain"].shape == (
            len(bvis.time),
            6,
            len(bvis.frequency),
            1,
            1,
        )
        assert (result_gain_table["gain"].data == 1.0 + 0j).all()

        if self.persist is True:
            self.persist_data_files()

if __name__ == "__main__":
    unittest.main()


from absl.testing import absltest
from absl.testing import parameterized
import jax
import numpy as np

from tracr.compiler import compiling
from tracr.compiler import lib
from tracr.compiler import test_cases
from tracr.craft import tests_common
from tracr.rasp import rasp

_COMPILER_BOS = "rasp_to_transformer_integration_test_BOS"
_COMPILER_PAD = "rasp_to_transformer_integration_test_PAD"

# Force float32 precision on TPU, which otherwise defaults to float16.
jax.config.update("jax_default_matmul_precision", "float32")

class CompilerIntegrationTest(tests_common.VectorFnTestCase):

  def assertSequenceEqualWhenExpectedIsNotNone(self, actual_seq, expected_seq):
    for actual, expected in zip(actual_seq, expected_seq):
      if expected is not None and actual != expected:
        self.fail(f"{actual_seq} does not match (ignoring Nones) "
                  f"expected_seq={expected_seq}")

  @parameterized.named_parameters(
      dict(
          testcase_name="map",
          program=rasp.Map(lambda x: x + 1, rasp.tokens)),
      dict(
          testcase_name="sequence_map",
          program=rasp.SequenceMap(lambda x, y: x + y, rasp.tokens,
                                   rasp.indices)),
      dict(
          testcase_name="sequence_map_with_same_input",
          program=rasp.SequenceMap(lambda x, y: x + y, rasp.tokens,
                                   rasp.indices)),
      dict(
          testcase_name="select_aggregate",
          program=rasp.Aggregate(
              rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.EQ),
              rasp.Map(lambda x: 1, rasp.tokens))))
  def test_rasp_program_and_transformer_produce_same_output(self, program):
    vocab = {0, 1, 2}
    max_seq_len = 3
    assembled_model = compiling.compile_rasp_to_model(
        program, vocab, max_seq_len, compiler_bos=_COMPILER_BOS)

    test_outputs = {}
    rasp_outputs = {}
    for val in vocab:
      test_outputs[val] = assembled_model.apply([_COMPILER_BOS, val]).decoded[1]
      rasp_outputs[val] = program([val])[0]

    with self.subTest(val=0):
      self.assertEqual(test_outputs[0], rasp_outputs[0])
    with self.subTest(val=1):
      self.assertEqual(test_outputs[1], rasp_outputs[1])
    with self.subTest(val=2):
      self.assertEqual(test_outputs[2], rasp_outputs[2])

  @parameterized.named_parameters(*test_cases.TEST_CASES)
  def test_compiled_models_produce_expected_output(self, program, vocab,
                                                   test_input, expected_output,
                                                   max_seq_len, **kwargs):
    del kwargs
    assembled_model = compiling.compile_rasp_to_model(
        program, vocab, max_seq_len, compiler_bos=_COMPILER_BOS)
    test_output = assembled_model.apply([_COMPILER_BOS] + test_input)

    if isinstance(expected_output[0], (int, float)):
      np.testing.assert_allclose(
          test_output.decoded[1:], expected_output, atol=1e-7, rtol=0.005)
    else:
      self.assertSequenceEqualWhenExpectedIsNotNone(test_output.decoded[1:],
                                                    expected_output)

  @parameterized.named_parameters(*test_cases.CAUSAL_TEST_CASES)
  def test_compiled_causal_models_produce_expected_output(
      self, program, vocab, test_input, expected_output, max_seq_len, **kwargs):
    del kwargs
    assembled_model = compiling.compile_rasp_to_model(
        program,
        vocab,
        max_seq_len,
        causal=True,
        compiler_bos=_COMPILER_BOS,
        compiler_pad=_COMPILER_PAD)
    test_output = assembled_model.apply([_COMPILER_BOS] + test_input)

    if isinstance(expected_output[0], (int, float)):
      np.testing.assert_allclose(
          test_output.decoded[1:], expected_output, atol=1e-7, rtol=0.005)
    else:
      self.assertSequenceEqualWhenExpectedIsNotNone(test_output.decoded[1:],
                                                    expected_output)

  @parameterized.named_parameters(
      dict(
          testcase_name="reverse_1",
          program=lib.make_reverse(rasp.tokens),
          vocab={"a", "b", "c", "d"},
          test_input=list("abcd"),
          expected_output=list("dcba"),
          max_seq_len=5),
      dict(
          testcase_name="reverse_2",
          program=lib.make_reverse(rasp.tokens),
          vocab={"a", "b", "c", "d"},
          test_input=list("abc"),
          expected_output=list("cba"),
          max_seq_len=5),
      dict(
          testcase_name="reverse_3",
          program=lib.make_reverse(rasp.tokens),
          vocab={"a", "b", "c", "d"},
          test_input=list("ad"),
          expected_output=list("da"),
          max_seq_len=5),
      dict(
          testcase_name="reverse_4",
          program=lib.make_reverse(rasp.tokens),
          vocab={"a", "b", "c", "d"},
          test_input=["c"],
          expected_output=["c"],
          max_seq_len=5),
      dict(
          testcase_name="length_categorical_1",
          program=rasp.categorical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=list("abc"),
          expected_output=[3, 3, 3],
          max_seq_len=5),
      dict(
          testcase_name="length_categorical_2",
          program=rasp.categorical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=list("ad"),
          expected_output=[2, 2],
          max_seq_len=5),
      dict(
          testcase_name="length_categorical_3",
          program=rasp.categorical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=["c"],
          expected_output=[1],
          max_seq_len=5),
      dict(
          testcase_name="length_numerical_1",
          program=rasp.numerical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=list("abc"),
          expected_output=[3, 3, 3],
          max_seq_len=5),
      dict(
          testcase_name="length_numerical_2",
          program=rasp.numerical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=list("ad"),
          expected_output=[2, 2],
          max_seq_len=5),
      dict(
          testcase_name="length_numerical_3",
          program=rasp.numerical(lib.make_length()),
          vocab={"a", "b", "c", "d"},
          test_input=["c"],
          expected_output=[1],
          max_seq_len=5),
  )
  def test_compiled_models_produce_expected_output_with_padding(
      self, program, vocab, test_input, expected_output, max_seq_len, **kwargs):
    del kwargs
    assembled_model = compiling.compile_rasp_to_model(
        program,
        vocab,
        max_seq_len,
        compiler_bos=_COMPILER_BOS,
        compiler_pad=_COMPILER_PAD)

    pad_len = (max_seq_len - len(test_input))
    test_input = test_input + [_COMPILER_PAD] * pad_len
    test_input = [_COMPILER_BOS] + test_input
    test_output = assembled_model.apply(test_input)
    output = test_output.decoded
    output_len = len(output)
    output_stripped = test_output.decoded[1:output_len - pad_len]

    self.assertEqual(output[0], _COMPILER_BOS)
    if isinstance(expected_output[0], (int, float)):
      np.testing.assert_allclose(
          output_stripped, expected_output, atol=1e-7, rtol=0.005)
    else:
      self.assertEqual(output_stripped, expected_output)

if __name__ == "__main__":
  absltest.main()


import re
import numpy as np
from networkx import from_scipy_sparse_matrix, pagerank
from nltk.tokenize.punkt import PunktSentenceTokenizer
from sklearn.feature_extraction.text import TfidfTransformer, CountVectorizer
from Sastrawi.StopWordRemover.StopWordRemoverFactory import ArrayDictionary, StopWordRemoverFactory, StopWordRemover

class TextSummarization():
    def __init__(self):
        with open('./stopwords.txt') as f:
            more_stopword=f.read().split('\n')

        SWfactory = StopWordRemoverFactory()
        stopword_data = ArrayDictionary(more_stopword+SWfactory.get_stop_words())
        self.stopword = StopWordRemover(stopword_data)

    def Preprocessing(self, text):
        clean = re.sub("#[^\W]+|@[^\W]+|http[^*\s]+|<[^>]*>|[0-9]", '', text) #cleansing data
        emoticons = re.findall('(?::|;|=)()(?:-)?(?:\)|\(|D|P)', clean)
        text = (re.sub('[\W]+', ' ', clean.lower()) +  #Case folding
                ' '.join(emoticons).replace('-', ''))
        result=''
        for kata in text.split():
            stop = self.stopword.remove(kata) #Stopword
            result += f"{stop} " if stop else ''
        return result

    def Summary(self, doc, preprocess=False):
        doc_tokenizer = PunktSentenceTokenizer()
        sentences_list = doc_tokenizer.tokenize(doc)

        clean_sentences_list=[]
        for sentence in sentences_list:
            clean_sentences_list.append(self.Preprocessing(sentence))

        cv = CountVectorizer()
        cv_matrix = cv.fit_transform(clean_sentences_list if preprocess else sentences_list)
        normal_matrix = TfidfTransformer().fit_transform(cv_matrix)

        tfidf=normal_matrix.toarray()
        res_graph = normal_matrix * normal_matrix.T # similaritas /adjacency matrix

        nx_graph= from_scipy_sparse_matrix(res_graph)
        pageranks = pagerank(nx_graph)

        sentence_array = sorted(((pageranks[i], s) for i, s in enumerate(sentences_list)), reverse=True)
        sentence_array = np.asarray(sentence_array)

        rank_max = float(sentence_array[0][0])
        rank_min = float(sentence_array[len(sentence_array) - 1][0])

        temp_array = []

        # Jika semua rank sama
        # taking any sentence will give the summary, say the first sentence
        flag = 0
        if rank_max - rank_min == 0:
            temp_array.append(0)
            flag = 1

        # If the sentence has different ranks
        if flag != 1:
            for i in range(0, len(sentence_array)):
                temp_array.append((float(sentence_array[i][0]) - rank_min) / (rank_max - rank_min))

        # Calculation of threshold:
        # We take the mean value of normalized scores
        # any sentence with the normalized score 0.2 more than the mean value is considered to be
        threshold = (sum(temp_array) / len(temp_array))# + 0.2

        # Separate out the sentences that satiasfy the criteria of having a score above the threshold
        sentence_list = []
        if len(temp_array) > 1:
            for i in range(0, len(temp_array)):
                if temp_array[i] > threshold:
                        sentence_list.append(sentence_array[i][1])
        else:
            sentence_list.append(sentence_array[0][1])

        summary = " ".join(str(x) for x in sentence_list)

        return summary


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# df = pd.read_csv ('output.csv', header = None)
# df.columns = ['Name', 'ELO', 'K']

# df.plot(x="Name", y=["ELO", "K"], kind="bar", figsize=(9, 8))

# plt.show()
# print(df)

df = pd.read_csv ('output.csv', header = None)
df.columns = ['Name', 'ELO', 'K']

ax = df.plot(x="Name", y=["ELO", "K"], kind="bar", figsize=(9, 8))

s = df["K"]

for index in range((df.shape[0])):
  ax.text(index, s[index], s[index], size=8)

plt.show()
print(s)


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

#本代码用于生成并保存最后一个轮次的全局模型和各个client本地模型

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
import copy
from torch.optim import lr_scheduler

# 载入总的训练集和测试集
x_train_all = np.load('../dataset/cifar10/x_train_cifar10_all.npy')
y_train_all = np.load('../dataset/cifar10/y_train_cifar10_all.npy')

x_test_all = np.load('../dataset/cifar10/x_test_cifar10_all.npy')
y_test_all = np.load('../dataset/cifar10/y_test_cifar10_all.npy')

##############
# 定义网络
class Net(nn.Module):
# 网络组成子模块的定义
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

        self.grad_list = []
        self.layer_list = []
        self.layer_output_list = []

    # 前向过程的定义 网络流程的定义
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5) #相当于 reshape 成 16 * 5 * 5 列， 行数自适应
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    # 用于测试集时，记录，避免在训练时占用过多计算机内存
    def eval_layer_output(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        self.layer_output_list.append(x)
        x = self.pool(F.relu(self.conv2(x)))
        self.layer_output_list.append(x)
        x = x.view(-1, 16 * 5 * 5)
        self.layer_output_list.append(x)
        x = F.relu(self.fc1(x))
        self.layer_output_list.append(x)
        x = F.relu(self.fc2(x))
        self.layer_output_list.append(x)
        x = self.fc3(x)
        self.layer_output_list.append(x)
        return self.layer_output_list, x

    # 根据泄露的层数确定攻击模型实际获取的特征数据
    def get_grad(self, list=[]):
        return self.grad_list

net_g = Net()
#################
# 联邦学习过程

### 参数设置：
# client 数量
num_client = 10

# 全局训练轮次
g_iteration_count = 10
# 本地训练轮次
l_iteration_count = 10

for g_epoch in range(g_iteration_count):
    # 在每个大的全局轮次中
    # 将当前全局轮次训练好的本地训练放入该数组中
    cur_all_client = []

    # 首先每个client下载当前的全局模型
    for client_i in range(num_client):
         print('######each client start to download global net_g ######')
         print('######this is client_{} in g_count_{}######'.format(client_i,g_epoch))
         cur_net = copy.deepcopy(net_g)

         #载入当前client的数据集
         x_train = np.load('../dataset/cifar10/x_train_cifar10_{}.npy'.format(client_i))
         y_train = np.load('../dataset/cifar10/y_train_cifar10_{}.npy'.format(client_i))

         x_test = np.load('../dataset/cifar10/x_test_cifar10_{}.npy'.format(client_i))
         y_test = np.load('../dataset/cifar10/y_test_cifar10_{}.npy'.format(client_i))

         # 转成Tensor形式
         x_train = np.array(x_train)
         x_train = torch.FloatTensor(x_train)

         y_train = np.array(y_train)
         y_train = torch.LongTensor(y_train).squeeze()

         # 进行batch划分
         torch_dataset = torch.utils.data.TensorDataset(x_train, y_train)
         train_loader = torch.utils.data.DataLoader(torch_dataset,
                                                   batch_size=64,
                                                  shuffle=False)
         len_cur_x_train = len(x_train)
         # 开始本地轮次的训练
         loss_func = torch.nn.CrossEntropyLoss()
         model_opt = optim.SGD(params=cur_net.parameters(),lr=0.005)
         # sgd,lr=0.005 10个全局轮次，聚合模型的预测准确率为0.46

         scheduler = lr_scheduler.LambdaLR(optimizer=model_opt, lr_lambda=lambda epoch:0.95**epoch)

        # 本地训练
         for epoch in range(l_iteration_count):
             # 基于划分的batch训练
             for step, data in enumerate(train_loader, start=0):
                 x_feature, y_feature = data # 解构出特征和标签
                 output = cur_net(x_feature)     # input x and predict based on x
                 loss = loss_func(output, y_feature)     # must be (1. nn output, 2. target)

                 model_opt.zero_grad()   # clear gradients for next train
                 loss.backward()         # backpropagation, compute gradients
                 model_opt.step()        # apply gradients

                 if step % (int(len_cur_x_train/64)) == (int(len_cur_x_train/64))-1: #
                     prediction = torch.max(output, 1)[1]#只返回最大值的每个索引,标签
                     pred = prediction.data.numpy()
                     target = y_feature.data.numpy()
                     accuracy = float((pred == target).astype(int).sum()) / float(target.size)
                     print("epoch={},step={} train_result model{} = {}".format(epoch, step, client_i, accuracy))
             scheduler.step()
         print('save current client_{}'.format(client_i))
         #测试一下此时的测试准确率
         # 转成Tensor形式
         x_test = np.array(x_test)
         x_test = torch.FloatTensor(x_test)

         y_test = np.array(y_test)
         y_test = torch.LongTensor(y_test)

         out_test = cur_net(x_test)     # input x and predict based on x
         prediction = torch.max(out_test, 1)[1]
         pred_y = prediction.data.numpy()
         target_y = y_test.data.numpy()
         accuracy = float((pred_y == target_y).astype(int).sum()) / float(target_y.size)
         print('test_result_currrent_client_model =', accuracy)

          ## 保存模板模型
         torch.save(cur_net, '../model/cifar10/client_{}_g_count_{}.pkl'.format(client_i,g_epoch))  # save entire net
         # torch.save(cur_net.state_dict(),'../model/cifar10/client_{}_g_count_{}_params.pkl'.format(client_i, g_epoch))#网络的参数

         # 将训练好的client model保存在数据中，用于后续的第一轮聚合
         cur_all_client.append(cur_net)

    # 进行全局模型的聚合
    print('##### start global_model_once_in_g_count_{}'.format(g_epoch))

    # 参数平均聚合,更新net_g
    # 根据模型层数，确定需要聚合的每层参数的个数
    item_name_arr = []
    # 前面的卷积层
    item_name_arr.append('conv1.weight')
    item_name_arr.append('conv1.bias')
    item_name_arr.append('conv2.weight')
    item_name_arr.append('conv2.bias')
    # 全连接层有三层
    num_model_layer =  3
    for i in range(num_model_layer):
        item_name_arr.append('fc{}.weight'.format(i+1))
    for i in range(num_model_layer):
        item_name_arr.append('fc{}.bias'.format(i+1))

    # 赋值net_g参数,这样才能更新修改net_g的参数，因为net_g.state_dict()返回的是一个深copy的副本
    model_dict  = net_g.state_dict()

    # 进行参数平均
    for i in range(len(item_name_arr)):
        cur_name = item_name_arr[i]
    #    print('net_g avg before')
    #    test_before = model_dict[cur_name]
    #    print(test_before)

        for j in range(len(cur_all_client)):
    #        print(j)
            net_cur_ele = cur_all_client[j]
    #        print(net_cur_ele.state_dict()[cur_name])
            if j == 0 :
                sum_item = net_cur_ele.state_dict()[cur_name]
            else:
                sum_item = sum_item + net_cur_ele.state_dict()[cur_name]
        # 参数平均
        model_dict[cur_name] = sum_item/(num_client*1.0)

    # 更新聚合的net_g
    net_g.load_state_dict(model_dict)
    print('##### finish global_model_once_in_g_count_{}'.format(g_epoch))

    ## 保存每次聚合的global model
    torch.save(net_g, '../model/cifar10/global_model_g_count_{}.pkl'.format(g_epoch))  # save entire net
    # torch.save(net_g.state_dict(),'../model/cifar10/global_model_g_count_{}_params.pkl'.format(g_epoch))#网络的参数

    # 测试一下此时聚合模型的测试准确度
    # 转为torch.tensor的格式
    x_test_g = np.array(x_test_all)
    x_test = torch.FloatTensor(x_test_g)
    y_test_g = np.array(y_test_all)
    y_test = torch.LongTensor(y_test_g)

    out_test = net_g(x_test)     # input x and predict based on x
    prediction = torch.max(out_test, 1)[1]
    pred_y = prediction.data.numpy()
    target_y = y_test.data.numpy()
    accuracy = float((pred_y == target_y).astype(int).sum()) / float(target_y.size)
    print("################# test_result_net_g_model_in_the_g_count_{} = ".format(g_epoch),accuracy)
#    print(net_g.state_dict())


from django.contrib import admin

from .models import Track, TrackTranslation

class TrackTranslationInline(admin.TabularInline):
    model = TrackTranslation
    fields = ('locale', 'title', 'description')
    ordering = ('locale',)
    min_num = 2
    extra = 0

@admin.register(Track)
class TrackAdmin(admin.ModelAdmin):
    list_display = ('slug', 'order', 'soon')
    list_filter = ('soon', )
    inlines = [
        TrackTranslationInline,
    ]


from collections import Counter
import bisect

w, h = map(int, input().split())

n = int(input())

points = []

for _ in range(n):
    x, y = map(int, input().split())
    points.append((x, y))

nxcuts = int(input())

xcuts = [0] + list(map(int, input().split())) + [w]

nycuts = int(input())

ycuts = [0] + list(map(int, input().split())) + [h]

ctr = Counter()

minval = float('inf')
maxval = float('-inf')

for x, y in points:
    xx = bisect.bisect_left(xcuts, x)
    yy = bisect.bisect_left(ycuts, y)
    ctr[(xx, yy)] += 1
    maxval = max(maxval, ctr[(xx, yy)])

for xx, yy in ctr:
    minval = min(minval, ctr[(xx, yy)])

if len(ctr) < (nxcuts + 1) * (nycuts + 1) or n < (nxcuts + 1) * (nycuts + 1):
    minval = 0

print(minval, maxval)


import os
import cv2
import tqdm
from datetime import datetime
import argparse

from datasets.prepare_custom import get_custom_dataset, get_target_labels
from detectron2.engine import DefaultTrainer, DefaultPredictor
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2 import model_zoo
from detectron2.utils.logger import setup_logger
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.checkpoint import DetectionCheckpointer

logger = setup_logger()

parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, default='', help='ex) train')
parser.add_argument('--model', type=str, default='', help='ex) output_20200624_1655')

target = 'multi'

def set_dataset():
    for d in ["train", "test"]:
        DatasetCatalog.register("custom/" + d, lambda d=d: get_custom_dataset(target, "custom/" + d))
        MetadataCatalog.get("custom/" + d).set(thing_classes=get_target_labels()[target])

def set_config_custom(output_folder, mode, now_date):
    logger.debug(f" Set configs")
    config_custom = {
        'DATASET': '20200625_103328_8posture',
        'AUG_IN_DATASET': 'X',
        'CLASS_LIST': str(get_target_labels()[target]),
        'NUM_WORKERS': 2,
        # 2
        # Number of data loading threads

        'IMS_PER_BATCH': 2,
        # 2
        # Number of images per batch across all machines.
        # If we have 16 GPUs and IMS_PER_BATCH = 32,
        # each GPU will see 2 images per batch.

        'BASE_LR': 0.00025,
        # 0.00025

        'MAX_ITER': 30000,
        # 40000

        'BATCH_SIZE_PER_IMAGE': 512,
        # 256
        # RoI minibatch size *per image* (number of regions of interest [ROIs])
        # Total number of RoIs per training minibatch =
        #   ROI_HEADS.BATCH_SIZE_PER_IMAGE * SOLVER.IMS_PER_BATCH
        # E.g., a common configuration is: 512 * 16 = 8192
        'MODEL_WEIGHTS': 'COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml',
        'TEST_EVAL_PERIOD': 100
    }
    os.makedirs(os.path.join(output_folder, f'{now_date}_{mode}'), exist_ok=True)
    seq_path = os.path.join(output_folder, f'{now_date}_{mode}', f'config_custom.txt')
    wf = open(seq_path, 'w', newline='', encoding='utf-8')
    for k, v in config_custom.items():
        wf.write(f'{str(k)}: {str(v)}\n')
    wf.write(f'\n')
    wf.write(f'TARGET_OBJECT: {str(target)}\n')
    wf.write(f'START_TIME: {str(datetime.now())}\n')
    wf.close()

    return config_custom

def set_config_all(cfg, config_custom, output_folder, mode, now_date):
    ## 1. https://rosenfelder.ai/Instance_Image_Segmentation_for_Window_and_Building_Detection_with_detectron2/
    cfg.merge_from_file(f'./configs/{config_custom["MODEL_WEIGHTS"]}')
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_custom["MODEL_WEIGHTS"])

    ## 2. https://towardsdatascience.com/detectron2-the-basic-end-to-end-tutorial-5ac90e2f90e3
    # cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
    # cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"

    cfg.DATASETS.TRAIN = ("custom/train",)
    cfg.DATASETS.TEST = ()
    # cfg.DATASETS.TEST = ()  # no metrics implemented for this dataset
    cfg.DATALOADER.NUM_WORKERS = config_custom['NUM_WORKERS']
    cfg.SOLVER.IMS_PER_BATCH = (
        config_custom['IMS_PER_BATCH']
    )
    cfg.SOLVER.BASE_LR = (
        config_custom['BASE_LR']
    )
    # cfg.SOLVER.WARMUP_ITERS = 1000 # the learning rate starts from 0 and goes to the preset one for this number of iterations
    cfg.SOLVER.MAX_ITER = (
        config_custom['MAX_ITER']
    )
    # cfg.SOLVER.STEPS = (1000, 1500) # the checkpoints (number of iterations) at which the learning rate will be reduced by GAMMA
    # cfg.SOLVER.GAMMA = 0.05
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (
        config_custom['BATCH_SIZE_PER_IMAGE']
    )
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(get_target_labels()[target])  # 4 action class (none,eating,peeing,pooping)
    # cfg.MODEL.DEVICE = "cuda:0,1"
    cfg.TEST.EVAL_PERIOD = config_custom['TEST_EVAL_PERIOD']

    cfg.OUTPUT_DIR=output_folder
    ## write config list
    seq_path = os.path.join(output_folder, f'{now_date}_{mode}', f'config_all.txt')
    wf = open(seq_path, 'w', newline='', encoding='utf-8')
    wf.write(f'{str(cfg)}')
    wf.close()

    return cfg

def train_model(cfg, output_folder, mode, now_date):
    logger.debug(f" Start train model")
    trainer = DefaultTrainer(cfg)
    trainer.resume_or_load(resume=False)
    trainer.train()

    # trainer.register_hooks(
    #     [hooks.EvalHook(0, lambda: test_with_TTA(cfg, trainer.model))]
    # )
    # 실시간 로스율 계산하기 참고: https://medium.com/@apofeniaco/training-on-detectron2-with-a-validation-set-and-plot-loss-on-it-to-avoid-overfitting-6449418fbf4e

    seq_path = os.path.join(output_folder, f"{now_date}_{mode}", f'config_custom.txt')
    wf = open(seq_path, 'a', newline='', encoding='utf-8')
    wf.write(f'END_TIME_TRAIN: {str(datetime.now())}\n')
    logger.debug(f" End to train:{str(datetime.now())}")
    wf.close()

    return True

def make_prediction_image(cfg, output_folder, mode, now_date):
    logger.debug(f" Start evaluating image")
    os.makedirs(os.path.join(output_folder, f"{now_date}_{mode}", f"prediction_{target}"), exist_ok=True)

    cfg.MODEL.WEIGHTS = os.path.join(output_folder, "model_final.pth")
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set the testing threshold for this model
    cfg.DATASETS.TEST = ("custom/test",)

    # Evaluate test datasets
    model = DefaultTrainer(cfg).build_model(cfg)
    DetectionCheckpointer(model, save_dir=output_folder).resume_or_load(cfg.MODEL.WEIGHTS, resume=False)

    evaluator = COCOEvaluator("custom/test", cfg, False,
                              output_dir=os.path.join(output_folder, f"{now_date}_{mode}", 'COCO_eval'))
    val_loader = build_detection_test_loader(cfg, "custom/test")
    order_dicts = inference_on_dataset(DefaultTrainer(cfg).model, val_loader, evaluator)

    seq_path = os.path.join(output_folder, f"{now_date}_{mode}", f'eval_result.txt')
    wf = open(seq_path, 'w', newline='', encoding='utf-8')
    wf.write(f'{str(order_dicts)}:\n')
    wf.close()

    # Predict test dataset
    predictor = DefaultPredictor(cfg)
    dataset_dicts = get_custom_dataset(target, "custom/test")
    for i, d in enumerate(dataset_dicts):
        logger.debug(f"predicting...{d['file_name']}")
        im = cv2.imread(d["file_name"])
        file_name = os.path.basename(d["file_name"])
        outputs = predictor(im)

        ## 1 https://rosenfelder.ai/Instance_Image_Segmentation_for_Window_and_Building_Detection_with_detectron2/
        v = Visualizer(
            im[:, :, ::-1],
            metadata=MetadataCatalog.get("custom/test"),
            scale=1,
            instance_mode=ColorMode.IMAGE_BW,  # remove the colors of unsegmented pixels
        )

        v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        cv2.imwrite(os.path.join(output_folder, f"{now_date}_{mode}", f"prediction_{target}", file_name),
                    v.get_image())

    logger.debug(f" End evaluating image")

    seq_path = os.path.join(output_folder, f"{now_date}_{mode}", f'config_custom.txt')
    wf = open(seq_path, 'a', newline='', encoding='utf-8')
    wf.write(f'END_TIME_EVAL_IMAGE: {str(datetime.now())}\n')
    logger.debug(f" End to train:{str(datetime.now())}")
    wf.close()

    return True

# https://stackoverflow.com/questions/60663073/how-can-i-properly-run-detectron2-on-videos
def make_prediction_video(cfg, output_folder, mode, now_date):
    logger.debug(f" Start evaluating video")

    # Extract video properties
    video = cv2.VideoCapture('./test_video.mp4')
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frames_per_second = video.get(cv2.CAP_PROP_FPS)
    num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

    # Initialize video writer
    video_writer = cv2.VideoWriter(os.path.join(output_folder, f"{now_date}_{mode}", f'test_result.mp4'),
                                   fourcc=cv2.VideoWriter_fourcc(*"mp4v"), fps=float(frames_per_second),
                                   frameSize=(width, height), isColor=True)

    # Initialize predictor
    cfg.MODEL.WEIGHTS = os.path.join(output_folder, "model_final.pth")
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set the testing threshold for this model
    cfg.DATASETS.TEST = ("custom/test",)
    predictor = DefaultPredictor(cfg)

    # Initialize visualizer
    v = VideoVisualizer(MetadataCatalog.get("custom/test"), ColorMode.IMAGE)

    # Create a cut-off for debugging
    # num_frames = 5000

    # Enumerate the frames of the video
    for visualization in tqdm.tqdm(run_on_video(video, v, num_frames, predictor), total=num_frames):
        # Write test image
        # cv2.imwrite(os.path.join(output_folder, 'test_result.png'), visualization)

        # Write to video file
        video_writer.write(visualization)

    # Release resources
    video.release()
    video_writer.release()
    # cv2.destroyAllWindows()
    logger.debug(f" End evaluating video")

    seq_path = os.path.join(output_folder, f"{now_date}_{mode}", f'config_custom.txt')
    wf = open(seq_path, 'a', newline='', encoding='utf-8')
    wf.write(f'END_TIME_EVAL_VIDEO: {str(datetime.now())}\n')
    logger.debug(f" End to train:{str(datetime.now())}")
    wf.close()

    return True

def run_on_video(video, v, maxFrames, predictor):
    """ Runs the predictor on every frame in the video (unless maxFrames is given),
    and returns the frame with the predictions drawn.
    """
    readFrames = 0
    while True:
        hasFrame, frame = video.read()

        if not hasFrame:
            break

        # Get prediction results for this frame
        outputs = predictor(frame)
        # Make sure the frame is colored
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        # Draw a visualization of the predictions using the video visualizer
        visualization = v.draw_instance_predictions(frame, outputs["instances"].to("cpu"))
        # Convert Matplotlib RGB format to OpenCV BGR format
        visualization = cv2.cvtColor(visualization.get_image(), cv2.COLOR_RGB2BGR)
        yield visualization
        readFrames += 1
        if readFrames > maxFrames:
            break

def init(mode, model=None):
    cfg_origin = get_cfg()
    today_date = datetime.now().strftime("%Y%m%d_%H%M")
    if model is not None and model != '':
        output_folder = os.path.splitext(cfg_origin.OUTPUT_DIR)[0].replace('output', model)
    else:
        output_folder = os.path.join(f'{cfg_origin.OUTPUT_DIR}_{today_date}_{target}')
        os.makedirs(output_folder, exist_ok=True)

    config_custom = set_config_custom(output_folder, mode, now_date)
    cfg = set_config_all(cfg_origin, config_custom, output_folder, mode, now_date)

    return cfg, output_folder

if __name__ == "__main__":
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
    args = parser.parse_args()
    mode = args.mode
    model = args.model
    now_date = datetime.now().strftime("%Y%m%d_%H%M")
    cfg, output_folder = init(mode, model)

    set_dataset()
    if mode == 'train':
        train_model(cfg, output_folder, mode, now_date)
        make_prediction_image(cfg, output_folder, mode, now_date)
        make_prediction_video(cfg, output_folder, mode, now_date)
    elif mode == 'predict_image':
        make_prediction_image(cfg, output_folder, mode, now_date)
    elif mode == 'predict_video':
        make_prediction_video(cfg, output_folder, mode, now_date)


class Solution:
    def mostCommonWord(self, paragraph: str, banned: List[str]) -> str:
        from collections import Counter

        chars = "!?',;."
        for c in chars:
            paragraph = paragraph.replace(c, ' ')
        arr = paragraph.lower().split()

        counter = Counter(arr)
        mylist = counter.most_common()

        for elem in mylist:
            if elem[0] in banned:
                continue
            return elem[0]


from discord.ext.commands import Cog
from discord.ext.commands import command
from discord import Embed

class Tiktok(Cog):
  def __init__(self, bot):
    self.bot = bot

  @command()
  async def tiktok(self, ctx, link, *args):
    title = str(" ".join(args))

    self.ctx = ctx
    embed = Embed(title=f"{title}",colour=0x03A4DB)

    fields = [(f"link:", f"{link}", False)]

    for name, value, inline in fields:
        embed.add_field(name=name, value=value, inline=inline)

    embed.set_author(name='pececito', icon_url=self.ctx.guild.icon_url)
    embed.set_image(url="https://cdn.discordapp.com/attachments/859880030015062047/862374250526539806/image0.png")

    await ctx.channel.send(embed=embed)

  @Cog.listener()
  async def on_ready(self):
    if not self.bot.ready:
      self.bot.cogs_ready.ready_up("tiktok")

def setup(bot):
  bot.add_cog(Tiktok(bot))


#!/usr/bin/env python

from __future__ import print_function

import argparse
import torch
import torch.nn.functional as F
from torch.nn import Parameter
from nn.conv import HGATConv
from utils import parse_feat

class UnsupervisedHGAT(torch.jit.ScriptModule):

    def __init__(self, in_u_dim, in_i_dim, u_embedding_dim, i_embedding_dim,
                 u_field_num, i_field_num, hidden, out_dim, heads, dropout,
                 negative_size, encode):
        super(UnsupervisedHGAT, self).__init__()
        # bipartite graph, two types of nodes, namely u,i. Edges: (u, i)

        self.user_input_dim = in_u_dim
        self.item_input_dim = in_i_dim
        self.user_embedding_dim = u_embedding_dim
        self.item_embedding_dim = i_embedding_dim
        self.user_field_num = u_field_num
        self.item_field_num = i_field_num
        self.heads = heads
        self.dropout = dropout
        self.negative_size = negative_size
        self.encode = encode

        # the transformed dense feature dimension for user and item
        in_u_dim = u_embedding_dim * u_field_num if u_embedding_dim > 0 else in_u_dim
        in_i_dim = i_embedding_dim * i_field_num if i_embedding_dim > 0 else in_i_dim

        self.convs_u = torch.nn.ModuleList()
        self.convs_i = torch.nn.ModuleList()

        # u-i-u
        self.convs_u.append(HGATConv(hidden, hidden, hidden, heads=heads, dropout=dropout))
        self.convs_u.append(HGATConv(hidden * self.heads, hidden * self.heads,
                                     hidden, heads=heads, dropout=dropout, act=True))

        # i-u-i
        self.convs_i.append(HGATConv(hidden, hidden, hidden, heads=heads, dropout=dropout))
        self.convs_i.append(HGATConv(hidden * self.heads, hidden * self.heads,
                                     hidden, heads=heads, dropout=dropout, act=True))

        # weight
        self.weight_u1 = Parameter(torch.zeros(in_u_dim, hidden))
        self.weight_i1 = Parameter(torch.zeros(in_i_dim, hidden))

        self.weight_u2 = Parameter(torch.zeros(hidden + hidden * self.heads * 2, out_dim))
        self.weight_i2 = Parameter(torch.zeros(hidden + hidden * self.heads * 2, out_dim))

        self.weight = Parameter(torch.zeros(out_dim, 2))
        self.bias = Parameter(torch.zeros(2))

        self.reset_parameters()

    def reset_parameters(self):
        torch.nn.init.xavier_uniform_(self.weight_u1)
        torch.nn.init.xavier_uniform_(self.weight_i1)
        torch.nn.init.xavier_uniform_(self.weight_u2)
        torch.nn.init.xavier_uniform_(self.weight_i2)
        torch.nn.init.xavier_uniform_(self.weight)
        if self.bias.dim() > 1:
            torch.nn.init.xavier_uniform_(self.bias)

    @torch.jit.script_method
    def forward_(self, pos_u, pos_i, first_u_edge_index, second_u_edge_index,
                 first_i_edge_index, second_i_edge_index,
                 u_batch_ids=torch.Tensor([]), i_batch_ids=torch.Tensor([]),
                 u_field_ids=torch.Tensor([]), i_field_ids=torch.Tensor([])):

        # encode u, i feat
        pos_u, pos_i = self.encode_feat(pos_u, pos_i, u_batch_ids,
                                        i_batch_ids, u_field_ids, i_field_ids)

        pos_u_embedding = self.embedding__(pos_u, pos_i,
                                           first_u_edge_index,
                                           first_i_edge_index,
                                           "u")
        pos_i_embedding = self.embedding__(pos_i, pos_u,
                                           second_u_edge_index,
                                           second_i_edge_index,
                                           "i")

        u_index, i_index = first_u_edge_index[0], first_u_edge_index[1]

        u_emb, i_emb = pos_u_embedding[u_index], pos_i_embedding[i_index]

        # out size: (edge_num, out_dim)
        return u_emb, i_emb

    @torch.jit.script_method
    def encode_feat(self, u, i,
                    u_batch_ids=torch.Tensor([]), i_batch_ids=torch.Tensor([]),
                    u_field_ids=torch.Tensor([]), i_field_ids=torch.Tensor([])):

        if self.user_field_num > 0:
            u = parse_feat(u, u_batch_ids, u_field_ids, self.user_field_num, self.encode)
            i = parse_feat(i, i_batch_ids, i_field_ids, self.item_field_num, self.encode)

        u = torch.matmul(u, self.weight_u1)
        i = torch.matmul(i, self.weight_i1)

        return u, i

    @torch.jit.script_method
    def loss(self, pos_u, pos_i):
        # (edge_num, out_dim)
        i_emb_origin = pos_i.repeat(1, 1)
        i_emb = pos_i
        # ((edge_num * (negative_size + 1)), out_dim)
        u_emb = pos_u.repeat(self.negative_size + 1, 1)

        for _ in range(self.negative_size):
            neg_index = torch.randperm(i_emb_origin.size(0)).view(-1).to(torch.long)
            neg_sample_emb = torch.index_select(i_emb_origin, 0, neg_index)
            i_emb = torch.cat([i_emb, neg_sample_emb], dim=0)

        # ((edge_num * (negative_size + 1)), 1)
        prod = torch.sum(torch.mul(u_emb, i_emb), 1, True) / 0.05

        # (edge_num, (negative_size + 1))
        logits = torch.transpose(torch.transpose(prod, 0, 1)
                                 .reshape(self.negative_size + 1, -1), 0, 1)

        true_label = torch.zeros(pos_u.size(0)).view(-1).to(torch.long)

        loss = F.cross_entropy(logits, true_label)

        acc_per_example = torch.eq(
            torch.argmax(logits, dim=1),
            torch.zeros([pos_u.size(0)], dtype=torch.int64)).view(-1).to(torch.float)

        loss1 = loss.mean()
        accuracy = acc_per_example.mean()
        print("loss: ", loss1, ", acc: ", accuracy)

        return loss

    @torch.jit.script_method
    def embedding__(self, u, i, first_u_edge_index, first_i_edge_index, conv="u"):
        # type: (Tensor, Tensor, Tensor, Tensor, str) -> Tensor

        # gat conv encode
        graph_user_emb = self.encode_conv(u, i, first_u_edge_index, first_i_edge_index, conv)

        user_emb = torch.cat([u[0: graph_user_emb.size(0)], graph_user_emb], dim=1)

        # conv for u, i
        if conv == "u":
            user_emb = torch.matmul(user_emb, self.weight_u2)
        else:
            user_emb = torch.matmul(user_emb, self.weight_i2)

        user_emb = F.normalize(user_emb, p=2.0)

        # out size: (node_num, out_dim)
        return user_emb

    @torch.jit.script_method
    def encode_conv(self, u, i, first_edge_index, second_edge_index, conv):
        # type: (Tensor, Tensor, Tensor, Tensor, str) -> Tensor

        # convs = self.convs_u
        if conv == "u":
            convs = self.convs_u
            z1 = convs[0](i, u, second_edge_index)
            z2 = convs[0](u, i, first_edge_index)
            z3 = convs[1](z2, z1, first_edge_index)
        else:
            convs = self.convs_i
            z1 = convs[0](i, u, second_edge_index)
            z2 = convs[0](u, i, first_edge_index)
            z3 = convs[1](z2, z1, first_edge_index)

        midh0 = torch.cat([z2, z3], dim=1)

        # out size: (node_num, (hidden * heads * 2))
        return midh0

    @torch.jit.script_method
    def user_embedding_(self, u, i, first_u_edge_index, first_i_edge_index,
                        u_batch_ids=torch.Tensor([]), i_batch_ids=torch.Tensor([]),
                        u_field_ids=torch.Tensor([]), i_field_ids=torch.Tensor([])):

        if self.user_field_num > 0:
            u, i = self.encode_feat(u, i, u_batch_ids, i_batch_ids, u_field_ids, i_field_ids)

        user_emb = self.embedding__(u, i, first_u_edge_index, first_i_edge_index, "u")

        # out size: (node_num, out_dim)
        return user_emb

    @torch.jit.script_method
    def item_embedding_(self, u, i, first_u_edge_index, first_i_edge_index,
                        u_batch_ids=torch.Tensor([]), i_batch_ids=torch.Tensor([]),
                        u_field_ids=torch.Tensor([]), i_field_ids=torch.Tensor([])):

        if self.user_field_num > 0:
            u, i = self.encode_feat(u, i, u_batch_ids, i_batch_ids, u_field_ids, i_field_ids)

        item_emb = self.embedding__(i, u, first_u_edge_index, first_i_edge_index, "i")

        # out size: (node_num, out_dim)
        return item_emb

FLAGS = None

def main():
    bisage = UnsupervisedHGAT(FLAGS.input_user_dim, FLAGS.input_item_dim,
                              FLAGS.input_user_embedding_dim,
                              FLAGS.input_item_embedding_dim,
                              FLAGS.input_user_field_num,
                              FLAGS.input_item_field_num,
                              FLAGS.hidden_dim, FLAGS.output_dim,
                              FLAGS.heads, FLAGS.dropout, FLAGS.negative_size,
                              FLAGS.encode)
    bisage.save(FLAGS.output_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_user_dim",
        type=int,
        default=-1,
        help="input dimention of user node features")
    parser.add_argument(
        "--input_item_dim",
        type=int,
        default=-1,
        help="input dimention of item node features")
    parser.add_argument(
        "--input_user_embedding_dim",
        type=int,
        default=-1,
        help="embedding dim of user node features")
    parser.add_argument(
        "--input_item_embedding_dim",
        type=int,
        default=-1,
        help="embedding dim of item node features")
    parser.add_argument(
        "--input_user_field_num",
        type=int,
        default=-1,
        help="field num of user node features")
    parser.add_argument(
        "--input_item_field_num",
        type=int,
        default=-1,
        help="field num of item node features")
    parser.add_argument(
        "--hidden_dim",
        type=int,
        default=-1,
        help="hidden dimension of bipartite graphsage convolution layer")
    parser.add_argument(
        "--output_dim",
        type=int,
        default=-1,
        help="output embedding dimension")
    parser.add_argument(
        "--heads",
        type=int,
        default=1,
        help="the number of heads")
    parser.add_argument(
        "--dropout",
        type=float,
        default=0.0,
        help="the percentage of dropout")
    parser.add_argument(
        "--negative_size",
        type=int,
        default=1,
        help="the negative sample size")
    parser.add_argument(
        "--encode",
        type=str,
        default="dense",
        help="data encode, could be one-hot, multi-hot or dense")
    parser.add_argument(
        "--output_file",
        type=str,
        default="unsupervised_bipartite_graphsage.pt",
        help="output file name")
    FLAGS, unparsed = parser.parse_known_args()
    main()


#! /usr/bin/env python3

"""
To run this demo program, hook up an ADXL345 to the Analog Discovery as follows:

   Ground <----> Ground
   V+     <----> VCC
   DIO0   -----> CS
   (NC)   ------ INT1
   (NC)   ------ INT2
   DIO3   -----> SDO
   DIO2   -----> SDA
   DIO1   -----> SCL
"""

import time
import argparse

from pydwf import DigilentWaveformsLibrary, DwfDigitalOutIdle
from demo_utilities import find_demo_device, DemoDeviceNotFoundError

def set_positive_supply_voltage(analogIO, voltage: float):
    analogIO.channelNodeSet(0, 0, 1)
    analogIO.channelNodeSet(0, 1, voltage)
    analogIO.enableSet(1)

def demo_spi_protocol_api(spi):

    """Demonstrate the SPI protocol functionality by connecting to an ADXL345 accellerometer sensor.

    The Digital SPI protocol API has 19 methods.
    """

    # We use the ADXL345 in four-wire SPI mode (MISO and MOSI are separate lines).

    SPI_CSn  = 0  # SPI chip-select [CS ] (DIO channel 0)
    SPI_SCLK = 1  # SPI clock       [SCL] (DIO channel 1)
    SPI_MOSI = 2  # SPI MOSI        [SDA] (DIO channel 2)
    SPI_MISO = 3  # SPI MISO        [SDO] (DIO channel 3)

    SPI_TRANSFERTYPE_MOSI_MISO = 1
    SPI_BITS_PER_WORD = 8

    SPI_CSn_START = 0
    SPI_CSn_STOP  = 1

    SPI_MODE = 3

    SPI_BITORDER = 1 # MSB first

    spi.reset()

    spi.frequencySet(5000000.0) # Maximum clock frequency according to datasheet.

    spi.clockSet(SPI_SCLK)      # Select clock pin
    spi.dataSet(0, SPI_MOSI)    # Select MOSI pin
    spi.dataSet(1, SPI_MISO)    # Select MISO pin

    spi.idleSet(0, DwfDigitalOutIdle.High) # Set MOSI pin idle mode
    spi.idleSet(1, DwfDigitalOutIdle.High) # Set MISO pin idle mode

    spi.modeSet(SPI_MODE)       # CPOL=1, CPHA=1

    spi.orderSet(SPI_BITORDER)  # Send MSB first

    time.sleep(0.100)

    print()

    # Enable measurements (set 'Measure' bit in POWER_CTL register).
    # First byte, high two bits indicate a 'write' operation of a single byte.
    spi.select(SPI_CSn, 0)   # Set chip-select to 0
    response = spi.writeRead(SPI_TRANSFERTYPE_MOSI_MISO, SPI_BITS_PER_WORD, [0x00 | 0x2d, 8])
    spi.select(SPI_CSn, 1)   # Set chip-select to 1

    # Loop until interrupted.
    while True:

        # Perform readout of registers 0x32 .. 0x37.
        # First byte, high two bits indicate a 'read' operation of multiple bytes.
        spi.select(SPI_CSn, 0)   # Set chip-select to 0
        response = spi.writeRead(SPI_TRANSFERTYPE_MOSI_MISO, SPI_BITS_PER_WORD, [0xc0 | 0x32, 0, 0, 0, 0, 0, 0])
        spi.select(SPI_CSn, 1)   # Set chip-select to 1

        axisdata = response[1:7]

        ax = (axisdata[0] + axisdata[1] * 256 + 32768) % 65536 - 32768
        ay = (axisdata[2] + axisdata[3] * 256 + 32768) % 65536 - 32768
        az = (axisdata[4] + axisdata[5] * 256 + 32768) % 65536 - 32768

        print("\r[SPI] ADXL345: ax {:6} ay {:6} az {:6}".format(ax, ay, az), end="", flush=True)

def main():

    parser = argparse.ArgumentParser(description="Demonstrate usage of the SPI protocol API with an ADXL345 accellerometer.")
    parser.add_argument('serial_number', nargs='?', help="serial number of the Digilent device")

    args = parser.parse_args()

    try:
        dwf = DigilentWaveformsLibrary()
        with find_demo_device(dwf, args.serial_number) as device:
            set_positive_supply_voltage(device.analogIO, 3.3)
            demo_spi_protocol_api(device.digitalSpi)
    except DemoDeviceNotFoundError:
        print("Could not find demo device, exiting.")
    except KeyboardInterrupt:
        print("Keyboard interrupt, ending demo.")

if __name__ == "__main__":
    main()


from src.examples.bls.bls_utils import *
from py_ecc import bn128

def key_gen(privkey=9078879759034877109383407430680751358246, ec=bn128) -> (int, (FQ2, FQ2)):
    pubkey = ec.multiply(ec.G2, privkey)
    return privkey, pubkey

def sign(message, privkey: int, ec=bn128) -> (FQ, FQ):
    digest = hash_to_ec_g1(message)
    return ec.multiply(digest, privkey)

def verify(message, signature: (FQ, FQ), pubkey: (FQ2, FQ2), ec=bn128) -> bool:
    digest = hash_to_ec_g1(message)
    lhs = ec.pairing(pubkey, digest)
    rhs = ec.pairing(ec.G2, signature)
    return lhs == rhs

if __name__ == '__main__':

    # KeyGen() -> (sk, pk)
    sk, pk = key_gen()

    # Sign(msg, sk) -> sig
    sig = sign("hello, world", sk)

    # Verify(msg, sig, pk) -> True/False
    validity = verify("hello, world", sig, pk)

    print(validity)


from typing import List

class Solution:
    def sortArrayByParity(self, nums: List[int]) -> List[int]:
        e = []
        odd = []

        for n in nums:
            if(n%2 == 0):
                e.append(n)
            else:
                odd.append(n)
        for o in odd:
            e.append(o)

        return e

nums = [3,1,2,4]
#nums = [0]

print(Solution.sortArrayByParity(Solution, nums))