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

'''
thin_vcf

Remove unneeded rows from vcf. Equivalent function to merge bedfile and
vcftools recode with keep individuals
'''

import argparse
import numpy as np
import sys
from typing import List, TextIO

def main():
    args = read_args()

    # read in and merge bedfile
    if args.bed_file is not None:
        with open(args.bed_file) as bed_file:
            bed = read_bed(bed_file, args.merge)
    else:
        bed = None

    if args.individuals_file is not None:
        with open(args.individuals_file) as indiv_file:
            indivs = read_indivs(indiv_file)
    else:
        indivs = None

    parser = line_parser(bed, indivs)

    for line in sys.stdin:
        print(parser.process_line(line), end='')

def read_args(args: List[str] = None) -> argparse.Namespace:
    '''
    read in command line arguments, returning namespace object
    '''
    parser = argparse.ArgumentParser(description='Filter vcfs from stdin')
    parser.add_argument('--bed_file',
                        default=None,
                        help='If set, retain only the sites within bed '
                        'file regions.  Assumes indexing is identical and '
                        'inclusive with start and end. Must be sorted.')

    parser.add_argument('--merge',
                        default=None,
                        type=int,
                        help='If set to an integer, merge bed regions within '
                        '<merge> away from eachother, inclusive.')

    parser.add_argument('--individuals_file',
                        default=None,
                        help='If set, only individuals from the file are kept.'
                        ' One individual per line of input file.')

    args = parser.parse_args(args)
    return args

def read_bed(reader: TextIO, merge: int = None) -> np.array:
    '''
    read in the opened file as a numpy array.
    If merge is provided, adjacent sites are merged if
    start_i - end_i-1 <= merge.
    Return only start and end sites, chromosome is assumed constant and equal
    '''
    result = np.loadtxt(reader,
                        dtype=int,
                        delimiter='\t',
                        usecols=(1, 2))
    if merge is not None:
        # find locations to keep (diff between end and start < merge)
        tokeep = np.concatenate(
            ([True],
             np.diff(result.flatten())[1::2] > merge))

        swap_inds = np.where(tokeep[:-1] != tokeep[1:])[0]
        # nothing to merge/swap
        if swap_inds.size > 0:
            # this happens when last ind needs to swap
            if swap_inds.size % 2 == 1:
                swap_inds = np.append(swap_inds, result.shape[0] - 1)
            result[swap_inds[0::2], 1] = result[swap_inds[1::2], 1]
            result = result[tokeep, :]
    return result

def read_indivs(reader: TextIO) -> List[str]:
    '''
    read in each line of the file, return as list with newline stripped
    '''
    return [line.rstrip() for line in reader]

def match_indivs(header: List[str], indivs: List[str]) -> List[int]:
    header = np.array(header)
    indivs = np.array(indivs)

    matches = np.char.equal(header[:, None], indivs[None, :])
    return np.where(matches)[0]

class line_parser():
    def __init__(self, bed: np.array, indivs: List[str]):
        self.bed = bed
        self.indivs = indivs
        if self.bed is None:
            self.parser = self.no_op
        else:
            self.bed += 1
            self.parser = self.filter_bed
        self.retain = 9
        self.bed_ind = 0

    def process_line(self, line: str) -> str:
        if line[1] == '#':
            return line

        if line[0] == '#':
            if self.indivs is not None:
                header = line.split('\t')
                self.indivs = match_indivs(header, self.indivs)
                self.indivs = np.concatenate((np.array(range(self.retain)),
                                             self.indivs))
                if self.bed is None:
                    self.parser = self.filter_indiv
                else:
                    self.parser = self.filter_both

                return self.filter_indiv(line)
            else:
                return line

        return self.parser(line)

    def no_op(self, line: str):
        return line

    def filter_indiv(self, line: str):
        line = line.split('\t')
        return '\t'.join([line[i] for i in self.indivs]).rstrip() + '\n'

    def filter_bed(self, line: str):
        pos = int(line.split('\t')[1])
        if self.bed_ind >= self.bed.shape[0]:
            return ''
        while pos >= self.bed[self.bed_ind, 1]:
            self.bed_ind += 1
            if self.bed_ind >= self.bed.shape[0]:
                return ''
        # by the time we are here, pos < end
        if pos >= self.bed[self.bed_ind, 0]:
            return line
        else:
            return ''

    def filter_both(self, line: str):
        tokens = line.split('\t')
        pos = int(tokens[1])
        if self.bed_ind >= self.bed.shape[0]:
            return ''
        while pos >= self.bed[self.bed_ind, 1]:
            self.bed_ind += 1
            if self.bed_ind >= self.bed.shape[0]:
                return ''
        # by the time we are here, pos < end
        if pos >= self.bed[self.bed_ind, 0]:
            return '\t'.join([tokens[i] for i in self.indivs]).rstrip() + '\n'
        else:
            return ''

if __name__ == "__main__":
    main()


from keras.models import load_model
import numpy as np
import pickle
import discord
from dotenv import load_dotenv
import os
from keras_preprocessing.sequence import pad_sequences
load_dotenv()

client = discord.Client()

rating_index = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
max_input_length = 1403
model = load_model("model/model.h5")
with open("model/tokenizer.pkl", "rb") as t:
    tokenizer = pickle.load(t)

def process_message_for_data(sentence):
    tokens = []
    for word in sentence.lower().split():
        try:
            tokens.append(tokenizer.word_index[word])
        except KeyError:
            tokens.append(0)
    transformed_sentence = pad_sequences([tokens], maxlen=max_input_length, padding="post")
    prediction = model.predict(transformed_sentence)[0]
    return np.argmax(prediction), np.amax(prediction)

@client.event
async def on_ready():
    print(f"We have logged in as {client.user}")

@client.event
async def on_message(message):
    if message.author == client.user:
        return

    index, val = process_message_for_data(message.content)
    rating = rating_index[index]
    if val > 0.5:
        try:
            await message.delete()
        except:
            print("ERROR DELETING!")
            return
        await message.channel.send(f"Hey! Toxic Waste Handler here! {message.author} said something really toxic, so I deleted it!")
    return

TOKEN = os.getenv("TOKEN")
client.run(TOKEN)


from __future__ import print_function
import sys
import os
import os.path
import csv
import re
import random
import datetime
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

import msmarco_eval

DEVICE                          = torch.device("cuda:0")    # torch.device("cpu"), if you want to run on CPU instead
ARCH_TYPE                       = 2
MAX_QUERY_TERMS                 = 20
MAX_DOC_TERMS                   = 200
NUM_HIDDEN_NODES                = 512
TERM_WINDOW_SIZE                = 3
POOLING_KERNEL_WIDTH_QUERY      = MAX_QUERY_TERMS - TERM_WINDOW_SIZE + 1 # 20 - 3 + 1 = 18
POOLING_KERNEL_WIDTH_DOC        = 100
NUM_POOLING_WINDOWS_DOC         = (MAX_DOC_TERMS - TERM_WINDOW_SIZE + 1) - POOLING_KERNEL_WIDTH_DOC + 1 # (200 - 3 + 1) - 100 + 1 = 99
NUM_NGRAPHS                     = 0
DROPOUT_RATE                    = 0.5
MB_SIZE                         = 256
EPOCH_SIZE                      = 1000
NUM_EPOCHS                      = 100
LEARNING_RATE                   = 1e-5
DATA_DIR                        = 'data'
DATA_FILE_NGRAPHS               = os.path.join(DATA_DIR, "ngraphs.txt")
DATA_FILE_IDFS                  = os.path.join(DATA_DIR, "idf.norm.tsv")
DATA_FILE_TRAIN                 = os.path.join(DATA_DIR, "triples.train.full.tsv")
DATA_FILE_PREDICT                = os.path.join(DATA_DIR, "top1000.dev.tiny.tsv")
#DATA_FILE_PREDICT                = os.path.join(DATA_DIR, "top1000.dev.tsv")
DATA_FILE_PREDICTOUT              = os.path.join(DATA_DIR, "ranked.dev.tsv")
QRELS                       = os.path.join(DATA_DIR, "qrels.dev.tsv")

#DATA_FILE_PREDICTOUT               = os.path.join(DATA_DIR, "ranked.eval.tsv")
#QRELS                      = os.path.join(DATA_DIR, "qrels.eval.tsv")
#DATA_FILE_PREDICT                   = os.path.join(DATA_DIR, "top1000.eval.tsv")

