from collections import OrderedDict
import time, os
curdir = os.path.expanduser("~")
import numpy as np
np.set_printoptions(precision=8, suppress=True, linewidth=400, threshold=100)
import tensorflow as tf
tf.keras.backend.set_floatx('float64')
# tf.config.run_functions_eagerly(True)
tf.keras.backend.set_epsilon(tf.experimental.numpy.finfo(tf.keras.backend.floatx()).eps) # 1e-7 default
import matplotlib.pyplot as plt
import gymnasium as gym
import gym_util, model_util as util, model_nets as nets
# CUDA 11.8.0_522.06, CUDNN 8.6.0.163, tensorflow-gpu==2.10.0, tensorflow_probability==0.18.0
physical_devices_gpu = tf.config.list_physical_devices('GPU')
for i in range(len(physical_devices_gpu)): tf.config.experimental.set_memory_growth(physical_devices_gpu[i], True)

class GeneralAI(tf.keras.Model):
    def __init__(self, arch, env, trader, env_render, save_model, chkpts, max_episodes, max_steps, learn_rates, latent_size, latent_dist, mixture_multi, net_lstm, net_attn, aio_max_latents, aug_data_step, aug_data_pos):
        super(GeneralAI, self).__init__()
        compute_dtype = tf.dtypes.as_dtype(self.compute_dtype)
        self.float_max = tf.constant(compute_dtype.max, compute_dtype)
        self.float_maxroot = tf.constant(tf.math.sqrt(self.float_max), compute_dtype)
        self.float_eps = tf.constant(tf.experimental.numpy.finfo(compute_dtype).eps, compute_dtype)
        self.float_eps_max = tf.constant(1.0 / self.float_eps, compute_dtype)
        self.loss_scale = tf.math.exp(tf.math.log(self.float_eps_max) * (1/2))
        self.float64_zero = tf.constant(0, tf.float64)

        self.arch, self.env, self.trader, self.env_render, self.save_model = arch, env, trader, env_render, save_model
        self.chkpts, self.max_episodes, self.max_steps, self.learn_rates = tf.constant(chkpts, tf.int32), tf.constant(max_episodes, tf.int32), tf.constant(max_steps, tf.int32), {}
        for k,v in learn_rates.items(): self.learn_rates[k] = tf.constant(v, compute_dtype)
        self.initializer = tf.keras.initializers.GlorotUniform(time.time_ns())

        self.obs_spec, self.obs_zero, _ = gym_util.get_spec(env.observation_space, space_name='obs', compute_dtype=self.compute_dtype, net_attn_io=net_attn['io'], aio_max_latents=aio_max_latents, mixture_multi=mixture_multi)
        self.action_spec, _, self.action_zero_out = gym_util.get_spec(env.action_space, space_name='actions', compute_dtype=self.compute_dtype, mixture_multi=mixture_multi)
        self.obs_spec_len, self.action_spec_len = len(self.obs_spec), len(self.action_spec)
        self.action_total_size = tf.constant(np.sum([np.prod(feat['step_shape']) for feat in self.action_spec]),compute_dtype)
        self.gym_step_shapes = [feat['step_shape'] for feat in self.obs_spec] + [tf.TensorShape((1,1)), tf.TensorShape((1,1)), tf.TensorShape((1,2)) if trader else tf.TensorShape((1,1))]
        self.gym_step_dtypes = [feat['dtype'] for feat in self.obs_spec] + [tf.float64, tf.bool, tf.float64]
        self.rewards_zero, self.dones_zero = tf.constant([[0]],tf.float64), tf.constant([[False]],tf.bool)
        self.step_zero, self.step_size_one = tf.constant([[0]]), tf.constant([[1]])

        latent_spec = {'dtype':compute_dtype, 'latent_size':latent_size, 'num_latents':1, 'max_latents':aio_max_latents}
        latent_spec.update({'inp':512, 'midp':256, 'outp':512, 'evo':64})
        if latent_dist == 'd': latent_spec.update({'dist_type':'d', 'num_components':latent_size, 'event_shape':(latent_size,)}) # deterministic
        if latent_dist == 'c': latent_spec.update({'dist_type':'c', 'num_components':0, 'event_shape':(latent_size, latent_size)}) # categorical # TODO https://keras.io/examples/generative/vq_vae/
        if latent_dist == 'mx': latent_spec.update({'dist_type':'mx', 'num_components':int(latent_size/16), 'event_shape':(latent_size,)}) # continuous

        if aug_data_step: self.obs_spec += [{'space_name':'step', 'name':'', 'event_shape':(1,), 'event_size':1, 'channels':1, 'step_shape':tf.TensorShape((1,1)), 'num_latents':1}]
        self.obs_spec += [{'space_name':'reward_prev', 'name':'', 'event_shape':(1,), 'event_size':1, 'channels':1, 'step_shape':tf.TensorShape((1,1)), 'num_latents':1}]
        inputs = {'obs':self.obs_zero, 'step':[self.step_zero], 'reward_prev':[self.rewards_zero], 'return_goal':[self.rewards_zero]}

        if arch in ('PG',):
            opt_spec = [{'name':'action', 'type':'a', 'schedule_type':'', 'learn_rate':self.learn_rates['action'], 'float_eps':self.float_eps}]
            stats_spec = [{'name':'rwd', 'b1':0.99, 'b2':0.99, 'dtype':tf.float64}, {'name':'loss', 'b1':0.99, 'b2':0.99, 'dtype':compute_dtype}, {'name':'delta', 'b1':0.99, 'b2':0.99, 'dtype':compute_dtype}]
            self.action = nets.ArchFull('A', inputs, opt_spec, stats_spec, self.obs_spec, self.action_spec, latent_spec, obs_latent=False, net_blocks=2, net_lstm=net_lstm, net_attn=net_attn, num_heads=4, memory_size=max_steps, latent_multi=1, aug_data_pos=aug_data_pos)
            outputs = self.action(inputs)
            self.action.optimizer_weights = util.optimizer_build(self.action.optimizer['action'], self.action.trainable_variables)
            util.net_build(self.action, self.initializer)

        self.metrics_spec()
        self.reset_states = tf.function(self.reset_states, experimental_autograph_options=tf.autograph.experimental.Feature.LISTS)
        self.reset_states()
        arch_run = getattr(self, arch); arch_run = tf.function(arch_run, experimental_autograph_options=tf.autograph.experimental.Feature.LISTS); setattr(self, arch, arch_run)

    def metrics_spec(self):
        metrics_loss = OrderedDict()
        metrics_loss['2rewards*'] = {'-rewards_ma':np.float64, '-rewards_total+':np.float64, 'rewards_final=':np.float64}
        metrics_loss['1steps'] = {'steps+':np.int64}
        if self.arch == 'PG':
            metrics_loss['1nets*'] = {'-loss_ma':np.float64, '-loss_action':np.float64}
            metrics_loss['1extras'] = {'loss_action_returns':np.float64}
            metrics_loss['1extras2*'] = {'actlog0':np.float64, 'actlog1':np.float64}

        for loss_group in metrics_loss.values():
            for k in loss_group.keys():
                if k.endswith('=') or k.endswith('+'): loss_group[k] = [0 for i in range(self.max_episodes)]
                else: loss_group[k] = [[] for i in range(self.max_episodes)]
        self.metrics_loss = metrics_loss

    def metrics_update(self, *args):
        args = list(args)
        log_metrics, episode, idx = args[0], args[1], 2
        for loss_group in self.metrics_loss.values():
            for k in loss_group.keys():
                if log_metrics[idx-2]:
                    if k.endswith('='): loss_group[k][episode] = args[idx]
                    elif k.endswith('+'): loss_group[k][episode] += args[idx]
                    else: loss_group[k][episode] += [args[idx]]
                idx += 1
        return np.asarray(0, np.int32) # dummy

    def env_reset(self, dummy):
        obs, info = self.env.reset(); reward, done = 0.0, False
        if self.env_render: self.env.render()
        if hasattr(self.env,'np_struc'): rtn = gym_util.struc_to_feat(obs)
        else: rtn = gym_util.space_to_feat(obs, self.env.observation_space)
        metrics = info['metrics'] if 'metrics' in info else [0]
        rtn += [np.asarray([[reward]], np.float64), np.asarray([[done]], bool), np.asarray([metrics], np.float64)]
        return rtn
    def env_step(self, *args): # args = tuple of ndarrays
        if hasattr(self.env,'np_struc'): action = gym_util.out_to_struc(list(args), self.env.action_dtype)
        else: action = gym_util.out_to_space(args, self.env.action_space, [0])
        obs, reward, terminated, truncated, info = self.env.step(action); done = (terminated or truncated)
        if self.env_render: self.env.render()
        if hasattr(self.env,'np_struc'): rtn = gym_util.struc_to_feat(obs)
        else: rtn = gym_util.space_to_feat(obs, self.env.observation_space)
        metrics = info['metrics'] if 'metrics' in info else [0]
        rtn += [np.asarray([[reward]], np.float64), np.asarray([[done]], bool), np.asarray([metrics], np.float64)]
        return rtn

    def checkpoints(self, *args):
        model_files = ""
        for net in self.layers:
            model_file = self.model_files[net.name]
            net.save_weights(model_file)
            model_files += ' '+model_file.split('/')[-1]
        print("SAVED{}".format(model_files))
        return np.asarray(0, np.int32) # dummy

    def reset_states(self, use_img=False):
        for net in self.layers:
            if hasattr(net, 'reset_states'): net.reset_states(use_img=use_img)

    def PG_actor(self, inputs, return_goal):
        print("tracing -> GeneralAI PG_actor")
        obs, actions = [None]*self.obs_spec_len, [None]*self.action_spec_len
        for i in range(self.obs_spec_len): obs[i] = tf.TensorArray(self.obs_spec[i]['dtype'], size=1, dynamic_size=True, infer_shape=False, element_shape=self.obs_spec[i]['step_shape'][1:])
        for i in range(self.action_spec_len): actions[i] = tf.TensorArray(self.action_spec[i]['dtype_out'], size=1, dynamic_size=True, infer_shape=False, element_shape=self.action_spec[i]['step_shape'][1:])
        rewards = tf.TensorArray(tf.float64, size=1, dynamic_size=True, infer_shape=False, element_shape=(1,))
        dones = tf.TensorArray(tf.bool, size=1, dynamic_size=True, infer_shape=False, element_shape=(1,))
        returns = tf.TensorArray(tf.float64, size=0, dynamic_size=True, infer_shape=False, element_shape=(1,))

        step = tf.constant(0)
        while not inputs['dones'][-1][0]:
            for i in range(self.obs_spec_len): obs[i] = obs[i].write(step, inputs['obs'][i][-1])

            action = [None]*self.action_spec_len
            inputs_step = {'obs':inputs['obs'], 'step':[tf.reshape(step,(1,1))], 'reward_prev':[inputs['rewards']], 'return_goal':[return_goal]}
            action_logits = self.action(inputs_step)
            for i in range(self.action_spec_len):
                action_dist = self.action.dist[i](action_logits[i])
                action[i] = action_dist.sample()

            action_dis = [None]*self.action_spec_len
            for i in range(self.action_spec_len):
                actions[i] = actions[i].write(step, action[i][0])
                action_dis[i] = util.discretize(action[i][0], self.action_spec[i])

            np_in = tf.numpy_function(self.env_step, action_dis, self.gym_step_dtypes)
            for i in range(len(np_in)): np_in[i].set_shape(self.gym_step_shapes[i])
            inputs['obs'], inputs['rewards'], inputs['dones'] = np_in[:-3], np_in[-3], np_in[-2]

            rewards = rewards.write(step, inputs['rewards'][-1])
            dones = dones.write(step, inputs['dones'][-1])
            returns = returns.write(step, [self.float64_zero])
            returns_updt = returns.stack()
            returns_updt = returns_updt + inputs['rewards'][-1]
            returns = returns.unstack(returns_updt)
            step += 1

        outputs = {}
        out_obs, out_actions = [None]*self.obs_spec_len, [None]*self.action_spec_len
        for i in range(self.obs_spec_len): out_obs[i] = obs[i].stack()
        for i in range(self.action_spec_len): out_actions[i] = actions[i].stack()
        outputs['obs'], outputs['actions'], outputs['rewards'], outputs['dones'], outputs['returns'] = out_obs, out_actions, rewards.stack(), dones.stack(), returns.stack()
        return outputs, inputs

    def PG_learner_onestep(self, inputs, training=True):
        print("tracing -> GeneralAI PG_learner_onestep")
        loss = {}
        loss_actions_lik = tf.TensorArray(self.compute_dtype, size=1, dynamic_size=True, infer_shape=False, element_shape=(1,))
        loss_actions = tf.TensorArray(self.compute_dtype, size=1, dynamic_size=True, infer_shape=False, element_shape=(1,))
        metric_actlog = tf.TensorArray(self.compute_dtype, size=1, dynamic_size=True, infer_shape=False, element_shape=(2,))

        inputs_rewards = tf.concat([self.rewards_zero, inputs['rewards']], axis=0)
        returns = inputs['returns'][0:1]
        for step in tf.range(tf.shape(inputs['dones'])[0]):
            obs = [None]*self.obs_spec_len
            for i in range(self.obs_spec_len): obs[i] = inputs['obs'][i][step:step+1]; obs[i].set_shape(self.obs_spec[i]['step_shape'])
            action = [None]*self.action_spec_len
            for i in range(self.action_spec_len): action[i] = inputs['actions'][i][step:step+1]; action[i].set_shape(self.action_spec[i]['step_shape'])
            returns_calc = tf.squeeze(tf.cast(returns,self.compute_dtype),axis=-1)

            inputs_step = {'obs':obs, 'step':[tf.reshape(step,(1,1))], 'reward_prev':[inputs_rewards[step:step+1]], 'return_goal':[returns]}
            with tf.GradientTape() as tape_action:
                action_logits = self.action(inputs_step)
                action_dist = [None]*self.action_spec_len
                for i in range(self.action_spec_len):
                    action_dist[i] = self.action.dist[i](action_logits[i])
                loss_action_lik = util.loss_likelihood(action_dist, action)
                loss_action = loss_action_lik * returns_calc
            if loss_action_lik > self.float_eps:
                gradients = tape_action.gradient(loss_action, self.action.trainable_variables)
                self.action.optimizer['action'].apply_gradients(zip(gradients, self.action.trainable_variables))
            loss_actions_lik = loss_actions_lik.write(step, loss_action_lik / self.action_total_size)
            loss_actions = loss_actions.write(step, loss_action)
            metric_actlog = metric_actlog.write(step, action_logits[0][0][0:2])

        loss['action_lik'], loss['action'], loss['actlog'] = loss_actions_lik.concat(), loss_actions.concat(), metric_actlog.stack()
        return loss

    def PG(self):
        print("tracing -> GeneralAI PG"); tf.print("RUNNING")
        return_goal, ma, ma_loss, snr_loss = tf.constant([[-self.loss_scale.numpy()]], tf.float64), tf.constant(0,tf.float64), self.float_maxroot, tf.constant(1,self.compute_dtype)
        episode = tf.constant(0)
        while episode < self.max_episodes:
            tf.autograph.experimental.set_loop_options(parallel_iterations=1)
            np_in = tf.numpy_function(self.env_reset, [tf.constant(0)], self.gym_step_dtypes)
            for i in range(len(np_in)): np_in[i].set_shape(self.gym_step_shapes[i])
            inputs = {'obs':np_in[:-3], 'rewards':np_in[-3], 'dones':np_in[-2]}

            self.reset_states(); outputs, inputs = self.PG_actor(inputs, return_goal)
            rewards_total = outputs['returns'][0][0]
            util.stats_update(self.action.stats['rwd'], rewards_total); avg, ma, ema, snr, std = util.stats_get(self.action.stats['rwd'])

            self.reset_states(); loss = self.PG_learner_onestep(outputs)
            util.stats_update(self.action.stats['loss'], tf.math.reduce_mean(loss['action_lik'])); avg_loss, ma_loss, ema_loss, snr_loss, std_loss = util.stats_get(self.action.stats['loss'])
            self.action.optimizer['action'].learning_rate = self.learn_rates['action'] * snr_loss

            log_metrics = [True,True,True,True,True,True,True,True,True,True,True,True,True,True]
            metrics = [log_metrics, episode, ma, tf.math.reduce_sum(outputs['rewards']), outputs['rewards'][-1][0], tf.shape(outputs['rewards'])[0],
                ma_loss, tf.math.reduce_mean(loss['action_lik']),
                tf.math.reduce_mean(loss['action']),
                tf.math.reduce_mean(loss['actlog'][:,0]), tf.math.reduce_mean(loss['actlog'][:,1]),
            ]
            dummy = tf.numpy_function(self.metrics_update, metrics, [tf.int32])

            if self.save_model:
                if episode > tf.constant(0) and episode % self.chkpts == tf.constant(0): tf.numpy_function(self.checkpoints, [tf.constant(0)], [tf.int32])
            episode += 1

