import os
from copy import deepcopy

import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
import torchvision
from PIL import Image
from skimage import color
from skimage.measure import compare_psnr, compare_ssim
from torch.autograd import Variable

import models
import utils
from models.downsampler import Downsampler

class DGP(object):

    def __init__(self, config):
        self.rank, self.world_size = 0, 1
        if config['dist']:
            self.rank = dist.get_rank()
            self.world_size = dist.get_world_size()
        self.config = config
        self.mode = config['dgp_mode']
        self.update_G = config['update_G']
        self.update_embed = config['update_embed']
        self.iterations = config['iterations']
        self.ftr_num = config['ftr_num']
        self.ft_num = config['ft_num']
        self.lr_ratio = config['lr_ratio']
        self.G_lrs = config['G_lrs']
        self.z_lrs = config['z_lrs']
        self.use_in = config['use_in']
        self.select_num = config['select_num']
        self.factor = 2 if self.mode == 'hybrid' else 4  # Downsample factor

        # create model
        self.G = models.Generator(**config).cuda()
        self.D = models.Discriminator(
            **config).cuda() if config['ftr_type'] == 'Discriminator' else None
        self.G.optim = torch.optim.Adam(
            [{'params': self.G.get_params(i, self.update_embed)}
                for i in range(len(self.G.blocks) + 1)],
            lr=config['G_lr'],
            betas=(config['G_B1'], config['G_B2']),
            weight_decay=0,
            eps=1e-8)

        # load weights
        if config['random_G']:
            self.random_G()
        else:
            utils.load_weights(
                self.G if not (config['use_ema']) else None,
                self.D,
                config['weights_root'],
                name_suffix=config['load_weights'],
                G_ema=self.G if config['use_ema'] else None,
                strict=False)

        self.G.eval()
        if self.D is not None:
            self.D.eval()
        self.G_weight = deepcopy(self.G.state_dict())

        # prepare latent variable and optimizer
        self._prepare_latent()
        # prepare learning rate scheduler
        self.G_scheduler = utils.LRScheduler(self.G.optim, config['warm_up'])
        self.z_scheduler = utils.LRScheduler(self.z_optim, config['warm_up'])

        # loss functions
        self.mse = torch.nn.MSELoss()
        if config['ftr_type'] == 'Discriminator':
            self.ftr_net = self.D
            self.criterion = utils.DiscriminatorLoss(ftr_num=config['ftr_num'][0])
        else:
            vgg = torchvision.models.vgg16(pretrained=True).cuda().eval()
            self.ftr_net = models.subsequence(vgg.features, last_layer='20')
            self.criterion = utils.PerceptLoss()

        # Downsampler for producing low-resolution image
        self.downsampler = Downsampler(
            n_planes=3,
            factor=self.factor,
            kernel_type='lanczos2',
            phase=0.5,
            preserve_size=True).type(torch.cuda.FloatTensor)

    def _prepare_latent(self):
        self.z = torch.zeros((1, self.G.dim_z)).normal_().cuda()
        self.z = Variable(self.z, requires_grad=True)
        self.z_optim = torch.optim.Adam(
            [{'params': self.z, 'lr': self.z_lrs[0]}],
            betas=(self.config['G_B1'], self.config['G_B2']),
            weight_decay=0,
            eps=1e-8
        )
        self.y = torch.zeros(1).long().cuda()

    def reset_G(self):
        self.G.load_state_dict(self.G_weight, strict=False)
        self.G.reset_in_init()
        if self.config['random_G']:
            self.G.train()
        else:
            self.G.eval()

    def random_G(self):
        self.G.init_weights()

    def set_target(self, target, category, img_path):
        self.target_origin = target
        # apply degradation transform to the original image
        self.target = self.pre_process(target, True)
        self.y.fill_(category.item())
        self.img_name = img_path[img_path.rfind('/') + 1:img_path.rfind('.')]

    def run(self, save_interval=None):
        save_imgs = self.target.clone()
        save_imgs2 = save_imgs.cpu().clone()
        loss_dict = {}
        curr_step = 0
        count = 0
        for stage, iteration in enumerate(self.iterations):
            # setup the number of features to use in discriminator
            self.criterion.set_ftr_num(self.ftr_num[stage])

            for i in range(iteration):
                curr_step += 1
                # setup learning rate
                self.G_scheduler.update(curr_step, self.G_lrs[stage],
                                        self.ft_num[stage], self.lr_ratio[stage])
                self.z_scheduler.update(curr_step, self.z_lrs[stage])

                self.z_optim.zero_grad()
                if self.update_G:
                    self.G.optim.zero_grad()
                x = self.G(self.z, self.G.shared(self.y), use_in=self.use_in[stage])
                # apply degradation transform
                x_map = self.pre_process(x, False)

                # calculate losses in the degradation space
                ftr_loss = self.criterion(self.ftr_net, x_map, self.target)
                mse_loss = self.mse(x_map, self.target)
                # nll corresponds to a negative log-likelihood loss
                nll = self.z**2 / 2
                nll = nll.mean()
                l1_loss = F.l1_loss(x_map, self.target)
                loss = ftr_loss * self.config['w_D_loss'][stage] + \
                    mse_loss * self.config['w_mse'][stage] + \
                    nll * self.config['w_nll']
                loss.backward()

                self.z_optim.step()
                if self.update_G:
                    self.G.optim.step()

                # These losses are calculated in the [-1,1] image scale
                # We record the rescaled MSE and L1 loss, corresponding to [0,1] image scale
                loss_dict = {
                    'ftr_loss': ftr_loss,
                    'nll': nll,
                    'mse_loss': mse_loss / 4,
                    'l1_loss': l1_loss / 2
                }

                # calculate losses in the non-degradation space
                if self.mode in ['reconstruct', 'colorization', 'SR', 'inpainting']:
                    # x2 is to get the post-processed result in colorization
                    metrics, x2 = self.get_metrics(x)
                    loss_dict = {**loss_dict, **metrics}

                if i == 0 or (i + 1) % self.config['print_interval'] == 0:
                    if self.rank == 0:
                        print(', '.join(
                            ['Stage: [{0}/{1}]'.format(stage + 1, len(self.iterations))] +
                            ['Iter: [{0}/{1}]'.format(i + 1, iteration)] +
                            ['%s : %+4.4f' % (key, loss_dict[key]) for key in loss_dict]
                        ))
                    # save image sheet of the reconstruction process
                    save_imgs = torch.cat((save_imgs, x), dim=0)
                    torchvision.utils.save_image(
                        save_imgs.float(),
                        '%s/images_sheet/%s_%s.jpg' %
                        (self.config['exp_path'], self.img_name, self.mode),
                        nrow=int(save_imgs.size(0)**0.5),
                        normalize=True)
                    if self.mode == 'colorization':
                        save_imgs2 = torch.cat((save_imgs2, x2), dim=0)
                        torchvision.utils.save_image(
                            save_imgs2.float(),
                            '%s/images_sheet/%s_%s_2.jpg' %
                            (self.config['exp_path'], self.img_name, self.mode),
                            nrow=int(save_imgs.size(0)**0.5),
                            normalize=True)

                if save_interval is not None:
                    if i == 0 or (i + 1) % save_interval[stage] == 0:
                        count += 1
                        save_path = '%s/images/%s' % (self.config['exp_path'],
                                                      self.img_name)
                        if not os.path.exists(save_path):
                            os.makedirs(save_path)
                        img_path = os.path.join(
                            save_path, '%s_%03d.jpg' % (self.img_name, count))
                        utils.save_img(x[0], img_path)

                # stop the reconstruction if the loss reaches a threshold
                if mse_loss.item() < self.config['stop_mse'] or ftr_loss.item(
                ) < self.config['stop_ftr']:
                    break

        # save images
        utils.save_img(
            self.target[0], '%s/images/%s_%s_target.png' %
            (self.config['exp_path'], self.img_name, self.mode))
        utils.save_img(
            self.target_origin[0],
            '%s/images/%s_%s_target_origin.png' %
            (self.config['exp_path'], self.img_name, self.mode))
        utils.save_img(
            x[0], '%s/images/%s_%s.png' %
            (self.config['exp_path'], self.img_name, self.mode))
        if self.mode == 'colorization':
            utils.save_img(
                x2[0], '%s/images/%s_%s2.png' %
                (self.config['exp_path'], self.img_name, self.mode))

        if self.mode == 'jitter':
            # conduct random jittering
            self.jitter(x)
        if self.config['save_G']:
            torch.save(
                self.G.state_dict(), '%s/G_%s_%s.pth' %
                (self.config['exp_path'], self.img_name, self.mode))
            torch.save(
                self.z, '%s/z_%s_%s.pth' %
                (self.config['exp_path'], self.img_name, self.mode))
        return loss_dict

    def select_z(self, select_y=False):
        with torch.no_grad():
            if self.select_num == 0:
                self.z.zero_()
                return
            elif self.select_num == 1:
                self.z.normal_()
                return
            z_all, y_all, loss_all = [], [], []
            if self.rank == 0:
                print('Selecting z from {} samples'.format(self.select_num))
            # only use last 3 discriminator features to compare
            self.criterion.set_ftr_num(3)
            for i in range(self.select_num):
                self.z.normal_(mean=0, std=self.config['sample_std'])
                z_all.append(self.z.cpu())
                if select_y:
                    self.y.random_(0, self.config['n_classes'])
                    y_all.append(self.y.cpu())
                x = self.G(self.z, self.G.shared(self.y))
                x = self.pre_process(x)
                ftr_loss = self.criterion(self.ftr_net, x, self.target)
                loss_all.append(ftr_loss.view(1).cpu())
                if self.rank == 0 and (i + 1) % 100 == 0:
                    print('Generating {}th sample'.format(i + 1))
            loss_all = torch.cat(loss_all)
            idx = torch.argmin(loss_all)
            self.z.copy_(z_all[idx])
            if select_y:
                self.y.copy_(y_all[idx])
            self.criterion.set_ftr_num(self.ftr_num[0])

    def pre_process(self, image, target=True):
        if self.mode in ['SR', 'hybrid']:
            # apply downsampling, this part is the same as deep image prior
            if target:
                image_pil = utils.np_to_pil(
                    utils.torch_to_np((image.cpu() + 1) / 2))
                LR_size = [
                    image_pil.size[0] // self.factor,
                    image_pil.size[1] // self.factor
                ]
                img_LR_pil = image_pil.resize(LR_size, Image.ANTIALIAS)
                image = utils.np_to_torch(utils.pil_to_np(img_LR_pil)).cuda()
                image = image * 2 - 1
            else:
                image = self.downsampler((image + 1) / 2)
                image = image * 2 - 1
            # interpolate to the orginal resolution via bilinear interpolation
            image = F.interpolate(
                image, scale_factor=self.factor, mode='bilinear')
        n, _, h, w = image.size()
        if self.mode in ['colorization', 'hybrid']:
            # transform the image to gray-scale
            r = image[:, 0, :, :]
            g = image[:, 1, :, :]
            b = image[:, 2, :, :]
            gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
            image = gray.view(n, 1, h, w).expand(n, 3, h, w)
        if self.mode in ['inpainting', 'hybrid']:
            # remove the center part of the image
            hole = min(h, w) // 3
            begin = (h - hole) // 2
            end = h - begin
            self.begin, self.end = begin, end
            mask = torch.ones(1, 1, h, w).cuda()
            mask[0, 0, begin:end, begin:end].zero_()
            image = image * mask
        return image

    def get_metrics(self, x):
        with torch.no_grad():
            l1_loss_origin = F.l1_loss(x, self.target_origin) / 2
            mse_loss_origin = self.mse(x, self.target_origin) / 4
            metrics = {
                'l1_loss_origin': l1_loss_origin,
                'mse_loss_origin': mse_loss_origin
            }
            # transfer to numpy array and scale to [0, 1]
            target_np = (self.target_origin.detach().cpu().numpy()[0] + 1) / 2
            x_np = (x.detach().cpu().numpy()[0] + 1) / 2
            target_np = np.transpose(target_np, (1, 2, 0))
            x_np = np.transpose(x_np, (1, 2, 0))
            if self.mode == 'colorization':
                # combine the 'ab' dim of x with the 'L' dim of target image
                x_lab = color.rgb2lab(x_np)
                target_lab = color.rgb2lab(target_np)
                x_lab[:, :, 0] = target_lab[:, :, 0]
                x_np = color.lab2rgb(x_lab)
                x = torch.Tensor(np.transpose(x_np, (2, 0, 1))) * 2 - 1
                x = x.unsqueeze(0)
            elif self.mode == 'inpainting':
                # only use the inpainted area to calculate ssim and psnr
                x_np = x_np[self.begin:self.end, self.begin:self.end, :]
                target_np = target_np[self.begin:self.end,
                                      self.begin:self.end, :]
            ssim = compare_ssim(target_np, x_np, multichannel=True)
            psnr = compare_psnr(target_np, x_np)
            metrics['psnr'] = torch.Tensor([psnr]).cuda()
            metrics['ssim'] = torch.Tensor([ssim]).cuda()
            return metrics, x

    def jitter(self, x):
        save_imgs = x.clone().cpu()
        z_rand = self.z.clone()
        stds = [0.3, 0.5, 0.7]
        save_path = '%s/images/%s_jitter' % (self.config['exp_path'],
                                             self.img_name)
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        with torch.no_grad():
            for std in stds:
                for i in range(30):
                    # add random noise to the latent vector
                    z_rand.normal_()
                    z = self.z + std * z_rand
                    x_jitter = self.G(z, self.G.shared(self.y))
                    utils.save_img(
                        x_jitter[0], '%s/std%.1f_%d.jpg' % (save_path, std, i))
                    save_imgs = torch.cat((save_imgs, x_jitter.cpu()), dim=0)

        torchvision.utils.save_image(
            save_imgs.float(),
            '%s/images_sheet/%s_jitters.jpg' %
            (self.config['exp_path'], self.img_name),
            nrow=int(save_imgs.size(0)**0.5),
            normalize=True)


