import argparse
import os
import sys
import pickle
import time
sys.path.append(os.getcwd())

from khrylib.utils import *
from khrylib.rl.core.policy_mcp_gaussian import MCPPolicyGaussian
from khrylib.rl.core.policy_additive_gaussian import AdditivePolicyGaussian
from khrylib.rl.core.policy_split import SplitPolicyGaussian
from khrylib.rl.core.policy_cond_gaussian import CondPolicyGaussian
from khrylib.rl.core.critic import GoalValue
from khrylib.rl.agents import AgentPPO
from khrylib.models.mlp import MLP
from torch.utils.tensorboard import SummaryWriter
from motion_imitation.envs.humanoid_im import HumanoidEnv
#from motion_imitation.envs.humanoid_im_nimble import HumanoidNimbleEnv
from motion_imitation.utils.config import Config
from motion_imitation.reward_function import reward_func
from motion_imitation.expert_scheduler import ExpertScheduler

policy_dict = {"multiplicative":MCPPolicyGaussian,
               "additive":AdditivePolicyGaussian,
               "split": SplitPolicyGaussian,
               "cond": CondPolicyGaussian}

parser = argparse.ArgumentParser()
parser.add_argument('--cfg', default=None)
parser.add_argument('--render', action='store_true', default=False)
parser.add_argument('--test', action='store_true', default=False)
parser.add_argument('--num_threads', type=int, default=24)
parser.add_argument('--gpu_index', type=int, default=0)
parser.add_argument('--iter', type=int, default=0)
parser.add_argument('--policy',type=str, default="multiplicative")
parser.add_argument('--show_noise', action='store_true', default=False)
parser.add_argument('--change_goal_freq',type=int, default=3)
parser.add_argument('--simulator',type=str, default="mujoco")
parser.add_argument('--exp_name', type=str, default=None, required=True)
parser.add_argument('--load_name', type=str, default=None)
parser.add_argument('--load_iter', type=int, default=None)
parser.add_argument('--freeze', action="store_true", default=False)
parser.add_argument('--collect_data', action='store_true', default=False)
parser.add_argument('--use_wandb',action="store_true", default=False)
args = parser.parse_args()
if args.render:
    args.num_threads = 1

cfg = Config(args.cfg, args.test, create_dirs=not (args.render or args.iter > 0))

dtype = torch.float64
torch.set_default_dtype(dtype)
device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu')
if torch.cuda.is_available():
    torch.cuda.set_device(args.gpu_index)
np.random.seed(cfg.seed)
torch.manual_seed(cfg.seed)
tb_logger = SummaryWriter(cfg.tb_dir) if not args.render else None
logger = create_logger(os.path.join(cfg.log_dir, 'log.txt'), file_handle=not args.render)

"""environment"""
if args.simulator == "mujoco":
    env = HumanoidEnv(cfg, filename=args.exp_name if args.collect_data else None)
elif args.simulator == "nimble":
    raise NotImplementedError
    env = HumanoidNimbleEnv(cfg)
env.seed(cfg.seed)
#actuators = env.model.actuator_names
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
running_state = ZFilter((state_dim,), clip=5)

exp_sch = ExpertScheduler(n_experts = env.get_num_expert(), expected_freq=args.change_goal_freq)
"""define actor and critic"""
#policy_net = PolicyGaussian(MLP(state_dim, cfg.policy_hsize, cfg.policy_htype), action_dim, log_std=cfg.log_std, fix_std=cfg.fix_std)
policy_net = policy_dict[args.policy](MLP(state_dim + (0 if args.policy!= 'cond' else 8), cfg.policy_hsize, cfg.policy_htype),
                                      action_dim, log_std=cfg.log_std, fix_std=cfg.fix_std, goal_dim=39)
if args.freeze:
    policy_net.freeze()
value_net = GoalValue(MLP(state_dim+39, cfg.value_hsize, cfg.value_htype)) # 39 is goal dim
policy_net.set_goal(env.get_goal())
value_net.set_goal(env.get_goal())
if args.load_iter is not None:
    cp_path = f"{cfg.model_dir}/iter_{args.load_name}_{args.load_iter:04}.p"
    logger.info('loading model from checkpoint: %s' % cp_path)
    model_cp = pickle.load(open(cp_path, "rb"))
    policy_net.load_state_dict(model_cp['policy_dict'])
    value_net.load_state_dict(model_cp['value_dict'])
    running_state = model_cp['running_state']
to_device(device, policy_net, value_net)

if cfg.policy_optimizer == 'Adam':
    optimizer_policy = torch.optim.Adam(policy_net.parameters(), lr=cfg.policy_lr, weight_decay=cfg.policy_weightdecay)
else:
    optimizer_policy = torch.optim.SGD(policy_net.parameters(), lr=cfg.policy_lr, momentum=cfg.policy_momentum, weight_decay=cfg.policy_weightdecay)
if cfg.value_optimizer == 'Adam':
    optimizer_value = torch.optim.Adam(value_net.parameters(), lr=cfg.value_lr, weight_decay=cfg.value_weightdecay)
else:
    optimizer_value = torch.optim.SGD(value_net.parameters(), lr=cfg.value_lr, momentum=cfg.value_momentum, weight_decay=cfg.value_weightdecay)

# reward functions
expert_reward = reward_func[cfg.reward_id]

"""create agent"""
agent = AgentPPO(env=env, dtype=dtype, device=device, running_state=running_state,
                 custom_reward=expert_reward, mean_action=args.render and not args.show_noise,
                 render=args.render, num_threads=args.num_threads,
                 policy_net=policy_net, value_net=value_net,
                 optimizer_policy=optimizer_policy, optimizer_value=optimizer_value, opt_num_epochs=cfg.num_optim_epoch,
                 gamma=cfg.gamma, tau=cfg.tau, clip_epsilon=cfg.clip_epsilon,
                 policy_grad_clip=[(policy_net.parameters(), 40)], end_reward=cfg.end_reward,
                 use_mini_batch=cfg.mini_batch_size < cfg.min_batch_size, mini_batch_size=cfg.mini_batch_size)

def pre_iter_update(i_iter):
    cfg.update_adaptive_params(i_iter)
    agent.set_noise_rate(cfg.adp_noise_rate)
    set_optimizer_lr(optimizer_policy, cfg.adp_policy_lr)
    if cfg.fix_std:
        if args.policy != "split":
            policy_net.action_log_std.fill_(cfg.adp_log_std)
        else:
            policy_net.action_log_std_lb.fill_(cfg.adp_log_std)
            policy_net.action_log_std_ub.fill_(cfg.adp_log_std)
    return

# Initialize wandb
if args.use_wandb:
    import wandb
    wandb.init()

def main_loop():

    if args.render:
        pre_iter_update(args.iter)
        agent.sample(1e8)
    else:
        average_length = np.zeros(agent.env.get_num_expert())
        for i_iter in range(args.iter, cfg.max_iter_num):
            """generate multiple trajectories that reach the minimum batch_size"""
            pre_iter_update(i_iter)
            batch, log = agent.sample(cfg.min_batch_size)
            if cfg.end_reward:
                agent.env.end_reward = log.avg_c_reward * cfg.gamma / (1 - cfg.gamma)

            """update networks"""
            t0 = time.time()
            agent.update_params(batch)
            t1 = time.time()

            """logging"""
            c_info = log.avg_c_info
            logger.info(
                '{}\tT_sample {:.2f}\tT_update {:.2f}\tETA {}\texpert_R_avg {:.4f} {}'
                '\texpert_R_range ({:.4f}, {:.4f})\teps_len {:.2f}'
                .format(i_iter, log.sample_time, t1 - t0, get_eta_str(i_iter, cfg.max_iter_num, t1 - t0 + log.sample_time), log.avg_c_reward,
                        np.array2string(c_info, formatter={'all': lambda x: '%.4f' % x}, separator=','),
                        log.min_c_reward, log.max_c_reward, log.avg_episode_len))
            if args.use_wandb:
                wandb.log({
                    "iter":i_iter,
                    "eps_len": log.avg_episode_len,
                    "expert_R_avg": log.avg_c_reward,
                    "goal":agent.env.get_expert_id()
                })

            tb_logger.add_scalar('total_reward', log.avg_c_reward, i_iter)
            tb_logger.add_scalar('episode_len', log.avg_episode_reward, i_iter)
            for i in range(c_info.shape[0]):
                tb_logger.add_scalar('reward_%d' % i, c_info[i], i_iter)
                tb_logger.add_scalar('eps_reward_%d' % i, log.avg_episode_c_info[i], i_iter)

            """ Update goal from time to time """
            # Adaptive scheduling.s
            exp_sch.record_eps_len(log.avg_episode_len, agent.env.get_expert_id())
            print("Now expert is: ", agent.env.get_expert_id())
            if exp_sch.update(agent.env.get_expert_id(),i_iter):
                agent.env.switch_expert()
                goal = agent.env.get_goal()
                agent.policy_net.set_goal(goal)
                agent.value_net.set_goal(goal)

            """ Save model from time to time """
            if cfg.save_model_interval > 0 and (i_iter+1) % cfg.save_model_interval == 0:
                tb_logger.flush()
                with to_cpu(policy_net, value_net):
                    cp_path = f"{cfg.model_dir}/iter_{args.exp_name}_{i_iter+1:04}.p"
                    model_cp = {'policy_dict': policy_net.state_dict(),
                                'value_dict': value_net.state_dict(),
                                'running_state': running_state}
                    pickle.dump(model_cp, open(cp_path, 'wb'))

            """clean up gpu memory"""
            torch.cuda.empty_cache()
        #env.save_data()
        logger.info('training done!')

main_loop()


# preRecognition.py aims to print out the number of TRs for the current recognition run.

'''
input:
    cfg
    curr_run

output:
    maxTR=xxx

usage example:
    python expScripts/recognition/preRecognition.py -c pilot_sub001.ses1.toml -r 3
'''
import sys,os
if 'watts' in os.getcwd():
    sys.path.append("/home/watts/Desktop/ntblab/kailong/rtSynth_rt/")
elif 'kailong' in os.getcwd():
    sys.path.append("/Users/kailong/Desktop/rtEnv/rtSynth_rt/")
elif 'milgram' in os.getcwd():
    sys.path.append('/gpfs/milgram/project/turk-browne/projects/rtSynth_rt/')

import argparse
import numpy as np
import pandas as pd
if 'watts' in os.getcwd():
    main_dir = "/home/watts/Desktop/ntblab/kailong/rtcloud_rt/" # main_dir = "/home/watts/Desktop/ntblab/kailong/rtcloud_kp/"