def params(): pass
load_model, save_model, chkpts = False, False, 5000
max_episodes = 100
latent_size = 16
latent_dist = 'd' # 'd' = deterministic, 'c' = categorical, 'mx' = continuous(mixture)
mixture_multi = 4 # mixture distribution size,multiply num components
net_lstm = False
net_attn = {'net':True, 'io':True, 'out':True, 'ar':True}
aio_max_latents = 16
aug_data_step, aug_data_pos = True, True

device_type = 'GPU' # use GPU for large networks (over 8 total net blocks?) or output data (512 bytes?)
device_type = 'CPU'

machine, device, extra = 'dev', 0, ''

trader, env_async, env_async_clock, env_async_speed, env_reconfig, chart_lim = False, False, 0.001, 160.0, False, 0.003
env_name, max_steps, env_render, env_reconfig, env = 'CartPole', 512, False, True, gym.make('CartPole-v1', max_episode_steps=500) # (4) float32    ()2 int64    500  475.0
# env_name, max_steps, env_render, env_reconfig, env = 'LunarLand', 1024, False, True, gym.make('LunarLander-v2') # (8) float32    ()4 int64    1000  200.0

arch = 'PG'; learn_rates = {'action':4e-6} # Policy Gradient agent, PG loss

if __name__ == '__main__':
    if env_async: import envs_local.async_wrapper as envaw_; env_name, env = env_name+'-asyn', envaw_.AsyncWrapperEnv(env, env_async_clock, env_async_speed, env_render)
    if env_reconfig: import envs_local.reconfig_wrapper as envrw_; env_name, env = env_name+'-r', envrw_.ReconfigWrapperEnv(env)
    with tf.device("/device:{}:{}".format(device_type,(device if device_type=='GPU' else 0))):
        model = GeneralAI(arch, env, trader, env_render, save_model, chkpts, max_episodes, max_steps, learn_rates, latent_size, latent_dist, mixture_multi, net_lstm, net_attn, aio_max_latents, aug_data_step, aug_data_pos)
        name = "gym-{}-{}-{}-a{}{}-{}".format(arch, env_name, machine, device, extra, time.strftime("%y-%m-%d-%H-%M-%S"))

        ## load models
        model_files, name_arch = {}, ""
        for net in model.layers:
            model_name = "{}-{}-a{}".format(net.arch_desc, machine, device)
            model_file = "{}/tf-data-models-local/{}.h5".format(curdir, model_name); loaded_model = False
            model_files[net.name] = model_file
            if (load_model or net.name == 'M') and tf.io.gfile.exists(model_file):
                net.load_weights(model_file, by_name=True, skip_mismatch=True)
                print("LOADED {} weights from {}".format(net.name, model_file)); loaded_model = True
            name_opt = "-O{}{}".format(net.opt_spec['type'], ('' if net.opt_spec['schedule_type']=='' else '-S'+net.opt_spec['schedule_type'])) if hasattr(net, 'opt_spec') else ''
            name_arch += "   {}{}-{}".format(net.arch_desc, name_opt, 'load' if loaded_model else 'new')
        model.model_files = model_files

        ## run
        print("RUN {}\n{}".format(name, name_arch))
        arch_run = getattr(model, arch)
        t1_start = time.perf_counter_ns()
        arch_run()
        total_time = (time.perf_counter_ns() - t1_start) / 1e9 # seconds
        env.close()

        ## metrics
        metrics_loss = model.metrics_loss
        for loss_group in metrics_loss.values():
            for k in loss_group.keys():
                for j in range(len(loss_group[k])):
                    loss_group[k][j] = 0 if loss_group[k][j] == [] else np.mean(loss_group[k][j])

        total_steps = int(np.nansum(metrics_loss['1steps']['steps+']))
        step_time = total_time/total_steps
        learn_rates_txt, attn_txt = "", ""
        for k,v in learn_rates.items(): learn_rates_txt += "  {}:{:.0e}".format(k,v)
        for k,v in net_attn.items(): attn_txt += " {}".format(k) if v else ''
        title = "{}    [{}-{}]\n{}\ntime:{}    steps:{}    t/s:{:.8f}    ms:{}".format(name, device_type, tf.keras.backend.floatx(), name_arch, util.print_time(total_time), total_steps, step_time, max_steps)
        title += "     |     attn:{}    al:{}".format(attn_txt, aio_max_latents)
        title += "     |     a-clk:{}    a-spd:{}    aug:{}{}".format(env_async_clock, env_async_speed, ('S' if aug_data_step else ''), ('P' if aug_data_pos else ''))
        title += "     |   {}".format(learn_rates_txt); print(title)

        import matplotlib as mpl
        mpl.rcParams['axes.prop_cycle'] = mpl.cycler(color=['blue','lightblue','green','lime','red','lavender','turquoise','cyan','magenta','salmon','yellow','gold','black','brown','purple','pink','orange','teal','coral','darkgreen','tan'])
        plt.figure(num=name, figsize=(34, 18), tight_layout=True)
        xrng, i, vplts = np.arange(0, max_episodes, 1), 0, 0
        for loss_group_name in metrics_loss.keys(): vplts += int(loss_group_name[0])

        for loss_group_name, loss_group in metrics_loss.items():
            rows, col, m_min, m_max, combine, yscale = int(loss_group_name[0]), 0, [0]*len(loss_group), [0]*len(loss_group), loss_group_name.endswith('*'), ('log' if loss_group_name[1] == '~' else 'linear')
            if combine: spg = plt.subplot2grid((vplts, 1), (i, 0), rowspan=rows, xlim=(0, max_episodes), yscale=yscale); plt.grid(axis='y',alpha=0.3)
            for metric_name, metric in loss_group.items():
                metric = np.asarray(metric, np.float64); m_min[col], m_max[col] = np.nanquantile(metric, chart_lim), np.nanquantile(metric, 1.0-chart_lim)
                if not combine: spg = plt.subplot2grid((vplts, len(loss_group)), (i, col), rowspan=rows, xlim=(0, max_episodes), ylim=(m_min[col], m_max[col]), yscale=yscale); plt.grid(axis='y',alpha=0.3,which='both')
                if metric_name.startswith('-'): plt.plot(xrng, metric, alpha=1.0, label=metric_name)
                else: plt.plot(xrng, util.ewma(metric, window=max_episodes//10+2), alpha=1.0, label=metric_name); plt.plot(xrng, metric, alpha=0.3)
                plt.ylabel('value'); plt.legend(loc='upper left'); col+=1
            if combine: spg.set_ylim(np.min(m_min), np.max(m_max))
            if i == 0: plt.title(title)
            i+=rows
        out_file = "output/{}.png".format(name)
        plt.savefig(out_file)
        print("SAVED {}   (run \033[94mpython serve.py\033[00m to access webserver on port 8080)".format(out_file))

        ## save models
        if save_model:
            for net in model.layers:
                model_file = model.model_files[net.name]
                net.save_weights(model_file)
                print("SAVED {} weights to {}".format(net.name, model_file))


import sys
import argparse

from collections import defaultdict

from PolTools.utils.build_counts_dict import build_counts_dict
from PolTools.utils.constants import hg38_chrom_sizes_file

def get_possible_tss_dict(five_prime_counts_dict, count_threshold):
    # Filter all positions that are below the threshold
    filtered_dict = defaultdict(
        lambda: {
            "+": [],
            "-": []
        }
    )

    for chromosome in five_prime_counts_dict:
        for strand in five_prime_counts_dict[chromosome]:
            for position, counts in five_prime_counts_dict[chromosome][strand].items():
                if counts >= count_threshold:
                    filtered_dict[chromosome][strand].append(
                        [chromosome, position, position + 1, "TSS", counts, strand]
                    )

            # Sort by the number of counts
            filtered_dict[chromosome][strand].sort(key=lambda x: x[-2], reverse=True)

    return filtered_dict

def find_tsrs_per_chrom(possible_tss_dict, radius):
    # Go through each chromosome and pick the largest. Go +- radius bp to make a region size of 1 + (2*radius).
    # Remove any TSSs within that that region and then pick the next max
    identified_tsrs = []
    tsr_number = 0

    for chrom in possible_tss_dict:
        for strand in possible_tss_dict[chrom]:

            tss_heights = possible_tss_dict[chrom][strand]
            while tss_heights:
                # Choose the first one since it is the max
                curr_max_tss = tss_heights[0]
                chrom, left, right, name, height, strand = curr_max_tss

                new_tsr = [chrom, left - radius, right + radius, "TSR" + str(tsr_number), height, strand]
                new_tsr_chrom, new_tsr_left, new_tsr_right, new_tsr_name, height, new_tsr_strand = new_tsr
                identified_tsrs.append(new_tsr)

                # Eliminate overlapping TSSs
                non_overlapping_tss = []
                for tss in tss_heights:
                    tss_chrom, tss_left, tss_right, tss_name, tss_score, tss_strand = tss

                    if tss_left < new_tsr_left or tss_left >= new_tsr_right:
                        # If there is no overlap between the newly defined TSR and this TSS, add it to the non overlapping
                        non_overlapping_tss.append(tss)

                tss_heights = non_overlapping_tss
                tsr_number += 1

    return identified_tsrs

def remove_invalid_tsrs(found_tsrs):
    # First load the chromosome sizes
    valid_tsrs = []

    chrom_sizes_dict = {}
    with open(hg38_chrom_sizes_file) as file:
        for line in file:
            chrom, size = line.split()
            size = int(size)
            chrom_sizes_dict[chrom] = size

    for tsr in found_tsrs:
        chrom, left, right, name, score, strand = tsr

        if chrom in chrom_sizes_dict and int(right) < chrom_sizes_dict[chrom] and int(left) >= 0:
            valid_tsrs.append(tsr)

    return valid_tsrs

def output_tsrs(found_tsrs, count_threshold, dataset):
    # Write the TSRs to a file
    output_filename = dataset.rsplit(".")[0] + "_min_" + str(count_threshold) + "-TSR.bed"

    with open(output_filename, 'w') as file:
        for tsr in found_tsrs:
            file.write(
                "\t".join(
                    [str(val) for val in tsr]
                ) + "\n"
            )

def positive_int(num):
    try:
        val = int(num)
        if val <= 0:
            raise Exception("Go to the except")
    except:
        raise argparse.ArgumentTypeError(num + " must be positive")

    return val

def parse_args(args):
    parser = argparse.ArgumentParser(prog='PolTools tsrPicker',
                                     description="Generate transcription start regions\n" +
                                                 "More information can be found at " +
                                                 "https://geoffscollins.github.io/PolTools/tsrPicker.html")

    parser.add_argument('seq_file', metavar='seq_file', type=str, help='Bed formatted sequencing file to find the TSRs')
    parser.add_argument('min_seq_depth', metavar='min_seq_depth', type=positive_int,
                        help="Minimum number of 5' ends to be considered a TSR")
    parser.add_argument('-r', '--radius', dest='radius', metavar='radius', type=positive_int, nargs='?', default=5)

    parsed_args = parser.parse_args(args)
    return parsed_args.seq_file, parsed_args.min_seq_depth, parsed_args.radius

def main(args):
    dataset, tss_strength_threshold, radius = parse_args(args)

    # Convert the thresholds to integers
    tss_strength_threshold = int(tss_strength_threshold)

    five_prime_counts_dict = build_counts_dict(dataset, 'five')

    # Get the counts for each of those positions in each of the tbp datasets
    possible_tss_dict = get_possible_tss_dict(five_prime_counts_dict, tss_strength_threshold)

    # Go through each chromosome and find tsrs
    found_tsrs = find_tsrs_per_chrom(possible_tss_dict, radius)

    # Remove tsrs that go past the genome
    final_tsrs = remove_invalid_tsrs(found_tsrs)

    output_tsrs(final_tsrs, tss_strength_threshold, dataset)

if __name__ == '__main__':
    main(sys.argv[1:])


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from tensor2tensor.utils import metrics

import tensorflow.compat.v1 as tf

class MetricsTest(tf.test.TestCase):

  def testAccuracyMetric(self):
    predictions = np.random.randint(1, 5, size=(12, 12, 12, 1))
    targets = np.random.randint(1, 5, size=(12, 12, 12, 1))
    expected = np.mean((predictions == targets).astype(float))
    with self.test_session() as session:
      scores, _ = metrics.padded_accuracy(
          tf.one_hot(predictions, depth=5, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32))
      a = tf.reduce_mean(scores)
      session.run(tf.global_variables_initializer())
      actual = session.run(a)
    self.assertAlmostEqual(actual, expected)

  def testAccuracyTopKMetric(self):
    predictions = np.random.randint(1, 5, size=(12, 12, 12, 1))
    targets = np.random.randint(1, 5, size=(12, 12, 12, 1))
    expected = np.mean((predictions == targets).astype(float))
    with self.test_session() as session:
      predicted = tf.one_hot(predictions, depth=5, dtype=tf.float32)
      scores1, _ = metrics.padded_accuracy_topk(
          predicted, tf.constant(targets, dtype=tf.int32), k=1)
      scores2, _ = metrics.padded_accuracy_topk(
          predicted, tf.constant(targets, dtype=tf.int32), k=7)
      a1 = tf.reduce_mean(scores1)
      a2 = tf.reduce_mean(scores2)
      session.run(tf.global_variables_initializer())
      actual1, actual2 = session.run([a1, a2])
    self.assertAlmostEqual(actual1, expected)
    self.assertAlmostEqual(actual2, 1.0)

  def testPrefixAccuracy(self):
    vocab_size = 10
    predictions = tf.one_hot(
        tf.constant([[[1], [2], [3], [4], [9], [6], [7], [8]],
                     [[1], [2], [3], [4], [5], [9], [7], [8]],
                     [[1], [2], [3], [4], [5], [9], [7], [0]]]),
        vocab_size)
    labels = tf.expand_dims(
        tf.constant([[[1], [2], [3], [4], [5], [6], [7], [8]],
                     [[1], [2], [3], [4], [5], [6], [7], [8]],
                     [[1], [2], [3], [4], [5], [6], [7], [0]]]),
        axis=-1)
    expected_accuracy = np.average([4.0 / 8.0,
                                    5.0 / 8.0,
                                    5.0 / 7.0])
    accuracy, _ = metrics.prefix_accuracy(predictions, labels)
    with self.test_session() as session:
      accuracy_value = session.run(accuracy)
      self.assertAlmostEqual(expected_accuracy, accuracy_value)

  def testSequenceAccuracyMetric(self):
    predictions = np.random.randint(4, size=(12, 12, 12, 1))
    targets = np.random.randint(4, size=(12, 12, 12, 1))
    expected = np.mean(
        np.prod((predictions == targets).astype(float), axis=(1, 2)))
    with self.test_session() as session:
      scores, _ = metrics.padded_sequence_accuracy(
          tf.one_hot(predictions, depth=4, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32))
      a = tf.reduce_mean(scores)
      session.run(tf.global_variables_initializer())
      actual = session.run(a)
    self.assertEqual(actual, expected)

  def testTwoClassAccuracyMetric(self):
    predictions = tf.constant([0.0, 0.2, 0.4, 0.6, 0.8, 1.0], dtype=tf.float32)
    targets = tf.constant([0, 0, 1, 0, 1, 1], dtype=tf.int32)
    expected = 2.0 / 3.0
    with self.test_session() as session:
      accuracy, _ = metrics.two_class_accuracy(predictions, targets)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      actual = session.run(accuracy)
    self.assertAlmostEqual(actual, expected)

  def testTwoClassLogLikelihood(self):
    predictions = np.array([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
    targets = np.array([0, 0, 1, 0, 1, 1])
    expected = (2.0 * np.log(0.8) + 2.0 * np.log(0.4)) / 6.0
    with self.test_session() as session:
      avg_log_likelihood, _ = metrics.two_class_log_likelihood(
          predictions, targets)
      actual = session.run(avg_log_likelihood)
    self.assertAlmostEqual(actual, expected)

  def testTwoClassLogLikelihoodVersusOldImplementation(self):
    def alt_two_class_log_likelihood_impl(predictions, labels):
      float_labels = tf.cast(labels, dtype=tf.float64)
      float_predictions = tf.cast(tf.squeeze(predictions), dtype=tf.float64)
      # likelihood should be just p for class 1, and 1 - p for class 0.
      # signs is 1 for class 1, and -1 for class 0
      signs = 2 * float_labels - tf.ones_like(float_labels)
      # constant_term is 1 for class 0, and 0 for class 1.
      constant_term = tf.ones_like(float_labels) - float_labels
      likelihoods = constant_term + signs * float_predictions
      log_likelihoods = tf.log(likelihoods)
      avg_log_likelihood = tf.reduce_mean(log_likelihoods)
      return avg_log_likelihood
    predictions = np.random.rand(1, 10, 1)
    targets = np.random.randint(2, size=10)
    with self.test_session() as session:
      new_log_likelihood, _ = metrics.two_class_log_likelihood(
          predictions, targets)
      alt_log_likelihood = alt_two_class_log_likelihood_impl(
          predictions, targets)
      new_impl, alt_impl = session.run([new_log_likelihood, alt_log_likelihood])
    self.assertAlmostEqual(new_impl, alt_impl)

  def testRMSEMetric(self):
    predictions = np.full((10, 1), 1)  # All 1's
    targets = np.full((10, 1), 3)  # All 3's
    expected = np.sqrt(np.mean((predictions - targets)**2))  # RMSE = 2.0
    with self.test_session() as session:
      rmse, _ = metrics.padded_rmse(
          tf.constant(predictions, dtype=tf.int32),
          tf.constant(targets, dtype=tf.int32))
      session.run(tf.global_variables_initializer())
      actual = session.run(rmse)
    self.assertEqual(actual, expected)

  def testUnpaddedRMSEMetric(self):
    predictions = np.full((10, 1), 1)  # All 1's
    targets = np.full((10, 1), 3)  # All 3's
    expected = np.mean((predictions - targets)**2)  # MSE = 4.0
    with self.test_session() as session:
      mse, _ = metrics.unpadded_mse(
          tf.constant(predictions, dtype=tf.int32),
          tf.constant(targets, dtype=tf.int32))
      session.run(tf.global_variables_initializer())
      actual = session.run(mse)
    self.assertEqual(actual, expected)

  def testSequenceEditDistanceMetric(self):
    predictions = np.array([[3, 4, 5, 1, 0, 0],
                            [2, 1, 3, 4, 0, 0],
                            [2, 1, 3, 4, 0, 0]])
    # Targets are just a bit different:
    #  - first sequence has a different prediction
    #  - second sequence has a different prediction and one extra step
    #  - third sequence is identical
    targets = np.array([[5, 4, 5, 1, 0, 0],
                        [2, 5, 3, 4, 1, 0],
                        [2, 1, 3, 4, 0, 0]])
    # Reshape to match expected input format by metric fns.
    predictions = np.reshape(predictions, [3, 6, 1, 1])
    targets = np.reshape(targets, [3, 6, 1, 1])
    with self.test_session() as session:
      scores, weight = metrics.sequence_edit_distance(
          tf.one_hot(predictions, depth=6, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32))
      session.run(tf.global_variables_initializer())
      actual_scores, actual_weight = session.run([scores, weight])
    self.assertAlmostEqual(actual_scores, 3.0 / 13)
    self.assertEqual(actual_weight, 13)

  def testWordErrorRateMetric(self):

    # Ensure availability of the WER metric function in the dictionary.
    assert metrics.Metrics.WORD_ERROR_RATE in metrics.METRICS_FNS

    # Test if WER is computed correctly.
    ref = np.asarray([
        # a b c
        [97, 34, 98, 34, 99],
        [97, 34, 98, 34, 99],
        [97, 34, 98, 34, 99],
        [97, 34, 98, 34, 99],
    ])

    hyp = np.asarray([
        [97, 34, 98, 34, 99],  # a b c
        [97, 34, 98, 0, 0],  # a b
        [97, 34, 98, 34, 100],  # a b d
        [0, 0, 0, 0, 0]  # empty
    ])

    labels = np.reshape(ref, ref.shape + (1, 1))
    predictions = np.zeros((len(ref), np.max([len(s) for s in hyp]), 1, 1, 256))

    for i, sample in enumerate(hyp):
      for j, idx in enumerate(sample):
        predictions[i, j, 0, 0, idx] = 1

    with self.test_session() as session:
      actual_wer, unused_actual_ref_len = session.run(
          metrics.word_error_rate(predictions, labels))

    expected_wer = 0.417
    places = 3
    self.assertAlmostEqual(round(actual_wer, places), expected_wer, places)

  def testNegativeLogPerplexity(self):
    predictions = np.random.randint(4, size=(12, 12, 12, 1))
    targets = np.random.randint(4, size=(12, 12, 12, 1))
    with self.test_session() as session:
      scores, _ = metrics.padded_neg_log_perplexity(
          tf.one_hot(predictions, depth=4, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32))
      a = tf.reduce_mean(scores)
      session.run(tf.global_variables_initializer())
      actual = session.run(a)
    self.assertEqual(actual.shape, ())

  def testNegativeLogPerplexityMasked(self):
    predictions = np.random.randint(4, size=(12, 12, 12, 1))
    targets = np.random.randint(4, size=(12, 12, 12, 1))
    features = {
        'targets_mask': tf.to_float(tf.ones([12, 12]))
    }
    with self.test_session() as session:
      scores, _ = metrics.padded_neg_log_perplexity_with_masking(
          tf.one_hot(predictions, depth=4, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32),
          features)
      a = tf.reduce_mean(scores)
      session.run(tf.global_variables_initializer())
      actual = session.run(a)
    self.assertEqual(actual.shape, ())

  def testNegativeLogPerplexityMaskedAssert(self):
    predictions = np.random.randint(4, size=(12, 12, 12, 1))
    targets = np.random.randint(4, size=(12, 12, 12, 1))
    features = {}

    with self.assertRaisesRegexp(
        ValueError,
        'masked_neg_log_perplexity requires targets_mask feature'):
      with self.test_session() as session:
        scores, _ = metrics.padded_neg_log_perplexity_with_masking(
            tf.one_hot(predictions, depth=4, dtype=tf.float32),
            tf.constant(targets, dtype=tf.int32),
            features)
        a = tf.reduce_mean(scores)
        session.run(tf.global_variables_initializer())
        _ = session.run(a)

  def testSigmoidAccuracyOneHot(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [-1., 1.],
        [1., -1.]
    ])
    labels = np.array([
        [0, 1],
        [1, 0],
        [1, 0],
        [0, 1]
    ])
    logits = np.expand_dims(np.expand_dims(logits, 1), 1)
    labels = np.expand_dims(np.expand_dims(labels, 1), 1)

    with self.test_session() as session:
      score, _ = metrics.sigmoid_accuracy_one_hot(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertEqual(s, 0.5)

  def testSigmoidAccuracy(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [-1., 1.],
        [1., -1.]
    ])
    labels = np.array([1, 0, 0, 1])

    with self.test_session() as session:
      score, _ = metrics.sigmoid_accuracy(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertEqual(s, 0.5)

  def testSigmoidPrecisionOneHot(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [1., -1.],
        [1., -1.]
    ])
    labels = np.array([
        [0, 1],
        [0, 1],
        [0, 1],
        [0, 1]
    ])
    logits = np.expand_dims(np.expand_dims(logits, 1), 1)
    labels = np.expand_dims(np.expand_dims(labels, 1), 1)

    with self.test_session() as session:
      score, _ = metrics.sigmoid_precision_one_hot(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertEqual(s, 0.25)

  def testSigmoidRecallOneHot(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [1., -1.],
        [1., -1.]
    ])
    labels = np.array([
        [0, 1],
        [0, 1],
        [0, 1],
        [0, 1]
    ])
    logits = np.expand_dims(np.expand_dims(logits, 1), 1)
    labels = np.expand_dims(np.expand_dims(labels, 1), 1)

    with self.test_session() as session:
      score, _ = metrics.sigmoid_recall_one_hot(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertEqual(s, 0.25)

  def testSigmoidCrossEntropyOneHot(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [1., -1.],
        [1., -1.]
    ])
    labels = np.array([
        [0, 1],
        [1, 0],
        [0, 0],
        [0, 1]
    ])
    logits = np.expand_dims(np.expand_dims(logits, 1), 1)
    labels = np.expand_dims(np.expand_dims(labels, 1), 1)

    with self.test_session() as session:
      score, _ = metrics.sigmoid_cross_entropy_one_hot(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertAlmostEqual(s, 0.688, places=3)

  def testRocAuc(self):
    logits = np.array([
        [-1., 1.],
        [1., -1.],
        [1., -1.],
        [1., -1.]
    ])
    labels = np.array([
        [1],
        [0],
        [1],
        [0]
    ])
    logits = np.expand_dims(np.expand_dims(logits, 1), 1)
    labels = np.expand_dims(np.expand_dims(labels, 1), 1)

    with self.test_session() as session:
      score, _ = metrics.roc_auc(logits, labels)
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      s = session.run(score)
    self.assertAlmostEqual(s, 0.750, places=3)

  def testMultilabelMatch3(self):
    predictions = np.random.randint(1, 5, size=(100, 1, 1, 1))
    targets = np.random.randint(1, 5, size=(100, 10, 1, 1))
    weights = np.random.randint(0, 2, size=(100, 1, 1, 1))
    targets *= weights

    predictions_repeat = np.repeat(predictions, 10, axis=1)
    expected = (predictions_repeat == targets).astype(float)
    expected = np.sum(expected, axis=(1, 2, 3))
    expected = np.minimum(expected / 3.0, 1.)
    expected = np.sum(expected * weights[:, 0, 0, 0]) / weights.shape[0]
    with self.test_session() as session:
      scores, weights_ = metrics.multilabel_accuracy_match3(
          tf.one_hot(predictions, depth=5, dtype=tf.float32),
          tf.constant(targets, dtype=tf.int32))
      a, a_op = tf.metrics.mean(scores, weights_)
      session.run(tf.local_variables_initializer())
      session.run(tf.global_variables_initializer())
      _ = session.run(a_op)
      actual = session.run(a)
    self.assertAlmostEqual(actual, expected, places=6)

  def testPearsonCorrelationCoefficient(self):
    predictions = np.random.rand(12, 1)
    targets = np.random.rand(12, 1)

    expected = np.corrcoef(np.squeeze(predictions), np.squeeze(targets))[0][1]
    with self.test_session() as session:
      pearson, _ = metrics.pearson_correlation_coefficient(
          tf.constant(predictions, dtype=tf.float32),
          tf.constant(targets, dtype=tf.float32))
      session.run(tf.global_variables_initializer())
      session.run(tf.local_variables_initializer())
      actual = session.run(pearson)
    self.assertAlmostEqual(actual, expected)

if __name__ == '__main__':
  tf.test.main()


__lic__ = '''
/**
 * AS - the open source Automotive Software on https://github.com/parai
 *
 * Copyright (C) 2021  AS <parai@foxmail.com>
 *
 * This source code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 as published by the
 * Free Software Foundation; See <http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt>.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * for more details.
 */
 '''
import os
CWD = os.path.dirname(__file__)
if(CWD==''): CWD=os.path.abspath('.')
ASROOT = os.path.abspath('%s/../../../../'%(CWD))
PYAS = '%s/build/%s/any/pyas'%(ASROOT, os.name)
if(not os.path.exists('%s/AS.%s'%(CWD, 'pyd' if os.name == 'nt' else 'so'))):
    if(os.name == 'nt'):
        cmd = 'cd %s & set BOARD=any& set ANY=pyas& scons'%(ASROOT)
        cmd += '& cp -v %s/pyas.dll %s/AS.pyd'%(PYAS, CWD)
    else:
        cmd = 'cd %s && export BOARD=any && export ANY=pyas && scons --verbose'%(ASROOT)
        cmd += ' && cp -v %s/libpyas.so %s/AS.so'%(PYAS, CWD)
    os.system(cmd)

try:
    from .AS import *
except:
    from AS import *

__all__ = ['lin_open','lin_write','lin_read','lin_close']

dev__ = asdev()

__fdMap = {}
__lastPid = {}

def __bit(v, pos):
    return (v>>pos) & 0x01

def __get_pid(pid):
    pid = pid & 0x3F
    p0 = __bit(pid, 0) ^ __bit(pid, 1) ^ __bit(pid, 2) ^ __bit(pid, 4)
    p1 = ~(__bit(pid, 1) ^ __bit(pid, 3) ^ __bit(pid, 4) ^ __bit(pid, 5))
    pid = pid | (p0<<6) | (p1<<7);
    return pid&0xFF

def __get_checksum(pid, data, enhanced):
    checksum = 0
    if(enhanced):
        checksum += __get_pid(pid)
    for d in data:
        checksum += d
    return ((~checksum)&0xFF)

def lin_open(busid, device='socket', port='0', baudrate=115200, enhanced=True):
    if(busid not in __fdMap):
        device = 'lin/%s/%s'%(device, port)
        option = '%s'%(baudrate)
        fd = dev__.open(device.encode('utf-8'), option.encode('utf-8'))
        if(fd >= 0):
            __fdMap[busid] = (fd, enhanced)
        else:
            return False
    return True

def lin_write(busid, pid=None, data=None):
    fd, enhanced = __fdMap[busid]
    if(pid is None):
        ds = [ord('D')]
        ds.extend(data)
        checksum = __get_checksum(__lastPid[busid], data, enhanced)
        ds.append(checksum)
    elif(data is None):
        pid = __get_pid(pid)
        ds = [ord('H'), pid]
        __lastPid[busid] = pid
    else:
        pid = __get_pid(pid)
        ds = [ord('F'), pid]
        ds.extend(data)
        checksum = __get_checksum(pid, data, enhanced)
        ds.append(checksum)
        __lastPid[busid] = pid
    r = dev__.write(fd, bytes(bytearray(ds)), len(ds))
    if(r == len(ds)):
        return True
    else:
        return False

def lin_read(busid):
    fd, enhanced = __fdMap[busid]
    size,data= dev__.read(fd)
    if(size > 0):
        if((data[0] == ord('H')) and (size == 2)):
            pid = __get_pid(data[1])
            __lastPid[busid] = pid
            if(pid == data[1]):
                return True, pid&0x3F, None
            else:
                print('warning: invalid Pid', data)
        elif((data[0] == ord('F')) and (size > 3)):
            pid = __get_pid(data[1])
            __lastPid[busid] = pid
            if(pid == data[1]):
                d = []
                for i in range(size-3):
                    d.append(data[2+i])
                checksum = __get_checksum(pid, d, enhanced)
                if(checksum == data[size-1]):
                    return True, pid&0x3F, d
                else:
                    print('warning: invalid checksum', data)
            else:
                print('warning: invalid Pid', data)
        elif((data[0] == ord('D')) and (size > 2)):
            pid = __lastPid[busid]
            d = []
            for i in range(size-2):
                d.append(data[1+i])
            checksum = __get_checksum(pid, d, enhanced)
            if(checksum == data[size-1]):
                return True, pid&0x3F, d
            else:
                print('warning: invalid checksum', data)
        else:
            print('warning: invalid data', data)
    return False, None, None

def lin_close(busid):
    fd, _ = __fdMap[busid]
    dev__.close(fd)

if __name__ == '__main__':
    import time
    lin_open(0)
    lin_write(0, 0x3C, [0x11, 0x02, 0x10, 0x02, 0x55, 0x55, 0x55, 0x55])
    time.sleep(0.020)
    lin_write(0, 0x3d)
    time.sleep(1)
    ercd, pid, data = lin_read(0)
    print(ercd, pid, data)
    lin_close(0)


#!/usr/bin/python3
'''
some docmentation
'''

def isWinner(x, nums):
    '''
some docmentation
'''
    wins = {'Maria': 0, 'Ben': 0}
    for n in nums:
        primes = get_primes(n)
        player = 'Maria'
        while True:
            move = get_optimal_move(primes, player)
            if move is None:
                break
            primes = remove_multiples(primes, move)
            player = switch_player(player)
        winner = switch_player(player)
        wins[winner] += 1
    if wins['Maria'] > wins['Ben']:
        return 'Maria'
    elif wins['Maria'] < wins['Ben']:
        return 'Ben'
    else:
        return None

def get_primes(n):
    '''
some docmentation
'''
    primes = [True] * (n + 1)
    primes[0] = primes[1] = False
    for i in range(2, int(n ** 0.5) + 1):
        if primes[i]:
            for j in range(i * i, n + 1, i):
                primes[j] = False
    return [i for i in range(2, n + 1) if primes[i]]

def get_optimal_move(primes, player):
    '''
some docmentation
'''
    if primes == []:
        return None
    if player == 'Maria':
        # pick the smallest prime
        return primes[0]
    else:
        # pick a prime that leaves the fewest primes for Maria
        min_primes_left = len(primes) + 1
        optimal_move = None
        for p in primes:
            new_primes = remove_multiples(primes, p)
            if len(new_primes) < min_primes_left:
                min_primes_left = len(new_primes)
                optimal_move = p
        return optimal_move

def remove_multiples(primes, p):
    '''
some docmentation
'''
    return [n for n in primes if n % p != 0]

def switch_player(player):
    '''
some docmentation
'''
    return 'Ben' if player == 'Maria' else 'Maria'


from __future__ import division
import torchvision.transforms as transforms_old
import torch
import math
import sys
import random
from PIL import Image

try:
    import accimage
except ImportError:
    accimage = None
import numpy as np
import numbers
import types
import collections
import warnings
import cv2
from torchvision.transforms import functional as F

if sys.version_info < (3, 3):
    Sequence = collections.Sequence
    Iterable = collections.Iterable
else:
    Sequence = collections.abc.Sequence
    Iterable = collections.abc.Iterable

__all__ = ["Compose", "ToTensor",
           "Lambda", "RandomApply", "RandomResizedCrop",
           "ColorJitter", "RandomGrayscale"]

_pil_interpolation_to_str = {
    Image.NEAREST: 'PIL.Image.NEAREST',
    Image.BILINEAR: 'PIL.Image.BILINEAR',
    Image.BICUBIC: 'PIL.Image.BICUBIC',
    Image.LANCZOS: 'PIL.Image.LANCZOS',
    Image.HAMMING: 'PIL.Image.HAMMING',
    Image.BOX: 'PIL.Image.BOX',
}

class ToPIL(object):
    def __call__(self, sample_rgb, sample_aux):
        img1 = sample_rgb
        img2 = sample_aux
        img1 = transforms_old.functional.to_pil_image(img1)
        img2 = transforms_old.functional.to_pil_image(img2)
        return img1, img2

class RandomResizedCrop(object):
    """Crop the given PIL Image to random size and aspect ratio.

    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
    is finally resized to given size.
    This is popularly used to train the Inception networks.

    Args:
        size: expected output size of each edge
        scale: range of size of the origin size cropped
        ratio: range of aspect ratio of the origin aspect ratio cropped
        interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), interpolation=Image.BILINEAR):
        if isinstance(size, tuple):
            self.size = size
        else:
            self.size = (size, size)
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
            warnings.warn("range should be of kind (min, max)")

        self.interpolation = interpolation
        self.scale = scale
        self.ratio = ratio

    @staticmethod
    def get_params(img, scale, ratio):
        """Get parameters for ``crop`` for a random sized crop.

        Args:
            img (PIL Image): Image to be cropped.
            scale (tuple): range of size of the origin size cropped
            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
                sized crop.
        """
        area = img.size[0] * img.size[1]

        for attempt in range(10):
            target_area = random.uniform(*scale) * area
            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
            aspect_ratio = math.exp(random.uniform(*log_ratio))

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if w <= img.size[0] and h <= img.size[1]:
                i = random.randint(0, img.size[1] - h)
                j = random.randint(0, img.size[0] - w)
                return i, j, h, w

        # Fallback to central crop
        in_ratio = img.size[0] / img.size[1]
        if (in_ratio < min(ratio)):
            w = img.size[0]
            h = int(round(w / min(ratio)))
        elif (in_ratio > max(ratio)):
            h = img.size[1]
            w = int(round(h * max(ratio)))
        else:  # whole image
            w = img.size[0]
            h = img.size[1]
        i = (img.size[1] - h) // 2
        j = (img.size[0] - w) // 2
        return i, j, h, w

    def __call__(self, img_rgb, img_aux):
        """
        Args:
            img (PIL Image): Image to be cropped and resized.

        Returns:
            PIL Image: Randomly cropped and resized image.
        """
        i, j, h, w = self.get_params(img_rgb, self.scale, self.ratio)
        return F.resized_crop(img_rgb, i, j, h, w, self.size, self.interpolation), F.resized_crop(img_aux, i, j, h, w,
                                                                                self.size,
                                                                                self.interpolation)

    def __repr__(self):
        interpolate_str = _pil_interpolation_to_str[self.interpolation]
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
        format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string

class RandomHorizontalFlip(object):
    """Horizontally flip the given PIL Image randomly with a given probability.

    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img_rgb, img_aux):
        """
        Args:
            img (PIL Image): Image to be flipped.

        Returns:
            PIL Image: Randomly flipped image.
        """
        if random.random() < self.p:
            return F.hflip(img_rgb), F.hflip(img_aux)
        return img_rgb, img_aux

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)

class RandomTransforms(object):
    """Base class for a list of transformations with randomness

    Args:
        transforms (list or tuple): list of transformations
    """

    def __init__(self, transforms):
        assert isinstance(transforms, (list, tuple))
        self.transforms = transforms

    def __call__(self, *args, **kwargs):
        raise NotImplementedError()

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string

class Compose(object):
    """Composes several transforms together.

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
    """

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

    def __call__(self, img_rgb, img_aux):
        for t in self.transforms:
            img_rgb, img_aux = t(img_rgb, img_aux)
        return img_rgb, img_aux

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string

class RandomApply(RandomTransforms):
    """Apply randomly a list of transformations with a given probability

    Args:
        transforms (list or tuple): list of transformations
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
        super(RandomApply, self).__init__(transforms)
        self.p = p

    def __call__(self, img_rgb, img_aux):
        if self.p < random.random():
            return img_rgb, img_aux
        for t in self.transforms:
            img_rgb, img_aux = t(img_rgb, img_aux)
        return img_rgb, img_aux

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        format_string += '\n    p={}'.format(self.p)
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string

class ColorJitter(object):
    """Randomly change the brightness, contrast and saturation of an image.

    Args:
        brightness (float or tuple of float (min, max)): How much to jitter brightness.
            brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
            or the given [min, max]. Should be non negative numbers.
        contrast (float or tuple of float (min, max)): How much to jitter contrast.
            contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]
            or the given [min, max]. Should be non negative numbers.
        saturation (float or tuple of float (min, max)): How much to jitter saturation.
            saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]
            or the given [min, max]. Should be non negative numbers.
        hue (float or tuple of float (min, max)): How much to jitter hue.
            hue_factor is chosen uniformly from [-hue, hue] or the given [min, max].
            Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5.
    """

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)

    def _check_input(self, value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
            value = [center - value, center + value]
            if clip_first_on_zero:
                value[0] = max(value[0], 0)
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError("{} values should be between {}".format(name, bound))
        else:
            raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name))

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value

    @staticmethod
    def get_params(brightness, contrast, saturation, hue):
        """Get a randomized transform to be applied on image.

        Arguments are same as that of __init__.

        Returns:
            Transform which randomly adjusts brightness, contrast and
            saturation in a random order.
        """
        transforms = []

        if brightness is not None:
            brightness_factor = random.uniform(brightness[0], brightness[1])
            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))

        if contrast is not None:
            contrast_factor = random.uniform(contrast[0], contrast[1])
            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))

        if saturation is not None:
            saturation_factor = random.uniform(saturation[0], saturation[1])
            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))

        if hue is not None:
            hue_factor = random.uniform(hue[0], hue[1])
            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))

        random.shuffle(transforms)
        transform = transforms_old.Compose(transforms)

        return transform

    def __call__(self, img_rgb, img_aux):
        """
        Args:
            img (PIL Image): Input image.

        Returns:
            PIL Image: Color jittered image.
        """
        transform = self.get_params(self.brightness, self.contrast,
                                    self.saturation, self.hue)
        return transform(img_rgb),transform(img_aux)

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        format_string += 'brightness={0}'.format(self.brightness)
        format_string += ', contrast={0}'.format(self.contrast)
        format_string += ', saturation={0}'.format(self.saturation)
        format_string += ', hue={0})'.format(self.hue)
        return format_string

class Lambda(object):
    """Apply a user-defined lambda as a transform.

    Args:
        lambd (function): Lambda/function to be used for transform.
    """

    def __init__(self, lambd):
        assert callable(lambd), repr(type(lambd).__name__) + " object is not callable"
        self.lambd = lambd

    def __call__(self, img):
        return self.lambd(img)

    def __repr__(self):
        return self.__class__.__name__ + '()'

class RandomGrayscale(object):
    """Randomly convert image to grayscale with a probability of p (default 0.1).

    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
        PIL Image: Grayscale version of the input image with probability p and unchanged
        with probability (1-p).
        - If input image is 1 channel: grayscale version is 1 channel
        - If input image is 3 channel: grayscale version is 3 channel with r == g == b

    """

    def __init__(self, p=0.1):
        self.p = p

    def __call__(self, img_rgb, img_aux):
        """
        Args:
            img (PIL Image): Image to be converted to grayscale.

        Returns:
            PIL Image: Randomly grayscaled image.
        """
        num_output_channels = 1 if img_rgb.mode == 'L' else 3
        if random.random() < self.p:
            return F.to_grayscale(img_rgb, num_output_channels=num_output_channels), F.to_grayscale(img_aux,
                                                                                num_output_channels=num_output_channels)
        return img_rgb, img_aux

    def __repr__(self):
        return self.__class__.__name__ + '(p={0})'.format(self.p)

class GaussianBlur(object):
    # Implements Gaussian blur as described in the SimCLR paper
    def __init__(self, kernel_size, min=0.1, max=2.0):
        self.min = min
        self.max = max
        # kernel size is set to be 10% of the image height/width
        self.kernel_size = kernel_size

    def __call__(self, sample_rgb, sample_aug):
        sample_rgb = np.array(sample_rgb)
        sample_aug = np.array(sample_aug)

        # blur the image with a 50% chance
        prob = np.random.random_sample()

        if prob < 0.5:
            #            print(self.kernel_size)
            sigma = (self.max - self.min) * np.random.random_sample() + self.min
            sample_rgb = cv2.GaussianBlur(sample_rgb, (self.kernel_size, self.kernel_size), sigma)
            sample_aug = cv2.GaussianBlur(sample_aug, (self.kernel_size, self.kernel_size), sigma)
        return sample_rgb, sample_aug

class ToTensor(object):
    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.

    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
    or if the numpy.ndarray has dtype = np.uint8

    In the other cases, tensors are returned without scaling.
    """

    def __call__(self, pic, pic_aug):
        """
        Args:
            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.

        Returns:
            Tensor: Converted image.
        """
        return F.to_tensor(pic), F.to_tensor(pic_aug)

    def __repr__(self):
        return self.__class__.__name__ + '()'


