import cv2
import matplotlib as mpl
import numpy as np
from sklearn.preprocessing import normalize
from tqdm import tqdm

mpl.use('agg')

import matplotlib.pyplot as plt

def rank_vis(query_mat, gallery_mat, n_match, save_dir):
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)

    query_feat = query_mat['feat']
    query_ids = query_mat['ids'].squeeze()
    query_cam_ids = query_mat['cam_ids'].squeeze()
    query_img_paths = query_mat['img_path'].tolist()

    gallery_feat = gallery_mat['feat']
    gallery_ids = gallery_mat['ids'].squeeze()
    gallery_cam_ids = gallery_mat['cam_ids'].squeeze()
    gallery_img_paths = gallery_mat['img_path'].tolist()

    dist = -np.dot(normalize(query_feat), normalize(gallery_feat).T)

    rank_list = np.argsort(dist, axis=1)

    num_query = len(query_ids)
    num_gallery = len(gallery_ids)

    for i in tqdm(range(num_query)):
        index = np.arange(num_gallery)

        junk_flag = np.equal(gallery_ids, -1)
        intra_flag = np.equal(gallery_cam_ids, query_cam_ids[i]),
        inter_flag = np.not_equal(gallery_cam_ids, query_cam_ids[i]),

        junk_index = index[junk_flag]
        intra_index = index[intra_flag]
        inter_index = index[inter_flag]

        sort_index = rank_list[i]
        flag = np.logical_not(np.isin(sort_index, junk_index))
        intra_flag = np.isin(sort_index, intra_index)
        inter_flag = np.isin(sort_index, inter_index)

        sort_index = sort_index[flag]
        intra_sort_index = sort_index[np.logical_and(flag, intra_flag)]
        inter_sort_index = sort_index[np.logical_and(flag, inter_flag)]

        plt.figure(figsize=(4 * n_match, 30))

        q_img_path = query_img_paths[i].strip()
        q_img = cv2.imread(q_img_path)
        q_img = cv2.cvtColor(q_img, cv2.COLOR_BGR2RGB)
        q_img = cv2.resize(q_img, (128, 256))

        if not os.path.basename(q_img_path) == '0115_c1_f0073267.jpg':
            continue

        # --------------------------------- #
        plt.subplot(3, n_match + 1, 1)

        plt.imshow(q_img)
        plt.text(30, -8, 'cam %d' % query_cam_ids[i], size=45)
        plt.xticks([])
        plt.yticks([])

        # --------------------------------- #
        plt.subplot(3, n_match + 1, 1 + n_match + 1)

        plt.imshow(q_img)
        plt.text(30, -8, 'cam %d' % query_cam_ids[i], size=45)
        plt.xticks([])
        plt.yticks([])

        # --------------------------------- #
        plt.subplot(3, n_match + 1, 1 + (n_match + 1) * 2)

        plt.imshow(q_img)
        plt.text(30, -8, 'cam %d' % query_cam_ids[i], size=45)
        plt.xticks([])
        plt.yticks([])

        for j in range(n_match):
            plt.subplot(3, n_match + 1, j + 2)

            g_img_path = gallery_img_paths[sort_index[j]].strip()
            g_img = cv2.imread(g_img_path)
            g_img = cv2.cvtColor(g_img, cv2.COLOR_BGR2RGB)

            plt.text(30, -8, 'cam %d' % gallery_cam_ids[sort_index[j]], size=45)

            plt.imshow(cv2.resize(g_img, (128, 256)))
            plt.xticks([])
            plt.yticks([])

            if query_ids[i] != gallery_ids[sort_index[j]]:
                color = "red"
            else:
                color = "green"

            ax = plt.gca()
            ax.spines['right'].set_linewidth(8)
            ax.spines['top'].set_linewidth(8)
            ax.spines['left'].set_linewidth(8)
            ax.spines['bottom'].set_linewidth(8)
            ax.spines['right'].set_color(color)
            ax.spines['top'].set_color(color)
            ax.spines['left'].set_color(color)
            ax.spines['bottom'].set_color(color)

            # --------------------------------- #
            plt.subplot(3, n_match + 1, j + 2 + n_match + 1)

            g_img_path = gallery_img_paths[intra_sort_index[j]].strip()
            g_img = cv2.imread(g_img_path)
            g_img = cv2.cvtColor(g_img, cv2.COLOR_BGR2RGB)

            plt.text(30, -8, 'cam %d' % gallery_cam_ids[intra_sort_index[j]], size=45)
            # plt.text(30, -8, '%.3f' % dist[i][intra_sort_index[j]], size=45)

            plt.imshow(cv2.resize(g_img, (128, 256)))
            plt.xticks([])
            plt.yticks([])

            if query_ids[i] != gallery_ids[intra_sort_index[j]]:
                color = "red"
            else:
                color = "green"

            ax = plt.gca()
            ax.spines['right'].set_linewidth(8)
            ax.spines['top'].set_linewidth(8)
            ax.spines['left'].set_linewidth(8)
            ax.spines['bottom'].set_linewidth(8)
            ax.spines['right'].set_color(color)
            ax.spines['top'].set_color(color)
            ax.spines['left'].set_color(color)
            ax.spines['bottom'].set_color(color)

            # --------------------------------- #
            plt.subplot(3, n_match + 1, j + 2 + (n_match + 1) * 2)

            g_img_path = gallery_img_paths[inter_sort_index[j]].strip()
            g_img = cv2.imread(g_img_path)
            g_img = cv2.cvtColor(g_img, cv2.COLOR_BGR2RGB)

            plt.text(30, -8, 'cam %d' % gallery_cam_ids[inter_sort_index[j]], size=45)
            # plt.text(30, -8, '%.3f' % dist[i][inter_sort_index[j]], size=45)

            plt.imshow(cv2.resize(g_img, (128, 256)))
            plt.xticks([])
            plt.yticks([])

            if query_ids[i] != gallery_ids[inter_sort_index[j]]:
                color = "red"
            else:
                color = "green"

            ax = plt.gca()
            ax.spines['right'].set_linewidth(8)
            ax.spines['top'].set_linewidth(8)
            ax.spines['left'].set_linewidth(8)
            ax.spines['bottom'].set_linewidth(8)
            ax.spines['right'].set_color(color)
            ax.spines['top'].set_color(color)
            ax.spines['left'].set_color(color)
            ax.spines['bottom'].set_color(color)

        plt.tight_layout()
        plt.savefig("{}/{}.jpg".format(save_dir, os.path.basename(q_img_path)))
        plt.close()

if __name__ == '__main__':
    import os
    import argparse

    import scipy.io as sio

    parser = argparse.ArgumentParser()
    parser.add_argument("dataset", type=str, choices=["market", "duke", "msmt", "cuhk"])
    parser.add_argument("prefix", type=str)
    args = parser.parse_args()

    q_mat_path = os.path.join("features", args.dataset, "query-%s.mat" % args.prefix)
    g_mat_path = os.path.join("features", args.dataset, "gallery-%s.mat" % args.prefix)

    g_mat = sio.loadmat(g_mat_path)
    q_mat = sio.loadmat(q_mat_path)

    rank_vis(q_mat, g_mat, n_match=40, save_dir="/home/chuanchen_luo")


#!/usr/bin/env python3
# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.

import torch
from collections import OrderedDict
import math
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
import torch.distributed as dist
from tqdm import tqdm
import json
import argparse
import logging
from torch import nn
import random
import os
import numpy as np
import pandas as pd
from transformers import BertConfig, BertModel, BertTokenizer, AutoTokenizer, AutoModelWithLMHead, AutoModel
from data.data import *
from models.siamese_model import *
from transformers import AdamW, get_linear_schedule_with_warmup
from torch.cuda.amp import autocast, GradScaler
from multiprocessing import Pool, cpu_count
from functools import partial
import utils.logger_utils
logger = utils.logger_utils.get_logger()

def set_seed():
    random.seed(5)
    np.random.seed(5)
    torch.manual_seed(5)
    torch.cuda.manual_seed_all(5)

# configurations
model_name = 'bert-base-uncased'
tokenizer_class = BertTokenizer
model_class = BertModel
config_class = BertConfig

# arguments
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', default=False, action='store_true')
parser.add_argument('--init_model', default=None, type=str)
parser.add_argument('--query_len', default=100, type=int)
parser.add_argument('--passage_len', default=100, type=int)
parser.add_argument('--epochs', default=1, type=int)
parser.add_argument('--warmup', default=0.1, type=float)
parser.add_argument('--lr', default=3e-5, type=float)
parser.add_argument('--wd', default=0, type=float)
parser.add_argument('--sample_size', default=24, type=int)
parser.add_argument('--batch_size', default=24, type=int)
parser.add_argument('--print_iter', default=100, type=int)
parser.add_argument('--eval_iter', default=1000, type=int)
parser.add_argument('--log', default=None, type=str)
parser.add_argument('--save_model', default='best_model.pt', type=str)
parser.add_argument('--passagefile', default=None, type=str)
parser.add_argument('--train_queryfile', default=None, type=str)
parser.add_argument('--train_answerfile', default=None, type=str)
parser.add_argument('--train_negativefile', default=None, type=str)
parser.add_argument('--val_passagefile', default=None, type=str)
parser.add_argument('--val_queryfile', default=None, type=str)
parser.add_argument('--val_answerfile', default=None, type=str)
parser.add_argument('--share', default=False, action='store_true')
parser.add_argument('--embed_size', default=128, type=int)
parser.add_argument('--mode', default='CLS', type=str)
parser.add_argument('--normalize', default=False, action='store_true')
parser.add_argument('--train', default=False, action='store_true')
parser.add_argument('--doc_embed', default=False, action='store_true')
parser.add_argument('--doc_embed_file', default=None, type=str)
parser.add_argument('--query_embed', default=False, action='store_true')
parser.add_argument('--query_embed_file', default=None, type=str)
parser.add_argument('--generate_retrieval', default=False, action='store_true')
parser.add_argument('--test_queryfile', default=None, type=str)
parser.add_argument('--topk', default=500, type=int)
parser.add_argument('--retrieval_outputfile', default=None, type=str)
parser.add_argument('--data_mode', default='WikiGQ', type=str)
parser.add_argument('--ICT_sent_keep', default=0.1, type=float)
parser.add_argument('--random', default=False, action='store_true')
parser.add_argument('--hard_negative_rate', default=0.0, type=float)
parser.add_argument('--neg_q', default=0, type=int)
parser.add_argument('--biobert', default=False, action='store_true')
# checkpoints
parser.add_argument('--ckpt_dir', default=None, type=str)
parser.add_argument('--resume', default=None, type=str, help='ckpt-epoch-file-steps')

# distributed training
parser.add_argument('--num_nodes', default=1, type=int)
parser.add_argument('--num_gpus', default=1, type=int)
parser.add_argument('--local_rank', default=0, type=int)
parser.add_argument('--node_rank', default=0, type=int)
parser.add_argument('--num_workers', default=4, type=int)
parser.add_argument('--dist_url', default=None, type=str)

args = parser.parse_args()

set_seed()
if args.log:
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    fhandler = logging.FileHandler(args.log)
    fhandler.setFormatter(formatter)
    logger.addHandler(fhandler)
logger.info(args)
n_gpus = args.num_gpus
num_nodes = args.num_nodes
world_size = n_gpus * num_nodes
distributed_training = True if world_size > 1 and args.local_rank > -1 else False
rank = args.node_rank * n_gpus + args.local_rank

if args.train and distributed_training:
    logger.info('setting up distributed traning... ')
    if num_nodes > 1:
        dist_url = args.dist_url
        dist.init_process_group(backend="nccl", init_method=dist_url, world_size=world_size, rank=rank)
    else:
        dist_url = "tcp://127.0.0.1:8080"
        dist.init_process_group(backend="nccl", init_method=dist_url, world_size=world_size, rank=rank)
    logger.info(
        "Initialized the distributed environment: NCCL backed on %d nodes." % (num_nodes))
    logger.info(
        "World size: %d, Node rank %d, GPU Rank %d, Local rank %d" % (
        world_size, args.node_rank, rank, args.local_rank))
    torch.cuda.set_device(args.local_rank)