else:
    main_dir="/Users/kailong/Desktop/rtEnv/rtSynth_rt/"

sys.path.append(main_dir)

import rtCommon.utils as utils
from rtCommon.utils import loadConfigFile
argParser = argparse.ArgumentParser()
argParser.add_argument('--config', '-c', default='sub001.ses2.toml', type=str, help='experiment file (.json or .toml)')
argParser.add_argument('--run', '-r', default='1', type=str, help='current run')
args = argParser.parse_args()
from rtCommon.cfg_loading import mkdir,cfg_loading
cfg = cfg_loading(args.config)

curr_run=int(args.run)

choose = np.load(f"{cfg.subjects_dir}/{cfg.subjectName}/ses{cfg.session}/recognition/choose.npy")
for curr_run in range(1,1+len(choose)):
    order = f'{cfg.orderFolder}/recognitionOrders_{choose[curr_run - 1]}.csv'
    trial_list = pd.read_csv(order)
    maxTR=trial_list['time'].iloc[-1]/cfg.TR
    print(f"curr_run={curr_run} maxTR={maxTR}")


from plasTeX import Command
from plasTeX.Packages.report import *
from plasTeX.Packages.hyperref import *

def ProcessOptions(options, document):
    from plasTeX.Packages import report
    report.ProcessOptions(options, document)
    document.context['thesection'].format = '${section}'
    document.context['theindex'].counter = 'section'
    document.context['theindex'].level = Environment.SECTION_LEVEL
    document.context['printindex'].counter = 'section'
    document.context['printindex'].level = Command.SECTION_LEVEL
    document.context['bibliography'].counter = 'section'
    document.context['bibliography'].level = Command.SECTION_LEVEL

class appendices(Command):

    class thesection(TheCounter):
        format = '${section.Alph}'

    def invoke(self, tex):
        self.ownerDocument.context.counters['section'].setcounter(0)
        self.ownerDocument.context['thesection'] = type(self).thesection

class separator(Command):
    pass

class fakeqed(Command):
    def invoke(self, tex):
        return []


"""
Menu

options is a start menu.
first_menu is reading the whole file
second_menu is reading line and column
third_menu adds line at the end of the file
fourth_menu loads information from file, adds one line in asked index and writes back to the file
fifth_menu loads information from file and deletes asked line with pop()
sixth_menu reads - writes (csv - txt, txt - csv)
"""

import check
import read_mod
import write_mod

def options():
    print("Choose one option: ")
    choic = input("1. Show all \n2. Find information \n3. Add information \n4. Correct information \n5. Delete information \n6. Export/Import \n")
    ticket = check.check_in(choic, "1234567")
    if choic == "7":
        end_menu("Bye!")
    elif ticket == False:
        end_menu("Sorry, your phrase is incorrect, try again.")
    else: return choic

def first_menu():
    everything = read_mod.reader("0")
    return everything

def second_menu():
    in_str = input("What kind of information would you like? \n1. Line \n2. Column \n")
    if check.check_in(in_str, "12"):
        if in_str == "1":
            line = input("Which line would you like to see?\n")
            inf = in_str + line
            answer = read_mod.reader(inf)
        else:
            column = input("Whick column would you like to see?\n")
            inf = in_str + column
            answer = read_mod.reader(inf)

        end_menu(answer)
    else: end_menu("Incorrect input!")

def third_menu():
    lst = first_menu()
    print("What would you like to add? ")
    add_lst = [input() for i in range(4)]
    lst.append(add_lst)
    write_mod.writer(lst)
    end_menu(first_menu())

def fourth_menu():
    all_inf = first_menu()
    end_menu(all_inf)
    num_line = input("Which line would you like to correct?\n")
    num_line = check.check_in(num_line, "1234567890")
    if num_line == True:
        lst = [input() for i in range(4)]
        all_inf.insert(num_line, lst)
        all_inf.pop(num_line + 1)
        write_mod.writer(all_inf)
        end_menu(first_menu())
    else: end_menu("Sorry, try again.")

def fifth_menu():
    all_inf = first_menu()
    end_menu(all_inf)
    num_line = input("Which line would you like to delete?\n")
    num_line = check.check_in(num_line, "1234567890")
    all_inf.pop(num_line)
    write_mod.writer(all_inf)
    end_menu(first_menu())

def sixth_menu():
    in_str = input("1. Export\n2. Import\n")
    checked_str = check.check_in(in_str, "12")

    if checked_str == True:
        if in_str == "1":
            all_inf = first_menu()
            with open("eximp.txt", "w") as f:
                for elem in all_inf:
                    elem = ",".join([ch for ch in elem if ch != "[" or ch != "]"])
                    f.write(elem + "\n")
        else:
            ls = []
            with open("eximp.txt", "r") as f:
                for line in f:
                    ls.append(line.rstrip().split(','))
            # lst = first_menu()
            # lst.append(ls)
            write_mod.writer(ls)
    else:
        print("Try again.")

    end_menu(first_menu())

def end_menu(output):
    if isinstance(output, list):
        for i in range(len(output)):
            print(*output[i])
    elif isinstance(output, str):
        print(output)
    elif isinstance(output, int):
        print(output)


from django.db import models
from django.db.models.deletion import CASCADE, SET_NULL

MANSIONI_TIPI = (
    ('VE', 'Verificatore'),
    ('VI', 'Vigilantes'),
    ('VO', 'Volontario'),
)

EVENTI_TIPI = (
    ('IM',      'In servizio di mattina'),
    ('IS',      'In servizio di sera'),
    ('I2S',     'In servizio di seconda serata'),
    ('ASS',     'Assente'),
    ('VIME',    'Visita medica'),
    ('A33',     'Articolo 33'),
    ('LIC',     'In licenza'),
    ('MAL',     'In malattia'),
    ('ALTRO',   'Altro'),
)

POLO_OPZ = (
    ('R1',      'Rimessa 1'),
    ('R8',      'Rimessa 8'),
    ('AL',      'P.za Borsellino'),
    ('ST',      'Stazione Centrale'),
    ('RP',      'P.za Repubblica'),
    ('PB',      'P.za Borsa'),
    ('SZ',      'P.za Sanzio'),
    ('OB',      'Due Obelischi'),
    ('NE',      'Nesima'),
    ('FN',      'Fontanarossa'),
    ('VA',      'Villaggio S\'Agata'),
)

FASCIA_OPZ = (
    ('M',        'Mattinale'),
    ('S',        'Serale'),
    ('S2',       'Seconda Serata'),
)

GIORNO_RIPOSO = (
    (7,        '1° Domenica'),
    (14,       '2° Domenica'),
    (21,       '3° Domenica'),
    (28,       '4° Domenica'),
    (35,       '5° Domenica'),
    (42,       '6° Domenica'),
    (49,       '7° Domenica'),
    (56,       '8° Domenica'),
    (63,       '9° Domenica'),
    (70,       '10° Domenica'),
    (77,       '11° Domenica'),
    (83,       'Sabato'),
    (89,       'Venerdi'),
    (95,       'Giovedi'),
    (101,      'Mercoledi'),
    (107,      'Martedi'),
    (1,        'Lunedi'),
)

class SequenzaRiposi(models.Model):
    nome                        = models.CharField(max_length=50, unique=True)
    sequenza                    = models.TextField()

    class Meta:
        verbose_name_plural     = "Sequenze Riposi"
        ordering                = ["nome"]

    def __str__(self):
        return self.nome

#TODO: Rinomina campi: data_inizio... e indice_in...
class Lavoratore(models.Model):
    matricola           = models.CharField(max_length=5, unique=True, blank=True)
    nominativo          = models.CharField(max_length=50)
    mansione            = models.CharField(max_length=2, choices=MANSIONI_TIPI)
    sequenza_riposi     = models.ForeignKey(SequenzaRiposi, on_delete=models.SET_NULL, null=True, blank=True)
    #data_di_riferimento_per_sequenza_riposi
    data_per_sequenza   = models.DateField("data di riferimento ...", null=True, blank=True)
    #indice_giorno_di_riferimento_per_sequenza_riposi
    indice_per_sequenza = models.IntegerField("corrisponde alla ...", null=True, blank=True, choices=GIORNO_RIPOSO)

    class Meta:
        verbose_name_plural     = "Lavoratori"
        ordering                = ["nominativo"]

    def __str__(self):
        return self.nominativo

class Linea(models.Model):
    nome                        = models.CharField(primary_key=True, unique=True, max_length=5)
    polo                        = models.CharField(max_length=2, blank=True, null=True, default='R8', choices=POLO_OPZ,)
    descrizione                 = models.CharField(max_length=255, blank=True, null=True)

    class Meta:
        verbose_name_plural     = "Linee"
        ordering                = ['nome']

    def __str__(self):
        return self.nome

class Turno(models.Model):
    empty_value_display         = '---'
    mansione                    = models.CharField(max_length=2, choices=MANSIONI_TIPI)
    priorita                    = models.IntegerField(default=99)
    linea                       = models.ForeignKey(Linea, to_field='nome', on_delete=models.SET_NULL, blank=True, null=True)
    fascia                      = models.CharField(max_length=2, blank=True, null=True, choices=FASCIA_OPZ,)
    ora_inizio                  = models.TimeField()
    ora_fine                    = models.TimeField()
    polo_monta                  = models.CharField(max_length=2, blank=True, null=True, choices=POLO_OPZ,)
    note                        = models.TextField(max_length=100, blank=True)

    class Meta:
        verbose_name_plural     = "Turni"
        ordering                = ['ora_inizio', 'polo_monta']
        #constraints             = [models.UniqueConstraint(fields=['mansione', 'linea', 'ora_inizio', 'ora_fine', 'polo_monta'], name='unique_turno')]

    def __str__(self):
        return ('%s %s %s %s %s %s %s' % (self.mansione, self.priorita, self.linea, self.ora_inizio.strftime("%H:%M"), self.ora_fine.strftime("%H:%M"), self.polo_monta, self.note))

class Matrice(models.Model):
    data                        = models.DateField(auto_now_add=True)
    nome                        = models.CharField(unique=True, max_length=20)
    turni                       = models.ManyToManyField(Turno)

    class Meta:
        verbose_name_plural     = "Matrici"
        ordering                = ['nome']

    def __str__(self):
        return ('%s del %s' % (self.nome, self.data.strftime("%d %B %Y")))