"""
Problem Description:
Cody was once understanding numbers, their squares and perfect squares from his teacher.
A perfect square is a number that can be expressed as square of an integer.
To check how much Cody understood the concept his teacher kept a test. He has to find the nearest perfect square of the given number N.

Input:
The first line of input consists of a single integer T denoting the number of test cases.
The first line of each test case consists of single integer N.

Output:
For each test case print the nearest perfect square.

Constraints:
1<=T<=100.
1<=N<=10^4.

Sample Input:
2
1602
2

Sample Output:
40
1
"""

t = int(input())

for i in range(t):
    n = int(input())
    l = []
    boo = True
    i = 1
    while boo:
        i += 1
        if i ** 2 <= n:
            l.append(i)

        else:
            l.append(i)
            boo = False

    if n == 2 or n == 1: print(1)
    elif l[-2] == n: print(n)
    else:
        if len(l) >= 3:
            print(l[-1]) if l[-1] ** 2 - n <= n - l[-2] ** 2 else print(l[-2])


from tkinter import *
from tkinter.ttk import *

MAIN_THEME = 'radiance'
character_limit_red = "#ff0000"
character_limit_yellow = '#b3b300'
character_limit_green = '#009900'
t_ent_font = ('Helvetica', 13)
treeview_even_row = '#99ccff'
treeview_odd_row = '#ffffff'