class DataReader:
    def __init__(self, data_file, num_meta_cols, multi_pass, mb_size):
        self.num_meta_cols                  = num_meta_cols
        self.multi_pass                     = multi_pass
        self.regex_drop_char                = re.compile('[^a-z0-9\s]+')
        self.regex_multi_space              = re.compile('\s+')
        self.mb_size                        = mb_size
        self.__load_ngraphs()
        self.__load_idfs()
        self.__init_data(data_file)
        self.__allocate_minibatch()
    def __tokenize(self, s, max_terms):
        return self.regex_multi_space.sub(' ', self.regex_drop_char.sub(' ', s.lower())).strip().split()[:max_terms]
    def __load_ngraphs(self):
        global NUM_NGRAPHS
        self.ngraphs                        = {}
        self.max_ngraph_len                 = 0
        with open(DATA_FILE_NGRAPHS, mode='r', encoding="utf-8") as f:
            reader                          = csv.reader(f, delimiter='\t')
            for row in reader:
                self.ngraphs[row[0]]        = int(row[1]) - 1
                self.max_ngraph_len         = max(self.max_ngraph_len, len(row[0]))
        NUM_NGRAPHS                         = len(self.ngraphs)
    def __load_idfs(self):
        self.idfs                           = {}
        with open(DATA_FILE_IDFS, mode='r', encoding="utf-8") as f:
            reader                          = csv.reader(f, delimiter='\t')
            for row in reader:
                self.idfs[row[0]]           = float(row[1])
    def __init_data(self, file_name):
        self.reader                         = open(file_name, mode='r', encoding="utf-8")
        self.num_docs                       = len(self.reader.readline().split('\t')) - self.num_meta_cols - 1
        self.reader.seek(0)
    def __allocate_minibatch(self):
        self.features                       = {}
        if ARCH_TYPE != 1:
            self.features['local']          = []
        if ARCH_TYPE > 0:
            self.features['dist_q']         = np.zeros((self.mb_size, NUM_NGRAPHS, MAX_QUERY_TERMS), dtype=np.float32)
            self.features['dist_d']         = []
        for i in range(self.num_docs):
            if ARCH_TYPE != 1:
                self.features['local'].append(np.zeros((self.mb_size, MAX_DOC_TERMS, MAX_QUERY_TERMS), dtype=np.float32))
            if ARCH_TYPE > 0:
                self.features['dist_d'].append(np.zeros((self.mb_size, NUM_NGRAPHS, MAX_DOC_TERMS), dtype=np.float32))
        self.features['labels']             = np.zeros((self.mb_size), dtype=np.int64)
        self.features['meta']               = []
    def __clear_minibatch(self):
        if ARCH_TYPE > 0:
            self.features['dist_q'].fill(np.float32(0))
        for i in range(self.num_docs):
            if ARCH_TYPE != 1:
                self.features['local'][i].fill(np.float32(0))
            if ARCH_TYPE > 0:
                self.features['dist_d'][i].fill(np.float32(0))
        self.features['meta'].clear()
    def get_minibatch(self):
        self.__clear_minibatch()
        for i in range(self.mb_size):
            row                             = self.reader.readline()
            if row == '':
                if self.multi_pass:
                    self.reader.seek(0)
                    row                     = self.dataFile.readline()
                else:
                    break
            cols                            = row.split('\t')
            q                               = self.__tokenize(cols[self.num_meta_cols], MAX_QUERY_TERMS)
            ds                              = [self.__tokenize(cols[self.num_meta_cols + i + 1], MAX_DOC_TERMS) for i in range(self.num_docs)]
            if ARCH_TYPE != 1:
                for d in range(self.num_docs):
                    for j in range(len(ds[d])):
                        for k in range(len(q)):
                            if ds[d][j] == q[k]:
                                if q[k] in self.idfs:
                                    self.features['local'][d][i, j, k] = self.idfs[q[k]]
                                else:
                                    self.features['local'][d][i, j, k] = 0
            if ARCH_TYPE > 0:
                for j in range(self.num_docs + 1):
                    terms = q if j == 0 else ds[j - 1]
                    for t in range(len(terms)):
                        term = '#' + terms[t] + '#'
                        term_len = len(term)
                        for k in range(term_len):
                            for l in range(1, self.max_ngraph_len + 1):
                                if k + l < term_len:
                                    ngraph_idx = self.ngraphs.get(term[k : k + l])
                                    if ngraph_idx != None:
                                        if j == 0:
                                            self.features['dist_q'][i, ngraph_idx, t] += 1
                                        else:
                                            self.features['dist_d'][j - 1][i, ngraph_idx, t] += 1
            self.features['meta'].append(tuple(cols[:self.num_meta_cols]))
        return self.features
    def get_minibatch2(self):
        self.__clear_minibatch()
        prev_d = []
        first_q = []
        for i in range(self.mb_size):
            row                             = self.reader.readline()
            if row == '':
                if self.multi_pass:
                    self.reader.seek(0)
                    row                     = self.dataFile.readline()
                else:
                    break
            cols                            = row.split('\t')
            q                               = self.__tokenize(cols[self.num_meta_cols], MAX_QUERY_TERMS)
            d                               = self.__tokenize(cols[self.num_meta_cols + 1], MAX_DOC_TERMS)
            if i == 0:
                first_q = q
            if ARCH_TYPE != 1:
                for j in range(len(d)):
                    for k in range(len(q)):
                        if d[j] == q[k]:
                            if q[k] in self.idfs:
                                self.features['local'][0][i, j, k] = self.idfs[q[k]]
                            else:
                                self.features['local'][0][i, j, k] = 0
                for j in range(len(prev_d)):
                    for k in range(len(q)):
                        if prev_d[j] == q[k]:
                            if q[k] in self.idfs:
                                self.features['local'][1][i, j, k] = self.idfs[q[k]]
                            else:
                                self.features['local'][1][i, j, k] = 0
            if ARCH_TYPE > 0:
                for j in range(self.num_docs + 1):
                    terms = q if j == 0 else (d if j == 1 else prev_d)
                    for t in range(len(terms)):
                        term = '#' + terms[t] + '#'
                        term_len = len(term)
                        for k in range(term_len):
                            for l in range(1, self.max_ngraph_len + 1):
                                if k + l < term_len:
                                    ngraph_idx = self.ngraphs.get(term[k : k + l])
                                    if ngraph_idx != None:
                                        if j == 0:
                                            self.features['dist_q'][i, ngraph_idx, t] += 1
                                        else:
                                            self.features['dist_d'][j - 1][i, ngraph_idx, t] += 1
            prev_d = d
            self.features['meta'].append(tuple(cols[:self.num_meta_cols]))
        if ARCH_TYPE != 1:
            for j in range(len(prev_d)):
                for k in range(len(first_q)):
                    if prev_d[j] == first_q[k]:
                        if q[k] in self.idfs:
                            self.features['local'][1][0, j, k] = self.idfs[q[k]]
                        else:
                            self.features['local'][1][0, j, k] = 0
        if ARCH_TYPE > 0:
            terms = prev_d
            for t in range(len(terms)):
                term = '#' + terms[t] + '#'
                term_len = len(term)
                for k in range(term_len):
                    for l in range(1, self.max_ngraph_len + 1):
                        if k + l < term_len:
                            ngraph_idx = self.ngraphs.get(term[k : k + l])
                            if ngraph_idx != None:
                                    self.features['dist_d'][1][0, ngraph_idx, t] += 1
        return self.features
    def reset(self):
        self.reader.seek(0)

class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()
    def forward(self, x):
        return x.view(x.size(0), -1)

