import os, sys
from time import time

import numpy as np
import pandas as pd
import cv2
import torch

import slowfast.utils.checkpoint as cu
import slowfast.utils.distributed as du
from slowfast.utils import logging
from slowfast.utils import misc
from slowfast.datasets import cv2_transform
from slowfast.models import model_builder
from slowfast.datasets.cv2_transform import scale
import math

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

import slowfast.utils.checkpoint as cu
import slowfast.utils.multiprocessing as mpu
from slowfast.config.defaults import get_config

from detector.detect import YOLO
import argparse

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="1"

def detector(object_predictor , image, backbone, cfg , display_height, display_width ):
    if backbone == 'yolo':
        boxes = object_predictor.detect_image(image)
        boxes = torch.as_tensor(boxes).float().cuda()
        return boxes
    else:
        outputs = object_predictor(image)
        fields = outputs["instances"]._fields
        pred_classes = fields["pred_classes"]
        selection_mask = pred_classes == 0
        # acquire person boxes
        pred_classes = pred_classes[selection_mask]
        pred_boxes = fields["pred_boxes"].tensor[selection_mask]
        scores = fields["scores"][selection_mask]
        boxes = cv2_transform.scale_boxes(cfg.DATA.TEST_CROP_SIZE,
                                            pred_boxes,
                                            display_height,
                                            display_width)
        return boxes