def style_widget(root):
    style = Style(root)

    # Label Widgets
    style.configure('TLabel', font=('Times', 20, 'underline'), anchor=CENTER)
    style.configure('Username.TLabel', font=('Times', 13, 'italic'))
    style.configure('MainPage.TLabel', font=('Times', 30, 'italic'), anchor=CENTER)
    style.configure('QuestionHead.TLabel', font=('Helvetica', 20, 'underline'), anchor=W)
    style.configure('Questions.TLabel', font=('Helvetica', 15, 'underline'), anchor=W)
    style.configure('Options.TLabel', font=('arial', 12))
    style.configure('CorrectAns.TLabel', font=('Helvetica', 15, 'bold'), foreground='#33ff33')
    style.configure('WrongAns.TLabel', font=('Helvetica', 15, 'bold'), foreground='red')
    style.configure('Typical.TLabel', font=('Helvetica', 15, 'bold'))
    style.configure('Char.TLabel', font=('Helvetica', 15, 'roman'))

    # Entry Widgets
    style.configure('Search.TEntry', )

    # RadioButtons
    style.configure('TRadiobutton', font=("arial", 14))

    # Button Widgets
    style.configure('TButton', font=('Helvetica', 13, 'italic', 'bold'))
    style.map('TButton',
              foreground=[('pressed', 'blue'), ('active', 'red'), ('disabled', 'grey')])

    # Frame Widgets
    style.configure('TopButtons.TFrame', borderwidth=2, relief='raised')

    # Treeview Widgets
    style.configure('Treeview',
                    background='#99ffcc',
                    foreground='#333300',
                    font=('Helvetica', 12),
                    rowheight=25,
                    fieldbackground='#99ffcc')
    style.map('Treeview', background=[('selected', '#6600ff')])
    style.configure('Treeview.Heading', font=('Helvetica', 15), background='#80ffff')
    style.map("Treeview.Heading", relief=[('active', 'groove'), ('pressed', 'sunken')])