class Duet(torch.nn.Module):
    def __init__(self):
        super(Duet, self).__init__()
        self.duet_local             = nn.Sequential(nn.Conv1d(MAX_DOC_TERMS, NUM_HIDDEN_NODES, kernel_size=1),
                                        nn.Tanh(),
                                        Flatten(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES * MAX_QUERY_TERMS, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE))
        self.duet_dist_q            = nn.Sequential(nn.Conv1d(NUM_NGRAPHS, NUM_HIDDEN_NODES, kernel_size=3),
                                        nn.Tanh(),
                                        nn.MaxPool1d(POOLING_KERNEL_WIDTH_QUERY),
                                        Flatten(),
                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),
                                        nn.Tanh())
        self.duet_dist_d            = nn.Sequential(nn.Conv1d(NUM_NGRAPHS, NUM_HIDDEN_NODES, kernel_size=3),
                                        nn.Tanh(),
                                        nn.MaxPool1d(POOLING_KERNEL_WIDTH_DOC, stride=1),
                                        nn.Conv1d(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES, kernel_size=1),
                                        nn.Tanh())
        self.duet_dist              = nn.Sequential(Flatten(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES * NUM_POOLING_WINDOWS_DOC, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE))
        self.duet_comb              = nn.Sequential(nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),
                                        nn.Tanh(),
                                        nn.Dropout(p=DROPOUT_RATE),
                                        nn.Linear(NUM_HIDDEN_NODES, 1),
                                        nn.Tanh())
    def forward(self, x_local, x_dist_q, x_dist_d):
        if ARCH_TYPE != 1:
            h_local                 = self.duet_local(x_local)
        if ARCH_TYPE > 0:
            h_dist_q                = self.duet_dist_q(x_dist_q)
            h_dist_d                = self.duet_dist_d(x_dist_d)
            h_dist                  = self.duet_dist(h_dist_q.unsqueeze(-1)*h_dist_d)
        y_score                     = self.duet_comb((h_local + h_dist) if ARCH_TYPE == 2 else (h_dist if ARCH_TYPE == 1 else h_local))
        return y_score
    def parameter_count(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

def print_message(s):
    print("[{}] {}".format(datetime.datetime.utcnow().strftime("%b %d, %H:%M:%S"), s), flush=True)

def train(READER_TRAIN, net, optimizer, criterion):
    net.train()
    train_loss          = 0.0
    for mb_idx in range(EPOCH_SIZE):
        features        = READER_TRAIN.get_minibatch()
        acc = []
        x = torch.from_numpy(features['labels']).to(DEVICE)
        x2 = Variable(torch.from_numpy(features['dist_q'])).to(DEVICE)
        for i in range(READER_TRAIN.num_docs):
            x1 = Variable(torch.from_numpy(features['local'][i])).to(DEVICE)
            x3 = Variable(torch.from_numpy(features['dist_d'][i])).to(DEVICE)
            y = net(x1,x2,x3).to(DEVICE)
            acc.append(y)
        out = torch.cat(tuple(acc),1)
        loss = criterion(out, x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss     += loss.item()
        del y, x, x1,x2,x3,acc, out,loss, features
        torch.cuda.empty_cache()
    return net, train_loss

def evaluate(net, qrels, scores, to_rerank):
    MB_SIZE = 2560
    READER_PREDICT = DataReader(DATA_FILE_PREDICT, 2, False, MB_SIZE)
    for docs in scores.values():
        docs.clear()
    is_complete         = False
    READER_PREDICT.reset()
    net.eval()
    while not is_complete:
        features        = READER_PREDICT.get_minibatch()
        meta_cnt        = len(features['meta'])
        is_complete     = (meta_cnt < MB_SIZE)
        x1 = Variable(torch.from_numpy(features['local'][0])).to(DEVICE)
        x2 = Variable(torch.from_numpy(features['dist_q'])).to(DEVICE)
        x3 = Variable(torch.from_numpy(features['dist_d'][0])).to(DEVICE)
        out         = net(x1,x2,x3).to(DEVICE).data.cpu().numpy()
        for i in range(meta_cnt):
            q           = int(features['meta'][i][0])
            d           = int(features['meta'][i][1])
            if q in scores:
                scores[q][d]= out[i][0]
        del x1,x2,x3, out, q, d ,features
        torch.cuda.empty_cache()
    mrr                 = 0
    for qid, docs in scores.items():
        ranked          = sorted(docs, key=docs.get, reverse=True)
        if len(ranked) > 0:
            for i in range(len(ranked)):
                if ranked[i] in qrels[qid]:
                    mrr    += 1 / (i + 1)
                    break
    mrr                /= len(to_rerank)
    with open(DATA_FILE_PREDICTOUT,'w') as w:
        for qid in to_rerank:
            ranked = sorted(scores[qid], reverse=False)
            for i in range(0, min(len(ranked),10)):
                w.write('{}\t{}\t{}\n'.format(qid,ranked[i],i+1))
    return mrr

def load_file_to_rerank(filename, offset):
    target,scores = {}, {}
    with open(filename,'r', encoding="utf-8") as f:
        reader = csv.reader(f, delimiter='\t')
        for row in reader:
            qid = int(row[0])
            pid = int(row[1+offset])
            if qid not in target:
                target[qid]          = []
            target[qid].append(pid)
            scores[qid] = {}
    return target,scores

def main():
    READER_TRAIN  = DataReader(DATA_FILE_TRAIN, 0, True, 256)
    qrels, scores = load_file_to_rerank(QRELS, 1)
    _, to_rerank = load_file_to_rerank(DATA_FILE_PREDICT,0)
    torch.manual_seed(1)
    print_message('Starting')
    try:
        net = torch.load('tensors.duet').to(DEVICE)
        print_message("Previous Model Found and Loaded")
    except:
        print_message("No Previous Model Found. Creating new")
        net = Duet().to(DEVICE)
    criterion               = nn.CrossEntropyLoss()
    optimizer               = optim.Adam(net.parameters(), lr=LEARNING_RATE)
    print_message('Number of learnable parameters: {}'.format(net.parameter_count()))
    print_message('Training for {} epochs'.format(NUM_EPOCHS))
    for ep_idx in range(NUM_EPOCHS):
        net, train_loss  = train(READER_TRAIN,net, optimizer,criterion)
        print_message('epoch:{}, loss: {}'.format(ep_idx + 1, train_loss / EPOCH_SIZE))
        mrr = evaluate(net,qrels,scores, to_rerank)
        print_message('MRR @1000:{}'.format(mrr))
        torch.cuda.empty_cache()
        torch.save(net,'tensors.duet')
    print_message('Done Training')
    print_message('Evaluating')
    metrics = msmarco_eval.compute_metrics_from_files(QRELS,DATA_FILE_PREDICTOUT)
    print_message('#####################')
    for metric in sorted(metrics):
        print_message('{}: {}'.format(metric, metrics[metric]))
    print_message('#####################')

if __name__ == '__main__':
    main()


import numpy as np
import torch
from torch.distributions import Distribution

def gaussian_processes(data=10):
    return NotImplementedError


import argparse
import os
import random
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from dataset.eval_dataset_nocs import Dataset
from libs.network import KeyNet
from libs.loss import Loss
import copy

choose_cate_list = [1,2,3,4,5,6]
resume_models = ['model_112_0.184814792342484_bottle.pth',
                 'model_120_0.10268432162888348_bowl.pth',
                 'model_118_0.2008235973417759_camera.pth',
                 'model_107_0.18291547849029302_can.pth',
                 'model_117_0.12762234719470145_laptop.pth',
                 'model_102_0.1468337191492319_mug.pth']

parser = argparse.ArgumentParser()
parser.add_argument('--dataset_root', type=str, default = 'My_NOCS', help='dataset root dir')
parser.add_argument('--eval_id', type=int, default = 1, help='the evaluation id')
parser.add_argument('--ite', type=int, default=10, help='first frame fix iteration')
parser.add_argument('--num_kp', type=int, default = 8, help='num of kp')
parser.add_argument('--num_points', type=int, default = 500, help='num of input points')
parser.add_argument('--num_cates', type=int, default = 6, help='number of categories')
parser.add_argument('--outf', type=str, default = 'models/', help='load model dir')
opt = parser.parse_args()

if not os.path.exists('eval_results'):
    os.makedirs('eval_results')

if not os.path.exists('eval_results/TEST_{0}'.format(opt.eval_id)):
    os.makedirs('eval_results/TEST_{0}'.format(opt.eval_id))
    for item in choose_cate_list:
        os.makedirs('eval_results/TEST_{0}/temp_{1}'.format(opt.eval_id, item))

for choose_cate in choose_cate_list:
    model = KeyNet(num_points = opt.num_points, num_key = opt.num_kp, num_cates = opt.num_cates)
    model.cuda()
    model.eval()

    model.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, resume_models[choose_cate-1])))  # 导入models

    pconf = torch.ones(opt.num_kp) / opt.num_kp
    pconf = Variable(pconf).cuda()

    test_dataset = Dataset('val', opt.dataset_root, False, opt.num_points, choose_cate, 1000)  # 类似于打开dataset？
    criterion = Loss(opt.num_kp, opt.num_cates)

    eval_list_file = open('dataset/eval_list/eval_list_{0}.txt'.format(choose_cate), 'r')
    while 1:
        input_line = eval_list_file.readline()
        if not input_line:
            break
        if input_line[-1:] == '\n': # [-1:] 表示列表中最后一个元素构成的切片
            input_line = input_line[:-1] # [:-1]表示除了最后一个元素构成的切片（从头开始到最后一个元素之前）
        _, choose_obj, choose_video = input_line.split(' ')

        try:
            current_r, current_t = test_dataset.getfirst(choose_obj, choose_video)
            rad_t = np.array([random.uniform(-0.02, 0.02) for i in range(3)]) * 1000.0
            current_t += rad_t

            if opt.ite != 0:
                min_dis = 1000.0
                for iterative in range(opt.ite):
                    img_fr, choose_fr, cloud_fr, anchor, scale = test_dataset.getone(current_r, current_t)
                    img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
                                                         Variable(choose_fr).cuda(), \
                                                         Variable(cloud_fr).cuda(), \
                                                         Variable(anchor).cuda(), \
                                                         Variable(scale).cuda()
                    Kp_fr, att_fr = model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, 0.0, True)
                    new_t, att, kp_dis = criterion.ev_zero(Kp_fr[0], att_fr[0])

                    if min_dis > kp_dis:
                        min_dis = kp_dis
                        best_current_r = copy.deepcopy(current_r)
                        best_current_t = copy.deepcopy(current_t)
                        best_att = copy.deepcopy(att)
                        print(min_dis)

                    current_t = current_t + np.dot(new_t, current_r.T)
                current_r, current_t, att = best_current_r, best_current_t, best_att

            img_fr, choose_fr, cloud_fr, anchor, scale = test_dataset.getone(current_r, current_t)
            img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
                                                 Variable(choose_fr).cuda(), \
                                                 Variable(cloud_fr).cuda(), \
                                                 Variable(anchor).cuda(), \
                                                 Variable(scale).cuda()
            Kp_fr, att_fr = model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, 0.0, True)

            test_dataset.projection('eval_results/TEST_{0}/temp_{1}/{2}_{3}'.format(opt.eval_id, choose_cate, choose_obj, choose_video), Kp_fr[0], current_r, current_t, scale, att_fr[0], True, 0.0)

            min_dis = 0.0005
            while 1:
                img_fr, choose_fr, cloud_fr, anchor, scale = test_dataset.getone(current_r, current_t)
                img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
                                                     Variable(choose_fr).cuda(), \
                                                     Variable(cloud_fr).cuda(), \
                                                     Variable(anchor).cuda(), \
                                                     Variable(scale).cuda()
                Kp_to, att_to = model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, min_dis, False) # Kp_to 应该是点集

                min_dis = 1000.0
                lenggth = len(Kp_to) # 一共有多少个keypoint?
                for idx in range(lenggth):
                    Kp_real, new_r, new_t, kp_dis, att = criterion.ev(Kp_fr[0], Kp_to[idx], att_to[idx])

                    if min_dis > kp_dis:
                        best_kp = Kp_to[idx]
                        min_dis = kp_dis
                        best_r = new_r
                        best_t = new_t
                        best_att = copy.deepcopy(att)
                print(min_dis)

                current_t = current_t + np.dot(best_t, current_r.T)
                current_r = np.dot(current_r, best_r)

                test_dataset.projection('eval_results/TEST_{0}/temp_{1}/{2}_{3}'.format(opt.eval_id, choose_cate, choose_obj, choose_video), Kp_real, current_r, current_t, scale, best_att, True, min_dis)

                print("NEXT FRAME!!!")

        except:
            continue
    print("FINISH!!!")