def demo(cfg, backbone):
    # Set random seed from configs.
    np.random.seed(cfg.RNG_SEED)
    torch.manual_seed(cfg.RNG_SEED)

    # Build the video model and print model statistics.
    model = model_builder.build_model(cfg)
    model.eval()
    misc.log_model_info(model)

   # Load a checkpoint to test if applicable.
    if cfg.TEST.CHECKPOINT_FILE_PATH != "":
        ckpt = cfg.TEST.CHECKPOINT_FILE_PATH
    elif cu.has_checkpoint(cfg.OUTPUT_DIR):
        ckpt = cu.get_last_checkpoint(cfg.OUTPUT_DIR)
    elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "":
        ckpt = cfg.TRAIN.CHECKPOINT_FILE_PATH
    else:
        raise NotImplementedError("Unknown way to load checkpoint.")

    cu.load_checkpoint(
        ckpt,
        model,
        cfg.NUM_GPUS > 1,
        None,
        inflation=False,
        convert_from_caffe2= "caffe2" in [cfg.TEST.CHECKPOINT_TYPE, cfg.TRAIN.CHECKPOINT_TYPE],
    )

    darknetlib_path = '/home/ubuntu/hanhbd/SlowFast/detector/libdarknet.so'
    config_path = '/home/ubuntu/hanhbd/SlowFast/detector/yolov4.cfg'
    meta_path = '/home/ubuntu/hanhbd/SlowFast/detector/coco.data'
    classes_path = '/home/ubuntu/hanhbd/SlowFast/detector/coco.names'
    weight_path = '/home/ubuntu/hanhbd/SlowFast/detector/yolov4.weights'

    if backbone == 'yolo':
        object_predictor =  YOLO.get_instance(darknetlib_path, config_path, meta_path, classes_path, weight_path)
    else:
        dtron2_cfg_file = cfg.DEMO.DETECTRON2_OBJECT_DETECTION_MODEL_CFG
        dtron2_cfg = get_cfg()
        dtron2_cfg.merge_from_file(model_zoo.get_config_file(dtron2_cfg_file))
        dtron2_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = .5
        dtron2_cfg.MODEL.WEIGHTS = cfg.DEMO.DETECTRON2_OBJECT_DETECTION_MODEL_WEIGHTS
        object_predictor = DefaultPredictor(dtron2_cfg)

    with open(cfg.DEMO.LABEL_FILE_PATH) as f:
        labels = f.read().split('\n')[:-1]
    palette = np.random.randint(64, 128, (len(labels), 3)).tolist()
    count_xxx = 0
    seq_len = cfg.DATA.NUM_FRAMES*cfg.DATA.SAMPLING_RATE

    frames = []
    org_frames = []
    mid_frame = None
    pred_labels = []
    draw_imgs = []

    cap  = cv2.VideoCapture(cfg.DEMO.DATA_SOURCE)
    was_read, frame = cap.read()
    display_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    display_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    fourcc = cv2.VideoWriter_fourcc(*'DIVX')
    videowriter = cv2.VideoWriter('./result/testset_fighting_05.avi',fourcc, fps, (display_width,display_height))

    while was_read :
        was_read, frame = cap.read()
        if not was_read:
            videowriter.release()
            break

        if len(frames) != seq_len:
            frame_processed = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_processed = scale(cfg.DATA.TEST_CROP_SIZE, frame_processed)
            frames.append(frame_processed)
            org_frames.append(frame)

        else:
            #predict all person box in all frame
            start = time()
            mid_frame = org_frames[seq_len//2 - 2]
            # just draw half of number frame because we will use slide = 1/2 length of sequence
            if cfg.DETECTION.ENABLE and len(draw_imgs) == 0:
                for idx in range(seq_len//2 - 1):
                    image = org_frames[idx]
                    boxes = detector(object_predictor , image, backbone, cfg , display_height, display_width )
                    # boxes = object_predictor.detect_image(img)
                    # boxes = torch.as_tensor(boxes).float().cuda()

                    boxes = torch.cat([torch.full((boxes.shape[0], 1), float(0)).cuda(), boxes], axis=1)
                    boxes = boxes.cpu().detach().numpy()
                    if backbone == 'yolo':
                        boxes = boxes[:, 1:]
                    else:
                        ratio = np.min(
                            [display_height, display_width]
                        ) / cfg.DATA.TEST_CROP_SIZE
                        boxes = boxes[:, 1:] * ratio

                    for box in boxes:
                        xmin, ymin, xmax, ymax = box
                        cv2.rectangle(image, (xmin, ymin), (xmax , ymax), (0, 255, 0), thickness=2)

                    draw_imgs.append(image)

            # detect box in mid frame
            if cfg.DETECTION.ENABLE:
                boxes = detector(object_predictor , mid_frame, backbone, cfg , display_height, display_width )
                boxes = torch.cat([torch.full((boxes.shape[0], 1), float(0)).cuda(), boxes], axis=1)

            inputs = torch.from_numpy(np.array(frames)).float()
            inputs = inputs / 255.0
            # Perform color normalization.
            inputs = inputs - torch.tensor(cfg.DATA.MEAN)
            inputs = inputs / torch.tensor(cfg.DATA.STD)
            # T H W C -> C T H W.
            inputs = inputs.permute(3, 0, 1, 2)

            # 1 C T H W.
            inputs = inputs.unsqueeze(0)

            # Sample frames for the fast pathway.
            index = torch.linspace(0, inputs.shape[2] - 1, cfg.DATA.NUM_FRAMES).long()
            fast_pathway = torch.index_select(inputs, 2, index)

            # Sample frames for the slow pathway.
            index = torch.linspace(0, fast_pathway.shape[2] - 1,
                                    fast_pathway.shape[2]//cfg.SLOWFAST.ALPHA).long()
            slow_pathway = torch.index_select(fast_pathway, 2, index)
            inputs = [slow_pathway, fast_pathway]

            # Transfer the data to the current GPU device.
            if isinstance(inputs, (list,)):
                for i in range(len(inputs)):
                    inputs[i] = inputs[i].cuda(non_blocking=True)
            else:
                inputs = inputs.cuda(non_blocking=True)

            # use a dummy variable to disable all computations below.
            if not len(boxes):
                preds = torch.tensor([])
            else:
                preds = model(inputs, boxes)

            # Gather all the predictions across all the devices to perform ensemble.
            if cfg.NUM_GPUS > 1:
                preds = du.all_gather(preds)[0]

            # post processing
            preds = preds.cpu().detach().numpy()
            pred_masks = preds > .1
            label_ids = [np.nonzero(pred_mask)[0] for pred_mask in pred_masks]
            pred_labels = [
                [labels[label_id] for label_id in perbox_label_ids]
                for perbox_label_ids in label_ids
            ]
            print(pred_labels)
            boxes = boxes.cpu().detach().numpy()
            if backbone == 'yolo':
                boxes = boxes[:, 1:]
            else:
                ratio = np.min(
                    [display_height, display_width]
                ) / cfg.DATA.TEST_CROP_SIZE
                boxes = boxes[:, 1:] * ratio

            # draw result on mid frame
            if pred_labels and boxes.any():
                for box, box_labels in zip(boxes.astype(int), pred_labels):
                    xmin, ymin, xmax, ymax = box
                    cv2.rectangle(mid_frame, (xmin, ymin), (xmax , ymax), (0, 255, 0), thickness=2)

                    label_origin = box[:2]
                    for label in box_labels:
                        label_origin[-1] -= 5
                        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, .5, 2)
                        cv2.rectangle(
                            mid_frame,
                            (label_origin[0], label_origin[1] + 5),
                            (label_origin[0] + label_width, label_origin[1] - label_height - 5),
                            palette[labels.index(label)], -1
                        )
                        cv2.putText(
                            mid_frame, label, tuple(label_origin),
                            cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1
                        )
                        label_origin[-1] -= label_height + 5

            # append mid frame to the draw array
            draw_imgs.append(mid_frame)

            # write image to videos
            for img_ in draw_imgs:
                videowriter.write(img_)
            print("time process", (time() - start) /64 )
            # clean the buffer of frames and org_frames with slide 1/2 seq_len
            # frames = frames[seq_len//2 - 1:]
            # org_frames = org_frames[seq_len//2 - 1:]

            frames = frames[1:]
            org_frames = org_frames[1:]
            draw_imgs = draw_imgs[-1:]

            count_xxx += 1

def parse_args():
    """
    Parse the following arguments for the video training and testing pipeline.
    Args:
        shard_id (int): shard id for the current machine. Starts from 0 to
            num_shards - 1. If single machine is used, then set shard id to 0.
        num_shards (int): number of shards using by the job.
        init_method (str): initialization method to launch the job with multiple
            devices. Options includes TCP or shared file-system for
            initialization. details can be find in
            https://pytorch.org/docs/stable/distributed.html#tcp-initialization
        cfg (str): path to the config file.
        opts (argument): provide addtional options from the command line, it
            overwrites the config loaded from file.
        """
    parser = argparse.ArgumentParser(
        description="Provide SlowFast video training, testing, and demo pipeline."
    )
    parser.add_argument(
        "--shard_id",
        help="The shard id of current node, Starts from 0 to num_shards - 1",
        default=0,
        type=int,
    )
    parser.add_argument(
        "--num_shards",
        help="Number of shards using by the job",
        default=1,
        type=int,
    )
    parser.add_argument(
        "--init_method",
        help="Initialization method, includes TCP or shared file-system",
        default="tcp://localhost:9999",
        type=str,
    )
    parser.add_argument(
        "--cfg",
        dest="cfg_file",
        help="Path to the config file",
        default="configs/Kinetics/SLOWFAST_4x16_R50.yaml",
        type=str,
    )
    parser.add_argument(
        "--backbone",
        help="help backbone",
        default="yolo",
        type=str,
    )
    parser.add_argument(
        "opts",
        help="See slowfast/config/defaults.py for all options",
        default=None,
        nargs=argparse.REMAINDER,
    )
    if len(sys.argv) == 1:
        parser.print_help()
    return parser.parse_args()

def load_config(args):
    """
    Given the arguemnts, load and initialize the configs.
    Args:
        args (argument): arguments includes `shard_id`, `num_shards`,
            `init_method`, `cfg_file`, and `opts`.
    """
    # Setup cfg.
    cfg = get_config()
    # Load config from cfg.
    if args.cfg_file is not None:
        cfg.merge_from_file(args.cfg_file)
    # Load config from command line, overwrite config from opts.
    if args.opts is not None:
        cfg.merge_from_list(args.opts)

    # Inherit parameters from args.
    if hasattr(args, "num_shards") and hasattr(args, "shard_id"):
        cfg.NUM_SHARDS = args.num_shards
        cfg.SHARD_ID = args.shard_id
    if hasattr(args, "rng_seed"):
        cfg.RNG_SEED = args.rng_seed
    if hasattr(args, "output_dir"):
        cfg.OUTPUT_DIR = args.output_dir

    # Create the checkpoint dir.
    cu.make_checkpoint_dir(cfg.OUTPUT_DIR)
    return cfg

def main():

    args = parse_args()
    cfg = load_config(args)

    if cfg.NUM_GPUS > 1:
        torch.multiprocessing.spawn(
            mpu.run,
            nprocs=cfg.NUM_GPUS,
            args=(
                cfg.NUM_GPUS,
                demo,
                args.init_method,
                cfg.SHARD_ID,
                cfg.NUM_SHARDS,
                cfg.DIST_BACKEND,
                cfg,
            ),
            daemon=False,
        )
    else:
        demo(cfg=cfg, backbone = args.backbone)

if __name__ == "__main__":
    torch.multiprocessing.set_start_method("forkserver")
    main()


import time
import cv2
import mss
import numpy as np
import keylogger
import pickle
from pynput import keyboard

state = {
    "space": False,
    "save": False,
    "reward": 1,
}

def rgb2gray(rgb):
    return np.dot(rgb[..., :3], [0.299, 0.587, 0.144])

def save(data):
    with open('train-data.pkl', 'wb') as output:
        pickle.dump(data, output, pickle.HIGHEST_PROTOCOL)

def previewScreen():

    with mss.mss() as sct:
        # Part of the screen to capture
        monitor = {'top': 200, 'left': 200, 'width': 600, 'height': 200}
        while 'Screen capturing':
            last_time = time.time()
            # Get raw pixels from the screen, save it to a Numpy array
            img = np.array(sct.grab(monitor))
            cv2.imshow('Preview', img)

            # Display the picture in grayscale
            # cv2.imshow('OpenCV/Numpy grayscale',
            # cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY))
            print('fps: {0}'.format(1 / (time.time()-last_time)))
            # Press "q" to quit
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                break

def recordScreen(sleep=0.01, max_stack=3):
    keylogger.keyListener(state)
    stack = 0
    imgstack = {
        "stack": [],
        "value": [1, 0]
    }
    train_data = []
    with mss.mss() as sct:
        # Part of the screen to capture
        monitor = {'top': 200, 'left': 200, 'width': 600, 'height': 200}

        while 'Screen capturing':
            last_time = time.time()
            time.sleep(sleep)
            # Get raw pixels from the screen, save it to a Numpy array
            img = np.array(sct.grab(monitor))
            gray = rgb2gray(img)
            if state["save"]:
                if stack == max_stack:
                    print(state["space"])
                    if state["space"]:
                        imgstack["value"] = [0, 1]
                    train_data.append(imgstack)
                    state["space"] = False
                    imgstack = {
                        "stack": [],
                        "value": [1, 0]
                    }
                    stack = 0

                gray = cv2.resize(gray, dsize=(120, 40),
                                  interpolation=cv2.INTER_CUBIC)
                print(np.average(gray))
                gray *= (255.0/gray.max())
                imgstack["stack"].append(gray)
                stack = stack + 1

            # Display the picture
            cv2.imshow('Preview', img)

            # Display the picture in grayscale
            # cv2.imshow('OpenCV/Numpy grayscale',
            # cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY))
            print('fps: {0}'.format(1 / (time.time()-last_time)))
            # Press "q" to quit
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                save(train_data)
                break

def predictScreen(model, max_stack=3):
    stack = 0
    imgstack = {
        "stack": [],
        "value": [1, 0]
    }
    keylogger.keyListener(state)
    handler = keylogger.init_controller()
    with mss.mss() as sct:
        # Part of the screen to capture
        monitor = {'top': 200, 'left': 200, 'width': 600, 'height': 200}

        while 'Screen capturing':
            last_time = time.time()
            # Get raw pixels from the screen, save it to a Numpy array
            img = np.array(sct.grab(monitor))
            gray = rgb2gray(img)
            if state["save"]:
                if stack == max_stack:
                    pred = model.predict_classes(
                        np.array([np.stack(imgstack["stack"], axis=2)]))
                    if pred[0] == 1:
                        print("jump")
                        keylogger.pressSpace(handler)
                    state["space"] = False
                    imgstack = {
                        "stack": [],
                        "value": [1, 0]
                    }
                    stack = 0

                gray = cv2.resize(gray, dsize=(120, 40),
                                  interpolation=cv2.INTER_CUBIC)
                gray *= (255.0/gray.max())
                imgstack["stack"].append(np.array(gray)/255)
                stack = stack + 1

            # Press "q" to quit
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                break


import argparse
import os
import random
import logging
import json
import math
from tqdm import tqdm, trange
import copy
import numpy as np
import torch
import torch.nn.functional as F
import time

from eval import eval
from dataloader import CheckerDevSupportEpisode, CheckerDevTestEpisode, TrainDataset, traincollate, read_data, DevDataset, DevSupportEpisode, DevTestEpisode, devcollate, checker_devcollate, test_checker_collate
from torch.utils.data import (DataLoader, RandomSampler, TensorDataset)
from transformers import (GPT2LMHeadModel, GPT2Tokenizer, GPT2Config, AdamW, WEIGHTS_NAME, CONFIG_NAME,
                          get_linear_schedule_with_warmup)
import pickle
from opt import get_args
from model import GPT2GenModel

def generate(model,tokenizer,batch_prompt_in,encoded_raw_sent,each_in_length,num_samples = 1,length = 1,temperature = 1,top_k=1,top_p=0.0):

    encoded_raw_sent = torch.tensor(encoded_raw_sent,device='cuda' if torch.cuda.is_available() else 'cpu')
    original_each_in_length = [i for i in each_in_length]

    generated = batch_prompt_in

    with torch.no_grad():
        end_flags=torch.tensor([1]*batch_prompt_in.shape[0],device='cuda' if torch.cuda.is_available() else 'cpu')
        for current_length in range(length):
            inputs = {'input_ids': generated}

            outputs = model(
                **inputs)
            logits = outputs[1]

            t = []
            for i, j in enumerate(each_in_length):

                t.append(logits[i].narrow(dim=0, start=j - 1, length=1))
            next_token_logits = torch.cat(t, dim=0)
            next_token_logits = torch.gather(next_token_logits, 1, encoded_raw_sent)
            next_token = torch.argmax(next_token_logits, dim=-1)
            next_token.unsqueeze_(-1)
            next_token = torch.gather(encoded_raw_sent, 1, next_token)
            torch_next_token = next_token.squeeze(1)
            next_token=torch_next_token.tolist()

            #collect '.',stop generating when all sentences in the batch meet '.'
            this_end = torch.where(torch_next_token==764,0,1)
            end_flags*=this_end

            g = []
            for i, j in enumerate(each_in_length):
                p1 = generated[i].narrow(dim=0, start=0, length=j).unsqueeze(0)
                p2 = torch.tensor([next_token[i]], device='cuda' if torch.cuda.is_available() else 'cpu').unsqueeze(0)
                p3 = generated[i].narrow(dim=0, start=j, length=generated.shape[-1] - j).unsqueeze(0)
                g.append(torch.cat([p1, p2, p3], dim=1))
                each_in_length[i] += 1
            generated = torch.cat(g, dim=0)
            if torch.sum(end_flags).item()==0:
                break

    #split generated slots and input sentence
    generated = generated.cpu().tolist()
    texts = []
    slots = []
    for i, j in enumerate(generated):
        if 50256 in j:
            idx = j.index(50256)
            text = j[:idx]
        else:
            text = j
        texts.append(tokenizer.decode(text, clean_up_tokenization_spaces=False))
        slot = text[original_each_in_length[i]:]
        slots.append(tokenizer.decode(slot, clean_up_tokenization_spaces=False))
    return texts, slots

def finetune_on_support(args,test_model,optimizer,scheduler,dataloader,logger):
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError(
                "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        test_model, optimizer1 = amp.initialize(test_model, optimizer, opt_level=args.fp16_opt_level)

    test_model.train()
    for idx_epochs in range(int(args.num_finetune_epochs)):

        all_loss = 0.0
        nb_steps = 0

        optimizer.zero_grad()

        total_step, update_step = 0, 0
        for id, data in enumerate(dataloader):
            batch_prompt_in, batch_prompt_out, batch_masked_prompt_out, batch_attention_mask, batch_raw_sent, batch_in_length = data
            outputs = test_model(batch_prompt_out, labels=batch_masked_prompt_out,
                                attention_mask=batch_attention_mask)
            loss = outputs[0]
            if args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            all_loss += loss.mean().item()
            nb_steps += 1

            if (total_step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(test_model.parameters(), args.max_grad_norm)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                update_step += 1

            total_step += 1
        all_loss = all_loss / nb_steps

        logger.info(
            "Epoch: {}, Step: {}, Finetuning stage 1 loss: {:.2e} update_step: {:.2e}\n".format(
                idx_epochs, total_step, all_loss, update_step))

def generate_procedure(args,test_model,test_model_2,tokenizer,test_test_dataloader):
    #stage 1 gen
    last_sent = ''
    all_raw_lines,all_label_lines,all_results_lines = [],[],[]
    domain_name =''
    for id, data in enumerate(test_test_dataloader):
        batch_prompt_in, batch_prompt_out, batch_masked_prompt_out, batch_attention_mask, batch_raw_sent, batch_in_length = data
        outs, slots = generate(model=test_model, tokenizer=tokenizer, batch_prompt_in=batch_prompt_in,
                                encoded_raw_sent=batch_raw_sent, each_in_length=batch_in_length, length=args.gen_length)
        count =0
        with open(args.test_path+'verb_label', 'r') as f:
            label_map = json.load(f)

        for out, slot in zip(outs, slots):
            if args.dataset=='mit':
                domain_name="atis"
            else:
                domain_name = out.split('"')[0].split(':')[0].strip()
            label_num = len(list(label_map[domain_name].items()))

            if out.split('"')[1] != last_sent:
                count=1
            else:
                count+=1
            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                all_label_lines.append('\n')

            label = label_map[domain_name][out.split('"')[-1].split('refers to')[0].strip()]
            all_label_lines.append(label)

            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                raw = out.split('"')[1].strip()
                all_raw_lines.append(raw)

            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                all_results_lines.append('\n')
                last_sent = out.split('"')[1]

            if not slot.endswith('.'):
                    slot+=' .'
            all_results_lines.append(slot)

    labels = []
    one_sent_label = []
    for line in all_label_lines:
        if line == '\n':
            one_sent_label = []
            labels.append(one_sent_label)
        else:
            one_sent_label.append(line.strip())

    results = []
    one_sent_results=[]
    for line in all_results_lines:
        if line=='\n':
            one_sent_results = []
            results.append(one_sent_results)
        else:
            one_sent_results.append(line.strip())

    raws = [i.strip() for i in all_raw_lines]

    #write stage1 generation result
    with open(args.pred_path+args.test_file+'/raw1.txt', 'a') as f:
        for i in raws:
            f.write(i)
            f.write('\n')

    with open(args.pred_path+args.test_file+'/label1.txt', 'a') as f:
        for i in labels:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    with open(args.pred_path+args.test_file+'/result1.txt', 'a') as f:
        for i in results:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    #stage2 gen
    with open(args.test_path+'label_verb','r')as f:
        label2verb=json.load(f)

    with open(args.test_path+'verb_label', 'r') as f:
            label_map = json.load(f)

    #construct stage2 generation inputs based on results from stage1
    checker_test_test_dataset = CheckerDevTestEpisode(args.dataset,raws,labels,results,label2verb[domain_name],domain_name,tokenizer)
    checker_test_test_dataloader = DataLoader(dataset=checker_test_test_dataset,batch_size=args.gen_batch_size,collate_fn=test_checker_collate)

    last_sent = ''
    all_label_lines_2,all_raw_lines_2,all_results_lines_2 = [],[], []
    for id, data in enumerate(checker_test_test_dataloader):
        batch_prompt_in , batch_raw_sent, batch_in_length = data
        outs, slots = generate(model=test_model_2, tokenizer=tokenizer, batch_prompt_in=batch_prompt_in,
                                encoded_raw_sent=batch_raw_sent, each_in_length=batch_in_length, length=args.gen_length)
        count =0

        for out, slot in zip(outs, slots):
            if args.dataset=='mit':
                domain_name="atis"
            else:
                domain_name = out.split('"')[0].split(':')[0].strip()
            label_num = len(list(label_map[domain_name].items()))

            if out.split('"')[1] != last_sent:
                count=1
            else:
                count+=1
            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                all_label_lines_2.append('\n')
            label_name = out.split('"')[-1].split('refers to')[-2].strip()
            if '.'in label_name:
                label_name = label_name.split('.')[-1].strip()
            label_2 = label_map[domain_name][label_name]
            all_label_lines_2.append(label_2)

            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                raw_2 = out.split('"')[1].strip()
                all_raw_lines_2.append(raw_2)

            if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                all_results_lines_2.append('\n')
                last_sent = out.split('"')[1]
            if not slot.endswith('.'):
                slot+=' .'
            all_results_lines_2.append(slot)

    labels_2 = []
    one_sent_label_2 = []
    for line in all_label_lines_2:
        if line == '\n':
            one_sent_label_2 = []
            labels_2.append(one_sent_label_2)
        else:
            one_sent_label_2.append(line.strip())

    results_2 = []
    one_sent_results_2=[]
    for line in all_results_lines_2:
        if line=='\n':
            one_sent_results_2 = []
            results_2.append(one_sent_results_2)
        else:
            one_sent_results_2.append(line.strip())

    raws_2 = [i.strip() for i in all_raw_lines_2]

    for id,data in enumerate(zip(labels,results,labels_2,results_2)):
        label1,result1,label2,result2 = data
        unrecognized_label = []
        for l,r in zip(label1,result1):
            if r.strip().startswith('none'):
                unrecognized_label.append(l)
        checked_label = []
        for l,r in zip(label2,result2):
            if not r.strip().startswith('none'):
                checked_label.append(l)
        for change_label in checked_label:
            if change_label in unrecognized_label:
                change_idx_1 = label1.index(change_label)
                change_idx_2 = label2.index(change_label)
                results[id][change_idx_1]=result2[change_idx_2]

    #write stage2 and final results
    with open(args.pred_path+args.test_file+'/raw.txt', 'a') as f:
        for i in raws:
            f.write(i)
            f.write('\n')

    with open(args.pred_path+args.test_file+'/label.txt', 'a') as f:
        for i in labels:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    with open(args.pred_path+args.test_file+'/result.txt', 'a') as f:
        for i in results:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    with open(args.pred_path+args.test_file+'/raw2.txt', 'a') as f:
        for i in raws_2:
            f.write(i)
            f.write('\n')

    with open(args.pred_path+args.test_file+'/label2.txt', 'a') as f:
        for i in labels_2:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    with open(args.pred_path+args.test_file+'/result2.txt', 'a') as f:
        for i in results_2:
            f.write('\n')
            for j in i:
                f.write(j)
                f.write('\n')

    return

def test(args,model,tokenizer,logger,test_data):

    # prepare dataset
    test_dataset = DevDataset(test_data,tokenizer)
    bar = tqdm(test_dataset,desc='Testing')

    # make output dir
    if not os.path.exists(args.pred_path):
        os.mkdir(args.pred_path)
    if not os.path.exists(args.pred_path+args.test_file):
        os.mkdir(args.pred_path+args.test_file)

    for eid, test_ep in enumerate(bar):
        #copy first-round generation model for each episode
        test_model = copy.deepcopy(model)
        test_model.to(args.device)

        #construct all dataloaders
        support_encoded_prompt_in, support_encoded_prompt_out, support_encoded_raw_sent, support_mask_encoded_prompt_out, test_encoded_prompt_in, test_encoded_prompt_out, test_encoded_raw_sent, test_mask_encoded_prompt_out,support_encoded_checker_prompt_in, support_encoded_checker_prompt_out, support_encoded_checker_raw_sent, support_mask_encoded_checker_prompt_out, test_encoded_checker_prompt_in, test_encoded_checker_prompt_out, test_encoded_checker_raw_sent, test_mask_encoded_checker_prompt_out = test_ep
        test_support_episode = DevSupportEpisode(support_encoded_prompt_in, support_encoded_prompt_out,
                                                support_encoded_raw_sent, support_mask_encoded_prompt_out,
                                                tokenizer)
        test_test_episode = DevTestEpisode(test_encoded_prompt_in, test_encoded_prompt_out,
                                          test_encoded_raw_sent, test_mask_encoded_prompt_out, tokenizer)
        checker_test_support_episode = CheckerDevSupportEpisode(support_encoded_checker_prompt_in, support_encoded_checker_prompt_out,
                                          support_encoded_checker_raw_sent, support_mask_encoded_checker_prompt_out, tokenizer)
        test_support_dataloader = DataLoader(dataset=test_support_episode,batch_size=2,shuffle=True,collate_fn=devcollate)
        test_test_dataloader = DataLoader(dataset=test_test_episode,batch_size=args.gen_batch_size,collate_fn=devcollate)
        checker_test_support_dataloader = DataLoader(dataset=checker_test_support_episode,batch_size=2,collate_fn=checker_devcollate)

        # prepare first-round optimizer and scheduler
        if args.max_finetune_steps > 0:
            total = args.max_finetune_steps
            args.num_finetune_epochs = args.max_finetune_steps // (
                    len(test_support_dataloader) // args.gradient_accumulation_steps) + 1
        else:
            total = len(test_support_dataloader) // args.gradient_accumulation_steps * args.num_finetune_epochs

        param_optimizer = list(test_model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay': args.weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer1 = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
        scheduler1 = get_linear_schedule_with_warmup(optimizer1, num_warmup_steps=args.warmup_rate * total,num_training_steps=total)

        #first-round finetune on the support set
        finetune_on_support(args=args,test_model=test_model,optimizer=optimizer1,scheduler=scheduler1,dataloader=test_support_dataloader,logger=logger)

        #copy second-round generation model for each episode
        test_model_2 = copy.deepcopy(test_model)
        test_model_2.to(args.device)
        test_model_2.train()

        # prepare second-round optimizer and scheduler
        if args.max_finetune_steps > 0:
            total = args.max_finetune_steps
            args.num_finetune_epochs = args.max_finetune_steps // (
                    len(checker_test_support_dataloader) // args.gradient_accumulation_steps) + 1
        else:
            total = len(checker_test_support_dataloader) // args.gradient_accumulation_steps * args.num_finetune_epochs
        param_optimizer = list(test_model_2.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay': args.weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer2 = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
        scheduler2 = get_linear_schedule_with_warmup(optimizer2, num_warmup_steps=args.warmup_rate * total,num_training_steps=total)

        #second-round finetune on the support set ,model copied from first round finetuned model.
        finetune_on_support(args=args,test_model=test_model_2,optimizer=optimizer2,scheduler=scheduler2,dataloader=checker_test_support_dataloader,logger=logger)
        #whole generation procedure
        generate_procedure(args,test_model,test_model_2,tokenizer,test_test_dataloader)
    _,_,test_score1 = eval(result_path=args.pred_path+args.test_file+'/result1.txt',label_path =args.pred_path+args.test_file+'/label1.txt',tar_path = args.pred_path+args.test_file+'/tar1.txt',raw_path=args.pred_path+args.test_file+'/raw1.txt',mode='test',args=args)
    _,_,test_score2 = eval(result_path=args.pred_path+args.test_file+'/result.txt',label_path =args.pred_path+args.test_file+'/label.txt',tar_path = args.pred_path+args.test_file+'/tar.txt',raw_path=args.pred_path+args.test_file+'/raw.txt',mode='test',args=args)
    return test_score1,test_score2

def model_selection(args,model,cpt_file,tokenizer,logger,best_score,dev_data):

    dev_dataset = DevDataset(dev_data, tokenizer)
    bar = tqdm(dev_dataset, desc="Dev")
    if not os.path.exists(args.model_selection_path):
        os.mkdir(args.model_selection_path)

    if not os.path.exists(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'):
        os.mkdir(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/')
    for eid, dev_ep in enumerate(bar):
        dev_model = copy.deepcopy(model)
        dev_model.to(args.device)
        support_encoded_prompt_in, support_encoded_prompt_out, support_encoded_raw_sent, support_mask_encoded_prompt_out, test_encoded_prompt_in, test_encoded_prompt_out, test_encoded_raw_sent, test_mask_encoded_prompt_out,support_encoded_checker_prompt_in, support_encoded_checker_prompt_out, support_encoded_checker_raw_sent, support_mask_encoded_checker_prompt_out, test_encoded_checker_prompt_in, test_encoded_checker_prompt_out, test_encoded_checker_raw_sent, test_mask_encoded_checker_prompt_out = dev_ep
        dev_support_episode = DevSupportEpisode(support_encoded_prompt_in, support_encoded_prompt_out,
                                                support_encoded_raw_sent, support_mask_encoded_prompt_out,
                                                tokenizer)
        dev_test_episode = DevTestEpisode(test_encoded_prompt_in, test_encoded_prompt_out,
                                          test_encoded_raw_sent, test_mask_encoded_prompt_out, tokenizer)
        dev_support_dataloader = DataLoader(dataset=dev_support_episode, batch_size=2, shuffle=True,
                                            collate_fn=devcollate)
        dev_test_dataloader = DataLoader(dataset=dev_test_episode, batch_size=args.gen_batch_size,
                                         collate_fn=devcollate)

        if args.max_finetune_steps > 0:
            total = args.max_finetune_steps
            args.num_finetune_epochs = args.max_finetune_steps // (
                        len(dev_support_dataloader) // args.gradient_accumulation_steps) + 1
        else:
            total = len(dev_support_dataloader) // args.gradient_accumulation_steps * args.num_finetune_epochs

        param_optimizer = list(dev_model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay': args.weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_rate*total,
                                                    num_training_steps=total)
        if args.fp16:
            try:
                from apex import amp
            except ImportError:
                raise ImportError(
                    "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
            dev_model, optimizer = amp.initialize(dev_model, optimizer, opt_level=args.fp16_opt_level)

        dev_model.train()

        for idx_epochs in trange(int(args.num_finetune_epochs), desc="Epoch"):
            all_loss = 0.0
            nb_steps = 0

            ft_tdqm_bar = tqdm(dev_support_dataloader, desc="Modelselection_FineTuning")
            optimizer.zero_grad()

            total_step, update_step=0, 0
            for id,data in enumerate(ft_tdqm_bar):
                batch_prompt_in, batch_prompt_out, batch_masked_prompt_out, batch_attention_mask, batch_raw_sent,batch_in_length  = data
                outputs = dev_model(batch_prompt_out, labels=batch_masked_prompt_out,attention_mask=batch_attention_mask)
                loss = outputs[0]
                if args.n_gpu > 1:
                    loss = loss.mean()  # mean() to average on multi-gpu parallel training
                if args.gradient_accumulation_steps > 1:
                    loss = loss / args.gradient_accumulation_steps

                if args.fp16:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()
                all_loss += loss.mean().item()
                nb_steps += 1

                if (total_step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16:
                        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                    else:
                        torch.nn.utils.clip_grad_norm_(dev_model.parameters(), args.max_grad_norm)
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad()
                    update_step += 1

                total_step += 1
            all_loss = all_loss / nb_steps
            ft_tdqm_bar.desc = "ModelSelection_Finetuning loss: {:.2e} lr: {:.2e}".format(all_loss, scheduler.get_lr()[0])

            logger.info(
                "Epoch: {}, Step: {}, Finetuning loss: {:.2e} update_step: {:.2e}\n".format(
                    idx_epochs, total_step, all_loss, update_step))

        gen_tdqm_bar = tqdm(dev_test_dataloader, desc="Generating")
        last_sent = ''
        for id, data in enumerate(gen_tdqm_bar):
            batch_prompt_in, batch_prompt_out, batch_masked_prompt_out, batch_attention_mask, batch_raw_sent, batch_in_length = data

            outs, slots = generate(model=dev_model, tokenizer=tokenizer, batch_prompt_in=batch_prompt_in,
                                   encoded_raw_sent=batch_raw_sent, each_in_length=batch_in_length, length=20)
            count = 0
            with open(args.dev_path+'verb_label', 'r') as f:
                label_map = json.load(f)

            for out, slot in zip(outs, slots):
                if args.dataset=='mit':
                    domain_name="atis"
                else:
                    domain_name = out.split('"')[0].split(':')[0].strip()
                label_num = len(list(label_map[domain_name].items()))
                if not os.path.exists(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/'):
                    os.mkdir(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/')

                with open(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/label_'+cpt_file+'.txt', 'a') as f:
                    if out.split('"')[1] != last_sent:
                        count = 1
                    else:
                        count += 1
                    if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                        f.write('\n')
                    f.write(label_map[domain_name][out.split('"')[-1].split('refers to')[0].strip()])
                    f.write('\n')
                with open(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/raw_'+cpt_file+'.txt', 'a') as f:
                    if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                        f.write(out.split('"')[1].strip())
                        f.write(('\n'))
                with open(args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/result_'+cpt_file+'.txt', 'a') as f:
                    if out.split('"')[1] != last_sent or count > label_num and count % label_num==1:
                        f.write('\n')
                        last_sent = out.split('"')[1]

                    f.write(slot)
                    f.write('\n')

    _,_,dev_score = eval(result_path=args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/result_'+cpt_file+'.txt',
                        label_path = args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/label_'+cpt_file+'.txt',
                        tar_path = args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/tar_'+cpt_file+'.txt',
                        raw_path = args.model_selection_path+str(args.dataset)+str(args.dev_file)+'/'+cpt_file+'/raw_'+cpt_file+'.txt',
                        mode='dev',args=args)

    best_model = None

    if dev_score > best_score:
        logger.info(" === Found new best!! "+cpt_file+"=== ")

        ''' store new best model  '''
        best_model = copy.deepcopy(model)  # copy model to avoid writen by latter training
        ''' save model file '''

        if os.path.exists(args.ft_model_output_dir+'_'+str(args.dev_file)+'/best'):
            os.system(r'rm -rf ' + args.ft_model_output_dir+'_'+str(args.dev_file)+'/best')
            save_model(finetune=True, args=args, logger=logger, model=best_model, tokenizer=tokenizer,train_file=args.train_file, time='best')
        else:
            save_model(finetune=True, args=args, logger=logger, model=best_model, tokenizer=tokenizer,train_file=args.train_file, time='best')
    logger.info("dev_score :" + str(dev_score))

    return dev_score, best_model

def select_model_from_checkpoint(args,tokenizer,logger,dev_data):
    all_cpt_file = sorted(list(os.listdir(args.model_output_dir+str(args.train_file)+'/')))
    best_score = 0
    best_model = None
    all_cpt_file = sorted(all_cpt_file, key=lambda x: int(x.split('_')[-1]))
    for cpt_file in all_cpt_file:

        model = GPT2GenModel.from_pretrained(pretrained_model_name_or_path =args.model_output_dir+str(args.train_file)+'/'+cpt_file)
        model.resize_token_embeddings(len(tokenizer))
        model.to(args.device)
        dev_score,  copied_model = model_selection(args,model, cpt_file,tokenizer,logger,best_score, dev_data)
        if dev_score > best_score:
            best_score = dev_score
            best_model = copied_model
    return best_model,best_score

def save_model(finetune,args,logger,model, tokenizer,train_file, time):
    if not os.path.exists(args.ft_model_output_dir):
        os.mkdir(args.ft_model_output_dir)
    if finetune:
        model_output_dir = os.path.join(args.ft_model_output_dir+train_file,str(time))
    else:
        model_output_dir = os.path.join(args.model_output_dir+train_file, str(time))

    if not os.path.exists(args.ft_model_output_dir):
        os.mkdir(args.ft_model_output_dir)
    if not os.path.exists(model_output_dir):
        os.makedirs(model_output_dir)
    logger.info("\nSaving the model to {}\n".format(os.path.join(model_output_dir)))

    # Save a trained model, configuration and tokenizer
    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self

    # If we save using the predefined names, we can load using `from_pretrained`
    output_model_file = os.path.join(model_output_dir, WEIGHTS_NAME)
    output_config_file = os.path.join(model_output_dir, CONFIG_NAME)

    torch.save(model_to_save.state_dict(), output_model_file)
    model_to_save.config.to_json_file(output_config_file)
    tokenizer.save_vocabulary(model_output_dir)

def main():
    args = get_args()
    print('args: ', args)
    if not os.path.exists(args.pred_path):
        os.mkdir(args.pred_path)
    if not os.path.exists(args.pred_path+args.test_file):
        os.mkdir(args.pred_path+args.test_file)
    # Get ready...
    logging.basicConfig(filename=args.log_output_path,
                        format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO)
    logger = logging.getLogger(__name__)
    logger.info('\nargs: {}'.format(args))

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    n_gpu = torch.cuda.device_count()
    logger.info("device: {}, n_gpu {}".format(args.device, n_gpu))

    #tokenizer & model
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    bos_token, eos_token, unk_token = tokenizer.special_tokens_map['bos_token'], \
                                      tokenizer.special_tokens_map['eos_token'], \
                                      tokenizer.special_tokens_map['unk_token']
    bos_token_id, eos_token_id, unk_token_id = tokenizer.convert_tokens_to_ids([bos_token, eos_token, unk_token])

    if args.wo_pretrained:
        config = GPT2Config.from_pretrained("gpt2", n_layer=args.wo_pretrained_layer)
        model = GPT2GenModel(config)
        logger.info("GPT2 model without pretrained: {}".format(model))
    else:
        model = GPT2GenModel.from_pretrained("gpt2")
    model.resize_token_embeddings(len(tokenizer))
    model.to(args.device)

    train_data, dev_data, test_data = read_data(args)
    # Load encoded data
    if args.do_train:
        train_dataset = TrainDataset(train_data,tokenizer)
        train_sampler = RandomSampler(train_dataset)
        train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.train_batch_size, collate_fn=traincollate)
        # Prepare optimizer

        if args.max_steps > 0:
            t_total = args.max_steps
            args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
        else:
            t_total = len(train_dataloader)  // args.gradient_accumulation_steps * args.num_train_epochs

        param_optimizer = list(model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay': args.weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_rate*t_total, num_training_steps=t_total)

        if args.fp16:
            try:
                from apex import amp
            except ImportError:
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
            model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

        if n_gpu > 1:
            model = torch.nn.DataParallel(model)
        model.train()
        update_step, total_step, exp_average_loss = 0, 0, None
        # Train!
        logger.info("***** Running training *****")
        logger.info("  Num examples = %d", len(train_dataset))
        logger.info("  Num Epochs = %d", args.num_train_epochs)
        logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                    args.train_batch_size * args.gradient_accumulation_steps)
        logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
        logger.info("  Total optimization steps = %d", t_total)

        for idx_epochs in trange(int(args.num_train_epochs),desc="Epoch"):
            total_step = 0
            num_save_checkpoint = len(train_dataloader) // args.n_save_per_epoch
            tqdm_bar = tqdm(train_dataloader, desc="Training")
            optimizer.zero_grad()
            for idx_batch, batch in enumerate(tqdm_bar):
                batch_encoded_prompt_in, batch_encoded_prompt_out,batch_mask_encoded_prompt_out, batch_encoded_raw_sent, batch_attention_mask, batch_slot_poses= batch
                batch_size = batch_encoded_prompt_in.shape[0]
                outputs = model(batch_encoded_prompt_out, labels=batch_slot_poses,attention_mask = batch_attention_mask,gen_tokens_list=batch_encoded_raw_sent ,slot_poses=batch_slot_poses)
                loss =outputs[0]

                if args.n_gpu > 1:
                    loss = loss.mean()  # mean() to average on multi-gpu parallel training
                if args.gradient_accumulation_steps > 1:
                    loss = loss / args.gradient_accumulation_steps

                if args.fp16:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                if (total_step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16:
                        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                    else:
                        torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad()
                    update_step += 1

                exp_average_loss = loss.item() if exp_average_loss is None else 0.7 * exp_average_loss + 0.3 * loss.item()
                total_step += 1
                tqdm_bar.desc = "Training loss: {:.2e} lr: {:.2e}".format(exp_average_loss, scheduler.get_lr()[0])
                logger.info(
                    "Epoch: {}, Step: {}, Training loss: {:.2e} current loss: {:.2e} "
                    "lr: {:.2e} update_step: {:.2e}\n".format(
                        idx_epochs, total_step, exp_average_loss, loss.item(),
                        scheduler.get_lr()[0], update_step))
            save_model(finetune=False,args=args,logger=logger,model=model, tokenizer=tokenizer,train_file=args.train_file,time=
                           'epoch_' + str(idx_epochs) + '_step_' + str(update_step))
    best_model = None
    best_score = 0.
    if args.do_dev:
        best_model,best_score = select_model_from_checkpoint(args,tokenizer,logger,dev_data)

    if args.do_test:
        if args.dataset=='snips':
            best_model=GPT2GenModel.from_pretrained(pretrained_model_name_or_path=args.ft_model_output_dir+str(args.train_file)+'/best')
        else:
            best_model=GPT2GenModel.from_pretrained('gpt2')
        test_score1,test_score2 = test(args,best_model,tokenizer,logger,test_data)
        logger.info("test_score_1 : "+str(test_score1))
        logger.info("test_score_2 : "+str(test_score2))

if __name__ == '__main__':
    main()


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pickle

import numpy as np
from tqdm import tqdm

from npybvh import bvh
import definitions
import torch

class BvhConverter(bvh.Bvh):
    root_path = ""
    joint_mappings = {}
    toe_joint = "b_r_ball"
    scale = 1
    ignore_files = []

    Keyframe = definitions.Keyframe
    KeyframeIdx = definitions.KeyframeIdx
    Skeleton = definitions.Skeleton
    save_separately = False
    file_occlusions = None
    file_out_of_views = None
    file_occlusions_joints = []

    def _resize_occlusions(self, target_size, occlusions, out_of_views, too_close):
        if len(occlusions) > target_size:
            print(f"occlusions are longer than animation ({len(occlusions)} vs {target_size})")
            occlusions = occlusions[len(occlusions) - target_size:]
            out_of_views = out_of_views[len(out_of_views) - target_size:]
            too_close = too_close[len(too_close) - target_size:]

        if len(occlusions) < target_size:
            print(f"occlusions are shorter than animation ({len(occlusions)} vs {target_size})")
            new_occ = np.zeros(target_size, dtype=np.int)
            new_oov = np.zeros(target_size, dtype=np.int)
            new_too_close = np.zeros(target_size, dtype=np.int)
            new_occ[:len(occlusions)] = occlusions
            new_oov[:len(occlusions)] = out_of_views
            new_too_close[:len(occlusions)] = too_close
            new_occ[len(occlusions):] = new_occ[len(occlusions) - 1]
            new_oov[len(occlusions):] = new_oov[len(occlusions) - 1]
            new_too_close[len(occlusions):] = too_close[len(occlusions) - 1]
            occlusions = new_occ
            out_of_views = new_oov
            too_close = new_too_close

        return occlusions, out_of_views, too_close

    def convert_all(self, file="", folder=""):
        p, r = self.batch_all_frame_poses()
        p = p.numpy()
        r = r.numpy()
        occlusions, out_of_views, too_close = self.get_occlusions_and_out_of_views(file)

        occlusions, out_of_views, too_close = self._resize_occlusions(len(p), occlusions, out_of_views)

        keyframes = []
        keyframes.extend(self.batch_convert_frames(p, r, occlusions, out_of_views, too_close))

        return np.array(keyframes)

    def batch_convert_frames(self, p, r, occlusions, out_of_views, too_close):
        p_converted = np.empty((len(p), len(self.KeyframeIdx.all), 3))
        r_converted = np.empty((len(p), len(self.KeyframeIdx.all), 3))

        for our_joint_name, bvh_joint_name in self.joint_mappings.items():
            id_keyframe = self.KeyframeIdx.all[our_joint_name]
            id_bvh = list(self.joints.keys()).index(bvh_joint_name)
            p_converted[:, id_keyframe] = p[:, id_bvh] / self.scale
            r_converted[:, id_keyframe] = r[:, id_bvh]

        r_local = self.Keyframe.batch_torch_recover_local_rotations(self.Skeleton,
                                                                    torch.from_numpy(r_converted)).numpy()

        result = np.zeros((len(p), len(self.KeyframeIdx.all), 10), dtype=np.float32)
        result[:, :, :3] = p_converted
        result[:, :, 3:6] = r_converted
        result[:, :, 6:9] = r_local
        result[:, 0, 9] = np.linspace(0.0, len(p) / float(self.fps), len(p))

        use_out_of_view = True
        use_too_close = True
        occ_mask = np.array(occlusions)
        if use_out_of_view and len(occlusions) > 0:
            occ_mask |= np.array(out_of_views)
        if use_too_close and len(occlusions) > 0:
            occ_mask |= np.array(too_close)

        if len(occlusions) > 0:
            for occlusion_joint_idx, joint_name in enumerate(self.file_occlusions_joints):
                if joint_name not in self.KeyframeIdx.all:
                    continue

                joint_bit_mask = 1 << occlusion_joint_idx
                result_joint_index = self.KeyframeIdx.all[joint_name]
                result[:, result_joint_index, 9][occ_mask & joint_bit_mask == joint_bit_mask] += 1.
        return result

    def get_occlusions_and_out_of_views(self, file=""):
        return [], [], []

    def skeleton(self):
        offset = np.zeros(3)
        poses = []
        self._add_pose_recursive(self.root, offset, poses)
        root_height = -poses[list(self.joints.keys()).index(self.toe_joint)][1]
        offset_pos = np.array([0, root_height, 0])

        positions = np.zeros((len(self.joint_mappings), 3))
        for key in self.joint_mappings.keys():
            positions_index = self.KeyframeIdx.all[key]
            bvh_index = list(self.joints.keys()).index(self.joint_mappings[key])
            positions[positions_index] = poses[bvh_index].reshape(3) + offset_pos

        return self.Skeleton(positions)

    @classmethod
    def get_all_file_paths(cls):
        result = []
        for folder in cls.get_all_folders():
            result.extend(cls.get_all_files_of_folder(folder))
        return

    @classmethod
    def get_normalized_path(cls):
        raise NotImplemented()

    @classmethod
    def convert_all_files(cls, normalize_skeleton=False):
        keyframes = []
        skeletons = []
        file_count = 0

        for file in tqdm(cls.get_all_file_paths(), f"Converting bvh to numpy"):
            parser = cls()
            # anim_name = file.split('/')[-1][:-4]
            anim_name = os.path.split(file)[-1][:-4]
            folder = file.split(anim_name)[0]
            parser.get_occlusions_and_out_of_views(folder)
            parser.parse_file(file)

            if normalize_skeleton:
                override_parser = cls()
                override_parser.parse_file(cls.get_normalized_path())
                for key, val in parser.joints.items():
                    val.offset = override_parser.joints[key].offset

            skeleton = parser.skeleton()
            if skeleton is None:
                continue
            skeletons.append(skeleton)
            keyframes.append(parser.convert_all(anim_name, folder))

            if cls.save_separately:
                suffix = f"{cls.Keyframe.__name__}{file_count}{'_normalized' if normalize_skeleton else ''}"
                with open(os.path.join(cls.root_path, f"skeletons_{suffix}.pickle"), "wb+") as outfile:
                    pickle.dump(skeletons, outfile)

                np.save(os.path.join(cls.root_path, f"keyframes_{suffix}.npy"), keyframes, allow_pickle=True)
                file_count += 1
                skeletons.clear()
                keyframes.clear()

        if not cls.save_separately:
            suffix = f"{cls.Keyframe.__name__}{'_normalized' if normalize_skeleton else ''}"
            with open(os.path.join(cls.root_path, f"skeletons{suffix}.pickle"), "wb+") as outfile:
                pickle.dump(skeletons, outfile)

            np.save(os.path.join(cls.root_path, f"keyframes{suffix}.npy"), keyframes, allow_pickle=True)

    @classmethod
    def get_all_folders(cls):
        return list(os.walk(cls.root_path))[0][1]

    @classmethod
    def get_all_files_of_folder(cls, folder_path):
        file_paths = []
        count_frames = False
        frames = 0
        for file_name in list(os.walk(folder_path))[0][2]:
            if file_name in cls.ignore_files:
                continue
            if file_name.endswith(".bvh"):
                file_paths.append(os.path.join(folder_path, file_name))

                if count_frames:
                    with open(file_paths[-1]) as f:
                        line = ""
                        while "Frames: " not in line:
                            line = f.readline()
                        frames += int(line.split(":    ")[1])

        if count_frames:
            print(f"frames: {frames}")

        return file_paths

    def _apply_postprocessing(self, keyframes):
        return keyframes

    @classmethod
    def get_all_keyframe_folders(cls):
        return cls.get_all_folders()

    def _check_and_convert(self, normalized):
        folders = self.__class__.get_all_keyframe_folders()
        suffix = f"{self.Keyframe.__name__}{0}{'_normalized' if normalized else ''}"
        if not os.path.exists(os.path.join(list(folders)[0], f"skeletons_{suffix}.pickle")):
            self.__class__.convert_all_files(normalize_skeleton=normalized)

    def load_numpy(self, normalized=False):
        # check if dataset was already converted. Otherwise, do now
        self._check_and_convert(normalized)

        folders = self.__class__.get_all_keyframe_folders()
        skeletons = []
        keyframes = []
        for folder in folders:
            file_count = 0
            _suffix = ("_normalized" if normalized else "")
            if self.save_separately:
                suffix = f"{self.Keyframe.__name__}{file_count}{_suffix}"
                while os.path.exists(os.path.join(folder, f"skeletons_{suffix}.pickle")):
                    try:
                        with open(os.path.join(folder, f"skeletons_{suffix}.pickle"), "rb+") as infile:
                            _skeletons = pickle.load(infile)
                        _keyframes = np.load(os.path.join(folder, f"keyframes_{suffix}.npy"), allow_pickle=True)
                    except:
                        print(f"could not import {folder}{suffix}")
                    if _keyframes.shape[0] == 1:
                        _keyframes = _keyframes.reshape(_keyframes.shape[1:])
                    _keyframes = self._apply_postprocessing(_keyframes)
                    _keyframes = list(_keyframes.reshape((1, *_keyframes.shape)))
                    yield _skeletons, _keyframes
                    file_count += 1
                    suffix = f"{self.Keyframe.__name__}{file_count}{_suffix}"
                return
            else:
                suffix = _suffix
                try:
                    with open(os.path.join(folder, f"skeletons{suffix}.pickle"), "rb+") as infile:
                        _skeletons = pickle.load(infile)
                    _keyframes = np.load(os.path.join(folder, f"keyframes{suffix}.npy"), allow_pickle=True)
                except:
                    print(f"could not import {folder}{suffix}")
                _keyframes = self._apply_postprocessing(_keyframes)
                skeletons.extend(_skeletons)
                keyframes.extend(list(_keyframes))

        if not self.save_separately:
            return zip(skeletons, keyframes)

    def get_bone_lengths(self, parent_idx, normalized=True):
        from utils import bone_lengths_np

        for skeleton, tracks in self.load_numpy(normalized=normalized):
            for _track in tracks:
                track = _track[0] if len(_track.shape) == 4 else _track
                return bone_lengths_np(track[0, :, :3], parent_idx)

    @classmethod
    def name(cls):
        return "none"


#!/usr/bin/env python3
#
# Copyright (C) 2021 Learning Algorithms and Systems Laboratory, EPFL, Switzerland
# Authors:
# 	Vaibhav Gupta (vaibhav.gupta@epfl.ch)
# 	Saurav Aryan (saurav.aryan@epfl.ch) [Maintainer]
# Website: lasa.epfl.ch
#
# This file is part of `allegro_hand_gazebo`.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
from collections import deque
import numpy as np

import rospy
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import JointState

# ----- Joint Specs -----
JOINT_SPECS = {
    "min": np.deg2rad([
        -26.92, -11.22, -9.97, -13.0,
        -26.92, -11.22, -9.97, -13.0,
        -26.92, -11.22, -9.97, -13.0,
        -1.00, -10.00, -10.80, -5.00,
    ]),
    "max": np.deg2rad([
        26.92, 92.25, 97.92, 92.70,
        26.92, 92.25, 97.92, 92.70,
        26.92, 92.25, 97.92, 92.70,
        85.00, 72.00, 94.00, 98.00,
    ]),
}

# ----- Global Variables -----
TARGET_UPDATED = False
TARGET = deque(maxlen=1)    # Change `maxlen` for window averaging
TARGET.append(np.zeros(16))

# ----- Callbacks -----
def joint_state_callback(data):
    """Callback for `JointState` message
    Updates target to `JointState` while ensuring joint limits.

    Args:
        data (JointState): ROS `JointState` message
    """
    global TARGET
    TARGET.append(np.minimum(JOINT_SPECS['max'], np.maximum(JOINT_SPECS['min'], data.position)))

# ----- Main Script -----
if __name__ == '__main__':
    rospy.init_node('gazebo_joint_state_cmd', anonymous=True)
    rate = rospy.Rate(30)

    rospy.Subscriber(
        "joint_cmd",
        JointState,
        joint_state_callback,
        queue_size=1
    )

    publisher = rospy.Publisher(
        "PositionController/command",
        Float64MultiArray,
        queue_size=1
    )

    while not rospy.is_shutdown():
        msg = Float64MultiArray()
        msg.data = np.mean(TARGET, axis=0)
        publisher.publish(msg)
        rate.sleep()


#! /usr/bin/env python
#
# ILE_data_handling_probe.py
#
# Like ILE, but only providing the Qlm and inverse spectrum filters
#

from __future__ import print_function

import sys
import functools
from optparse import OptionParser, OptionGroup

import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as plt

import numpy
import numpy as np
try:
  import cupy
  xpy_default=cupy
  identity_convert = cupy.asnumpy
except:
  print(' no cupy')
#  import numpy as cupy  # will automatically replace cupy calls with numpy!
  xpy_default=numpy  # just in case, to make replacement clear and to enable override
  identity_convert = lambda x: x  # trivial return itself

import lal
from glue.ligolw import utils, lsctables, table, ligolw
from glue.ligolw.utils import process
import glue.lal

import lalsimutils
import factored_likelihood
import RIFT.integrators.mcsampler as mcsampler
import xmlutils

#from lalinference.bayestar import fits as bfits

__author__ = "Evan Ochsner <evano@gravity.phys.uwm.edu>, Chris Pankow <pankow@gravity.phys.uwm.edu>, R. O'Shaughnessy"

#
# Pinnable parameters -- for command line processing
#
LIKELIHOOD_PINNABLE_PARAMS = ["right_ascension", "declination", "psi", "distance", "phi_orb", "t_ref", "inclination"]

def get_pinned_params(opts):
    """
    Retrieve a dictionary of user pinned parameters and their pin values.
    """
    return dict([(p,v) for p, v in opts.__dict__.items() if p in LIKELIHOOD_PINNABLE_PARAMS and v is not None])

def get_unpinned_params(opts, params):
    """
    Retrieve a set of unpinned parameters.
    """
    return params - set([p for p, v in opts.__dict__.items() if p in LIKELIHOOD_PINNABLE_PARAMS and v is not None])

#
# Option parsing
#

optp = OptionParser()
optp.add_option("-c", "--cache-file", default=None, help="LIGO cache file containing all data needed.")
optp.add_option("-C", "--channel-name", action="append", help="instrument=channel-name, e.g. H1=FAKE-STRAIN. Can be given multiple times for different instruments.")
optp.add_option("-p", "--psd-file", action="append", help="instrument=psd-file, e.g. H1=H1_PSD.xml.gz. Can be given multiple times for different instruments.")
optp.add_option("-k", "--skymap-file", help="Use skymap stored in given FITS file.")
optp.add_option("-x", "--coinc-xml", help="gstlal_inspiral XML file containing coincidence information.")
optp.add_option("-I", "--sim-xml", help="XML file containing mass grid to be evaluated")
optp.add_option("-E", "--event", default=0,type=int, help="Event number used for this run")
optp.add_option("--n-events-to-analyze", default=1,type=int, help="Number of events to analyze from this XML")
optp.add_option("--soft-fail-event-range",action='store_true',help='Soft failure (exit 0) if event ID is out of range. This happens in pipelines, if we have pre-built a DAG attempting to analyze more points than we really have')
optp.add_option("-f", "--reference-freq", type=float, default=100.0, help="Waveform reference frequency. Required, default is 100 Hz.")
optp.add_option("--fmin-template", dest='fmin_template', type=float, default=40, help="Waveform starting frequency.  Default is 40 Hz. Also equal to starting frequency for integration")
optp.add_option("--fmin-ifo", action='append' , help="Minimum frequency for each IFO. Implemented by setting the PSD=0 below this cutoff. Use with care.")
#optp.add_option("--nr-params",default=None, help="List of specific NR parameters and groups (and masses?) to use for the grid.")
#optp.add_option("--nr-index",type=int,default=-1,help="Index of specific NR simulation to use [integer]. Mass used: mtot= m1+m2")
optp.add_option('--nr-group', default=None,help="If using a *ssingle specific simulation* specified on the command line, provide it here")
optp.add_option('--nr-param', default=None,help="If using a *ssingle specific simulation* specified on the command line, provide it here")
optp.add_option("--nr-lookup",action='store_true', help=" Look up parameters from an NR catalog, instead of using the approximant specified")
optp.add_option("--nr-lookup-group",action='append', help="Restriction on 'group' for NR lookup")
optp.add_option("--nr-hybrid-use",action='store_true',help="Enable use of NR (or ROM!) hybrid, using --approx as the default approximant and with a frequency fmin")
optp.add_option("--nr-hybrid-method",default="taper_add",help="Hybridization method for NR (or ROM!).  Passed through to LALHybrid. pseudo_aligned_from22 will provide ad-hoc higher modes, if the early-time hybridization model only includes the 22 mode")
optp.add_option("--rom-group",default=None)
optp.add_option("--rom-param",default=None)
optp.add_option("--rom-use-basis",default=False,action='store_true',help="Use the ROM basis for inner products.")
optp.add_option("--rom-limit-basis-size-to",default=None,type=int)
optp.add_option("--rom-integrate-intrinsic",default=False,action='store_true',help='Integrate over intrinsic variables. REQUIRES rom_use_basis at present. ONLY integrates in mass ratio as present')
optp.add_option("--nr-perturbative-extraction",default=False,action='store_true')
optp.add_option("--nr-use-provided-strain",default=False,action='store_true')
optp.add_option("--no-memory",default=False,action='store_true', help="At present, turns off m=0 modes. Use with EXTREME caution only if requested by model developer")
optp.add_option("--restricted-mode-list-file",default=None,help="A list of ALL modes to use in likelihood. Incredibly dangerous. Only use when comparing with models which provide restricted mode sets, or otherwise to isolate the effect of subsets of modes on the whole")
optp.add_option("--use-external-EOB",default=False,action='store_true')
optp.add_option("--maximize-only",default=False, action='store_true',help="After integrating, attempts to find the single best fitting point")
optp.add_option("--dump-lnL-time-series",default=False, action='store_true',help="(requires --sim-xml) Dump lnL(t) at the injected parameters")
optp.add_option("-a", "--approximant", default="TaylorT4", help="Waveform family to use for templates. Any approximant implemented in LALSimulation is valid.")
optp.add_option("-A", "--amp-order", type=int, default=0, help="Include amplitude corrections in template waveforms up to this e.g. (e.g. 5 <==> 2.5PN), default is Newtonian order.")
optp.add_option("--l-max", type=int, default=2, help="Include all (l,m) modes with l less than or equal to this value.")
optp.add_option("-s", "--data-start-time", type=float, default=None, help="GPS start time of data segment. If given, must also give --data-end-time. If not given, sane start and end time will automatically be chosen.")
optp.add_option("-e", "--data-end-time", type=float, default=None, help="GPS end time of data segment. If given, must also give --data-start-time. If not given, sane start and end time will automatically be chosen.")
optp.add_option("--data-integration-window-half",default=50*1e-3,type=float,help="Only change this window size if you are an expert. The window for time integration is -/+ this quantity around the event time")
optp.add_option("-F", "--fmax", type=float, help="Upper frequency of signal integration. Default is use PSD's maximum frequency.")
optp.add_option("--srate",default=16384,type=int,help="Sampling rate. Change ONLY IF YOU ARE ABSOLUTELY SURE YOU KNOW WHAT YOU ARE DOING.")
optp.add_option("-t", "--event-time", type=float, help="GPS time of the event --- probably the end time. Required if --coinc-xml not given.")
optp.add_option("-i", "--inv-spec-trunc-time", type=float, default=8., help="Timescale of inverse spectrum truncation in seconds (Default is 8 - give 0 for no truncation)")
optp.add_option("-w", "--window-shape", type=float, default=0, help="Shape of Tukey window to apply to data (default is no windowing)")
optp.add_option("-m", "--time-marginalization", action="store_true", help="Perform marginalization over time via direct numerical integration. Default is false.")
optp.add_option("--vectorized", action="store_true", help="Perform manipulations of lm and timeseries using numpy arrays, not LAL data structures.  (Combine with --gpu to enable GPU use, where available)")
optp.add_option("--gpu", action="store_true", help="Perform manipulations of lm and timeseries using numpy arrays, CONVERTING TO GPU when available. You MUST use this option with --vectorized (otherwise it is a no-op). You MUST have a suitable version of cupy installed, your cuda operational, etc")
optp.add_option("--force-xpy", action="store_true", help="Use the xpy code path.  Use with --vectorized --gpu to use the fallback CPU-based code path. Useful for debugging.")
optp.add_option("-o", "--output-file", help="Save result to this file.")
optp.add_option("-O", "--output-format", default='xml', help="[xml|hdf5]")
optp.add_option("--verbose",action='store_true')
#
# Add the integration options
#
integration_params = OptionGroup(optp, "Integration Parameters", "Control the integration with these options.")
# Default is actually None, but that tells the integrator to go forever or until n_eff is hit.
integration_params.add_option("--interpolate-time", default=False,help="If using time marginalization, compute using a continuously-interpolated array. (Default=false)")
optp.add_option_group(integration_params)

#
# Add the intrinsic parameters
#
intrinsic_params = OptionGroup(optp, "Intrinsic Parameters", "Intrinsic parameters (e.g component mass) to use.")
#intrinsic_params.add_option("--pin-distance-to-sim",action='store_true', help="Pin *distance* value to sim entry. Used to enable source frame reconstruction with NR.")
intrinsic_params.add_option("--mass1", type=float, help="Value of first component mass, in solar masses. Required if not providing coinc tables.")
intrinsic_params.add_option("--mass2", type=float, help="Value of second component mass, in solar masses. Required if not providing coinc tables.")
intrinsic_params.add_option("--eff-lambda", type=float, help="Value of effective tidal parameter. Optional, ignored if not given.")
intrinsic_params.add_option("--deff-lambda", type=float, help="Value of second effective tidal parameter. Optional, ignored if not given")
optp.add_option_group(intrinsic_params)

#
# Add options to integrate over intrinsic parameters.  Same conventions as util_ManualOverlapGrid.py.
# Parameters have special names, and we adopt priors that use those names.
# NOTE: Only 'q' implemented
#
intrinsic_int_params = OptionGroup(optp, "Intrinsic integrated parameters", "Intrinsic parameters to integrate over. ONLY currently used with ROM version")
intrinsic_int_params.add_option("--parameter",action='append')
intrinsic_int_params.add_option("--parameter-range",action='append',type=str)
intrinsic_int_params.add_option("--adapt-intrinsic",action='store_true')
optp.add_option_group(intrinsic_int_params)

#
# Add the pinnable parameters
#
pinnable = OptionGroup(optp, "Pinnable Parameters", "Specifying these command line options will pin the value of that parameter to the specified value with a probability of unity.")
for pin_param in LIKELIHOOD_PINNABLE_PARAMS:
    option = "--" + pin_param.replace("_", "-")
    pinnable.add_option(option, type=float, help="Pin the value of %s." % pin_param)
optp.add_option_group(pinnable)

opts, args = optp.parse_args()

if opts.gpu and xpy_default is numpy:
    print(" Override --gpu  (not available);  use --force-xpy to require the identical code path is used (with xpy =[np|cupy]")
    opts.gpu=False
    if opts.force_xpy:
        opts.gpu=True

deltaT = None
fSample= opts.srate # change sampling rate
if not(fSample is None):
  deltaT =1./fSample

# Load in restricted mode set, if available
restricted_mode_list=None
if not(opts.restricted_mode_list_file is None):
    modes =numpy.loadtxt(opts.restricted_mode_list_file,dtype=int) # columns are l m.  Must contain all. Only integers obviously
    restricted_mode_list = [ (l,m) for l,m in modes]
    print(" RESTRICTED MODE LIST target :", restricted_mode_list)

intrinsic_param_names = opts.parameter
valid_intrinsic_param_names = ['q']
if intrinsic_param_names:
 for param in intrinsic_param_names:

    # Check if in the valid list
    if not(param in valid_intrinsic_param_names):
            print(' Invalid param ', param, ' not in ', valid_intrinsic_param_names)
            sys.exit(0)
    param_ranges = []
    if len(intrinsic_param_names) == len(opts.parameter_range):
        param_ranges = numpy.array(map(eval, opts.parameter_range))
        # Rescale mass-dependent ranges to SI units
        for p in ['mc', 'm1', 'm2', 'mtot']:
          if p in intrinsic_param_names:
            indx = intrinsic_param_names.index(p)
            param_ranges[indx]= numpy.array(param_ranges[indx])* lal.MSUN_SI

# Check both or neither of --data-start/end-time given
if opts.data_start_time is None and opts.data_end_time is not None:
    raise ValueError("You must provide both or neither of --data-start-time and --data-end-time.")
if opts.data_end_time is None and opts.data_start_time is not None:
    raise ValueError("You must provide both or neither of --data-start-time and --data-end-time.")

#
# Import NR grid
#
NR_template_group=None
NR_template_param=None
if opts.nr_group and opts.nr_param:
    import NRWaveformCatalogManager3 as nrwf
    NR_template_group = opts.nr_group
    if nrwf.internal_ParametersAreExpressions[NR_template_group]:
        NR_template_param = eval(opts.nr_param)
    else:
        NR_template_param = opts.nr_param

#
# Hardcoded variables
#
template_min_freq = opts.fmin_template # minimum frequency of template
#t_ref_wind = 50e-3 # Interpolate in a window +/- this width about event time.
t_ref_wind = opts.data_integration_window_half
T_safety = 2. # Safety buffer (in sec) for wraparound corruption

#
# Inverse spectrum truncation control
#
T_spec = opts.inv_spec_trunc_time
if T_spec == 0.: # Do not do inverse spectrum truncation
    inv_spec_trunc_Q = False
    T_safety += 8. # Add a bit more safety buffer in this case
else:
    inv_spec_trunc_Q = True

if opts.event_time is not None:
    event_time = glue.lal.LIGOTimeGPS(opts.event_time)
    print("Event time from command line: %s" % str(event_time))
else:
    print(" Error! ")
    sys.exit(0)

#
# Template descriptors
#

fiducial_epoch = lal.LIGOTimeGPS()
fiducial_epoch = event_time.seconds + 1e-9*event_time.nanoseconds   # no more direct access to gpsSeconds

# Struct to hold template parameters
P_list = None
if opts.sim_xml:
    print("====Loading injection XML:", opts.sim_xml, opts.event, " =======")
    P_list = lalsimutils.xml_to_ChooseWaveformParams_array(str(opts.sim_xml))
    if  len(P_list) < opts.event+opts.n_events_to_analyze:
        print(" Event list of range; soft exit")
        sys.exit(0)
    P_list = P_list[opts.event:(opts.event+opts.n_events_to_analyze)]
    for P in P_list:
        P.radec =False  # do NOT propagate the epoch later
        P.fref = opts.reference_freq
        P.fmin = template_min_freq
        P.tref = fiducial_epoch  # the XML table
        m1 = P.m1/lal.MSUN_SI
        m2 =P.m2/lal.MSUN_SI
        lambda1, lambda2 = P.lambda1,P.lambda2
        P.dist = factored_likelihood.distMpcRef * 1.e6 * lal.PC_SI   # use *nonstandard* distance
        P.phi=0.0
        P.psi=0.0
        P.incl = 0.0       # only works for aligned spins. Be careful.
        P.fref = opts.reference_freq
        if opts.approximant != "TaylorT4": # not default setting
            P.approx = lalsimutils.lalsim.GetApproximantFromString(opts.approximant)  # allow user to override the approx setting. Important for NR followup, where no approx set in sim_xml!

    P = P_list[0]  # Load in the physical parameters of the injection.

# User requested bounds for data segment
if not (opts.data_start_time == None) and  not (opts.data_end_time == None):
    start_time =  opts.data_start_time
    end_time =  opts.data_end_time
    print("Fetching data segment with start=", start_time)
    print("                             end=", end_time)

# Automatically choose data segment bounds so region of interest isn't corrupted
# FIXME: Use estimate, instead of painful full waveform generation call here.
else:
    approxTmp = P.approx
    T_tmplt = lalsimutils.estimateWaveformDuration(P) + 4  # Much more robust than the previous (slow, prone-to-crash approach)
#     if (m1+m2)<30:
#         P.approx = lalsimutils.lalsim.TaylorT4  # should not impact length much.  IMPORTANT because EOB calls will fail at the default sampling rate
#         htmplt = lalsimutils.hoft(P)   # Horribly wasteful waveform gneeration solely to estimate duration.  Will also crash if spins are used.
#         P.approx = approxTmp
#         T_tmplt = - float(htmplt.epoch)    # ASSUMES time returned by hlm is a RELATIVE time that does not add tref. P.radec=False !
#         print fiducial_epoch, event_time, htmplt.epoch
#     else:
#         print " Using estimate for waveform length in a high mass regime: beware!"
#         T_tmplt = lalsimutils.estimateWaveformDuration(P) + 4
    T_seg = T_tmplt + T_spec + T_safety # Amount before and after event time
#    if opts.use_external_EOB:
#        T_seg *=2   # extra safety factor
    start_time = float(event_time) - T_seg
    end_time = float(event_time) + T_seg
    print("Fetching data segment with start=", start_time)
    print("                             end=", end_time)
    print("\t\tEvent time is: ", float(event_time))
    print("\t\tT_seg is: ", T_seg)

#
# Load in data and PSDs
#
data_dict, psd_dict = {}, {}

for inst, chan in map(lambda c: c.split("="), opts.channel_name):
    print("Reading channel %s from cache %s" % (inst+":"+chan, opts.cache_file))
    data_dict[inst] = lalsimutils.frame_data_to_non_herm_hoff(opts.cache_file,
            inst+":"+chan, start=start_time, stop=end_time,
            window_shape=opts.window_shape,deltaT=deltaT)
    print("Frequency binning: %f, length %d" % (data_dict[inst].deltaF,
            data_dict[inst].data.length))

flow_ifo_dict = {}
if opts.fmin_ifo:
 for inst, freq_str in map(lambda c: c.split("="), opts.fmin_ifo):
    freq_low_here = float(freq_str)
    print("Reading low frequency cutoff for instrument %s from %s" % (inst, freq_str), freq_low_here)
    flow_ifo_dict[inst] = freq_low_here

for inst, psdf in map(lambda c: c.split("="), opts.psd_file):
    print("Reading PSD for instrument %s from %s" % (inst, psdf))
    psd_dict[inst] = lalsimutils.get_psd_series_from_xmldoc(psdf, inst)

    deltaF = data_dict[inst].deltaF
    psd_dict[inst] = lalsimutils.resample_psd_series(psd_dict[inst], deltaF)
    print("PSD deltaF after interpolation %f" % psd_dict[inst].deltaF)

    # implement cutoff.
    if inst in flow_ifo_dict.keys():
        if isinstance(psd_dict[inst], lal.REAL8FrequencySeries):
            psd_fvals = psd_dict[inst].f0 + deltaF*numpy.arange(psd_dict[inst].data.length)
            psd_dict[inst].data.data[ psd_fvals < flow_ifo_dict[inst]] = 0 #
        else:
            print('FAIL')
            sys.exit(0)
#        elif isinstance(psd_dict[inst], pylal.xlal.datatypes.real8frequencyseries.REAL8FrequencySeries):  # for backward compatibility
#            psd_fvals = psd_dict[inst].f0 + deltaF*numpy.arange(len(psd_dict[inst].data))
#            psd_dict[inst].data[psd_fvals < ifo_dict[inst]] =0

    assert psd_dict[inst].deltaF == deltaF

    # Highest freq. at which PSD is defined
    # if isinstance(psd_dict[inst],
    #         pylal.xlal.datatypes.real8frequencyseries.REAL8FrequencySeries):
    #     fmax = psd_dict[inst].f0 + deltaF * (len(psd_dict[inst].data) - 1)
    if isinstance(psd_dict[inst], lal.REAL8FrequencySeries):
        fmax = psd_dict[inst].f0 + deltaF * (psd_dict[inst].data.length - 1)

    # Assert upper limit of IP integral does not go past where PSD defined
    assert opts.fmax is None or opts.fmax<= fmax
    # Allow us to target a smaller upper limit than provided by the PSD. Important for numerical PSDs that turn over at high frequency
    if opts.fmax and opts.fmax < fmax:
        fmax = opts.fmax # fmax is now the upper freq. of IP integral

# Ensure data and PSDs keyed to same detectors
if sorted(psd_dict.keys()) != sorted(data_dict.keys()):
    print("Got a different set of instruments based on data and PSDs provided.", file=sys.stderr)

# Ensure waveform has same sample rate, padded length as data
#
# N.B. This assumes all detector data has same sample rate, length
#
# data_dict holds 2-sided FrequencySeries, so their length is the same as
# that of the original TimeSeries that was FFT'd = Nsamples
# Also, deltaF = 1/T, with T = the duration (in sec) of the original TimeSeries
# Therefore 1/(data.length*deltaF) = T/Nsamples = deltaT
P.deltaT = 1./ (data_dict[data_dict.keys()[0]].data.length * deltaF)
P.deltaF = deltaF
for Psig in P_list:
  Psig.deltaT = P.deltaT
  Psig.deltaf = P.deltaF

def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_trunc_Q=inv_spec_trunc_Q, T_spec=T_spec):
    nEvals=0
    P = P_list[indx_event]
    # Precompute
    t_window = 0.15
    rholms_intp, cross_terms, cross_terms_V,  rholms, rest=factored_likelihood.PrecomputeLikelihoodTerms(
            fiducial_epoch, t_window, P, data_dict, psd_dict, opts.l_max, fmax,
            False, inv_spec_trunc_Q, T_spec,
            NR_group=NR_template_group,NR_param=NR_template_param,
            use_external_EOB=opts.use_external_EOB,nr_lookup=opts.nr_lookup,nr_lookup_valid_groups=opts.nr_lookup_group,perturbative_extraction=opts.nr_perturbative_extraction,use_provided_strain=opts.nr_use_provided_strain,hybrid_use=opts.nr_hybrid_use,hybrid_method=opts.nr_hybrid_method,ROM_group=opts.rom_group,ROM_param=opts.rom_param,ROM_use_basis=opts.rom_use_basis,verbose=opts.verbose,quiet=not opts.verbose,ROM_limit_basis_size=opts.rom_limit_basis_size_to,no_memory=opts.no_memory,skip_interpolation=opts.vectorized)

    # Only allocate the FFT plans ONCE
    ifo0 = psd_dict.keys()[0]
    npts = data_dict[ifo0].data.length
    print(" Allocating FFT forward, reverse plans ", npts)
    fwdplan = lal.CreateForwardREAL8FFTPlan(npts, 0)
    revplan = lal.CreateReverseREAL8FFTPlan(npts, 0)

    for ifo in psd_dict:
        # Plot PSDs
        plt.figure(99) # PSD figure)
        fvals = psd_dict[ifo].f0 + np.arange(psd_dict[ifo].data.length)*psd_dict[ifo].deltaF
        plt.plot(fvals,np.log10(np.sqrt(psd_dict[ifo].data.data)),label=ifo)
        plt.figure(88)
        plt.plot(fvals, fvals**(-14./6.) /psd_dict[ifo].data.data)
        plt.figure(1)

        # Plot inverse spectrum filters
        #     - see lalsimutils code for inv_spec_trunc_Q , copied verbatim
        plt.figure(98)
        IP_here = lalsimutils.ComplexOverlap(P.fmin, fmax, 1./2/P.deltaT, P.deltaF, psd_dict[ifo], False,
                                             inv_spec_trunc_Q, T_spec)
        WFD = lal.CreateCOMPLEX16FrequencySeries('FD root inv. spec.',
                                                 lal.LIGOTimeGPS(0.), 0., IP_here.deltaF,
                                                 lal.DimensionlessUnit, IP_here.len1side)
        WTD = lal.CreateREAL8TimeSeries('TD root inv. spec.',
                    lal.LIGOTimeGPS(0.), 0., IP_here.deltaT,
                    lal.DimensionlessUnit, IP_here.len2side)
        print(ifo, IP_here.len2side)
#        if not(fwdplan is None):
        WFD.data.data[:] = np.sqrt(IP_here.weights) # W_FD is 1/sqrt(S_n(f))
        WFD.data.data[0] = WFD.data.data[-1] = 0. # zero 0, f_Nyq bins
        lal.REAL8FreqTimeFFT(WTD, WFD, revplan) # IFFT to TD
        tvals = IP_here.deltaT*np.arange(IP_here.len2side)
        plt.plot(tvals, np.log10(np.abs(WTD.data.data)),label=ifo)

        # Plot Q's
        plt.figure(1);
        plt.clf()
        for mode in rholms[ifo]:
            Q_here = rholms[ifo][mode]
            tvals = lalsimutils.evaluate_tvals(Q_here) - P.tref
            plt.plot(tvals, np.abs(Q_here.data.data),label=ifo+str(mode[0])+","+str(mode[1]))
        plt.legend();
        plt.xlabel("t (s)"); plt.xlim(-0.1,0.1); plt.xlabel("$Q_{lm}$")
        plt.title(ifo)
        plt.savefig(ifo+ "_Q.png")

    plt.figure(99); plt.legend(); plt.xlabel("f(Hz)"); plt.ylabel(r'$\sqrt{S_h}$'); plt.ylim(-24,-21)
    plt.savefig("PSD_plots.png")
    plt.xlim(5,800)
    plt.savefig("PSD_plots_detail.png")
    plt.xlim(10,400); plt.ylim(-23.5,-22)
    plt.savefig("PSD_plots_detail_fine.png")
    plt.figure(98); plt.legend(); plt.xlabel("t (s)")
    plt.savefig("PSD_inv_plots.png")
    plt.figure(88); plt.legend(); plt.xlabel("f (Hz)"); plt.ylabel(r'$(S_h f^{15/6})^{-1}$')
    plt.savefig("PSD_weight_f.png")

lnL_sofar = -np.inf
no_adapt_sky = False
for indx in numpy.arange(len(P_list)):
  res = analyze_event(P_list, indx, data_dict, psd_dict, fmax, opts)


import argparse

import numpy as np

from mtools.mplotqueries.plottypes.base_type import BasePlotType
from mtools.util.log2code import Log2CodeConverter

try:
    from matplotlib.dates import date2num, num2date
except ImportError:
    raise ImportError("Can't import matplotlib. See "
                      "github.com/rueckstiess/mtools/blob/master/INSTALL.md "
                      "for instructions how to install matplotlib or try "
                      "mlogvis instead, which is a simplified version of "
                      "mplotqueries that visualizes the logfile in a "
                      "web browser.")

class HistogramPlotType(BasePlotType):
    """
    Plot a histogram plot over all logevents.

    The bucket size can be specified with the --bucketsize or -b parameter.
    Unit is in seconds.
    """

    plot_type_str = 'histogram'
    timeunits = {'sec': 1, 's': 1, 'min': 60, 'm': 1, 'hour': 3600, 'h': 3600,
                 'day': 86400, 'd': 86400}
    sort_order = 1
    default_group_by = 'namespace'
    l2cc = Log2CodeConverter()

    def __init__(self, args=None, unknown_args=None):
        BasePlotType.__init__(self, args, unknown_args)

        # parse arguments further to get --bucketsize argument
        argparser = argparse.ArgumentParser("mplotqueries --type histogram")
        argparser.add_argument('--bucketsize', '-b', action='store',
                               metavar='SIZE',
                               help="histogram bucket size in seconds",
                               default=60)
        argparser.add_argument('--no-stacked', action='store_true',
                               help=("switch graph mode from stacked "
                                     "histogram (default) to side-by-side "
                                     "histograms."), default=False)
        sub_args = vars(argparser.parse_args(unknown_args))

        self.logscale = args['logscale']
        # get bucket size, either as int (seconds) or as string
        # (see timeunits above)
        bs = sub_args['bucketsize']
        try:
            self.bucketsize = int(bs)
        except ValueError:
            self.bucketsize = self.timeunits[bs]
        self.barstacked = not sub_args['no_stacked']

        self.ylabel = "# lines per %i second bin" % self.bucketsize

    def accept_line(self, logevent):
        """
        Return True for each line.

        We bucket everything. Filtering has to be done before passing to this
        type of plot.
        """
        return True

    def log2code(self, logevent):
        codeline = self.l2cc(logevent.line_str)
        if codeline:
            return ' ... '.join(codeline.pattern)
        else:
            return None

    def plot_group(self, group, idx, axis):
        raise NotImplementedError("Not implemented for histogram plots.")

    def plot(self, axis, ith_plot, total_plots, limits):
        """
        Plot the histogram as a whole over all groups.

        Do not plot as individual groups like other plot types.
        """
        print(self.plot_type_str.upper() + " plot")
        print("%5s %9s  %s" % ("id", " #points", "group"))

        for idx, group in enumerate(self.groups):
            print("%5s %9s  %s" % (idx + 1, len(self.groups[group]), group))

        print('')

        datasets = []
        colors = []
        minx = np.inf
        maxx = -np.inf

        for idx, group in enumerate(self.groups):
            x = date2num([logevent.datetime
                          for logevent in self.groups[group]])
            minx = min(minx, min(x))
            maxx = max(maxx, max(x))
            datasets.append(x)
            color, marker = self.color_map(group)
            colors.append(color)

        if total_plots > 1:
            # if more than one plot, move histogram to twin axis on the right
            twin_axis = axis.twinx()
            twin_axis.set_ylabel(self.ylabel)
            axis.set_zorder(twin_axis.get_zorder() + 1)  # put ax ahead of ax2
            axis.patch.set_visible(False)  # hide the 'canvas'
            axis = twin_axis

        n_bins = max(1, int((maxx - minx) * 24. * 60. * 60. / self.bucketsize))
        if n_bins > 1000:
            # warning for too many buckets
            print("warning: %i buckets, will take a while to render. "
                  "consider increasing --bucketsize." % n_bins)

        n, bins, artists = axis.hist(datasets, bins=n_bins, align='mid',
                                     log=self.logscale,
                                     histtype="barstacked"
                                     if self.barstacked else "bar",
                                     color=colors, edgecolor="none",
                                     linewidth=0, alpha=0.8, picker=True,
                                     label=map(str, self.groups.keys()))

        # scale current y-axis to match min and max values
        axis.set_ylim(np.min(n), np.max(n))

        # add meta-data for picking
        if len(self.groups) > 1:
            for g, group in enumerate(self.groups.keys()):
                for i in range(len(artists[g])):
                    artists[g][i]._mt_plot_type = self
                    artists[g][i]._mt_group = group
                    artists[g][i]._mt_n = n[g][i]
                    if self.barstacked:
                        artists[g][i]._mt_n -= (n[g - 1][i] if g > 0 else 0)

                    artists[g][i]._mt_bin = bins[i]
        else:
            for i in range(len(artists)):
                artists[i]._mt_plot_type = self
                artists[i]._mt_group = group
                artists[i]._mt_n = n[i]
                artists[i]._mt_bin = bins[i]

        return artists

    def clicked(self, event):
        """Print group name and number of items in bin."""
        group = event.artist._mt_group
        n = event.artist._mt_n
        dt = num2date(event.artist._mt_bin)
        print("%4i %s events in %s sec beginning at %s"
              % (n, group, self.bucketsize, dt.strftime("%b %d %H:%M:%S")))