#!/usr/bin/env python
# coding: utf-8

# In[2]:

import cv2
from matplotlib import pyplot as plt

# In[16]:

# algorithms

speed_limit_signs = cv2.CascadeClassifier('E:/Traffic Sign Detection/haarCascade.xml')
stop_signs = cv2.CascadeClassifier('E:/Traffic Sign Detection/Stopsigns.xml')
bump_signs = cv2.CascadeClassifier('E:/Traffic Sign Detection/bump_signs.xml')

# In[53]:

# open image

img  = cv2.imread('E:/Traffic Sign Detection/archive/images/data of traffic signs/p/road328.png')
img = cv2.cvtColor(img  , cv2.COLOR_BGR2RGB)

# convert to gray
gray = cv2.cvtColor(img  , cv2.COLOR_RGB2GRAY)

# signs detection
dic = {speed_limit_signs:'SPEED' , stop_signs:'STOP' , bump_signs:'BUMP'}

for key_of_dic in dic:

    signs = key_of_dic.detectMultiScale(gray  ,  minSize =(20, 20) , minNeighbors =2)

    # create a rectangle
    for (x,y,w,h) in signs:
        cv2.rectangle(img  , (x,y) , (x+w , y+h) , (255,0,0) ,3)
        # create a text

        cv2.putText(img  , str(dic[key_of_dic]), (x-10,y-10) ,cv2.FONT_HERSHEY_COMPLEX_SMALL , 1 ,(100,0,0) , 2 , cv2.LINE_AA )

plt.imshow(img)
plt.show()
# cv2.imshow('img ' , img)

# In[54]:

# open Video
cap = cv2.VideoCapture('E:/Traffic Sign Detection/Ford Focus - Traffic Sign Recognition_Trim.mp4')