import argparse
import datetime
import math
import os
import random
import sys
import time
from pathlib import Path
from tqdm import tqdm

# Ranger Optimizer
from ranger import Ranger  # this is from ranger.py
from ranger import RangerVA  # this is from ranger913A.py
from ranger import RangerQH  # this is from rangerqh.py

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.init as init
import torch.optim as optim
import torch.utils.data as data
from torch.autograd import Variable
from torch.nn.functional import interpolate
from torch.optim.lr_scheduler import ReduceLROnPlateau

# Oof
import eval as eval_script
from data import *
from layers.modules import MultiBoxLoss
from utils import timer
from utils.augmentations import BaseTransform, SSDAugmentation
from utils.functions import MovingAverage, SavePath
from utils.logger import Log
from yolact import Discriminator, Discriminator_StandFord, Discriminator_Dcgan, Yolact

def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")

parser = argparse.ArgumentParser(
    description='Yolact Training Script')
parser.add_argument('--batch_size', default=2, type=int,
                    help='Batch size for training')
parser.add_argument('--resume', default=None, type=str,
                    help='Checkpoint state_dict file to resume training from. If this is "interrupt"'\
                         ', the model will resume training from the interrupt file.')
parser.add_argument('--start_iter', default=-1, type=int,
                    help='Resume training at this iter. If this is -1, the iteration will be'\
                         'determined from the file name.')
parser.add_argument('--num_workers', default=os.cpu_count(), type=int,
                    help='Number of workers used in dataloading')
parser.add_argument('--cuda', default=True, type=str2bool,
                    help='Use CUDA to train model')
parser.add_argument('--lr', '--learning_rate', default=None, type=float,
                    help='Initial learning rate. Leave as None to read this from the config.')
parser.add_argument('--momentum', default=None, type=float,
                    help='Momentum for SGD. Leave as None to read this from the config.')
parser.add_argument('--decay', '--weight_decay', default=None, type=float,
                    help='Weight decay for SGD. Leave as None to read this from the config.')
parser.add_argument('--gamma', default=None, type=float,
                    help='For each lr step, what to multiply the lr by. Leave as None to read this from the config.')
parser.add_argument('--save_folder', default='weights/',
                    help='Directory for saving checkpoint models.')
parser.add_argument('--log_folder', default='logs/',
                    help='Directory for saving logs.')
parser.add_argument('--config', default=None,
                    help='The config object to use.')
parser.add_argument('--save_interval', default=10000, type=int,
                    help='The number of iterations between saving the model.')
parser.add_argument('--validation_size', default=5000, type=int,
                    help='The number of images to use for validation.')
parser.add_argument('--validation_epoch', default=2, type=int,
                    help='Output validation information every n iterations. If -1, do no validation.')
parser.add_argument('--keep_latest', dest='keep_latest', action='store_true',
                    help='Only keep the latest checkpoint instead of each one.')
parser.add_argument('--keep_latest_interval', default=100000, type=int,
                    help='When --keep_latest is on, don\'t delete the latest file at these intervals. This should be a multiple of save_interval or 0.')
parser.add_argument('--dataset', default=None, type=str,
                    help='If specified, override the dataset specified in the config with this one (example: coco2017_dataset).')
parser.add_argument('--no_log', dest='log', action='store_false',
                    help='Don\'t log per iteration information into log_folder.')
parser.add_argument('--log_gpu', dest='log_gpu', action='store_true',
                    help='Include GPU information in the logs. Nvidia-smi tends to be slow, so set this with caution.')
parser.add_argument('--no_interrupt', dest='interrupt', action='store_false',
                    help='Don\'t save an interrupt when KeyboardInterrupt is caught.')
parser.add_argument('--batch_alloc', default=None, type=str,
                    help='If using multiple GPUS, you can set this to be a comma separated list detailing which GPUs should get what local batch size (It should add up to your total batch size).')
parser.add_argument('--no_autoscale', dest='autoscale', action='store_false',
                    help='YOLACT will automatically scale the lr and the number of iterations depending on the batch size. Set this if you want to disable that.')

parser.set_defaults(keep_latest=False, log=True, log_gpu=False, interrupt=True, autoscale=True)
args = parser.parse_args()

if args.config is not None:
    set_cfg(args.config)

if args.dataset is not None:
    set_dataset(args.dataset)