import base64

# @param img_file_name 传入文件路径
# @return base64_String
# img_file_name: ./dog.png
def img_to_base64(img_file_name):
    f = open(img_file_name,'rb')
    ls_f = base64.b64encode(f.read())
    f.close()
    img_name = img_file_name[:img_file_name.index('.')]
    return dict({"img_name":img_name,"base64_string":ls_f});

# 将转化后的base64存在txt文件中
# path = '/user/bobtang/desktop/'
# info['img_name'] = 'dog'
# txt文件: /user/bobtang/desktop/dog.txt
def save_as_txt(path,info):
    file_name = path + info['img_name']+'txt'
    with open(file_name,"w") as file:   #”w"代表着每次运行都覆盖内容
        file.write(str(info['base64_string'],encoding='utf8'))
        file.close()

if __name__ == '__main__':
    # img_file_name
    img_file_name = '/Users/bobtang/desktop/deskfile/vue-command/note/v-show-display:none.png'
    # 当前文件夹
    path = ''
    info = img_to_base64(img_file_name)
    save_as_txt(path,info)


from django.db import models
# Create your models here.

class Author(models.Model):
    name = models.CharField('Nombre', max_length=200, blank=False)
    nationality = models.CharField('Nacionalidad', max_length=3)

    def __str__(self):
        """Return string representation of object."""
        return self.name

