#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by Kleber Noel for disseration:
    inversion mapping using a Local Linear embedding
June 2019
"""
import numpy as np
from numpy import mean

#import pandas as pd
import os, pdb, argparse, re
from random import shuffle
from tqdm import tqdm
from collections import defaultdict, OrderedDict
from time import time
import torch, torch.nn as nn
import torch.optim as optim
from torch import FloatTensor
from torch.autograd import Variable
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence

get_seq_lens = lambda l : [x[-1] for x in l]
get_names    = lambda l : [x[0] for x in l]
dim1_list    = lambda l : [x.shape[1] for x in l]
report = 'count:{}, batch_size:{}, loss:{:.4}; mm:{:.4}, frames:{}, cor:{:.4}'

# GET CUDA STATUS:

def get_args():
    """
    Args define program usage.
    """
    info = 'Script preprocesses data and saves, given i/o directories.'
    parser = argparse.ArgumentParser(info)
    parser.add_argument('--train', action='store_true', default=False)
    parser.add_argument('--test',  action='store_true', default=False)

    parser.add_argument('--lsf-dir',  help='path to lsf directory')
    parser.add_argument('--mfcc-dir', help='path to mfcc directory')
    parser.add_argument('--ema-dir',  help='path to ema directory')
    parser.add_argument('--bilstm-layers', type=int, default=2, help='number of bi-directional LSTM layers')

    parser.add_argument('--shuffle', default=False, action='store_true', help='shuffle batches')
    parser.add_argument('--batch-size', type=int, default=6, help='shuffle batches')

    parser.add_argument('--cuda', action='store_true', default=False, help='GPU support')

    parser.add_argument('--load-dir', help='directory to load a saved model from')
    parser.add_argument('--save-dir', help='directory to save a final model in')

    parser.add_argument('--checkpoint-dir', default = '', help='checkpoint directory to load/save model from/in')
    parser.add_argument('--patience', type=int, default=10, help='save model after n epochs without improvement')
    parser.add_argument('--save-internal', default=5, help='epoch arg to save the model parameters every *n* times')

    parser.add_argument('--graph-data', action='store_true', help='graph data')

    return parser.parse_args()

class LoadInputOutput():
    def __init__(self, args, d):
        if args.lsf_dir: self.lsf_dict = self.load_data(args.lsf_dir)
        if args.mfcc_dir: self.mfcc_dict= self.load_data(args.mfcc_dir)
        self.ema_dict = self.load_data(args.ema_dir)
        self.device   = d
        self.shuffle  = args.shuffle
        self.input_dim, self.output_dim = self.get_io_dims()
        print('input dimensionality of lsfs:\t\t{} '.format(self.input_dim))
        print('output dimensionality of emas:\t\t{}'.format(self.output_dim))

    def get_io_dims(self):
        n_input_dims  = set(dim1_list(self.lsf_dict['channel_lsfd'].values()))
        n_output_dims = set(dim1_list(self.ema_dict['channel_ema'].values()))
        assert {len(n_output_dims),len(n_input_dims)}.pop()==1 # sanity check
        return n_input_dims.pop(), n_output_dims.pop()

    def load_means(self):
        return np.array([float(i[0])*10**int(i[1]) for i in self.norm])

    def load_4xSTDs(self):
        return np.array([4*(float(i[0])*10**int(i[1])) for i in self.norm])

    def load_data(self, directory):
        ema_norm_str, data_dict = '(\d+\.\d+)e([\-\+]\d+)', defaultdict(dict)
        for d1 in os.listdir("./{}".format(directory)):
            print('loading data to dict from {} : {}'.format(directory, d1))
            if not "channel" in d1: continue
            for d2 in os.listdir("./{}/{}".format(directory,d1)):
                key = d2.split('.')[0]
                if re.match('ema_.*.txt',d2):
                    with open(os.path.join(directory, d1, d2), 'r') as f:
                        self.norm = re.findall(ema_norm_str,f.readlines()[0])
                    if 'means' in d2: data_dict.update({key:FloatTensor(self.load_means())})
                    elif 'stds' in d2:data_dict.update({key:FloatTensor(self.load_4xSTDs())})
                    continue
                np_array = np.load(os.path.join(directory, d1, d2),
                                   allow_pickle = True)
                if d1=='channel_name': data_dict.update({key:np_array})
                else:  data_dict[d1].update({key:FloatTensor(np_array)})
        return data_dict

    def get_batches(self):
        total_samples = len(self.lsf_dict['channel_lsfd'].values())
        # currently just choose a validation index: 10%, training 90%
        self.val_idx = int(total_samples * (1 - 0.1))
        self.sample_names = sorted([k for k in self.lsf_dict['channel_lsfd'].keys()])
        report = 'minimum batch target size = {}\nsegmenting data into batches...'
        print(report.format(args.batch_size))
        # define batch TRAINING (t) and VALIDATION (v) for...
        #   SOURCE set (s) & TARGET set (t)
        self.bts, self.bvs = self.make_batches(self.lsf_dict['channel_lsfd'])
        self.btt, self.bvt = self.make_batches(self.ema_dict['channel_ema'])
        # Normalize outputs to millimeters (for calculating RMSE in mm)
        ### DO FOR WHOLE BATCH
        batch_list_tt = [i for i in self.btt[0]]
        batch_list_vt = [i for i in self.bvt[0]]
        self.bttm = self.unnormalize_output(batch_list_tt, True)
        self.bvtm = self.unnormalize_output(batch_list_vt, True)

    def reorder(self, t, i, tup=False):
        if tup: return ([t[0][n] for n in i],[t[1][n] for n in i])
        if not tup: return [t[n] for n in i]

    def shuffle_batches(self):
        """
        Shuffles batches based on numpy:
            get len of training/validation batches -> 0,1, ... len(batch).
        """
        t, v = np.arange(len(self.bts[0])), np.arange(len(self.bvs[0]))
        np.random.shuffle(t), np.random.shuffle(v) # new order for t & v.
        t, v = t.tolist(), v.tolist()
        self.bts  = self.reorder(self.bts, t, tup=True)
        self.bvs  = self.reorder(self.bvs, v, tup=True)
        self.btt  = self.reorder(self.btt, t, tup=True)
        self.bvt  = self.reorder(self.bvt, v, tup=True)
        self.bttm = self.reorder(self.bttm, t, tup=False)
        self.bvtm = self.reorder(self.bvtm, v, tup=False)

    def make_batches(self, in_dict, split=True):
        """
        Makes and pads a batch using data dictionary
        """
        if split: #split into train, validation sets
            t, v = self.return_train_valid(in_dict) #INSERT NEW make_batch var to give random sequence for train/val
            return self.make_batches(t, False), self.make_batches(v, False)
        batches, batch, =  [], []
        frames_to_names, mxfr2nmfr = defaultdict(set), defaultdict(int)
        for k,s in [(k,v.shape[0]) for k,v in in_dict.items()]: frames_to_names[s].add(k)
        # Run through k:frames --> v:names dict in order (lowest to highest)
        for frames, names in sorted(frames_to_names.items(), key=lambda x:x[0]):
            batch.extend([(n,frames) for n in names])
            ### Aim to get up to or above batch size... issue: dif batch sizes.
            if len(batch) >= args.batch_size and batch!=0: ## why batch=0? FIX!
                mxfr2nmfr.update({frames:sorted(batch,reverse=True, key=lambda x:x[1])})
                batch = []
        if len(batch) > 0: # remainder that didn't fit into target batch sizes
            mxfr2nmfr.update({frames:sorted(batch,reverse=True, key=lambda x:x[1])})
        n2f=[]
        for batch in mxfr2nmfr.values():
            if batch==0: continue
            t = [FloatTensor(in_dict[k]).to(self.device) for k, f in batch]
            batches.append(pad_sequence(t, batch_first=True))
            n2f.append([(k,in_dict[k].shape[0]) for k, f in batch])
        return batches, n2f

    def return_train_valid(self, in_dict):
        t, v, = defaultdict(str), defaultdict(str)
        t = {k:in_dict[k] for k in self.sample_names[:self.val_idx]}
        v = {k:in_dict[k] for k in self.sample_names[self.val_idx:]}
        return t,v

    def unnormalize_output(self, batches, batches_set=True):
        unnormalized_data = []
        stds_tensor = FloatTensor(self.ema_dict['ema_stds']).to(self.device)
        if batches_set==True: # A WHOLE BATCH
            for b in batches:
                norm = stds_tensor.repeat(b.shape[0],b.shape[1],1)
                unnormalized_data.append(torch.mul(norm, b))
            return unnormalized_data
        else:
            b = batches
            norm = stds_tensor.repeat(b.shape[0],b.shape[1],1)
            return torch.mul(norm, b)

    def zip_validation_batches(self):
        """
        Val args in zip: 1. src; 2. tgt; 3. tgt (mm), 4. names & sizes
        """
        return zip(self.bvs[0], self.bvt[0], self.bvtm, self.bvt[1])

    def zip_training_batches(self):
        """
        Train args in zip: 1. src; 2. tgt; 3. tgt (mm), 4. names & sizes
        """
        return zip(self.bts[0], self.btt[0], self.bttm, self.btt[1])

    def get_prog_bar(self, zipped_batches, batches_type, epoch):
        t = len(self.bts[0]) if batches_type=='training' else len(self.bvs[0])
        bar = tqdm(iterable = zipped_batches,
                   total = t,
                   desc = '| {} Epoch {:03d}'.format(batches_type, epoch),
                   leave = False,
                   disable = False)
        return bar

class BuildModel(nn.Module):
    def __init__(self, args, d, input_dim, output_dim):
        super(BuildModel, self).__init__()
        self.args          = args
        self.device        = d
        self.input_dim     = input_dim
        self.output_dim    = output_dim
        self.hidden_dim    = 256
        self.fc_hidden_dim = 512
        self.bilstm_layers = args.bilstm_layers
        self.batch_size    = args.batch_size # default: 6
        # L1 lstm cell: should use sigmoid as R activation and tanh as final
        self.bidirlstm = nn.LSTM(input_size    = self.input_dim,
                                 hidden_size   = self.hidden_dim,
                                 bidirectional = True,
                                 num_layers    = self.bilstm_layers)

        self.fc = nn.Linear(in_features  = self.fc_hidden_dim,
                            out_features = self.output_dim)
        self.hc = None
        self.criterion, self.criterion_mm = nn.MSELoss(), nn.MSELoss()
        self.optimizer = optim.Adam(self.parameters(), lr=0.001)
        self.is_new = False # assume model loaded from chkpnt til proven otherwise

    def forward(self, padded_seq_batch):
        """
            padded_seq_batch: padded sequence
            hc: hidden and cell states
            tuple of hidden and cell state
        """
        self.get_len(padded_seq_batch.shape[0])
        self.out_seq = torch.empty((self.sequence_len,
                                       self.batch_size,
                                       self.output_dim))
        packed_seq_batch =  pack_padded_sequence(padded_seq_batch,
                                                 lengths=self.seq_lens,
                                                 batch_first=True)

        ## Reinitialize hidden state, WRT to the dimensions of the LSTM
        ## as LSTM input dim1 changes due to packed sequence dim1 changes in
        ## batch. Some batches are dif sizes, hidden state size changes
        ## Before, I was only passing the packed_seq_batch to bidirlstm.

        self.hc = self.init_hidden(2*self.bilstm_layers,
                                   max(packed_seq_batch[1].tolist()),
                                   self.hc)
        packed_outputs, self.hc = self.bidirlstm(packed_seq_batch, self.hc)
        lstm_out, _ = pad_packed_sequence(packed_outputs, batch_first=True)
        self.out_seq = self.fc(lstm_out)

    def init_hidden(self, a, b, x = None):
        return (Variable(torch.randn(a, b, self.hidden_dim,device=self.device)),
                Variable(torch.randn(a, b, self.hidden_dim,device=self.device)))

    def init_history(self):
        if self.args.checkpoint_dir:
            print('checkpoint "{}" not found'.format(args.checkpoint_dir))
            print('initialising history and creating new model...')
        self.history, self.is_new = defaultdict(dict), True
        self.history['training'], self.history['validation'] = OrderedDict(), OrderedDict()

    def get_len(self, length):
        self.sequence_len = length

    def get_cor(self, output, target):
        """
        function for correlations between target and output in a batch
        normalized over the number of frames samples
        """
        self.av_cor = float(0)
        for frames, o, t in zip(self.seq_lens, output, target):
            vo = o - torch.mean(o[:frames,:])
            vt = t - torch.mean(o[:frames,:])
            numerator = torch.sum(vo * vt)
            denominator = (torch.sqrt(torch.sum(vo ** 2)) * torch.sqrt(torch.sum(vt ** 2)))
            cor = numerator/denominator
            self.av_cor += cor.item()*frames/sum(self.seq_lens)

    def get_lens_names_frames(self, names_to_frames):
        """
        computed batch-wise
        """
        self.seq_lens = get_seq_lens(names_to_frames)
        self.names = get_names(names_to_frames)
        self.frames = sum(self.seq_lens)

    def get_losses(self, out_seqm, pad_tgt_seq, pad_tgt_seqm):
        self.loss    = self.criterion(self.out_seq, pad_tgt_seq)
        self.loss_mm = self.criterion_mm(out_seqm, pad_tgt_seqm).item()*1000

    def propagate_loss(self):
        self.loss.backward()
        self.optimizer.step()

    def update_info(self, epoch, phase):
        """
        phase is either training or testing.
        info during epoch is updated as a set of loss values (mm and norm),
        frame number, correlation, sample number for each minibatch
        at end, average loss is computed.
        """
        self.epoch = epoch
        sample_number = len(self.names)
        s = tuple([sample_number, self.frames, self.loss.item(), self.loss_mm, self.av_cor])
        if self.history[phase].get(epoch): self.history[phase][epoch].append(s)
        else: self.history[phase].update({epoch:[s]})

    def set_prog_bar_desc(self, prog_bar, c):
        batch_size = len(self.names)
        prog_bar.set_description(report.format(c, batch_size, self.loss.item(),
                                               self.loss_mm, self.frames,
                                               self.av_cor))

    def get_epoch(self):
        """
        since checkpoints/models are only saved after validation,
        we only look into epoch of validation set here.
        If process was cancelled in 1st train, start new history
        """
        if dict(self.history.get('validation'))=={}:
            self.init_history()
            self.__init__(self.args, self.input_dim, self.output_dim)
            return 0
        else:
            last_epoch = max(dict(self.history.get('validation')).keys())
            return last_epoch

    def compute_batch_averages(self, phase):
        """
        average loss (norm & mm), and cor calculated over batch for phase
        normalize over entire batch instead of retaining single batch samples
        """
        b_lm, b_l, b_c = float(0), float(0), float(0)
        total_in_batch = sum([i[0] for i in dict(self.history[phase])[self.epoch]])
        for n, f, l, lm, c in dict(self.history[phase])[self.epoch]:
            b_c  += c * n  / total_in_batch
            b_l  += l * n  / total_in_batch
            b_lm += lm * n / total_in_batch
        del self.history[phase][self.epoch]
        self.history[phase].update({self.epoch:{b_l, b_c, b_lm}})
        print('\nOverall {0}: lossmm = {1:.3f} cor = {2:.3f}'.format(phase, b_lm, b_c))
        if phase=='validation': return b_l

def save_model(model, model_path):
    torch.save(model.state_dict(), model_path)

def load_model(model, model_path, device):
    model.load_state_dict(torch.load(model_path, device.type))

def save_checkpoint(model, history, filepath):
    state = {'epoch': model.epoch,
             'state_dict': model.state_dict(),
             'optimizer': model.optimizer.state_dict(),
             'history': model.history}
    torch.save(state, filepath)

def load_checkpoint(model, args, device):
    checkpoint = torch.load(args.checkpoint_dir, map_location = device.type)
    model.epoch = checkpoint['epoch']
    model.load_state_dict(checkpoint['state_dict'])
    model.to(device)
    model.optimizer.load_state_dict(checkpoint['optimizer'])
    model.history = checkpoint['history']

def get_knn(batch_out_seq, seq_lens, kneighbors):
    from sklearn.manifold import locally_linear_embedding as lle
    from sklearn.neighbors import NearestNeighbors
    import numpy.matlib as matlib
    from numpy.linalg import inv

    # By default the knn function below will return the same vector.
    # e.g. k = 4 vector: 0, neighbors [0, 1, 2, 3]. We want [1, 2, 3, 4]
    # therefore add one to get distinct from the reference.
    knn = NearestNeighbors(n_neighbors = kneighbors + 1)
    pdb.set_trace()
    for sample, l in zip(batch_out_seq, seq_lens):
        W = np.zeros((kneighbors,sample.shape[0]))
        knn.fit(sample[:l,:])
        k_nn, k_ni = knn.kneighbors(sample[:l,:], return_distance=True)
        k_nn, k_ni = k_nn[:,1:], k_ni[:,1:]
        for i in range(l):
            Z = sample[:l,][k_ni[i]] - matlib.repmat(sample[:l,][i], kneighbors, 1)
            C = Z @ Z.T
            W[:,i] = inv(C)*np.ones((kneighbors, 1))
            W[:,i] = W[:,i]/W[:,i].sum(axis=1,keepdims=1)

def get_device(cuda):
    if torch.cuda.is_available() and cuda: return torch.device('cuda')
    else: return torch.device('cpu')

def validate(model, data, epoch, count):
    val_prog_bar = data.get_prog_bar(data.zip_validation_batches(),
                                     'validation', epoch)
    torch.no_grad(), model.eval()
    for pad_src_seq, pad_tgt_seq, pad_tgt_seqm, n2f in val_prog_bar:
        count += 1

        #if count > 4:  break # CODE TO STOP AND TEST
            #k_nearest_neighbors = get_knn(model.out_seq.detach().numpy(), model.seq_lens)
        model.get_lens_names_frames(n2f)
        model.forward(Variable(pad_src_seq)) # pass thru model
        model.get_cor(model.out_seq, pad_tgt_seq) # get correlation
        out_seq_mm = data.unnormalize_output(model.out_seq, False) #MSE(mm)
        model.get_losses(out_seq_mm, pad_tgt_seq, pad_tgt_seqm)

        model.update_info(epoch, 'validation')
        model.set_prog_bar_desc(val_prog_bar, count)

def train(model, data, epoch, count):
    train_prog_bar = data.get_prog_bar(data.zip_training_batches(),
                                       'training', epoch)
    model.train()
    for pad_src_seq, pad_tgt_seq, pad_tgt_seqm, n2f in train_prog_bar:
        count += 1
        #if count > 2: break # CODE TO STOP AND TEST
        model.get_lens_names_frames(n2f)
        model.forward(Variable(pad_src_seq)) # pass thru model
        model.get_cor(model.out_seq, pad_tgt_seq) # get correlation
        out_seq_mm = data.unnormalize_output(model.out_seq, False) #MSE(mm)
        model.get_losses(out_seq_mm, pad_tgt_seq, pad_tgt_seqm)
        model.propagate_loss()
        model.optimizer.zero_grad()
        model.update_info(epoch, 'training')
        model.set_prog_bar_desc(train_prog_bar, count)

def train_val_loop(model, data):
    bad_epochs = 0
    best_valid_loss = float('inf')
    print('Model Overview:\n\n{}'.format(model))
    while bad_epochs < args.patience:
        count = 0
        epoch += 1
        if args.shuffle: data.shuffle_batches()
        train(model, data, epoch, count)
        validate(model, data, epoch, count)
        model.compute_batch_averages('training')
        bv_l = model.compute_batch_averages('validation')
        if bv_l < best_valid_loss: # Decide whether to terminate training while loop
            best_valid_loss, bad_epochs = bv_l, 0
            if epoch % args.save_internal ==0: save_checkpoint(model, model.history, args.checkpoint_dir)
        else: bad_epochs += 1
    print('No validation set improvements observed for {:d} epochs. Early stop!'.format(args.patience))
    print('Saving Model...')
    try: save_model(model, args.save_model)
    except: save_checkpoint(model, model.history, args.checkpoint_dir)

def test(model, data, count, epoch):
    val_prog_bar = data.get_prog_bar(data.zip_validation_batches(),
                                     'validation', epoch)
    torch.no_grad(), model.eval()
    for pad_src_seq, pad_tgt_seq, pad_tgt_seqm, n2f in val_prog_bar:
        count += 1
        model.get_lens_names_frames(n2f)
        model.forward(Variable(pad_src_seq)) # pass thru model
        model.get_cor(model.out_seq, pad_tgt_seq) # get correlation
        out_seq_mm = data.unnormalize_output(model.out_seq, False) #MSE(mm)
        model.get_losses(out_seq_mm, pad_tgt_seq, pad_tgt_seqm)
        if count > 0: # break # CODE TO STOP AND TEST
            k_nearest_neighbors = get_knn(model.out_seq.detach().numpy(),
                                          model.seq_lens,
                                          kneighbors = 5)
            pdb.set_trace()
            model.out_seq.detach()
            k_nearest_neighbors

        model.set_prog_bar_desc(val_prog_bar, count)

def run():
    """
    Run Training and Validation over a number of epochs
    """
    d = get_device(args.cuda)
    data  = LoadInputOutput(args, d)
    model = BuildModel(args, d, data.input_dim, data.output_dim)
    data.get_batches()
    if args.cuda: model = model.cuda()
    if args.load_dir: load_model(model, args.load_dir, d)
    elif os.path.isfile(args.checkpoint_dir): load_checkpoint(model, args, d)
    else: model.init_history()

    epoch = 0 if model.is_new else model.get_epoch()
    if args.train: train_val_loop(model, data, epoch)
    if args.test: test(model, data, 0, epoch)

def main(args):
    run()

if __name__ == '__main__':
    args = get_args()
    main(args)


from __future__ import print_function
import abc

import numpy as np
import logging
import torch
import torch.utils.data
from torch import nn
from torch.autograd import Variable
from torch.nn import functional as F

from .nearest_embed import NearestEmbed

class AbstractAutoEncoder(nn.Module):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def encode(self, x):
        return

    @abc.abstractmethod
    def decode(self, z):
        return

    @abc.abstractmethod
    def forward(self, x):
        """model return (reconstructed_x, *)"""
        return

    @abc.abstractmethod
    def sample(self, size):
        """sample new images from model"""
        return

    @abc.abstractmethod
    def loss_function(self, **kwargs):
        """accepts (original images, *) where * is the same as returned from forward()"""
        return

    @abc.abstractmethod
    def latest_losses(self):
        """returns the latest losses in a dictionary. Useful for logging."""
        return

class VAE(nn.Module):
    """Variational AutoEncoder for MNIST
       Taken from pytorch/examples: https://github.com/pytorch/examples/tree/master/vae"""
    def __init__(self, kl_coef=1, **kwargs):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.kl_coef = kl_coef
        self.bce = 0
        self.kl = 0

    def encode(self, x):
        h1 = self.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h3 = self.relu(self.fc3(z))
        return self.tanh(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

    def sample(self, size):
        sample = Variable(torch.randn(size, 20))
        if self.cuda():
            sample = sample.cuda()
        sample = self.decode(sample).cpu()
        return sample

    def loss_function(self, x, recon_x, mu, logvar):
        self.bce = F.binary_cross_entropy(recon_x, x.view(-1, 784), size_average=False)
        batch_size = x.size(0)

        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        self.kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

        return self.bce + self.kl_coef*self.kl

    def latest_losses(self):
        return {'bce': self.bce, 'kl': self.kl}

class VQ_VAE(nn.Module):
    """Vector Quantized AutoEncoder for mnist"""
    def __init__(self, hidden=200, k=10, vq_coef=0.2, comit_coef=0.4, **kwargs):
        super(VQ_VAE, self).__init__()

        self.emb_size = k
        self.fc1 = nn.Linear(784, 400)
        self.fc2 = nn.Linear(400, hidden)
        self.fc3 = nn.Linear(hidden, 400)
        self.fc4 = nn.Linear(400, 784)

        self.emb = NearestEmbed(k, self.emb_size)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.vq_coef = vq_coef
        self.comit_coef = comit_coef
        self.hidden = hidden
        self.ce_loss = 0
        self.vq_loss = 0
        self.commit_loss = 0

    def encode(self, x):
        h1 = self.relu(self.fc1(x))
        h2 = self.fc2(h1)
        return h2.view(-1, self.emb_size, int(self.hidden / self.emb_size))

    def decode(self, z):
        h3 = self.relu(self.fc3(z))
        return self.tanh(self.fc4(h3))

    def forward(self, x):
        z_e = self.encode(x.view(-1, 784))
        z_q, _ = self.emb(z_e, weight_sg=True).view(-1, self.hidden)
        emb, _ = self.emb(z_e.detach()).view(-1, self.hidden)
        return self.decode(z_q), z_e, emb

    def sample(self, size):
        sample = Variable(torch.randn(size, self.emb_size, int(self.hidden / self.emb_size)))
        if self.cuda():
            sample = sample.cuda()
        emb, _ = self.emb(sample)
        sample = self.decode(emb(sample).view(-1, self.hidden)).cpu()
        return sample

    def loss_function(self, x, recon_x, z_e, emb):
        self.ce_loss = F.binary_cross_entropy(recon_x, x.view(-1, 784))
        self.vq_loss = F.mse_loss(emb, z_e.detach())
        self.commit_loss = F.mse_loss(z_e, emb.detach())

        return self.ce_loss + self.vq_coef*self.vq_loss + self.comit_coef*self.commit_loss

    def latest_losses(self):
        return {'cross_entropy': self.ce_loss, 'vq': self.vq_loss, 'commitment': self.commit_loss}

class ResBlock(nn.Module):
    def __init__(self, in_channels, channels, bn=False):
        super(ResBlock, self).__init__()

        layers = [
            nn.ReLU(),
            nn.Conv2d(in_channels, channels, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, channels, kernel_size=1, stride=1, padding=0)]
        if bn:
            layers.insert(2, nn.BatchNorm2d(channels))
        self.convs = nn.Sequential(*layers)

    def forward(self, x):
        return x + self.convs(x)

class CVAE(AbstractAutoEncoder):
    def __init__(self, d, kl_coef=0.1, **kwargs):
        super(CVAE, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(3, d // 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d // 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(d // 2, d, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),
            ResBlock(d, d, bn=True),
        )
        self.decoder = nn.Sequential(
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),

            nn.ConvTranspose2d(d, d // 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d//2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(d // 2, 3, kernel_size=4, stride=2, padding=1, bias=False),
        )
        self.f = 8
        self.d = d
        self.fc11 = nn.Linear(d * self.f ** 2, d * self.f ** 2)
        self.fc12 = nn.Linear(d * self.f ** 2, d * self.f ** 2)
        self.kl_coef = kl_coef
        self.kl_loss = 0
        self.mse = 0

    def encode(self, x):
        h1 = self.encoder(x)
        h1 = h1.view(-1, self.d * self.f ** 2)
        return self.fc11(h1), self.fc12(h1)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        z = z.view(-1, self.d, self.f, self.f)
        h3 = self.decoder(z)
        return F.tanh(h3)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

    def sample(self, size):
        sample = Variable(torch.randn(size, self.d * self.f ** 2), requires_grad=False)
        if self.cuda():
            sample = sample.cuda()
        return self.decode(sample).cpu()

    def loss_function(self, x, recon_x, mu, logvar):
        self.mse = F.mse_loss(recon_x, x)
        batch_size = x.size(0)

        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        self.kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        # Normalise by same number of elements as in reconstruction
        self.kl_loss /= batch_size * 3 * 1024

        # return mse
        return self.mse + self.kl_coef * self.kl_loss

    def latest_losses(self):
        return {'mse': self.mse, 'kl': self.kl_loss}

class VQ_CVAE(nn.Module):
    def __init__(self, d, k=10, bn=True, vq_coef=1, commit_coef=0.5, num_channels=3, **kwargs):
        super(VQ_CVAE, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(num_channels, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            nn.Conv2d(d, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            ResBlock(d, d, bn),
            nn.BatchNorm2d(d),
            ResBlock(d, d, bn),
            nn.BatchNorm2d(d),
        )
        self.decoder = nn.Sequential(
            ResBlock(d, d),
            nn.BatchNorm2d(d),
            ResBlock(d, d),
            nn.ConvTranspose2d(d, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(d, num_channels, kernel_size=4, stride=2, padding=1),
        )
        self.d = d
        self.emb = NearestEmbed(k, d)
        self.vq_coef = vq_coef
        self.commit_coef = commit_coef
        self.mse = 0
        self.vq_loss = Variable(torch.zeros(1))
        self.commit_loss = 0

        for l in self.modules():
            if isinstance(l, nn.Linear) or isinstance(l, nn.Conv2d):
                l.weight.detach().normal_(0, 0.02)
                torch.fmod(l.weight, 0.04)
                nn.init.constant_(l.bias, 0)

        self.encoder[-1].weight.detach().fill_(1 / 40)

        self.emb.weight.detach().normal_(0, 0.02)
        torch.fmod(self.emb.weight, 0.04)

    def encode(self, x):
        return self.encoder(x)

    def decode(self, x):
        return F.tanh(self.decoder(x))

    def forward(self, x):
        z_e = self.encode(x)
        self.f = z_e.shape[-1]
        z_q, argmin = self.emb(z_e, weight_sg=True)
        emb, _ = self.emb(z_e.detach())
        return self.decode(z_q), z_e, emb, argmin

    def sample(self, size):
        sample = Variable(torch.randn(size, self.d, self.f, self.f), requires_grad=False)
        if self.cuda():
            sample = sample.cuda()
        emb, _ = self.emb(sample)
        return self.decode(emb.view(size, self.d, self.f, self.f)).cpu()

    def loss_function(self, x, recon_x, z_e, emb, argmin):
        self.mse = F.mse_loss(recon_x, x)

        self.vq_loss = torch.mean(torch.norm((emb - z_e.detach())**2, 2, 1))
        self.commit_loss = torch.mean(torch.norm((emb.detach() - z_e)**2, 2, 1))

        return self.mse + self.vq_coef*self.vq_loss + self.commit_coef*self.commit_loss

    def latest_losses(self):
        return {'mse': self.mse, 'vq': self.vq_loss, 'commitment': self.commit_loss}

    def print_atom_hist(self, argmin):

        argmin = argmin.detach().cpu().numpy()
        unique, counts = np.unique(argmin, return_counts=True)
        logging.info(counts)
        logging.info(unique)


"""Example to perform the combination between ANTARES and Super-Kamiokande.

This is taking benefit of the the specific formats described in ANTARES and Super-Kamiokande examples.
All provided values are dummy.
"""

import healpy as hp
import numpy as np

import jang.gw
import jang.limits
import jang.results
from jang.parameters import Parameters
from jang.neutrinos import BackgroundFixed, Detector, SuperDetector

from examples.superkamiokande import EffectiveAreaSK

def single_event(
    gwname: str,
    gwdbfile: str,
    ant_results: dict,
    sk_results: dict,
    pars: Parameters,
    dbfile: str = None,
):

    database_gw = jang.gw.Database(gwdbfile)
    database_res = jang.results.Database(dbfile)

    antares = Detector("examples/input_files/detector_antares.yaml")
    sk = Detector("examples/input_files/detector_superk.yaml")
    effarea_sk = [
        EffectiveAreaSK(filename=sk_results["effarea"], sample=s) for s in sk.samples
    ]

    # Combination
    ant_sk = SuperDetector("ANTARES + Super-Kamiokande")
    ant_sk.add_detector(antares)
    ant_sk.add_detector(sk)

    gw = database_gw.find_gw(gwname)

    antares.set_acceptances(ant_results["acceptances"], pars.spectrum, pars.nside)
    bkg_ant = [BackgroundFixed(b) for b in ant_results["nbkg"]]
    antares.set_observations(ant_results["nobs"], bkg_ant)
    pars.nside = antares.get_acceptances(pars["spectrum"])[0].nside

    accs = [
        effarea.to_acceptance(sk, pars.nside, gw.jd, pars.spectrum)
        for effarea in effarea_sk
    ]
    sk.set_acceptances(accs, pars.spectrum, pars.nside)
    bkg_sk = [BackgroundFixed(b) for b in sk_results["nbkg"]]
    sk.set_observations(sk_results["nobs"], bkg_sk)

    limit_flux = jang.limits.get_limit_flux(ant_sk, gw, pars)
    limit_etot = jang.limits.get_limit_etot(ant_sk, gw, pars)
    limit_fnu = jang.limits.get_limit_fnu(ant_sk, gw, pars)
    database_res.add_entry(
        ant_sk, gw, pars, limit_flux, limit_etot, limit_fnu, None, None, None
    )

    if dbfile is not None:
        database_res.save()

if __name__ == "__main__":

    pars = Parameters("examples/input_files/config.yaml")
    pars.set_models("x**-2", jang.conversions.JetIsotropic())

    gw_db_file = "examples/input_files/gw_catalogs/database_example.csv"
    npix = hp.nside2npix(8)
    ant_results = {
        "nobs": [0, 0, 0, 0],
        "nbkg": [0, 0, 0, 0],
        "acceptances": [np.ones(npix), np.zeros(npix), np.ones(npix), np.zeros(npix)],
    }
    sk_results = {
        "nobs": [0, 0, 0],
        "nbkg": [0, 0, 0],
        "effarea": "examples/input_files/effarea_superk.root",
    }
    single_event("GW190412", gw_db_file, ant_results, sk_results, pars)


#!/usr/bin/env python3
#--coding:utf-8 --
"""
callTransLoops.py
"""

#sys
import os
import sys
import json
from glob import glob
from datetime import datetime
from collections import Counter

#3rd
import joblib
import numpy as np
import pandas as pd
from tqdm import tqdm
from joblib import Parallel, delayed
from sklearn import linear_model
from scipy.stats import hypergeom, binom, poisson, combine_pvalues

#cLoops
from cLoops2.ds import Loop, XY
from cLoops2.io import parseIxy, ixy2pet, loops2juiceTxt, loops2washuTxt, ixy2pet, updateJson, loops2txt
from cLoops2.geo import checkLoopOverlap, combineLoops
from cLoops2.settings import *
from cLoops2.blockDBSCAN import blockDBSCAN as DBSCAN
from cLoops2.callCisLoops import getPerRegions, selSigLoops, estAnchorSig

def runTransDBSCANLoops(fixy, eps, minPts):
    """
    Run DBSCAN to detect interactions for one .ixy file.
    @param fixy: str, .ixy file name
    @param eps: int, eps for DBSCAN
    @param minPts: int, minPts for DBSCAN
    """
    loops, loopReads = [], []
    key, mat = parseIxy(fixy, cut=0)
    #gave mat each PET a id
    mat2 = np.zeros((mat.shape[0],3))
    mat2[:,0] = np.arange(mat.shape[0])
    mat2[:,1] = mat[:,0]
    mat2[:,2] = mat[:,1]
    mat = mat2

    #data for interaction records, read for readId
    report = "%s \t Clustering %s and %s using eps %s, minPts %s\n" % (
        datetime.now(), key[0], key[1], eps, minPts)
    sys.stderr.write(report)
    db = DBSCAN(mat, eps, minPts)
    labels = pd.Series(db.labels)
    mat = pd.DataFrame(mat[:, 1:].astype("int"),
                       index=mat[:, 0],
                       columns=["X", "Y"])
    nlabels = set(labels.values)
    #collect clusters
    for label in nlabels:
        los = list(labels[labels == label].index)
        loopReads.extend(los)
        sub = mat.loc[los, :]
        if int(np.min(sub["X"])) == int(np.max(sub["X"])) or int(
                np.min(sub["Y"])) == int(np.max(sub["Y"])):
            continue
        #define loops
        loop = Loop()
        loop.rab = sub.shape[0]
        loop.chromX = key[0]
        loop.chromY = key[1]
        loop.x_start = int(np.min(sub["X"]))
        loop.x_end = int(np.max(sub["X"]))
        loop.x_center = (loop.x_start + loop.x_end) / 2
        loop.y_start = int(np.min(sub["Y"]))
        loop.y_end = int(np.max(sub["Y"]))
        loop.y_center = (loop.y_start + loop.y_end) / 2
        loop.distance = -1
        loop.cis = False
        loops.append(loop)
    report = "%s \t Clustering %s and %s finished. Estimated %s reads for %s candidate loops. \n" % (
        datetime.now(), key[0], key[1], len(loopReads), len(loops))
    sys.stderr.write(report)
    return "-".join(key), loops

#related
def parallelRunTransDBSCANLoops(meta, eps, minPts, cpu=1):
    """
    Paralle version of runCisDBSCANLoops
    @param meta: meta information parsed form petMeta.json
    @param eps: int, eps for DBSCAN
    @param minPts: int, minPts for DBSCAN
    """
    ds = Parallel(n_jobs=cpu,backend="multiprocessing")(delayed(runTransDBSCANLoops)(
        meta["data"]["trans"][key]["ixy"], eps, minPts)
                              for key in meta["data"]["trans"].keys())
    loops = {}
    for d in ds:
        if d is not None and len(d[1]) > 0:
            key, di = d[0], d[1]
            loops[key] = di
    return loops

def estLoopSig(key,
               loops,
               fixy,
               minPts=5,
               pseudo=1,
               peakPcut=1e-5,
               peakFccut=2,
               countDiffCut=10):
    """
    Estimate the loop statstical significance for one chromosomal.
    @param loops: list of Loop object
    @param fixy: cLoops2 pre generated .ixy file
    @param minPts: int, minPts
    """
    xy = ixy2pet(fixy, cut=0)
    N = xy.number
    print("%s \t Estimate significance for %s candidate interactions in %s." %
          (datetime.now(), len(loops), key))
    nloops = []
    for loop in tqdm(loops):
        ra, rb, rab = xy.queryLoop(loop.x_start, loop.x_end, loop.y_start,
                                   loop.y_end)
        ra, rb, rab = len(ra), len(rb), len(rab)
        if rab < minPts:
            continue
        loop.ra = ra
        loop.rb = rb
        loop.rab = rab
        #unbalanced anchor density, to avoid lines, unknow reason for lines, maybe stripes
        if ra / float(rb) > countDiffCut or rb / float(ra) > countDiffCut:
            continue
        if (loop.x_end -
                loop.x_start) / (loop.y_end - loop.y_start) > countDiffCut or (
                    loop.y_end - loop.y_start) / (loop.x_end -
                                                  loop.x_start) > countDiffCut:
            continue
        lowerra, lowerrb, lowerrab = xy.queryLoop(
            loop.x_start - (loop.x_end - loop.x_start), loop.x_start,
            loop.y_start - (loop.y_end - loop.y_start), loop.y_start)  #p2ll
        loop.P2LL = float(rab) / max(len(lowerrab), pseudo)
        px, esx = estAnchorSig(xy, loop.x_start, loop.x_end)
        py, esy = estAnchorSig(xy, loop.y_start, loop.y_end)
        loop.x_peak_poisson_p_value = px
        loop.x_peak_es = esx
        loop.y_peak_poisson_p_value = py
        loop.y_peak_es = esy
        #hypergeometric p-value
        hyp = max([1e-300, hypergeom.sf(rab - 1.0, N, ra, rb)])
        #start caculate the permutated background
        rabs, nbps = [], []
        nas, nbs = getPerRegions(loop, xy)
        for na in nas:
            nac = float(len(na))
            for nb in nbs:
                nbc = float(len(nb))
                nrab = float(len(na.intersection(nb)))
                #collect the value for poisson test
                rabs.append(nrab)
                #collect the possibility for following binomial test
                if nac > 0 and nbc > 0:
                    den = nrab / (nac * nbc)
                    nbps.append(den)
                else:
                    nbps.append(0.0)
        rabs, nbps = np.array(rabs), np.array(nbps)
        mrabs = float(np.mean(rabs))
        mbps = np.mean(nbps)
        #local fdr
        fdr = len(rabs[rabs > rab]) / float(len(rabs))
        #enrichment score
        es = rab / max(mrabs, pseudo)
        #simple possion test
        pop = max([1e-300, poisson.sf(rab - 1.0, mrabs)])
        #simple binomial test
        nbp = max([
            1e-300, binom.sf(rab - 1.0, ra * rb, mbps)
        ])  #the p-value is quit similar to that of cLoops 1 binomial test
        loop.FDR = fdr
        loop.ES = es
        loop.density = float(
            loop.rab) / (loop.x_end - loop.x_start + loop.y_end -
                         loop.y_start) / N * 10.0**9
        loop.hypergeometric_p_value = hyp
        loop.poisson_p_value = pop
        loop.binomial_p_value = nbp
        nloops.append(loop)
        #print(ra,rb,rab,mrabs,es,fdr,hyp,pop,nbp,n,nbp2)
    return key, nloops

def markSigLoops(key, loops):
    """
    Mark the significance of different loops.
    """
    sig = lambda x: True if x.binomial_p_value <= 1e-10 and x.FDR <= 0.05 and loop.ES >= 2 else False
    for loop in loops:
        if sig(loop):
            loop.significant = 1
        else:
            loop.significant = 0
    return key, loops

def callTransLoops(
        predir,
        fout,
        logger,
        eps=[2000, 5000],
        minPts=[5, 10],
        cpu=1,
        filter=False,
        washU=False,
        juicebox=False,
):
    """
    Call inter-chromosomal loops parallel.
    @param metaf: str, petMeta.json file for calling peaks
    @param eps: list
    @param minPts: list
    """

    metaf = predir + "/petMeta.json"
    meta = json.loads(open(metaf).read())
    ## step 1 find the candidate loops by running multiple times of clustering
    loops = {}  #candidate loops
    for ep in eps:
        for minPt in minPts:
            loops_2 = parallelRunTransDBSCANLoops(meta, ep, minPt, cpu=cpu)
            loops = combineLoops(loops, loops_2)
    ## step 2 determine the statstical significance of candidate loops
    logger.info("Estimating loop statstical significance.")
    ds = Parallel(n_jobs=cpu,backend="multiprocessing")(delayed(estLoopSig)(
        key,
        loops[key],
        meta["data"]["trans"][key]["ixy"],
        minPts=max(minPts),
    ) for key in loops.keys())
    nds = {}
    for d in ds:
        nds[d[0]] = d[1]

    #mark the significant loops
    ds = Parallel(n_jobs=cpu,backend="multiprocessing")(delayed(markSigLoops)(key, nds[key])
                              for key in nds.keys())
    nds = {}
    for d in ds:
        nds[d[0]] = d[1]

    ## step 4 for the overlapped loops, output the most significant one
    logger.info("Selecting the most significant loops of overlapped ones. ")
    ds = Parallel(n_jobs=cpu,backend="multiprocessing")(delayed(selSigLoops)(key, nds[key])
                              for key in nds.keys())
    nds = {}
    for d in ds:
        nds[d[0]] = d[1]
    ds = Parallel(n_jobs=cpu,backend="multiprocessing")(delayed(selSigLoops)(key, nds[key])
                              for key in nds.keys())
    nds = {}
    for d in ds:
        nds[d[0]] = d[1]
    loops = []
    for d in ds:
        loops.extend(d[1])
    ## step 5 output
    logger.info("Output %s loops to %s_loops.txt" % (len(loops), fout))
    loops2txt(loops, fout + "_trans_loops.txt")
    if washU:
        loops2washuTxt(loops, fout + "_trans_loops_washU.txt")
    if juicebox:
        loops2juiceTxt(loops, fout + "_trans_loops_juicebox.txt")


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 10 10:32:04 2018

@author: smrak
"""
import numpy as np
from datetime import datetime
import apexpy as ap
import igrf12
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.feature.nightshade import Nightshade
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from sunrise import terminator as ter

projection_dict = {'stereo': ccrs.Stereographic(),
              'merc': ccrs.Mercator(),
              'plate': ccrs.PlateCarree(),
              'lambert': ccrs.LambertConformal(),
              'mollweide': ccrs.Mollweide(),
              'north': ccrs.NorthPolarStereo(),
              'south': ccrs.NorthPolarStereo(),
              'ortographic': ccrs.Orthographic()
                  }

def plotCartoMap(latlim=[0, 75], lonlim=[-40, 40], parallels=None, meridians=None,
                 pole_center_lon=0,figsize=(12, 8), terrain=False, ax=False,
                 projection='stereo', title='', resolution='110m', lon0=None,lat0=None,
                 states=True, grid_linewidth=0.5, grid_color='black',
                 grid_linestyle='--', background_color=None, border_color='k',
                 figure=False, nightshade=False, ns_alpha=0.1,
                 apex=False, igrf=False, date=None,
                 mlat_levels=None, mlon_levels=None, alt_km=0.0,
                 mlon_colors='blue', mlat_colors='red', mgrid_width=1,
                 mgrid_labels=True, mgrid_fontsize=12, mlon_cs='mlon',
                 incl_levels=None, decl_levels=None, igrf_param='incl',
                 mlon_labels=True, mlat_labels=True, mgrid_style='--',
                 label_colors='k', apex_alt=0,
                 decl_colors='k', incl_colors='k',
                 terminator=False, terminator_altkm=350,
                 polarization_terminator=False, pt_hemisphere='north',
                 ter_color='red', ter_style='--', ter_width=2,
                 midnight=False, midnight_colors='m', midnight_width=2,
                 midnight_style='--'):
    if lonlim is None or lonlim == []:
        lonlim = [-180, 180]
    STATES = cfeature.NaturalEarthFeature(
        category='cultural',
        name='admin_1_states_provinces_lines',
        scale='50m',
        facecolor='none')

    if not ax:
        if figsize is None:
            fig = plt.figure()
        else:
            fig = plt.figure(figsize=figsize)

        if projection == 'stereo':
            ax = plt.axes(projection=ccrs.Stereographic(central_longitude=(sum(lonlim)/2)))
        elif projection == 'merc':
            ax = plt.axes(projection=ccrs.Mercator())
        elif projection == 'plate':
            ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=(sum(lonlim)/2)))
        elif projection == 'lambert':
            ax = plt.axes(projection=ccrs.LambertConformal(central_longitude=(sum(lonlim)/2),
                                                           central_latitude=(sum(latlim)/2)))
        elif projection == 'mollweide':
            ax = plt.axes(projection=ccrs.Mollweide(central_longitude=(sum(lonlim)/2)))
        elif projection == 'north':
            if lon0 is None:
                lon0 = 0
            ax = plt.axes(projection=ccrs.NorthPolarStereo(central_longitude=lon0))
        elif projection == 'south':
            if lon0 is None:
                lon0 = 0
            ax = plt.axes(projection=ccrs.SouthPolarStereo(central_longitude=lon0))
        elif projection == 'ortographic':
            if lon0 is None:
                lon0 = 0
            if lat0 is None:
                lat0 = 0
            ax = plt.axes(projection=ccrs.Orthographic(central_longitude=lon0, central_latitude=lat0))
        else:
            print ("Projection is invalid. Please enter the right one. \n \
                   'stereo', 'merc', 'plate', 'lambret', mollweide', 'north, 'south', 'ortographic'")
            return 0
    if background_color is not None:
        ax.background_patch.set_facecolor(background_color)
    ax.set_title(title)
    ax.coastlines(color=border_color, resolution=resolution)  # 110m, 50m or 10m
    if states:
        ax.add_feature(STATES, edgecolor=border_color)
    ax.add_feature(cfeature.BORDERS, edgecolor=border_color)
    if terrain:
        ax.stock_img()
    if nightshade:
        assert date is not None
        assert ns_alpha is not None
        ax.add_feature(Nightshade(date, ns_alpha))
    # Draw Parralels
    if projection == 'merc' or projection == 'plate':
        if isinstance(meridians, np.ndarray):
            meridians = list(meridians)
        if isinstance(parallels, np.ndarray):
            parallels = list(parallels)

        gl = ax.gridlines(crs=ccrs.PlateCarree(), color=grid_color, draw_labels=False,
                          linestyle=grid_linestyle, linewidth=grid_linewidth)
        if meridians is None:
            gl.xlines = False
        else:
            if len(meridians) > 0:
                gl.xlocator = mticker.FixedLocator(meridians)
                gl.xlabels_bottom = True
            else:
                gl.ylines = False

        if parallels is None:
            gl.ylines = False
        else:
            if len(parallels) > 0:
                gl.ylocator = mticker.FixedLocator(parallels)
                gl.ylabels_left = True
            else:
                gl.ylines = False
    else:
        gl = ax.gridlines(crs=ccrs.PlateCarree(), color=grid_color, draw_labels=False,
                          linestyle=grid_linestyle, linewidth=grid_linewidth)
        if meridians is None:
            gl.xlines = False
        else:
            gl.xlocator = mticker.FixedLocator(meridians)
        if parallels is not None:
            if isinstance(parallels, np.ndarray):
                parallels = list(parallels)
            gl.ylocator = mticker.FixedLocator(parallels)
        else:
            gl.ylines = False

    # Geomagnetic coordinates @ Apex
    if apex:
        if date is None:
            date = datetime(2017, 12, 31, 0, 0, 0)
        assert isinstance(date, datetime)

        A = ap.Apex(date = date)
        # Define levels and ranges for conversion
        if mlon_cs == 'mlt':
            if mlon_levels is None:
                mlon_levels = np.array([])
                mlon_range = np.arange(0, 24.01, 0.01)
            elif isinstance(mlon_levels, bool):
                if mlon_levels == False:
                    mlon_levels = np.array([])
                    mlon_range = np.arange(0, 24.01, 0.01)
            else:
                mlon_range = np.arange(mlon_levels[0], mlon_levels[-1]+0.1, 0.01)
        else:
            if mlon_levels is None:
                mlon_levels = np.array([])
                mlon_range = np.arange(-180,180,0.5)
            elif isinstance(mlon_levels, bool):
                if mlon_levels == False:
                    mlon_levels = np.array([])
                    mlon_range = np.arange(-180, 181, 0.1)
            else:
                mlon_range = np.arange(mlon_levels[0], mlon_levels[0]+362, 0.1)
        if mlat_levels is None:
            mlat_levels = np.arange(-90, 90.1, 1)
        mlat_range = np.arange(mlat_levels[0], mlat_levels[-1]+0.1, 0.1)

        # Do meridans
        for mlon in mlon_levels:
            MLON = mlon * np.ones(mlat_range.size)
            if mlon_cs == 'mlt':
                y, x = A.convert(mlat_range, MLON, 'mlt', 'geo', datetime=date, height=apex_alt)
            else:
                y, x  = A.convert(mlat_range, MLON, 'apex', 'geo', height=apex_alt)
            mlat_mask_extent = (y >=  latlim[0] + 0.2) & (y <=  latlim[1] - 0.2)
            # Plot meridian
            inmap = np.logical_and(x >= lonlim[0], x <= lonlim[1])
            if np.sum(inmap) > 10:
                ax.plot(np.unwrap(x[mlat_mask_extent], 180),
                        np.unwrap(y[mlat_mask_extent], 90), c=mlon_colors,
                         lw=mgrid_width, linestyle=mgrid_style, zorder=90,
                         transform=ccrs.PlateCarree())

            # Labels
            if mlon_labels:
                ix = abs(y-np.mean(latlim)).argmin()
                mx = x[ix] - 1 if mlon >=10 else x[ix] - 0.5
                my = np.mean(latlim)
                if np.logical_and(mx >= lonlim[0], mx <= lonlim[1]):
                    if mlon_cs == 'mlt' and mlon != 0:
                        ax.text(mx, my, str(int(mlon)), color=label_colors,
                                 fontsize=14, backgroundcolor='white',
                                 transform=ccrs.PlateCarree())
                    elif mlon_cs != 'mlt' and mlon != 360:
                        ax.text(mx, my, str(int(mlon)), color=label_colors,
                                 fontsize=14, backgroundcolor='white',
                                 transform=ccrs.PlateCarree())
        # Do parallels
        for mlat in mlat_levels:
            MLAT = mlat * np.ones(mlon_range.size)
            if mlon_cs == 'mlt':
                gy, gx = A.convert(MLAT, mlon_range, 'mlt', 'geo', datetime=date, height=apex_alt)
            else:

                gy, gx = A.convert(MLAT, mlon_range, 'apex', 'geo', datetime=date, height=apex_alt)
            inmap = np.logical_and(gy >= latlim[0], gy <= latlim[1])
            if np.sum(inmap) > 20:
                ax.plot(np.unwrap(gx, 180), np.unwrap(gy, 90), c=mlat_colors,
                         lw=mgrid_width, linestyle=mgrid_style, zorder=90,
                         transform=ccrs.PlateCarree())

            # Labels
            if mlat_labels:
                ix = abs(gx-np.mean(lonlim)).argmin()
                mx = np.mean(lonlim)
                my = gy[ix] - 0.5
                if np.logical_and(mx >= lonlim[0], mx <= lonlim[1]) and \
                np.logical_and(my >= latlim[0], my <= latlim[1]):
                    ax.text(mx, my, str(int(mlat)), color=label_colors,
                             fontsize=14, backgroundcolor='white',
                             transform=ccrs.PlateCarree())
    if igrf:
        glon = np.arange(lonlim[0]-40, lonlim[1] + 40.1, 0.5)
        glat = np.arange(-90, 90 + 0.1, 0.5)
        longrid, latgrid = np.meshgrid(glon, glat)
        mag = igrf12.gridigrf12(t=date, glat=latgrid, glon=longrid, alt_km=alt_km)
        if decl_levels is not None:
            z = mag.decl.values
            for declination in decl_levels:
                ax.contour(longrid, latgrid, z, levels=declination,  zorder=90,
                             colors=decl_colors, transform=ccrs.PlateCarree())
        if incl_levels is not None:
            z = mag.incl.values
            for inclination in incl_levels:
                ax.contour(longrid, latgrid, z, levels=declination,  zorder=90,
                                 colors=incl_colors, transform=ccrs.PlateCarree())
    # Terminators
    if terminator:
        assert date is not None
        if not isinstance(terminator_altkm, list):
            terminator_altkm = [terminator_altkm]
        for takm in terminator_altkm:
            try:
                glon_ter, glat_ter = ter.get_terminator(date, alt_km = takm)
                if glon_ter is not None and glat_ter is not None:
                    if isinstance(glon_ter, list):
                        for i in range(len(glon_ter)):
                            ax.plot(np.unwrap(glon_ter[i], 180), np.unwrap(glat_ter[i], 90),
                                c=ter_color, lw=ter_width, ls=ter_style, zorder=90,
                                transform=ccrs.PlateCarree())
                    else:
                        ax.plot(np.unwrap(glon_ter, 180), np.unwrap(glat_ter, 90),
                                c=ter_color, lw=ter_width, ls=ter_style, zorder=90,
                                transform=ccrs.PlateCarree())
            except:
                pass
    if midnight:
        mlat_range = np.arange(-89.9, 90.1, 0.1)
        MLON = 0 * np.ones(mlat_range.size)
        if mlon_cs == 'mlt':
            y, x = A.convert(mlat_range, MLON, 'mlt', 'geo', datetime=date, height=apex_alt)
        else:
            y, x  = A.convert(mlat_range, MLON, 'apex', 'geo', height=apex_alt)
        mlat_mask_extent = (y >=  latlim[0] + 0.2) & (y <=  latlim[1] - 0.2)
        # Plot meridian
        inmap = np.logical_and(x >= lonlim[0], x <= lonlim[1])
        if np.sum(inmap) > 10:
            ax.plot(np.unwrap(x[mlat_mask_extent], 180),
                    np.unwrap(y[mlat_mask_extent], 90), c=midnight_colors,
                     lw=midnight_width, linestyle=midnight_style, zorder=90,
                     transform=ccrs.PlateCarree())
    # Set Extent
    if projection == 'north' or projection == 'south':
        import matplotlib.path as mpath
        ax.set_extent([-180, 181, latlim[0], latlim[1]], crs=ccrs.PlateCarree())
        theta = np.linspace(0, 2*np.pi, 100)
        center, radius = [0.5, 0.5], 0.5
        verts = np.vstack([np.sin(theta), np.cos(theta)]).T
        circle = mpath.Path(verts * radius + center)
        ax.set_boundary(circle, transform=ax.transAxes)
    elif lonlim[0] != -180 and lonlim[1] != 180:
        ax.set_extent([lonlim[0], lonlim[1], latlim[0], latlim[1]], crs=ccrs.PlateCarree())#ccrs.PlateCarree())

    if 'fig' in locals():
        return fig, ax
    else:
        return ax


# WaveformToVolume Paraview filter

# Load the `WaveformDataReader` to access its VTK keys
# sys.path.append(os.path.dirname(__file__))
# from WaveformDataReader import WaveformDataReader

import numpy as np
import logging
import time
from vtkmodules.vtkCommonDataModel import vtkUniformGrid
from vtkmodules.vtkCommonCore import vtkDataArraySelection
from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase
from vtkmodules.numpy_interface import dataset_adapter as dsa
from paraview.util.vtkAlgorithm import smproxy, smproperty, smdomain
from paraview import util
from paraview.vtk.util import numpy_support as vtknp
import gwpv.plugin_util.timesteps as timesteps_util
import gwpv.plugin_util.data_array_selection as das_util
from gwpv import swsh_cache
import h5py
from os.path import exists
logger = logging.getLogger(__name__)

# A grid is created  here, and if analytical data created by creator.py
# is recognized iti is loaded into the grid...

gb_path = None

if exists("timeseparated.h5"):
    gb_path = "timeseparated.h5"
elif exists("gwpv/timeseparated.h5"):
    gb_path = "gwpv/timeseparated.h5"
else:
    gb_path = None

logger.warning("Retrieving data from "+str(gb_path)+".")

if gb_path != None:
    with h5py.File(gb_path, "r") as f:
        h = f['Extrapolated_N2.dir']
        tick = np.real(h['t_values.dir'])
        strains = np.zeros((len(tick), 1000000), dtype=np.complex)
        for i in range(len(tick)):
            strains[i] += np.real(h['t_{}.dir'.format(tick[i])]) + 1j*np.imag(h['t_{}.dir'.format(tick[i])])
else:
     logger.warning('No Bremsstrahlung data found, proceeding with merger visualization.')

def get_mode_name(l, abs_m):
    return "({}, {}) Mode".format(l, abs_m)

# Reproduces `spherical_functions.LM_index` so we don't need to import the
# `spherical_functions` module when using a cached SWSH grid
def LM_index(ell, m, ell_min):
    return ell * (ell + 1) - ell_min ** 2 + m

def smoothstep(x):
    return np.where(x < 0, 0, np.where(x <= 1, 3 * x**2 - 2 * x**3, 1))

def activation(x, width):
    return smoothstep(x / width)

def deactivation(x, width, outer):
    return smoothstep((outer - x) / width)

# Caching
# When using SwshGrid input this is not necessary anymore
# These are global variables because setting them on the filter object appears
# to trigger a "ModifiedEvent" so the data is recomputed.
_cached_swsh_grid = None
_cached_r = None
_cached_grid_id = None
def cached_swsh_grid(size, radial_scale, activation_offset, activation_width, deactivation_width, add_one_over_r_scaling, **swsh_grid_kwargs):
    global _cached_swsh_grid, _cached_r, _cached_grid_id
    grid_id = dict(size=size,
                   radial_scale=radial_scale,
                   activation_offset=activation_offset,
                   activation_width=activation_width,
                   deactivation_width=deactivation_width,
                   add_one_over_r_scaling=add_one_over_r_scaling)
    grid_id.update(swsh_grid_kwargs)
    if _cached_grid_id == grid_id:
        logger.debug("Using cached SWSHs grid from memory.")
        return _cached_swsh_grid, _cached_r
    else:
        logger.debug("No SWSH grid in memory, retrieving from disk cache.")
        swsh_grid, r = swsh_cache.cached_swsh_grid(size=size,
                                                   **swsh_grid_kwargs)
        # Apply screening
        screen = activation(r - activation_offset,
                            activation_width) * deactivation(
                                r, deactivation_width, size)
        swsh_grid *= screen.reshape(screen.shape + (1, ))
        # Apply radial scale
        r *= radial_scale
        if add_one_over_r_scaling:
            swsh_grid /= (r + 1.e-30).reshape(r.shape + (1,))
        # Cache and return
        _cached_swsh_grid = swsh_grid
        _cached_r = r
        _cached_grid_id = grid_id
        return swsh_grid, r

has_shown_warning_nonuniformly_sampled = False

@smproxy.filter(label="Waveform To Volume")
@smproperty.input(name="WaveformData", port_index=0)
@smdomain.datatype(dataTypes=["vtkTable"])
# TODO: We should be able to use a `SwshGrid` as a second input to this filter
# and use this class to compute the volume data for the waveform without having
# to generate the grid. Multiple inputs work fine, but for some reason
# `pv.LoadPlugin` doesn't load `SwshGrid`.
# @smproperty.input(name="GridData", port_index=1)
# @smdomain.datatype(dataTypes=["vtkUniformGrid"])
class WaveformToVolume(VTKPythonAlgorithmBase):
    def __init__(self):
        VTKPythonAlgorithmBase.__init__(
            self,
            # nInputPorts=2,
            nInputPorts=1,
            nOutputPorts=1,
            # Choosing `vtkUniformGrid` for the output for the following reasons:
            # - `vtkRectilinearGrid` doesn't support volume rendering
            #   (in Paraview v5.7.0 at least)
            # - The unstructured grids don't support the 'GPU Based'
            #   volume rendering mode, which can do shading and looks nice
            outputType='vtkUniformGrid')
        self.modes_selection = vtkDataArraySelection()
        # TODO: We should really retrieve the available modes from the input
        # info in `RequestInformation`, but the `WAVEFORM_MODES` keys is not
        # propagating downstream for some reason...
        # modes_arrays_key = WaveformDataReader.MODES_ARRAYS_KEY
        # if waveform_data_info.Has(modes_arrays_key):
        #     for i in range(waveform_data_info.Length(modes_arrays_key)):
        #         self.modes_selection.AddArray(waveform_data_info.Get(
        #             modes_arrays_key, i))
        for l in range(2, 5 + 1):
            for m in range(0, l + 1):
                self.modes_selection.AddArray(get_mode_name(l, m))
        self.modes_selection.AddObserver(
            "ModifiedEvent", das_util.create_modified_callback(self))
        self.polarizations_selection = vtkDataArraySelection()
        self.polarizations_selection.AddArray("Plus")
        self.polarizations_selection.AddArray("Cross")
        self.polarizations_selection.AddObserver(
            "ModifiedEvent", das_util.create_modified_callback(self))

    def FillInputPortInformation(self, port, info):
        # When using multiple inputs we (may) have to set their data types here
        # info.Set(self.INPUT_REQUIRED_DATA_TYPE(),
        #          'vtkTable' if port == 0 else 'vtkUniformGrid')
        info.Set(self.INPUT_REQUIRED_DATA_TYPE(), 'vtkTable')

    def _get_waveform_data(self):
        return dsa.WrapDataObject(self.GetInputDataObject(0, 0))

    # def _get_grid_data(self):
    #     return dsa.WrapDataObject(self.GetInputDataObject(1, 0))

    @smproperty.dataarrayselection(name="Modes")
    def GetModes(self):
        return self.modes_selection

    @smproperty.intvector(name="StoreIndividualModes", default_values=False)
    def SetStoreIndividualModes(self, value):
        self.store_individual_modes = value
        self.Modified()

    @smproperty.intvector(name="NormalizeEachMode", default_values=False)
    def SetNormalizeEachMode(self, value):
        self.normalize_each_mode = value
        self.Modified()

    @smproperty.dataarrayselection(name="Polarizations")
    def GetPolarizations(self):
        return self.polarizations_selection

    # Not needed when using SwshGrid input
    @smproperty.doublevector(name="Size", default_values=100)
    def SetSize(self, value):
        self.size = value
        self.Modified()

    # Not needed when using SwshGrid input
    @smproperty.intvector(name="SpatialResolution", default_values=100)
    def SetSpatialResolution(self, value):
        self.num_points_per_dim = value
        self.Modified()

    @smproperty.intvector(name="KeepEveryNthTimestep", default_values=1)
    def SetKeepEveryNthTimestep(self, value):
        self.keep_every_n_timestep = value
        self.Modified()

    # Not needed when using SwshGrid input
    @smproperty.intvector(name="EllMax", default_values=2)
    def SetEllMax(self, value):
        self.ell_max = value
        self.Modified()

    @smproperty.doublevector(name="RadialScale", default_values=10)
    def SetRadialScale(self, value):
        self.radial_scale = value
        self.Modified()

    @smproperty.intvector(name="ClipYNormal", default_values=False)
    @smdomain.xml('<BooleanDomain name="bool"/>')
    def SetClipYNormal(self, value):
        self.clip_y_normal = value
        self.Modified()

    @smproperty.intvector(name="ClipZNormal", default_values=False)
    @smdomain.xml('<BooleanDomain name="bool"/>')
    def SetClipZNormal(self, value):
        self.clip_z_normal = value
        self.Modified()

    @smproperty.intvector(name="OneOverRScaling", default_values=False)
    @smdomain.xml('<BooleanDomain name="bool"/>')
    def SetOneOverRScaling(self, value):
        self.add_one_over_r_scaling = value
        self.Modified()

    @smproperty.intvector(name="InvertRotationDirection", default_values=False)
    @smdomain.xml('<BooleanDomain name="bool"/>')
    def SetInvertRotationDirection(self, value):
        self.invert_rotation_direction = value
        self.Modified()

    @smproperty.doublevector(name="ActivationOffset", default_values=10)
    def SetActivationOffset(self, value):
        self.activation_offset = value
        self.Modified()

    @smproperty.doublevector(name="ActivationWidth", default_values=10)
    def SetActivationWidth(self, value):
        self.activation_width = value
        self.Modified()

    @smproperty.doublevector(name="DeactivationWidth", default_values=10)
    def SetDeactivationWidth(self, value):
        self.deactivation_width = value
        self.Modified()

    @smproperty.stringvector(name="SwshCacheDirectory", default_values="")
    def SetSwshCacheDirectory(self, value):
        self.swsh_cache_dir = value
        self.Modified()

    def _get_timesteps(self):
        logger.debug("Getting time range from data...")
        waveform_data = self._get_waveform_data()
        ts = waveform_data.RowData['Time']
        # Using a few timesteps within the data range so we can animate through
        # them in the GUI
        return np.linspace(ts[0], ts[-1], 100)

    @smproperty.doublevector(name="TimestepValues",
                             information_only="1",
                             si_class="vtkSITimeStepsProperty")
    def GetTimestepValues(self):
        return self._get_timesteps().tolist()

    def RequestInformation(self, request, inInfo, outInfo):
        logger.debug("Requesting information...")
        waveform_data_info = inInfo[0].GetInformationObject(0)
        # Careful with printing these information objects, their stream operator
        # may randomly crash...
        # logger.debug("Waveform data info: {}".format(waveform_data_info))
        # grid_info = inInfo[1].GetInformationObject(0)
        info = outInfo.GetInformationObject(0)

        # For the `vtkUniformGrid` output we need to provide extents
        # so that it gets rendered at all.
        # When using the SwshGrid input we can retrieve them from the
        # information object and pass them on.
        # grid_extents = grid_info.Get(self.GetExecutive().WHOLE_EXTENT())
        N = self.num_points_per_dim
        N_y = N // 2 if self.clip_y_normal else N
        N_z = N // 2 if self.clip_z_normal else N
        grid_extents = [0, N - 1, 0, N_y - 1, 0, N_z - 1]
        util.SetOutputWholeExtent(self, grid_extents)

        # This needs the time data from the waveform file, so we may have to
        # set the `TIME_RANGE` and `TIME_STEPS` already in the
        # WaveformDataReader.
        timesteps_util.set_timesteps(self,
                                     self._get_timesteps(),
                                     logger=logger)

        # logger.debug("Information object: {}".format(info))
        return 1

    def RequestData(self, request, inInfo, outInfo):
        logger.debug("Requesting data...")
        waveform_data = self._get_waveform_data()
        # grid_data = self._get_grid_data()
        output = dsa.WrapDataObject(vtkUniformGrid.GetData(outInfo))

        t = timesteps_util.get_timestep(self, logger=logger)
        N = self.num_points_per_dim
        D = self.size

        # We may have to forward the grid data here when using SwshGrid input
        # output.SetDimensions(*grid_data.GetDimensions())
        # output.SetOrigin(*grid_data.GetOrigin())
        # output.SetSpacing(*grid_data.GetSpacing())
        dx = 2. * D / N
        N_y = N // 2 if self.clip_y_normal else N
        N_z = N // 2 if self.clip_z_normal else N
        output.SetDimensions(N, N_y, N_z)
        output.SetOrigin(-D, -D, -D)
        output.SetSpacing(dx, dx, dx)

        # Compute the SWSHs on the grid
        # This section can be deleted when using SwshGrid input
        spin_weight = -2
        ell_max = self.ell_max

        # We skip this section of the code if non-simulation data is handled
        # since it becomes simply unnecessary.
        if type(waveform_data.RowData['Y_l2_m2'][5]) is not dsa.VTKNoneArray:
            swsh_grid, r = cached_swsh_grid(
                size=D,
                num_points=N,
                spin_weight=spin_weight,
                ell_max=ell_max,
                radial_scale=self.radial_scale,
                clip_y_normal=self.clip_y_normal,
                clip_z_normal=self.clip_z_normal,
                activation_offset=self.activation_offset,
                activation_width=self.activation_width,
                deactivation_width=self.deactivation_width,
                add_one_over_r_scaling=self.add_one_over_r_scaling,
                cache_dir=self.swsh_cache_dir)

            logger.info("Computing volume data at t={}...".format(t))
        start_time = time.time()

        # We skip this section of the code if non-simulation data is handled
        # since it becomes simply unnecessary.
        if type(waveform_data.RowData['Y_l2_m2'][5]) is not dsa.VTKNoneArray:

            # Compute scaled waveform phase on the grid
            # r = vtknp.vtk_to_numpy(grid_data.GetPointData()['RadialCoordinate'])
            phase = t - r

            # Invert rotation direction
            rotation_direction = -1. if self.invert_rotation_direction else 1.

            # Compute strain in the volume from the input waveform data
            skip_timesteps = self.keep_every_n_timestep
            waveform_timesteps = waveform_data.RowData['Time'][::skip_timesteps]
            strain = np.zeros(len(r), dtype=np.complex)
            # Optimization for when the waveform is sampled uniformly
            # TODO: Cache this

        # Here the code checks if non-simulation data was passed and visualizes the
        # analytical data loaded erlier

        if type(waveform_data.RowData['Y_l2_m2'][5]) is dsa.VTKNoneArray:
            strain = np.zeros(1000000, dtype=np.complex)
            # generate the index of the current timestep...
            indexx = list(map(abs, list(waveform_data.RowData['Time'] - t + 1000))).index(
                    np.min(list(map(abs, list(waveform_data.RowData['Time'] - t + 1000)))))
            # ...and pass the associated column of the strains grid to teh starin value
            # which is then visualized at the bottom
            strain += 100*strains[indexx]

        else:

            dt = np.diff(waveform_timesteps)
            waveform_uniformly_sampled = np.allclose(dt, dt[0])
            global has_shown_warning_nonuniformly_sampled
            if waveform_uniformly_sampled:
                dt = dt[0]
                logger.debug("Waveform sampled uniformly with dt={:.2e}, using optimized interpolation:".format(dt))
                waveform_start_time = waveform_timesteps[0]
                waveform_start_index = min(len(waveform_timesteps) - 2, max(
                    0, int(np.floor((np.min(phase) - waveform_start_time) / dt))))
                waveform_stop_index = max(waveform_start_index + 1, min(
                    len(waveform_timesteps),
                    int(np.ceil((np.max(phase) - waveform_start_time) / dt))))
                if waveform_stop_index == len(waveform_timesteps):
                    waveform_stop_index = -1
                logger.debug(
                    "Restricting interpolation to waveform indices {}, that's between waveform times {}. We will interpolate to times between {} (should be contained in restricted waveform range except for boundary effects)."
                    .format((waveform_start_index, waveform_stop_index),
                            (waveform_timesteps[waveform_start_index],
                             waveform_timesteps[waveform_stop_index]),
                            (np.min(phase), np.max(phase))))
                waveform_timesteps = waveform_timesteps[waveform_start_index:waveform_stop_index]
            elif not has_shown_warning_nonuniformly_sampled:
                logger.warning("Waveform is not sampled uniformly so interpolation is slightly more expensive.")
                has_shown_warning_nonuniformly_sampled = True
            # for i in range(self.modes_selection.GetNumberOfArrays()):
            #     mode_name = self.modes_selection.GetArrayName(i)
            for l in range(abs(spin_weight), ell_max + 1):
                for abs_m in range(0, l + 1):
                    mode_name = get_mode_name(l, abs_m)
                    strain_mode = np.zeros(len(r), dtype=np.complex)
                    if not self.modes_selection.ArrayIsEnabled(mode_name):
                        continue
                    for sign_m in (-1, 1):
                        m = abs_m * sign_m
                        dataset_name = "Y_l{}_m{}".format(l, m)
                        mode_profile = swsh_grid[:, LM_index(l, m, 0)]
                        # mode_profile = vtknp.vtk_to_numpy(grid_data.GetPointData()[dataset_name])
                        waveform_mode_data = waveform_data.RowData[dataset_name][::skip_timesteps]
                        if isinstance(waveform_mode_data, dsa.VTKNoneArray):
                            logger.warning(
                                "Dataset '{}' for mode {} not available in waveform data, skipping."
                                .format(dataset_name, (l, m)))
                            continue
                        # TODO: Make sure inverting the rotation direction like this
                        # is correct.
                        waveform_mode_data = waveform_mode_data[:, 0] + rotation_direction * 1j * waveform_mode_data[:, 1]
                        if self.normalize_each_mode:
                            waveform_mode_data /= np.max(np.abs(waveform_mode_data))
                        if waveform_uniformly_sampled:
                            waveform_mode_data = waveform_mode_data[waveform_start_index:waveform_stop_index]
                        mode_data = np.interp(phase,
                                              waveform_timesteps,
                                              waveform_mode_data,
                                              left=0.,
                                              right=0.)
                        strain_mode += mode_data * mode_profile
                    strain += strain_mode
                    # Expose individual modes in output
                    if self.store_individual_modes:
                        if self.polarizations_selection.ArrayIsEnabled("Plus"):
                            strain_mode_real_vtk = vtknp.numpy_to_vtk(
                                np.real(strain_mode), deep=True)
                            strain_mode_real_vtk.SetName(mode_name + ' Plus')
                            output.GetPointData().AddArray(strain_mode_real_vtk)
                        if self.polarizations_selection.ArrayIsEnabled("Cross"):
                            strain_mode_imag_vtk = vtknp.numpy_to_vtk(
                                np.imag(strain_mode), deep=True)
                            strain_mode_imag_vtk.SetName(mode_name + ' Cross')
                            output.GetPointData().AddArray(strain_mode_imag_vtk)
        if self.polarizations_selection.ArrayIsEnabled("Plus"):
            strain_real_vtk = vtknp.numpy_to_vtk(np.real(strain), deep=True)
            strain_real_vtk.SetName('Plus strain')
            output.GetPointData().AddArray(strain_real_vtk)
        if self.polarizations_selection.ArrayIsEnabled("Cross"):
            strain_imag_vtk = vtknp.numpy_to_vtk(np.imag(strain), deep=True)
            strain_imag_vtk.SetName('Cross strain')
            output.GetPointData().AddArray(strain_imag_vtk)

        logger.info("Volume data computed in {:.3f}s.".format(time.time() -
                                                              start_time))
        return 1


# all lifetimes in years

import numpy as np

CO2         = np.nan # dummy - to preserve indexing order
CH4         = 9.3   # for constant lifetime runs
N2O         = 121.
CF4         = 50000.
C2F6        = 10000.
C3F8        = 2600.
C4F10       = 2600.
C5F12       = 4100.
C6F14       = 3100.
C7F16       = 3000.
C8F18       = 3000.
C_C4F8      = 3200.
HFC23       = 222.
HFC32       = 5.2
HFC43_10    = 16.1
HFC43_10MEE = HFC43_10
HFC125      = 28.2
HFC134A     = 13.4
HFC143A     = 47.1
HFC152A     = 1.5
HFC227EA    = 38.9
HFC236FA    = 242.
HFC245FA    = 7.7
HFC365MFC   = 8.7
SF6         = 3200.
NF3         = 500.
SO2F2       = 36.
CFC11       = 45.
CFC12       = 100.
CFC113      = 85.
CFC114      = 190.
CFC115      = 1020.
CARB_TET    = 26.
CCL4        = CARB_TET
MCF         = 5.
CH3CCL3     = MCF
HCFC22      = 11.9
HCFC141B    = 9.2
HCFC142B    = 17.2
HALON1211   = 16.
HALON1202   = 2.9
HALON1301   = 65.
HALON2402   = 20.
CH3BR       = 0.8
CH3CL       = 1.
CH2CL2      = 0.3945  # from Hodnebrog et al., 2013
CHCL3       = 0.4082  # from Hodnebrog et al., 2013

# This is the list of gases included in the RCPs/AR5/CMIP5.
aslist   = [CO2, CH4, N2O, CF4, C2F6, C6F14, HFC23, HFC32, HFC43_10, HFC125,
            HFC134A, HFC143A, HFC227EA, HFC245FA, SF6, CFC11, CFC12, CFC113,
            CFC114, CFC115, CARB_TET, MCF, HCFC22, HCFC141B, HCFC142B,
            HALON1211, HALON1202, HALON1301, HALON2402, CH3BR, CH3CL]


# This script computes the relative hydration free energy with a single topology protocol.
# It checks to make sure the estimated dG is in agreement between the following three methods:
# 1) Two seperate absolute free energy calculations
# 2) Relative free energy with full atom-mapping
# 3) Relative free energy with partial atom-mapping (4D-decoupling)

import os
import argparse
import numpy as np

from rdkit import Chem
from rdkit.Chem import rdFMCS
from rdkit.Chem import AllChem

from fe import topology
from md import builders
from md import minimizer

import functools

from ff import Forcefield
from ff.handlers.deserialize import deserialize_handlers

import multiprocessing

from fe import free_energy

def wrap_method(args, fn):
    gpu_idx = args[0]
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_idx)
    return fn(*args[1:])

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Relative Hydration Free Energy Consistency Testing",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--num_gpus",
        type=int,
        help="number of gpus",
        required=True
    )

    parser.add_argument(
        "--num_vacuum_windows",
        type=int,
        help="number of vacuum lambda windows",
        required=True
    )

    parser.add_argument(
        "--num_solvent_windows",
        type=int,
        help="number of solvent lambda windows",
        required=True
    )

    parser.add_argument(
        "--num_equil_steps",
        type=int,
        help="number of equilibration steps for each lambda window",
        required=True
    )

    parser.add_argument(
        "--num_prod_steps",
        type=int,
        help="number of production steps for each lambda window",
        required=True
    )

    parser.add_argument(
        "--num_absolute_windows",
        type=int,
        help="number of absolute lambda windows",
        required=True
    )

    cmd_args = parser.parse_args()

    multiprocessing.set_start_method('spawn') # CUDA runtime is not forkable
    pool = multiprocessing.Pool(cmd_args.num_gpus)

    suppl = Chem.SDMolSupplier('tests/data/benzene_fluorinated.sdf', removeHs=False)
    all_mols = [x for x in suppl]
    mol_a = all_mols[0]
    mol_b = all_mols[1]

    ff_handlers = deserialize_handlers(open('ff/params/smirnoff_1_1_0_ccc.py').read())
    ff = Forcefield(ff_handlers)

    # the water system first.
    solvent_system, solvent_coords, solvent_box, omm_topology = builders.build_water_system(4.0)
    solvent_box += np.eye(3)*0.1 # BFGS this later

    print("Minimizing the host structure to remove clashes.")
    minimized_solvent_coords = minimizer.minimize_host_4d([mol_a], solvent_system, solvent_coords, ff, solvent_box)

    absolute_lambda_schedule = np.concatenate([
        np.linspace(0.0, 0.333, cmd_args.num_absolute_windows - cmd_args.num_absolute_windows//3, endpoint=False),
        np.linspace(0.333, 1.0, cmd_args.num_absolute_windows//3),
    ])

    abs_dGs = []

    for idx, mol in enumerate([mol_a, mol_b]):

        afe = free_energy.AbsoluteFreeEnergy(mol, ff)
        absolute_args = []

        for lambda_idx, lamb in enumerate(absolute_lambda_schedule):
            gpu_idx = lambda_idx % cmd_args.num_gpus
            absolute_args.append((gpu_idx, lamb, solvent_system, minimized_solvent_coords, solvent_box, cmd_args.num_equil_steps, cmd_args.num_prod_steps))

        results = pool.map(functools.partial(wrap_method, fn=afe.host_edge), absolute_args, chunksize=1)

        for lamb, (bonded_du_dl, nonbonded_du_dl) in zip(absolute_lambda_schedule, results):
            print("final absolute", idx, "lambda", lamb, "bonded:", bonded_du_dl[0], bonded_du_dl[1], "nonbonded:", nonbonded_du_dl[0], nonbonded_du_dl[1])

        dG = np.trapz([x[0][0]+x[1][0] for x in results], absolute_lambda_schedule)
        print("mol", idx, "dG absolute:", dG)
        abs_dGs.append(dG)

    print("Absolute Difference", abs_dGs[0] - abs_dGs[1])

    # relative free energy, compare two different core approaches

    core_full = np.stack([
        np.arange(mol_a.GetNumAtoms()),
        np.arange(mol_b.GetNumAtoms())
    ], axis=1)

    core_part = np.stack([
        np.arange(mol_a.GetNumAtoms() - 1),
        np.arange(mol_b.GetNumAtoms() - 1)
    ], axis=1)

    for core_idx, core in enumerate([core_full, core_part]):
        single_topology = topology.SingleTopology(mol_a, mol_b, core, ff)

        rfe = free_energy.RelativeFreeEnergy(single_topology)

        vacuum_lambda_schedule = np.linspace(0.0, 1.0, cmd_args.num_vacuum_windows)
        solvent_lambda_schedule = np.linspace(0.0, 1.0, cmd_args.num_solvent_windows)

        # vacuum leg
        vacuum_args = []
        for lambda_idx, lamb in enumerate(vacuum_lambda_schedule):
            gpu_idx = lambda_idx % cmd_args.num_gpus
            vacuum_args.append((gpu_idx, lamb, cmd_args.num_equil_steps, cmd_args.num_prod_steps))

        results = pool.map(functools.partial(wrap_method, fn=rfe.vacuum_edge), vacuum_args, chunksize=1)

        # TODO: update this to reflect new return type of rfe.vacuum_edge
        for lamb, (bonded_du_dl, nonbonded_du_dl) in zip(vacuum_lambda_schedule, results):
            print("final vacuum lambda", lamb, "bonded:", bonded_du_dl[0], bonded_du_dl[1], "nonbonded:", nonbonded_du_dl[0], nonbonded_du_dl[1])

        dG_vacuum = np.trapz([x[0][0]+x[1][0] for x in results], vacuum_lambda_schedule)
        print("dG vacuum:", dG_vacuum)

        # solvent leg
        solvent_args = []
        for lambda_idx, lamb in enumerate(solvent_lambda_schedule):
            gpu_idx = lambda_idx % cmd_args.num_gpus
            solvent_args.append((gpu_idx, lamb, solvent_system, minimized_solvent_coords, solvent_box, cmd_args.num_equil_steps, cmd_args.num_prod_steps))

        results = pool.map(functools.partial(wrap_method, fn=rfe.host_edge), solvent_args, chunksize=1)
        # TODO: update this to reflect new return type of rfe.vacuum_edge

        for lamb, (bonded_du_dl, nonbonded_du_dl) in zip(solvent_lambda_schedule, results):
            print("final solvent lambda", lamb, "bonded:", bonded_du_dl[0], bonded_du_dl[1], "nonbonded:", nonbonded_du_dl[0], nonbonded_du_dl[1])

        dG_solvent = np.trapz([x[0][0]+x[1][0] for x in results], solvent_lambda_schedule)
        print("dG solvent:", dG_solvent)

        print("Core map", core_idx, "Difference", dG_solvent - dG_vacuum)


#!/usr/bin/env python
#-*- coding:utf-8 -*-
"""
    Author:             Benjamin Rumeau (intern at CLS in 2014 under supervision
                        of Philippe Gaspard and Anne-Cecile Dragon.
    Module description: Library containing all the geographical classes and fun-
                        ctions.
"""
import numpy as np

# ||==========================================
# ||                                        ||
# ||           Objets géographiques         ||
# ||             et biologiques             ||
# ||                                        ||
# ==========================================||

class NestingPeriod :
    """
    class nesting_period used as a structure
    contains the begin, end and peak of the nesting season
    """
    def __init__(self,first=0.,last=0.,peak=0.,incubation_time=60):
        self.begin = np.float32(first)
        self.end = np.float32(last)
        self.peak = np.float32(peak)
        self.incubation_time=incubation_time

    def incubation(self):
        """
        Corrects nesting period to take into account incubation time.
        Incubation time must be expressed in days.
        """
        print("\n=> correcting nesting periods with incubation times")
        self.begin+=self.incubation_time
        self.end+=self.incubation_time
        self.peak+=self.incubation_time

class Beach :
    """
    class beach is used as a structure
    contains all information needed about the nesting beach
    => launching area
    => nesting season
    => geographic coordinates of the beach
    => name of the beach
    => eventually the reference of those information
    """
    def __init__(self,
                 nesting_period=NestingPeriod(),
                 north_pt=[1,1],
                 south_pt=[0,0],
                 orientation="east",
                 d_min=30,
                 d_max=50,
                 ref="no reference",
                 beach_name="no name",
                 file_name="no name"):
        self.nesting_period=nesting_period
        self.north_pt=north_pt
        self.south_pt=south_pt
        if orientation=="east":
            self.orientation=1
        elif orientation=="west":
            self.orientation=-1
        else:
            print("Error: orientation must be east or west")
        self.d_min=d_min
        self.d_max=d_max
        self.ref=str(ref)
        self.beach_name=str(beach_name)
        self.nesting_period.incubation()

class Geolocation:
    """
    class position_4D contains the a 4 uplet x,y,z,t giving the time ans space
    localisation of a particle, as well as an information about the coordinate
    system.
    """
    def __init__(self,x=0.,y=0.,t=0.,z=-1.0,coord_sys="geo"):
        self.x=np.float32(x)
        self.y=np.float32(y)
        self.z=np.float32(z)
        self.t=np.float32(t)
        self.coord_sys=coord_sys

    def geo_to_grid(self,lon_mat,lat_mat):
        """
        converts (lon,lat) geographical coordinate to grid coordinates (conti-
        nuous conversion as coordinates of the closest neighbour are computed
        in Ariane.
        - lon,lat are the geographical coordinate of the point.
        - lon_mat, lat are matrices containing the longitudes (resp. latitudes) of the centers of the cells of the Arakawa C grid.
        NB : this function is only appropriate for regular, square grids.
        """
        if self.coord_sys=="grid":
            print("\n Warning: already in grid coordinates")
        else:
            # Read extreme coordinates
            lon_max=np.max(lon_mat)
            lon_min=np.min(lon_mat)
            lat_max=np.max(lat_mat)
            lat_min=np.min(lat_mat)
            # Converts to grid coordinate
            if ((self.x> lon_max) or (self.x< lon_min) or (self.y>lat_max) or (self.y<lat_min)):
                print("\n Longitude or latitude out of range, must be between ", lon_min, " & ", lon_max, " for longitude and ", lat_min, " & ", lat_max, " for latitude \n")
            else:
                h_x=(lon_max-lon_min)/len(lon_mat)
                h_y=(lat_max-lat_min)/len(lat_mat)
                self.x=(self.x-lon_min)/h_x
                self.y=(self.y-lat_min)/h_y
                self.coord_sys="grid"

    def grid_to_geo(self, lon_mat, lat_mat) :
        if self.coord_sys=="grid":
            print("\n Warning: already in geographical coordinates")
        # Read extreme coordinates
        lon_max=np.max(lon_mat)
        lon_min=np.min(lon_mat)
        lat_max=np.max(lat_mat)
        lat_min=np.min(lat_mat)
        # Converts to grid coordinate
        h_x=(lon_max-lon_min)/len(lon_mat)
        h_y=(lat_max-lat_min)/len(lat_mat)
        self.x=h_x * self.x + lon_min
        self.y=h_y * self.y + lat_min
        self.coord_sys="grid"

    def write_line(self,f):
        """
        write as a proper line readable for Ariane the position in a given file
        """
        line=' '+"%.3f" % self.x+' '+\
             "%.3f" % self.y+'    '+\
             "%.1f" % self.z+'  '+\
             "%.2f" % self.t+'     '+\
             '1.0\n'

        line=line.replace(',','.')
        f.write(line)

# ||==========================================
# ||                                        ||
# ||       Transformations géométriques     ||
# ||                                        ||
# ==========================================||

def rotation(x,y,Ox,Oy,theta):
    """
        Computes coordinates X,Y in referential R' from a poin (x,y) in a referential R
        where R' is obtained by a rotation of R of angle theta and center O(Ox,Oy)
    """
    X=Ox+(x-Ox)*np.cos(theta)+(y-Oy)*np.sin(theta)
    Y=Oy-(x-Ox)*np.sin(theta)+(y-Oy)*np.cos(theta)
    return(X,Y)

def geo_to_grid(lon,lat,lon_mat,lat_mat):
    """
    converts (lon,lat) geographical coordinate to grid coordinates (conti-
    nuous conversion as coordinates of the closest neighbour are computed
    in Ariane.
    - lon,lat are the geographical coordinate of the point.
    - lon_mat, lat are matrices containing the longitudes (resp. latitudes) of the centers of the cells of the Arakawa C grid.
    NB : this function is only appropriate for regular, square grids.
    """
    # Read extreme coordinates
    lon_max=np.max(lon_mat)
    lon_min=np.min(lon_mat)
    lat_max=np.max(lat_mat)
    lat_min=np.min(lat_mat)
    # Converts to grid coordinate
    if ((lon> lon_max) or (lon< lon_min) or (lat>lat_max) or (lat<lat_min)):
        print("\n Longitude or latitude out of range, must be between ", lon_min, " & ", lon_max, " for longitude and ", lat_min, " & ", lat_max, " for latitude \n")
    else:
        h_x=(lon_max-lon_min)/len(lon_mat)
        h_y=(lat_max-lat_min)/len(lat_mat)
        i1 = int((lon-lon_min)/h_x)
        i2 = i1 + 1
        I = [i1,i2]
        i0 = np.argmin([abs(lon - lon_mat[i1]), abs(lon - lon_mat[i2])])
        i = I[i0]
        j1 = int((lat-lat_min)/h_y)
        j2 = j1 + 1
        J = [j1,j2]
        j0 = np.argmin([abs(lat - lat_mat[j1]), abs(lat - lat_mat[j2])])
        j = J[j0]
    return i,j