if args.autoscale and args.batch_size != 8:
    factor = args.batch_size / 8
    if __name__ == '__main__':
        print('Scaling parameters by %.2f to account for a batch size of %d.' % (factor, args.batch_size))

    cfg.lr *= factor
    cfg.max_iter //= factor
    cfg.lr_steps = [x // factor for x in cfg.lr_steps]

# Update training parameters from the config if necessary
def replace(name):
    if getattr(args, name) == None: setattr(args, name, getattr(cfg, name))
replace('lr')
replace('decay')
replace('gamma')
replace('momentum')

# This is managed by set_lr
cur_lr = args.lr

if torch.cuda.device_count() == 0:
    print('No GPUs detected. Exiting...')
    exit(-1)

if args.batch_size // torch.cuda.device_count() < 6:
    if __name__ == '__main__':
        print('Per-GPU batch size is less than the recommended limit for batch norm. Disabling batch norm.')
    cfg.freeze_bn = True

loss_types = ['B', 'C', 'M', 'P', 'D', 'E', 'S', 'I']

if torch.cuda.is_available():
    if args.cuda:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    if not args.cuda:
        print("WARNING: It looks like you have a CUDA device, but aren't " +
              "using CUDA.\nRun with --cuda for optimal training speed.")
        torch.set_default_tensor_type('torch.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

class NetLoss(nn.Module):
    """
    A wrapper for running the network and computing the loss
    This is so we can more efficiently use DataParallel.
    """

    def __init__(self, net:Yolact, criterion:MultiBoxLoss, pred_seg:bool=False):
        super().__init__()

        #NOTE
        self.pred_seg = pred_seg
        self.net = net
        self.criterion = criterion

    # one can also check the output format of the dataset
    def forward(self, images, targets, masks, num_crowds):
        preds = self.net(images)
        if self.pred_seg:
            losses, seg_list, pred_list = self.criterion(self.net, preds, targets, masks, num_crowds)
            return losses, seg_list, pred_list
        else:
            losses = self.criterion(self.net, preds, targets, masks, num_crowds)
            return losses

class CustomDataParallel(nn.DataParallel):
    """
    This is a custom version of DataParallel that works better with our training data.
    It should also be faster than the general case.
    """

    def scatter(self, inputs, kwargs, device_ids):
        # More like scatter and data prep at the same time. The point is we prep the data in such a way
        # that no scatter is necessary, and there's no need to shuffle stuff around different GPUs.
        devices = ['cuda:' + str(x) for x in device_ids]
        splits = prepare_data(inputs[0], devices, allocation=args.batch_alloc)

        return [[split[device_idx] for split in splits] for device_idx in range(len(devices))], \
            [kwargs] * len(devices)

    def gather(self, outputs, output_device):
        out = {}

        for k in outputs[0]:
            out[k] = torch.stack([output[k].to(output_device) for output in outputs])

        return out
# NOTE
def gan_init(mod):
    # Since the output value (before final activation) of discriminator
    # with possible negative initial values will easily be became zero while pass though
    # the ReLU. This will make the BCE loss explode
    # By default, while BCE loss explode, torch will clamp the value
    # to 100, in order to keep the gradient calculation vaild.
    if isinstance(mod, nn.Conv2d):
        # advice for relu activation
        # nn.init.kaiming_uniform_(mod.weight)
        # nn.init.zeros_(mod.bias)

        # advice for leakyReLU
        nn.init.normal_(mod.weight, 0.0, 0.02)
    elif isinstance(mod, nn.BatchNorm2d):
        nn.init.normal_(mod.weight, 1.0, 0.02)
        nn.init.constant_(mod.bias, 0)

def freeze_pretrain(model, freeze=True):
    if freeze:
        for par in model.parameters():
            par.requires_grad = False
    else:
        for par in model.parameters():
            par.requires_grad = True

def seg_mask_clas(seg_list, pred_list, datum):
    # need to transform seg_list to [b, classes, h, w]
    seg_list = [v.permute(2,0,1).contiguous() for v in seg_list]
    b = len(seg_list) # batch size
    _, seg_h, seg_w = seg_list[0].size()

    # empty tensors to be pasted on later (neglact the background class)
    seg_clas    = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)
    mask_clas   = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)

    # TODO: might change to upsample the prediction map
    # seg_clas = interpolate(seg_clas, size = (550,550),
    #                        mode='bilinear',align_corners=True)

    # get the gt target label
    target_list = [target for target in datum[1][0]]
    # input ground truth mask is [b, instances, 550, 550]
    # need to interpolate to     [b, classes, seg_h, seg_w]
    mask_list   = [interpolate(mask.unsqueeze(0), size = (seg_h,seg_w),mode='bilinear', \
                                        align_corners=False).squeeze() for mask in datum[1][1]]
    # all prepared, let's paste them all!
    # NOTE
    # === same class will be pasted together ===
    for idx in range(b):
        for i, (pred, i_target) in enumerate(zip(pred_list[idx], target_list[idx])):
            seg_clas[idx,  pred, ...]                += seg_list[idx][i,...]
            mask_clas[idx, i_target[-1].long(), ...] += mask_list[idx][i,...]

    seg_clas = torch.clamp(seg_clas, 0, 1)

    return seg_clas, mask_clas, b, (seg_h, seg_w)