class Evento(models.Model):
    data                        = models.DateField()
    lavoratore                  = models.ForeignKey(Lavoratore, on_delete=models.CASCADE)
    tipologia                   = models.CharField(max_length=5, choices=EVENTI_TIPI)
    turno                       = models.ForeignKey(Turno, on_delete=models.SET_NULL, null=True, blank=True)

    class Meta:
        verbose_name_plural     = "Eventi"
        ordering                = ['data', 'lavoratore']


# coding: utf-8

import argparse
from hbconfig import Config
from src.data.corpus import Corpus
from src.data.corpusMeta import CorpusMeta
from src.modelHelper import ModelHelper
from src.layerHelper import LayerUtils, LayerHelper
import os
import sys
import time
import random
from src.utils.gazetteer import Gazetteer
import torch
import numpy as np
import random
import torch.nn as nn
from src.utils.gpu_tools import move2cuda
from src.utils.gpu_tools import move2cpu
from src.utils.globalVariable import GLOBAL_VARIABLE
import sys

def setSeed(seed, useGpu):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)
    if useGpu:
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

def parseTag(s):
    if s != '<UNK>' and s != '<START>' and s != '<EOS>' and s != '<PAD>':
        if s == 'O':
            return ('O', '<UNK>')
        return (s.split('-'))
    else:
        return ('<UNK>', '<UNK>')

def getEntities(tags, tag2Idx, idx2Tag):
    entityStart = -1
    entities = []
    current = 0
    while current < len(tags):
        currentLabel, currentType = parseTag(idx2Tag[tags[current]])
        if currentLabel.startswith('O'):
            if entityStart != -1:
                entities.append((entityStart, current - 1, idx2Tag[tags[entityStart]].split('-')[1]))
                entityStart = -1
        elif currentLabel.startswith('B'):
            if entityStart != -1:
                entities.append((entityStart, current - 1, idx2Tag[tags[entityStart]].split('-')[1]))
            entityStart = current
        elif currentLabel.startswith('S'):
            if entityStart != -1:
                entities.append((entityStart, current - 1, idx2Tag[tags[entityStart]].split('-')[1]))
            entities.append((current, current, currentType))
            entityStart = -1
        elif currentLabel.startswith('E'):
            if entityStart != -1:
                entities.append((entityStart, current, currentType))
            entityStart = -1
        current += 1

    if entityStart != -1:
        entities.append((entityStart, len(tags) - 1, idx2Tag[tags[entityStart]].split('-')[1]))
    return entities

def evaluate(model, batchInput, tag2Idx, idx2Tag, dumpFile=None, id2Word=None, rawSentenceBatch=None, useGpu=False):
    model.eval()
    s = 0
    g = 0
    sIgStrict = 0
    sIgLoose = 0
    batchId = -1
    for batchItem in batchInput:
        batchId += 1
        if useGpu:
            batchItem = move2cuda(batchItem)
        wordSeqTensor =  batchItem[0]
        tagSeqTensor = batchItem[1]
        wordSeqLengths = batchItem[2]
        bestScores, decodeIdx = model(batchItem)
        tagSeq = tagSeqTensor.data
        predSeq = decodeIdx.data
        for batchPos in range(len(wordSeqTensor)):
            sentLength = wordSeqLengths[batchPos]
            gold = tagSeq[batchPos][: sentLength]
            pred = reversed(predSeq[batchPos][: sentLength])

            sentenceString = ' '.join(rawSentenceBatch[batchId][batchPos])

            if dumpFile is not None:
                for wordIdx in range(len(pred)):
                    dumpFile[2].write(rawSentenceBatch[batchId][batchPos][wordIdx] + '\t' + idx2Tag[pred[wordIdx]] + '\n')
                dumpFile[2].write('\n')
                for wordIdx in range(len(gold)):
                    dumpFile[5].write(rawSentenceBatch[batchId][batchPos][wordIdx] + '\t' + idx2Tag[gold[wordIdx]] + '\n')
                dumpFile[5].write('\n')

            goldEntities = getEntities(gold, tag2Idx, idx2Tag)
            predEntities = getEntities(pred, tag2Idx, idx2Tag)

            if dumpFile is not None:
                dumpFile[3].write(sentenceString + '\n')
                dumpFile[4].write(sentenceString + '\n')
                for entity in goldEntities:
                    dumpFile[3].write(' '.join(rawSentenceBatch[batchId][batchPos][entity[0]: entity[1] + 1]) + '\t' + str(entity[0]) + '\t'
                                      + str(entity[1]) + '\t' + entity[2] + '\n')
                for entity in predEntities:
                    dumpFile[4].write(' '.join(rawSentenceBatch[batchId][batchPos][entity[0]: entity[1] + 1]) + '\t' + str(entity[0]) + '\t'
                                      + str(entity[1]) + '\t' + entity[2] + '\n')
                dumpFile[3].write('\n')
                dumpFile[4].write('\n')

            s += len(predEntities)
            g += len(goldEntities)
            if len(predEntities) == 0:
                if dumpFile is not None:
                    dumpFile[0].write(sentenceString)
                    for entity in goldEntities:
                        dumpFile[0].write('\t'.join([' '.join(rawSentenceBatch[batchId][batchPos][entity[0]: entity[1] + 1]), str(entity[0]), str(entity[1]), entity[2]]) + '\n')
                continue
            predIdx = 0
            strictFalse = []
            looseFalse = []
            for entityIdx in range(len(goldEntities)):
                losseMatch = False
                strictMath = False
                entityStart, entityEnd, entityType = goldEntities[entityIdx]
                while predIdx < len(predEntities) - 1 and predEntities[predIdx][1] < entityStart:
                    predIdx += 1
                if entityType == predEntities[predIdx][2]:
                    if entityStart == predEntities[predIdx][0] and entityEnd == predEntities[predIdx][1]:
                        sIgStrict += 1
                        sIgLoose += 1
                        strictMath = True
                    elif max(predEntities[predIdx][0], entityStart) <= min(entityEnd, predEntities[predIdx][1]):
                        sIgLoose += 1
                        losseMatch = True
                if dumpFile is not None and not strictMath:
                    if not losseMatch:
                        strictFalse.append(goldEntities[entityIdx])
                    else:
                        looseFalse.append([goldEntities[entityIdx], predEntities[predIdx]])
            if dumpFile is not None:
                if len(strictFalse) > 0:
                    dumpFile[0].write(sentenceString + '\n')
                    for entity in strictFalse:
                        dumpFile[0].write('\t'.join([''.join(rawSentenceBatch[batchId][batchPos][entity[0]: entity[1] + 1]), str(entity[0]), str(entity[1]), entity[2]]) + '\n')
                if len(looseFalse) > 0:
                    dumpFile[1].write(sentenceString + '\n')
                    for entity in looseFalse:
                        dumpFile[1].write('\t'.join([''.join(rawSentenceBatch[batchId][batchPos][entity[0][0]: entity[0][1] + 1]), str(entity[0][0]), str(entity[0][1]), entity[0][2], ''.join(rawSentenceBatch[batchId][batchPos][entity[1][0]: entity[1][1] + 1]), str(entity[1][0]), str(entity[1][1]), entity[1][2]]) + '\n')
    ps = sIgStrict / s if s > 0 else 0
    pl = sIgLoose / s if s > 0 else 0
    rs = sIgStrict / g if g > 0 else 0
    rl = sIgLoose / g if g > 0 else 0
    fs = 2 * ps * rs / (ps + rs) if ps + rs > 0 else 0
    fl = 2 * pl * rl / (pl + rl) if pl + rl > 0 else 0
    return ps * 100, pl * 100, rs * 100, rl * 100, fs * 100, fl * 100