class Books(models.Model):
    title = models.CharField('Título', max_length=250)
    year = models.IntegerField('Año', default=1900, )
    language = models.CharField('Idioma', max_length=2, help_text='ISO 639-1 Language code (2 chars)')
    cover_url = models.URLField('Portada', help_text='Guarda el url de una imagen', blank=True, null=True)
    price = models.DecimalField('Precio', max_digits=10, decimal_places=2, default='')
    sellable = models.BooleanField('Disponible', default=True, help_text='True disponible, False no disponible')
    copies = models.IntegerField('Copias', default=1)
    description = models.TextField('Descripción', blank=True, max_length=500)
    # blank = True hace el campo opcional, y blank = False hace el campo obligatorio
    author = models.ForeignKey(Author, on_delete=models.CASCADE)

    def __str__(self):
        """Return string representation of object."""
        return self.title


import sys
from collections import deque
#deque 가져오기 from collections from deque
#sys.stdin = open("input.txt")
n , m = map(int, input().split())
que = [(pos, val) for pos, val in enumerate(list(map(int, input().split())))]
#1차원 리스트 * 1차원 튜플 = 2차원 <= 튜플형태로 받기
#1차원 배열로 접근한다.
cnt=1
#deque선언
que = deque(que)
while True:
    #큐의 pop은 대입이 가능하다.
    cur = que.popleft()
    # if cur[0] == m:
    if any(x[1] > cur[1] for x in que):
        que.append(cur)
    else:
        if cur[0] == m:
            print(cnt)
            break
        else:
            cnt += 1

    #C style
    # for x in que:
    #     if x[1]>cur[1]:
    #         que.append(cur)

    # for pos, val in que:
    #     if
    # print(1)

응급실_중요하다(맵,튜플.큐).py


import collections, util, math, random
import numpy as np

############################################################

############################################################
# Problem 2a

class ValueIteration(util.MDPAlgorithm):

    # Implement value iteration.  First, compute V_opt using the methods
    # discussed in class.  Once you have computed V_opt, compute the optimal
    # policy pi.  Note that ValueIteration is an instance of util.MDPAlgrotithm,
    # which means you will need to set pi and V (see util.py).
    def solve(self, mdp, epsilon=0.001):
        mdp.computeStates()

        v_opt_old = dict.fromkeys(mdp.states, 0)
        self.pi = dict.fromkeys(mdp.states, None)

        max_change = float('inf')
        while max_change > epsilon: #while v_opt is changing, keep looping
            self.V = dict.fromkeys(mdp.states, -float('inf'))

            for state in mdp.states: #loop through states
                for action in mdp.actions(state): #loop though actions

                    q_opt = 0 #for adding up the possibilities from successor states
                    for succesor in mdp.succAndProbReward(state, action): #loop through sucessor states
                        q_opt += succesor[1] * (succesor[2] + mdp.discount()*v_opt_old[succesor[0]])

                    if q_opt > self.V[state]: #take best action and set optimal policy
                        self.V[state] = q_opt
                        self.pi[state] = action

            #compute stopping criterion
            delta_dict = {key: abs(self.V[key] - v_opt_old.get(key, 0)) for key in self.V.keys()}
            max_change = delta_dict[max(delta_dict, key=delta_dict.get)]
            v_opt_old = self.V.copy()