def train():
    if not os.path.exists(args.save_folder):
        os.mkdir(args.save_folder)

    dataset = COCODetection(image_path=cfg.dataset.train_images,
                            info_file=cfg.dataset.train_info,
                            transform=SSDAugmentation(MEANS))

    if args.validation_epoch > 0:
        setup_eval()
        val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
                                    info_file=cfg.dataset.valid_info,
                                    transform=BaseTransform(MEANS))

    # Parallel wraps the underlying module, but when saving and loading we don't want that
    yolact_net = Yolact()
    net = yolact_net
    net.train()
    print('\n--- Generator created! ---')

    # NOTE
    # I maunally set the original image size and seg size as 138
    # might change in the future, for example 550
    if cfg.pred_seg:
        dis_size = 138
        dis_net  = Discriminator_Dcgan(i_size = dis_size, s_size = dis_size)
        # set the dis net's initial parameter values
        dis_net.apply(gan_init)
        dis_net.train()
        print('--- Discriminator created! ---\n')

    if args.log:
        log = Log(cfg.name, args.log_folder, dict(args._get_kwargs()),
            overwrite=(args.resume is None), log_gpu_stats=args.log_gpu)

    # I don't use the timer during training (I use a different timing method).
    # Apparently there's a race condition with multiple GPUs, so disable it just to be safe.
    timer.disable_all()

    # Both of these can set args.resume to None, so do them before the check
    if args.resume == 'interrupt':
        args.resume = SavePath.get_interrupt(args.save_folder)
    elif args.resume == 'latest':
        args.resume = SavePath.get_latest(args.save_folder, cfg.name)

    if args.resume is not None:
        print('Resuming training, loading {}...'.format(args.resume))
        yolact_net.load_weights(args.resume)

        if args.start_iter == -1:
            args.start_iter = SavePath.from_str(args.resume).iteration
    else:
        print('Initializing weights...')
        yolact_net.init_weights(backbone_path=args.save_folder + cfg.backbone.path)

    # optimizer_gen = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
    #                       weight_decay=args.decay)
    # if cfg.pred_seg:
    #     optimizer_dis = optim.SGD(dis_net.parameters(), lr=cfg.dis_lr, momentum=args.momentum,
    #                         weight_decay=args.decay)
    #     schedule_dis  = ReduceLROnPlateau(optimizer_dis, mode = 'min', patience=6, min_lr=1E-6)

    # NOTE: Using the Ranger Optimizer for the generator
    optimizer_gen     = Ranger(net.parameters(), lr = args.lr, weight_decay=args.decay)
    if cfg.pred_seg:
        optimizer_dis = optim.SGD(dis_net.parameters(), lr=cfg.dis_lr)
        schedule_dis  = ReduceLROnPlateau(optimizer_dis, mode = 'min', patience=6, min_lr=1E-6)

    criterion     = MultiBoxLoss(num_classes=cfg.num_classes,
                                pos_threshold=cfg.positive_iou_threshold,
                                neg_threshold=cfg.negative_iou_threshold,
                                negpos_ratio=cfg.ohem_negpos_ratio, pred_seg=cfg.pred_seg)

    criterion_dis = nn.BCELoss()

    if args.batch_alloc is not None:
        args.batch_alloc = [int(x) for x in args.batch_alloc.split(',')]
        if sum(args.batch_alloc) != args.batch_size:
            print('Error: Batch allocation (%s) does not sum to batch size (%s).' % (args.batch_alloc, args.batch_size))
            exit(-1)

    net = CustomDataParallel(NetLoss(net, criterion, pred_seg=cfg.pred_seg))
    if args.cuda:
        net     = net.cuda()
        if cfg.pred_seg:
            dis_net = dis_net.cuda()

    # Initialize everything
    if not cfg.freeze_bn: yolact_net.freeze_bn() # Freeze bn so we don't kill our means
    yolact_net(torch.zeros(1, 3, cfg.max_size, cfg.max_size).cuda())

    if not cfg.freeze_bn: yolact_net.freeze_bn(True)

    # loss counters
    loc_loss = 0
    conf_loss = 0
    iteration = max(args.start_iter, 0)
    last_time = time.time()

    epoch_size = len(dataset) // args.batch_size
    num_epochs = math.ceil(cfg.max_iter / epoch_size)

    # Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
    step_index = 0

    data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True)
    # NOTE
    val_loader  = data.DataLoader(val_dataset, args.batch_size,
                                  num_workers=args.num_workers*2,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True)

    save_path = lambda epoch, iteration: SavePath(cfg.name, epoch, iteration).get_path(root=args.save_folder)
    time_avg = MovingAverage()

    global loss_types # Forms the print order
                      # TODO: global command can modify global variable inside of the function.
    loss_avgs  = { k: MovingAverage(100) for k in loss_types }

    print('Begin training!')
    print()
    # try-except so you can use ctrl+c to save early and stop training
    try:
        for epoch in range(num_epochs):
            # Resume from start_iter
            if (epoch+1)*epoch_size < iteration:
                continue

            for datum in data_loader:
                # Stop if we've reached an epoch if we're resuming from start_iter
                if iteration == (epoch+1)*epoch_size:
                    break

                # Stop at the configured number of iterations even if mid-epoch
                if iteration == cfg.max_iter:
                    break

                # Change a config setting if we've reached the specified iteration
                changed = False
                for change in cfg.delayed_settings:
                    if iteration >= change[0]:
                        changed = True
                        cfg.replace(change[1])

                        # Reset the loss averages because things might have changed
                        for avg in loss_avgs:
                            avg.reset()

                # If a config setting was changed, remove it from the list so we don't keep checking
                if changed:
                    cfg.delayed_settings = [x for x in cfg.delayed_settings if x[0] > iteration]

                # Warm up by linearly interpolating the learning rate from some smaller value
                if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
                    set_lr(optimizer_gen, (args.lr - cfg.lr_warmup_init) * (iteration / cfg.lr_warmup_until) + cfg.lr_warmup_init)

                # Adjust the learning rate at the given iterations, but also if we resume from past that iteration
                while step_index < len(cfg.lr_steps) and iteration >= cfg.lr_steps[step_index]:
                    step_index += 1
                    set_lr(optimizer_gen, args.lr * (args.gamma ** step_index))

                # Zero the grad to get ready to compute gradients
                optimizer_gen.zero_grad()
                # NOTE
                if cfg.pred_seg:
                    # ====== GAN Train ======
                    optimizer_dis.zero_grad()
                    dis_net.zero_grad()
                    # train the gen and dis in different iteration count
                    it_alter_period = iteration % (cfg.gen_iter + cfg.dis_iter)
                    if it_alter_period >= cfg.gen_iter:
                        freeze_pretrain(yolact_net, freeze=True)
                        freeze_pretrain(net, freeze=True)
                        freeze_pretrain(dis_net, freeze=False)
                        if it_alter_period == (cfg.gen_iter+1):
                            print('--- Generator     freeze   ---')
                            print('--- Discriminator training ---')

                        # ----- Discriminator part -----
                        # seg_list  is the prediction mask
                        # can be regarded as generated images from YOLACT
                        # pred_list is the prediction label
                        # seg_list  dim: list of (138,138,instances)
                        # pred_list dim: list of (instances)
                        losses, seg_list, pred_list = net(datum)
                        seg_clas, mask_clas, b, seg_size = seg_mask_clas(seg_list, pred_list, datum)

                        # input image size is [b, 3, 550, 550]
                        # downsample to       [b, 3, seg_h, seg_w]
                        image    = interpolate(torch.stack(datum[0]), size = seg_size,
                                                    mode='bilinear',align_corners=False)

                        # Because in the discriminator training, we do not
                        # want the gradient flow back to the generator part
                        # we detach seg_clas (mask_clas come the data, does not have grad)
                        seg_input   = seg_clas.clone().detach()
                        output_pred = dis_net(img = image, seg = seg_input)
                        output_grou = dis_net(img = image, seg = mask_clas)

                        if iteration % (cfg.gen_iter + cfg.dis_iter) == 0:
                            print(f'Probability of fake is fake: {output_pred.mean().item():.2f}')
                            print(f'Probability of real is real: {output_grou.mean().item():.2f}')

                        # 0 for Fake/Generated
                        # 1 for True/Ground Truth
                        fake_label = torch.zeros(b)
                        real_label = torch.ones(b)

                        # Advice of practical implementation
                        # from https://arxiv.org/abs/1611.08408
                        # loss_pred = -criterion_dis(output_pred,target=real_label)
                        loss_pred = criterion_dis(output_pred,target=fake_label)
                        loss_grou = criterion_dis(output_grou,target=real_label)
                        loss_dis  = loss_pred + loss_grou

                        # TODO: Grid Search this one
                        lambda_dis = cfg.lambda_dis
                        loss_dis   = lambda_dis*loss_dis
                        # Backprop of the discriminator
                        loss_dis.backward()
                        optimizer_dis.step()

                    else:
                        freeze_pretrain(yolact_net, freeze=False)
                        freeze_pretrain(net, freeze=False)
                        freeze_pretrain(dis_net, freeze=True)
                        if it_alter_period == 0:
                            print('--- Generator     training ---')
                            print('--- Discriminator freeze   ---')

                        # ----- Generator part -----
                        net.zero_grad()
                        losses, seg_list, pred_list = net(datum)
                        seg_clas, mask_clas, b, seg_size = seg_mask_clas(seg_list, pred_list, datum)
                        image    = interpolate(torch.stack(datum[0]), size = seg_size,
                                                    mode='bilinear',align_corners=False)
                        # Perform forward pass of all-fake batch through D
                        # NOTE that seg_clas CANNOT detach, in order to flow the
                        # gradient back to the generator
                        output = dis_net(img = image, seg = seg_clas)
                        # Since the log(1-D(G(x))) not provide sufficient gradients
                        # We want log(D(G(x)) instead, this can be achieve by
                        # use the real_label as target.
                        # This step is crucial for the information of discriminator
                        # to go into the generator.
                        # Calculate G's loss based on this output
                        real_label = torch.ones(b)
                        loss_gen = criterion_dis(output,target=real_label)
                        if not it_alter_period >= cfg.gen_iter:
                            # since the dis is already freeze, the gradients will only
                            # record the YOLACT
                            loss_gen.backward()
                            # Do this to free up vram even if loss is not finite
                            losses = { k: (v).mean() for k,v in losses.items() } # Mean here because Dataparallel
                            loss = sum([losses[k] for k in losses])
                            if torch.isfinite(loss).item():
                                # since the optimizer_gen is for YOLACT only
                                # only the gen will be updated
                                optimizer_gen.step()

                else:
                    # ====== Normal YOLACT Train ======
                    # Forward Pass + Compute loss at the same time (see CustomDataParallel and NetLoss)
                    losses = net(datum)
                    losses = { k: (v).mean() for k,v in losses.items() } # Mean here because Dataparallel
                    loss = sum([losses[k] for k in losses])
                    # no_inf_mean removes some components from the loss, so make sure to backward through all of it
                    # all_loss = sum([v.mean() for v in losses.values()])

                    # Backprop
                    loss.backward() # Do this to free up vram even if loss is not finite
                    if torch.isfinite(loss).item():
                        optimizer_gen.step()

                # Add the loss to the moving average for bookkeeping
                for k in losses:
                    loss_avgs[k].add(losses[k].item())

                cur_time  = time.time()
                elapsed   = cur_time - last_time
                last_time = cur_time

                # Exclude graph setup from the timing information
                if iteration != args.start_iter:
                    time_avg.add(elapsed)

                if iteration % 10 == 0:
                    eta_str = str(datetime.timedelta(seconds=(cfg.max_iter-iteration) * time_avg.get_avg())).split('.')[0]

                    total = sum([loss_avgs[k].get_avg() for k in losses])
                    loss_labels = sum([[k, loss_avgs[k].get_avg()] for k in loss_types if k in losses], [])
                    if cfg.pred_seg and (it_alter_period >= cfg.gen_iter):
                        print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) + ' T: %.3f || Dis: %.2f || ETA: %s || timer: %.3f')
                                % tuple([epoch, iteration] + loss_labels + [total, loss_dis, eta_str, elapsed]), flush=True)
                    else:
                        print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) + ' T: %.3f ||  ETA: %s || timer: %.3f')
                                % tuple([epoch, iteration] + loss_labels + [total, eta_str, elapsed]), flush=True)
                    # Loss Key:
                    #  - B: Box Localization Loss
                    #  - C: Class Confidence Loss
                    #  - M: Mask Loss
                    #  - P: Prototype Loss
                    #  - D: Coefficient Diversity Loss
                    #  - E: Class Existence Loss
                    #  - S: Semantic Segmentation Loss
                    #  - T: Total loss
                    #  -Dis:Discriminator Loss

                if args.log:
                    precision = 5
                    loss_info = {k: round(losses[k].item(), precision) for k in losses}
                    loss_info['T'] = round(loss.item(), precision)

                    if args.log_gpu:
                        log.log_gpu_stats = (iteration % 10 == 0) # nvidia-smi is sloooow

                    log.log('train', loss=loss_info, epoch=epoch, iter=iteration,
                        lr=round(cur_lr, 10), elapsed=elapsed)

                    log.log_gpu_stats = args.log_gpu

                iteration += 1

                if iteration % args.save_interval == 0 and iteration != args.start_iter:
                    if args.keep_latest:
                        latest = SavePath.get_latest(args.save_folder, cfg.name)

                    print('Saving state, iter:', iteration)
                    yolact_net.save_weights(save_path(epoch, iteration))

                    if args.keep_latest and latest is not None:
                        if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
                            print('Deleting old save...')
                            os.remove(latest)

            # This is done per epoch
            if args.validation_epoch > 0:
                # NOTE: Validation loss
                # if cfg.pred_seg:
                #     net.eval()
                #     dis_net.eval()
                #     cfg.gan_eval = True
                #     with torch.no_grad():
                #         for datum in tqdm(val_loader, desc='GAN Validation'):
                #             losses, seg_list, pred_list = net(datum)
                #             losses, seg_list, pred_list = net(datum)
                #             # TODO: warp below as a function
                #             seg_list = [v.permute(2,1,0).contiguous() for v in seg_list]
                #             b = len(seg_list) # batch size
                #             _, seg_h, seg_w = seg_list[0].size()

                #             seg_clas    = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)
                #             mask_clas   = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)
                #             target_list = [target for target in datum[1][0]]
                #             mask_list   = [interpolate(mask.unsqueeze(0), size = (seg_h,seg_w),mode='bilinear', \
                #                             align_corners=False).squeeze() for mask in datum[1][1]]

                #             for idx in range(b):
                #                 for i, (pred, i_target) in enumerate(zip(pred_list[idx], target_list[idx])):
                #                     seg_clas[idx, pred, ...]                 += seg_list[idx][i,...]
                #                     mask_clas[idx, i_target[-1].long(), ...] += mask_list[idx][i,...]

                #             seg_clas = torch.clamp(seg_clas, 0, 1)
                #             image    = interpolate(torch.stack(datum[0]), size = (seg_h,seg_w),
                #                                         mode='bilinear',align_corners=False)
                #             real_label  = torch.ones(b)
                #             output_pred = dis_net(img = image, seg = seg_clas)
                #             output_grou = dis_net(img = image, seg = mask_clas)
                #             loss_pred   = -criterion_dis(output_pred,target=real_label)
                #             loss_grou   =  criterion_dis(output_grou,target=real_label)
                #             loss_dis    = loss_pred + loss_grou
                #         losses = { k: (v).mean() for k,v in losses.items() }
                #         loss = sum([losses[k] for k in losses])
                #         val_loss = loss - cfg.lambda_dis*loss_dis
                #         schedule_dis.step(loss_dis)
                #         lr = [group['lr'] for group in optimizer_dis.param_groups]
                #         print(f'Discriminator lr: {lr[0]}')
                #     net.train()
                if epoch % args.validation_epoch == 0 and epoch > 0:
                    cfg.gan_eval = False
                    dis_net.eval()
                    compute_validation_map(epoch, iteration, yolact_net, val_dataset, log if args.log else None)

        # Compute validation mAP after training is finished
        compute_validation_map(epoch, iteration, yolact_net, val_dataset, log if args.log else None)
    except KeyboardInterrupt:
        if args.interrupt:
            print('Stopping early. Saving network...')

            # Delete previous copy of the interrupted network so we don't spam the weights folder
            SavePath.remove_interrupt(args.save_folder)

            yolact_net.save_weights(save_path(epoch, repr(iteration) + '_interrupt'))
        exit()

    yolact_net.save_weights(save_path(epoch, iteration))