def train(Config):
    setSeed(Config.train.seed, Config.use_gpu)

    # load data
    dataPath = os.path.join(Config.data.data_base_path)
    if Config.eval.do_eval:
        testData = Corpus(os.path.join(dataPath, 'test.txt'), Config.data.use_normalized_word)
    trainData = Corpus(os.path.join(dataPath, 'train.txt'), Config.data.use_normalized_word)
    if Config.use_dev:
        devData = Corpus(os.path.join(dataPath, 'dev.txt'), Config.data.use_normalized_word)
    if Config.data.get('word_embedding', False):
        trainData.words.update(testData.words)
        trainData.words.update(devData.words)
        trainData.bwbigrams.update(testData.bwbigrams)
        trainData.bwbigrams.update(devData.bwbigrams)
        trainData.fwbigrams.update(testData.fwbigrams)
        trainData.fwbigrams.update(devData.fwbigrams)

    # gazetter
    gazetters = []
    if Config.model.get("graph_emb", None) is not None and Config.model.graph_emb.get('gazetter', None) is not None:
        for name in Config.model.graph_emb.gazetter.to_dict():
            if name != 'get_tag':
                gaItem = Config.model.graph_emb.gazetter.get(name)
                gazetters.append(Gazetteer(name, gaItem['path'], Config.data.use_normalized_word, gaItem['emb_dim'], gaItem['method'], gaItem['space'], gaItem['match_ignore_case'], embedding=gaItem.get('embedding', None)))
                gazetters[-1].matchCorpus(testData)
                gazetters[-1].matchCorpus(devData)
                gazetters[-1].matchCorpus(trainData)

    # generate train data meta
    print("Generating corpus meta...")
    trainMeta = CorpusMeta(trainData, Config)
    trainMeta.updateMaxSentLength(devData)
    trainMeta.updateMaxSentLength(testData)
    trainMeta.gazetters = gazetters

    # initialize model
    print("Initializing model...")
    layerUtils = LayerUtils(Config, trainMeta)
    layerHelper = LayerHelper(Config, layerUtils)

    modelHelper = ModelHelper(Config, layerHelper)
    if Config.model.get('load_from_pretrain', False):
        model = modelHelper.loadModel()
    else:
        model = modelHelper.getModel()
    if Config.use_gpu:
        model.cuda()
    trainer = modelHelper.getTrainer(model)
    if len(Config.gpu_num) > 1:
        trainer = nn.DataParallel(trainer, device_ids=device_ids)

    # genenerate batch input
    print("Generating batch input...")

    if Config.use_dev:
        if len(Config.gpu_num) > 1:
            devRawSentenceBatch = model.module.getRawSentenceBatches(devData, trainMeta, Config.train.batch_size)
        else:
            devRawSentenceBatch = model.getRawSentenceBatches(devData, trainMeta, Config.train.batch_size)
    if Config.eval.do_eval:
        if len(Config.gpu_num) > 1:
            testRawSentenceBatch = model.module.getRawSentenceBatches(testData, trainMeta, Config.train.batch_size)
        else:
            testRawSentenceBatch = model.getRawSentenceBatches(testData, trainMeta, Config.train.batch_size)
    logFolder = os.path.join(Config.log_folder, Config.model.model_name)
    if not os.path.exists(logFolder):
        os.makedirs(logFolder)

    # train
    trainLog = open(os.path.join(logFolder, "train.log"), 'w', encoding='utf-8')
    trainLog.write(Config.__str__()+'\n')
    print(Config.__str__())
    trainStart = time.time()
    bestF1 = -1
    bestEpoch = -1
    for epoch in range(1, Config.train.epoch + 1):
        random.shuffle(trainData.utterances)
        if len(Config.gpu_num) > 1:
            batchInput = model.module.generateBatchInput(trainData, trainMeta, Config.train.batch_size)
        else:
            batchInput = model.generateBatchInput(trainData, trainMeta, Config.train.batch_size)
        if Config.train.optimizer == "sgd":
            trainer = modelHelper.lrDecay(trainer, epoch)
        model.train()
        model.zero_grad()
        sampleLoss = 0
        sampleCount = 0
        tempStart = time.time()
        epochStart = time.time()
        epochLoss = 0
        for batchItem in batchInput:
            if Config.use_gpu:
                batchItem = move2cuda(batchItem)
            if len(Config.gpu_num) > 1:
                loss = model(batchItem, negMode=True).sum()
            else:
                loss = model.negLogLikelihoodLoss(batchItem)
            sampleLoss += loss.data
            epochLoss += float(loss.data)
            sampleCount += len(batchItem[0])
            loss.backward()
            trainer.step()
            model.zero_grad()
            if sampleCount >= Config.train.report_frequence:
                tempTime = time.time()
                tempCost = tempTime - tempStart
                tempStart = tempTime
                print("Process {} sentences. Loss: {:.2f}. Time: {:.2f}".format(sampleCount, loss/sampleCount, tempCost))
                trainLog.write("Process {} sentences. Loss: {:.2f}. Time: {:.2f}".format(sampleCount, loss/sampleCount, tempCost) + '\n')
                if sampleLoss > 1e8 or str(sampleLoss) == "nan":
                    print("ERROR: LOSS EXPLOSION (>1e8) ! PLEASE SET PROPER PARAMETERS AND STRUCTURE! EXIT....")
                    trainLog.write("ERROR: LOSS EXPLOSION (>1e8) ! PLEASE SET PROPER PARAMETERS AND STRUCTURE! EXIT...." + '\n')
                    exit(1)
                sys.stdout.flush()
                sampleLoss = 0
                sampleCount = 0
        epochCost = time.time() - epochStart
        if Config.use_dev:
            if len(Config.gpu_num) > 1:
                devBatchInput = model.module.generateBatchInput(devData, trainMeta, Config.train.batch_size)
            else:
                devBatchInput = model.generateBatchInput(devData, trainMeta, Config.train.batch_size)
        if Config.eval.do_eval:
            if len(Config.gpu_num) > 1:
                testBatchInpit = model.module.generateBatchInput(testData, trainMeta, Config.train.batch_size)
            else:
                testBatchInpit = model.generateBatchInput(testData, trainMeta, Config.train.batch_size)
        if Config.use_dev and epoch % Config.train.dev_epoch_frequence == 0:
            model.eval()
            ps, pl, rs, rl, fs, fl = evaluate(model, devBatchInput, trainMeta.tag2Idx, trainMeta.idx2Tag, rawSentenceBatch=devRawSentenceBatch, useGpu=Config.use_gpu)
            if fs >= bestF1:
                bestF1 = fs
                bestEpoch = epoch
                modelHelper.saveModel(model, epoch)
            print("Epoch {}. Loss: {:.2f}. Time: {:.2f}".format(epoch, epochLoss, epochCost))
            print('Dev P: {:.2f} R: {:.2f} F1: {:.2f}'.format(ps, rs, fs))
            ps, pl, rs, rl, fs, fl = evaluate(model, testBatchInpit, trainMeta.tag2Idx, trainMeta.idx2Tag, rawSentenceBatch=testRawSentenceBatch, useGpu=Config.use_gpu)
            print('Test P: {:.2f} R: {:.2f} F1: {:.2f}'.format(ps, rs, fs))
            trainLog.write("Epoch {}. Loss: {:.2f}. Dev P: {:.2f} R: {:.2f} F1: {:.2f}. Time: {:.2f}".format(
                    epoch, epochLoss, ps, rs, fs, epochCost) + '\n')
        else:
            print("Epoch {}. Loss: {:.2f}. Time: {:.2f}".format(epoch, epochLoss, epochCost))
            trainLog.write("Epoch {}. Loss: {:.2f}. Time: {:.2f}".format(epoch, epochLoss, epochCost) + '\n')
        sys.stdout.flush()

    if bestF1 == 0:
        modelHelper.saveModel(model, 0)

    print("Finish training. Best epoch {}. F1 {:.2f}. Time {:.2f}".format(bestEpoch, bestF1, time.time() - trainStart))
    trainLog.write("Finish training. Time {:.2f}".format(time.time() - trainStart) + '\n')
    trainLog.close()

    if Config.eval.do_eval:
        model = modelHelper.loadModel(bestEpoch)
        if Config.use_gpu:
            model.cuda()
        testBatchInpit = model.generateBatchInput(testData, trainMeta, 20)
        testRawSentenceBatch = model.getRawSentenceBatches(testData, trainMeta, 20)
        testLog = open(os.path.join(logFolder, 'test.log'), 'w', encoding='utf-8')
        strictFalse = open(os.path.join(logFolder, 'strictFalse'), 'w', encoding='utf-8')
        looseFalse = open(os.path.join(logFolder, 'looseFalse'), 'w', encoding='utf-8')
        rawOutput = open(os.path.join(logFolder, 'testOutput'), 'w', encoding='utf-8')
        goldOutput = open(os.path.join(logFolder, 'testGold'), 'w', encoding='utf-8')
        goldEntityOutput = open(os.path.join(logFolder, 'testGoldEntities'), 'w', encoding='utf-8')
        predEntityOutput = open(os.path.join(logFolder, 'testPredEntities'), 'w', encoding='utf-8')
        ps, pl, rs, rl, fs, fl = evaluate(model, testBatchInpit, trainMeta.tag2Idx, trainMeta.idx2Tag, (
        strictFalse, looseFalse, rawOutput, goldEntityOutput, predEntityOutput, goldOutput), trainMeta.idx2Word,
                                          testRawSentenceBatch, useGpu=Config.use_gpu)
        print("Test P: {:.2f} R: {:.2f} F1: {:.2f}".format(ps, rs, fs))
        testLog.write("Test P: {:.2f} R: {:.2f} F1: {:.2f}".format(ps, rs, fs) + '\n')
        strictFalse.close()
        looseFalse.close()
        testLog.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config', type=str, default='baseline')
    args = parser.parse_args()

    Config(args.config)
    if Config.mode == 'train':
        train(Config)


import torch
import csv
import spacy
import re
import pickle
import random
import csv
from nltk import word_tokenize
import nltk
nltk.download('punkt')
import time

import argparse
import os
import pprint
import numpy as np
from nltk.tokenize import word_tokenize
from io import open
import sys
import json
import pickle
from torch import nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm, trange
from operator import itemgetter
from sklearn.model_selection import train_test_split
from sklearn.metrics import average_precision_score
from sklearn.metrics import accuracy_score
import transformers as tf
from transformers import BertTokenizer
from transformers import PYTORCH_PRETRAINED_BERT_CACHE
from transformers import BertConfig, WEIGHTS_NAME, CONFIG_NAME
from transformers import AdamW, get_linear_schedule_with_warmup

import torch.utils.data as data
from transformers import BertModel
from transformers import BertPreTrainedModel

from model import BERT_GenderClassifier

from string import punctuation

import pickle
import pandas as pd
import random
from collections import namedtuple

# import age variables from utils
from age_utils import (
  young_words,
  old_words,
  age_words
)
# import age functiond from utils
from age_utils import (
  gender_pickle_generator,
  race_pickle_generator,
  label_human_caption,
  label_human_annotations,
  match_labels,
  make_train_test_split,
  save_leak_model
)
from age_dataset import BERT_ANN_leak_data, BERT_MODEL_leak_data

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", default='captioning', type=str)
    parser.add_argument("--cap_model", default='sat', type=str)
    parser.add_argument("--gender_or_race", default='gender', type=str)
    parser.add_argument("--calc_ann_leak", default=False, type=bool)
    parser.add_argument("--calc_model_leak", default=False, type=bool)
    parser.add_argument("--calc_mw_acc", default=True, type=bool)
    parser.add_argument("--test_ratio", default=0.1, type=float)
    parser.add_argument("--balanced_data", default=True, type=bool)
    parser.add_argument("--mask_age_words", default=True, type=bool)
    parser.add_argument("--freeze_bert", default=False, type=bool)
    parser.add_argument("--store_topk_age_pred", default=False, type=bool)
    parser.add_argument("--topk_age_pred", default=50, type=int)
    parser.add_argument("--calc_score", default=True, type=bool)
    parser.add_argument("--align_vocab", default=True, type=bool)

    parser.add_argument("--batch_size", default=64, type=int)
    parser.add_argument("--seed", default=0, type=int)
    parser.add_argument("--num_epochs", default=5, type=int)
    parser.add_argument("--learning_rate", default=1e-5, type=float)
    parser.add_argument("--optimizer", default='adamw', type=str, help="adamw or adam")
    parser.add_argument("--adam_correct_bias", default=True, type=bool)
    parser.add_argument("--beta1", default=0.9, type=float)
    parser.add_argument("--beta2", default=0.98, type=float, help="0.999:huggingface, 0.98:RoBERTa paper")
    parser.add_argument("--adam_epsilon", default=1e-6, type=float, help="Epsilon for Adam optimizer. 1e-8:first, 1e-6:RoBERTa paper")
    parser.add_argument("--weight_decay", default=0.1, type=float, help="Weight deay if we apply some. 0.001:first, 0.01:RoBERTa")
    parser.add_argument("--coco_lk_model_dir", default='/Bias/leakage/', type=str)
    parser.add_argument("--workers", default=1, type=int)

    parser.add_argument("--max_seq_length", default=64, type=int)
    parser.add_argument("--hidden_dim", default=256, type=int)
    parser.add_argument("--output_dim", default=1, type=int)
    parser.add_argument("--dropout", default=0.5, type=float)

    # Add an argument to save the model after training is finished.
    parser.add_argument('--save_model', default = True, type = bool)
    parser.add_argument('--every', default = 5, type = int, help = 'Saving period of the model in epochs')

    return parser