############################################################
class ser_fd(util.MDP):
    def __init__(self):
        """
        cardValues: array of card values for each card type
        multiplicity: number of each card type
        """

    # Return the start state. First is a indicator of whether on upper or lower trajectory, second is step in the task
    def startState(self):
        return (0, 0)

    # Return set of actions possible from |state|.
    def actions(self, state):
        A, stage = state
        if stage == 0 or stage == 1:
            return ['Next']
        else:
            return ['Go', 'Dont']

    # Return a list of (newState, prob, reward) tuples corresponding to edges
    # coming out of |state|.
    def succAndProbReward(self, state, action):
        succ_prob_reward = []
        stim, stage = state

        #stage 0, A
        #stage 1, B which leads to dot
        #stage 2, B which doesnt lead to dot

        if stage is None:
            return []
        elif stage == 0:
            succ_prob_reward.append(((0,1),.5,0))
            succ_prob_reward.append(((1,1),.5,0))
        elif stage == 1 and stim == 0: #A trial
            succ_prob_reward.append(((0,2),.95,0))
            succ_prob_reward.append(((1,2),.05,0))
        elif stage == 1 and stim == 1:
            succ_prob_reward.append(((0,2),.05,0))
            succ_prob_reward.append(((1,2),.95,0))
        elif stage == 2 and stim ==  0 and action == 'Go': #B+ and respond
            succ_prob_reward.append(((None,None),.95,1))
            succ_prob_reward.append(((None,None),.05,-1))
        elif stage == 2 and stim ==  0 and action == 'Dont': #B+ and dont respond
            succ_prob_reward.append(((None,None),.95,-.1))
            succ_prob_reward.append(((None,None),.05,.1))
        elif stage == 2 and stim ==  1 and action == 'Go': #B- and respond
            succ_prob_reward.append(((None,None),.95,-1))
            succ_prob_reward.append(((None,None),.05,1))
        elif stage == 2 and stim ==  1 and action == 'Dont': #B- and dont respond
            succ_prob_reward.append(((None,None),.95,.1))
            succ_prob_reward.append(((None,None),.05,-.1))

        return succ_prob_reward

    def discount(self):
        return 1

############################################################
class ser_fd_features(util.MDP):
    def __init__(self):
        """
        cardValues: array of card values for each card type
        multiplicity: number of each card type
        """

    # Return the start state. First is a indicator of whether on upper or lower trajectory, second is step in the task
    def startState(self):
        return (0,0,0)

    # Return set of actions possible from |state|.
    def actions(self, state):
        A,B,stage = state
        if stage == 0:
            return ['Next']
        else:
            return ['Go', 'Dont']

    # Return a list of (newState, prob, reward) tuples corresponding to edges
    # coming out of |state|.
    def succAndProbReward(self, state, action):
        succ_prob_reward = []
        A,B,stage = state

        if stage is None:
            return []
        elif stage == 0:
            succ_prob_reward.append(((1,100,1),.5,0))
            succ_prob_reward.append(((0,100,1),.5,0))
        elif stage == 1 and A == 1 and action == 'Dont': #A trial
            succ_prob_reward.append(((2,1,2),1,0))
        elif stage == 1 and A == 1 and action == 'Go': #A trial
            succ_prob_reward.append(((2,1,2),1,-1))
        elif stage == 1 and A == 0 and action == 'Dont':
            succ_prob_reward.append(((3,1,2),1,0))
        elif stage == 1 and A == 0 and action == 'Go':
            succ_prob_reward.append(((3,1,2),1,-1))

        # elif stage == 2 and A == 2 and action == 'Go': #A -> B+ and respond
        #     succ_prob_reward.append(((None,None,None),.95,1))
        #     succ_prob_reward.append(((None,None,None),.05,-1))
        # elif stage == 2 and A == 2 and action == 'Dont': #A -> B+ and dont respond
        #     succ_prob_reward.append(((None,None,None),.95,-.1))
        #     succ_prob_reward.append(((None,None,None),.05,.1))
        # elif stage == 2 and A == 3 and action == 'Go': #B- and respond
        #     succ_prob_reward.append(((None,None,None),.95,-1))
        #     succ_prob_reward.append(((None,None,None),.05,1))
        # elif stage == 2 and A == 3 and action == 'Dont': #B- and dont respond
        #     succ_prob_reward.append(((None,None,None),.95,.1))
        #     succ_prob_reward.append(((None,None,None),.05,-.1))

        elif stage == 2 and A == 2 and action == 'Go': #A -> B+ and respond
            succ_prob_reward.append(((None,None,None),.95,-1))
            succ_prob_reward.append(((None,None,None),.05,1))
        elif stage == 2 and A == 2 and action == 'Dont': #A -> B+ and dont respond
            succ_prob_reward.append(((None,None,None),.95,.1))
            succ_prob_reward.append(((None,None,None),.05,-.1))
        elif stage == 2 and A == 3 and action == 'Go': #B- and respond
            succ_prob_reward.append(((None,None,None),.95,1))
            succ_prob_reward.append(((None,None,None),.05,-1))
        elif stage == 2 and A == 3 and action == 'Dont': #B- and dont respond
            succ_prob_reward.append(((None,None,None),.95,-.1))
            succ_prob_reward.append(((None,None,None),.05,.1))
        return succ_prob_reward

    def discount(self):
        return 1

############################################################
class sim_fd_features(util.MDP):
    def __init__(self):
        """
        cardValues: array of card values for each card type
        multiplicity: number of each card type
        """

    # Return the start state. First is a indicator of whether on upper or lower trajectory, second is step in the task
    def startState(self):
        return (0,0,0)

    # Return set of actions possible from |state|.
    def actions(self, state):
        A,B,stage = state
        if stage == 0:
            return ['Next']
        else:
            return ['Go', 'Dont']

    # Return a list of (newState, prob, reward) tuples corresponding to edges
    # coming out of |state|.
    def succAndProbReward(self, state, action):
        succ_prob_reward = []
        A,B,stage = state

        if stage is None:
            return []
        elif stage == 0:
            succ_prob_reward.append(((1,1,1),.5,0))
            succ_prob_reward.append(((0,1,1),.5,0))
        elif stage == 1 and A == 1 and action == 'Go': #AB+ and respond
            succ_prob_reward.append(((None,None,None),.95,1))
            succ_prob_reward.append(((None,None,None),.05,-1))
        elif stage == 1 and A == 1 and action == 'Dont': #A -> B+ and dont respond
            succ_prob_reward.append(((None,None,None),.95,-.1))
            succ_prob_reward.append(((None,None,None),.05,.1))
        elif stage == 1 and A == 0 and action == 'Go': #B- and respond
            succ_prob_reward.append(((None,None,None),.95,-1))
            succ_prob_reward.append(((None,None,None),.05,1))
        elif stage == 1 and A == 0 and action == 'Dont': #B- and dont respond
            succ_prob_reward.append(((None,None,None),.95,.1))
            succ_prob_reward.append(((None,None,None),.05,-.1))
        # elif stage == 1 and A == 1 and action == 'Go': #AB+ and respond
        #     succ_prob_reward.append(((None,None,None),.95,-1))
        #     succ_prob_reward.append(((None,None,None),.05,1))
        # elif stage == 1 and A == 1 and action == 'Dont': #A -> B+ and dont respond
        #     succ_prob_reward.append(((None,None,None),.95,.1))
        #     succ_prob_reward.append(((None,None,None),.05,-.1))
        # elif stage == 1 and A == 0 and action == 'Go': #B- and respond
        #     succ_prob_reward.append(((None,None,None),.95,1))
        #     succ_prob_reward.append(((None,None,None),.05,-1))
        # elif stage == 1 and A == 0 and action == 'Dont': #B- and dont respond
        #     succ_prob_reward.append(((None,None,None),.95,-.1))
        #     succ_prob_reward.append(((None,None,None),.05,.1))

        return succ_prob_reward

    def discount(self):
        return 1

class phase_2_blocking(util.MDP):
    def __init__(self):
        """
        cardValues: array of card values for each card type
        multiplicity: number of each card type
        """

    # Return the start state. First is a indicator of whether on upper or lower trajectory, second is step in the task
    def startState(self):
        return (0,0,0)

    # Return set of actions possible from |state|.
    def actions(self, state):
        A,B,stage = state
        if stage == 0:
            return ['Next']
        else:
            return ['Go', 'Dont']

    # Return a list of (newState, prob, reward) tuples corresponding to edges
    # coming out of |state|.
    def succAndProbReward(self, state, action):
        succ_prob_reward = []
        A,B,stage = state

        if stage is None:
            return []
        elif stage == 0:
            succ_prob_reward.append(((1,10,1),.5,0))
            succ_prob_reward.append(((11,12,1),.5,0))
        elif stage == 1 and A == 1 and action == 'Go': #AC+ and respond
            succ_prob_reward.append(((None,None,None),.95,1))
            succ_prob_reward.append(((None,None,None),.05,-1))
        elif stage == 1 and A == 1 and action == 'Dont': #AC+ and dont respond
            succ_prob_reward.append(((None,None,None),.95,-.1))
            succ_prob_reward.append(((None,None,None),.05,.1))
        elif stage == 1 and A == 11 and action == 'Go': #DE- and respond
            succ_prob_reward.append(((None,None,None),.95,-1))
            succ_prob_reward.append(((None,None,None),.05,1))
        elif stage == 1 and A == 11 and action == 'Dont': #DE- and dont respond
            succ_prob_reward.append(((None,None,None),.95,.1))
            succ_prob_reward.append(((None,None,None),.05,-.1))

        return succ_prob_reward

    def discount(self):
        return 1

############################################################
# Performs Q-learning.  Read util.RLAlgorithm for more information.
# actions: a function that takes a state and returns a list of actions.
# discount: a number between 0 and 1, which determines the discount factor
# featureExtractor: a function that takes a state and action and returns a list of (feature name, feature value) pairs.
# explorationProb: the epsilon value indicating how frequently the policy
# returns a random action
class QLearningAlgorithm(util.RLAlgorithm):
    def __init__(self, actions, discount, featureExtractor, init_weights = collections.Counter(), explorationProb=0.2):
        self.actions = actions
        self.discount = discount
        self.featureExtractor = featureExtractor
        self.explorationProb = explorationProb
        self.weights = init_weights
        self.numIters = 0

    # Return the Q function associated with the weights and features
    def getQ(self, state, action):
        score = 0
        for f, v in self.featureExtractor(state, action):
            score += self.weights[f] * v
        return score

    # This algorithm will produce an action given a state.
    # Here we use the epsilon-greedy algorithm: with probability
    # |explorationProb|, take a random action.
    def getAction(self, state):
        self.numIters += 1
        if random.random() < self.explorationProb:
            return random.choice(self.actions(state))
        else:
            return max((self.getQ(state, action), action) for action in self.actions(state))[1]

    # Call this function to get the step size to update the weights.
    def getStepSize(self):
        return 1.0 / math.sqrt(self.numIters)

    # We will call this function with (s, a, r, s'), which you should use to update |weights|.
    # Note that if s is a terminal state, then s' will be None.  Remember to check for this.
    # You should update the weights using self.getStepSize(); use
    # self.getQ() to compute the current estimate of the parameters.
    def incorporateFeedback(self, state, action, reward, newState):
        if newState == None:
            maxQ = 0
        else:
            maxQ = max((self.getQ(newState, act), act) for act in self.actions(newState))[0]
        r = reward + self.discount*maxQ - self.getQ(state, action)

        for f, v in self.featureExtractor(state, action):
            self.weights[f] = self.weights[f] + self.getStepSize() * r * v

# Return a singleton list containing indicator feature for the (state, action)
# pair.  Provides no generalization.
def identityFeatureExtractor(state, action):
    featureKey = (state, action)
    featureValue = 1
    return [(featureKey, featureValue)]

############################################################
def serFeatureExtractor(state,action):
    A,B,stage = state
    if stage == None:
        return []
    featureValue = 1 #indicator features
    features = []
    if A == 0:
        featureKey = ('noA',action)
    elif A == 1:
        featureKey = ('A',action)
    elif A == 2:
        featureKey = ('pastA',action)
    elif A==3:
        featureKey = ('no_pastA',action)
    elif A == 11:
        featureKey = ('D',action)

    features.append((featureKey,featureValue))

    if B == 0:
        featureKey = ('noB',action)
    elif B==1:
        featureKey = ('B',action)
    elif B==10:
        featureKey = ('C',action)
    elif B==12:
        featureKey = ('E',action)
    features.append((featureKey,featureValue))

    #conjunctive features
    if A == 2 and B == 1:
        featureKey = ('pastA','B',action)
    if A == 1 and B == 1:
        featureKey = ('A','B',action)
    if A == 1 and B == 10:
        featureKey = ('A','C',action)
    if A == 11 and B == 12:
        featureKey = ('D','E',action)
    features.append((featureKey,featureValue))

    return features