device = torch.device('cuda') if args.gpu else torch.device('cpu')

# initialize tokenizer, configs, and bert model
logger.info('setting up model from configs')
if args.biobert:
    tokenizer = AutoTokenizer.from_pretrained("monologg/biobert_v1.1_pubmed")
    bert = AutoModel.from_pretrained("monologg/biobert_v1.1_pubmed")
else:
    tokenizer = tokenizer_class.from_pretrained(model_name, do_lower_case=True, cache_dir=None)
    config = config_class.from_pretrained(model_name)
    bert = model_class.from_pretrained(model_name, config=config)

if args.random:
    bert.init_weights()

doc_bert = BertEmbed(bert, embedding_size=args.embed_size, padding_id=tokenizer.pad_token_id,
                     mode=args.mode, normalize=args.normalize)

if args.share:
    logger.info('using tied weights')
    query_bert = doc_bert
else:
    logger.info('not sharing parameters')
    if args.biobert:
        tokenizer = AutoTokenizer.from_pretrained("monologg/biobert_v1.1_pubmed")
        bert = AutoModel.from_pretrained("monologg/biobert_v1.1_pubmed")
    else:
        tokenizer = tokenizer_class.from_pretrained(model_name, do_lower_case=True, cache_dir=None)
        config = config_class.from_pretrained(model_name)
        bert = model_class.from_pretrained(model_name, config=config)
    query_bert = BertEmbed(bert, embedding_size=args.embed_size, padding_id=tokenizer.pad_token_id,
                           mode=args.mode, normalize=args.normalize)

if args.init_model:
    try:
        missing_keys, unexpected_keys = doc_bert.load_state_dict(torch.load(args.init_model + '.doc', map_location='cpu'), strict=False)
        logger.info(f"Loading init Doc model {args.init_model}.doc")
        assert len(unexpected_keys) == 0
        assert len(missing_keys) == 0 or missing_keys == ["bert.embeddings.position_ids"]
        missing_keys, unexpected_keys = query_bert.load_state_dict(torch.load(args.init_model + '.query', map_location='cpu'), strict=False)
        logger.info(f"Loading init Query model {args.init_model}.query")
        assert len(unexpected_keys) == 0
        assert len(missing_keys) == 0 or missing_keys == ["bert.embeddings.position_ids"]
    except:
        logger.info("Loading init model failed!!!")

if args.resume:
    ckpt_dir = os.path.join(args.ckpt_dir, f'{args.resume}')
    doc_bert.load_state_dict(torch.load(os.path.join(ckpt_dir, 'ckpt.doc'), map_location='cpu'))
    logger.info(f"Resume Doc model {os.path.join(ckpt_dir, 'ckpt.doc')}")
    query_bert.load_state_dict(torch.load(os.path.join(ckpt_dir, 'ckpt.query'), map_location='cpu'))
    logger.info(f"Resume Query model {os.path.join(ckpt_dir, 'ckpt.query')}")
else:
    ckpt_dir = None

query_bert.to(device)
doc_bert.to(device)