while cap.isOpened():

    __ , frame = cap.read()

    # convert to gray
    gray = cv2.cvtColor(frame  , cv2.COLOR_BGR2GRAY)

    dic = {speed_limit_signs:'SPEED' , stop_signs:'STOP' , bump_signs:'BUMP'}

    # signs detection
    for key_of_dic in dic:
        signs = key_of_dic.detectMultiScale(gray  ,  minSize =(20, 20) ,  minNeighbors =10)

        for (x,y,w,h) in signs:
            # create a rectangle
            cv2.rectangle(frame  , (x,y) , (x+w , y+h) , (255,0,0) ,3)

            # create a text

            cv2.putText(frame  ,str(dic[key_of_dic]) , (x-10,y-10) ,cv2.FONT_HERSHEY_COMPLEX_SMALL, 1 ,(100,0,0) , 2 , cv2.LINE_AA )

    cv2.imshow('Traffic Signs Detection' , frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import robosuite
import imageio
import numpy as np
import os
from glob import glob
from copy import deepcopy
import pickle

import json
from imageio import mimwrite

from replay_buffer import ReplayBuffer, compress_frame

from torch.utils.tensorboard import SummaryWriter
import torch
import robosuite.utils.macros as macros
torch.set_num_threads(3)
import TD3_kinematic

from dh_utils import seed_everything, normalize_joints, skip_state_keys
from robosuite.wrappers import JacoSim2RealWrapper
from utils import build_replay_buffer, build_env, build_model, plot_replay, get_rot_mat
from IPython import embed

"""
eef_rot_offset?
https://github.com/ARISE-Initiative/robosuite/blob/fc3738ca6361db73376e4c9d8a09b0571167bb2d/robosuite/models/robots/manipulators/manipulator_model.py
https://github.com/ARISE-Initiative/robosuite/blob/65d3b9ad28d6e7a006e9eef7c5a0330816483be4/robosuite/environments/manipulation/single_arm_env.py#L41
"""
#def reacher_kinematic_fn(action, state, body, next_body):
#    bs,fs = action.shape
#    n_joints = len(robot_dh.npdh['DH_a'])
#    # turn relative action to abs action
#    # env.obs_keys = ['position', to_target', 'velocity']
#    joint_position = action+torch.FloatTensor(state[:,:2])
#    eef_rot = robot_dh.torch_angle2ee(robot_dh.base_matrix, joint_position)
#    eef_pos = eef_rot[:,:2,3]
#    st_target = n_joints+19+3
#    target_pos = next_body[:,st_target:st_target+16].reshape(bs, 4, 4)[:,:2,3]
#    target_pos = torch.FloatTensor(target_pos)
#    return eef_pos, target_pos
#
#def jaco_kinematic_fn(action, state, body, next_body):
#    # last dim is gripper
#    bs = action.shape[0]
#    n_joints = len(robot_dh.npdh['DH_a'])
#    # turn relative action to abs action
#    joint_position = action[:, :n_joints] + torch.FloatTensor(body[:, :n_joints])
#    eef_rot = robot_dh.torch_angle2ee(robot_dh.base_matrix, joint_position)
#    eef_pos = eef_rot[:,:3,3]
#
#    # second body n_joints + 3 + 16 + 3 = 29
#    handle_rot = next_body[:,29:].reshape(bs, 4, 4)
#    handle_pos = torch.FloatTensor(handle_rot[:,:3,3])
#    return eef_pos, handle_pos

def eval_policy(eval_env, policy, kwargs, eval_episodes=10):

    total_rewards = []
    for _ in range(eval_episodes):
        done = False
        state, body = eval_env.reset()
        ep_reward = 0
        while not done:
            action = policy.select_action(np.array(state)).clip(-kwargs['max_action'], kwargs['max_action'])
            state, body, reward, done, _ = eval_env.step(action)
            ep_reward += reward
        print(ep_reward)
        total_rewards.append(ep_reward)
    avg_reward = np.mean(total_rewards)

    print("---------------------------------------")
    print(f"Evaluation over {eval_episodes} episodes: {avg_reward:.3f}")
    print("---------------------------------------")
    return total_rewards

def run_train(env, eval_env, model, replay_buffer, kwargs, savedir, exp_name, start_timesteps, eval_freq, num_steps=0, max_timesteps=2000, use_frames=False, expl_noise=0.1, batch_size=128, num_eval_episodes=10):
    tb_writer = SummaryWriter(savedir)
    steps = 0
    evaluations = []
    while num_steps < max_timesteps:
        #ts, reward, d, o = env.reset()
        done = False
        state, body =  env.reset()
        if use_frames:
            frame_compressed = compress_frame(env.render(camera_name=args.camera, height=h, width=w))
        ep_reward = 0
        e_step = 0
        while not done:
            if num_steps < start_timesteps:
                action = random_state.uniform(low=kwargs['min_action'], high=kwargs['max_action'], size=kwargs['action_dim'])
            else:
                # Select action randomly or according to policy
                action = (
                        policy.select_action(state)
                        + random_state.normal(0, kwargs['max_action'] * expl_noise, size=kwargs['action_dim'])
                    ).clip(-kwargs['max_action'], kwargs['max_action'])

            if env_type == 'dm_control':
                # we are working on joint position and don't have a joint position controller
                target_joint_position = body[:len(action)] + action
                env.sim.data.qpos[:len(action)] = target_joint_position
                next_state, next_body, reward, done, info = env.step(0*action)
            else:
                next_state, next_body, reward, done, info = env.step(action)
            ep_reward += reward

            if use_frames:
                next_frame_compressed = compress_frame(env.render(camera_name=args.camera, height=h, width=w))

                replay_buffer.add(state, body, action, reward, next_state, next_body, done,
                              frame_compressed=frame_compressed,
                              next_frame_compressed=next_frame_compressed)
                frame_compressed = next_frame_compressed
            else:
                replay_buffer.add(state, body, action, reward, next_state, next_body, done)

            state = next_state
            body = next_body
            if num_steps > start_timesteps:
            #if num_steps > 2000:
                loss_dict = policy.train(num_steps, replay_buffer, batch_size=batch_size)
                tb_writer.add_scalars('DRLloss', loss_dict, num_steps)
            if not num_steps % eval_freq:
                step_filepath = os.path.join(savedir, '{}_{:010d}'.format(exp_name, num_steps))
                policy.save(step_filepath+'.pt')
                # evaluate
                eval_rewards = eval_policy(eval_env, policy, kwargs, num_eval_episodes)
                evaluations.append(eval_rewards)
                tb_writer.add_scalar('eval', np.mean(eval_rewards), num_steps)
                np.save(f"{savedir}/evaluations", evaluations)

            num_steps+=1
            e_step+=1
        tb_writer.add_scalar('train_reward', ep_reward, num_steps)

    step_filepath = os.path.join(savedir, '{}_{:010d}'.format(exp_name, num_steps))
    pickle.dump(replay_buffer, open(step_filepath+'.pkl', 'wb'))
    policy.save(step_filepath+'.pt')

def make_savedir(cfg):
    cnt = 0

    savedir = os.path.join(cfg['experiment']['log_dir'], "%s_%s_%05d_%s_%s_%02d"%(cfg['experiment']['exp_name'],
                                                                        cfg['robot']['env_name'],  cfg['experiment']['seed'],
                                                                        cfg['robot']['robots'][0], cfg['robot']['controller'],  cnt))
    while len(glob(os.path.join(savedir, '*.pt'))):
        cnt +=1
        savedir = os.path.join(cfg['experiment']['log_dir'], "%s_%s_%05d_%s_%s_%02d"%(cfg['experiment']['exp_name'],
                                                                        cfg['robot']['env_name'],  cfg['experiment']['seed'],
                                                                        cfg['robot']['robots'][0], cfg['robot']['controller'],  cnt))
    if not os.path.exists(savedir):
        os.makedirs(savedir)

    os.system('cp -r %s %s'%(args.cfg, os.path.join(savedir, 'cfg.txt')))
    return savedir

def run_eval(env, policy, replay_buffer, kwargs, cfg, cam_dim, savebase):
    robot_name = cfg['robot']['robots'][0]

    env_type = cfg['experiment']['env_type']
    num_steps = 0
    total_steps = replay_buffer.max_size-1
    use_frames = cam_dim[0] > 0
    if use_frames:
        print('recording camera: %s'%args.camera)

    h, w, c = cam_dim
    torques = []
    rewards = []
    while num_steps < total_steps:
        done = False
        state, body =  env.reset()
        if use_frames:
            frame_compressed = compress_frame(env.render(camera_name=args.camera, height=h, width=w))
        ep_reward = 0
        e_step = 0
        while not done:
            # Select action randomly or according to policy
            action = (
                    policy.select_action(state)
                ).clip(-kwargs['max_action'], kwargs['max_action'])

            print(action)
            if robot_name == 'reacher':
                target_joint_position = body[:len(action)] + action
                env.sim.data.qpos[:len(action)] = target_joint_position
                next_state, next_body, reward, done, info = env.step(action)
            else:
                next_state, next_body, reward, done, info = env.step(action) # take a random action
            ep_reward += reward
            if e_step+1 == args.max_eval_timesteps:
                done = True
            if use_frames:
                next_frame_compressed = compress_frame(env.render(camera_name=args.camera, height=h, width=w))
                replay_buffer.add(state, body, action, reward, next_state, next_body, done,
                              frame_compressed=frame_compressed,
                              next_frame_compressed=next_frame_compressed)
                frame_compressed = next_frame_compressed
            else:
                replay_buffer.add(state, body, action, reward, next_state, next_body, done)
            #torques.append(env.env.robots[0].torques)
            #print(action)
            #print(torques[-1])
            state = next_state
            body = next_body
            num_steps+=1
            e_step+=1
        rewards.append(ep_reward)
    #replay_buffer.torques = torques
    return rewards, replay_buffer

def rollout():
    if os.path.isdir(args.load_model):
        load_model = sorted(glob(os.path.join(args.load_model, '*.pt')))[-1]
        cfg_path = os.path.join(args.load_model, 'cfg.cfg')
    else:
        assert args.load_model.endswith('.pt')
        load_model = args.load_model
        load_dir, model_name = os.path.split(args.load_model)

    load_dir, model_name = os.path.split(load_model)
    print('loading model: %s'%load_model)
    cfg_path = os.path.join(load_dir, 'cfg.txt')
    if not os.path.exists(cfg_path):
        cfg_path = os.path.join(load_dir, 'cfg.cfg')
    print('loading cfg: %s'%cfg_path)
    cfg = json.load(open(cfg_path))
    print(cfg)
    #if "kinematic_function" in cfg['experiment'].keys():
    #    kinematic_fn = cfg['experiment']['kinematic_function']
    #    print("setting kinematic function", kinematic_fn)
    #    robot_name = cfg['robot']['robots'][0]
    #    if 'robot_dh' in cfg['robot'].keys():
    #        robot_dh_name = cfg['robot']['robot_dh']
    #    else:
    #        robot_dh_name = cfg['robot']['robots'][0]

    env_type = cfg['experiment']['env_type']
    # TODO find skip_state_keys -

    env = build_env(cfg['robot'], cfg['robot']['frame_stack'], skip_state_keys=skip_state_keys, env_type=env_type, default_camera=args.camera)
    if 'eval_seed' in cfg['experiment'].keys():
        eval_seed = cfg['experiment']['eval_seed'] + 1000
    else:
        eval_seed = cfg['experiment']['seed'] + 1000
    if args.frames: cam_dim = (240,240,3)
    else:
       cam_dim = (0,0,0)

    #if 'eval_replay_buffer_size' in cfg['experiment'].keys():
    #    eval_replay_buffer_size = cfg['experiment']['eval_replay_buffer_size']
    #else:
    eval_replay_buffer_size =  args.max_eval_timesteps*args.num_eval_episodes
    print('running eval for %s steps'%eval_replay_buffer_size)

    policy, kwargs = build_model(cfg['experiment']['policy_name'], env, cfg)

    #if 'kinematic' in cfg['experiment']['policy_name']:
    #    policy.kinematic_fn = eval(kinematic_fn)
    #    policy.kine_loss_weight = cfg['experiment']['kine_loss_weight']
    #    policy.kine_loss_stop = cfg['experiment']['kine_loss_stop']
    savebase = load_model.replace('.pt','_eval_%06d_S%06d'%(eval_replay_buffer_size, eval_seed))
    replay_file = savebase+'.pkl'
    movie_file = savebase+'_%s.mp4' %args.camera
    if not os.path.exists(replay_file):
        policy.load(load_model)
        replay_buffer = build_replay_buffer(cfg, env, eval_replay_buffer_size, cam_dim, eval_seed)

        rewards, replay_buffer = run_eval(env, policy, replay_buffer, kwargs, cfg, cam_dim, savebase)
        pickle.dump(replay_buffer, open(replay_file, 'wb'))
        plt.figure()
        plt.plot(rewards)
        plt.title('eval episode rewards')
        plt.savefig(savebase+'.png')

    else:
        replay_buffer = pickle.load(open(replay_file, 'rb'))
    plot_replay(env, replay_buffer, savebase, frames=args.frames)

def rollout_real():
    if os.path.isdir(args.load_model):
        load_model = sorted(glob(os.path.join(args.load_model, '*.pt')))[-1]
        cfg_path = os.path.join(args.load_model, 'cfg.cfg')
    else:
        assert args.load_model.endswith('.pt')
        load_model = args.load_model
        load_dir, model_name = os.path.split(args.load_model)

    load_dir, model_name = os.path.split(load_model)
    print('loading model: %s'%load_model)
    cfg_path = os.path.join(load_dir, 'cfg.txt')
    if not os.path.exists(cfg_path):
        cfg_path = os.path.join(load_dir, 'cfg.cfg')
    print('loading cfg: %s'%cfg_path)
    cfg = json.load(open(cfg_path))
    print(cfg)
    #if "kinematic_function" in cfg['experiment'].keys():
    #    kinematic_fn = cfg['experiment']['kinematic_function']
    #    print("setting kinematic function", kinematic_fn)
    #    robot_name = cfg['robot']['robots'][0]
    #    if 'robot_dh' in cfg['robot'].keys():
    #        robot_dh_name = cfg['robot']['robot_dh']
    #    else:
    #        robot_dh_name = cfg['robot']['robots'][0]

    env_type = cfg['experiment']['env_type']
    # TODO find skip_state_keys -

    env = build_env(cfg['robot'], cfg['robot']['frame_stack'], skip_state_keys=skip_state_keys, env_type=env_type, default_camera=args.camera)
    env.env = JacoSim2RealWrapper(env.env)
    if 'eval_seed' in cfg['experiment'].keys():
        eval_seed = cfg['experiment']['eval_seed'] + 1000
    else:
        eval_seed = cfg['experiment']['seed'] + 1000
    if args.frames: cam_dim = (240,240,3)
    else:
       cam_dim = (0,0,0)

    #if 'eval_replay_buffer_size' in cfg['experiment'].keys():
    #    eval_replay_buffer_size = cfg['experiment']['eval_replay_buffer_size']
    #else:
    eval_replay_buffer_size =  args.max_eval_timesteps*args.num_eval_episodes
    print('running eval for %s steps'%eval_replay_buffer_size)

    policy,  kwargs = build_model(cfg['experiment']['policy_name'], env, cfg)

    #if 'kinematic' in cfg['experiment']['policy_name']:
    #    policy.kinematic_fn = eval(kinematic_fn)
    #    policy.kine_loss_weight = cfg['experiment']['kine_loss_weight']
    #    policy.kine_loss_stop = cfg['experiment']['kine_loss_stop']
    savebase = load_model.replace('.pt','_eval_%06d_S%06d'%(eval_replay_buffer_size, eval_seed))
    replay_file = savebase+'.pkl'
    movie_file = savebase+'_%s.mp4' %args.camera
    if not os.path.exists(replay_file):
        policy.load(load_model)
        replay_buffer = build_replay_buffer(cfg, env, eval_replay_buffer_size, cam_dim, eval_seed)

        rewards, replay_buffer = run_eval(env, policy, replay_buffer, kwargs, cfg, cam_dim, savebase)
        pickle.dump(replay_buffer, open(replay_file, 'wb'))
        plt.figure()
        plt.plot(rewards)
        plt.title('eval episode rewards')
        plt.savefig(savebase+'.png')

    else:
        replay_buffer = pickle.load(open(replay_file, 'rb'))
    plot_replay(env, replay_buffer, savebase, frames=args.frames)

if __name__ == '__main__':
    import argparse
    from glob import glob
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', default='experiments/base_robosuite.cfg')
    parser.add_argument('--eval', action='store_true', default=False)
    parser.add_argument('--real', action='store_true', default=False)
    parser.add_argument('--frames', action='store_true', default=False)
    parser.add_argument('--camera', default='', choices=['default', 'frontview', 'sideview', 'birdview', 'agentview'])
    parser.add_argument('--load_model', default='')
    parser.add_argument('--num_eval_episodes', default=10, type=int)
    parser.add_argument('--max_eval_timesteps', default=100, type=int)
    args = parser.parse_args()
    # keys that are robot specific

    if args.eval and args.real:
        rollout_real()
    elif args.eval:
        rollout()
    else:
        cfg = json.load(open(args.cfg))
        print(cfg)
        env_type = cfg['experiment']['env_type']
        seed_everything(cfg['experiment']['seed'])
        random_state = np.random.RandomState(cfg['experiment']['seed'])
        env = build_env(cfg['robot'], cfg['robot']['frame_stack'], skip_state_keys=skip_state_keys, env_type=env_type, default_camera=args.camera)
        eval_env = build_env(cfg['robot'], cfg['robot']['frame_stack'], skip_state_keys=skip_state_keys,
                             env_type=cfg['experiment']['env_type'], default_camera=args.camera)
        savedir = make_savedir(cfg)
        policy, kwargs = build_model(cfg['experiment']['policy_name'], env, cfg)
        replay_buffer = build_replay_buffer(cfg, env, cfg['experiment']['replay_buffer_size'], cam_dim=(0,0,0), seed=cfg['experiment']['seed'])
        #robot_name = cfg['robot']['robots'][0]
        #if 'robot_dh' in cfg['robot'].keys():
        #    robot_dh_name = cfg['robot']['robot_dh']
        #else:
        #    robot_dh_name = cfg['robot']['robots'][0]
        #robot_dh = robotDH(robot_name=robot_name, device='cpu')
        #robot_dh.base_matrix = torch.FloatTensor(replay_buffer.base_matrix)

        #if "kinematic_function" in cfg['experiment'].keys():
        #    kinematic_fn = cfg['experiment']['kinematic_function']
        #    policy.kine_loss_weight = cfg['experiment']['kine_loss_weight']
        #    policy.kine_loss_stop = cfg['experiment']['kine_loss_stop']
        #    print("setting kinematic function", kinematic_fn)
        #    policy.kinematic_fn = eval(kinematic_fn)

        run_train(env, eval_env, policy, replay_buffer, kwargs, savedir, cfg['experiment']['exp_name'], cfg['experiment']['start_training'], cfg['experiment']['eval_freq'], num_steps=0, max_timesteps=cfg['experiment']['max_timesteps'], expl_noise=cfg['experiment']['expl_noise'], batch_size=cfg['experiment']['batch_size'], num_eval_episodes=cfg['experiment']['n_eval_episodes'])


import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.autograd as autograd
from .modelutils import SN, gcn_layer, init_weights
from loguru import logger
import NCE.toolbox as tb
from scipy.linalg import block_diag
import random
from utils.datautils import treeDataset, collate_wrapper, hdf5Dataset, hdf5_jsonDataset
from .RL_utils import get_advantages, ppo_loss
from scipy import stats
from sklearn  import metrics

class Crit(nn.Module):
    def __init__(self, config, sn=False, outlayer=True):
        """
        The initial func
        :param config: The dict containing all needed parameters
        """
        super(Crit, self).__init__()
        self.atom_size=config["atom_size"]
        self.bond_size=config["bond_size"]
        self.layers=config["dis_layers"]
        self.config=config
        self.grammar=config["grammar"]
        self.bonds=config["bonds"]
        self.atoms=config["atoms"]
        self.outlayer=outlayer
        self.clock=100

        self.emb_size=32
        self.bond_emb_size=8

        self.atom_linear=nn.Linear(self.atom_size, self.emb_size, bias=True)
        self.bond_linear=nn.Linear(self.bond_size, self.bond_emb_size, bias=True)
        self.layer=[]
        for _ in range(self.layers):
            if sn:
                self.layer.append(SN(gcn_layer(self.emb_size, self.emb_size, self.bond_emb_size, sn=True)))
            else:
                self.layer.append(gcn_layer(self.emb_size, self.emb_size, self.bond_emb_size, bias=True))

        self.gcn=nn.Sequential(*self.layer)
        if sn:
            self.out_layer=SN(nn.Linear(self.emb_size, 1))
        else:
            self.out_layer = nn.Linear(self.emb_size, 1)
        self.act=nn.Sigmoid()

        self.cl=nn.NLLLoss()
        self.acc=[0,0,0]
        if "dis_optimizer" in config:
            self.optimizer=config['dis_optimizer']
        else:
            lr=config["dis_lr"]
            self.optimizer=torch.optim.Adam(self.parameters(), lr=lr,
                                           weight_decay=1e-4)
        self.apply(init_weights)

    def forward(self, inp):
        mol, adj=inp
        maskmol=mol.clamp(0.0, 1.0).float().unsqueeze(2)
        mask=(maskmol.unsqueeze(1)*maskmol.unsqueeze(2))-torch.eye(adj.shape[1]).cuda(mol.device).unsqueeze(0).unsqueeze(3)
        mask=mask.clamp(0,1)
        adj=nn.functional.one_hot(adj, self.bond_size).float()
        adj=self.bond_linear(adj)*mask # not very sure to remain
        mol=nn.functional.one_hot(mol, self.atom_size).float()
        mol=self.atom_linear(mol)*maskmol
        assert (adj*(1-mask)).sum()==0
        adj=adj.permute(0,3,1,2)
        assert (mol*(1-maskmol)).sum()==0

        inp=[mol, adj, maskmol, mask]
        mol, adj, _, _=self.gcn(inp) #the mol is in the shape of n*self.embdsize
        mol=mol.mean(1)#torch.cumsum(mol, 0)

        x=self.out_layer(mol)
        if self.outlayer:
            x=nn.Sigmoid()(x)
        return x

    def reg_onestep(self, data):
        self.train()
        self.zero_grad()
        self.optimizer.zero_grad()
        mol, adj, score=data
        critdevice=next(self.parameters()).device
        mol=mol.cuda(critdevice)
        adj=adj.cuda(critdevice)
        score=(score.cuda(critdevice)<0.2).float().reshape(-1)
        pred=self.forward([mol, adj]).reshape(-1)
        loss=score*torch.log(pred+1e-12)+(1-score)*torch.log(1-pred+1e-12)
        loss=-loss.mean()
        loss.backward()
        self.optimizer.step()
        return loss

    def reg_eval(self, data):
        self.eval()
        Pred=[]
        Score=[]
        for x in data:
            mol, adj, score=x
            mol=mol.cuda()
            adj=adj.cuda()
            with torch.no_grad():
                pred=self.forward([mol, adj]).reshape(-1).cpu().numpy().tolist()
            Pred.extend(pred)
            Score.extend((score<0.2).float().numpy().tolist())
        fpr, tpr, _=metrics.roc_curve(Score, Pred)
        acc=[int(int(y>0.5)==x) for x,y in zip(Score, Pred)]
        acc=sum(acc)*1.0/len(acc)
        logger.info("Datasize is {} Pred size is {} Reg Evaluation auc is {}, acc is {}".format(len(Score), len(Pred), metrics.auc(fpr, tpr), acc))

    def get_score(self, mol):
        with torch.no_grad():
            critdevice = next(self.parameters()).device
            mol, adj = mol.get_mol_adjs([mol], [0])
            ret= self.forward(
                [torch.tensor(mol).to(critdevice), torch.tensor(adj).to(critdevice)])[0]

            ret=float(ret.cpu().numpy()[0])
            return ret, ret
    def get_single_score(self, mol):
        return self.get_score(mol)[0]
    def ft_onestep(self):
        fake=self.config["fake"]
        fakelabel=torch.ones(len(fake[0]))
        loss1=self.reg_onestep([fake[0], fake[1], fakelabel])
        try:
            pos_data=self.config["antib_train_data_pos"].next()
        except:
            self.config["antib_train_data_pos"]=iter(self.config["raw_atd_pos"])
            pos_data=self.config["antib_train_data_pos"].next()
        try:
            neg_data=self.config["antib_train_data_neg"].next()
        except:
            self.config["antib_train_data_neg"]=iter(self.config["raw_atd_neg"])
            neg_data=self.config["antib_train_data_neg"].next()
        loss2=self.reg_onestep(pos_data)
        loss3=self.reg_onestep(neg_data)


import sys
sys.setrecursionlimit(1000000)

c.Spawner.default_url = '/lab'

c.JupyterHub.admin_access = True
c.JupyterHub.authenticator_class = 'nativeauthenticator.NativeAuthenticator'
c.Authenticator.admin_users = {"clamorz"}
c.Authenticator.open_signup = True
c.Authenticator.ask_email_on_signup = True

import pwd
import subprocess

def pre_spawn_hook(spawner):
    username = spawner.user.name
    try:
        pwd.getpwnam(username)
    except KeyError:
        subprocess.check_call(['useradd', '-ms', '/bin/bash', username])

c.Spawner.pre_spawn_hook = pre_spawn_hook

# import shutil
# c.LanguageServerManager.language_servers = {
#     "C": {
#         "argv": [shutil.which("clangd-10")],
#         "languages": ["c", "C", "c++", "C++17", "xcpp17", "c++src", "text/x-c++src"],
#         "version": 1
#     }
# }


import tkinter

class MyGUI:
    def __init__(self):
        # Создать элемент интерфейса главного окна.
        self.main_window = tkinter.Tk()
        self.top_frame = tkinter.Frame()
        self.bottom_frame = tkinter.Frame()

        self.calc_button1 = tkinter.Button(self.bottom_frame,
                                          text='sinister',
                                          command=self.get_info1)
        self.calc_button2 = tkinter.Button(self.bottom_frame,
                                          text='dexter',
                                          command=self.get_info2)
        self.calc_button3 = tkinter.Button(self.bottom_frame,
                                          text='medium',
                                          command=self.get_info3)

        self.calc_button1.pack(side='left')
        self.calc_button2.pack(side='left')
        self.calc_button3.pack(side='left')

        self.top_frame.pack()
        self.bottom_frame.pack()
        # Войти в главный цикл tkinter.
        tkinter.mainloop()

    def get_info1(self):
           self.label1 = tkinter.Label(self.top_frame,
                                    text='левый')

           self.label1.pack(side='top')

    def get_info2(self):
           self.label2 = tkinter.Label(self.top_frame,
                                    text='правый')

           self.label2.pack(side='top')

    def get_info3(self):
           self.label3 = tkinter.Label(self.top_frame,
                                    text='средний')

           self.label3.pack(side='top')

# Создать экземпляр класса MyGUI.
my_gui = MyGUI()


#!/bin/bash/python3
"""As part three of your homework, you are asked to write a program that will ask a user to
enter <filename> as an input argument. Secondly, your program should read the file and saves
each line of the file as an element in a list and returns it (Note, the retuned value should not
have “\n”). The input file might look like the following:
Mazda, 21, 6, 160, 110, 3.9, 2.62
Valiant, 18.1, 6, 225, 105, 2.1, 3, 46
Cadillac, 10.4, 8, 472, 205, 2.93, 5.25
Full version of the file can be downloaded here
(https://gist.github.com/ssheff/9817ab169a7aa6f9143ef7868b3ddda7)
Your program(function) should return a list as the following:
[‘Mazda, 21, 6, 160, 110, 3.9, 2.62’, ‘Valiant, 18.1, 6, 225, 105, 2.1, 3, 46’, ‘Cadillac,
10.4, 8, 472, 205, 2.93, 5.25’]
 """
import csv
import sys

def my_read_function(cars):

    print(cars)
    # Preparing a file
    file_prompt = open(cars, mode='r')

    # Redaing file using realines() function
    car_data = file_prompt.readlines()
    #print(car_data)
    #sys.exit()

    #creating empty list file
    car_output_list =[]

    #stripping "\" from each line"
    for line in car_data:

      #print(line)

      line_with_no_new_line=line.strip()
      car_output_list.append(line_with_no_new_line)

      #print(car_output_list)
      #sys.exit()

    return car_output_list

if __name__ == "__main__":
   print_out_car=my_read_function("cars.csv")
   print(print_out_car)


#  Guess Number Reversed Game
#  Author: Kajetan Mroske
#
#  This program plays a game where the computer tries to guess a number that player has chosen
#  between 0 and 1000. The game ends when the computer correctly guesses the number or the player
#  switches the number during the game.

print('''\nThink about a number in between 0 and 1000...
I'm going to guess it in 10 moves! \n
Pick your number and start the game with enter.''')

start_ = input()  # Player initiates the game with enter.

min_ = 0  # Minimum margin for guess number is set.
max_ = 1000  # Maximum margin for guess number is set.

#  The variable option_ stores the player's feedback on each guess made by the computer.
#  There are three possible options:
#  1. "too big" if the player's number is smaller than the computer's guess.
#  2. "too small" if the player's number is larger than the computer's guess.
#  3. "you are correct" if the computer's guess is correct.
option_ = ""

#  The computer will make at most 10 guesses, which is enough given the range of marigns.
#  The game ends either when the player confirms the computer's guess or when the maximum
#  number of guesses is reached.
for i in range(11):
    if option_ == "3":
        # Player picked the option: 3 - you are correct!
        print("\n Thanks for your time! It was a good game!")
        break  # Loop breaks and program stops.
    if i == 10:
        # Player must have switched the selected number during the game.
        print("\n Don't cheat! You have switched the number! \n")
        break

    #  The computer makes a guess by selecting a number in the middle of the current margins.
    guess_ = int((max_ - min_) / 2) + min_

    #  The computer prints its guess and asks the player for feedback.
    print(f'\n I pick: {guess_}!')
    print('\n Is it the number you chose? \n')

    while True:
        #  The player enters feedback on the computer's guess.
        option_ = input('Pick one of following options: 1 - too big, 2 - too small, 3 - you are correct!')

        if option_ == "3":
            #  If the player confirms the computer's guess, the game ends.
            print("\n Yeah! I won!")
            break  # While loop breaks. Leads to program stop.
        elif option_ == "2":
            #  If the player's number is smaller than the computer's guess,
            #  the minimum margin is replaced with the guess made by the computer.
            min_ = guess_
            break
        elif option_ == "1":
            #  If the player's number is larger than the computer's guess,
            #  the maximum margin is replaced with the guess made by the computer.
            max_ = guess_
            break
        else:
            #  If the player does not select an option, the program asks again.
            continue


"""
Домашнее задание №1
Функции и структуры данных
"""

def power_numbers(*args):
    """
    функция, которая принимает N целых чисел,
    и возвращает список квадратов этих чисел
    >>> power_numbers(1, 2, 5, 7)
    <<< [1, 4, 25, 49]
    """
    return[val**2 for val in args]

# filter types
ODD = "odd"
EVEN = "even"
PRIME = "prime"

def is_prime(num):
    if num ==1:
        return False
    if num % 2 == 0:
        return num == 2
    i = 3
    while i * i <= num and num % i != 0:
        i += 2
    return i * i > num

def filter_numbers(lst, param):
    if param == ODD:
        return list(filter(lambda x: x % 2 != 0, lst))
    if param == EVEN:
        return list(filter(lambda x: x % 2 == 0, lst))
    if param == PRIME:
        return list(filter(lambda x: is_prime(x), lst))


import setuptools

with open("README.md", "r") as fh:
    long_description = fh.read()

setuptools.setup(
    name="tidbits",
    version="0.0.1",
    author="Tang U-Liang",
    author_email="uliang6482@gmail.com",
    description="Utility functions and classes",
    long_description_content_type="text/markdown",
    packages=setuptools.find_packages(),
    url="https://github.com/uliang/MyUtilities",
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: MIT License",
        "Operating System :: MacOS"
    ],
    python_requires=">=3.8"

)