def binary_accuracy(preds, y):
    """
    Returns accuracy per batch, i.e. if you get 8/10 right, this returns 0.8, NOT 8
    """
    #round predictions to the closest integer
    rounded_preds = torch.round(torch.sigmoid(preds))
    correct = (rounded_preds == y).float() #convert into float for division
    acc = correct.sum() / len(correct)
    return acc

def calc_random_acc_score(args, model, test_dataloader):
    print("--- Random guess --")
    model = model.cuda()
    optimizer = None
    epoch = None
    val_loss, val_acc, val_young_acc, val_old_acc, avg_score = calc_leak_epoch_pass(epoch, test_dataloader, model, optimizer, False, print_every=500)

    return val_acc, val_loss, val_young_acc, val_old_acc, avg_score

def calc_leak(args, model, train_dataloader, test_dataloader):
    model = model.cuda()
    print("Num of Trainable Parameters:", sum(p.numel() for p in model.parameters() if p.requires_grad))
    if args.optimizer == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay = 1e-5)
    elif args.optimizer == 'adamw':
        param_optimizer = list(model.named_parameters())
        param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]
        #no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        no_decay = ['bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            #{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, betas=(args.beta1, args.beta2), correct_bias=args.adam_correct_bias, eps=args.adam_epsilon)

    train_loss_arr = list()
    train_acc_arr = list()

    # training
    for epoch in range(args.num_epochs):
        # train
        train_loss, train_acc, _, _, _ = calc_leak_epoch_pass(epoch, train_dataloader, model, optimizer, True, print_every=500)
        train_loss_arr.append(train_loss)
        train_acc_arr.append(train_acc)
        if epoch % 5 == 0:
            print('train, {0}, train loss: {1:.2f}, train acc: {2:.2f}'.format(epoch, \
                train_loss*100, train_acc*100))
        if args.save_model:
              if epoch % args.every == 0:
                path = "saved_models/{}"
                if args.freeze_bert:
                  file_name = "age_annotation_{}_model_pretrained_bert_{}_seed_{}_epoch_{}.pt"
                else:
                  file_name = "age_annotation_{}_model_bert_{}_seed_{}_epoch_{}.pt"
                save_leak_model(model,epoch, file_name,path, args)

    print("Finish training")
    print('{0}: train acc: {1:2f}'.format(epoch, train_acc))
    # We also save the model after training is finished
    if args.freeze_bert:
      file_name = "age_annotation_{}_model_pretrained_bert_{}_seed_{}_epoch_{}.pt"
    else:
      file_name = "age_annotation_{}_model_bert_{}_seed_{}_epoch_{}.pt"
    path = "saved_models/{}"
    save_leak_model(model,epoch, file_name,path, args)
    # validation
    val_loss, val_acc, val_young_acc, val_old_acc, avg_score = calc_leak_epoch_pass(epoch, test_dataloader, model, optimizer, False, print_every=500)
    print('val, {0}, val loss: {1:.2f}, val acc: {2:.2f}'.format(epoch, val_loss*100, val_acc *100))
    if args.calc_mw_acc:
        print('val, {0}, val loss: {1:.2f}, Young val acc: {2:.2f}'.format(epoch, val_loss*100, val_young_acc *100))
        print('val, {0}, val loss: {1:.2f}, Old val acc: {2:.2f}'.format(epoch, val_loss*100, val_old_acc *100))

    return val_acc, val_loss, val_young_acc, val_old_acc, avg_score

def calc_leak_epoch_pass(epoch, data_loader, model, optimizer, training, print_every):
    t_loss = 0.0
    n_processed = 0
    preds = list()
    truth = list()
    young_preds_all, old_preds_all = list(), list()
    young_truth_all, old_truth_all = list(), list()

    if training:
        model.train()
    else:
        model.eval()

    if args.store_topk_age_pred:
        all_young_pred_values, all_old_pred_values = [], []
        all_young_inputs, all_old_inputs = [], []

    total_score = 0 # for calculate scores

    cnt_data = 0
    for ind, (input_ids, attention_mask, token_type_ids, age_target, img_id) in enumerate(data_loader): # images are not provided
        input_ids = input_ids.cuda()
        attention_mask = attention_mask.cuda()
        token_type_ids = token_type_ids.cuda()

        age_target = torch.squeeze(age_target).cuda()
        predictions = model(input_ids, attention_mask, token_type_ids)
        cnt_data += predictions.size(0)

        loss = F.cross_entropy(predictions, age_target, reduction='mean')

        if not training and args.store_topk_age_pred:
            pred_values = np.amax(F.softmax(predictions, dim=1).cpu().detach().numpy(), axis=1).tolist()
            pred_ages = np.argmax(F.softmax(predictions, dim=1).cpu().detach().numpy(), axis=1)

            for pv, pg, imid, ids in zip(pred_values, pred_ages, img_id, input_ids):
                tokens = model.tokenizer.convert_ids_to_tokens(ids)
                text = model.tokenizer.convert_tokens_to_string(tokens)
                text = text.replace('[PAD]', '')
                if pg == 0:
                    all_young_pred_values.append(pv)
                    all_young_inputs.append({'img_id': imid, 'text': text})
                else:
                    all_old_pred_values.append(pv)
                    all_old_inputs.append({'img_id': imid, 'text': text})

        if not training and args.calc_score:
            pred_ages = np.argmax(F.softmax(predictions, dim=1).cpu().detach(), axis=1)
            age_target = age_target.cpu().detach()
            correct = torch.eq(pred_ages, age_target)
            #if ind == 0:
            #    print('correct:', correct, correct.shape)

            pred_score_tensor = torch.zeros_like(correct, dtype=float)
            for i in range(pred_score_tensor.size(0)):
                young_score = F.softmax(predictions, dim=1).cpu().detach()[i,0]
                old_score = F.softmax(predictions, dim=1).cpu().detach()[i,1]
                if young_score >= old_score:
                    pred_score = young_score
                else:
                    pred_score = old_score

                pred_score_tensor[i] = pred_score

            scores_tensor = correct.int() * pred_score_tensor
            correct_score_sum = torch.sum(scores_tensor)
            total_score += correct_score_sum.item()

        predictions = np.argmax(F.softmax(predictions, dim=1).cpu().detach().numpy(), axis=1)
        preds += predictions.tolist()
        truth += age_target.cpu().numpy().tolist()

        if training:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        t_loss += loss.item()
        n_processed += len(age_target)

        if (ind + 1) % print_every == 0 and training:
            print('{0}: task loss: {1:4f}'.format(ind + 1, t_loss / n_processed))

        if args.calc_mw_acc and not training:
            young_target_ind = [i for i, x in enumerate(age_target.cpu().numpy().tolist()) if x == 0]
            old_target_ind = [i for i, x in enumerate(age_target.cpu().numpy().tolist()) if x == 1]
            young_pred = [*itemgetter(*young_target_ind)(predictions.tolist())]
            old_pred = [*itemgetter(*old_target_ind)(predictions.tolist())]
            young_target = [*itemgetter(*young_target_ind)(age_target.cpu().numpy().tolist())]
            old_target = [*itemgetter(*old_target_ind)(age_target.cpu().numpy().tolist())]
            young_preds_all += young_pred
            young_truth_all += young_target
            old_preds_all += old_pred
            old_truth_all += old_target

    acc = accuracy_score(truth, preds)

    if args.calc_mw_acc and not training:
        young_acc = accuracy_score(young_truth_all, young_preds_all)
        old_acc = accuracy_score(old_truth_all, old_preds_all)
    else:
        young_acc, old_acc = None, None

    return t_loss / n_processed, acc, young_acc, old_acc, total_score / cnt_data

def main(args):
    '''
    Modified version of main from gender_bert_leakage.

    Arguments
    ---------
    args : NamedTuple
        The parsed arguments. See get_parser() for more details.
    '''
    # remove previously generated train/val/test files
    if os.path.exists('bias_data/train.csv'):
        os.remove('bias_data/train.csv')
    if os.path.exists('bias_data/val.csv'):
        os.remove('bias_data/val.csv')
    if os.path.exists('bias_data/test.csv'):
        os.remove('bias_data/test.csv')

    # set up seeds, gpu learning
    torch.backends.cudnn.deterministic = True
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    n_gpu = torch.cuda.device_count()
    print("device: {} n_gpu: {}".format(device, n_gpu))
    if n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

    age_val_obj_cap_entries = label_human_annotations(gender_pickle_generator('human'),young_words,old_words) # Human captions
    selected_cap_age_entries = match_labels(age_val_obj_cap_entries,gender_pickle_generator(args.cap_model))
    print("Generated using {}".format(args.cap_model))

    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    ##################### ANN LIC score #######################
    if args.calc_ann_leak:
        print('--- calc ANN LIC score ---')
        ## Captioning ##
        if args.task == 'captioning':
            print('-- Task is Captioning --')
            d_train, d_test = make_train_test_split(args, age_val_obj_cap_entries)
            val_acc_list = []
            score_list = []
            young_acc_list, old_acc_list = [], []
            rand_acc_list = []
            rand_score_list = []
            for caption_ind in range(5):
                trainANNCAPobject = BERT_ANN_leak_data(d_train, d_test, args, age_val_obj_cap_entries, age_words, tokenizer,
                                                args.max_seq_length, split='train', caption_ind=caption_ind)
                testANNCAPobject = BERT_ANN_leak_data(d_train, d_test, args, age_val_obj_cap_entries, age_words, tokenizer,
                                                args.max_seq_length, split='test', caption_ind=caption_ind)
                train_dataloader = torch.utils.data.DataLoader(trainANNCAPobject, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.workers)
                test_dataloader = torch.utils.data.DataLoader(testANNCAPobject, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers)
                # initialize gender classifier (works for age too)
                model = BERT_GenderClassifier(args, tokenizer)
                # calculate random predictions
                val_acc, val_loss, val_young_acc, val_old_acc, avg_score = calc_random_acc_score(args, model, test_dataloader)
                rand_acc_list.append(val_acc)
                rand_score_list.append(avg_score)
                # train and test
                val_acc, val_loss, val_young_acc, val_old_acc, avg_score = calc_leak(args, model, train_dataloader, test_dataloader)
                val_acc_list.append(val_acc)
                young_acc_list.append(val_young_acc)
                old_acc_list.append(val_old_acc)
                score_list.append(avg_score)

            old_avg_acc = sum(old_acc_list) / len(old_acc_list)
            young_avg_acc = sum(young_acc_list) / len(young_acc_list)
            avg_score = sum(score_list) / len(score_list)
            print('########### Reluts ##########')
            print(f"LIC score (LIC_D): {avg_score*100:.2f}%")
            #print(f"\t Old Accuracy: {old_avg_acc*100:.2f}%")
            #print(f"\t Young Accuracy: {young_avg_acc*100:.2f}%")
            print('#############################')

    ##################### MODEL LIC score #######################
    if args.calc_model_leak:
        print('--- calc MODEL LIC score---')
        ## Captioning ##
        if args.task == 'captioning':
            print('-- Task is Captioning --')
            d_train, d_test = make_train_test_split(args, selected_cap_age_entries)
            trainMODELCAPobject = BERT_MODEL_leak_data(d_train, d_test, args, selected_cap_age_entries, age_words, tokenizer,
                                                args.max_seq_length, split='train')
            testMODELCAPobject = BERT_MODEL_leak_data(d_train, d_test, args, selected_cap_age_entries, age_words, tokenizer,
                                                args.max_seq_length, split='test')
            train_dataloader = torch.utils.data.DataLoader(trainMODELCAPobject, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.workers)
            test_dataloader = torch.utils.data.DataLoader(testMODELCAPobject, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers)
            # initialize gender classifier
            model = BERT_GenderClassifier(args, tokenizer)
            # calculate random predictions
            rand_val_acc, rand_val_loss, rand_val_young_acc, rand_val_old_acc, rand_avg_score = calc_random_acc_score(args, model, test_dataloader)
            # train and test
            val_acc, val_loss, val_young_acc, val_old_acc, avg_score = calc_leak(args, model, train_dataloader, test_dataloader)

            print('########### Reluts ##########')
            print(f'LIC score (LIC_M): {avg_score*100:.2f}%')
            #print(f'\t Young. Acc: {val_young_acc*100:.2f}%')
            #print(f'\t Old. Acc: {val_old_acc*100:.2f}%')
            print('#############################')