# # run algorithms
# mdp = sim_fd_features()
# QLearning = QLearningAlgorithm(mdp.actions,mdp.discount(),serFeatureExtractor)
# rew_q = util.simulate(mdp, QLearning,numTrials=1000, maxIterations=1000)
#
# valueIter = ValueIteration()
# valueIter.solve(mdp)
# print valueIter.V
# print QLearning.weights
#
#
# #run subsequent
# mdp_phase2 = phase_2_blocking()
# QLearning_phase2 =  QLearningAlgorithm(mdp.actions,mdp.discount(),serFeatureExtractor,init_weights = QLearning.weights)
# rew_q2= util.simulate(mdp_phase2, QLearning_phase2,numTrials=25, maxIterations=1000)
#
# #get state q values
# for state in valueIter.pi:
#     print state
#     # print valueIter.pi[state]
#
#     actions = QLearning.actions(state)
#     if len(actions) > 1:
#        for action in actions:
#            print action
#            print QLearning.getQ(state,action)
#     else:
#         print actions
#         print QLearning.getQ(state,actions[0])
#
# ##get feature values
# print QLearning_phase2.weights
# print QLearning_phase2.weights[('C','Go')]
# print QLearning_phase2.weights[('C','Dont')]

c_go_out = '/Users/ianballard/Dropbox/fd/c_go_ser_fn.txt'
c_dont_out = '/Users/ianballard/Dropbox/fd/c_dont_ser_fn.txt'
num_iters = 100
c_go_weights = []
c_dont_weights = []
for i in range(0,num_iters):
    mdp = ser_fd_features()
    QLearning = QLearningAlgorithm(mdp.actions,mdp.discount(),serFeatureExtractor)
    rew_q = util.simulate(mdp, QLearning,numTrials=1000, maxIterations=1000)

    #run subsequent
    mdp_phase2 = phase_2_blocking()
    QLearning_phase2 =  QLearningAlgorithm(mdp.actions,mdp.discount(),serFeatureExtractor,init_weights = QLearning.weights)
    rew_q2= util.simulate(mdp_phase2, QLearning_phase2,numTrials=25, maxIterations=1000)

    c_go_weights.append(QLearning_phase2.weights[('C','Go')])
    c_dont_weights.append(QLearning_phase2.weights[('C','Dont')])

np.savetxt(c_go_out,c_go_weights)
np.savetxt(c_dont_out,c_dont_weights)


import csv
from urllib.request import urlopen
from bs4 import BeautifulSoup
from datetime import datetime

#add the stocks you want into this list. You should use 'https://finance.yahoo.com/quote/XXX' where XXX is the stock identifier.
stocks = ['https://finance.yahoo.com/quote/VTI/','https://finance.yahoo.com/quote/VGK/','https://finance.yahoo.com/quote/VPL/','https://finance.yahoo.com/quote/VWO/','https://finance.yahoo.com/quote/BWX/','https://finance.yahoo.com/quote/IEF/']
#update your holdings correspondingly.
currentHoldings = [1,1,1,1,1,1]
"""Vanguard total Stock Market (VTI) = [1]
Vanguard FTSE Europe (VGK) = [2]
Vanguard FTSE Pacific (VPL) = [3]
Vangaurd FTSE Emerging Market (VWO) = [4]
SPDR Bloomberg Barclays Int (BWX) = [5]
iShares 7-10Y T-BOND (IEF) = [6]
"""

totalAssetListUSD = []
totalAssetListSGD = []

def getUrl(stockURL):
    url = urlopen(stockURL)
    return url

def beautifulSoup(url):
    soup = BeautifulSoup(url, 'html.parser')
    return soup

def getName(soup):
    name_box = soup.find('h1',attrs = {'class':'D(ib) Fz(16px) Lh(18px)'})
    name = name_box.text.strip()
    return name

def getPrice(soup):
    price_box = soup.find('span', attrs = {'class':'Trsdu(0.3s) Trsdu(0.3s) Fw(b) Fz(36px) Mb(-4px) D(b)'})
    price = price_box.text.strip()
    return round(float(price),2)

def getExchangeRate():
    url = urlopen('https://finance.yahoo.com/quote/SGD%3DX/')
    soup = BeautifulSoup(url, 'html.parser')
    rate_box = soup.find('span', attrs = {'class':'Trsdu(0.3s) Trsdu(0.3s) Fw(b) Fz(36px) Mb(-4px) D(b)'})
    rate = float(rate_box.text.strip())
    return rate

def exportCSV(name, price, now):
    with open('record.csv', 'a') as csv_file:
        write = csv.writer(csv_file)
        write.writerow([name, price, now])

def main():
    now =  datetime.now()
    timeDate = now.strftime("%m/%d/%Y, %H:%M:%S")
    sgdRate = getExchangeRate()
    totalAssetUSD = 0
    totalAssetSGD = 0
    for i in range(0,len(stocks)):
        url = getUrl(stocks[i])
        soup = beautifulSoup(url)
        name = getName(soup)
        price = getPrice(soup)
        sgdPrice = round((price*sgdRate),2)
        print("@",timeDate)
        print(name, "\nCurrent price in USD:", price, "\nCurren price in SGD:", sgdPrice)
        addUSD = price * currentHoldings[i]
        addSGD = sgdPrice * currentHoldings[i]
        totalAssetListUSD.append(addUSD)
        totalAssetListSGD.append(addSGD)
        exportCSV(name, price, timeDate)
    for i in range(0, len(totalAssetListUSD)):
        totalAssetUSD = totalAssetUSD + totalAssetListUSD[i]
        totalAssetSGD = totalAssetSGD + totalAssetListSGD[i]
    print("Total Asset in USD: ", totalAssetUSD)
    print("Total Asset in SGD: ", totalAssetSGD)
    exportCSV("Total Asset in USD", totalAssetUSD, timeDate)

if __name__ == "__main__":
    main()


def countApplesAndOranges(s, t, a, b, apples, oranges):
    apple,orange = 0,0
    x,y = 0,0
    for i in range(len(apples)):
        x = a + apples[i]
        if x >= s and x <= t:
            apple += 1

    for j in range(len(oranges)):
        y = b + oranges[j]
        if y >= s and y <= t:
            orange += 1

    print(apple)
    print(orange)


import os
import glob
import numpy as np
import pandas as pd

from sortasurvey import observing

class Sample:
    """
    Creates a filtered sample based on the provided survey program. The most common
    use of the sample is during the target selection process, which calls the
    get_highest_priority method at each iteration. All other methods are useful
    for post-target selection filtering, shown in the provided Jupyter notebook example (TODO).

    Attributes
    ----------
    df : pandas.DataFrame
        vetted sample from the survey.Survey
    programs : pandas.DataFrame
        survey programs from the survey.Survey

    Parameters
    ----------
    survey : Optional[survey.Survey]
        will save the science programs and vetted sample of a given survey to the
        Sample.programs and Sample.df attributes, respectively if not `None`.
    path_init : Optional[str]
        path to vetted sample (default = 'info/TOIs_perfect.csv')
    path_final : Optional[str]
        root path to results (default = `None`)
    program : str
        selected program of interest within a survey

    """

    def __init__(self, program, survey=None, path_init='info/TOIs_perfect.csv', path_final=None):
        self.path_init = path_init
        self.program = program
        if survey is not None:
            self.df = survey.candidates.copy()
            self.programs = survey.sciences.copy()
            self.get_vetted_science(program)

    def get_vetted_science(self, drop_dup=True):
        """
        Similar to the vetted sample method above, it loads the same sample which was (is)
        available for target selection but filters the sample to address a specific science
        case's (or program) needs.

        Attributes
        ----------
        program : str
            the program used to filter the selected Sample
        query : pandas.DataFrame
            the original vetted survey sample filtered on the specific program's selection criteria

        Parameters
        ----------
        drop_dup : Optional[bool]
            option to drop duplicate targets in the sample query, since the target will only be observed once. As a result, the default is `True`.

        """
        if not hasattr(self, 'df'):
            self.df = pd.read_csv(self.path_init)
        query = self.df.query(self.programs.loc[self.program,'filter'])
        if drop_dup:
            query = query.drop_duplicates(subset='tic')
        query = observing.get_actual_costs(self.program, self.programs, query)
        query = query.sort_values(by=self.programs.loc[self.program,'prioritize_by'], ascending=self.programs.loc[self.program,'ascending_by'])
        if self.programs.loc[self.program,'high_priority'] != []:
            top = pd.DataFrame(columns = query.columns.values.tolist())
            for toi in self.programs.loc[self.program,'high_priority']:
                new_filter = 'toi == %.2f'%toi
                new = self.df.query(new_filter)
                new = observing.get_actual_costs(self.program, self.programs, new)
                top = pd.concat([top,new])
            query = pd.concat([top,query])
        query.reset_index(drop=True, inplace=True)
        self.query = query.copy()

    def get_highest_priority(self):
        """
        Returns the highest priority target for a selected program that has not
        yet been selected by that program.

        Attributes
        ----------
        program : str
            the program used to filter the selected Sample
        query : pandas.DataFrame
            the original vetted survey sample filtered on the specific program's selection criteria

        Returns
        -------
        pick : pandas.DataFrame
            the selected's program's highest priority pick

        """
        for i in self.query.index.values.tolist():
            if not int(self.query.loc[i,'in_%s'%self.program]):
                pick = self.query.loc[i]
                break
        return pick

    def get_vetted_sample(self):
        """
        Loads the vetted sample that was available during the survey selection process.
        Using the final csv path, it selects the sample from the most recent output directory.

        Parameters
        ----------
        final_path : Optional[str]
            path of output csv file to read in

        Returns
        -------
        df : pandas.DataFrame
            the full vetted sample used for the target selection

        """
        list_of_files = glob.glob(self.final_path)
        latest_file = max(list_of_files, key=os.path.getctime)
        df = pd.read_csv(latest_file)
        return df

    def get_selected_sample(self, final_path=None):
        """
        Loads the final sample selected from the target prioritization process. If
        no path is provided, it will automatically load in the most recent output.

        Parameters
        ----------
        final_path : Optional[str]
            path of output csv file to read in

        Returns
        -------
        df : pandas.DataFrame
            the sample selected by the algorithm for a given Survey

        """
        if final_path is None:
            list_of_files = glob.glob('results/other/*/*/TOIs_perfect_final.csv')
            latest_file = max(list_of_files, key=os.path.getctime)
        else:
            latest_file = final_path
        df = pd.read_csv(latest_file)
        df.query("in_other_programs != 0", inplace=True)
        return df

    def get_science_sample(self, program, path_final=None):
        """
        Loads the final sample selected from the target prioritization process for a
        specific science case or program from the Survey. If no path is provided, it
        will automatically load in the most recent output.

        Parameters
        ----------
        program : str
            specific Survey program to downselect the sample for
        path_final : Optional[str]
            path of output csv file to read in

        Returns
        -------
        df : pandas.DataFrame
            the sample selected by the algorithm for a specific Survey program

        """
        list_of_files = glob.glob(self.final_path)
        latest_file = max(list_of_files, key=os.path.getctime)
        df = pd.read_csv(latest_file)
        df.query("in_%s == 1"%program, inplace=True)
        return df