import torch.utils.data as data
import os
from PIL import Image
import torch
import torchvision.transforms as transforms
import random
from transform import *
import numpy as np

class Kinetics(data.Dataset):
    def __init__(self, video_path, of_path, slice_num, batch_size, frame_num=9, img_size=112, overlap_rate=0.25):
        self.samples = []
        actions = sorted(os.listdir(video_path))

        self.action_emb = {action: actions.index(action) for action in actions}
        for action in actions:
            train_videos = os.listdir(video_path + '/' + action)
            train_videos.sort()

            for video in train_videos:
                self.samples.append(action + '/' + video)

        self.video_path = video_path
        self.of_path = of_path
        self.frame_num = frame_num
        self.img_size = img_size
        self.transform = transforms.Compose([transforms.ToTensor(),
                                             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
        self.argumentation = transforms.Compose([
            RandomHorizontalFlip(p=0.5),
            #RandomVerticalFlip(p=0.2),
            #RandomGrayscale(p=0.4),
            ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
            RandomCenterCrop(size=img_size, p=0.2)
            ])

        self.to_pil_image = transforms.ToPILImage()
        self.slice_num = slice_num
        self.batch_size = batch_size
        self.overlap_rate = overlap_rate

        random.shuffle(self.samples)
        num = len(self.samples)
        if -1 * (num % batch_size) == 0:
            self.samples = self.samples
        else:
            self.samples = self.samples[0: -1 * (num % batch_size)]

        self.samples = np.array(self.samples)
        self.samples = self.samples.reshape(-1, batch_size)

        print("Samples: " + str(self.samples.shape))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        video_names = self.samples[idx]
        length = self.frame_num
        inputs = torch.zeros(self.slice_num, self.batch_size, length, 3, self.img_size, self.img_size)
        ofs = torch.zeros(self.slice_num, self.batch_size, length, 2, self.img_size, self.img_size)
        actions = torch.FloatTensor(self.slice_num, self.batch_size, length)
        overlap_frame_num = torch.FloatTensor(self.slice_num, self.batch_size, 1)

        for b, video_name in enumerate(video_names):
            frames_ = os.listdir(self.video_path + '/' + video_name)
            frames = []
            for frame in frames_:
                if '.jpg' in frame:
                    frames.append(frame)
            frames.sort()

            need_frame = float(length * (1 - self.overlap_rate) * (self.slice_num - 1)) + length
            self.space = max(1, int(round(float(len(frames)) / float(need_frame))))
            imgs = []
            of_imgs = []
            v_length = len(frames)
            for i in range(v_length):
                if i % self.space == self.space - 1:
                    try:
                        img = Image.open(self.video_path + '/' + video_name + '/' + frames[i])
                        of_img = Image.open(self.of_path + '/' + video_name + '/' + frames[i])
                        imgs.append(img)
                        of_imgs.append(of_img)
                    except:
                        pass

            sample = []
            of_sample = []
            sample_area = max(0, len(imgs) - length)
            stride = min(int(length * (1 - self.overlap_rate)), int(sample_area / self.slice_num))

            for i in range(self.slice_num):
                start_point = stride * i

                for j in range(length):
                    n_frame = start_point + j
                    if n_frame < len(imgs):
                        sample.append(imgs[n_frame])
                        of_sample.append(of_imgs[n_frame])
                try:
                    while len(sample) < length * (i + 1):
                        sample.append(sample[-1])
                    while len(of_sample) < length * (i + 1):
                        of_sample.append(of_sample[-1])
                except:
                    print(video_name)

            input = torch.zeros(self.slice_num, length, 3, self.img_size, self.img_size)
            of = torch.zeros(self.slice_num, length, 2, self.img_size, self.img_size)
            sample_with_of_sample = []
            sample_with_of_sample.extend(sample)
            sample_with_of_sample.extend(of_sample)
            sample_with_of_sample = self.argumentation(sample_with_of_sample)
            sample = sample_with_of_sample[: len(sample_with_of_sample) // 2]
            of_sample = sample_with_of_sample[len(sample_with_of_sample) // 2:]

            for i in range(self.slice_num):
                for j in range(length):
                    img = sample[i * length + j].resize((self.img_size, self.img_size))
                    of_img = of_sample[i * length + j].resize((self.img_size, self.img_size))
                    transformed_img = self.transform(img)
                    transformed_of_img = self.transform(of_img)
                    img.close()
                    of_img.close()
                    input[i, j, :, :, :] = transformed_img
                    of[i, j, :, :, :] = transformed_of_img[0:2]

            action = torch.FloatTensor(self.slice_num, length)
            for batch in range(self.slice_num):
                for fn in range(0, length):
                    action[batch, fn] = self.action_emb[video_name.split('/')[0]]

            inputs[:, b] = input
            ofs[:, b] = of
            actions[:, b] = action

            overlap_frame_num[:, b] = torch.ones(self.slice_num, 1) * (length - stride)

        return inputs, ofs, 0, overlap_frame_num, actions


#!/usr/bin/env python
"""
 Copyright 2019 Jesus Villalba (Johns Hopkins University)
 Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
"""

import sys
import os
from jsonargparse import (
    ArgumentParser,
    ActionConfigFile,
    ActionParser,
    namespace_to_dict,
)
import time
import logging

import numpy as np

import torch

from hyperion.hyp_defs import config_logger, float_cpu
from hyperion.utils import Utt2Info
from hyperion.io import DataWriterFactory as DWF
from hyperion.io import SequentialDataReaderFactory as DRF
from hyperion.io import VADReaderFactory as VRF
from hyperion.feats import MeanVarianceNorm as MVN

from hyperion.torch.utils import open_device
from hyperion.torch import TorchModelLoader as TML

def init_device(use_gpu):
    set_float_cpu("float32")
    num_gpus = 1 if use_gpu else 0
    logging.info("initializing devices num_gpus={}".format(num_gpus))
    device = open_device(num_gpus=num_gpus)
    return device

def init_mvn(device, **kwargs):
    mvn_args = MVN.filter_args(**kwargs["mvn"])
    logging.info("mvn args={}".format(mvn_args))
    mvn = MVN(**mvn_args)
    if mvn.norm_mean or mvn.norm_var:
        return mvn
    return None

def load_model(model_path, device):
    logging.info("loading model {}".format(model_path))
    model = TML.load(model_path)
    logging.info("xvector-model={}".format(model))
    model.to(device)
    model.eval()
    return model

def extract_xvectors(
    input_spec,
    output_spec,
    vad_spec,
    write_timestamps_spec,
    slidwin_params_path,
    vad_path_prefix,
    model_path,
    chunk_length,
    embed_layer,
    win_length,
    win_shift,
    snip_edges,
    feat_frame_length,
    feat_frame_shift,
    feat_snip_edges,
    use_gpu,
    **kwargs
):

    logging.info("initializing")
    rng = np.random.RandomState(seed=1123581321 + kwargs["part_idx"])
    device = init_device(use_gpu)
    mvn = init_mvn(device, **kwargs)
    model = load_model(model_path, device)

    if write_timestamps_spec is not None:
        time_writer = DWF.create(write_timestamps_spec, scp_sep=scp_sep)

    dr_args = DRF.filter_args(**kwargs)
    logging.info("opening output stream: %s" % (output_spec))
    with DWF.create(output_spec) as writer:

        logging.info("opening input stream: %s" % (output_spec))
        with DRF.create(input_spec, **dr_args) as reader:
            if vad_spec is not None:
                logging.info("opening VAD stream: %s" % (vad_spec))
                v_reader = VRF.create(vad_spec, path_prefix=vad_path_prefix)

            while not reader.eof():
                t1 = time.time()
                key, data = reader.read(1)
                if len(key) == 0:
                    break
                t2 = time.time()
                logging.info("processing utt %s" % (key[0]))
                x = data[0]
                if mvn is not None:
                    x = mvn.normalize(x)
                t3 = time.time()
                tot_frames = x.shape[0]
                if vad_spec is not None:
                    vad = v_reader.read(key, num_frames=x.shape[0])[0].astype(
                        "bool", copy=False
                    )
                    x = x[vad]

                logging.info(
                    "utt %s detected %d/%d (%.2f %%) speech frames"
                    % (key[0], x.shape[0], tot_frames, x.shape[0] / tot_frames * 100)
                )

                t4 = time.time()
                if x.shape[0] == 0:
                    y = np.zeros(
                        (
                            1,
                            model.embed_dim,
                        ),
                        dtype=float_cpu(),
                    )
                else:
                    xx = torch.tensor(x.T[None, :], dtype=torch.get_default_dtype())
                    with torch.no_grad():
                        y = (
                            model.extract_embed_slidwin(
                                xx,
                                win_length,
                                win_shift,
                                snip_edges=snip_edges,
                                feat_frame_length=feat_frame_length,
                                feat_frame_shift=feat_frame_shift,
                                chunk_length=chunk_length,
                                embed_layer=embed_layer,
                                detach_chunks=True,
                            )
                            .detach()
                            .cpu()
                            .numpy()[0]
                        )

                        # if np.any(np.isnan(y)):
                        #     y = y.T
                        #     idx = np.any(np.isnan(y), axis=1)
                        #     print(y[idx])
                        #     print(np.where(idx))
                        #     raise Exception()
                        # y1 = model.extract_embed(
                        #     xx[:,:,:148],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y.shape, y1.shape))
                        # logging.info('{} {}'.format(y[:20, 0], y1[:20]))
                        # y2 = model.extract_embed(
                        #     xx[:,:,25:173],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 1], y2[:20]))
                        # y3 = model.extract_embed(
                        #      xx[:,:,250:398],
                        #      chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 10], y3[:20]))

                        # win_length = 20
                        # y = model.extract_embed_slidwin(
                        #     xx, win_length, win_shift, snip_edges=True,
                        #     feat_frame_length=feat_frame_length, feat_frame_shift=feat_frame_shift,
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]

                        # y1 = model.extract_embed(
                        #     xx[:,:,:1999],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y.shape, y1.shape))
                        # logging.info('{} {}'.format(y[:20, 0], y1[:20]))
                        # y2 = model.extract_embed(
                        #     xx[:,:,25:2024],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 1], y2[:20]))
                        # y3 = model.extract_embed(
                        #      xx[:,:,250:2249],
                        #      chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 10], y3[:20]))

                        # win_length = 20
                        # y = model.extract_embed_slidwin(
                        #     xx, win_length, win_shift, snip_edges=False,
                        #     feat_frame_length=feat_frame_length, feat_frame_shift=feat_frame_shift,
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]

                        # y1 = model.extract_embed(
                        #     xx[:,:,:1112],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y.shape, y1.shape))
                        # logging.info('{} {}'.format(y[:20, 0], y1[:20]))
                        # y2 = model.extract_embed(
                        #     xx[:,:,25:1037],
                        #     chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 1], y2[:20]))
                        # y3 = model.extract_embed(
                        #      xx[:,:,250:1262],
                        #      chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 10], y3[:20]))

                        # y3 = model.extract_embed(
                        #      xx[:,:,250:1262],
                        #      chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 10], y3[:20]))

                        # y3 = model.extract_embed(
                        #      xx[:,:,2500:3512],
                        #      chunk_length=chunk_length,
                        #     embed_layer=embed_layer, detach_chunks=True).detach().cpu().numpy()[0]
                        # logging.info('{} {}'.format(y[:20, 100], y3[:20]))

                t5 = time.time()
                y = y.T
                writer.write(key, [y])

                if write_timestamps_spec is not None:
                    num_wins = y.shape[0]
                    timestamps = model.compute_slidwin_timestamps(
                        num_wins,
                        win_length,
                        win_shift,
                        snip_edges,
                        feat_frame_length,
                        feat_frame_length,
                        feat_snip_edges,
                    ).numpy()
                    logging.info("{}".format(timestamps))
                    time_writer.write(key, [timestamps])
                t6 = time.time()
                logging.info(
                    (
                        "utt %s total-time=%.3f read-time=%.3f mvn-time=%.3f "
                        "vad-time=%.3f embed-time=%.3f write-time=%.3f "
                        "rt-factor=%.2f"
                    )
                    % (
                        key[0],
                        t6 - t1,
                        t2 - t1,
                        t3 - t2,
                        t4 - t3,
                        t5 - t4,
                        t6 - t5,
                        x.shape[0] * 1e-2 / (t6 - t1),
                    )
                )

    if write_timestamps_spec is not None:
        time_writer.close()

    if slidwin_params_path is not None:
        params = {
            "padding": model.compute_slidwin_left_padding(
                win_length,
                win_shift,
                snip_edges,
                feat_frame_length,
                feat_frame_length,
                feat_snip_edges,
            ),
            "win_length": win_length,
            "win_shift": win_shift,
        }
        with open(slidwin_params_path, "w") as f:
            yaml.dump(params, f)

if __name__ == "__main__":

    parser = ArgumentParser(description="Extract x-vectors over a sliding window")

    parser.add_argument("--cfg", action=ActionConfigFile)
    parser.add_argument("--input", dest="input_spec", required=True)
    DRF.add_class_args(parser)
    parser.add_argument("--vad", dest="vad_spec", default=None)
    parser.add_argument(
        "--write-timestamps", dest="write_timestamps_spec", default=None
    )
    parser.add_argument("--slidwin-params-path", default=None)

    # parser.add_argument('--scp-sep', dest='scp_sep', default=' ',
    #                     help=('scp file field separator'))
    # parser.add_argument('--path-prefix', dest='path_prefix', default=None,
    #                     help=('scp file_path prefix'))
    parser.add_argument(
        "--vad-path-prefix",
        dest="vad_path_prefix",
        default=None,
        help=("scp file_path prefix for vad"),
    )

    MVN.add_class_args(parser, prefix="mvn")

    parser.add_argument("--model-path", required=True)
    parser.add_argument(
        "--win-length",
        type=float,
        default=1.5,
        help=("window length for x-vector extraction in seconds"),
    )
    parser.add_argument(
        "--win-shift",
        type=float,
        default=0.25,
        help=("window shift for x-vector extraction in seconds"),
    )
    parser.add_argument(
        "--snip-edges",
        default=False,
        action="store_true",
        help=(
            "If true, end effects will be handled by outputting "
            "only windows that completely fit in the file, "
            "and the number of windows depends on the window-length. "
            "If false, the number of windows depends only on "
            "the window-shift, and we reflect the data at the ends."
        ),
    )

    parser.add_argument(
        "--feat-frame-length",
        type=float,
        default=25,
        help=("frame-length used to compute the acoustic features in msecs"),
    )
    parser.add_argument(
        "--feat-frame-shift",
        type=float,
        default=10,
        help=("frame-shift used to compute the acoustic features in msecs"),
    )
    parser.add_argument(
        "--feat-snip-edges",
        default=False,
        action="store_true",
        help=(
            "If true, end effects will be handled by outputting only windows "
            "that completely fit in the file, and the number of windows "
            "depends on the feat-frame-length. "
            "If false, the number of feature frames depends only on the "
            "feat-frame-shift, and we reflect the waveform at the ends."
        ),
    )

    parser.add_argument(
        "--chunk-length",
        type=int,
        default=0,
        help=(
            "number of frames used in each forward pass of the x-vector encoder,"
            "if 0 the full utterance is used"
        ),
    )

    parser.add_argument(
        "--embed-layer",
        type=int,
        default=None,
        help=(
            "classifier layer to get the embedding from,"
            "if None the layer set in training phase is used"
        ),
    )

    parser.add_argument("--output", dest="output_spec", required=True)
    parser.add_argument(
        "--use-gpu", default=False, action="store_true", help="extract xvectors in gpu"
    )
    # parser.add_argument('--part-idx', dest='part_idx', type=int, default=1,
    #                     help=('splits the list of files in num-parts and process part_idx'))
    # parser.add_argument('--num-parts', dest='num_parts', type=int, default=1,
    #                     help=('splits the list of files in num-parts and process part_idx'))
    parser.add_argument(
        "-v", "--verbose", dest="verbose", default=1, choices=[0, 1, 2, 3], type=int
    )

    args = parser.parse_args()
    config_logger(args.verbose)
    del args.verbose
    logging.debug(args)

    extract_xvectors(**namespace_to_dict(args))


import plotnine as pn
from analysis import plot
from statsmodels.formula import api as smf
from rebar import parallel
from pavlov import stats
import pandas as pd
import torch
from .arena import common
from boardlaw import learning, main, sql
from rebar import arrdict, dotdict
import numpy as np
from tqdm.auto import tqdm
from logging import getLogger
from multiprocessing import set_start_method

log = getLogger(__name__)

# Smallest networks to make it to within -.5 of perfect play.
RUNS = {
    3: '2021-02-17 19-34-03 valid-ships',
    4: '2021-02-21 05-22-28 watery-drunks',
    5: '2021-02-19 11-15-16 skinny-tactic',
    6: '2021-02-21 08-52-28 these-plow',
    7: '2021-02-19 19-50-36 tan-buffer',
    8: '2021-02-21 22-55-52 double-spoon',
    9: '2021-02-20 21-11-32 intent-nets'}

def stored_agent(agent_id):
    info = sql.query('select * from agents_details where id == ?', int(agent_id)).iloc[0]
    agent = common.agent(info.run, info.idx, 'cuda')
    agent.kwargs['n_nodes'] = info.test_nodes
    agent.kwargs['c_puct'] = info.test_c
    return agent

def stored_worlds(agent_id, n_envs):
    info = sql.query('select * from agents_details where id == ?', int(agent_id)).iloc[0]
    return common.worlds(info.run, n_envs, 'cuda')

def collect(agent_id, n_envs=32*1024):
    agent = stored_agent(agent_id)
    worlds = stored_worlds(agent_id, n_envs)

    buffer = []
    while True:
        while len(buffer) < 64:
            with torch.no_grad():
                decisions = agent(worlds, value=True)
            new_worlds, transition = worlds.step(decisions.actions)

            buffer.append(arrdict.arrdict(
                worlds=worlds,
                decisions=decisions.half(),
                transitions=learning.half(transition)).detach())

            worlds = new_worlds

        chunk, buffer = main.as_chunk(buffer, n_envs)

        mixness = chunk.transitions.terminal.float().mean(1)
        mixness = (mixness.max() - mixness.min())/mixness.median()
        if mixness < .25:
            break

    return agent, chunk

def flat_grad(network, loss):
    for p in network.parameters():
        p.grad = None
    loss.backward(retain_graph=True)
    return torch.cat([p.grad.flatten() for p in network.parameters() if p.grad is not None])

def gradient(network, batch):
    d0 = batch.decisions
    d = network(batch.worlds)

    zeros = torch.zeros_like(d.logits)
    l = d.logits.where(d.logits > -np.inf, zeros)
    l0 = d0.logits.float().where(d0.logits > -np.inf, zeros)

    policy_loss = -(l0.exp()*l).sum(axis=-1).mean()

    target_value = batch.reward_to_go
    value_loss = (target_value - d.v).square().mean()

    return arrdict.arrdict(
        policy=flat_grad(network, policy_loss),
        value=flat_grad(network, value_loss),
        joint=flat_grad(network, policy_loss + value_loss))

def gradients(network, chunk):
    grads = []
    for t in range(chunk.reward_to_go.size(0)):
        grads.append(gradient(network, chunk[t]))
    return arrdict.stack(grads)

def noise_scale_components(chunk, gs, kind):
    T = chunk.reward_to_go.size(0)
    B = chunk.reward_to_go.size(1)
    return pd.Series({
        'kind': kind,
        'mean_sq': float(gs.mean(0).pow(2).mean()),
        'sq_mean': float(gs.pow(2).mean()),
        'variance': float((gs - gs.mean(0, keepdim=True)).pow(2).mean(0).mul(T/(T-1)).mean(0)),
        'n_params': float(gs.shape[1]),
        'batch_size': float(B),
        'batches': float(T)})

def noise_scale(result):
    return result.batch_size*result.variance/result.mean_sq

def evaluate_noise_scale(agent_id):
    extant = sql.query('select * from noise_scales where agent_id == ?', int(agent_id))
    if len(extant) == 0:
        agent, chunk = collect(agent_id)
        gs = gradients(agent.network, chunk)

        results = pd.DataFrame([noise_scale_components(chunk, gs[k], k) for k in gs])
        results['agent_id'] = agent_id
        log.info(f'{agent_id}: {noise_scale(results.iloc[0]):.0f}')
        sql.save_noise_scale(results)

def agents_opponent(agent_id):
    return sql.query('''
        select id from agents
        where snap == (
            select snap from agents where id == ?)
        and nodes == 64
        and c == 1./16''', int(agent_id)).id.iloc[0]

def evaluate_perf(agent_id, n_envs=1024):
    opponent_id = agents_opponent(agent_id)
    extant = sql.query('''
        select * from trials
        where ((black_agent == ?) and (white_agent == ?))
        or ((white_agent == ?) and (black_agent == ?))''',
        int(agent_id), int(opponent_id), int(agent_id), int(opponent_id))
    games = (extant.black_wins + extant.white_wins).sum()
    if games < n_envs:
        a = stored_agent(agent_id)
        o = stored_agent(opponent_id)
        w = stored_worlds(agent_id, n_envs)

        results = common.evaluate(w, [(agent_id, a), (opponent_id, o)])

        sql.save_trials(results)

def evaluate(run, idx, nodes, c_puct, perf=True):
    snap_id = sql.query_one('select id from snaps where run == ? and idx == ?', run, int(idx)).id
    extant = sql.query('select * from agents where snap == ? and nodes == ? and c == ?', int(snap_id), int(nodes), float(c_puct))
    if len(extant) == 0:
        log.info(f'Creating agent run="{run}", idx={idx}, nodes={nodes}, c_puct={c_puct:.3f}')
        sql.execute('insert into agents values (null, ?, ?, ?)', int(snap_id), int(nodes), float(c_puct))
        extant = sql.query('select * from agents where snap == ? and nodes == ? and c == ?', int(snap_id), int(nodes), float(c_puct))
    agent_id = extant.id.iloc[0]

    evaluate_noise_scale(agent_id)
    if perf:
        evaluate_perf(agent_id)

def sweep_trees(boardsize=None):
    if boardsize is None:
        for b in RUNS:
            sweep_trees(b)
        return
    run = RUNS[boardsize]
    snaps = sql.query('select * from snaps where run == ?', run)

    set_start_method('spawn', True)
    with parallel.parallel(evaluate, N=2, executor='cuda', desc=str(boardsize)) as pool:
        futures = {}
        for idx in snaps.idx[snaps.idx % 2 == 0].unique():
            for nodes in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
                for c in [1/64, 1/32, 1/16, 1/8, 1/4, 1/2, 1.]:
                    futures[idx, nodes, c] = pool(run, idx, nodes, c)
        pool.wait(futures)

def sweep_runs():
    runs = sql.query('select * from runs where description like "bee/%"')
    runs = runs[runs.width*runs.depth <= 1024] # Bigger than this blows out my memory. Think there's a leak somewhere.
    runs = runs.sample(frac=1)

    set_start_method('spawn', True)
    for i, run in enumerate(runs.run.unique()):
        snaps = sql.query('select * from snaps where run == ?', run)
        with parallel.parallel(evaluate, N=2, executor='cuda', desc=str(i)) as pool:
            pool.wait([pool(run, idx, 64, 1/16, perf=False) for idx in snaps.idx.unique()])

def node_sweep(run='2021-02-21 09-04-51 wavy-mills', idx=13):
    results = {}
    for c in [1/64, 1/32, 1/16, 1/8, 1/4, 1/2, 1., 2.]:
        for n in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
            results[c, n] = noise_scale_components(run, idx, n_nodes=n, c_puct=c)
    results = pd.concat(results, 0).unstack()
    return results

def relative_elo(g):
    left = g[g.idx == 0].elo.mean()
    g = g.copy()
    g['rel_elo'] = -g.elo/left
    return g

def load():
    from analysis import data

    ags = data.load()

    noise = (sql.query('select * from noise_scales')
                .set_index(['agent_id', 'kind'])
                .pipe(lambda df: df.batch_size*df.variance/df.mean_sq)
                .unstack())

    df = pd.merge(ags, noise, left_index=True, right_index=True, how='inner')
    df['uplift'] = df.groupby('snap_id').apply(lambda g: g.elo - g.query('test_nodes == 1').elo.mean()).reset_index(0, drop=True)
    df['tree_spec'] = df.test_c.astype(str) + '/' + df.test_nodes.astype(str)
    df['params'] = df.width**2 * df.depth
    df = df.groupby('boardsize', as_index=False).apply(relative_elo).reset_index(level=0, drop=True)

    expected = sql.query('select * from snaps').groupby('run').idx.count()
    actual = df.groupby(['run', 'idx']).idx.count().groupby(level=0).count()
    df['complete'] = actual.reindex_like(expected).eq(expected).reindex(df.run.values).values

    return df

def plot_policy(df):
    trunk = df.query('complete & test_nodes == 64 & test_c == 1/16')

    return (pn.ggplot(trunk, pn.aes(x='rel_elo', y='policy', group='run', color='factor(boardsize)'))
        + pn.geom_jitter(show_legend=False, size=.25, width=.02, shape='.')
        + pn.scale_y_continuous(trans='log10')
        + pn.scale_color_continuous(trans='log2')
        + pn.labs(
            x='Relative Elo (-1 random, 0 perfect play)',
            y='Policy Noise Scale')
        + plot.IEEE((5, 4)))

class NoiseScales:

    def __init__(self, agent, buffer_len):
        self._agent = agent
        self._count = 0
        self._buffer_len = buffer_len

    def step(self, chunk):
        if (self._count % self._buffer_len == 0):
            gs = gradients(self._agent.network, chunk)
            results = pd.DataFrame([noise_scale_components(chunk, gs[k], k) for k in gs])

            for k, v in results.set_index('kind').unstack().iteritems():
                stats.silent('noise.' + '.'.join(k), v)

            for _, row in results.iterrows():
                stats.mean(f'noise.{row.kind}', row.batch_size*row.variance/row.mean_sq)

        self._count += 1


from __future__ import annotations

import logging
import math
from typing import List

import numpy as np
from jinja2 import ChoiceLoader, Environment, PackageLoader

from scitbx.array_family import flex
from xfel.clustering.cluster import Cluster
from xfel.clustering.cluster_groups import unit_cell_info

from dials.algorithms.clustering import plots as cluster_plotter
from dials.util import tabulate

logger = logging.getLogger("dials.algorithms.indexing.ssx.analysis")

def generate_html_report(plots: dict, filename: str) -> None:
    loader = ChoiceLoader(
        [
            PackageLoader("dials", "templates"),
            PackageLoader("dials", "static", encoding="utf-8"),
        ]
    )
    env = Environment(loader=loader)
    template = env.get_template("simple_report.html")
    html = template.render(
        page_title="DIALS SSX indexing report",
        panel_title="Indexing plots",
        graphs=plots,
    )
    with open(filename, "wb") as f:
        f.write(html.encode("utf-8", "xmlcharrefreplace"))

def make_summary_table(results_summary: dict) -> tabulate:
    # make a summary table
    overall_summary_header = [
        "Image",
        "expt_id",
        "n_indexed",
        "RMSD X",
        "RMSD Y",
        "RMSD dPsi",
    ]

    rows = []
    total = 0
    if any(len(v) > 1 for v in results_summary.values()):
        show_lattices = True
        overall_summary_header.insert(1, "lattice")
    else:
        show_lattices = False
    for k in sorted(results_summary.keys()):
        for j, cryst in enumerate(results_summary[k]):
            if not cryst["n_indexed"]:
                continue
            n_idx, n_strong = (cryst["n_indexed"], cryst["n_strong"])
            frac_idx = f"{n_idx}/{n_strong} ({100*n_idx/n_strong:2.1f}%)"
            row = [
                cryst["Image"],
                str(total),
                frac_idx,
                cryst["RMSD_X"],
                cryst["RMSD_Y"],
                cryst["RMSD_dPsi"],
            ]
            if show_lattices:
                row.insert(1, j + 1)
            rows.append(row)
            total += 1

    summary_table = tabulate(rows, overall_summary_header)
    return summary_table

def combine_results_dicts(results_summaries: List[dict]) -> dict:
    """For a list of dictionaries, each with keys 0..n-1,
    combine into a single dictionary with keys 0..ntot-1"""
    combined_summary = {}
    n_overall = 0
    for d in results_summaries:
        n_this = len(d)
        for i in range(n_this):
            combined_summary[i + n_overall] = d.pop(i)
        n_overall += n_this
    return combined_summary

def make_cluster_plots(large_clusters: List[Cluster]) -> dict:
    cluster_plots = {}
    for n, cluster in enumerate(large_clusters):
        uc_params = [flex.double() for i in range(6)]
        for c in cluster.members:
            ucp = c.crystal_symmetry.unit_cell().parameters()
            for i in range(6):
                uc_params[i].append(ucp[i])
        d_this = cluster_plotter.plot_uc_histograms(uc_params)
        d_this["uc_scatter"]["layout"]["title"] += f" cluster {n+1}"
        d_this["uc_hist"]["layout"]["title"] += f" cluster {n+1}"
        d_this[f"uc_scatter_{n}"] = d_this.pop("uc_scatter")
        d_this[f"uc_hist_{n}"] = d_this.pop("uc_hist")
        cluster_plots.update(d_this)
    return cluster_plots

def report_on_crystal_clusters(crystal_symmetries, make_plots=True):
    ucs = Cluster.from_crystal_symmetries(crystal_symmetries)
    clusters, _ = ucs.ab_cluster(5000, log=None, write_file_lists=False, doplot=False)
    cluster_plots = {}
    min_cluster_pc = 5
    threshold = math.floor((min_cluster_pc / 100) * len(crystal_symmetries))
    large_clusters = [c for c in clusters if len(c.members) > threshold]
    large_clusters.sort(key=lambda x: len(x.members), reverse=True)

    if large_clusters:
        logger.info(
            f"""
Unit cell clustering analysis, clusters with >{min_cluster_pc}% of the number of crystals indexed
{unit_cell_info(large_clusters)}
"""
        )
        if make_plots:
            cluster_plots = make_cluster_plots(large_clusters)
    else:
        logger.info(
            f"No clusters found with >{min_cluster_pc}% of the number of crystals."
        )
    return cluster_plots, large_clusters

def generate_plots(summary_data: dict) -> dict:
    """Generate indexing plots from the summary data from index_all_concurrent"""
    # n_indexed_arrays are cumulative n_indexed for nth lattice
    n_indexed_arrays = [np.zeros(len(summary_data))]
    rmsd_x_arrays = [np.zeros(len(summary_data))]
    rmsd_y_arrays = [np.zeros(len(summary_data))]
    rmsd_z_arrays = [np.zeros(len(summary_data))]
    n_total_indexed = np.zeros(len(summary_data))
    n_strong_array = np.zeros(len(summary_data))
    images = np.arange(1, len(summary_data) + 1)
    n_lattices = 1

    for k in sorted(summary_data.keys()):
        n_lattices_this = len(summary_data[k])
        n_strong_array[k] = summary_data[k][0]["n_strong"]
        for j, cryst in enumerate(summary_data[k]):
            if not cryst["n_indexed"]:
                continue
            if n_lattices_this > n_lattices:
                for _ in range(n_lattices_this - n_lattices):
                    n_indexed_arrays.append(np.zeros(len(summary_data)))
                    rmsd_x_arrays.append(np.zeros(len(summary_data)))
                    rmsd_y_arrays.append(np.zeros(len(summary_data)))
                    rmsd_z_arrays.append(np.zeros(len(summary_data)))
                n_lattices = n_lattices_this
            n_indexed_arrays[j][k] = cryst["n_indexed"]
            rmsd_x_arrays[j][k] = cryst["RMSD_X"]
            rmsd_y_arrays[j][k] = cryst["RMSD_Y"]
            rmsd_z_arrays[j][k] = cryst["RMSD_dPsi"]
            n_total_indexed[k] += cryst["n_indexed"]

    n_indexed_data = [
        {
            "x": images.tolist(),
            "y": n_indexed_arrays[0].tolist(),
            "type": "scatter",
            "mode": "markers",
            "name": "N indexed",
        },
    ]
    rmsd_data = [
        {
            "x": images[rmsd_x_arrays[0] > 0].tolist(),
            "y": rmsd_x_arrays[0][rmsd_x_arrays[0] > 0].tolist(),
            "type": "scatter",
            "mode": "markers",
            "name": "RMSD X",
        },
        {
            "x": images[rmsd_y_arrays[0] > 0].tolist(),
            "y": rmsd_y_arrays[0][rmsd_y_arrays[0] > 0].tolist(),
            "type": "scatter",
            "mode": "markers",
            "name": "RMSD Y",
        },
    ]
    rmsdz_data = [
        {
            "x": images[rmsd_z_arrays[0] > 0].tolist(),
            "y": rmsd_z_arrays[0][rmsd_z_arrays[0] > 0].tolist(),
            "type": "scatter",
            "mode": "markers",
            "name": "RMSD dPsi",
        },
    ]
    if n_lattices > 1:
        n_indexed_data[0]["name"] += " (lattice 1)"
        rmsd_data[0]["name"] += " (lattice 1)"
        rmsd_data[1]["name"] += " (lattice 1)"
        rmsdz_data[0]["name"] += " (lattice 1)"
        for i, arr in enumerate(n_indexed_arrays[1:]):
            sub_images = images[arr > 0]
            sub_data = arr[arr > 0]
            n_indexed_data.append(
                {
                    "x": sub_images.tolist(),
                    "y": sub_data.tolist(),
                    "type": "scatter",
                    "mode": "markers",
                    "name": f"N indexed (lattice {i+2})",
                }
            )
        for i, arr in enumerate(rmsd_x_arrays[1:]):
            sub_images = images[arr > 0]
            sub_data_x = arr[arr > 0]
            sub_data_y = rmsd_y_arrays[i + 1][arr > 0]
            rmsd_data.append(
                {
                    "x": sub_images.tolist(),
                    "y": sub_data_x.tolist(),
                    "type": "scatter",
                    "mode": "markers",
                    "name": f"RMSD X (lattice {i+2})",
                },
            )
            rmsd_data.append(
                {
                    "x": sub_images.tolist(),
                    "y": sub_data_y.tolist(),
                    "type": "scatter",
                    "mode": "markers",
                    "name": f"RMSD Y (lattice {i+2})",
                },
            )
        for i, arr in enumerate(rmsd_z_arrays[1:]):
            sub_images = images[arr > 0]
            sub_data = arr[arr > 0]
            rmsdz_data.append(
                {
                    "x": sub_images.tolist(),
                    "y": sub_data.tolist(),
                    "type": "scatter",
                    "mode": "markers",
                    "name": f"RMSD dPsi (lattice {i+2})",
                },
            )
    percent_indexed = 100 * n_total_indexed / n_strong_array
    images = images.tolist()
    n_indexed_data.append(
        {
            "x": images,
            "y": n_strong_array.tolist(),
            "type": "scatter",
            "mode": "markers",
            "name": "N strong",
        },
    )

    percent_bins = np.linspace(0, 100, 51)
    percent_hist = np.histogram(percent_indexed, percent_bins)[0]

    def _generate_hist_data(rmsd_arrays, step=0.01):
        all_rmsd = np.concatenate(rmsd_arrays)
        all_rmsd = all_rmsd[all_rmsd > 0]
        mult = int(1 / 0.01)
        start = math.floor(np.min(all_rmsd) * mult) / mult
        stop = math.ceil(np.max(all_rmsd) * mult) / mult
        nbins = int((stop - start) / step)
        hist, bin_edges = np.histogram(
            all_rmsd,
            bins=nbins,
            range=(start, stop),
        )
        bin_centers = bin_edges[:-1] + np.diff(bin_edges) / 2
        return hist, bin_centers

    hist_x, bin_centers_x = _generate_hist_data(rmsd_x_arrays)
    hist_y, bin_centers_y = _generate_hist_data(rmsd_y_arrays)
    hist_z, bin_centers_z = _generate_hist_data(rmsd_z_arrays, 0.001)

    plots = {
        "n_indexed": {
            "data": n_indexed_data,
            "layout": {
                "title": "Number of indexed reflections per image",
                "xaxis": {"title": "image number"},
                "yaxis": {"title": "N reflections"},
            },
        },
        "percent_indexed": {
            "data": [
                {
                    "x": images,
                    "y": percent_indexed.tolist(),
                    "type": "scatter",
                    "mode": "markers",
                    "name": "Percentage of strong spots indexed",
                }
            ],
            "layout": {
                "title": "Percentage of strong spots indexed per image",
                "xaxis": {"title": "image number"},
                "yaxis": {"title": "Percentage"},
            },
        },
        "percent_indexed_hist": {
            "data": [
                {
                    "x": percent_bins.tolist(),
                    "y": percent_hist.tolist(),
                    "type": "bar",
                }
            ],
            "layout": {
                "title": "Distribution of percentage indexed",
                "xaxis": {"title": "Percentage indexed"},
                "yaxis": {"title": "Number of images"},
                "bargap": 0,
            },
        },
        "rmsds": {
            "data": rmsd_data,
            "layout": {
                "title": "RMSDs (x, y) per image",
                "xaxis": {"title": "image number"},
                "yaxis": {"title": "RMSD (px)"},
            },
        },
        "rmsdz": {
            "data": rmsdz_data,
            "layout": {
                "title": "RMSD (dPsi) per image",
                "xaxis": {"title": "image number"},
                "yaxis": {"title": "RMSD dPsi (deg)"},
            },
        },
        "rmsdxy_hist": {
            "data": [
                {
                    "x": bin_centers_x.tolist(),
                    "y": hist_x.tolist(),
                    "type": "bar",
                    "name": "RMSD X",
                    "opacity": 0.6,
                },
                {
                    "x": bin_centers_y.tolist(),
                    "y": hist_y.tolist(),
                    "type": "bar",
                    "name": "RMSD Y",
                    "opacity": 0.6,
                },
            ],
            "layout": {
                "title": "Distribution of RMSDs (x, y)",
                "xaxis": {"title": "RMSD (px)"},
                "yaxis": {"title": "Number of images"},
                "bargap": 0,
                "barmode": "overlay",
            },
        },
        "rmsdz_hist": {
            "data": [
                {
                    "x": bin_centers_z.tolist(),
                    "y": hist_z.tolist(),
                    "type": "bar",
                    "name": "RMSD dPsi",
                },
            ],
            "layout": {
                "title": "Distribution of RMSDs (dPsi)",
                "xaxis": {"title": "RMSD dPsi (deg)"},
                "yaxis": {"title": "Number of images"},
                "bargap": 0,
            },
        },
    }
    return plots


import os
import numpy as np
from subprocess import Popen, PIPE
from scipy.interpolate import griddata
import _pickle as cPickle
from sklearn.neighbors import KDTree

from Bio.PDB import *
from Bio.SeqUtils import seq1, seq3
from Bio.Seq import Seq
from Bio import SeqIO

from default_config.bin_path import bin_path
from default_config.dir_options import dir_opts

from compute_surface.protonate import protonate
from compute_surface.extractPDB import extractPDB
from compute_surface.extract_xyzrn import extract_xyzrn
from compute_surface.apply_msms import computeMSMS

from features.Hydrophobicity import kd_scale
from features.pKa import pKa_scale
from features.macc import macc
from features.Physicochemical import li_scale

class RBP():

    def __init__(self, protein_name):
        self.pdb_id, self.chain = protein_name.split('_') #complex, protein, rna

        #download pdb, call Reduce and MSMS
        self._download_pdb()

        self.model = self._load_pdb()

        self.RNA_space, self.RNA_chain = self._get_RNA_space()

        self.vertices, self.vert_info = self._get_surface() #ndarray (n*3), list(n) A_19_x_VAL_HG11

        #Extract the sequence of protien
        self.seq, self.index2resid, self.resid2Residue = self._get_pdb_seq(self.chain, kd_scale)

        #Get the coordinates of atom of all RNAs in complex
        self.geometric_center = self._get_geometric_center()

        #Get res_id on the surface
        self.surface_res = self._get_surface_res()

        #Get surface center of each residue to calculate the sampling density on the surface.
        self.res_to_vertice = self._get_res_to_vertice()

        #Calculate the label of each surface_res
        self.label, self.all_label = self._get_label()

    def _download_pdb(self):
        if not os.path.exists(dir_opts['raw_pdb_dir']):
            os.makedirs(dir_opts['raw_pdb_dir'])
        pdbl = PDBList()
        pdbl.retrieve_pdb_file(self.pdb_id, pdir=dir_opts['raw_pdb_dir'], file_format='mmCif')

    def _get_surface(self):
        extractPDB(self.pdb_id + '_' + self.chain)
        protonate(self.pdb_id+'_'+self.chain)
        extract_xyzrn(self.pdb_id, self.chain)
        vertices, faces1, normalv1, vert_info = computeMSMS(self.pdb_id, self.chain)
        return vertices, vert_info

    def _load_pdb(self):
        pdb_file = os.path.join(dir_opts['raw_pdb_dir'], self.pdb_id+'.cif')
        parser = MMCIFParser(QUIET=True)
        struct = parser.get_structure(self.pdb_id, pdb_file)
        model = Selection.unfold_entities(struct, "M")[0]
        return model

    def _get_pdb_seq(self, chain, scale):
        chain = self.model.child_dict[chain]
        res_seq = ''
        index2resid = {}
        resid2structure = {}
        index = 0
        for Residue in chain.child_list:
            if Residue.get_resname() in scale.keys():
                res_seq += (seq1(Residue.get_resname()))
                index2resid[index] = Residue.get_id()[1]
                resid2structure[Residue.get_id()[1]] = Residue
                index += 1
        return res_seq, index2resid, resid2structure

    def _get_RNA_space(self):
        RNA_list = ["A", "C", "G", "U"]
        RNA_chain = set()
        atom_list = []
        for chain in self.model.child_list:
            if chain.id == self.chain:
                continue
            for res in chain.child_list:
                atom_type = res.resname.strip()
                if atom_type in RNA_list:
                    for atom in res:
                        atom_list.append(atom.coord)
                        RNA_chain.add(chain.id)
        chain = '_'.join(RNA_chain)
        return atom_list, chain

    def _get_label(self):
        rna_tree = KDTree(self.RNA_space)
        label_dict = {}
        for res_id in self.surface_res:
            res = self.resid2Residue[res_id]
            res_coord = []
            for atom in res.child_list:
                res_coord.append(atom.coord)
            d, t = rna_tree.query(res_coord)
            if np.min(d) < 5.0:
                label_dict[res_id] = 1
            else:
                label_dict[res_id] = 0
        label_all = []
        for index in range(len(self.seq)):
            if self.index2resid[index] in self.surface_res:
                label_all.append(label_dict[self.index2resid[index]])
            else:
                label_all.append(0)
        return label_dict, np.array(label_all)

    def _get_surface_res(self):
        surface_res = set()
        for item in self.vert_info:
            resid = int(item.split('_')[1])
            surface_res.add(resid)
        return list(surface_res)

    def _get_res_to_vertice(self):
        res_to_vertices = {}
        for index, item in enumerate(self.vert_info):
            res_id = int(item.split('_')[1])
            if res_id in res_to_vertices.keys():
                res_to_vertices[res_id].append(self.vertices[index,:])
            else:
                res_to_vertices[res_id] = [self.vertices[index,:],]
        res_to_vertice = {}
        for item in res_to_vertices.keys():
            res_surface_array = np.array(res_to_vertices[item])
            res_to_vertice[item] = np.mean(res_surface_array,axis=0)
        return res_to_vertice

    def _get_geometric_center(self):
        return self.vertices.mean(axis=0)

    def _call_dssp(self):
        dssp_bin = bin_path['DSSP']
        pdb_file = os.path.join(dir_opts['chain_pdb_dir'], self.pdb_id + '_'+self.chain + '.pdb')
        if not os.path.exists(dir_opts['dssp']):
            os.mkdir(dir_opts['dssp'])
        out_file = os.path.join(dir_opts['dssp'], self.pdb_id+'_'+self.chain+'.dssp')

        args = [dssp_bin, '-i', pdb_file, '-o', out_file]
        p2 = Popen(args, stdout=PIPE, stderr=PIPE)
        stdout, stderr = p2.communicate()
        rASA = {}
        secondary = {}
        flag_1 = False
        flag_2 = False
        with open(out_file) as pid:
            for line in pid.readlines():
                if line[2] == '#':
                    flag_1 = True
                    continue
                if flag_1 and line[11] == self.chain[-1]:
                    flag_2 = True
                    res_name = line[13]
                    C_list = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
                    if res_name in C_list:
                        res_name = 'C'
                    if res_name == 'X':
                        continue
                    res_id = int(line[5:10])
                    rASA[res_id] = float(line[34:39]) / float(macc[res_name])
                    if line[16] in ('G', 'H', 'I'):
                        secondary[res_id] = [1, 0, 0]
                    if line[16] in ('E', 'B'):
                        secondary[res_id] = [0, 1, 0]
                    if line[16] in ('T', 'S', ' '):
                        secondary[res_id] = [0, 0, 1]
                if flag_1 and flag_2 and line[13:15] == '!*':
                    break
        return rASA, secondary

    def _call_psiblast(self):
        psiblast_bin = bin_path['PSIBLAST']
        uniprot_database = bin_path['PSIBLAST_DATABASE']

        if not os.path.exists(dir_opts['blast_fasta_dir']):
            os.mkdir(dir_opts['blast_fasta_dir'])
        if not os.path.exists(dir_opts['blast_pssm_dir']):
            os.mkdir(dir_opts['blast_pssm_dir'])

        #save fasta file
        seq = Seq(self.seq)
        seq_io = SeqIO.SeqRecord(seq, name = self.pdb_id + '_' + self.chain)
        fasta_dir = os.path.join(dir_opts['blast_fasta_dir'], self.pdb_id + '_' + self.chain + '.fasta')
        if not os.path.exists(fasta_dir):
            SeqIO.write(seq_io, fasta_dir, "fasta")

        #make pssm_file
        pssm_dir = os.path.join(dir_opts['blast_pssm_dir'],self.pdb_id+'_'+self.chain+'.pssm')
        args = [psiblast_bin, "-db", uniprot_database, "-query", fasta_dir, "-evalue", '0.001', "-num_iterations", '3', "-out_ascii_pssm", pssm_dir]
        if not os.path.exists(pssm_dir):
            p2 = Popen(args, stdout=PIPE, stderr=PIPE)
            print('calling psiblast start')
            stdout, stderr = p2.communicate()
            print('calling psiblast over')

        #sparse pssm_file
        RNA_order = 'ARNDCQEGHILKMFPSTWYV'
        pssm_feature = {}
        with open(pssm_dir, 'r') as pid:
            for index, line in enumerate(pid.readlines()[3:]):
                if line == '\n':
                    break
                pssm_acid, feature = self._sparse_pssm_line(line)
                # amino_acid_position = RNA_order.index(pssm_acid)
                pssm_feature[self.index2resid[index]] = feature
        return pssm_feature

    def _call_hhm(self):
        HHblits = bin_path['HHblits']
        HHblits_DB = bin_path['HHblits_DB']

        if not os.path.exists(dir_opts['blast_fasta_dir']):
            os.mkdir(dir_opts['blast_fasta_dir'])
        if not os.path.exists(dir_opts['hhm_dir']):
            os.mkdir(dir_opts['hhm_dir'])

        #save fasta file
        seq = Seq(self.seq)
        seq_io = SeqIO.SeqRecord(seq, name = self.pdb_id + '_' + self.chain)
        fasta_dir = os.path.join(dir_opts['blast_fasta_dir'], self.pdb_id + '_' + self.chain + '.fasta')
        if not os.path.exists(fasta_dir):
            SeqIO.write(seq_io, fasta_dir, "fasta")

        hhm_dir = os.path.join(dir_opts['hhm_dir'], self.pdb_id+'_'+self.chain+'.hhm')

        args = [HHblits, '-d', HHblits_DB, '-i', fasta_dir, '-ohhm', hhm_dir]
        if not os.path.exists(hhm_dir):
            p2 = Popen(args, stdout=PIPE, stderr=PIPE)
            print('calling hhblits start')
            stdout, stderr = p2.communicate()
            print('calling hhblits over')
        hhm_feature = {}
        hhm_dir = os.path.join(dir_opts['hhm_dir'], self.pdb_id+'_'+self.chain+'.hhm')

        with open(hhm_dir, 'r') as f:
            text=f.readlines()
        hhm_begin_line=0
        hhm_end_line=0
        for i in range(len(text)):
            if '#' in text[i]:
                hhm_begin_line=i + 5
            elif '//' in text[i]:
                hhm_end_line=i
        hhm=np.zeros([int((hhm_end_line - hhm_begin_line) / 3), 30])

        axis_x=0
        for i in range(hhm_begin_line, hhm_end_line, 3):
            line1=text[i].split()[2:-1]
            line2=text[i + 1].split()
            axis_y=0
            for j in line1:
                if j == '*':
                    hhm[axis_x][axis_y]=9999 / 10000.0
                else:
                    hhm[axis_x][axis_y]=float(j) / 10000.0
                axis_y+=1
            for j in line2:
                if j == '*':
                    hhm[axis_x][axis_y]=9999 / 10000.0
                else:
                    hhm[axis_x][axis_y]=float(j) / 10000.0
                axis_y+=1
            axis_x+=1
        hhm=(hhm - np.min(hhm)) / (np.max(hhm) - np.min(hhm))
        for index in range(int((hhm_end_line - hhm_begin_line) / 3)):
            resid = self.index2resid[index]
            hhm_feature[resid] = hhm[index,:].tolist()

        return hhm_feature

    def _sparse_pssm_line(self, line):
        a = line.split()
        pssm_str = a[2:22]
        pssm_value = []
        for item in pssm_str:
            # pssm_value.append(1/(1+np.exp(float(item))))
            pssm_value.append(float(item))
        amino_acid = a[1]
        return amino_acid, pssm_value

    def _get_pKa_feature(self):
        #used to calculate pKa feature
        pKa_mapping = {}
        for item in pKa_scale.keys():
            pKa_mapping[item] = ((pKa_scale[item][0] - 1.82) / (2.38-1.82), (pKa_scale[item][1] - 8.80) / (10.96 - 8.80),
                                 (pKa_scale[item][2] - 3.65) / (12.48-3.65))
        return pKa_mapping

    def _get_Physicochemical_feature(self):
        physio_scale = {}
        for item in li_scale.keys():
            physio_scale[seq3(item).upper()] = ((li_scale[item][0] - (-1))/(1-(-1)), ((li_scale[item][1] - 2)/(4-2)))
        return physio_scale

    def _get_hydrophobicity_feature(self):
        #used to calculate hydrophily
        hydrophobicity_scale = {}
        for item in kd_scale.keys():
            hydrophobicity_scale[item] = (kd_scale[item] - (-4.5)) /(4.5 - (-4.5))
        return hydrophobicity_scale

    def _get_fea_mapping(self):
        #call dssp and psiblast
        self.feature_rASA, self.feature_secondary = self._call_dssp()
        self.feature_pssm = self._call_psiblast()
        self.feature_hhm = self._call_hhm()
        self.feature_pKa = self._get_pKa_feature() #shape: (length_seq, 3)
        self.feature_physicochemical = self._get_Physicochemical_feature() #(length_seq, 2)
        self.feature_hydophobicity = self._get_hydrophobicity_feature() #(length_seq, 1)

    def _get_rotation_matrix(self, res_id):
        vertex_center = self.res_to_vertice[res_id]
        vertex_center = vertex_center - self.geometric_center
        rotation_x = np.arctan2(vertex_center[2], vertex_center[1])
        if rotation_x >= 0:
            rotation_x = np.pi * 2 - rotation_x
        else:
            rotation_x = - rotation_x
        rotation_matrix_x = np.array([[1, 0, 0],
                                    [0, np.cos(rotation_x), np.sin(rotation_x)],
                                    [0, -np.sin(rotation_x), np.cos(rotation_x)],
                                    ])
        vertex_center_rotation_x = np.matmul(vertex_center, rotation_matrix_x)
        rotation_z = np.arctan2(vertex_center_rotation_x[1], vertex_center_rotation_x[0])
        if rotation_z >= 0:
            rotation_z = np.pi * 2 - rotation_z
        else:
            rotation_z = -rotation_z
        rotation_matrix_z = np.array([[np.cos(rotation_z), np.sin(rotation_z), 0],
                                    [-np.sin(rotation_z), np.cos(rotation_z), 0],
                                    [0, 0, 1],
                                    ])
        return rotation_matrix_x, rotation_matrix_z

    def _get_polar_coord(self, res_id):
        rotation_x, rotation_z = self._get_rotation_matrix(res_id)
        coord = self.vertices - self.geometric_center
        coord_rotation_x = np.matmul(coord,rotation_x)
        coord = np.matmul(coord_rotation_x, rotation_z)
        #get all lat lon
        x, y, z, radius = coord[:,0], coord[:,1], coord[:,2], np.sqrt(np.sum((coord ** 2), axis=1))
        lat = np.arcsin(z / radius) * 180 / np.pi  # Latitudes
        lon = np.arctan2(y, x) * 180 / np.pi  # Longitudes
        lon_lat = np.hstack([lon[:, None], lat[:, None]])
        #get min and max lat and lon
        center_point = self.res_to_vertice[res_id] - self.geometric_center
        center_point = np.matmul(center_point, rotation_x)
        center_point = np.matmul(center_point, rotation_z)
        kdtree = KDTree(coord)
        (ind, dist) = kdtree.query_radius(center_point[None,:], 12, return_distance = True)
        coord = coord[ind.tolist()[0],:]
        x, y, z, radius = coord[:,0], coord[:,1], coord[:,2], np.sqrt(np.sum((coord ** 2), axis=1))
        lat = np.arcsin(z / radius) * 180 / np.pi  # Latitudes
        lon = np.arctan2(y, x) * 180 / np.pi  # Longitudes
        min_max_lon_lat = (np.min(lon), np.max(lon), np.min(lat), np.max(lat))
        return lon_lat, min_max_lon_lat

    def _interpolation(self, polar_pix, min_max_lon_lat, feature):
        x_lons = np.linspace(min_max_lon_lat[0], min_max_lon_lat[1], 32)
        y_lats = np.linspace(min_max_lon_lat[2], min_max_lon_lat[3], 32)
        MhpS = griddata(polar_pix, feature, (x_lons[None, :], y_lats[:, None]), method='nearest')
        return MhpS

    def _get_graph(self):
        self._get_fea_mapping()
        coord = []
        graph = []
        for index, vert in enumerate(self.vert_info):
            try:
                res_name = vert.split('_')[3]
                res_id = int(vert.split('_')[1])
                fea = []
                fea.append(self.feature_rASA[res_id]) #1
                fea.extend(self.feature_secondary[res_id]) #3
                fea.extend(self.feature_pssm[res_id])#20
                fea.extend(self.feature_hhm[res_id])#30
                fea.extend(self.feature_pKa[res_name])#3
                fea.extend(self.feature_physicochemical[res_name])#2
                fea.append(self.feature_hydophobicity[res_name])#1
                coord.append(self.vertices[index, :])
                graph.append(fea)
            except:
                pass
        self.vertices = coord
        return np.array(graph, dtype=np.float)

    def get_data(self):
        self.graph = self._get_graph()
        data = []
        label = []
        res_surface = []
        for index, res_id in enumerate(self.surface_res):
            try:
                lon_lat, min_max_lon_lat = self._get_polar_coord(res_id)
                features = self._interpolation(lon_lat, min_max_lon_lat, self.graph)
                data.append(features)
                label.append(self.label[res_id])
                res_surface.append(res_id)
            except:
                print(self.pdb_id + '_' + self.chain, index, res_id, 'wrong')
                continue
        batch_dict = {'data':np.array(data),  'label':np.array(label), 'res_id':res_surface}
        with open(os.path.join(dir_opts['data_label'], self.pdb_id + '_' + self.chain ), 'wb') as pid:
            cPickle.dump(batch_dict, pid)


#   GW ALERT FOR O3 RUN IN QSO SERVER
#   BASED ON https://emfollow.docs.ligo.org/userguide/tutorial/index.html
#	PLEASE RUN DIRECTORY /data3/gwshare/bayestar/IDLzone
#   2019.03.03 MADE		BY Gregory S.H. Paek
#	2019.03.28 UPDATED	BY Gregory S.H. Paek
#	2019.04.18 UPDATED	BY Gregory S.H. Paek
#	2019.04.27 UPDATED	BY Gregory S.H. Paek
#============================================================#
#	MODULE
#------------------------------------------------------------#
import gcn
import gcn.handlers
import gcn.notice_types
import healpy as hp
import numpy as np
#============================================================#
# Function to call every time a GCN is received.
# Run only for notices of type
# LVC_PRELIMINARY, LVC_INITIAL, or LVC_UPDATE.
@gcn.handlers.include_notice_types(
	gcn.notice_types.LVC_PRELIMINARY,
	gcn.notice_types.LVC_INITIAL,
	gcn.notice_types.LVC_UPDATE)
def process_gcn(payload, root):
	'''
	EXTRACT CONFIDENCE MAP FROM BAYESTAR FILE
	'''
	import gcn
	import gcn.handlers
	import gcn.notice_types
	import healpy as hp
	import numpy as np
	import time
	import os, glob
	from astropy.table import Table, Column, MaskedColumn
	import astropy.units as u
	from astropy.coordinates import SkyCoord
	from astropy.io import ascii
	import matplotlib.pyplot as plt
	from imsng import gw
	from imsng import tool
	from astropy.table import Table, vstack
	#------------------------------------------------------------#
	#	Read all of the VOEvent parameters from the "What" section.
	#------------------------------------------------------------#
	#	GIVE DELAY TO ESCAPE OVERAPPED COMMAND
	#time.sleep(10)
	starttime	= time.time()
	params	= {elem.attrib['name']:
			  elem.attrib['value']
			  for elem in root.iterfind('.//Param')}
	role	= root.attrib['role']
	#	Print all parameters.
	eventname	= params['GraceID']+'_'+params['AlertType']
	#save_healfix	= '/data3/gwshare/bayestar/'
	#save_path		= save_healfix+eventname
	save_healfix	= '/data3/gwshare/bayestar/IDLzone'
	save_path		= '/data3/gwshare/bayestar/'+eventname

	os.system('mkdir '+save_path)

	yymmdd, hhmmss	= tool.timename()
	yy	= yymmdd[0:2]
	yyyy= '20'+yy
	mm	= yymmdd[2:4]
	dd	= yymmdd[4:6]
	#------------------------------------------------------------#
	#	PROGENITOR
	try:
		prog		= ['BNS', 'NSBH', 'BBH', 'Terrestrial']
		prog_prob	= []
		for key in prog:
			value	= params[key]
			prog_prob.append(float(value))
		indx_prog	= np.where(np.max(prog_prob) == prog_prob)
		progenitor	= prog[indx_prog[0][0]]
	except:
		pass
	#------------------------------------------------------------#
	healpixfits	= params['skymap_fits']
	os.system('wget '+healpixfits)

	#	IDL PART (Joonho Kim)
	if root.attrib['role'] != 'test':
		gzlist	= glob.glob(save_healfix+'/*.gz')
		os.system('gzip -d '+gzlist[0])
		os.system('sh /data3/gwshare/idlpro/gwfollowup_kmtnet.sh')
		baylist	= glob.glob(save_healfix+'/*.fits')
		os.system('mv '+baylist[0]+' '+save_path+'/'+eventname+'-bayestar.fits')
		os.system('mv '+save_healfix+'/sc*.cat '+save_path+'/')
		os.system('mv '+save_healfix+'/gw*.txt '+save_path+'/')
		os.system('chmod 777 '+save_path+'/*')
	else:
		print('TEST SIGNAL')
		os.system('rm bayestar.fits.gz')
	deltime		= time.time() - starttime
	print('All PROCESS IS DONE.\t('+str(round(deltime, 1))+' sec)')
	print('IDL PART FOR GW ALERT IS STANDBY.')

'''
#	TEST ALERT
import lxml.etree
payload		= open('/home/gw/Research/test/MS181101ab-1-Preliminary.xml', 'rb').read()
root		= lxml.etree.fromstring(payload)

#	TEST BAYESTAR
confidence	= 0.9
healpixfits	= '/home/sonic/Research/cat/GWlocal/bayestar170814-refined.fits'
progenitor	= 'BBH'
'''

#------------------------------------------------------------#
#   Listen for GCNs until the program is interrupted
#   (killed or interrupted with control-C).
gcn.listen(handler=process_gcn)
#------------------------------------------------------------#

import numpy,import astropy,from astropy


import logging
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
from PIL import Image

from .palette import Palette
from .utils import open_img_file

class Master(Palette):
    @classmethod
    def from_file(
        cls,
        master_image_file: Path,
        *args,
        crop_ratio: Optional[float] = None,
        img_size: Optional[Tuple[int, int]] = None,
        convert: Optional[str] = None,
        **kwargs,
    ) -> "Master":
        """Create a master image from file.

        Args:
            master_image_file: path to image file.
            crop_ratio: width to height ratio to crop the master image to. 1 results in a square image.
            image_size: resize the image to the provided size, width followed by height.
            convert: convert the image to the provided mode. See PIL Modes.

        Returns:
            Master instance.

        Examples:
            For black and white square 1280x1280 image.

            >>> Master.from_file("master.png", crop_ratio=1, img_size=(1280, 1280), convert="L")
        """
        img = open_img_file(
            master_image_file, crop_ratio=crop_ratio, size=img_size, convert=convert
        )
        return cls.from_image(img, *args, **kwargs)

    @classmethod
    def from_image(cls, master_image: Image.Image, *args, **kwargs) -> "Master":
        """Create a master image from PIL.Image.Image

        Args:
            master_image: PIL.Image instance.

        Returns:
            Master instance.
        """
        array = np.asarray(master_image)
        # make sure the arrays have 3 channels even in black and white
        if array.ndim == 2:
            array = np.stack([array] * 3, -1)
        return cls(array, *args, **kwargs)

    def __init__(self, array: np.ndarray) -> None:
        """The master image.

        Args:
            array: numpy array of the image, should contain 3 channels.

        Returns:
            Master instance.
        """
        self.array = array
        self._log = logging.getLogger(__name__)
        self._log.info("master shape: %s", self.array.shape)

    @property
    def img(self):
        """PIL.Image of the Master image."""
        return Image.fromarray(self.array.round(0).astype("uint8"), mode="RGB")

    @property
    def pixels(self) -> np.ndarray:
        """Array containing the 3-channel pixel values of the Master image."""
        return self.array.reshape(-1, self.array.shape[-1])

    def __repr__(self) -> str:
        return f"""{self.__class__.__module__}.{self.__class__.__name__} at {hex(id(self))}:
    shape: {self.array.shape}"""


#!/usr/bin/env python
"""Guess the coordinates of captured regions from sample read depths.

Two approaches available:

- (Faster) Scan a given list of exons and/or other potentially targeted regions.
  The script checks each region and drops those with very low coverage
  indicating they were not captured.
- (Slower) Scan the entire genome, or the given sequencing-accessible regions,
  for regions with elevated coverage. Choose reasonable boundaries for each
  apparently captured region.

Use multiple BAMs for greater robustness in detecting targeted regions, as a
single sample may have poor coverage are some targets by chance.
Still, this script does not guarantee correctly detecting all targets.

See also: https://github.com/brentp/goleft
"""
import argparse
import collections
import logging
import subprocess
import sys

import numpy as np
import pandas as pd

import cnvlib
from cnvlib import parallel
from cnvlib.descriptives import modal_location
from skgenome import tabio, GenomicArray as GA

logging.basicConfig(level=logging.INFO, format="%(message)s")

# ___________________________________________
# Guided method: guess from potential targets

def filter_targets(target_bed, sample_bams, procs):
    """Check if each potential target has significant coverage."""
    try:
        baits = tabio.read(target_bed, 'bed4')
    except:
        raise RuntimeError("Targets must be in BED format; try skg_convert.py")
    logging.info("Loaded %d candidate regions from %s", len(baits), target_bed)
    # Loop over BAMs to calculate weighted averages of bin coverage depths
    total_depths = np.zeros(len(baits), dtype=np.float_)
    for bam_fname in sample_bams:
        logging.info("Evaluating targets in %s", bam_fname)
        sample = cnvlib.do_coverage(target_bed, bam_fname, processes=procs)
        assert len(sample) == len(baits), \
                "%d != %d" % (len(sample), len(baits))
        total_depths += sample['depth'].values
    baits['depth'] = total_depths / len(sample_bams)
    logging.info("Average candidate-target depth:\n%s",
                 baits['depth'].describe())
    return baits

# _________________________________________
# Unguided method: guess from raw depths

def scan_targets(access_bed, sample_bams, min_depth, min_gap, min_length,
                 procs):
    """Estimate baited regions from a genome-wide, per-base depth profile."""
    bait_chunks = []
    # ENH: context manager to call rm on bed chunks? with to_chunks as pool, ck?
    logging.info("Scanning for enriched regions in:\n  %s",
                 '\n  '.join(sample_bams))
    #  with futures.ProcessPoolExecutor(procs) as pool:
    with parallel.pick_pool(procs) as pool:
        args_iter = ((bed_chunk, sample_bams,
                    min_depth, min_gap, min_length)
                    for bed_chunk in parallel.to_chunks(access_bed))
        for bed_chunk_fname, bait_chunk in pool.map(_scan_depth, args_iter):
            bait_chunks.append(bait_chunk)
            parallel.rm(bed_chunk_fname)
    baits = GA(pd.concat(bait_chunks))
    baits['depth'] /= len(sample_bams)
    return baits

def _scan_depth(args):
    """Wrapper for parallel map"""
    bed_fname, bam_fnames, min_depth, min_gap, min_length = args
    regions = list(drop_small(merge_gaps(scan_depth(bed_fname, bam_fnames,
                                                    min_depth),
                                         min_gap),
                              min_length))
    result = pd.DataFrame.from_records(list(regions),
                                       columns=regions[0]._fields)
    return bed_fname, result

def scan_depth(bed_fname, bam_fnames, min_depth):
    """Locate sub-regions with enriched read depth in the given regions.

    Yields
    ------
    tuple
        Region coordinates (0-indexed, half-open): chromosome name, start, end
    """
    Region = collections.namedtuple('Region', 'chromosome start end depth')

    nsamples = len(bam_fnames)
    if nsamples == 1:
        def get_depth(depths):
            return int(depths[0])
    else:
        min_depth *= nsamples
        # NB: samtools emits additional BAMs' depths as trailing columns
        def get_depth(depths):
            return sum(map(int, depths))

    proc = subprocess.Popen(['samtools', 'depth',
                             '-Q',  '1',  # Skip pseudogenes
                             '-b', bed_fname,
                            ] + bam_fnames,
                            stdout=subprocess.PIPE,
                            shell=False)

    # Detect runs of >= min_depth; emit their coordinates
    chrom = start = depths = None
    for line in proc.stdout:
        fields = line.split('\t')
        depth = get_depth(fields[2:])
        is_enriched = (depth >= min_depth)
        if start is None:
            if is_enriched:
                # Entering a new captured region
                chrom = fields[0]
                start = int(fields[1]) - 1
                depths = [depth]
            else:
                # Still not in a captured region
                continue
        elif is_enriched and fields[0] == chrom:
            # Still in a captured region -- extend it
            depths.append(depth)
        else:
            # Exiting a captured region
            # Update target region boundaries
            darr = np.array(depths)
            half_depth = 0.5 * darr.max()
            ok_dp_idx = np.nonzero(darr >= half_depth)[0]
            start_idx = ok_dp_idx[0]
            end_idx = ok_dp_idx[-1] + 1
            yield Region(chrom,
                            start + start_idx,
                            start + end_idx,
                            darr[start_idx:end_idx].mean())
            chrom = start = depths = None

def merge_gaps(regions, min_gap):
    """Merge regions across small gaps."""
    prev = next(regions)
    for reg in regions:
        if reg.start - prev.end < min_gap:
            # Merge
            prev = prev._replace(end=reg.end)
        else:
            # Done merging; move on
            yield prev
            prev = reg
    # Residual
    yield prev

def drop_small(regions, min_length):
    """Merge small gaps and filter by minimum length."""
    return (reg for reg in regions
            if reg.end - reg.start >= min_length)

# ___________________________________________
# Shared

def normalize_depth_log2_filter(baits, min_depth, enrich_ratio=0.1):
    """Calculate normalized depth, add log2 column, filter by enrich_ratio."""
    # Normalize depths to a neutral value of 1.0
    dp_mode = modal_location(baits.data.loc[baits['depth'] > min_depth,
                                            'depth'].values)
    norm_depth = baits['depth'] / dp_mode
    # Drop low-coverage targets
    keep_idx = (norm_depth >= enrich_ratio)
    logging.info("Keeping %d/%d bins with coverage depth >= %f, modal depth %f",
                 keep_idx.sum(), len(keep_idx), dp_mode * enrich_ratio, dp_mode)
    return baits[keep_idx]

if __name__ == '__main__':
    AP = argparse.ArgumentParser(description=__doc__)
    AP.add_argument('sample_bams', nargs='+',
                    help="""Sample BAM file(s) to test for target coverage.""")
    AP.add_argument('-o', '--output', metavar='FILENAME',
                    help="""The inferred targets, in BED format.""")
    AP.add_argument('-c', '--coverage', metavar='FILENAME',
                    help="""Filename to output average coverage depths in .cnn
                    format.""")
    AP.add_argument('-p', '--processes', metavar='CPU',
                    nargs='?', type=int, const=0, default=1,
                    help="""Number of subprocesses to segment in parallel.
                    If given without an argument, use the maximum number
                    of available CPUs. [Default: use 1 process]""")

    AP_x = AP.add_mutually_exclusive_group(required=True)
    AP_x.add_argument('-t', '--targets', metavar='TARGET_BED',
                    help="""Potentially targeted genomic regions, e.g. all known
                    exons in the reference genome, in BED format. Each of these
                    regions will be tested as a whole for enrichment. (Faster
                    method)""")
    AP_x.add_argument('-a', '--access', metavar='ACCESS_BED',
                    # default="../data/access-5k-mappable.grch37.bed",
                    help="""Sequencing-accessible genomic regions (e.g. from
                    'cnvkit.py access'), or known genic regions in the reference
                    genome, in BED format. All bases will be tested for
                    enrichment. (Slower method)""")

    AP_target = AP.add_argument_group("With --targets only")
    AP_target.add_argument('-d', '--min-depth', metavar='DEPTH',
                    type=int, default=5,
                    help="""Minimum sequencing read depth to accept as captured.
                    [Default: %(default)s]""")

    AP_access = AP.add_argument_group("With --access only")
    AP_access.add_argument('-g', '--min-gap', metavar='GAP_SIZE',
                    type=int, default=25,
                    help="""Merge regions separated by gaps smaller than this.
                    [Default: %(default)s]""")
    AP_access.add_argument('-l', '--min-length', metavar='TARGET_SIZE',
                    type=int, default=50,
                    help="""Minimum region length to accept as captured.
                    [Default: %(default)s]""")

    args = AP.parse_args()

    # ENH: can we reserve multiple cores for htslib?
    if args.processes < 1:
        args.processes = None

    if args.targets:
        baits = filter_targets(args.targets, args.sample_bams, args.processes)
    else:
        baits = scan_targets(args.access, args.sample_bams,
                             0.5 * args.min_depth,  # More sensitive 1st pass
                             args.min_gap, args.min_length, args.processes)
    baits = normalize_depth_log2_filter(baits, args.min_depth)
    tabio.write(baits, args.output or sys.stdout, 'bed')
    if args.coverage:
        baits['log2'] = np.log2(baits['depth'] / baits['depth'].median())
        tabio.write(baits, args.coverage, 'tab')