if __name__ == "__main__":
    parser = get_parser()
    args, unknown = parser.parse_known_args()
    print()
    print("---Start---")
    print('Seed:', args.seed)
    print("Epoch:", args.num_epochs)
    print("Freeze BERT:", args.freeze_bert)
    print("Learning rate:", args.learning_rate)
    print("Batch size:", args.batch_size)
    print("Calculate score:", args.calc_score)
    print("Task:", args.task)
    if args.task == 'captioning' and args.calc_model_leak:
        print("Captioning model:", args.cap_model)
    print("Protected attribute: Age")
    print("Save Model : ", args.save_model)
    if args.save_model:
      print("Saving Every : ", args.every)
    if args.calc_ann_leak:
        print('Align vocab:', args.align_vocab)
        if args.align_vocab:
            print('Vocab of ', args.cap_model)

    print()

    main(args)


"""Contains landing page forms.
"""
from django import forms
from frontend.utils import get_essay_types_for_display

class CallToActionForm(forms.Form):
    """This form initiates user interaction. It collects
    their email and prompts them to make an order

    Args:
        forms (class): describes the logical structure of an object
    """
    error_messages = {
        'email': {
            'required': 'Please provide an email address',
            'invalid': 'The email format is invalid'
        },
        'number_of_pages': {
            'required': 'You must provide the total number of pages',
            'invalid': 'Please provide a whole number for number of pages',
            'min_value': 'Number of pages cannot be less than one'
        },
        'date_deadline': {
            'required': 'Please provide a deadline for your essay',
            'invalid': (
                'Please provide a valid date, format should be '
                '(Month day, year) eg. Oct 25, 2006'
            )
        },
        'time_deadline': {
            'invalid': (
                'The time deadline provided is invalid'
                'it should be in the format (H:M) eg. 6:30'
            ),
        },
        'essay_type': {
            'required': 'Please provide the type of essay',
            'invalid_choice': 'Please select one of the choices provided'
        },
        'terms_and_conditions': {
            'required': 'Please accepts our terms and conditions to proceed'
        }
    }
    essays = get_essay_types_for_display()
    CHOICES = ()
    if essays['code'] == 'success':
        CHOICES = tuple((k, v) for k, v in essays['payload'].items())
    email = forms.EmailField(error_messages=error_messages['email'])
    number_of_pages = forms.IntegerField(
        min_value=1, error_messages=error_messages['number_of_pages']
    )
    date_deadline = forms.DateField(
        input_formats=['%b %d, %Y'],
        error_messages=error_messages['date_deadline']
    )
    time_deadline = forms.TimeField(
        required=False,
        input_formats=['%H:%M'],
        error_messages=error_messages['time_deadline']
    )
    essay_type = forms.ChoiceField(
        choices=CHOICES,
        error_messages=error_messages['essay_type']
    )
    terms_and_conditions = forms.BooleanField(
        error_messages=error_messages['terms_and_conditions']
    )
    bonus_and_promotions = forms.BooleanField(required=False)


from collections import deque,namedtuple
import networkx as nx

_LockRequest = namedtuple('_LockRequest', ['transaction', 'lock_type'])
_LockTableResponse = namedtuple('_LockTableResponse', ['response', 'transactions'])

class LockTable(object):
    """LockTable for a single site.

    Attributes
    ----------
    lock_table : dict of dicts
        The lock table is a dict, keyed by the variable names. Each
        variable name maps to three objects:
            RL : set
                A set of Transactions currently holding read-locks on x
            WL : Transaction or None
                Transaction currently holding (exclusive) write-lock on x
            waiting : deque
                A deque of _LockRequests, each with a transaction and the
                requested lock type
    waits_for : networkx.DiGraph
        A digraph with waits-for edges for locks at this specific site.
    """
    def __init__(self,variables):

        self.lock_table = {
            x: {
                'RL': set(),
                'WL': None,
                'waiting': deque()
            } for x in variables
        }

        self.waits_for = nx.DiGraph()

    def RL_available(self,T,x):
        """ Check if transaction T can get a read lock on x. A transaction
        can get a read lock on x if:
            - It already holds a read lock or a write lock on x
            - It does not already hold a sufficient lock, and no
              other transactions are either (i) currently holding a
              write lock, or (ii) have write-lock requests waiting in the
              queue.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting RL on x
        x : Variable name
            Variable for which RL is requested

        Returns
        -------
        _LockTableResponse
            _LockTableResponse with two attributes:
                response : bool
                    True if RL available, False if not
                transactions : set
                    Set of transactions for which T would need to wait for RL.
                    Empty if lock is available.
        """
        holding_write_lock = self.lock_table[x]['WL']
        waiting_for_write_lock = set(other.transaction for other in self.lock_table[x]['waiting']
                                     if other.lock_type == 'WL').difference([T])

        # Case 1 -- already holding a sufficient lock
        if (holding_write_lock == T) or (T in self.lock_table[x]['RL']):
            return _LockTableResponse(response=True,
                                      transactions=set())

        # Case 2 -- not already holding a sufficient lock
        else:
            # Case 2a -- another transaction is currently holding a WL
            if (holding_write_lock is not None):
                return _LockTableResponse(
                    response=False,
                    transactions=set([holding_write_lock]).union(waiting_for_write_lock)
                )

            # Case 2b -- no other transaction is currently holding a write lock,
            # but some other transaction is waiting for a write lock
            elif (holding_write_lock is None) and (len(waiting_for_write_lock) > 0):
                return _LockTableResponse(
                    response=False,
                    transactions=waiting_for_write_lock
                )

            # Case 2c -- all other cases, T can have lock
            else:
                return _LockTableResponse(response=True,
                                          transactions=set())

    def WL_available(self,T,x):
        """ Check if transaction T can get write lock on x. A transaction can get a write
        lock on x if
            - It already holds a write-lock on x
            - It is the only transaction holding a read-lock on x (implying no T' has a WL),
              and there are also no other transactions waiting in the queue for a lock on x.
              Note it would be possible for T to hold a RL on x, and thus create a queue
              (if there is a T' requesting a WL, followed by any sequence of read
              or write lock requests)
            - No transactions are holding either a read or write lock on x and no
              transactions are waiting on a lock.
        In any other case, (i) some transaction has a lock on x, and/or (ii) other
        transactions are waiting for a lock -- so this lock request waits for
        all other queued transactions.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting WL on x
        x : Variable name
            Variable for which WL is requested

        Returns
        -------
        _LockTableResponse
            _LockTableResponse with two attributes:
                response : bool
                    True if WL available, False if not
                transactions : set
                    Set of transactions for which T would need to wait for WL.
                    Empty if lock is available.
        """

        waiting_for_lock = set(other.transaction for other in self.lock_table[x]['waiting']).difference([T])

        # Case 1 -- T already has WL
        if (self.lock_table[x]['WL'] == T):
            return _LockTableResponse(response=True,
                                     transactions=set())

        # Case 2 -- T is the only transaction with a RL and no T' are waiting for locks
        if ((T in self.lock_table[x]['RL']) and
            (len(self.lock_table[x]['RL']) == 1) and
            (len(waiting_for_lock) == 0)):
            return _LockTableResponse(response=True,
                                     transactions=set())

        # Case 3 -- no locks held and no waiters
        if ((self.lock_table[x]['WL'] is None) and
            (len(self.lock_table[x]['RL']) == 0) and
            (len(waiting_for_lock) == 0)):
            return _LockTableResponse(response=True,
                                     transactions=set())

        # All other cases, T needs to wait for any non-T transaction that is
        # currently holding or waiting for a lock on x

        waiting_for = set()

        # Waits for transaction holding write lock
        holding_write_lock = self.lock_table[x]['WL']
        if (holding_write_lock is not None):
            waiting_for = set([holding_write_lock])

        # Waits for any other transaction currently holding RL on x
        other_transactions_w_RLs = self.lock_table[x]['RL'].difference([T])
        if other_transactions_w_RLs is not None:
            waiting_for = waiting_for.union(other_transactions_w_RLs)

        # Waits for any other transactions waiting for a lock
        if waiting_for_lock is not None:
            waiting_for = waiting_for.union(waiting_for_lock)

        return _LockTableResponse(
            response=False,
            transactions=waiting_for
        )

    def give_transaction_WL(self,T,x):
        """ Gives transaction T a WL on x.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting WL on x
        x : Variable name
            Variable for which WL is requested

        Notes
        ------------
        - Updates both the LockTable, and the set of WL's held by T
        """
        self.lock_table[x]['WL'] = T
        T.write_locks[self].add(x)

    def request_write_lock(self,T,x):
        """Check if the transaction T is waiting on a WL
        on x at this site. If not, have T request a WL
        on this site, and give T this lock if it's available, or
        add it T to the queue if not. In either case, add
        this site to T's locks_needed set.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting WL on x
        x : Variable name
            Variable for which WL is requested

        Notes
        ------------
        - Updates both the lock table, and either T's locks_needed
          or the locks it is holding.
        """
        if x not in T.locks_needed[self]:
            # If not, then we need to try to give T the lock, or
            # add it to the queue, and update T's lock_needed dict
            T.locks_needed[self].add(x)
            lock_available, waiting_for = self.WL_available(T,x)
            if lock_available:
                self.give_transaction_WL(T,x)
            else:
                self.add_transaction_to_lock_queue(T,x,waiting_for,'WL')

    def give_transaction_RL(self,T,x):
        """ Gives transaction T a RL on x.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting RL on x
        x : Variable name
            Variable for which RL is requested

        Notes
        ------------
        - Updates both the LockTable, and the set of RL's held by T
        """
        self.lock_table[x]['RL'].add(T)
        T.read_locks[self].add(x)

    def add_transaction_to_lock_queue(self,T,x,waiting_for,lock_type):
        """ Adds transaction T to the queue for a lock on variable x,
        and updates the waits-for graph.

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting lock on x
        x : Variable name
            Variable for which lock is requested
        waiting_for : set
            Set of transactions blocking T's request for the x lock
        lock_type : str (RL or WL)
            Type of lock requested

        Notes
        -----
        - Adds the lock on x at this site to the set of locks T
          needs to proceed.
        """
        self.lock_table[x]['waiting'].append(_LockRequest(transaction=T,lock_type=lock_type))
        T.locks_needed[self].add(x)
        for other in waiting_for:
            self.waits_for.add_edge(T,other)

    def reassign_locks(self,x):
        """
        Reassign locks (as needed) for variable x. Called after a
        Transaction is ended (committed or aborted), when locks are released.

        Parameters
        ----------
        x : Variable name
            Variable to reassign locks for

        Notes
        ------------
        - If locks reassigned, updates the locks held by transactions
          given new locks
        """
        gave_new_lock = True
        waiters = self.lock_table[x]['waiting']
        while gave_new_lock and (len(waiters)>0):
            next_in_line = self.lock_table[x]['waiting'].popleft()
            next_lock_type = next_in_line.lock_type
            requesting_T = next_in_line.transaction

            if (next_lock_type == 'RL'):
                # Check if RL available
                # Case 1: No transaction holds a WL
                if (self.lock_table[x]['WL'] is None):
                    self.give_transaction_RL(requesting_T,x)
                # Case 2: some other transaction has WL
                else:
                    gave_new_lock = False
                    # Put this transaction back at front of queue
                    self.lock_table[x]['waiting'].appendleft(next_in_line)

            # Requesting WL
            else:
                # Case 1: No other transaction holds a WL and
                # no other transactions have read locks
                if ((self.lock_table[x]['WL'] is None) and
                    (len(self.lock_table[x]['RL'].difference([requesting_T]))==0)):
                    self.give_transaction_WL(requesting_T,x)
                # Case 2: All other cases can't give lock
                else:
                    gave_new_lock = False
                    # Put this transaction back at front of queue
                    self.lock_table[x]['waiting'].appendleft(next_in_line)

    def remove_T_from_lock_table_and_waitsfor(self,T):
        """
        Completely removes T from the lock table and waits
        for graph, and updates the lock table (reassigning locks).

        Parameters
        ----------
        T : ReadWriteTransaction
            Transaction requesting lock

        Notes
        ------------
        - Remove T from the lock table
        - Remove T from the waits_for graph
        - Reassigns locks (where there were blocked transactions
          waiting on the lock)
        """
        for x,row in self.lock_table.items():
            # Remove T from the RL set if it is present
            row['RL'].discard(T)
            # Reset the WL entry to None if T was holding WL
            if row['WL'] == T:
                row['WL'] = None
            # Remove T from the queue for locks on x
            new_wait_list = deque()
            for waiter in row['waiting']:
                if waiter.transaction != T:
                    new_wait_list.append(waiter)
            row['waiting'] = new_wait_list

            # Then reassign locks
            self.reassign_locks(x)

        # Finally remove node from waits_for graph
        if T in self.waits_for.nodes():
            self.waits_for.remove_node(T)