def set_lr(optimizer_gen, new_lr):
    for param_group in optimizer_gen.param_groups:
        param_group['lr'] = new_lr

    global cur_lr
    cur_lr = new_lr

def gradinator(x):
    x.requires_grad = False
    return x

def prepare_data(datum, devices:list=None, allocation:list=None):
    with torch.no_grad():
        if devices is None:
            devices = ['cuda:0'] if args.cuda else ['cpu']
        if allocation is None:
            allocation = [args.batch_size // len(devices)] * (len(devices) - 1)
            allocation.append(args.batch_size - sum(allocation)) # The rest might need more/less

        images, (targets, masks, num_crowds) = datum

        cur_idx = 0
        for device, alloc in zip(devices, allocation):
            # NOTE:
            # Since the allocation is for training
            # not for GAN validation
            try:
                for _ in range(alloc):
                    images[cur_idx]  = gradinator(images[cur_idx].to(device))
                    targets[cur_idx] = gradinator(targets[cur_idx].to(device))
                    masks[cur_idx]   = gradinator(masks[cur_idx].to(device))
                    cur_idx += 1
            except:
                print('GAN spile done')
                break

        if cfg.preserve_aspect_ratio:
            # Choose a random size from the batch
            _, h, w = images[random.randint(0, len(images)-1)].size()

            for idx, (image, target, mask, num_crowd) in enumerate(zip(images, targets, masks, num_crowds)):
                images[idx], targets[idx], masks[idx], num_crowds[idx] \
                    = enforce_size(image, target, mask, num_crowd, w, h)

        cur_idx = 0
        split_images, split_targets, split_masks, split_numcrowds \
            = [[None for alloc in allocation] for _ in range(4)]

        for device_idx, alloc in enumerate(allocation):
            split_images[device_idx]    = torch.stack(images[cur_idx:cur_idx+alloc], dim=0)
            split_targets[device_idx]   = targets[cur_idx:cur_idx+alloc]
            split_masks[device_idx]     = masks[cur_idx:cur_idx+alloc]
            split_numcrowds[device_idx] = num_crowds[cur_idx:cur_idx+alloc]

            cur_idx += alloc

        return split_images, split_targets, split_masks, split_numcrowds

def no_inf_mean(x:torch.Tensor):
    """
    Computes the mean of a vector, throwing out all inf values.
    If there are no non-inf values, this will return inf (i.e., just the normal mean).
    """

    no_inf = [a for a in x if torch.isfinite(a)]

    if len(no_inf) > 0:
        return sum(no_inf) / len(no_inf)
    else:
        return x.mean()

def compute_validation_loss(net, data_loader, criterion):
    global loss_types

    with torch.no_grad():
        losses = {}

        # Don't switch to eval mode because we want to get losses
        iterations = 0
        for datum in data_loader:
            images, targets, masks, num_crowds = prepare_data(datum)
            out = net(images)

            wrapper = ScatterWrapper(targets, masks, num_crowds)
            _losses = criterion(out, wrapper, wrapper.make_mask())

            for k, v in _losses.items():
                v = v.mean().item()
                if k in losses:
                    losses[k] += v
                else:
                    losses[k] = v

            iterations += 1
            if args.validation_size <= iterations * args.batch_size:
                break

        for k in losses:
            losses[k] /= iterations

        loss_labels = sum([[k, losses[k]] for k in loss_types if k in losses], [])
        print(('Validation ||' + (' %s: %.3f |' * len(losses)) + ')') % tuple(loss_labels), flush=True)

# compute_validation_map(epoch, iteration, yolact_net, val_dataset, log if args.log else None)
def compute_validation_map(epoch, iteration, yolact_net, dataset, log:Log=None):
    with torch.no_grad():
        yolact_net.eval()

        start = time.time()
        print('\n')
        print("Computing validation mAP (this may take a while)...", flush=True)
        val_info = eval_script.evaluate(yolact_net, dataset, train_mode=True)
        end = time.time()

        if log is not None:
            log.log('val', val_info, elapsed=(end - start), epoch=epoch, iter=iteration)

        yolact_net.train()

def setup_eval():
    eval_script.parse_args(['--no_bar', '--max_images='+str(args.validation_size)])

if __name__ == '__main__':
    train()


import torch
import numpy as np
from networks import get_network, set_requires_grad
from util.hausdorff import directed_hausdorff
from agent.base import GANzEAgent

class MainAgent(GANzEAgent):
    def __init__(self, config):
        super(MainAgent, self).__init__(config)

        self.weight_z_L1 = config.weight_z_L1
        self.weight_partial_rec = config.weight_partial_rec

    def build_net(self, config):
        # load pretrained pointAE
        self.pointAE = get_network(config, "pointAE")
        try:
            ae_weights = torch.load(config.pretrain_ae_path)['model_state_dict']
        except Exception as e:
            raise ValueError("Check the path for pretrained model of point AE. \n{}".format(e))
        self.pointAE.load_state_dict(ae_weights)
        self.pointAE = self.pointAE.eval().cuda()
        set_requires_grad(self.pointAE, False)

        # load pretrained pointVAE
        pointVAE = get_network(config, "pointVAE")
        try:
            vae_weights = torch.load(config.pretrain_vae_path)['model_state_dict']
        except Exception as e:
            raise ValueError("Check the path for pretrained model of point VAE. \n{}".format(e))
        pointVAE.load_state_dict(vae_weights)
        self.netE = pointVAE.encoder.eval().cuda()
        self.vaeD = pointVAE.decoder.eval().cuda()
        set_requires_grad(self.netE, False)  # netE remains fixed

        # build G, D
        self.netG = get_network(config, "G").cuda()
        self.netD = get_network(config, "D").cuda()

    def collect_loss(self):
        loss_dict = {"D_GAN": self.loss_D,
                     "G_GAN": self.loss_G_GAN,
                     "z_L1": self.loss_z_L1,
                     "partial_rec": self.loss_partial_rec}
        return loss_dict

    def forward(self, data):
        self.raw_pc = data['raw'].cuda()
        self.real_pc = data['real'].cuda()

        with torch.no_grad():
            self.raw_latent = self.pointAE.encode(self.raw_pc)
            self.real_latent = self.pointAE.encode(self.real_pc)

        self.forward_GE()

    def forward_GE(self):
        self.z_random = self.get_random_noise(self.raw_latent.size(0))

        self.fake_latent = self.netG(self.raw_latent, self.z_random)
        self.fake_pc = self.pointAE.decode(self.fake_latent)
        self.z_rec, z_mu, z_logvar = self.netE(self.fake_pc)

    def backward_D(self):
        # fake
        pred_fake = self.netD(self.fake_latent.detach())
        fake = torch.zeros_like(pred_fake).fill_(0.0).cuda()
        self.loss_D_fake = self.criterionGAN(pred_fake, fake)

        # real
        pred_real = self.netD(self.real_latent.detach())
        real = torch.ones_like(pred_real).fill_(1.0).cuda()
        self.loss_D_real = self.criterionGAN(pred_real, real)

        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()

    def update_D(self):
        set_requires_grad(self.netD, True)

        self.optimizer_D.zero_grad()
        self.backward_D()
        self.optimizer_D.step()

    def backward_EG(self):
        # 1. G(A) fool D
        pred_fake = self.netD(self.fake_latent)
        real = torch.ones_like(pred_fake).fill_(1.0).cuda()
        self.loss_G_GAN = self.criterionGAN(pred_fake, real)

        # 2. noise reconstruction |E(G(A, z)) - z_random|
        self.loss_z_L1 = self.criterionL1(self.z_rec, self.z_random) * self.weight_z_L1

        # 3. partial scan reconstruction
        self.loss_partial_rec = directed_hausdorff(self.raw_pc, self.fake_pc) * self.weight_partial_rec

        self.loss_EG = self.loss_G_GAN + self.loss_z_L1 + self.loss_partial_rec
        self.loss_EG.backward()

    def update_G_and_E(self):
        set_requires_grad(self.netD, False)
        self.optimizer_G.zero_grad()
        self.backward_EG()
        self.optimizer_G.step()

    def optimize_network(self):
        self.update_G_and_E()
        self.update_D()

    def get_point_cloud(self):
        """get real/fake/raw point cloud of current batch"""
        real_pts = self.real_pc.transpose(1, 2).detach().cpu().numpy()
        fake_pts = self.fake_pc.transpose(1, 2).detach().cpu().numpy()
        raw_pts = self.raw_pc.transpose(1, 2).detach().cpu().numpy()
        return real_pts, fake_pts, raw_pts

    def visualize_batch(self, data, mode, **kwargs):
        tb = self.train_tb if mode == 'train' else self.val_tb

        num = 2

        real_pts = data['real'][:num].transpose(1, 2).detach().cpu().numpy()
        fake_pts = self.fake_pc[:num].transpose(1, 2).detach().cpu().numpy()
        raw_pts = self.raw_pc[:num].transpose(1, 2).detach().cpu().numpy()

        fake_pts = np.clip(fake_pts, -0.999, 0.999)

        tb.add_mesh("real", vertices=real_pts, global_step=self.clock.step)
        tb.add_mesh("fake", vertices=fake_pts, global_step=self.clock.step)
        tb.add_mesh("input", vertices=raw_pts, global_step=self.clock.step)


"""
Author: Kelvin Hong
Training MobileNetV2-SSD model.
Reference: https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Object-Detection
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torchsummary
import argparse
from PIL import Image
import os
import glob
from utils import *
import model
import math
from math import sqrt
import random
import numpy as np
import datetime
import time

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

if __name__ == '__main__':
    folder = "001"
    parser = argparse.ArgumentParser(description='Accept number of epoch.')
    parser.add_argument("-e", "--epoch", type=int, default=50, help='Number of epochs for training. Default 50.')
    args = parser.parse_args()
    labels = open('coco.names').read().splitlines()
    num_classes = len(labels)
    # Creating dataset
    images = []
    locs_labels = []
    # Seed to make sure sampling is reproducible
    random.seed(0)
    raw_paths = list(glob.iglob(os.path.join(folder, "*.jpg")))
    random.shuffle(raw_paths)
    n = len(raw_paths)
    i = 0
    for image_path in raw_paths:

        loc_label = []
        try:
            f = open(image_path[:-3]+"txt")
        except:
            pass
        else:
            lines = f.read().splitlines()
            images.append(image_path)
            for line in lines:
                label, *locs = line.split(' ')
                label = int(label)
                locs = list(map(float, locs))
                loc_label.append([label]+locs)
            locs_labels.append(loc_label)
        i += 1
        print(i, '/', n)

    # Take 70% as training data
    whole_n = len(images)
    train_n = int(0.7*whole_n)
    train_images = images[:train_n]
    train_locs_labels = locs_labels[:train_n]
    test_images = images[train_n:]
    test_locs_labels = locs_labels[train_n:]

    # Configuration
    num_epochs = args.epoch
    batch_size = 4

    transform = transforms.Compose([
        transforms.Resize((300,300)),
        transforms.ColorJitter(brightness=0.5, contrast=0.5,saturation=0.5),
        transforms.ToTensor()
    ])
    train_dataset = model.CustomDataset(train_images, train_locs_labels, transform = transform)
    test_dataset = model.CustomDataset(test_images, test_locs_labels, transform=transform)
    trainloader = DataLoader(train_dataset, batch_size = batch_size, collate_fn=model.collate_fn, shuffle=True, drop_last=True)
    testloader = DataLoader(test_dataset, batch_size=batch_size, collate_fn=model.collate_fn, shuffle=False, drop_last=True)

    # Define model and start training
    net = model.MobileNetV2_SSD(num_classes).to(device)
    net.apply(model.weights_init_rule)
    torchsummary.summary(net, (3,300,300),1,"cuda")
    optimizer = optim.Adam(net.parameters(), lr=0.005)
    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[20,40,60,80], gamma=0.5)
    prior_boxes = net.create_prior_boxes()
    criterion = model.MultiBoxLoss(prior_boxes).to(device)
    # scheduler =
    step = 20

    net.eval()
    dummy_input = torch.randn(1,3,300,300).to(device)
    onnx_model_path = f"models/mobilenetv2_ssd_{datetime.datetime.now()}.onnx"
    torch.onnx.export(net, dummy_input, onnx_model_path, verbose=False)

    loader_n = len(trainloader)
    test_loader_n = len(testloader)
    for epoch in range(num_epochs):
        # Training
        net.train()
        accu_loss = 0
        accu_conf_loss = 0
        accu_loc_loss = 0
        start = time.time()
        for i, data in enumerate(trainloader, 0):

            batch_images = data['images']
            batch_locs = data['locs']
            batch_labels = data['labels']
            pred_locs, pred_scores = net(batch_images)
            # print("Predicted Locations: ", pred_locs)
            # print("Predicted Scores: ", pred_scores)
            conf_loss, loc_loss, loss = criterion(pred_locs, pred_scores, batch_locs, batch_labels)
            accu_conf_loss += conf_loss.item()
            accu_loc_loss += loc_loss.item()
            accu_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if i%step == step-1:
                print(f"Epoch {epoch+1} avg loss at step {str(i+1).rjust(4)} of {str(loader_n).rjust(4)}: {round(accu_loss/i,2)}. Conf loss: {round(accu_conf_loss/i,2)}. Loc loss: {round(accu_loc_loss/i,2)}. lr={round(get_lr(optimizer),6)}.")
            # Freeing memories
            del pred_scores, pred_locs, batch_images, batch_locs, batch_labels
        end = time.time()
        te = convert_sec(end-start)
        print(f"Epoch {epoch+1} training finished. Proceed to validation. Time Elapsed {te}.")

        # Evaluating
        net.eval()
        taccu_loss = 0
        taccu_conf_loss = 0
        taccu_loc_loss = 0
        start = time.time()
        with torch.no_grad():
            for i, data in enumerate(testloader, 0):
                batch_images = data['images']
                batch_locs = data['locs']
                batch_labels = data['labels']
                pred_locs, pred_scores = net(batch_images)
                conf_l, loc_l, loss = criterion(pred_locs, pred_scores, batch_locs, batch_labels)
                taccu_conf_loss += conf_l.item()
                taccu_loc_loss += loc_l.item()
                taccu_loss += loss.item()
                if i%step == step-1:
                    print(f"Epoch {epoch+1} validation avg loss at step {str(i+1).rjust(4)} of {str(test_loader_n).rjust(4)}: {round(taccu_loss/i,2)}. Conf loss: {round(taccu_conf_loss/i,2)}. Loc loss: {round(taccu_loc_loss/i,2)}. lr={round(get_lr(optimizer),6)}.")
                # Freeing memories
                del pred_scores, pred_locs, batch_images, batch_locs, batch_labels
        end = time.time()
        te = convert_sec(end-start)

        scheduler.step()
        print(f"Epoch {epoch+1} validation finished. Time Elapsed {te}.")
        if (epoch+1)%10 == 0 or epoch+1==num_epochs:
            print(f"Saving Model epoch {epoch+1}...")
            torch.save({
                'epoch': epoch+1,
                'model_state_dict': net.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': loss,
            }, f'models/model_mobilenetv2_ssd_Epoch{epoch+1}_{datetime.datetime.now()}.pth')
            print(f"Saved!")
            print("You can find the model in the models/ folder.")
        print("="*50)