if args.train:
    # prepare training data
    logger.info('preparing data')
    if args.data_mode == 'ICT':
        padding_func = lambda x: ICT_batchify(x, tokenizer.pad_token_id, args.query_len, args.passage_len)
    elif args.data_mode == 'WikiGQ':
        padding_func = lambda x: batchify(x, tokenizer.pad_token_id, args.query_len, args.passage_len)
    else:
        raise ValueError('Wrong Data Mode')

    if os.path.isdir(args.passagefile):
        # passage file is a directory (list of files). do streaming
        logger.info('filestream training')

        logger.info('computing number of query-answer pairs')
        num_pairs = 0
        if args.data_mode == 'wiki_gen_query':
            ans_files = [n for n in os.listdir(args.train_answerfile) if n[-4:] == 'json']
            for f in ans_files:
                ans_dict = json.load(open(os.path.join(args.train_answerfile, f), 'r'))
                num_answers = sum([len(ans) for ans in ans_dict.values()])
                num_pairs += num_answers
        elif args.data_mode == 'ICT':
            passage_files = [n for n in os.listdir(args.passagefile) if n[-5:] == 'jsonl']
            for f in passage_files:
                num_pairs += len(open(os.path.join(args.passagefile, f), 'r').readlines())

        else:
            raise Exception('Invalid task.')

        epoch_len = num_pairs // (args.batch_size * world_size)
        logger.info('Training data: %d samples, %d batches', num_pairs, epoch_len)
    else:
        logger.info('single file training')
        if args.data_mode == 'WikiGQ':
            train_data = SiameseDataset(tokenizer=tokenizer, query_path=args.train_queryfile, passage_path=args.passagefile,
                                        answer_path=args.train_answerfile, sampled_softmax_n=args.sample_size,
                                        neg_path=args.train_negativefile)

        elif args.data_mode == 'ICT':
            train_data = ICTDataset(tokenizer=tokenizer, passage_path=args.passagefile, sent_keep=args.ICT_sent_keep,
                                    sampled_softmax_n=args.sample_size, hard_negative_rate=args.hard_negative_rate)

        else:
            raise Exception('Invalid task.')

        train_sampler = DistributedSampler(train_data) if distributed_training else RandomSampler(train_data)
        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.batch_size,
                                      num_workers=args.num_workers, pin_memory=True,
                                      collate_fn=padding_func)
        logger.info('Training data: %d samples, %d batches', len(train_data), len(train_dataloader))
        train_dataloaders = [(train_dataloader, train_sampler)]
        epoch_len = len(train_dataloader)

    if args.eval_iter > 0:
        val_data = SiameseDataset(tokenizer=tokenizer, query_path=args.val_queryfile, passage_path=args.val_passagefile,
                              answer_path=args.val_answerfile, sampled_softmax_n=args.sample_size)
        val_sampler = SequentialSampler(val_data)
        val_dataloader = DataLoader(val_data, sampler=val_sampler, batch_size=args.batch_size,
                                    num_workers=args.num_workers, pin_memory=True,
                                    collate_fn=lambda x:
                                    batchify(x, tokenizer.pad_token_id, args.query_len, args.passage_len))
        logger.info('Validation data: %d samples, %d batches', len(val_data), len(val_dataloader))

    # prepare optimizer
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters1 = [
        {
            "params": [p for n, p in query_bert.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": args.wd,
        },
        {
            "params": [p for n, p in query_bert.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0},
    ]
    lr = args.lr
    query_optimizer = AdamW(optimizer_grouped_parameters1, lr=lr, eps=1e-8)
    t_total = epoch_len * args.epochs
    num_warmup_steps = int(args.warmup * t_total)
    query_scheduler = get_linear_schedule_with_warmup(
        query_optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=t_total
    )
    if (ckpt_dir and os.path.isfile(os.path.join(ckpt_dir, "query_optimizer.pt"))
            and os.path.isfile(os.path.join(ckpt_dir, "query_scheduler.pt"))):
        # Load in optimizer and scheduler states
        query_optimizer.load_state_dict(torch.load(os.path.join(ckpt_dir, "query_optimizer.pt"), map_location='cpu'))
        query_scheduler.load_state_dict(torch.load(os.path.join(ckpt_dir, "query_scheduler.pt"), map_location='cpu'))
        logger.info(f'Load query optimizer states from {os.path.join(ckpt_dir, "query_optimizer.pt")}')

    if not args.share:
        optimizer_grouped_parameters2 = [
            {
                "params": [p for n, p in doc_bert.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": args.wd,
            },
            {
                "params": [p for n, p in doc_bert.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0},
        ]
        doc_optimizer = AdamW(optimizer_grouped_parameters2, lr=lr, eps=1e-8)
        doc_scheduler = get_linear_schedule_with_warmup(
            doc_optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=t_total
        )
        if (ckpt_dir and os.path.isfile(os.path.join(ckpt_dir, "doc_optimizer.pt"))
                and os.path.isfile(os.path.join(ckpt_dir, "doc_scheduler.pt"))):
            # Load in optimizer and scheduler states
            doc_optimizer.load_state_dict(torch.load(os.path.join(ckpt_dir, "doc_optimizer.pt"), map_location='cpu'))
            doc_scheduler.load_state_dict(torch.load(os.path.join(ckpt_dir, "doc_scheduler.pt"), map_location='cpu'))
            logger.info(f'Load doc optimizer states from {os.path.join(ckpt_dir, "doc_optimizer.pt")}')

    else:
        doc_bert = query_bert
        doc_optimizer = None
        doc_scheduler = None
    if args.num_gpus > 1 and not distributed_training:
        doc_bert = torch.nn.DataParallel(doc_bert)
        query_bert = torch.nn.DataParallel(query_bert)

    # multi-gpu training (should be after apex fp16 initialization)
    if distributed_training:
        logger.info('using distributed training')
        query_bert = torch.nn.parallel.DistributedDataParallel(query_bert, device_ids=[args.local_rank],
                                                               find_unused_parameters=True)
        if not args.share:
            doc_bert = torch.nn.parallel.DistributedDataParallel(doc_bert, device_ids=[args.local_rank],
                                                                 find_unused_parameters=True)
    else:
        logger.info('using single machine training')

def save_model(query_bert, doc_bert, save_model):
    query_model_to_save = query_bert.module if hasattr(query_bert, "module") else query_bert
    torch.save(query_model_to_save.state_dict(), save_model+'.query')
    doc_model_to_save = doc_bert.module if hasattr(doc_bert, "module") else doc_bert
    torch.save(doc_model_to_save.state_dict(), save_model + '.doc')

def evaluation(query_bert, doc_bert, loss_fun, val_data):
    query_bert.eval()
    doc_bert.eval()
    total_loss = 0
    with torch.no_grad():
        for step, batch in enumerate(tqdm(val_data)):
            batch = tuple(t.to(device) for t in batch)
            query_tokens, passages_tokens, labels = batch
            query_embeds = query_bert(query_tokens)
            passages_embeds = doc_bert(passages_tokens)
            scores = torch.mm(query_embeds, passages_embeds.t())
            val_loss = loss_fun(scores, labels).mean()
            total_loss += val_loss
    avg_loss = total_loss / len(val_data)
    query_bert.train()
    doc_bert.train()
    return avg_loss

def train():

    # start training
    logger.info('starting training')
    loss_fun = nn.CrossEntropyLoss()
    query_bert.train()
    doc_bert.train()
    scaler = GradScaler()
    global_step = 0
    best_loss = 1e8
    epochs_trained = 0
    files_trained_in_current_epoch = 0

    if args.resume:
        # resume: ckpt-epoch-file_number-step:
        try:
            ckpt_suffix = args.resume.split('-')[1:]
            global_step = int(ckpt_suffix[-1])
            epochs_trained = int(ckpt_suffix[0])
            files_trained_in_current_epoch = int(ckpt_suffix[1]) + 1
            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info("  Continuing training from epoch %d", epochs_trained)
            logger.info("  Continuing training from global step %d", global_step)
            logger.info("  Will skip the first %d files in the first epoch", files_trained_in_current_epoch)
        except ValueError:
            logger.info("  Starting fine-tuning.")

    for epoch in tqdm(range(epochs_trained, args.epochs)):
        if os.path.isdir(args.passagefile):
            current_dataloaders = file_stream(tokenizer=tokenizer, data_mode=args.data_mode, passage_dir=args.passagefile,
                                            answer_dir=args.train_answerfile, query_dir=args.train_queryfile,
                                            distributed_training=distributed_training, num_workers=args.num_workers,
                                            padding_func=padding_func, batch_size=args.batch_size,
                                            sample_size=args.sample_size, sent_keep=args.ICT_sent_keep, hard_negative_rate=args.hard_negative_rate)
        else:
            current_dataloaders = train_dataloaders
        for num_file, (train_dataloader, train_sampler) in enumerate(current_dataloaders):
            if files_trained_in_current_epoch > 0:
                files_trained_in_current_epoch -= 1
                continue
            if distributed_training:
                train_sampler.set_epoch(epoch)

            for step, batch in enumerate(tqdm(train_dataloader)):
                query_bert.zero_grad()
                doc_bert.zero_grad()
                batch = tuple(t.to(device) for t in batch)
                query_tokens, passages_tokens, labels = batch
                with autocast():
                    query_embeds = query_bert(query_tokens)
                    passages_embeds = doc_bert(passages_tokens)
                    scores = torch.mm(query_embeds, passages_embeds.t())
                    loss = loss_fun(scores, labels).mean()
                scaler.scale(loss).backward()
                scaler.step(query_optimizer)
                query_scheduler.step()  # Update learning rate schedule

                if not args.share:
                    scaler.step(doc_optimizer)
                    doc_scheduler.step()  # Update learning rate schedule
                scaler.update()

                global_step += 1
                if global_step % args.print_iter == 0:
                    logger.info(f'Global step {global_step}:  Train loss: {loss.item():.3f}')
                if args.eval_iter > 0 and global_step % args.eval_iter == 0 and (rank == 0 or not distributed_training):
                    val_loss = evaluation(query_bert, doc_bert, loss_fun,  val_dataloader)
                    logger.info(f'Global step {global_step}:  val loss: {val_loss:.3f}')
                    if val_loss < best_loss:
                        best_loss = val_loss
                        save_model(query_bert, doc_bert, args.save_model)
        # if args.ckpt_dir and os.path.isdir(args.passagefile) and (rank == 0 or not distributed_training):
        if args.ckpt_dir and (rank == 0 or not distributed_training):
            output_dir = os.path.join(args.ckpt_dir, f'ckpt-{epoch}-{num_file}-{global_step}')
            os.makedirs(output_dir, exist_ok=True)
            save_model(query_bert, doc_bert, os.path.join(output_dir, 'ckpt'))
            logger.info(f'Save checkpoint to {output_dir}')
            torch.save(query_optimizer.state_dict(), os.path.join(output_dir, "query_optimizer.pt"))
            torch.save(query_scheduler.state_dict(), os.path.join(output_dir, "query_scheduler.pt"))
            if not args.share:
                torch.save(doc_optimizer.state_dict(), os.path.join(output_dir, "doc_optimizer.pt"))
                torch.save(doc_scheduler.state_dict(), os.path.join(output_dir, "doc_scheduler.pt"))
            logger.info("Saving optimizer and scheduler states to %s", output_dir)
        if args.eval_iter <= 0 and (rank == 0 or not distributed_training):
            save_model(query_bert, doc_bert, args.save_model+str(epoch))
    if args.eval_iter > 0 and (rank == 0 or not distributed_training):
        val_loss = evaluation(query_bert, doc_bert, loss_fun, val_dataloader)
        logger.info(f'Global step {global_step}:  val loss: {val_loss:.3f}')
        if val_loss < best_loss:
            best_loss = val_loss
            save_model(query_bert, doc_bert, args.save_model)
    logger.info(f'Training Done! Best model saved into {args.save_model}.')

def doc_embed():
    """Compute Passage Embeddings"""
    doc_bert = BertEmbed(bert, embedding_size=args.embed_size, padding_id=tokenizer.pad_token_id,
                         mode=args.mode, normalize=args.normalize)
    missing_keys, unexpected_keys = doc_bert.load_state_dict(torch.load(args.save_model + '.doc'), strict=False)
    logger.info(f'Load Passages Embedding model from {args.save_model}.doc')
    assert len(unexpected_keys) == 0
    assert len(missing_keys) == 0 or missing_keys == ["bert.embeddings.position_ids"]
    doc_bert.to(device)
    if args.num_gpus <=0:
        args.num_gpus = torch.cuda.device_count()

    if args.num_gpus > 1:
        doc_bert = torch.nn.DataParallel(doc_bert)
    doc_bert.eval()
    logger.info('Generate Passages Embeddings')
    passage_embedding_dict = {}

    with open(args.passagefile) as fp:
        for line in fp:
            items = json.loads(line.strip())
            if isinstance(items, dict): items = [items]
            for item in items:
                passage_embedding_dict[item['id']] = item['body']

    pids = list(passage_embedding_dict.keys())
    psgs = list(passage_embedding_dict.values())
    psg_data = SimpleDataset(psgs, transform=lambda x: tokenizer.encode(x, max_length=args.passage_len, truncation=True))
    num_psgs = len(psgs)
    psg_embeds = np.zeros((num_psgs, args.embed_size))

    padding_func = lambda x: padding_util(x, tokenizer.pad_token_id, args.passage_len)
    sampler = SequentialSampler(psg_data)
    psg_dataloader = DataLoader(psg_data, sampler=sampler, batch_size=args.batch_size, collate_fn=padding_func)
    count = 0

    with torch.no_grad():
        for i, batch in enumerate(tqdm(psg_dataloader, desc='Embedding Passages')):
            batch = batch.to(device)
            bz = batch.shape[0]
            embed = doc_bert(batch).data.cpu().numpy()
            psg_embeds[count:count+bz, :] = embed
            count += bz

    psg_df = pd.DataFrame(psg_embeds, index=pids)
    psg_df.to_csv(args.doc_embed_file)
    logger.info(f'Passage Embeddings saved into {args.doc_embed_file}.')

def generate_retrieval(retrieve=True):
    """Compute Query Embeddings and Retrieve passages"""
    try:
        query_df = pd.read_csv(args.query_embed_file, index_col=0)
        logger.info(f'Load Query Embeddings from {args.query_embed_file}.')
    except:
        missing_keys, unexpected_keys = query_bert.load_state_dict(torch.load(args.save_model + '.query'), strict=False)
        logger.info(f'Load Query Embedding model from {args.save_model}.query')
        assert len(unexpected_keys) == 0
        assert len(missing_keys) == 0 or missing_keys == ["bert.embeddings.position_ids"]
        query_bert.to(device)
        query_bert.eval()
        test_queries = json.load(open(args.test_queryfile))
        if isinstance(test_queries, list):
            test_queries = {q['query_id']: q['query'] for q in test_queries}
        qids = list(test_queries.keys())
        query_data = SimpleDataset(list(test_queries.values()), transform=tokenizer.encode)
        query_embeds = np.zeros((len(query_data), query_bert.embedding_size))
        sampler = SequentialSampler(query_data)
        padding_func = lambda x: padding_util(x, tokenizer.pad_token_id, args.query_len)
        query_dataloader = DataLoader(query_data, sampler=sampler, batch_size=args.batch_size, collate_fn=padding_func)
        if args.topk == 0: args.topk = len(psg_df)
        count = 0
        with torch.no_grad():
            for i, batch in enumerate(tqdm(query_dataloader, desc='Embedding Queries')):
                batch = batch.to(device)
                bz = batch.shape[0]
                embed = query_bert(batch).data.cpu().numpy()
                query_embeds[count:count + bz, :] = embed
                count += bz
        query_df = pd.DataFrame(query_embeds, index=qids)
        query_df.to_csv(args.query_embed_file)
        logger.info(f'Query Embeddings saved into {args.query_embed_file}')
    if retrieve:
        try:
            psg_df = pd.read_csv(args.doc_embed_file, index_col=0)
            logger.info(f'Load Passages Embeddings from {args.doc_embed_file}.')
        except:
            raise ValueError('Load Passage Embedding failed.')

        try:
            query_scores = psg_df.dot(query_df.to_numpy().transpose())
            results = {}
            for i, qid in enumerate(tqdm(query_scores.keys())):
                score = query_scores[qid].nlargest(args.topk)
                results[qid] = score.to_dict()

        except Exception as e:
            logger.info(e)
            # for very large files we will chunk the file for retrieval.
            results = {}
            for query_df in tqdm(pd.read_csv(args.query_embed_file, index_col=0, chunksize=1000)):
                query_scores = psg_df.dot(query_df.transpose())
                for i, qid in enumerate(tqdm(query_scores.keys())):
                    score = query_scores[qid].nlargest(args.topk)
                    results[qid] = score.to_dict()

        json.dump(results, open(args.retrieval_outputfile, 'w'))
        logger.info(f'Retrieval results saved into {args.retrieval_outputfile}')

if __name__== "__main__":
    if args.train:
        train()
    if args.doc_embed:
        doc_embed()
    if args.query_embed:
        generate_retrieval(args.generate_retrieval)


# pylint: skip-file

"""
 Copyright (c) 2018 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import os
import tempfile
import unittest

import caffe
import numpy as np

def python_net_file():
    with tempfile.NamedTemporaryFile(mode='w+', delete=False) as f:
        f.write("""name: 'sampling_extractor_layer_test_net'
                   force_backward: true
                   layer { name: "detections" type: "Input" top: "detections"
                           input_param { shape { dim: 1 dim: 1 dim: 2 dim: 11 } } }
                   layer { name: "anchor1" type: "Input" top: "anchor1"
                           input_param { shape { dim: 2 dim: 5 dim: 3 dim: 4 } } }
                   layer { name: "anchor2" type: "Input" top: "anchor2"
                           input_param { shape { dim: 2 dim: 5 dim: 3 dim: 4 } } }
                   layer { type: 'Python' name: 'sampling_extractor' bottom: 'detections'
                           bottom: 'anchor1' bottom: 'anchor2' top: 'embeddings' top: 'labels'
                           python_param { module: 'caffe.custom_layers.sampling_extractor_layer'
                                          layer: 'SamplingExtractorLayer'
                                          param_str: '{ "num_anchors": 2, "valid_action_ids": [0, 1], "num_steps": 7, "threshold": 0.0}' } }""")
        return f.name

@unittest.skipIf('Python' not in caffe.layer_type_list(), 'Caffe built without Python layer support')
class TestSamplingExtractorLayer(unittest.TestCase):
    def setUp(self):
        net_file = python_net_file()
        self.net = caffe.Net(net_file, caffe.TRAIN)
        os.remove(net_file)

        self.detections = np.array([[0, 1, 0, 0, 1, 1, 0, 0, 0, 2, 1],
                                    [1, 1, 0, 0, 1, 1, 1, 1, 0, 3, 2]], dtype=np.float32)

        anchor1 = np.random.normal(size=[2, 5, 3, 4])
        anchor1 /= np.sqrt(np.sum(np.square(anchor1), axis=1, keepdims=True))
        anchor2 = np.random.normal(size=[2, 5, 3, 4])
        anchor2 /= np.sqrt(np.sum(np.square(anchor2), axis=1, keepdims=True))
        self.anchors = [anchor1, anchor2]

    def test_forward(self):
        corner_embeddings = []
        for i in xrange(2):
            detection = self.detections[i]
            item = int(detection[0])
            anchor_id = int(detection[6])
            x_pos = int(detection[9])
            y_pos = int(detection[10])

            corner_embeddings.append(self.anchors[anchor_id][item, :, y_pos, x_pos])

        self.net.blobs['detections'].data[...] = self.detections
        self.net.blobs['anchor1'].data[...] = self.anchors[0]
        self.net.blobs['anchor2'].data[...] = self.anchors[1]

        net_outputs = self.net.forward()
        predicted_embeddings = net_outputs['embeddings']
        predicted_labels = net_outputs['labels']
        self.assertTupleEqual(predicted_embeddings.shape, (7, 5))
        self.assertTupleEqual(predicted_labels.shape, (7,))

        for i in xrange(7):
            self.assertEqual(int(predicted_labels[i]), 0)

            p = predicted_embeddings[i]
            norm = np.sqrt(np.sum(np.square(p)))
            self.assertAlmostEqual(norm, 1.0, places=5)


#!/usr/bin/env python3
"""A script to process gfa images with the SDSS-V FPS"""
import time
import tqdm
import fitsio
import sep
import argparse
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from pathlib import Path
from matplotlib import animation
from matplotlib.patches import Ellipse
from astropy.time import Time

from bin import sjd
from sdssobstools import sdss_paths

sns.set(style="darkgrid")

__author__ = "Dylan Gatlin"

camera_offsets = np.array([0.0, -50.2985, -41.05837499999998,
                           11.466562500000009, 52.425531914893625, 0.0])
camera_shifts = np.array([0., 0., +0.05, 0., -0.02])

def build_filt(obj_arr: np.ndarray):
    ecc = np.sqrt(obj_arr['a']**2 - obj_arr['b']**2) / obj_arr['a']
    # print(np.percentile(ecc, 75))
    # print(obj_arr["cpeak"], obj_arr["peak"].max())
    # print(ecc, np.mean(ecc), np.percentile(ecc, 95))
    filt = (
        # (obj_arr["flux"] > np.percentile(obj_arr["flux"], 80))
        ((obj_arr["a"] * 0.216) > 0.5)
        & (obj_arr["cpeak"] < 60000)
        # & (obj_arr["npix"] > 5)
        & (ecc < 0.7)
    )
    # print(obj_arr["a"] * 0.216, obj_arr["b"] * 0.216,
    # np.mean([obj_arr["a"], obj_arr["b"]]))
    return filt

def show_img(obj_arr: np.ndarray, obj_filt, data_sub: np.ndarray, plot_file=""):
    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
    m, s = np.mean(data_sub), np.std(data_sub)
    im = ax.imshow(data_sub, interpolation='nearest', cmap='gray',
                   vmin=m-s, vmax=m+s, origin='lower')

    # plot an ellipse for each object
    for i, obj in enumerate(obj_arr[obj_filt]):
        e = Ellipse(xy=(obj['x'], obj['y']),
                    width=6*obj['a'],
                    height=6*obj['b'],
                    angle=obj['theta'] * 180. / np.pi)
        e.set_facecolor('none')
        e.set_edgecolor('red')
        ax.add_artist(e)
    plt.show()
    if plot_file:
        fig.savefig(plot_file)

def get_img_path(mjd: int, cam_num: int, exp_num: int):
    """Resolves a path of a gfa image, if none can be found, it will return a
    path object 'Null'
    """
    p = (sdss_paths.gcam / f"{mjd}/proc-gimg-gfa{cam_num:.0f}n-"
         f"{exp_num:0>4.0f}.fits")
    try:
        p.exists()
    except OSError:
        time.sleep(1)
    if not p.exists():
        p = (sdss_paths.gcam / f"{mjd}/proc-gimg-gfa{cam_num:.0f}n-"
             f"{exp_num:0>4.0f}.fits.gz")
    if not p.exists():
        p = (sdss_paths.gcam / f"{mjd}/gimg-gfa{cam_num:.0f}n-"
             f"{exp_num:0>4.0f}.fits")
    if not p.exists():
        p = (sdss_paths.gcam / f"{mjd}/gimg-gfa{cam_num:.0f}n-"
             f"{exp_num:0>4.0f}.fits.gz")
    if not p.exists():
        p = Path("Null")
    return p

class GFASet:
    def __init__(self, verbose: bool = False, gfas=np.array([2, 3, 4, 5, 6])):
        self.paths = []
        self.filter = []
        self.isots = []
        self.fwhms = []
        self.focuses = []
        self.n_objs = []
        self.im_nums = []
        self.verbose = verbose
        self.ani_fig = None
        self.ani_ax = None
        self.gfas = gfas

    def add_index(self, paths_iter, im_num: int):
        paths = []
        filter = []
        focus = []
        isots = []
        fwhms = []
        n_objs = []
        for p in paths_iter:
            paths.append(p)
            if p.exists():
                filter.append(True)
                fwhm, n, objs = self.get_fwhm(p, verbose=self.verbose)
                hdr = fitsio.read_header(p, 1)
                focus.append(hdr["FOCUS"])
                isots.append(hdr["DATE-OBS"])
                fwhms.append(fwhm)
                n_objs.append(n)
            else:
                filter.append(False)
                focus.append(np.nan)
                fwhms.append(np.nan)
                isots.append("2000-01-01")
                n_objs.append(np.nan)
        self.im_nums.append(im_num)
        self.paths.append(paths)
        self.focuses.append(focus)
        self.filter.append(filter)
        self.fwhms.append(fwhms)
        self.n_objs.append(n_objs)
        self.isots.append(isots)

    def remove_first_index(self):
        self.paths.pop(0)
        self.filter.pop(0)
        self.isots.pop(0)
        self.fwhms.pop(0)
        self.focuses.pop(0)
        self.n_objs.pop(0)
        self.im_nums.pop(0)

    @staticmethod
    def get_fwhm(path: Path, verbose: bool = False):
        try:
            path.exists()
        except OSError:
            pass
        try:
            data = fitsio.read(path, 1)
        except OSError:
            # For cases when the file is actively being written
            time.sleep(1)
            data = fitsio.read(path, 1)
        bkg = sep.Background(data.astype(float))
        objs = sep.extract(data - bkg, 1.5, bkg.globalrms)
        filt = build_filt(objs)
        if verbose:
            print(f"Using {filt.sum()}/{filt.shape[0]} points")
        # fwhm = 0.216 * np.mean(2.0*np.sqrt(
            # 2.0*np.log(2)*(objs['a'][filt]**2 + objs['b'][filt]**2))
        # )
        fwhm = 0.216 * np.mean(
            2 * np.sqrt(np.log(2) * (objs['a'][filt]**2 + objs['b'][filt]**2))
        )
        if fwhm > 5:
            fwhm = np.nan
        return fwhm, filt.sum(), objs

    def sort(self):
        self.aim_nums = np.array(self.im_nums)
        self.apaths = np.array(self.paths)
        self.afocuses = np.array(self.focuses)
        self.afilter = np.array(self.filter)
        self.afwhms = np.array(self.fwhms)
        self.an_objs = np.array(self.n_objs)
        self.aisots = Time(self.isots)
        sorter = self.aim_nums.argsort()
        self.aim_nums = self.aim_nums[sorter]
        self.apaths = self.apaths[sorter]
        self.afocuses = self.afocuses[sorter]
        self.afilter = self.afilter[sorter]
        self.afwhms = self.afwhms[sorter]
        self.an_objs = self.an_objs[sorter]
        self.aisots = self.aisots[sorter]

    def print(self):
        if self.verbose:
            print("Measurements in arcseconds with a place scale of 0.216")
        print(f"{'Img #':<6} {'Focus':<5} {'GFA1':<6} {'GFA2':<6} {'GFA3':<6}"
              f" {'GFA4':<6}"
              f" {'GFA5':<6} {'GFA6':<6} {'N Objs':<6}")
        print('=' * 80)
        for i, (im, foc) in enumerate(zip(self.aim_nums, self.afocuses)):
            if np.isnan(foc[1]):
                foc = f"{'-':>5}"
            else:
                foc = f"{foc[1]:>5.0f}"
            if np.all(np.isnan(self.an_objs[i])):
                n_objs = f"{'-':>6}"
            else:
                n_objs = f"{np.nanmean(self.an_objs[i]):>6.0f}"
            print(f"{im:>6} {foc:>5} "
                  + " ".join([f"{f:>6.2f}"
                              if not np.isnan(f) else f"{'-':>6}"
                              for (j, f) in enumerate(self.afwhms[i])]
                             ) + f" {n_objs:6>}"
                  )

    @staticmethod
    def quadratic(x: float, a: float, b: float, c: float):
        return a * x**2 + x * b + c

    def plot(self, plot_file, fig=None, ax=None, dont_print=False):
        flat_foc = (self.afocuses + camera_offsets[self.gfas - 1]).flatten()
        flat_fwhms = self.afwhms.flatten()
        flat_nstars = self.an_objs.flatten()
        nan_filt = ~np.isnan(flat_fwhms)
        if flat_nstars[nan_filt].size == 0:
            if not dont_print:
                print("Nothing to plot")
            return
        weight_nstars = flat_nstars[nan_filt] / np.nanmax(flat_nstars)
        weight_times = (20 - (np.nanmax(self.aisots)
                              - self.aisots.flatten()[nan_filt]
                              ) * 24 * 60)
        weight_times[weight_times < 0] = 0
        weight_times = weight_times / 30
        weight = weight_nstars * weight_times
        weight = weight / np.nanmax(weight)
        a, b, c = np.polyfit(flat_foc[nan_filt],
                             flat_fwhms[nan_filt], deg=2,
                             w=weight)
        expected = self.quadratic(flat_foc[nan_filt], a, b, c)
        chi_squared = (expected - flat_fwhms[nan_filt])**2 / expected
        chi_rejects = chi_squared > np.percentile(chi_squared, 85)
        a, b, c = np.polyfit(flat_foc[nan_filt][~chi_rejects],
                             flat_fwhms[nan_filt][~chi_rejects],
                             deg=2, w=weight[~chi_rejects])
        focs = np.linspace(flat_foc[nan_filt].min(),
                           flat_foc[nan_filt].max(), 100)
        fit = -b / 2 / a
        fwhm = self.quadratic(fit, a, b, c)
        mum = "\u03BCm"
        if a < 0:  # Not a parabola with a minimum
            fit_found = False
            print("Optimal focus not found")
        else:
            fit_found = True
            if not dont_print:
                print(f'Optimal focus is at {fit:>4.0f}{mum} with FWHM'
                      f' {fwhm:>4.1f}"')
        if fig is None:
            for_ani = False
            fig, ax = plt.subplots(1, 1, figsize=(6, 4))
        else:
            for_ani = True

        if fit_found:
            ax.plot(focs, self.quadratic(focs, a, b, c),
                    alpha=0.8,
                    # label="Best Fit"
                    )
            ax.set_title(f"Best Focus is {-b / 2 / a:.0f}{mum} with FWHM"
                         f' {fwhm:.1f}"')
            ax.axvline(-b / 2 / a, c="r", linestyle="--", alpha=0.6)

        not_old = 10 - (np.nanmax(self.aisots) - self.aisots) > 0
        for i, cam in enumerate(self.gfas):
            ax.scatter(self.afocuses[not_old[:, i], i]
                       + camera_offsets[self.gfas - 1][i],
                       self.afwhms[not_old[:, i], i], s=6,
                       alpha=0.7,
                       label=f"{cam}")
        ax.scatter(flat_foc[nan_filt][chi_rejects],
                   flat_fwhms[nan_filt][chi_rejects], marker="x", alpha=0.7,
                   linewidth=0.5, c="k")
        ax.legend(ncol=6, loc=1)
        ax.set_xlabel("Focus ($\mu m$)")
        ax.set_ylabel("FWHM (arcseconds)")
        if not for_ani:
            plt.show()
            if plot_file:
                fig.savefig(plot_file, dpi=100)
        return fig

    def separate_plot(self, plot_file=None):
        fig, ax = plt.subplots(1, 1, figsize=(6, 4))
        xs = np.linspace(np.nanmin(self.afocuses), np.nanmax(self.afocuses))
        mum = "\u03BCm"
        for i, cam in enumerate(self.gfas):
            if np.all(np.isnan(self.afwhms[:, i])):
                continue
            a, b, c = np.polyfit(self.afocuses[:, i], self.afwhms[:, i], deg=2)
            best = -b / 2 / a
            print(f'GFA {i+1:<4.0f} {best:>6.2f}{mum}'
                  f' {self.quadratic(best, a, b, c):>6.2f}"')
            ax.scatter(self.afocuses[:, i], self.afwhms[:, i], alpha=0.6,
                       label=f"GFA {cam+1:.0f}")
            ax.plot(xs, self.quadratic(xs, a, b, c), alpha=0.8)
        ax.legend()
        ax.set_xlabel("Focus ($\mu m$)")
        ax.set_ylabel("FWHM (arcseconds)")
        plt.show()
        if plot_file:
            fig.savefig(plot_file, dpi=100)
        return fig

    def exp_num_plot(self, plot_file=None):
        fig, ax = plt.subplots(1, 1, figsize=(6, 4))
        for i, cam in enumerate(self.gfas):
            ax.plot(self.aim_nums, self.afwhms[:, i], linewidth=1,
                    alpha=0.8,
                    label=f"GFA {cam+1}")
            ax.scatter(self.aim_nums, self.afwhms[:, i], s=6,
                       alpha=0.8)
        ax.legend()
        ax.set_xlabel("Exposure Number")
        ax.set_ylabel("FWHM (arcseconds)")
        plt.show()
        if plot_file:
            fig.savefig(plot_file, dpi=100)
        return fig

    def init_continuous_plot(self, window="0-15"):
        low, high = window.split("-")
        n_images = int(high) - int(low)
        today = sjd.sjd()
        img_dir = sdss_paths.gcam / f"{today:.0f}/"
        latest = 0
        try:
            img_dir.exists()
        except OSError:
            pass
        for fil in img_dir.glob("gimg-gfa4n-*.fits"):
            if "snap" in fil.name:
                continue
            current_num = int(fil.name.split("gfa4n-")[-1].split(".fits")[0])
            latest = max(latest, current_num)
        for im_num in tqdm.tqdm(range(latest - n_images, latest + 1)):
            im_ps = []
            for n in self.gfas:
                p = get_img_path(today, n, im_num)
                im_ps.append(p)
            self.add_index(im_ps, im_num)
        self.sort()
        self.ani_fig, self.ani_ax = plt.subplots(1, 1, figsize=(6, 4))
        return self.plot(None, self.ani_fig, self.ani_ax,
                         dont_print=not self.verbose)

    def continuous_plot_update(self, i):
        today = sjd.sjd()
        im_num = self.im_nums[-1] + 1
        im_num_0 = im_num
        p = get_img_path(today, 4, im_num)
        current_range = self.im_nums[-1] - self.im_nums[0]
        if self.verbose:
            print(im_num, p.exists(), im_num_0)
            print(current_range)
        while p.exists() and (im_num - im_num_0 < current_range):
            im_ps = []
            for n in self.gfas:
                p = get_img_path(today, n, im_num)
                im_ps.append(p)
            self.add_index(im_ps, im_num)
            self.remove_first_index()
            im_num += 1
        self.sort()
        self.ani_ax.clear()
        if self.verbose:
            print(f"Plotting from {self.im_nums[0]:>3.0f}-{self.im_nums[-1]:>3.0f}"
                  ", ", end="")
        return self.plot(None, fig=self.ani_fig, ax=self.ani_ax,
                         dont_print=not self.verbose)

def parse_args():
    parser = argparse.ArgumentParser(description="Prints a table of FWHMs from"
                                     " GFAs, useful for focus sweeps")
    parser.add_argument("-c", "--continuous", action="store_true",
                        help="Update plot/fit continuously, overrides all other"
                        " plots")
    parser.add_argument("-m", "--mjd", default=sjd.sjd())
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Verbose debugging output")
    parser.add_argument("-f", "--file", nargs="+",
                        help="Specific files to process")
    parser.add_argument("-n", "--gfas", nargs="+", type=int, default=[2, 3, 4,
                                                                      5, 6])
    parser.add_argument("-w", "--window",
                        help="A dash separated beginning and end of the window."
                             " Always include a --master-field with your window."
                             " Ex: -w 15-50 -s 35")
    parser.add_argument("-i", "--ignore", nargs="+", type=int, default=[],
                        help="Image numbers to ignore in a window")
    parser.add_argument("-e", "--exp-num", action="store_true",
                        help="Plot fwhm relative to exposure number")
    parser.add_argument("-t", "--interval", default=15,
                        help="Gap between reading images when using -c")
    # parser.add_argument("-s", "--master-field", type=int,
    #                     help="The master field with all stars you want. Only"
    #                     " used in conjunction with --window. You need both"
    #                     " arguments.")
    parser.add_argument("-p", "--plot", action="store_true",
                        help="A plot to find best fit")
    parser.add_argument("-k", "--plot-file", help="Filename to save plot to."
                        " only works with --plot and --window together. If not"
                        " included, the plot will show up in a new window (or"
                        " maybe not depending on X11).")
    parser.add_argument("-s", "--separate", action="store_true",
                        help="Creates separate fits for each gfa in a plot")
    parser.add_argument("-d", "--plot-image", action="store_true",
                        help="Plot an the sky image with ellipses traced, only works with -f"
                        )
    args = parser.parse_args()

    args.gfas = np.array(args.gfas)

    return args

def main(args=None):
    if args is None:
        args = parse_args()
    gfas = GFASet(verbose=args.verbose, gfas=args.gfas)
    if args.file:
        print(f"{'File Name':<20} {'FWHM':<6} {'N Objects':<11} {'Mean Ecc':<6}")
        print('=' * 80)
        for file in args.file:
            if '*' not in file:
                file = Path(file)
                if not file.exists():
                    raise FileNotFoundError(f"Couldn't find {file.absolute()}")
                fwhm, n_objs, objs = gfas.get_fwhm(file)
                filt = build_filt(objs)
                ecc = (np.sqrt(objs['a'][filt]**2 - objs['b'][filt]**2)
                       / objs['a'][filt])
                print(f"{file.name:<20} {fwhm:>6.2f} {filt.sum():>5.0f}"
                      f"/{len(objs):<5.0f}  {np.mean(ecc):>6.2f}")
                if args.plot_image:
                    data = fitsio.read(file, 1)
                    bkg = sep.Background(data.astype(float))
                    show_img(objs, filt, data - bkg)
            else:
                file = Path(file)
                for p in file.parent.glob(file.name):
                    pass

    if args.continuous:
        anis = []
        print("Plotting continuously")
        if args.window is None:
            gfas.init_continuous_plot()
        else:
            gfas.init_continuous_plot(args.window)
        # I cannot explain why I must create a list I will never use to get an
        # animation to run. By rights, it should not be necessary. However,
        # it is apparently essential to the existence of an animation loop
        anis.append(animation.FuncAnimation(gfas.ani_fig,
                                            gfas.continuous_plot_update,
                                            interval=1000 * args.interval))
        plt.show()

    elif not args.window and not args.continuous:
        path_stem = sdss_paths.gcam / f"{args.mjd}"
        low = 1000
        high = 0
        for file in path_stem.glob("gimg-gfa5n-*.fits*"):
            print(file.name)
            ind = int(file.name.split("n-")[-1].rstrip(
                "-snap.fits.gz"))
            if ind < low:
                low = ind
            if ind > high:
                high = ind
        if high > low:
            args.window = f"{low:.0f}-{high:.0f}"

    if args.window:
        low, high = args.window.split("-")
        try:
            low = int(low)
            high = int(high)
        except ValueError:
            raise ValueError(f"Window range must be two ints with a - between"
                             f" them {args.window}")
        if args.verbose:
            print(f"Reading Images from {low} to {high} in MJD {args.mjd}")
        for im_num in tqdm.tqdm(range(low, high + 1)):
            if im_num in args.ignore:
                continue
            im_ps = []
            for n in args.gfas:
                p = get_img_path(args.mjd, n, im_num)
                im_ps.append(p)

            gfas.add_index(im_ps, im_num)
        gfas.sort()
        gfas.print()

        if args.separate:
            gfas.separate_plot(args.plot_file)
        elif args.exp_num:
            gfas.exp_num_plot(args.plot_file)
        elif args.plot:
            gfas.plot(args.plot_file)

    return 0

if __name__ == "__main__":
    main()


import copy
import os
import numpy as np
import torch
import torch.nn.functional as F
from common.buffers import OfflineBuffer
from utils.train_tools import soft_target_update, evaluate
from utils import log_tools
from common.networks import MLPSquashedReparamGaussianPolicy, CVAE, MLPQsaNet

class BEAR_Agent:
    """
    Implementation of Bootstrapping Error Accumulation Reduction (BEAR)
    https://arxiv.org/abs/1906.00949
    BEAR's MMD Loss's weight alpha_prime is tuned automatically by default.

    Actor Loss: alpha_prime * MMD Loss + -minQ(s,a)
    Critic Loss: Like BCQ
    Alpha_prime Loss: -(alpha_prime * (MMD Loss - threshold))
    """
    def __init__(self,
                 env,
                 data_buffer: OfflineBuffer,
                 policy_net: MLPSquashedReparamGaussianPolicy,  # actor
                 q_net1: MLPQsaNet,  # critic
                 q_net2: MLPQsaNet,
                 cvae_net: CVAE,
                 policy_lr=1e-4,
                 qf_lr=3e-4,
                 cvae_lr=3e-4,
                 gamma=0.99,
                 tau=0.05,

                 # BEAR
                 lmbda=0.75,  # used for double clipped double q-learning
                 mmd_sigma=20.0,  # the sigma used in mmd kernel
                 kernel_type='gaussian',  # the type of mmd kernel(gaussian or laplacian)
                 lagrange_thresh=0.05,  # the hyper-parameter used in automatic tuning alpha in cql loss
                 n_action_samples=100,  # the number of action samples to compute the best action when choose action
                 n_target_samples=10,  # the number of action samples to compute BCQ-like target value
                 n_mmd_action_samples=4,  # the number of action samples to compute MMD.
                 warmup_step=40000,  # do support matching with a warm start before policy(actor) train

                 max_train_step=1000000,
                 log_interval=1000,
                 eval_freq=5000,
                 train_id="bear_hopper-medium-v2_test",
                 resume=False,  # if True, train from last checkpoint
                 device='cpu',
                 ):

        self.env = env
        self.data_buffer = data_buffer

        self.device = torch.device(device)

        # the network and optimizers
        self.policy_net = policy_net.to(self.device)
        self.q_net1 = q_net1.to(self.device)
        self.q_net2 = q_net2.to(self.device)
        self.target_q_net1 = copy.deepcopy(self.q_net1).to(self.device)
        self.target_q_net2 = copy.deepcopy(self.q_net2).to(self.device)
        self.cvae_net = cvae_net.to(self.device)
        self.policy_optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=policy_lr)
        self.q_optimizer1 = torch.optim.Adam(self.q_net1.parameters(), lr=qf_lr)
        self.q_optimizer2 = torch.optim.Adam(self.q_net2.parameters(), lr=qf_lr)
        self.cvae_optimizer = torch.optim.Adam(self.cvae_net.parameters(), lr=cvae_lr)

        self.gamma = gamma
        self.tau = tau

        self.max_train_step = max_train_step
        self.eval_freq = eval_freq
        self.train_step = 0

        self.resume = resume  # whether load checkpoint start train from last time

        # BEAR
        self.lmbda = lmbda
        self.mmd_sigma = mmd_sigma
        self.kernel_type = kernel_type
        self.lagrange_thresh = lagrange_thresh
        self.n_action_samples = n_action_samples
        self.n_target_samples = n_target_samples
        self.n_mmd_action_samples = n_mmd_action_samples
        self.warmup_step = warmup_step

        # mmd loss's temperature
        self.log_alpha_prime = torch.zeros(1, requires_grad=True, device=self.device)
        self.alpha_prime_optimizer = torch.optim.Adam([self.log_alpha_prime], lr=1e-3)

        # log dir and interval
        self.log_interval = log_interval
        self.result_dir = os.path.join(log_tools.ROOT_DIR, "run/results", train_id)
        log_tools.make_dir(self.result_dir)
        self.checkpoint_path = os.path.join(self.result_dir, "checkpoint.pth")
        self.tensorboard_writer = log_tools.TensorboardLogger(self.result_dir)

    def choose_action(self, obs, eval=False):
        with torch.no_grad():
            obs = torch.FloatTensor(obs).reshape(1, -1).repeat(self.n_action_samples, 1).to(self.device)
            action, _, _ = self.policy_net(obs)
            q1 = self.q_net1(obs, action)
            ind = q1.argmax(dim=0)
        return action[ind].cpu().numpy().flatten()

    def mmd_loss_laplacian(self, samples1, samples2, sigma=0.2):
        """MMD constraint with Laplacian kernel for support matching"""
        # sigma is set to 10.0 for hopper, cheetah and 20 for walker/ant
        diff_x_x = samples1.unsqueeze(2) - samples1.unsqueeze(1)  # B x N x N x d
        diff_x_x = torch.mean((-(diff_x_x.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_x_y = samples1.unsqueeze(2) - samples2.unsqueeze(1)
        diff_x_y = torch.mean((-(diff_x_y.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_y_y = samples2.unsqueeze(2) - samples2.unsqueeze(1)  # B x N x N x d
        diff_y_y = torch.mean((-(diff_y_y.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        overall_loss = (diff_x_x + diff_y_y - 2.0 * diff_x_y + 1e-6).sqrt()
        return overall_loss

    def mmd_loss_gaussian(self, samples1, samples2, sigma=0.2):
        """MMD constraint with Gaussian Kernel support matching"""
        # sigma is set to 10.0 for hopper, cheetah and 20 for walker/ant
        diff_x_x = samples1.unsqueeze(2) - samples1.unsqueeze(1)  # B x N x N x d
        diff_x_x = torch.mean((-(diff_x_x.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_x_y = samples1.unsqueeze(2) - samples2.unsqueeze(1)
        diff_x_y = torch.mean((-(diff_x_y.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_y_y = samples2.unsqueeze(2) - samples2.unsqueeze(1)  # B x N x N x d
        diff_y_y = torch.mean((-(diff_y_y.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        overall_loss = (diff_x_x + diff_y_y - 2.0 * diff_x_y + 1e-6).sqrt()
        return overall_loss

    def train(self):
        # Sample
        batch = self.data_buffer.sample()
        obs = batch["obs"].to(self.device)
        acts = batch["acts"].to(self.device)
        rews = batch["rews"].to(self.device)
        next_obs = batch["next_obs"].to(self.device)
        done = batch["done"].to(self.device)

        """
        Train the Behaviour cloning policy to be able to take more than 1 sample for MMD.
        Conditional VAE is used as Behaviour cloning policy in BEAR.
        """
        recon_action, mu, log_std = self.cvae_net(obs, acts)
        cvae_loss = self.cvae_net.loss_function(recon_action, acts, mu, log_std)

        self.cvae_optimizer.zero_grad()
        cvae_loss.backward()
        self.cvae_optimizer.step()

        """
        Critic Training
        """
        with torch.no_grad():
            # generate 10 actions for every next_obs(Same as BCQ)
            next_obs = torch.repeat_interleave(next_obs, repeats=self.n_target_samples, dim=0).to(self.device)
            # compute target Q value of generated action
            target_q1 = self.target_q_net1(next_obs, self.policy_net(next_obs)[0])
            target_q2 = self.target_q_net2(next_obs, self.policy_net(next_obs)[0])
            # soft clipped double q-learning
            target_q = self.lmbda * torch.min(target_q1, target_q2) + (1. - self.lmbda) * torch.max(target_q1, target_q2)
            # take max over each action sampled from the generation and perturbation model
            target_q = target_q.reshape(obs.shape[0], self.n_target_samples, 1).max(1)[0].squeeze(1)
            target_q = rews + self.gamma * (1. - done) * target_q

        # compute current Q
        current_q1 = self.q_net1(obs, acts).squeeze(1)
        current_q2 = self.q_net2(obs, acts).squeeze(1)
        # compute critic loss
        critic_loss1 = F.mse_loss(current_q1, target_q)
        critic_loss2 = F.mse_loss(current_q2, target_q)

        self.q_optimizer1.zero_grad()
        critic_loss1.backward()
        self.q_optimizer1.step()

        self.q_optimizer2.zero_grad()
        critic_loss2.backward()
        self.q_optimizer2.step()

        # MMD Loss
        # sample actions from dataset and current policy(B x N x D)
        raw_sampled_actions = self.cvae_net.decode_multiple_without_squash(obs, decode_num=self.n_mmd_action_samples,
                                                                           z_device=self.device)
        raw_actor_actions = self.policy_net.sample_multiple_without_squash(obs, sample_num=self.n_mmd_action_samples)
        if self.kernel_type == 'gaussian':
            mmd_loss = self.mmd_loss_gaussian(raw_sampled_actions, raw_actor_actions, sigma=self.mmd_sigma)
        else:
            mmd_loss = self.mmd_loss_laplacian(raw_sampled_actions, raw_actor_actions, sigma=self.mmd_sigma)

        """
        Alpha prime training(lagrangian parameter update for MMD loss weight)
        """
        alpha_prime_loss = -(self.log_alpha_prime.exp() * (mmd_loss - self.lagrange_thresh)).mean()
        self.alpha_prime_optimizer.zero_grad()
        alpha_prime_loss.backward(retain_graph=True)
        self.alpha_prime_optimizer.step()

        self.log_alpha_prime.data.clamp_(min=-5.0, max=10.0)  # clip for stability

        """
        Actor Training
        Actor Loss = alpha_prime * MMD Loss + -minQ(s,a)
        """
        a, log_prob, _ = self.policy_net(obs)
        min_q = torch.min(self.q_net1(obs, a), self.q_net2(obs, a)).squeeze(1)
        # policy_loss = (self.alpha * log_prob - min_q).mean()  # SAC Type
        policy_loss = - (min_q.mean())

        # BEAR Actor Loss
        actor_loss = (self.log_alpha_prime.exp() * mmd_loss).mean()
        if self.train_step > self.warmup_step:
            actor_loss = policy_loss + actor_loss
        self.policy_optimizer.zero_grad()
        actor_loss.backward()  # the mmd_loss will backward again in alpha_prime_loss.
        self.policy_optimizer.step()

        soft_target_update(self.q_net1, self.target_q_net1, tau=self.tau)
        soft_target_update(self.q_net2, self.target_q_net2, tau=self.tau)

        self.train_step += 1

        return critic_loss1.cpu().item(), critic_loss2.cpu().item(), policy_loss.cpu().item(), alpha_prime_loss.cpu().item()

    def learn(self):
        if self.resume:
            self.load_agent_checkpoint()
        else:
            # delete tensorboard log file
            log_tools.del_all_files_in_dir(self.result_dir)

        while self.train_step < (int(self.max_train_step)):
            # train
            q_loss1, q_loss2, policy_loss, alpha_prime_loss = self.train()

            if self.train_step % self.eval_freq == 0:
                avg_reward, avg_length = evaluate(agent=self, episode_num=5)
                self.tensorboard_writer.log_eval_data({"eval_episode_length": avg_length,
                                                       "eval_episode_reward": avg_reward}, self.train_step)

            if self.train_step % self.log_interval == 0:
                self.store_agent_checkpoint()
                self.tensorboard_writer.log_train_data({"q_loss_1": q_loss1,
                                                        "q_loss_2": q_loss2,
                                                        "policy_loss": policy_loss,
                                                        "alpha_prime_loss": alpha_prime_loss}, self.train_step)

    def store_agent_checkpoint(self):
        checkpoint = {
            "q_net1": self.q_net1.state_dict(),
            "q_net2": self.q_net2.state_dict(),
            "policy_net": self.policy_net.state_dict(),
            "q_optimizer1": self.q_optimizer1.state_dict(),
            "q_optimizer2": self.q_optimizer2.state_dict(),
            "policy_optimizer": self.policy_optimizer.state_dict(),
            "log_alpha_prime": self.log_alpha_prime,
            "alpha_prime_optimizer": self.alpha_prime_optimizer.state_dict(),
            "train_step": self.train_step,
        }

        torch.save(checkpoint, self.checkpoint_path)

    def load_agent_checkpoint(self):
        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)  # can load gpu's data on cpu machine
        self.q_net1.load_state_dict(checkpoint["q_net1"])
        self.q_net2.load_state_dict(checkpoint["q_net2"])
        self.policy_net.load_state_dict(checkpoint["policy_net"])
        self.q_optimizer1.load_state_dict(checkpoint["q_optimizer1"])
        self.q_optimizer2.load_state_dict(checkpoint["q_optimizer2"])
        self.policy_optimizer.load_state_dict(checkpoint["policy_optimizer"])
        self.log_alpha_prime = checkpoint["log_alpha_prime"]
        self.alpha_prime_optimizer.load_state_dict(checkpoint["alpha_prime_optimizer"])
        self.train_step = checkpoint["train_step"]

        print("load checkpoint from \"" + self.checkpoint_path +
              "\" at " + str(self.train_step) + " time step")


import numpy as np
import torch as t
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from .decoder import Decoder
from .encoder import Encoder

from selfModules.embedding import Embedding

from utils.functional import kld_coef, parameters_allocation_check, fold
from beam_search import Beam

class RVAE(nn.Module):
    def __init__(self, params, params_2):
        super(RVAE, self).__init__()

        self.params = params
        self.params_2 = params_2  # Encoder-2 parameters

        self.embedding = Embedding(self.params, '')
        self.embedding_2 = Embedding(self.params_2, '', True)

        self.encoder = Encoder(self.params)
        self.encoder_2 = Encoder(self.params_2)

        #
        self.context_to_mu = nn.Linear(self.params.encoder_rnn_size * 2, self.params.latent_variable_size)
        self.context_to_logvar = nn.Linear(self.params.encoder_rnn_size * 2, self.params.latent_variable_size)

        # self.encoder_3 = Encoder(self.params)
        self.decoder = Decoder(self.params_2)  # change this to params_2

    def forward(self, drop_prob,
                encoder_word_input=None, encoder_character_input=None,
                encoder_word_input_2=None, encoder_character_input_2=None,
                decoder_word_input_2=None, decoder_character_input_2=None,
                z=None, initial_state=None):

        assert parameters_allocation_check(self)
        use_cuda = self.embedding.word_embed.weight.is_cuda

        assert z is None and fold(lambda acc, parameter: acc and parameter is not None,
                                  [encoder_word_input, encoder_character_input, decoder_word_input_2],
                                  True) \
               or (z is not None and decoder_word_input_2 is not None)

        if z is None:
            [batch_size, _] = encoder_word_input.size()

            encoder_input = self.embedding(encoder_word_input, encoder_character_input)

            [batch_size_2, _] = encoder_word_input_2.size()

            encoder_input_2 = self.embedding_2(encoder_word_input_2, encoder_character_input_2)

            context, h_0, c_0 = self.encoder(encoder_input, None)

            State = (h_0, c_0)
            context_2, _, _ = self.encoder_2(encoder_input_2, State)

            mu = self.context_to_mu(context_2)
            logvar = self.context_to_logvar(context_2)
            std = t.exp(0.5 * logvar)

            z = Variable(t.randn([batch_size, self.params.latent_variable_size]))
            if use_cuda:
                z = z.cuda()

            z = z * std + mu

            kld = (-0.5 * t.sum(logvar - t.pow(mu, 2) - t.exp(logvar) + 1, 1)).mean().squeeze()

            # encoder_input = self.embedding(encoder_word_input, encoder_character_input)
            # _ , h_0 , c_0 = self.encoder_3(encoder_input, None)
            initial_state = State

        else:
            kld = None
            mu = None
            std = None

        decoder_input_2 = self.embedding_2.word_embed(decoder_word_input_2)
        out, final_state = self.decoder(decoder_input_2, z, drop_prob, initial_state)

        return out, final_state, kld, mu, std

    def learnable_parameters(self):
        return [p for p in self.parameters() if p.requires_grad]

    def trainer(self, optimizer, batch_loader, batch_loader_2):
        def train(i, batch_size, use_cuda, dropout, start_index):
            input = batch_loader.next_batch(batch_size, 'train', start_index)
            input = [Variable(t.from_numpy(var)) for var in input]
            input = [var.long() for var in input]
            input = [var.cuda() if use_cuda else var for var in input]
            [encoder_word_input, encoder_character_input, decoder_word_input, decoder_character_input, target] = input

            input_2 = batch_loader_2.next_batch(batch_size, 'train', start_index)
            input_2 = [Variable(t.from_numpy(var)) for var in input_2]
            input_2 = [var.long() for var in input_2]
            input_2 = [var.cuda() if use_cuda else var for var in input_2]
            [encoder_word_input_2, encoder_character_input_2, decoder_word_input_2, decoder_character_input_2,
             target] = input_2

            logits, _, kld, _, _ = self(dropout,
                                        encoder_word_input, encoder_character_input,
                                        encoder_word_input_2, encoder_character_input_2,
                                        decoder_word_input_2, decoder_character_input_2,
                                        z=None)

            # logits = logits.view(-1, self.params.word_vocab_size)
            logits = logits.view(-1, self.params_2.word_vocab_size)
            target = target.view(-1)
            cross_entropy = F.cross_entropy(logits, target)
            loss = 79 * cross_entropy + kld_coef(i) * kld
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            return cross_entropy, kld, kld_coef(i)

        return train

    def validater(self, batch_loader, batch_loader_2):
        def validate(batch_size, use_cuda, start_index):
            input = batch_loader.next_batch(batch_size, 'valid', start_index)
            input = [Variable(t.from_numpy(var)) for var in input]
            input = [var.long() for var in input]
            input = [var.cuda() if use_cuda else var for var in input]

            [encoder_word_input, encoder_character_input, decoder_word_input, decoder_character_input, target] = input

            input_2 = batch_loader_2.next_batch(batch_size, 'valid', start_index)
            input_2 = [Variable(t.from_numpy(var)) for var in input_2]
            input_2 = [var.long() for var in input_2]
            input_2 = [var.cuda() if use_cuda else var for var in input_2]
            [encoder_word_input_2, encoder_character_input_2, decoder_word_input_2, decoder_character_input_2,
             target] = input_2

            logits, _, kld, _, _ = self(0.,
                                        encoder_word_input, encoder_character_input,
                                        encoder_word_input_2, encoder_character_input_2,
                                        decoder_word_input_2, decoder_character_input_2,
                                        z=None)
            logits = logits.view(-1, self.params_2.word_vocab_size)
            target = target.view(-1)
            cross_entropy = F.cross_entropy(logits, target)

            return cross_entropy, kld

        return validate

    def sample(self, batch_loader, seq_len, seed, use_cuda, State):
        if use_cuda:
            seed = seed.cuda()

        decoder_word_input_np, decoder_character_input_np = batch_loader.go_input(1)

        decoder_word_input = Variable(t.from_numpy(decoder_word_input_np).long())
        decoder_character_input = Variable(t.from_numpy(decoder_character_input_np).long())

        if use_cuda:
            decoder_word_input, decoder_character_input = decoder_word_input.cuda(), decoder_character_input.cuda()

        result = ''

        initial_state = State

        for i in range(seq_len):
            logits, initial_state, _, _, _ = self(0., None, None,
                                                  None, None,
                                                  decoder_word_input, decoder_character_input,
                                                  seed, initial_state)

            # forward(self, drop_prob,
            #           encoder_word_input=None, encoder_character_input=None,
            #           encoder_word_input_2=None, encoder_character_input_2=None,
            #           decoder_word_input_2=None, decoder_character_input_2=None,
            #           z=None, initial_state=None):

            # logits = logits.view(-1, self.params.word_vocab_size)
            # logits = logits.view(-1, self.params.word_vocab_size)
            logits = logits.view(-1, self.params_2.word_vocab_size)
            # print '---------------------------------------'
            # print 'Printing logits'
            # print logits
            # print '------------------------------------------'

            prediction = F.softmax(logits)

            word = batch_loader.sample_word_from_distribution(prediction.data.cpu().numpy()[-1])

            if word == batch_loader.end_token:
                break

            result += ' ' + word

            decoder_word_input_np = np.array([[batch_loader.word_to_idx[word]]])
            decoder_character_input_np = np.array([[batch_loader.encode_characters(word)]])

            decoder_word_input = Variable(t.from_numpy(decoder_word_input_np).long())
            decoder_character_input = Variable(t.from_numpy(decoder_character_input_np).long())

            if use_cuda:
                decoder_word_input, decoder_character_input = decoder_word_input.cuda(), decoder_character_input.cuda()

        return result

    def sampler(self, batch_loader, batch_loader_2, seq_len, seed, use_cuda, i, beam_size, n_best):
        input = batch_loader.next_batch(1, 'valid', i)
        input = [Variable(t.from_numpy(var)) for var in input]
        input = [var.long() for var in input]
        input = [var.cuda() if use_cuda else var for var in input]
        [encoder_word_input, encoder_character_input, decoder_word_input, decoder_character_input, target] = input

        encoder_input = self.embedding(encoder_word_input, encoder_character_input)

        _, h0, c0 = self.encoder(encoder_input, None)
        State = (h0, c0)

        results, scores = self.sample_beam(batch_loader_2, seq_len, seed, use_cuda, State, beam_size, n_best)

        return results, scores

    def sample_beam(self, batch_loader, seq_len, seed, use_cuda, State, beam_size, n_best):
        # seed = Variable(t.from_numpy(seed).float())
        if use_cuda:
            seed = seed.cuda()

        decoder_word_input_np, decoder_character_input_np = batch_loader.go_input(1)

        decoder_word_input = Variable(t.from_numpy(decoder_word_input_np).long())
        decoder_character_input = Variable(t.from_numpy(decoder_character_input_np).long())

        if use_cuda:
            decoder_word_input, decoder_character_input = decoder_word_input.cuda(), decoder_character_input.cuda()

        dec_states = State

        dec_states = [
            dec_states[0].repeat(1, beam_size, 1),
            dec_states[1].repeat(1, beam_size, 1)
        ]

        drop_prob = 0.0
        beam_size = beam_size
        batch_size = 1

        beam = [Beam(beam_size, batch_loader, cuda=True) for k in range(batch_size)]

        batch_idx = list(range(batch_size))
        remaining_sents = batch_size

        for i in range(seq_len):

            input = t.stack(
                [b.get_current_state() for b in beam if not b.done]
            ).t().contiguous().view(1, -1)

            trg_emb = self.embedding_2.word_embed(Variable(input).transpose(1, 0))

            # print trg_emb.size()
            # print seed.size()

            trg_h, dec_states = self.decoder.only_decoder_beam(trg_emb, seed, drop_prob, dec_states)

            dec_out = trg_h.squeeze(1)

            # print "dec_out:", dec_out.size()

            out = F.softmax(self.decoder.fc(dec_out)).unsqueeze(0)

            word_lk = out.view(
                beam_size,
                remaining_sents,
                -1
            ).transpose(0, 1).contiguous()

            active = []
            for b in range(batch_size):
                if beam[b].done:
                    continue

                idx = batch_idx[b]
                if not beam[b].advance(word_lk.data[idx]):
                    active += [b]

                for dec_state in dec_states:  # iterate over h, c
                    # layers x beam*sent x dim
                    sent_states = dec_state.view(
                        -1, beam_size, remaining_sents, dec_state.size(2)
                    )[:, :, idx]
                    sent_states.data.copy_(
                        sent_states.data.index_select(
                            1,
                            beam[b].get_current_origin()
                        )
                    )

            if not active:
                break

            active_idx = t.cuda.LongTensor([batch_idx[k] for k in active])
            batch_idx = {beam: idx for idx, beam in enumerate(active)}

            def update_active(t):
                view = t.data.view(
                    -1, remaining_sents,
                    self.params.decoder_rnn_size
                )
                new_size = list(t.size())
                new_size[-2] = new_size[-2] * len(active_idx) \
                               // remaining_sents
                return Variable(view.index_select(
                    1, active_idx
                ).view(*new_size))

            dec_states = (
                update_active(dec_states[0]),
                update_active(dec_states[1])
            )
            dec_out = update_active(dec_out)

            remaining_sents = len(active)

        allHyp, allScores = [], []

        for b in range(batch_size):
            scores, ks = beam[b].sort_best()
            allScores += [scores[:n_best]]
            hyps = zip(*[beam[b].get_hyp(k) for k in ks[:n_best]])
            allHyp += [hyps]

        return allHyp, allScores


# attack.py -- generate audio adversarial examples
#
# Copyright (C) 2018, Hiromu Yakura <hiromu1996@gmail.com>.
# Copyright (C) 2017, Nicholas Carlini <nicholas@carlini.com>.
#
# This program is licenced under the BSD 2-Clause licence,
# contained in the LICENCE file in this directory.

from __future__ import print_function

import argparse
import math
import os
import random
import shutil
import string
import struct
import sys
import tempfile
import time

import numpy as np
import tensorflow as tf
import scipy.io.wavfile as wav

# Okay, so this is ugly. We don't want DeepSpeech to crash.
# So we're just going to monkeypatch TF and make some things a no-op.
# Sue me.
sys.path.append(os.path.join(os.path.dirname(__file__), 'DeepSpeech'))

tf.load_op_library = lambda x: x
tmp = os.path.exists
os.path.exists = lambda x: True
class Wrapper:
    def __init__(self, d):
        self.d = d
    def __getattr__(self, x):
        return self.d[x]
class HereBeDragons:
    d = {}
    FLAGS = Wrapper(d)
    def __getattr__(self, x):
        return self.do_define
    def do_define(self, k, v, *x):
        self.d[k] = v
tf.app.flags = HereBeDragons()
import DeepSpeech
os.path.exists = tmp

# More monkey-patching, to stop the training coordinator setup
DeepSpeech.TrainingCoordinator.__init__ = lambda x: None
DeepSpeech.TrainingCoordinator.start = lambda x: None

from util.text import ctc_label_dense_to_sparse
from tf_logits import get_logits
from weight_decay_optimizers import AdamWOptimizer

# These are the tokens that we're allowed to use.
# The - token is special and corresponds to the epsilon value in CTC decoding, and can not occur in the phrase.
toks = " abcdefghijklmnopqrstuvwxyz'-"

# Sampling rate of the input files
Fs = 16000

class Attack:
    def __init__(self, sess, restore_path, audio, impulse, phrase, freq_min, freq_max, batch_size, learning_rate, weight_decay):
        """
        Set up the attack procedure.

        Here we create the TF graph that we're going to use to actually generate the adversarial examples.
        """

        assert len(audio.shape) == 1
        assert len(impulse.shape) == 2

        self.sess = sess
        self.phrase = phrase

        self.impulse_size = impulse.shape[0]
        self.batch_size = batch_size

        # Store arguments as constant tensors.
        original = tf.constant(audio.astype(np.float32))

        # Create impulse filters in a frequency domain
        conv_length = audio.shape[0] + impulse.shape[1] - 1
        nfft = 2 ** int(math.ceil(math.log(conv_length, 2)))
        imp_filters = tf.constant(np.fft.rfft(impulse, nfft).astype(np.complex64))

        # Change filters to apply dynamically
        self.imp_indices = tf.Variable(np.zeros((batch_size, ), dtype=np.int32), name='qq_filters')
        apply_filters = tf.gather(imp_filters, self.imp_indices)

        # Create all the variables necessary they are prefixed with qq_ just so that we know which ones are ours
        # so when we restore the session we don't clobber them.
        self.delta = tf.Variable(np.random.normal(0, np.sqrt(np.abs(audio).mean()), audio.shape).astype(np.float32), name='qq_delta')

        # Create a band pass filter to be applied to the perturbation.
        freq = np.fft.rfftfreq(audio.shape[0], 1.0 / Fs)
        bp_filter = ((freq_min < freq) & (freq < freq_max)).astype(np.int32)
        bp_filter = tf.constant(bp_filter.astype(np.complex64))

        # Apply the filter for the delta to simulate the real-world and create an adversarial example.
        self.delta_filtered = tf.spectral.irfft(tf.spectral.rfft(self.delta) * bp_filter)
        self.ae_input = original + tf.pad(self.delta_filtered, [[0, original.get_shape().as_list()[0] - self.delta_filtered.get_shape().as_list()[0]]])

        # Convolve the impulse responses to the input
        fft_length = tf.constant(np.array([nfft], dtype=np.int32))
        ae_frequency = tf.spectral.rfft(self.ae_input, fft_length=[nfft]) * apply_filters
        ae_convolved = tf.spectral.irfft(ae_frequency, fft_length=[nfft])[:, :conv_length]

        # Normalize the convolved audio
        max_audio = tf.reduce_max(tf.abs(ae_convolved), axis=1, keepdims=True)
        self.ae_transformed = ae_convolved / max_audio * tf.reduce_max(tf.abs(self.ae_input))

        # Add a tiny bit of noise to help make sure that we can clip our values to 16-bit integers and not break things.
        self.noise_ratio = tf.Variable(np.ones((1, ), dtype=np.float32), name='qq_noise_ratio')
        small_noise = tf.random_normal(self.ae_transformed.get_shape().as_list(), stddev=2 ** 14) * ([1] - self.noise_ratio)
        final_input = tf.clip_by_value(self.ae_transformed + small_noise, -2 ** 15, 2 ** 15 - 1)

        # Feed this final value to get the logits.
        lengths = tf.constant(np.array([(conv_length - 1) // 320] * batch_size, dtype=np.int32))
        self.logits = get_logits(final_input, lengths)

        # And finally restore the graph to make the classifier actually do something interesting.
        saver = tf.train.Saver([x for x in tf.global_variables() if 'qq' not in x.name])
        saver.restore(sess, restore_path)

        # Calculate CTC loss.
        target_phrase = tf.constant(np.array([list(phrase) for _ in range(batch_size)], dtype=np.int32))
        target_phrase_lengths = tf.constant(np.array([len(phrase) for _ in range(batch_size)], dtype=np.int32))
        target = ctc_label_dense_to_sparse(target_phrase, target_phrase_lengths, batch_size)
        self.ctcloss = tf.nn.ctc_loss(labels=tf.cast(target, tf.int32), inputs=self.logits, sequence_length=lengths)

        # Set up the Adam optimizer to perform gradient descent for us
        start_vars = set(x.name for x in tf.global_variables())
        optimizer = AdamWOptimizer(weight_decay, learning_rate)

        gradients = optimizer.compute_gradients(self.ctcloss, [self.delta])
        self.train = optimizer.apply_gradients(gradients, decay_var_list=[self.delta]) # tf.sign(grad)?

        end_vars = tf.global_variables()
        new_vars = [x for x in end_vars if x.name not in start_vars]

        sess.run(tf.variables_initializer(new_vars + [self.delta]))

        # Decoder from the logits, to see how we're doing
        self.decoded, _ = tf.nn.ctc_beam_search_decoder(self.logits, lengths, merge_repeated=False, beam_width=100)

    def attack(self, outdir, num_iterations=5000):
        sess = self.sess

        # Initialize all of the variables
        sess.run(tf.variables_initializer([self.delta, self.noise_ratio]))

        # Record loss of each impulse response
        imp_losses = np.full(self.impulse_size, 1e10)

        # Create misc variables
        prefix = ''.join([random.choice(string.lowercase) for _ in range(3)])
        time_last, time_start = time.time(), time.time()

        # We'll make a bunch of iterations of gradient descent here
        for itr in xrange(num_iterations + 1):
            indice = np.random.choice(self.impulse_size, self.batch_size, p=(imp_losses / imp_losses.sum()))

            # Actually do the optimization step
            decoded, logits, ctcloss, ae_transformed, ae_input, delta_filtered, _ = sess.run([self.decoded, self.logits, self.ctcloss, self.ae_transformed, self.ae_input, self.delta_filtered, self.train], {self.imp_indices: indice})
            imp_losses[indice] = ctcloss

            # Report progress
            print('Iter: %d, Elapsed Time: %.3f, Iter Time: %.3f\n\tLosses: %s\n\t Delta: %s' % \
                  (itr, time.time() - time_start, time.time() - time_last, ' '.join('% 6.2f' % x for x in ctcloss), np.array_str(delta_filtered, max_line_width=120)))
            time_last = time.time()

            # Print out some debug information every 5 iterations.
            if itr % 5 == 0:
                res = np.zeros(decoded[0].dense_shape) + len(toks) - 1
                for j in xrange(len(decoded[0].values)):
                    x, y = decoded[0].indices[j]
                    res[x, y] = decoded[0].values[j]

                # Here we print the strings that are recognized.
                res = [''.join(toks[int(x)] for x in y).replace('-', '') for y in res]
                print('Recognition:\n\t' + '\n\t'.join(res))

                # And here we print the argmax of the alignment.
                res_al = np.argmax(logits, axis=2).T
                res_al = [''.join(toks[int(x)] for x in y) for y in res_al]
                print('Alignment:\n\t' + '\n\t'.join(res_al))

                # Check if we've succeeded then we should record our progress and decrease the rescale constant.
                matched = filter(lambda index: res[index] == ''.join([toks[x] for x in self.phrase]), range(self.batch_size))
                if len(matched) > self.batch_size * 0.5:
                    # Get the current constant
                    ratio = sess.run(self.noise_ratio)
                    print('=> It: %d, Noise Ratio: %.3f' % (itr, 1.0 - ratio[0]))

                    # Update with the new noise
                    sess.run(self.noise_ratio.assign(ratio * 0.99))

                if itr % 100 == 0 or len(matched) > self.batch_size * 0.5:
                    # Just for debugging, save the adversarial example so we can see it if we want
                    wav.write(os.path.join(outdir, '%s-adv-%d.wav' % (prefix, itr)), Fs, np.array(np.clip(np.round(ae_input), -2 ** 15, 2 ** 15 - 1), dtype=np.int16))
                    wav.write(os.path.join(outdir, '%s-delta-%d.wav' % (prefix, itr)), Fs, np.array(np.clip(np.round(delta_filtered), -2 ** 15, 2 ** 15 - 1), dtype=np.int16))
                    for i in xrange(ae_transformed.shape[0]):
                        wav.write(os.path.join(outdir, '%s-conv-%d-%d.wav' % (prefix, itr, i)), Fs, np.array(np.clip(np.round(ae_transformed[i]), -2 ** 15, 2 ** 15 - 1), dtype=np.int16))

                    # Save also the logits
                    np.save(os.path.join(outdir, '%s-logit-%d.npy' % (prefix, itr)), logits[:, matched, :])

def main():
    """
    Do the attack here.

    This is all just boilerplate; nothing interesting
    happens in this method.

    For now we only support using CTC loss and only generating
    one adversarial example at a time.
    """
    parser = argparse.ArgumentParser(description=None)
    parser.add_argument('--in', type=str, dest='input', required=True,
                        help='Input audio .wav file, at {fs}Hz'.format(fs=Fs))
    parser.add_argument('--imp', type=str, dest='impulse', nargs='+', required=True,
                        help='Input impulse response .wav file, at {fs}Hz'.format(fs=Fs))
    parser.add_argument('--target', type=str, required=True,
                        help='Target transcriptions')
    parser.add_argument('--out', type=str, required=True,
                        help='Directory for saving intermediate files')
    parser.add_argument('--batch_size', type=int, required=False, default=20,
                        help='Batch size for generation')
    parser.add_argument('--freq_min', type=int, required=False, default=1000,
                        help='Lower limit of band pass filter for adversarial noise')
    parser.add_argument('--freq_max', type=int, required=False, default=4000,
                        help='Higher limit of band pass filter for adversarial noise')
    parser.add_argument('--lr', type=float, required=False, default=200,
                        help='Learning rate for optimization')
    parser.add_argument('--decay', type=float, required=False, default=0.001,
                        help='Weight decay for optimization')
    parser.add_argument('--iterations', type=int, required=False, default=1000,
                        help='Maximum number of iterations of gradient descent')
    parser.add_argument('--session', type=str, required=False,
                        default=os.path.join(os.path.dirname(__file__), 'models/session_dump'),
                        help='Path for the session file taken from DeepSpeech')
    args = parser.parse_args()

    print('Command line:', args)

    with tf.Session() as sess:
        # Load the inputs that we're given
        fs, audio = wav.read(args.input)
        assert fs == Fs
        print('Source dB:', 20 * np.log10(np.max(np.abs(audio))))

        irs = []
        for i in range(len(args.impulse)):
            fs, ir = wav.read(args.impulse[i])
            assert fs == Fs
            irs.append(ir)

        # Pad the impulse responses
        maxlen = max(map(len, irs))
        for i in range(len(irs)):
            irs[i] = np.concatenate((irs[i], np.zeros(maxlen - irs[i].shape[0], dtype=irs[i].dtype)))
        irs = np.array(irs)

        # Set up the attack class and run it
        attack = Attack(sess, args.session, audio, irs, [toks.index(x) for x in args.target], freq_min=args.freq_min, freq_max=args.freq_max,
                        batch_size=args.batch_size, learning_rate=args.lr, weight_decay=args.decay)
        attack.attack(outdir=args.out, num_iterations=args.iterations)

if __name__ == '__main__':
    main()


from __future__ import unicode_literals
from flask import Flask, jsonify, request, send_from_directory
import logging
from logging import FileHandler
import time
import data_helper
import keras
import numpy as np
import os
import codecs
import random
import defs
import sys

t = time.time()
app = Flask(__name__)
application = app
model = None
modelFile = './save_tmp.h5'
try:
  if len(sys.argv) > 1:
    modelFile = sys.argv[1]
  file_handler = FileHandler("lang_detection.log", "a")
  file_handler.setLevel(logging.WARNING)
  app.logger.addHandler(file_handler)
  model = keras.models.load_model(modelFile)
  if __name__ == '__main__':
    print("Took {}".format(time.time() - t))
except Exception as e:
  print e.message

@app.route('/', methods=['GET'])
def browse_default():
  try:
    return send_from_directory('ui', 'index.html')
  except Exception as e:
    print e.message
    return e.message

@app.route('/<path:path>', methods=['GET'])
def staticx(path):
   return send_from_directory('ui', path)

@app.route('/api/v1/lang/detection', methods=['GET'])
def detect_url():

  f_url = request.args.get('file')
  x = np.array([data_helper.turn_url_to_vector(f_url, file_vector_size=2048)])
  y = model.predict(x)
  result = model.predict_proba(x)
  return jsonify({"result": result[0].tolist()})

@app.route('/api/v1/lang/detection', methods=['POST'])
def detect_text():
  try:
    text = request.get_json()['content']
    x = np.array([data_helper.turn_text_to_vector(text, file_vector_size=2048)])
    y = model.predict(x)
    result = [100*a for a in model.predict_proba(x)] # tunr to percentage
    return jsonify({"result": result[0].tolist()})
  except Exception as e:
    print e.message
    return e.message

@app.route('/api/v1/lang/random-snippet', methods=['GET'])
def random_snippet():
  try:
    root = './data/stackoverflow-snippets/'
    lang = request.args.get('lang')
    if lang is None or len(lang) == 0:
      lang = defs.langs[int(random.uniform(0, len(defs.langs)))]
    root = root + lang + '/'
    print root
    files = os.listdir(root)
    fn = root + files[int(random.uniform(0, len(files)))]
    print fn
    with codecs.open(fn) as f:
      return f.read()
  except Exception as e:
    print e.message
    return e.message

if __name__ == '__main__':
  app.run(debug=True, use_reloader=False, host='0.0.0.0')