from regexast import *
import pydot
import copy
import pdb

class FSM(object):
    def __init__(self, letter_cls):
        self.letter_cls = letter_cls
        self.initial_state = None
        self.states = []
        self.state_count = 1

    def match(self, letter_list):
        state = self.initial_state
        for letter in letter_list:
            state = state.transition(letter)
            if state is None:
                return False

        return state.terminal

    def get_state(self, id):
        for state in self.states:
            if state._id == id:
                return state
        else:
            return None

    def clear_visited(self):
        for state in self.states:
            state.clear_visited()

    def add_state(self, state):
        state._id = self.state_count
        self.state_count += 1
        self.states += [state]

    def draw(self, filename, show_dead=False):
        graph = pydot.Dot(graph_type='digraph')
        state_hash = {}
        def _draw(state):
            if not state.visited:
                node = pydot.Node(str(state._id), label=('"%s"' % str(state._id)))
                if state.terminal:
                    node.set_shape('"doublecircle"')
                else:
                    node.set_shape('"circle"')

                graph.add_node(node)
                state_hash[state] = node
                state.visit()

                for transition_set, new_state in state.edges:
                    if new_state is not self.dead_state or show_dead:
                        new_node = _draw(new_state)
                        letter_str_list = []
                        for letter_id in list(transition_set):
                            if letter_id == -1:
                                letter_str_list += ["0"]
                            else:
                                letter_str_list += [str(self.letter_cls(letter_id))]

                        edge_label = '"' +  ", ".join(letter_str_list) + '"'
                        graph_edge = pydot.Edge(node, new_node)
                        graph_edge.set_label(edge_label)
                        graph.add_edge(graph_edge)

                return node
            else:
                return state_hash[state]

        _draw(self.initial_state)
        graph.write_png(filename)
        self.clear_visited()

class NFSM(FSM):
    def __init__(self, letter_cls):
        super(NFSM, self).__init__(letter_cls)
        self.make_dead_state()

    def make_dead_state(self):
        self.dead_state = FSMState([])
        self.dead_state._id = 0
        self.states += [self.dead_state]

    def reverse_ids(self):
        def _reverse_ids(state):
            if not state.visited:
                state._id = self.state_count - state._id
                state.visit()
                for _, next_state in state.edges:
                    _reverse_ids(next_state)

        _reverse_ids(self.initial_state)
        self.clear_visited()

    def _connect_dead_state(self, state):
        dead_transitions = set()
        for letter in self.letter_cls.all():
            if letter.id() not in state.valid_transitions():
                dead_transitions = dead_transitions | {letter.id()}

        state.add_edges([(dead_transitions, self.dead_state)])

    def connect_dead_state(self):
        for state in self.states:
            if state is not self.dead_state:
                self._connect_dead_state(state)

    def is_impossible(self, state):
        if state is self.dead_state:
            return False

        if state.terminal:
            return False

        for letters, subs_state in state.edges:
            if subs_state is not self.dead_state and \
                    not (len(letters) == 1 and tuple(letters)[0] == -1):
                return False

        return True

    def remove_impossible(self, state):
        # pdb.set_trace()
        epsilon_edges = state.get_epsilon_edges()
        for other_state in self.states:
            if state is other_state:
                continue

            edges_to = other_state.get_edges_to(state)
            if len(edges_to) != 0:
                other_state.remove_edges_to(state)
                new_edges = []
                for incoming_letters, _ in edges_to:
                    for _, destination_state in epsilon_edges:
                        new_edges += [(incoming_letters, destination_state)]

                other_state.add_edges(new_edges)

    def remove_impossible_states(self):
        for state in self.states:
            if self.is_impossible(state):
                self.remove_impossible(state)

    def from_regex(self, ast):
        self.initial_transitions = self._from_regex(ast)
        self.initial_state = FSMState([self.initial_transitions])
        self.add_state(self.initial_state)
        self.connect_dead_state()
        self.reverse_ids()
        self.clear_visited()
        self.remove_impossible_states()

    def _from_regex(self, ast, next_transition=None):
        if ast.type == RegexASTNode.CAT_EXPR:
            transition_B = self._from_regex(ast.right, next_transition=next_transition)
            mid_state = FSMState([transition_B])
            self.add_state(mid_state)
            transition_A = self._from_regex(ast.left, next_transition=(-1, mid_state))
            return transition_A
        elif ast.type == RegexASTNode.LETTER:
            state_B = None
            if next_transition is None:
                state_B = FSMState([], terminal=True)
                self.add_state(state_B)
            else:
                _, state_B = next_transition

            return (ast.value.id(), state_B)
        elif ast.type == RegexASTNode.EPSILON:
            state_B = None
            if next_transition is None:
                state_B = FSMState([], terminal=True)
                self.add_state(state_B)
            else:
                _, state_B = next_transition

            return (-1, state_B)
        elif ast.type == RegexASTNode.OR_EXPR:
            if next_transition is None:
                final_state = FSMState([], terminal=True)
                self.add_state(final_state)
                next_transition = (-1, final_state)

            transitions_A = self._from_regex(ast.left, next_transition=next_transition)
            transitions_B = self._from_regex(ast.right, next_transition=next_transition)
            or_initial_state = FSMState([transitions_A, transitions_B])
            self.add_state(or_initial_state)
            return (-1, or_initial_state)
        elif ast.type == RegexASTNode.REP_EXPR:
            terminal = (next_transition is None)
            loop_state = FSMState([], terminal=terminal)
            loop_init_transition = self._from_regex(ast.left, next_transition=(-1, loop_state))

            if next_transition is not None:
                loop_state.add_edges([loop_init_transition, next_transition])
            else:
                loop_state.add_edges([loop_init_transition])

            self.add_state(loop_state)
            return (-1, loop_state)
        else:
            raise Exception("Unknown AST type.")