"""add packs

Revision ID: ec89740b6903
Revises: da7556b25d75
Create Date: 2022-08-02 14:55:41.971432

"""
import sqlalchemy as sa
import sqlmodel
from alembic import op

# revision identifiers, used by Alembic.
revision = "ec89740b6903"
down_revision = "da7556b25d75"
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "pack",
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("id", sa.Integer(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name"),
    )
    op.create_index(op.f("ix_pack_name"), "pack", ["name"], unique=False)
    op.create_table(
        "packsentence",
        sa.Column("sentence_id", sa.Integer(), nullable=False),
        sa.Column("pack_id", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(
            ["pack_id"],
            ["pack.id"],
        ),
        sa.ForeignKeyConstraint(
            ["sentence_id"],
            ["sentence.id"],
        ),
        sa.PrimaryKeyConstraint("sentence_id", "pack_id"),
    )
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("packsentence")
    op.drop_index(op.f("ix_pack_name"), table_name="pack")
    op.drop_table("pack")
    # ### end Alembic commands ###


from unittest import TestCase

from hummingbot.connector.exchange.woo_x import woo_x_constants as CONSTANTS, woo_x_web_utils as web_utils

class WebUtilsTests(TestCase):
    def test_rest_url(self):
        url = web_utils.public_rest_url(path_url=CONSTANTS.MARKET_TRADES_PATH, domain=CONSTANTS.DEFAULT_DOMAIN)
        self.assertEqual('https://api.woo.org/v1/public/market_trades', url)
        url = web_utils.public_rest_url(path_url=CONSTANTS.MARKET_TRADES_PATH, domain='woo_x_testnet')
        self.assertEqual('https://api.staging.woo.org/v1/public/market_trades', url)


"""
Various procedures used in the drawing
"""

from multiplicates import overlap, overlap_len
from collections import OrderedDict

class GFF_feature():
    """
    A GFF feature data class.
    It's only designed to work with a particular dialect of GFF and is by no
    means a general purpose GFF parser.
    """
    def __init__(self, **kwargs):
        self.id = kwargs['feat_id']
        for key in ('feature_class', 'start', 'end',
                    'strand', 'parent', 'source'):
                self.__setattr__(key, kwargs[key])

    def __str__(self):
        pass

    def get_id_prefix(self):
        """
        Return the first ':' or '-'-separated element of the ID. If neither
        char is present, return ID
        :return:
        """
        return self.id.split(':')[0].split('-')[0]

def parse_gff_line(line):
    """
    Create the GFF_feature object from a GFF line
    :param line:
    :return:
    """
    l = line.rstrip().split('\t')
    key_values = l[8].split(';')
    kv_dict = {}
    for x in key_values:
        pair = x.split('=')
        kv_dict[pair[0]] = pair[1]
    return GFF_feature(feat_id=kv_dict['ID'],
                       feature_class=l[2],
                       start=int(l[3]),
                       end=int(l[4]),
                       strand=l[6],
                       parent='Parent' in kv_dict and kv_dict['Parent'] or None,
                       source=l[0])

def convert_single_coord(exons, coord, strand='+'):
    """
    Convert a single coordinate using simplified exon representation
    Exons should be sorted by (source) coordinates, whatever their strand
    :param exons:
    :param coord:
    :return:
    """
    exon_lengths = [x[1] - x[0] + 1 for x in exons]
    # A pretty ugly, but working hack. Not sure if it's not missing a codon,
    # but whatever. I need a general feel anyway.
    if strand == '-':
        c = int(sum(exon_lengths)/3) - coord
    elif strand == '+':
        c = coord
    else:
        raise ValueError("Strand should be either \'+\' or \'-\'")
    x = 0
    for index, step in enumerate(exon_lengths):
        if x + step >= c * 3:
            return exons[index][0] - x + (c - 1) * 3 + 1
        else:
            x += step

def protein_coord_to_gene_coord(features, coord_set):
    """
    Take a set of gene features and a set of coordinates (in aminoacids),
    return these coordinates in nucleotides.

    Considers only 'exon' features for intron processing
    :param features:
    :param coord_set:
    :return:
    """
    exons = []
    for feature in features:
        if feature.feature_class == 'exon':
            exons.append((feature.start, feature.end))
    # Assuming all exons are on the same strand
    strand = feature.strand
    return [tuple((convert_single_coord(exons, x, strand) for x in y))
            for y in coord_set]

def convert_coord_dict(features, coords):
    """
    Same as `protein_coord_to_gene_coord` except that it takes an OrderedDict
    whose keys are coordinate pair and returns the OrderedDict where coordinates
    are converted and keys are preserved
    :param features:
    :param coords:
    :return:
    """
    r = OrderedDict()
    exons = []
    for feature in features:
        if feature.feature_class == 'exon':
            exons.append((feature.start, feature.end))
    strand = feature.strand
    for coord_pair in coords:
        r[(convert_single_coord(exons, coord_pair[0], strand),
           convert_single_coord(exons, coord_pair[1], strand))] =\
                    coords[coord_pair]
    return r

def reduce_coord_list(coord_set, initial_counts, overlap_cutoff):
    """
    Reduce a coordinate set, averaging overlapping coordinates.
    Coordinates are considered overlapping when at least `overlap_cutoff`
    of either pair of coordinates is overlapping with the other.
    :param coord_set:
    :param counts:
    :param overlap_cutoff:
    :return:
    """
    assert len(coord_set) == len(initial_counts)
    r = [coord_set[0]]
    counts = [initial_counts[0]]
    for hsp_index, hsp in enumerate(coord_set[1:]):
        merged = False
        for index, coord in enumerate(r):
            if overlap(coord, hsp):
                l = overlap_len(coord, hsp)
                hsp_len = hsp[1] - hsp[0]
                coord_len = coord[1] - coord[0]
                if l >= overlap_cutoff * hsp_len or \
                        l > overlap_cutoff * coord_len:
                    # Forget the smaller one if the other is at least twice
                    # bigger
                    if hsp_len >= coord_len * 2:
                        r[index] = hsp
                    elif coord_len >= hsp_len * 2:
                        break
                    # Average position if they are ~equal
                    else:
                        r[index] = round((hsp[0] + coord[0]) / 2), \
                                   round((hsp[1] + coord[1]) / 2)
                    counts[index] += initial_counts[hsp_index + 1]
                    merged = True
                    break
        if not merged:
            r.append(hsp)
            counts.append(initial_counts[hsp_index + 1])
    return OrderedDict((r[x], counts[x]) for x in range(len(r)))

def reduce_coords(coord_set, overlap_cutoff=0.8):
    """
    Wrapper around reduce_coord_list that converts BLAST hit data and runs
    a second reduction round
    :param coord_set: A list of lists of coord tuples
    :param overlap_cutoff:
    :return:
    """
    coords = [hsp for hit in coord_set for hsp in hit]
    c = [1 for x in coords]
    r = reduce_coord_list(coords, c, overlap_cutoff)
    # A se
    return reduce_coord_list(tuple(r.keys()), tuple(r.values()), overlap_cutoff)

    # r = list(coord_set[0])
    # counts = [1 for _ in r]
    # # A hit is a collection of HSP coordinate pairs
    # # Non-blast coordinates are treated similarly
    # for hit in coord_set[1:]:
    #     for hsp in hit:
    #         merged = False
    #         for index, coord in enumerate(r):
    #             if overlap(coord, hsp):
    #                 l = overlap_len(coord, hsp)
    #                 hsp_len = hsp[1] - hsp[0]
    #                 coord_len = coord[1] - coord[0]
    #                 if l >= overlap_cutoff * hsp_len or \
    #                         l > overlap_cutoff * coord_len:
    #                     # Forget the smaller one if the other is at least twice
    #                     # bigger
    #                     if hsp_len >= coord_len * 2:
    #                         r[index] = hsp
    #                     elif coord_len >= hsp_len * 2:
    #                         break
    #                     # Average position if they are ~equal
    #                     else:
    #                         r[index] = round((hsp[0] + coord[0]) / 2), \
    #                                    round((hsp[1] + coord[1]) / 2)
    #                     counts[index] += 1
    #                     merged = True
    #                     break
    #         if not merged:
    #             r.append(hsp)
    #             counts.append(1)
    # return {r[x]: counts[x] for x in range(len(r))}


import copy
from abc import ABC, abstractmethod
from functools import partial
from typing import Callable, Optional, Union
import warnings

import jax
import jax.numpy as jnp
import numpy as np
import numpyro
import pandas
import tqdm
from jax.random import PRNGKey
from numpyro import deterministic, infer, optim
from numpyro.handlers import condition, trace
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.svi import SVIRunResult
from pysersic.exceptions import ShapeMatchError,KernelError
from pysersic.priors import PySersicMultiPrior, PySersicSourcePrior, base_profile_params
from pysersic.rendering import BaseRenderer, HybridRenderer
from pysersic.results import PySersicResults

from .loss import gaussian_loss

ArrayLike = Union[np.array, jax.numpy.array]

class BaseFitter(ABC):
    """
    Base class for Pysersic Fitters
    """
    def __init__(self,
        data: ArrayLike,
        rms: ArrayLike,
        psf: ArrayLike,
        mask: Optional[ArrayLike] = None,
        loss_func: Optional[Callable] = gaussian_loss,
        renderer: Optional[BaseRenderer] =  HybridRenderer,
        renderer_kwargs: Optional[dict] = {}) -> None:
        """Initialze BaseFitter class

        Parameters
        ----------
        data : ArrayLike
            Science image to be fit
        weight_map : ArrayLike
            Weight map (one over the variance) corresponding to `data`, must be the same shape
        psf_map : ArrayLike
            Pixelized PSF
        mask : Optional[ArrayLike], optional
            Array specifying the mask, `True` or 1 signifies a pixel should be masked, must be same shape as `data`
        sky_model : Optional[str], optional
            One of None, 'flat' or 'tilted-plane' specifying how to model the sky background
        renderer : Optional[BaseRenderer], optional
            The renderer to be used to generate model images, by default HybridRenderer
        renderer_kwargs : Optional[dict], optional
            Any additional arguments to pass to the renderer, by default {}
        """

        self.loss_func = loss_func

        self.data = jnp.array(data.astype(np.float32))
        self.rms = jnp.array(rms.astype(np.float32))
        self.psf = jnp.array(psf.astype(np.float32))
        self.mask = parse_mask(mask,self.data)
        data_isgood = check_input_data(self.data,rms=self.rms,psf=self.psf,mask=jnp.logical_not(self.mask))
        self.renderer = renderer(data.shape, self.psf, **renderer_kwargs)
        self.prior = None
        self.prior_dict = {}

    def set_loss_func(self, loss_func: Callable) -> None:
        """Set loss function to be used for inference

        Parameters
        ----------
        loss_func : Callable
            Functions which takes samples the loss function, see utils/loss.py for some examples.
        """
        self.loss_func = loss_func

    def set_prior(self,parameter: str,
        distribution: numpyro.distributions.Distribution) -> None:
        """Set the prior for a specific parameter

        Parameters
        ----------
        parameter : str
            Parameter to be set
        distribution : numpyro.distributions.Distribution
            Numpyro distribution object corresponding to the prior
        """
        self.prior_dict[parameter] = distribution

    def sample(self,
                num_samples: int = 1000,
                num_warmup: int = 1000,
                num_chains: int = 2,
                init_strategy: Optional[Callable] = infer.init_to_sample,
                sampler_kwargs: Optional[dict] ={},
                mcmc_kwargs: Optional[dict] = {},
                return_model: Optional[bool] = True,
                reparam_func: Optional[Callable] = lambda x: x,
                rkey: Optional[jax.random.PRNGKey] = jax.random.PRNGKey(3)
        ) -> pandas.DataFrame:
        """ Perform inference using a NUTS sampler

        Parameters
        ----------
        num_samples : int, optional
            Number of samples to draw, by default 1000
        num_warmup : int, optional
            Number of warmup samples, by default 1000
        num_chains : int, optional
            Number of chains to run, by default 2
        init_strategy : Optional[Callable], optional
            Initialization strategy for the sampler, by default infer.init_to_sample. See numpyro.infer.initialization for more options
        sampler_kwargs : Optional[dict], optional
            Arguments to pass to the numpyro NUTS kernel
        mcmc_kwargs : Optional[dict], optional
            Arguments to pass to the numpyro MCMC sampler
        return_model : Optional[bool]
            Whether to return the model images but adds a small memory/time overhead, by default True
        rkey : Optional[jax.random.PRNGKey], optional
            PRNG key to use, by default jax.random.PRNGKey(3)

        Returns
        -------
        pandas.DataFrame
            ArviZ summary of posterior
        """

        model =  self.build_model(return_model = return_model)
        model = reparam_func(model)
        sampler =infer.MCMC(infer.NUTS(model,init_strategy=init_strategy, **sampler_kwargs),num_chains=num_chains, num_samples=num_samples, num_warmup=num_warmup,  **mcmc_kwargs)
        sampler.run(rkey)
        self.sampling_results = PySersicResults(data=self.data,rms=self.rms,psf=self.psf,mask=self.mask,loss_func=self.loss_func,renderer=self.renderer)
        self.sampling_results.add_prior(self.prior)
        self.sampling_results.injest_data(sampler = sampler)
        return self.sampling_results

    def _train_SVI(self,
            autoguide: numpyro.infer.autoguide.AutoContinuous,
            method:str,
            ELBO_loss: Optional[Callable] = infer.Trace_ELBO(1),
            lr_init: Optional[int] = 1e-2,
            num_round: Optional[int] = 3,
            SVI_kwargs: Optional[dict]= {},
            train_kwargs: Optional[dict] = {},
            return_model: Optional[bool] = True,
            num_sample: Optional[int] = 1_000,
            rkey: Optional[jax.random.PRNGKey] = jax.random.PRNGKey(6),
            )-> pandas.DataFrame:
        """
        Internal function to perform inference using stochastic variational inference.

        Parameters
        ----------
        autoguide : numpyro.infer.autoguide.AutoContinuous
            Function to build guide
        method: str
            name of method being used; for saving results
        Elbo_loss : Optional[Callable], optional
            Loss function to use, by default infer.Trace_ELBO(1), see numpyro.infer.elbo for more options
        lr_init : Optional[int], optional
            Initial learning rate, by default 1e-2
        num_round : Optional[int], optional
            Number of rounds for training, lr decreases each round, by default 3
        SVI_kwargs : Optional[dict], optional
            Additional arguments to pass to numpyro.infer.SVI, by default {}
        train_kwargs : Optional[dict], optional
            Additional arguments to pass to utils.train_numpyro_svi_early_stop, by default {}
        return_model : Optional[bool]
            Whether to return the model images but adds a small memory/time overhead, by default True
        num_sample: Optional[int]
            Number of samples to draw from trained SVI posterior
        rkey : Optional[jax.random.PRNGKey], optional
            PRNG key, by default jax.random.PRNGKey(6)

        Returns
        -------
        pandas.DataFrame
            ArviZ summary of posterior
        """
        model_cur = self.build_model(return_model=return_model)
        guide = autoguide(model_cur)

        svi_kernel = SVI(model_cur,guide, optim.Adam(0.01), loss = ELBO_loss, **SVI_kwargs)
        numpyro_svi_result = train_numpyro_svi_early_stop(svi_kernel,rkey=rkey,lr_init=lr_init,
                                                        num_round=num_round, **train_kwargs)

        svi_res_dict =  dict(guide = guide, model = model_cur, svi_result = numpyro_svi_result)
        self.svi_results = PySersicResults(data=self.data,rms=self.rms,psf=self.psf,mask=self.mask,loss_func=self.loss_func,renderer=self.renderer)
        self.svi_results.injest_data(svi_res_dict=svi_res_dict,purge_extra=True, num_sample = num_sample)
        self.svi_results.add_prior(self.prior)
        self.svi_results.add_method_used(method)
        return self.svi_results

    def find_MAP(self,
                return_model: Optional[bool] = True,
                rkey: Optional[jax.random.PRNGKey] = jax.random.PRNGKey(3),
                purge_extra: Optional[bool] = True):
        """Find the "best-fit" parameters as the maximum a-posteriori and return a dictionary with values for the parameters.

        Parameters
        ----------
        return_model : Optional[bool]
            Whether to return the model images but adds a small memory/time overhead, by default True
        rkey : Optional[jax.random.PRNGKey], optional
            rng key, by default jax.random.PRNGKey(3)

        Returns
        -------
        dict
            dictionary with fit parameters and their values.
        """
        model_cur = self.build_model(return_model=return_model)
        autoguide_map = infer.autoguide.AutoDelta(model_cur, init_loc_fn= infer.init_to_median)
        train_kwargs = dict(lr_init = 0.05, num_round = 3,frac_lr_decrease  = 0.1, patience = 250, optimizer = optim.Adam, max_train = int(2e4))
        svi_kernel = SVI(model_cur,autoguide_map, optim.Adam(0.01),loss=Trace_ELBO())

        res = train_numpyro_svi_early_stop(svi_kernel,rkey=rkey, **train_kwargs)

        use_dict = {}
        for key in res.params.keys():
            pref = key.split('_auto_loc')[0]
            use_dict[pref] = res.params[key]
        trace_out = trace(condition(model_cur, use_dict)).get_trace()
        real_out = {}

        if purge_extra:
            for key in trace_out:
                if 'Loss' in key:
                    continue
                elif key == 'model':
                    real_out[key] = np.asarray(trace_out[key]['value'])
                elif not ('base' in key or 'auto' in key or 'unwrapped' in key or 'factor' in key or 'loss' in key):
                    real_out[key] = np.round(trace_out[key]['value'], 5 )
        else:
            for k,v in trace_out.items():
                real_out[k] = v['value']

        return real_out

    def estimate_posterior(self,
                        method : str='laplace',
                        return_model: bool = True,
                        num_sample: Optional[int] = 1_000,
                        rkey: Optional[jax.random.PRNGKey] = jax.random.PRNGKey(6),
                        ) -> pandas.DataFrame:
        """Estimate the posterior using a method other than MCMC sampling. Generally faster than MCMC, but could be less accurate.
        Current Options are:
        - 'laplace'
            - Uses the Laplace approximation, which finds the MAP and then uses a Gaussian approximation to the posterior. The covariance matrix is calculated using the Hessian of the log posterior at the MAP. Generally the fastest method but can lead to numerical problems, especially for fitting several (~> 5) sources since this involves inverting a large matrix.
        - 'svi-mvn'
            - Use variational inference to fit a multivariate normal distribution to the posterior. In practice should give similar results to 'laplace', as both assume the posterior is a multivariate Gaussian but is trained differently and generally less susceptible to the numerical issues discussed above.
        - 'svi-flow'
            - Uses a neural flow (currently a BNAF, https://arxiv.org/abs/1904.04676) to approximate the posterior. This is more flexible than the two methods above as the flow is a more flexible representation of the posterior which can capture non-Gaussian behavior.However it is slower to train. Also optimization can be inconsistent so use and interpret with caution. Best to cross-reference with sample on tests cases.

        Parameters
        ----------
        method : str, optional
            method to use, by default 'laplace'
        return_model : Optional[bool]
            Whether to return the model images but adds a small memory/time overhead, by default True
                num_sample: Optional[int]
        Number of samples to draw from trained SVI posterior
            Number of samples to draw from trained SVI posterior
        rkey : Optional[jax.random.PRNGKey], optional
            rng key, by default jax.random.PRNGKey(6)

        """
        assert method in ['laplace','svi-flow','svi-mvn']
        if method=='laplace':
            train_kwargs = dict(patience = 250, max_train = 20000)
            guide_func = partial(infer.autoguide.AutoLaplaceApproximation, init_loc_fn = infer.init_to_median )
            results = self._train_SVI(guide_func,method=method, lr_init = 0.05, train_kwargs=train_kwargs, return_model = return_model, rkey=rkey, num_sample=num_sample)
        elif method=='svi-flow':
            train_kwargs = dict(patience = 500, max_train = 20000)
            guide_func = partial(infer.autoguide.AutoBNAFNormal, num_flows =4,hidden_factors = [5,], init_loc_fn = infer.init_to_median)
            results = self._train_SVI(guide_func,method='svi-flow',ELBO_loss= infer.Trace_ELBO(8),train_kwargs=train_kwargs,num_round=3,lr_init = 1e-2, rkey=rkey,return_model = return_model,num_sample=num_sample)
        elif method=='svi-mvn':
            train_kwargs = dict(patience = 200, max_train = 5000)
            guide_func = partial(infer.autoguide.AutoLowRankMultivariateNormal, init_scale = 5e-3, init_loc_fn = infer.init_to_median)
            results = self._train_SVI(guide_func,method='svi-mvn',ELBO_loss= infer.TraceMeanField_ELBO(16),train_kwargs=train_kwargs,num_round=3,lr_init = 1e-1, rkey=rkey,return_model = return_model,num_sample=num_sample)
        return results

    @abstractmethod
    def build_model(self,return_model: bool = True):
        raise NotImplementedError

class FitSingle(BaseFitter):
    """
    Class used to fit a single source
    """
    def __init__(self,
        data: ArrayLike,
        rms: ArrayLike,
        psf: ArrayLike,
        prior: PySersicSourcePrior,
        mask: Optional[ArrayLike] = None,
        loss_func: Optional[Callable] = gaussian_loss,
        renderer: Optional[BaseRenderer] =  HybridRenderer,
        renderer_kwargs: Optional[dict] = {}) -> None:
        """Initialze FitSingle class

        Parameters
        ----------
        data : ArrayLike
            Science image to be fit
        weight_map : ArrayLike
            Weight map (one over the variance) corresponding to `data`, must be the same shape
        psf_map : ArrayLike
            Pixelized PSF
        mask : Optional[ArrayLike], optional
            Array specifying the mask, `True` or 1 signifies a pixel should be masked, must be same shape as `data`
        sky_model : Optional[str], optional
            One of None, 'flat' or 'tilted-plane' specifying how to model the sky background
        profile_type : Optional[str], optional
            Must be one of: ['sersic','doublesersic','pointsource','exp','dev'] specifying how to paramaterize the source, default 'sersic'
        renderer : Optional[BaseRenderer], optional
            The renderer to be used to generate model images, by default HybridRenderer
        renderer_kwargs : Optional[dict], optional
            Any additional arguments to pass to the renderer, by default {}
        """

        super().__init__(data,rms,psf,loss_func = loss_func, mask = mask, renderer = renderer, renderer_kwargs = renderer_kwargs)

        if prior.profile_type not in self.renderer.profile_types:
            raise AssertionError('Profile must be one of:', self.renderer.profile_types)
        self.profile_type_number = 0
        self.prior = prior

    def build_model(self, return_model: bool = True) -> Callable:
        """ Generate Numpyro model for the specified image, profile and priors

        Returns
        -------
        model: Callable
            Function specifying the current model in Numpyro, can be passed to inference algorithms
        """
        # Make sure all variables have priors
        check_prior = self.prior.check_vars
        if not check_prior:
            raise AssertionError('Not all variables have priors, please run .autogenerate_priors')

        @numpyro.handlers.reparam(config = self.prior.reparam_dict)
        def model(return_model: bool = return_model):
            params = self.prior()
            out = self.renderer.render_source(params, self.prior.profile_type, suffix = self.prior.suffix)

            sky = self.prior.sample_sky(self.renderer.X, self.renderer.Y)

            obs = out + sky
            if return_model:
                deterministic(f'model{self.prior.suffix}', obs)
            self.loss_func(obs, self.data, self.rms, self.mask, suffix = self.prior.suffix)
        return model

class FitMulti(BaseFitter):
    """
    Class used to fit multiple sources within a single image
    """
    def __init__(self,
        data: ArrayLike,
        rms: ArrayLike,
        psf: ArrayLike,
        prior: PySersicMultiPrior,
        mask: Optional[ArrayLike] = None,
        loss_func: Optional[Callable] = gaussian_loss,
        renderer: Optional[BaseRenderer] =  HybridRenderer,
        renderer_kwargs: Optional[dict] = {}) -> None:
        """Initialze FitMulti class

        Parameters
        ----------
        data : ArrayLike
            Science image to be fit
        weight_map : ArrayLike
            Weight map (one over the variance) corresponding to `data`, must be the same shape
        psf_map : ArrayLike
            Pixelized PSF
        mask : Optional[ArrayLike], optional
            Array specifying the mask, `True` or 1 signifies a pixel should be masked, must be same shape as `data`
        sky_model : Optional[str], optional
            One of None, 'flat' or 'tilted-plane' specifying how to model the sky background
        profile_type : Optional[str], optional
            Must be one of: ['sersic','doublesersic','pointsource','exp','dev'] specifying how to paramaterize the source, default 'sersic'
        renderer : Optional[BaseRenderer], optional
            The renderer to be used to generate model images, by default HybridRenderer
        renderer_kwargs : Optional[dict], optional
            Any additional arguments to pass to the renderer, by default {}
        """
        super().__init__(data,rms,psf,mask = mask,loss_func = loss_func,renderer = renderer, renderer_kwargs = renderer_kwargs)
        self.prior = prior

    def build_model(self, return_model: bool = True) -> Callable:
        """Generate Numpyro model for the specified image, profile and priors

        Returns
        -------
        model: Callable
            Function specifying the current model in Numpyro, can be passed to inference algorithms
        """

        @numpyro.handlers.reparam(config = self.prior.reparam_dict)
        def model(return_model: bool = return_model):
            source_variables = self.prior()

            out = self.renderer.render_for_model(source_variables,self.prior.catalog['type'], suffix = self.prior.suffix)
            sky = self.prior.sample_sky(self.renderer.X, self.renderer.Y)

            obs = out + sky

            if return_model:
                deterministic(f'model{self.prior.suffix}', obs)
            self.loss_func(obs, self.data, self.rms, self.mask, suffix = self.prior.suffix)

        return model

    def find_MAP(self,
                return_model: Optional[bool] = True,
                rkey: Optional[jax.random.PRNGKey] = jax.random.PRNGKey(3),
                purge_extra: Optional[bool] = True):
        """Find the "best-fit" parameters as the maximum a-posteriori and return a dictionary with values for the parameters.

        Parameters
        ----------
        return_model : Optional[bool], optional
            whether to return the model image, adds a small time and memory overhead, by default True
        rkey : Optional[jax.random.PRNGKey], optional
            rng key, by default jax.random.PRNGKey(3)

        Returns
        -------
        dict
            dictionary with fit parameters and their values.
        """
        raw_dict = super().find_MAP(return_model=return_model, rkey=rkey, purge_extra= purge_extra)

        if purge_extra:
            results_dict = {}
            for i in range(self.prior.N_sources):
                results_dict[f'source_{i}'] = {}
                for pname in base_profile_params[self.prior.catalog['type'][i]]:
                    results_dict[f'source_{i}'][pname] = raw_dict.pop(pname + f'_{i:d}')
            results_dict.update(raw_dict)
        else:
            results_dict = raw_dict
        return results_dict

def train_numpyro_svi_early_stop(
        svi_class: SVI,
        num_round: Optional[int] = 3,
        max_train: Optional[int] = 5000,
        lr_init: Optional[float] = 0.01,
        frac_lr_decrease: Optional[float]  = 0.1,
        patience: Optional[int] = 100,
        optimizer: Optional[optim._NumPyroOptim] = optim.Adam,
        rkey: Optional[PRNGKey] = PRNGKey(10),
    )-> SVIRunResult:
    """Optimize a SVI model by training for multiple rounds with a deacreasing learning rate, and early stopping for each round

    Parameters
    ----------
    svi_class : SVI
        Initialized numpyo SVI class, note that the optimizer will be overwritten
    num_round : Optional[int], optional
        Number of training rounds, by default 3
    max_train : Optional[int], optional
        Max number of training epochs per ropund, by default 3000
    lr_init : Optional[float], optional
        Initial learning rate, by default 0.1
    frac_lr_decrease : Optional[float], optional
        Multiplicative factor to change learning rate each round, by default 0.1
    patience : Optional[int], optional
        Number of training epochs to wait for improvement, by default 100
    optimizer : Optional[optim._NumPyroOptim], optional
        Optimizer algorithm tro use, by default optim.Adam
    rkey : Optional[PRNGKey], optional
        Jax PRNG key, by default PRNGKey(10)

    Returns
    -------
    SVIRunResult
        SVI Result class containing trained model
    """
    optim_init = optimizer(lr_init)
    svi_class.__setattr__('optim', optim_init)

    init_state = svi_class.init(rkey)
    all_losses = []

    @partial(jax.jit, static_argnums = 1)
    def update_func(state,svi_class,lr):
        svi_class.__setattr__('optim', optimizer(lr))
        state,loss = svi_class.stable_update(state)
        return state,loss

    best_state, best_loss = update_func(init_state, svi_class, lr_init)

    for r in range(num_round):
        losses = []
        wait_counter = 0
        svi_state = copy.copy(best_state)

        if r>0:
            lr_cur = lr_init*frac_lr_decrease**r
            best_loss = jnp.inf
        else:
            lr_cur = lr_init

        with tqdm.trange(1, max_train + 1) as t:
            for j in t:
                svi_state, loss = update_func(svi_state, svi_class, lr_cur)
                if loss < best_loss:
                    best_loss = loss
                    best_state = copy.copy(svi_state)
                    wait_counter = 0
                elif wait_counter >= patience:
                    break
                else:
                    wait_counter += 1
                t.set_postfix_str(f'Round = {r:d},step_size = {lr_cur:.1e} loss: {best_loss:.3e}',refresh=False)
                losses.append(loss)

        all_losses.append(losses)

    return SVIRunResult(svi_class.get_params(best_state), svi_state,losses)

def parse_mask(mask:ArrayLike=None,data:ArrayLike=None):
    if mask is None:
        return jnp.ones_like(data).astype(jnp.bool_)
    else:
        return jnp.logical_not(jnp.array(mask.astype(int))).astype(jnp.bool_)

def check_input_data(data:ArrayLike,rms:ArrayLike,psf:ArrayLike,mask:ArrayLike=None):
    """Check input data for certain conditions and raise warnings or exceptions if needed

    Parameters
    ----------
    data : ArrayLike
        input image (galaxy/cutout)
    rms : ArrayLike
       rms/error map of the data.
    psf : ArrayLike
        pixelized PSF
    mask :ArrayLike, optional
       mask with True/1 indicating a pixel should be masked, by default None

    Raises
    ------
    RMSWarning
        If the rms map is highly discrepant from the rms of the data
    PSFNormalizationWarning
        if the pst normalization is not ~1
    KernelError
        if the provided PSF is larger than the input image (exception)
    MaskWarning
        If more than 50% of the image is masked.
    """
    rms = jnp.array(rms)
    data = jnp.array(data)
    psf = jnp.array(psf)
    if data.shape != rms.shape:
        raise ShapeMatchError('RMS map ndims must match input data ndims.')
    if jnp.median(rms) > 5*jnp.std(data):
        warnings.warn('Input RMS map appears to be highly offset (>5x) in magnitude from the rms of the pixels in the input image.')
    if not jnp.isclose(jnp.sum(psf),1.0,0.1):
        warnings.warn('PSF does not appear to be appropriately normalized; Sum(psf) is more than 0.1 away from 1.')
    if jnp.all(data.shape<psf.shape):
        raise KernelError('PSF pixel image size must be smaller than science image.')
    if mask is not None:
        mask = parse_mask(mask,data)
        if jnp.sum(mask)/jnp.prod(jnp.array(mask.shape))>0.8:
            warnings.warn('More than 80 percent of input image is masked. Is this correct? (Pysersic treats True/1 as masked; you may need to flip your boolean array.) ')
        if mask.shape !=data.shape:
            raise ShapeMatchError('Mask ndims must match input data ndims.')
    return True


if __name__ == '__main__':
    # ANSWER

    # define the age of A & B (period: this year) from the information that we get
    A = 4/14*49
    B = 10/14*49

    # print A & B age result in next 2 years
    print(f"Umur A 2 tahun lagi = {A+2}")
    print(f"Umur B 2 tahun lagi = {B+2}")

TUGAS/Modul 1/Day 2 - Assignment/Soal 4.py


"""Script de précorrection pour l'évaluation à mi-semestre du S1 2019.

TOUS LES RENDUS ONT SUBI UNE CORRECTION MANUELLE, CE SCRIPT N'EST QU'UNE AIDE, NE PAS S'EN SERVIR
SEUL POUR ESTIMER UNE NOTE.
"""

import ast
import math
import os
import pathlib
import shutil
import subprocess
import tempfile
import typing as ty

import tqdm

data_source = pathlib.Path("AI08.txt")

def extract_answers(answer_script: os.PathLike):
    with tempfile.TemporaryDirectory() as tempdir:
        temp_answer = shutil.copy(answer_script, tempdir)
        shutil.copy(data_source, tempdir)
        answers_proc: ty.Union[subprocess.TimeoutExpired, subprocess.CompletedProcess]
        try:
            answers_proc = subprocess.run(
                ["python", temp_answer],
                capture_output=True,
                text=True,
                cwd=tempdir,
                timeout=10,
            )
            run_returned = not answers_proc.returncode
        except subprocess.TimeoutExpired as e:
            answers_proc = e
            run_returned = False
    if answers_proc.stdout is None:
        answers = [""]
    else:
        answers = [l.strip() for l in answers_proc.stdout.splitlines()]
    return run_returned, answers, answers_proc

class AnswerValidation(ty.NamedTuple):
    validable: bool
    correct: ty.Optional[float]
    answer: ty.Optional[str]
    canonical: str

def validate_answer1(answer):
    """Quantité d'ammoniac (en t/an) émise en Alsace en 2004"""
    canonical = "180.13"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    if answer.isspace():
        return AnswerValidation(
            validable=True, correct=0.0, answer=None, canonical=canonical
        )
    try:
        val = ast.literal_eval(answer)
        validable = isinstance(val, float)
    except SyntaxError:
        validable = False
    if validable:
        correct = 1.0 if math.isclose(val, 180.13) else 0.0
    else:
        correct = None
    return AnswerValidation(
        validable=validable, correct=correct, answer=answer, canonical=canonical
    )

def validate_answer2(answer):
    """Combien de noms de régions mentionnées dans ce fichier comportent un nombre de
    lettres supérieur ou égal à 10 (espaces et tirets non compris) ?"""
    canonical = "13"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    if answer.isspace():
        return AnswerValidation(
            validable=True, correct=0.0, answer=None, canonical=canonical
        )
    try:
        val = ast.literal_eval(answer)
        validable = isinstance(val, int)
    except SyntaxError:
        validable = False
    if validable:
        correct = 1.0 if val == 13 else 0.0
    else:
        correct = None
    return AnswerValidation(
        validable=validable, correct=correct, answer=answer, canonical=canonical
    )

def validate_answer3(answer):
    """Quelle est la région pour laquelle les émissions de monoxyde de carbone on été
    maximales en 2008 ?"""
    canonical = "LORRAINE"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    return AnswerValidation(
        validable=True,
        correct=1.0 if answer.lower() == canonical.lower() else 0.0,
        answer=answer,
        canonical=canonical,
    )

def validate_answer4(answer):
    """Même question que 3., mais pour le total des émissions de monoxyde de carbone
    sur l'ensemble de la période 1999 — 2009"""
    canonical = "AQUITAINE"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    return AnswerValidation(
        validable=True,
        correct=1.0 if answer.lower() == canonical.lower() else 0.0,
        answer=answer,
        canonical=canonical,
    )

def validate_answer5(answer):
    """Donner le contenu de la précision donnée entre parenthèses la plus longue du
    fichier."""
    canonical = "et ses composants"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    accepted = ("et ses composants", "NOX y compris NO2")
    if answer in accepted:
        correct = 1.0
    elif answer[1:-1] in accepted:
        correct = 0.9
    else:
        correct = 0.0
    return AnswerValidation(
        validable=True, correct=correct, answer=answer, canonical=canonical
    )

def validate_answer6(answer):
    """Donner, sous forme de liste séparée par des barres `|` et triée par ordre
    alphabétique l'ensemble des noms de composés chimiques ou de catégories de
    composés mentionnés dans ce fichier."""
    canonical = "1,2-dichloroéthane|acide cyanhydrique|ammoniac|arsenic|benzène|cadmium|chloroforme|chlorure d'hydrogène|chlorure de vinyle monomère|composés organiques volatils non méthaniques|dichlorométhane|dioxyde de carbone|fluorure d'hydrogène|mercure|monoxyde de carbone|oxyde d'azote|oxydes de soufre|perchloroéthylène|plomb|poussières|protoxyde d'azote|sulfure d'hydrogène|trichloroéthylène"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    if answer.isspace():
        return AnswerValidation(
            validable=True, correct=0.0, answer=None, canonical=canonical
        )
    parsed = answer.split("|")
    validable = len(parsed) > 1 and parsed == sorted(parsed)
    if validable:
        if answer == canonical:
            correct = 1.0
        # Facile à louper et pas important: virer les parenthèses peut laisser des double espaces
        elif answer.replace("  ", " ") == canonical:
            correct = 1.0
        else:
            correct = 0.0
    else:
        correct = None
    return AnswerValidation(
        validable=validable, correct=correct, answer=answer, canonical=canonical
    )

def validate_answer7(answer):
    """Donner, sous forme de liste séparée par des barres `|` et triée par ordre
    alphabétique l'ensemble des noms de régions présents dans ce fichier en
    capitalisant les débuts de mots"""
    canonical = "Alsace|Aquitaine|Auvergne|Basse Normandie|Bourgogne|Bretagne|Centre-Val De Loire|Champagne-Ardenne|Corse|D.O.M.|Franche Comte|Haute Normandie|Ile De France|Languedoc-Roussillon|Limousin|Lorraine|Midi-Pyrenees|Nord-Pas-De-Calais|Pays De La Loire|Picardie|Poitou-Charentes|Provence-Alpes-Cote D Azur|Rhone-Alpes"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )

    if answer.isspace():
        return AnswerValidation(
            validable=True, correct=0.0, answer=None, canonical=canonical
        )
    parsed = answer.split("|")
    validable = len(parsed) > 1 and parsed == sorted(parsed)
    if validable:
        correct = 1.0 if answer == canonical else 0.0
    else:
        correct = None
    return AnswerValidation(
        validable=validable, correct=correct, answer=answer, canonical=canonical
    )

def validate_answer8(answer):
    """Quelle est le composé dont pour lequel la moyenne annuelle des émissions en
    région Basse Normandie a été la plus élevée sur la période 1999–2009 ?"""
    canonical = "Dioxyde de carbone (CO2)"
    if answer is None:
        return AnswerValidation(
            validable=False, correct=0.0, answer=None, canonical=canonical
        )
    return AnswerValidation(
        validable=True,
        correct=1.0 if answer == canonical else 0.0,
        answer=answer,
        canonical=canonical,
    )

def autocorrect(
    answers: ty.Sequence[ty.Optional[str]]
) -> ty.Generator[AnswerValidation, None, None]:
    yield validate_answer1(answers[0])
    yield validate_answer2(answers[1])
    yield validate_answer3(answers[2])
    yield validate_answer4(answers[3])
    yield validate_answer5(answers[4])
    yield validate_answer6(answers[5])
    yield validate_answer7(answers[6])
    yield validate_answer8(answers[7])

results_file = pathlib.Path("results.tsv")
answers_dir = pathlib.Path("answers")
debugged_answers_dir = pathlib.Path("answers_dbug")
outputs_dir = pathlib.Path("outputs")
outputs_dir.mkdir(exist_ok=True)

with open(results_file, "w") as out_stream:
    # Header
    out_stream.write("name\treturned\tfull\t")
    out_stream.write("\t".join(f"q{i}\t" for i in range(1, 9)))
    out_stream.write("\n")
    pbar = tqdm.tqdm(list(answers_dir.glob("*.py")), unit="file")
    for answer_script in pbar:
        pbar.set_description(f"validating {answer_script.name}")
        full = True
        run_returned, answers, proc = extract_answers(answer_script)
        if not run_returned:
            _, answers, proc = extract_answers(
                debugged_answers_dir / answer_script.name
            )
            answers = [None if a == "DEBUGGED" else a for a in answers]
        if len(answers) < 8:
            answers = [*answers, *("" for _ in range(8 - len(answers)))]
            full = False
        elif len(answers) > 8:
            answers = answers[:8]
            full = False
        out_stream.write(f"{answer_script.stem}\t{int(run_returned)}\t{int(full)}\t")
        for validation in autocorrect(answers):
            out_stream.write(f"{int(validation.validable)}\t{validation.correct}\t")
        out_stream.write("\n")
        with open(
            outputs_dir / f"{answer_script.stem}.txt", "w"
        ) as script_output_stream:
            script_output_stream.write("========== stderr ===========\n")
            script_output_stream.write(proc.stderr if proc.stderr is not None else "")
            script_output_stream.write("\n========= stdout ===========\n")
            script_output_stream.write(proc.stdout if proc.stdout is not None else "")


import cv2
import os
import sys

cam = cv2.VideoCapture(0, cv2.CAP_DSHOW)

start = False
counter = 4000
num_samples = int(sys.argv[1]) + counter
IMG_SAVE_PATH = 'images'
IMG_CLASS_PATH = ' '

try:
    os.mkdir(IMG_SAVE_PATH)
except FileExistsError:
    pass

while True:
    ret, frame = cam.read()
    if not ret:
        print("failed to grab frame")
        break
    if counter == num_samples:
       break

    cv2.rectangle(frame, (10, 30), (310, 330), (0, 255, 0), 2)

    k = cv2.waitKey(1)
    if k == ord('c'):
            name = 'OpenHand6'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('v'):
            name = 'ClosedFist7'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('b'):
            name = 'ThumbsUp7'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('m'):
            name = 'PeaceSignUpwards6'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('a'):
            name = 'PeaceSignDownwards6'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('d'):
            name = 'PointingLeft6'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)
    if k == ord('n'):
            name = 'Nothing6'
            IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, name)
            os.mkdir(IMG_CLASS_PATH)

    if start:
        roi = frame[25:335, 8:315]
        save_path = os.path.join(IMG_CLASS_PATH, '{}.jpg'.format(counter + 1))
        print(save_path)
        cv2.imwrite(save_path, roi)
        counter += 1
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame,"Collecting {}".format(counter),
            (10, 20), font, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
    cv2.imshow("Collecting images", frame)

    if k == ord('s'):
        start = not start

    if k == ord('q'):
            break

print("\n{} image(s) saved to {}".format(counter, IMG_CLASS_PATH))
cam.release()
cv2.destroyAllWindows()