class FSMState(object):
    def __init__(self, edges, id=None, terminal=False):
        # edges is expected to be of the form
        #    [ ({letter1_id1, letter1_id2}, new_state1), ... , (lettern_id, new_staten)]
        self.edges = []
        self._id = id
        self.add_edges(edges)
        self.visited = False

        self.terminal = terminal

    def __hash__(self):
        return hash(self._id)

    def clear_visited(self):
        self.visited = False

    def visit(self):
        self.visited = True

    def get_epsilon_edges(self):
        epsilon_edges = []
        for letters, next_state in self.edges:
            if len(letters) == 1 and tuple(letters)[0] == -1:
                epsilon_edges += [(letters, next_state)]

        return epsilon_edges

    def get_edges_to(self, state):
        edges_to = []
        for letters, next_state in self.edges:
            if next_state is state:
                edges_to += [(letters, next_state)]

        return edges_to

    def remove_edges_to(self, state):
        self.edges = [edge for edge in self.edges if not edge[1] is state]

    def add_edges(self, edges):
        for edge, new_state in edges:
            if type(edge) is set:
                self.edges += [(edge, new_state)]
            else:
                self.edges += [({edge}, new_state)]

    # only call on deterministic nodes
    def transition(self, letter):
        for edge_letter_set, state in self.edges:
            if letter.id() in edge_letter_set:
                return state

        return None

    def valid_transitions(self):
        letters = set()
        for transition_letters, next_state in self.edges:
            letters = letters | transition_letters

        return letters

    def nd_transition(self, letter):
        new_states = set()
        if not self.visited:
            self.visit()
            for edge_letter_set, state in self.edges:
                if letter in edge_letter_set:
                    new_states = new_states | {state}

                if -1 in edge_letter_set:
                    new_states = new_states | state.nd_transition(letter)

        return new_states

    def __eq__(self, other):
        return self._id == other._id

class DFSM(FSM):
    def from_regex(self, ast):
        nfsm = NFSM(self.letter_cls)
        nfsm.from_regex(ast)
        self.from_nfsm(nfsm)

    def from_nfsm(self, nfsm):
        self.dead_state = FSMSuperState({nfsm.dead_state}, [])
        self.add_state(self.dead_state)
        self.initial_state = FSMSuperState({nfsm.initial_state}, [])
        self.add_state(self.initial_state)
        self._from_nfsm(nfsm, self.initial_state)
        self.clear_visited()

    def get_by_ref_states(self, ref_states):
        for state in self.states:
            if ref_states == state.component_states:
                return state
        else:
            return None

    def make_state(self, component_states):
        ref_state = self.get_by_ref_states(component_states)
        if ref_state is None:
            ref_state = FSMSuperState(component_states, [])
            self.add_state(ref_state)

        return ref_state

    def _from_nfsm(self, nfsm, state):
        if not state.visited:
            state.visit()
            paths = []

            for letter in self.letter_cls.all():
                this_letter_states = state.nd_transition(letter.id(), nfsm)
                next_state = self.make_state(this_letter_states)

                paths += [(letter.id(), next_state)]

            state.add_edges(paths)
            for _, next_state in state.edges:
                self._from_nfsm(nfsm, next_state)

class FSMSuperState(FSMState):
    # component_states expected to be a set
    def __init__(self, component_states, edges, id=None, terminal=False):
        super(FSMSuperState, self).__init__(edges, id=id, terminal=terminal)
        self.component_states = component_states
        self.set_terminal()

    def set_terminal(self):
        for component_state in self.component_states:
            self.terminal |= component_state.terminal

    def __eq__(self, other):
        if type(other) is FSMState:
            if len(self.component_states) == 1:
                return self.component_states[0] == other

        if type(other) is not FSMSuperState:
            return False

        same = True
        for component_state in self.component_states:
            same = same and (component_state in other.component_states)

    def nd_transition(self, letter_id, nfsm):
        new_states = set()
        for component_state in self.component_states:
            new_states = new_states | component_state.nd_transition(letter_id)
            nfsm.clear_visited()

        return new_states


# -*- coding: utf-8 -*-
"""
Created on Fri Sep 14 14:54:03 2018

@author: Administrator
"""
from selenium import webdriver
#操作键盘鼠标
from selenium.webdriver.common.keys import Keys
import time

driver = webdriver.PhantomJS()
driver.get("http://www.baidu.com/")

driver.find_element_by_id("kw").send_keys(u"美女")
driver.save_screenshot("image/美女.png")
print("ok")
driver.find_element_by_id("su").click()
time.sleep(10)
driver.save_screenshot("image/nv.png")
print("ok1")


import sqlalchemy
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session
from sqlalchemy import create_engine, func, inspect

import numpy as np
import pandas as pd
import datetime as dt

from flask import Flask, jsonify

#################################################
# Database Setup
#################################################
engine = create_engine("sqlite:///Resources/hawaii.sqlite")

# reflect an existing database into a new model
Base = automap_base()
# reflect the tables
Base.prepare(engine, reflect=True)

# Save reference to the table
Measurement = Base.classes.measurement
Station = Base.classes.station

# Create our session (link) from Python to the DB
session = Session(engine)

#################################################
# Flask Setup
#################################################
app = Flask(__name__)

#################################################
# Flask Routes
#################################################

@app.route("/")
def index():
    """List all available api routes."""
    return (
        f"Available Routes:
"
        f"/api/v1.0/precipitation
"
        f" - list of rain totals from all stations
"
        f"/api/v1.0/stations
"
        f"- list of Station numbers and names
"
        f"/api/v1.0/tobs
"
        f"- list of prior year temperatures from all stations
"
        f"/api/v1.0/STARTTIME
"
        f"- When given the start date (yyyy-mm-dd), calculates the min/ave/max temp for all dates greate than and equal to teh start date.
"
        f"/api/v1.0/STARTTIME/ENDTIME
"
        f"- Please provide the start date (yyyy-mm-dd) and end date (yyyy-mm-dd)"
        )

@app.route("/api/v1.0/precipitation")
def percipitation():
    """Return percipitation"""
    # Query all date and prcp from Measurement
    results = session.query(Measurement).all()
    # Create a dictionary from the row data and append to a list of all_stations
    all_stations = []

    for result in results:
        stations_dict = {}
        stations_dict["date"] = result.date
        stations_dict["prcp"] = result.prcp
        all_stations.append(stations_dict)

    return jsonify(all_stations)

@app.route("/api/v1.0/stations")
def stations():
    """Return Station Names"""
    results = session.query(Station.station, Station.name,).all()
    # Create a dictionary from the row data and append to a list of all_stations
    all_stations=[]
    for result in results:
        stations_dict = {}
        stations_dict["station"] = result[0]
        stations_dict["name"] = result[1]
        all_stations.append(stations_dict)

    return jsonify(all_stations)

@app.route("/api/v1.0/tobs")
def tobs():
    #query for the dates and temps froma year fromt the last data point.
    query_date = dt.date(2017,8,23) - dt.timedelta(weeks=52)
    #create all data from 2016/8/23-2017/8/23
    lasttwelvemonths = (session.query(Measurement)
                   .filter(func.strftime(Measurement.date)>=query_date)
                   .order_by(Measurement.date)
                   .all())
    one_year_tobs=[]

    for last in lasttwelvemonths:
        tobs_dict = {}
        tobs_dict["date"] = last.date
        tobs_dict["tobs"] = last.tobs
        one_year_tobs.append(tobs_dict)

    return jsonify(one_year_tobs)

@app.route("/api/v1.0/<start>")
def trip1(start):

    trip_data = session.query(func.min(Measurement.tobs), func.avg(Measurement.tobs),
    func.max(Measurement.tobs)).filter(Measurement.date>=start).all()

    trip = list(np.ravel(trip_data))

    return jsonify(trip)

@app.route("/api/v1.0/<start>/<end>")
def trip2(start,end):

    trip_data = session.query(func.min(Measurement.tobs), func.avg(Measurement.tobs),
    func.max(Measurement.tobs)).filter(Measurement.date>=start).filter(Measurement.date<=end).all()

    trip = list(np.ravel(trip_data))

    return jsonify(trip)

if __name__ == '__main__':
    app.run(debug=True)


from random import randint

class Color():
    def __init__(self, r, g, b):
        r = min(255, r)
        g = min(255, g)
        b = min(255, b)
        self.r = r
        self.g = g
        self.b = b

    def __str__(self):
        r = hex(min(255, self.r))[2:]
        g = hex(min(255, self.g))[2:]
        b = hex(min(255, self.b))[2:]
        return '#{:0>2}{:0>2}{:0>2}'.format(r, g, b)

    def __add__(self, other):
        r = self.r + other.r
        g = self.g + other.g
        b = self.b + other.b
        return Color(r, g, b)

    def __sub__(self, other):
        return self + (other * -1)

    def __mul__(self, coef):
        r = int(self.r * coef)
        g = int(self.g * coef)
        b = int(self.b * coef)
        return Color(r, g, b)

    def __eq__(self, other):
        return self.r == other.r and self.g == other.g and self.b == other.b

    def to_rgb(self):
        return [self.r, self.g, self.b]

    def mean(self, other):
        return mean_color(self, other)

    def cp(self):
        return Color(self.r, self.g, self.b)

def color_from_rgb(rgb):
    return Color(rgb[0], rgb[1], rgb[2])

def random_color():
    r = randint(0, 255)
    g = randint(0, 255)
    b = randint(0, 255)
    return Color(r, g, b)

def color_from_limit(first, second):
    rgb_first = first.to_rgb()
    rgb_second = second.to_rgb()

    rgb_output = []
    for i in range(3):
        lim_1 = rgb_first[i]
        lim_2 = rgb_second[i]
        if lim_1 > lim_2:
            lim_1, lim_2 = lim_2, lim_1
        rgb_output.append(randint(lim_1, lim_2))

    return color_from_rgb(rgb_output)

def mean_color(first, second):
    rgb_first = first.to_rgb()
    rgb_second = second.to_rgb()

    rgb_mean = []
    for i in range(3):
        rgb_mean.append((rgb_first[i] + rgb_second[i]) // 2)

    return color_from_rgb(rgb_mean)

def modify_color(color, delta, coef):
    rgb  = color.to_rgb()
    for i in range(2):
        rgb[i] = int(rgb[i] * coef) + delta
    return color_from_rgb(rgb)

def colors_delta(first, second):
    rgb_first = first.to_rgb()
    rgb_second = second.to_rgb()
    delta = 0
    for i in range(3):
        delta += abs(rgb_first[i] - rgb_second[i])
    return delta

def colors_are_equal(first, second):
    delta = colors_delta(first, second)
    if delta <= 6:
        return True
    else:
        return False

def randstr(length, lvl=3):
    s = ""
    for i in range(length):
        roll = randint(1, lvl)
        if roll == 1:
            c = chr(randint(48, 57))
        elif roll == 2:
            c = chr(randint(65, 90))
        elif roll == 3:
            c = chr(randint(97, 122))
        elif roll == 4:
            c = choice(SYMBS)
        s = s + c
    return s

def randstrchoice(length, chars):
    s = []
    for i in range(length):
        s.append(choice(chars))
    return ''.join(s)