from __future__ import print_function
import os
import csv
import numpy as np
from scipy.signal import decimate
import soundfile as sf
import librosa
from keras.models import model_from_json

import argparse

__author__ = 'Jakob Abesser'
__copyright__ = 'J. Abesser, S. Balke, K. Frieler, M. Pfleiderer, M. Mueller, 2017'

class WalkingBassTranscription:
    """ Algorithm for walking bass transcription in jazz ensemble recordings
         [1] J. Abesser, S. Balke, K. Frieler, M. Pfleiderer, M. Mueller: Deep Learning for Jazz Walking Bass Transcription, AES conference
             on Semantic Audio, Erlangen, Germany, 2017
        Examples can be found here:
         [2] http://www. audiolabs- erlangen.de/resources/MIR/ 2017-AES-WalkingBassTranscription/
        This algorithm was used to create the beat-wise bass pitch values included in the Weimar Jazz Database
         [3] http://jazzomat.hfm-weimar.de/dbformat/dboverview.html
    """

    def __init__(self,
                 hopsize=1024,
                 blocksize=2048,
                 pitch_range=(28, 67),  # E1 - G4
                 bins_per_octave=12):
        """ Initialize transcriber
        Args:
            hopsize (int): Hopsize in samples
            blocksize (int): Blocksize in samples
            pitch_range (tuple of int): Lower and upper pitch range
            bins_per_octave (int): Frequency axis resolution (number of bins per octave)
        """
        self.hopsize = hopsize
        self.blocksize = blocksize
        self.pitch_range = pitch_range
        self.bins_per_octave = bins_per_octave

        # generate logarithmically spaced frequency axis
        delta_midi = 12. / self.bins_per_octave
        self.f_axis_midi = np.arange(self.pitch_range[0],
                                     self.pitch_range[1] + delta_midi,
                                     delta_midi, dtype=int)
        tuning_freq_hz = 440.
        self.f_axis_hz = tuning_freq_hz * 2 ** ((self.f_axis_midi - 69.) / 12.)

        # load DNN model
        self.model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'data')
        self.model = None
        self._load_model()

    def _load_model(self):
        """ Load DNN model by loading architecture and weights and initialize model accordingly
        """
        fn_model_architecture = os.path.join(self.model_path, 'model.yaml')
        fn_model_weights = os.path.join(self.model_path, 'weights.h5')
        with open(fn_model_architecture, 'r') as f:
            model_json = f.read()
        self.model = model_from_json(model_json)
        self.model.load_weights(fn_model_weights)

    def transcribe(self,
                   fn_wav,
                   dir_out,
                   beat_times=None,
                   tuning_frequency_hz=440.,
                   threshold=0.2,
                   aggregation='beat'):
        """ Transcribe audio file
        Args:
            fn_wav (string): WAV file name
            dir_out (string): Directory to store results (if None, same directory as fn_wav is used)
            beat_times (ndarray): Beat times in seconds (if None, only bass saliency is extracted)
            tuning_frequency_hz (float): Tuning frequency (Hz)
            threshold (float): Decision treshold
            aggregation (string): Aggregation method. Possible values are 'beat' (beat-wise aggregation)"
                                  "and 'flex-q' (dynamic estimation of most likely tatum per beat)
        Returns:
            pitch_saliency (2d ndarray): Bass pitch saliency (num_pitches x num_frames)
            midi_axis (ndarray): MIDI pitch values of pitch axis
            time_axis_sec (ndarray): Time frames [s]
        """
        tuning_dev_in_semitones = np.log2(tuning_frequency_hz/440.)*12

        # load audio file
        x, fs = sf.read(fn_wav)

        # convert to mono
        if x.ndim == 2:
            x = np.mean(x, axis=1)

        # signal decimation
        if fs != 22050.:
            decimation_factor = int(np.round(fs / 22050.))
            x = decimate(x, decimation_factor, zero_phase=True)
            fs /= decimation_factor

        # compute contant-Q spectrogram
        mag_spec = np.abs(librosa.cqt(x,
                                      sr=fs,
                                      hop_length=self.hopsize,
                                      fmin=440*2**((self.pitch_range[0]-69)/12.),
                                      n_bins=self.pitch_range[1] - self.pitch_range[0] + 1,
                                      bins_per_octave=self.bins_per_octave,
                                      tuning=tuning_dev_in_semitones))

        num_frames = mag_spec.shape[1]
        time_axis_sec = (np.arange(num_frames) + .5) * self.hopsize / fs

        # frame stacking
        features = frame_stacking(mag_spec.T, 2)

        # feature normalization
        features = normalize_euclidean(features)

        # model prediction
        pitch_saliency = self.model.predict(features)

        # save saliency matrix
        base_name = os.path.basename(fn_wav).replace('.wav', '')
        np.save(
            os.path.join(dir_out, '{}_bass_pitch_saliency.npy'.format(base_name)),
            pitch_saliency
        )

        # export most salient bass track as CSV file
        with open(os.path.join(dir_out, '{}_bass_f0.csv'.format(base_name)), 'w') as fhandle:
            writer = csv.writer(fhandle, delimiter=',')
            for t in range(pitch_saliency.shape[0]):
                i = np.argmax(pitch_saliency[t])
                time_val = time_axis_sec[t]
                if pitch_saliency[t, i] >= threshold:
                    freq_val = self.f_axis_hz[i]
                else:
                    freq_val = -1*self.f_axis_hz[i]

                writer.writerow([time_val, freq_val])

        # aggregate pitch saliency to note events
        if beat_times is not None:

            onset_sec, offset_sec, pitch = aggregate_saliency_to_notes(pitch_saliency,
                                                                       self.f_axis_midi,
                                                                       time_axis_sec,
                                                                       beat_times,
                                                                       method=aggregation,
                                                                       threshold=threshold)

            # export score to be imported to Sonic Visualiser as note layer
            score_mat = np.vstack((onset_sec, offset_sec, pitch)).T
            np.savetxt(os.path.join(dir_out, '{}_bass_line.csv'.format(base_name)),
                       score_mat,
                       fmt='%4.4f,%4.4f,%d')

        return pitch_saliency, self.f_axis_midi, time_axis_sec

def aggregate_saliency_to_notes(pitch_saliency,
                                freq_bins_midi,
                                frame_times_sec,
                                beat_times_sec,
                                method='beat',
                                num_tatums_per_beat=None,
                                threshold=0.2):
    """ Aggregate frame-wise pitch saliency values to note events based on given beat times
    Args:
        pitch_saliency (2d np.ndarray): Frame-wise pitch saliency (num_frames x num_pitches)
        freq_bins_midi (np.ndarray): MIDI pitch values (num_pitches)
        frame_times_sec (np.ndarray): Frame times in seconds (num_pitches)
        beat_times_sec (np.ndarray): Beat times in seconds (num_beats)
        method (string): Aggregation method
        aggregation (string): Aggregation method. Possible values are 'beat' (beat-wise aggregation)"
                              "and 'flex-q' (dynamic estimation of most likely tatum per beat)
        num_tatums_per_beat (tuple): Number of tatums per beats - different beat subdivisions that are tested if
                                     aggregation == 'flex-q' (default: (1, 2, 3))
        threshold (float): Minimum saliency threshold to detect notes
    Returns:
        onset (np.ndarray): Note-wise onset times in seconds
        offset (np.ndarray): Note-wise offset times in seconds
        pitch (np.ndarray): Note-wise pitch values
    """
    assert method in ('beat', 'flex-q'), "Non-valid value for method!"

    if num_tatums_per_beat is None:
        num_tatums_per_beat = (1, 3)
    num_subdivisions = len(num_tatums_per_beat)

    num_beats = len(beat_times_sec)
    pitch = []
    onset = []
    offset = []

    # map beat times from seconds to frames
    beat_frames = [closest_bin(frame_times_sec, _) for _ in beat_times_sec]

    # iterate over beats
    for b in range(num_beats-1):

        if method == 'beat':
            # take most likely pitch
            beat_pitch_saliency = np.mean(
                pitch_saliency[beat_frames[b]: beat_frames[b + 1], :], axis=0
            )

            # if saliency exceeds threshold > store note
            if np.max(beat_pitch_saliency) > threshold:
                best_pitch_idx = np.argmax(beat_pitch_saliency)
                onset.append(beat_times_sec[b])
                offset.append(beat_times_sec[b+1])
                pitch.append(freq_bins_midi[best_pitch_idx])

        elif method == 'flex-q':

            curr_onset = []
            curr_offset = []
            curr_pitch = []

            scores = np.zeros(num_subdivisions)
            # try different tatum subdivisions
            for s, sub_div in enumerate(num_tatums_per_beat):

                scores[s], _, _ = get_sub_beat_saliency(pitch_saliency,
                                                        beat_times_sec[b],
                                                        beat_times_sec[b + 1],
                                                        frame_times_sec,
                                                        sub_div)

            # get optimal subdivision from highest score
            sub_div_opt = num_tatums_per_beat[np.argmax(scores)]

            # use optimal subdivision to extract note events
            _, sub_beat_saliency, sub_beat_times_sec = get_sub_beat_saliency(pitch_saliency,
                                                                             beat_times_sec[b],
                                                                             beat_times_sec[b + 1],
                                                                             frame_times_sec,
                                                                             sub_div_opt)

            pitch_idx_opt = np.argmax(sub_beat_saliency, axis=1)
            saliency_opt = np.max(sub_beat_saliency, axis=1)

            # check in which tatum segments, the saliency exceeds the threshold
            is_valid = saliency_opt >= threshold

            for n in range(sub_div_opt):
                curr_onset.append(sub_beat_times_sec[n])
                curr_offset.append(sub_beat_times_sec[n+1])
                curr_pitch.append(freq_bins_midi[pitch_idx_opt[n]])

            curr_onset = np.array(curr_onset)
            curr_offset = np.array(curr_offset)
            curr_pitch = np.array(curr_pitch)

            if np.all(is_valid):
                # merge adjacent notes with same pitch as we can't do onset detection solely based on saliency
                if len(np.unique(curr_pitch)) == 1:
                    curr_pitch = np.array((curr_pitch[0],))
                    curr_onset = np.array((curr_onset[0],))
                    curr_offset = np.array((curr_offset[-1],))
            else:
                curr_onset = curr_onset[is_valid]
                curr_offset = curr_offset[is_valid]
                curr_pitch = curr_pitch[is_valid]

            onset.append(curr_onset)
            offset.append(curr_offset)
            pitch.append(curr_pitch)

    if method == 'beat':
        onset = np.array(onset)
        offset = np.array(offset)
        pitch = np.array(pitch)
    elif method == 'flex-q':
        onset = np.concatenate(onset)
        offset = np.concatenate(offset)
        pitch = np.concatenate(pitch).astype(int)

    return onset, offset, pitch

def get_sub_beat_saliency(pitch_saliency,
                          start_time_sec,
                          end_time_sec,
                          frame_times_sec,
                          sub_div):
    """ Get pitch saliency and likelihood score for subdivision of given segment in pitch saliency matrix
    Args:
        pitch_saliency (2d np.ndarray): Frame-wise pitch saliency (num_frames x num_pitches)
        start_time_sec (float): Start time in seconds
        end_time_sec (float): End time in seconds
        frame_times_sec (np.ndarray): Frame times in seconds
        sub_div (int, >= 1): Number of subdivisions (e.g. 2 -> given segment is divided into 2 subsegments of equal
                             duration)
    Returns:
        score (float): Likelihood score for current subdivision based on difference between highest and second
                       highest saliency value
        sub_beat_saliency (2d np.ndarray): Average pitch saliency vectors for each sub beat (tatum level) (num_sub_beats x num_pitches)
        sub_beat_times_sec (np.ndarray): Boundary times in seconds for subbeats (num_sub_beats + 1)
    """
    num_pitches = pitch_saliency.shape[1]
    beat_len_sec = end_time_sec - start_time_sec
    sub_beat_len_sec = beat_len_sec / sub_div
    # sub_beat times
    sub_beat_times_sec = np.arange(sub_div + 1) * sub_beat_len_sec + start_time_sec
    assert len(sub_beat_times_sec) == sub_div + 1
    # seconds to frames
    sub_beat_times_frames = [closest_bin(frame_times_sec, _) for _ in sub_beat_times_sec]
    sub_beat_saliency = np.zeros((sub_div, num_pitches))
    for sb in range(sub_div):
        sub_beat_saliency[sb, :] = get_segment_saliency(pitch_saliency,
                                                        sub_beat_times_frames[sb],
                                                        sub_beat_times_frames[sb + 1])

    # sort pitch saliency values in descending order accross pitches
    sub_beat_saliency_sorted = -np.sort(-sub_beat_saliency, axis=1)
    score = np.mean(sub_beat_saliency_sorted[:, 0] - sub_beat_saliency_sorted[:, 1])

    return score, sub_beat_saliency, sub_beat_times_sec

def get_segment_saliency(pitch_saliency, start_frame, end_frame):
    """ Get average saliency over segment
    Args:
        pitch_saliency (2d np.ndarray): Frame-wise pitch saliency (num_frames x num_pitches)
        start_frame (int): Start frame
        end_frame (int): End frame
    Returns:
        segment_pitch_saliency (np.ndarray): Averaged pitch saliency over segment (num_pitches)
    """
    return np.mean(pitch_saliency[start_frame:end_frame, :], axis=0)

def closest_bin(axis, val):
    """ Return closest bin in axis to value
    Args:
        axis (ndarray): Axis
        val (float / int): Value
    Return:
        idx (int): Closest index in axis to val
    """
    return np.argmin(np.abs(axis-val))

def frame_stacking(x, context_size):
    """ Stack frames to incorporate some temporal context
    Args:
        x (2d ndarray): Feature matrix (num_frames x num_features)
        context_size (int): Context size for frame stacking (e.g. context size of 2 means that we stack 5 frames)
    Return
        x_stacked (2d ndarray): Stacked feature matrix (num_frames x ((context_size*2 + 1)*num_features)
    """
    num_frames = x.shape[0]
    c = 2 * context_size + 1
    assert c < num_frames, "Context size is too big for number of frames! Stacking not possible."
    num_frames_after = num_frames - c
    x_new = []
    for start_frame in range(0, c):
        x_new.append(x[start_frame:start_frame + num_frames_after + 1, :])
    x = np.hstack(x_new)
    x = np.vstack((np.random.random((context_size, x.shape[1])), x, np.random.random((context_size, x.shape[1]))))

    return x

def normalize_euclidean(x):
    """ Frame-wise feature normalization to mean euclidean norm
    Args:
        x (2d ndarray): Feature matrix (num_frames x num_feat_dims)
    Returns:
        x (2d ndarray): Feature matrix (num_frames x num_feat_dims)
    """
    # avoid nans
    idx_zero = np.where(np.sum(x, axis=1) == 0)[0]
    for idx in idx_zero:
        x[idx, :] = .000001
    return x / np.sqrt(np.sum(np.square(x), axis=1, keepdims=True))

def main(args):
    """Main method to run transcription
    """

    # parse beats_file argument
    if args.beats_file == '':
        beat_times = None
    elif not os.path.exists(args.beats_file):
        beat_times = None
        print("[Warning] Could not find provided beats file.")
    else:
        beat_times = np.loadtxt(args.beat_file, delimiter=',', usecols=[0])

    transcriber = WalkingBassTranscription()

    transcriber.transcribe(args.input_wav,
                           dir_out=args.output_dir,
                           beat_times=beat_times,
                           threshold=args.threshold,
                           aggregation=args.aggregation)

    print('Finished bass line transcription! :)')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Predict walking bass salience or f0"
                    "from and audio file.")
    parser.add_argument("input_wav",
                        type=str,
                        help="Path to input wav file.")
    parser.add_argument("output_dir",
                        type=str,
                        help="Path to save location of bass transcription.")
    parser.add_argument("-b", "--beats_file",
                        type=str,
                        default='',
                        help="Path to beat annotation file. If not given, "
                        "does not produce note-level outputs.")
    parser.add_argument("-t", "--threshold",
                        type=float,
                        default=0.2,
                        help="Amplitude threshold. Only used when "
                        "output_format is singlef0 or multif0")
    parser.add_argument("-a", "--aggregation",
                        type=str,
                        default="beat",
                        help="Aggregation method. Possible values are 'beat' (beat-wise aggregation)"
                        "and 'flex-q' (dynamic estimation of most likely tatum per beat)")

    main(parser.parse_args())


# /usr/bin/python

from __future__ import print_function

import argparse
import torch
import pickle
import numpy as np
import os
import math
import random
import sys
import matplotlib.pyplot as plt
import data
import scipy.io
import pyLDAvis

from navec import Navec

from torch import nn, optim
from torch.nn import functional as F

from etm import ETM
from utils import (nearest_neighbors, get_topic_coherence, get_topic_diversity,
                   get_coherence_gensim, get_dictionary, get_topics)

parser = argparse.ArgumentParser(description='The Embedded Topic Model')

### data and file related arguments
parser.add_argument('--dataset', type=str, default='20ng', help='name of corpus')
parser.add_argument('--data_path', type=str, default='data/20ng', help='directory containing data')
parser.add_argument('--emb_path', type=str, default='data/20ng_embeddings.txt',
                    help='directory containing word embeddings')
parser.add_argument('--save_path', type=str, default='./results', help='path to save results')
parser.add_argument('--batch_size', type=int, default=1000, help='input batch size for training')

### model-related arguments
parser.add_argument('--num_topics', type=int, default=50, help='number of topics')
parser.add_argument('--rho_size', type=int, default=300, help='dimension of rho')
parser.add_argument('--emb_size', type=int, default=300, help='dimension of embeddings')
parser.add_argument('--t_hidden_size', type=int, default=800, help='dimension of hidden space of q(theta)')
parser.add_argument('--theta_act', type=str, default='relu',
                    help='tanh, softplus, relu, rrelu, leakyrelu, elu, selu, glu)')
parser.add_argument('--train_embeddings', type=int, default=0, help='whether to fix rho or train it')

### optimization-related arguments
parser.add_argument('--lr', type=float, default=0.005, help='learning rate')
parser.add_argument('--lr_factor', type=float, default=4.0, help='divide learning rate by this...')
parser.add_argument('--epochs', type=int, default=20, help='number of epochs to train...150 for 20ng 100 for others')
parser.add_argument('--mode', type=str, default='train', help='train or eval model')
parser.add_argument('--optimizer', type=str, default='adam', help='choice of optimizer')
parser.add_argument('--seed', type=int, default=2019, help='random seed (default: 1)')
parser.add_argument('--enc_drop', type=float, default=0.0, help='dropout rate on encoder')
parser.add_argument('--clip', type=float, default=0.0, help='gradient clipping')
parser.add_argument('--nonmono', type=int, default=10, help='number of bad hits allowed')
parser.add_argument('--wdecay', type=float, default=1.2e-6, help='some l2 regularization')
parser.add_argument('--anneal_lr', type=int, default=0, help='whether to anneal the learning rate or not')
parser.add_argument('--bow_norm', type=int, default=1, help='normalize the bows or not')

### evaluation, visualization, and logging-related arguments
parser.add_argument('--num_words', type=int, default=10, help='number of words for topic viz')
parser.add_argument('--log_interval', type=int, default=2, help='when to log training')
parser.add_argument('--visualize_every', type=int, default=10, help='when to visualize results')
parser.add_argument('--eval_batch_size', type=int, default=1000, help='input batch size for evaluation')
parser.add_argument('--load_from', type=str, default='', help='the name of the ckpt to eval from')
parser.add_argument('--tc', type=int, default=0, help='whether to compute topic coherence or not')
parser.add_argument('--td', type=int, default=0, help='whether to compute topic diversity or not')

args = parser.parse_args()

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

print('\n')
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(args.seed)

## get data
# 1. vocabulary
vocab, train, valid, test = data.get_data(os.path.join(args.data_path))
vocab_size = len(vocab)
args.vocab_size = vocab_size

# 1. training data
train_tokens = train['tokens']
train_counts = train['counts']
args.num_docs_train = len(train_tokens)

# 2. dev set
valid_tokens = valid['tokens']
valid_counts = valid['counts']
args.num_docs_valid = len(valid_tokens)

# 3. test data
test_tokens = test['tokens']
test_counts = test['counts']
args.num_docs_test = len(test_tokens)
test_1_tokens = test['tokens_1']
test_1_counts = test['counts_1']
args.num_docs_test_1 = len(test_1_tokens)
test_2_tokens = test['tokens_2']
test_2_counts = test['counts_2']
args.num_docs_test_2 = len(test_2_tokens)

embeddings = None
if not args.train_embeddings:
    emb_path = args.emb_path
    navec = Navec.load(emb_path)
    # vect_path = os.path.join(args.data_path.split('/')[0], 'embeddings.pkl')
    # vectors = {}
    # with open(emb_path, 'rb') as f:
    #    for l in f:
    #        line = l.decode().split()
    #        word = line[0]
    #        if word in vocab:
    #            vect = np.array(line[1:]).astype(np.float)
    #            vectors[word] = vect
    embeddings = np.zeros((vocab_size, args.emb_size))
    words_found = 0
    for i, word in enumerate(vocab):
        try:
            embeddings[i] = navec[word]
            words_found += 1
        except KeyError:
            embeddings[i] = np.random.normal(scale=0.6, size=(args.emb_size,))
    embeddings = torch.from_numpy(embeddings).to(device)
    args.embeddings_dim = embeddings.size()

print('=*' * 100)
print('Training an Embedded Topic Model on {} with the following settings: {}'.format(args.dataset.upper(), args))
print('=*' * 100)

## define checkpoint
if not os.path.exists(args.save_path):
    os.makedirs(args.save_path)

if args.mode in ['eval', 'ppx']:
    ckpt = args.load_from
else:
    ckpt = os.path.join(args.save_path,
                        'etm_{}_K_{}_Htheta_{}_Optim_{}_Clip_{}_ThetaAct_{}_Lr_{}_Bsz_{}_RhoSize_{}_trainEmbeddings_{}'.format(
                            args.dataset, args.num_topics, args.t_hidden_size, args.optimizer, args.clip,
                            args.theta_act,
                            args.lr, args.batch_size, args.rho_size, args.train_embeddings))

## define model and optimizer
model = ETM(args.num_topics, vocab_size, args.t_hidden_size, args.rho_size, args.emb_size,
            args.theta_act, embeddings, args.train_embeddings, args.enc_drop).to(device)

print('model: {}'.format(model))

if args.optimizer == 'adam':
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
elif args.optimizer == 'adagrad':
    optimizer = optim.Adagrad(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
elif args.optimizer == 'adadelta':
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
elif args.optimizer == 'rmsprop':
    optimizer = optim.RMSprop(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
elif args.optimizer == 'asgd':
    optimizer = optim.ASGD(model.parameters(), lr=args.lr, t0=0, lambd=0., weight_decay=args.wdecay)
else:
    print('Defaulting to vanilla SGD')
    optimizer = optim.SGD(model.parameters(), lr=args.lr)

train_losses = []
val_losses = []

def train(epoch):
    model.train()
    acc_loss = 0
    acc_kl_theta_loss = 0
    cnt = 0
    indices = torch.randperm(args.num_docs_train)
    indices = torch.split(indices, args.batch_size)
    for idx, ind in enumerate(indices):
        optimizer.zero_grad()
        model.zero_grad()
        data_batch = data.get_batch(train_tokens, train_counts, ind, args.vocab_size, device)
        sums = data_batch.sum(1).unsqueeze(1)
        if args.bow_norm:
            normalized_data_batch = data_batch / sums
        else:
            normalized_data_batch = data_batch
        recon_loss, kld_theta = model(data_batch, normalized_data_batch)
        total_loss = recon_loss + kld_theta
        total_loss.backward()

        if args.clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
        optimizer.step()

        acc_loss += torch.sum(recon_loss).item()
        acc_kl_theta_loss += torch.sum(kld_theta).item()
        cnt += 1

        if idx % args.log_interval == 0 and idx > 0:
            cur_loss = round(acc_loss / cnt, 2)
            cur_kl_theta = round(acc_kl_theta_loss / cnt, 2)
            cur_real_loss = round(cur_loss + cur_kl_theta, 2)

            # print('Epoch: {} .. batch: {}/{} .. LR: {} .. KL_theta: {} .. Rec_loss: {} .. NELBO: {}'.format(
            # epoch, idx, len(indices), optimizer.param_groups[0]['lr'], cur_kl_theta, cur_loss, cur_real_loss))

    cur_loss = round(acc_loss / cnt, 2)
    cur_kl_theta = round(acc_kl_theta_loss / cnt, 2)
    cur_real_loss = round(cur_loss + cur_kl_theta, 2)

    train_losses.append(cur_loss)

    print('*' * 100)
    print('Epoch----->{} .. LR: {} .. KL_theta: {} .. Rec_loss: {} .. NELBO: {}'.format(
        epoch, optimizer.param_groups[0]['lr'], cur_kl_theta, cur_loss, cur_real_loss))
    print('*' * 100)

def visualize(m, show_emb=True):
    if not os.path.exists('./results'):
        os.makedirs('./results')

    m.eval()

    queries = ['россия', 'человек', 'компьютер', 'спорт', 'религия', 'любовь',
               'доллар', 'правительство', 'здоровье', 'семья']

    ## visualize topics using monte carlo
    with torch.no_grad():
        print('#' * 100)
        print('Visualize topics...')
        topics_words = []
        gammas = m.get_beta()
        for k in range(args.num_topics):
            gamma = gammas[k]
            top_words = list(gamma.cpu().numpy().argsort()[-args.num_words + 1:][::-1])
            topic_words = [vocab[a] for a in top_words]
            topics_words.append(' '.join(topic_words))
            print('Topic {}: {}'.format(k, topic_words))

        if show_emb:
            ## visualize word embeddings by using V to get nearest neighbors
            print('#' * 100)
            print('Visualize word embeddings by using output embedding matrix')
            try:
                embeddings = m.rho.weight  # Vocab_size x E
            except:
                embeddings = m.rho  # Vocab_size x E
            neighbors = []
            for word in queries:
                print('word: {} .. neighbors: {}'.format(
                    word, nearest_neighbors(word, embeddings, vocab)))
            print('#' * 100)

def evaluate(m, source, tc=False, td=False):
    """Compute perplexity on document completion.
    """
    m.eval()
    with torch.no_grad():
        if source == 'val':
            indices = torch.split(torch.tensor(range(args.num_docs_valid)), args.eval_batch_size)
            tokens = valid_tokens
            counts = valid_counts
        else:
            indices = torch.split(torch.tensor(range(args.num_docs_test)), args.eval_batch_size)
            tokens = test_tokens
            counts = test_counts

        ## get \beta here
        beta = m.get_beta()

        ### do dc and tc here
        acc_loss = 0
        cnt = 0
        indices_1 = torch.split(torch.tensor(range(args.num_docs_test_1)), args.eval_batch_size)
        for idx, ind in enumerate(indices_1):
            ## get theta from first half of docs
            data_batch_1 = data.get_batch(test_1_tokens, test_1_counts, ind, args.vocab_size, device)
            sums_1 = data_batch_1.sum(1).unsqueeze(1)
            if args.bow_norm:
                normalized_data_batch_1 = data_batch_1 / sums_1
            else:
                normalized_data_batch_1 = data_batch_1
            theta, _ = m.get_theta(normalized_data_batch_1)

            ## get prediction loss using second half
            data_batch_2 = data.get_batch(test_2_tokens, test_2_counts, ind, args.vocab_size, device)
            sums_2 = data_batch_2.sum(1).unsqueeze(1)
            res = torch.mm(theta, beta)
            preds = torch.log(res)
            recon_loss = -(preds * data_batch_2).sum(1)

            loss = recon_loss / sums_2.squeeze()
            loss = loss.mean().item()
            acc_loss += loss
            cnt += 1
        cur_loss = acc_loss / cnt
        print('{} loss: {}'.format(source.upper(), cur_loss))
        val_losses.append(cur_loss)

        ppl_dc = round(math.exp(cur_loss), 1)
        print('*' * 100)
        print('{} Doc Completion PPL: {}'.format(source.upper(), ppl_dc))
        print('*' * 100)
        if tc or td:
            beta = beta.data.cpu().numpy()
            if tc:
                print('Computing topic coherence...')
                get_topic_coherence(beta, train_tokens, vocab)
            if td:
                print('Computing topic diversity...')
                get_topic_diversity(beta, 25)
        return ppl_dc

def prepare_viz_data(phi, theta, n_wd, vocab):
    theta = theta / theta.sum(axis=1, keepdims=1)
    data = {'topic_term_dists': phi,
            'doc_topic_dists': theta,
            'doc_lengths': n_wd.sum(axis=1).tolist(),
            'vocab': vocab,
            'term_frequency': n_wd.sum(axis=0).tolist()}
    return data

def show_viz(model, data_batch, vocab):
    phi = model.get_beta().cpu().numpy()
    theta, _ = model.get_theta(data_batch)
    theta = theta.cpu().numpy()
    model_data = prepare_viz_data(phi, theta, data_batch, vocab)
    model_viz = pyLDAvis.prepare(**model_data)
    pyLDAvis.save_html(model_viz, 'etm_vis.html')
    print('\n\nVisualization has been saved')

def plot_curves():
    fig, axs = plt.subplots(1, 2)
    axs[0].plot(train_losses)
    axs[0].set_title('train')
    axs[1].plot(val_losses)
    axs[1].set_title('val')
    plt.ylabel('loss values')
    plt.xlabel('epoch')
    plt.show()

if args.mode == 'train':
    ## train model on data
    best_epoch = 0
    best_val_ppl = 1e9
    all_val_ppls = []
    print('\n')
    print('Visualizing model quality before training...')
    visualize(model)
    print('\n')
    for epoch in range(1, args.epochs):
        train(epoch)
        val_ppl = evaluate(model, 'val')
        if val_ppl < best_val_ppl:
            with open(ckpt, 'wb') as f:
                torch.save(model, f)
            best_epoch = epoch
            best_val_ppl = val_ppl
        else:
            ## check whether to anneal lr
            lr = optimizer.param_groups[0]['lr']
            if args.anneal_lr and (
                    len(all_val_ppls) > args.nonmono and val_ppl > min(all_val_ppls[:-args.nonmono]) and lr > 1e-5):
                optimizer.param_groups[0]['lr'] /= args.lr_factor
        if epoch % args.visualize_every == 0:
            visualize(model)
        all_val_ppls.append(val_ppl)

    plot_curves()

    with open(ckpt, 'rb') as f:
        model = torch.load(f)
    model = model.to(device)
    val_ppl = evaluate(model, 'val')
elif args.mode == 'eval':
    with open(ckpt, 'rb') as f:
        model = torch.load(f)
    model = model.to(device)
    model.eval()

    with torch.no_grad():
        ## get document completion perplexities
        test_ppl = evaluate(model, 'test', tc=args.tc, td=args.td)

        ## get most used topics
        indices = torch.tensor(range(args.num_docs_train))
        indices = torch.split(indices, args.batch_size)
        thetaAvg = torch.zeros(1, args.num_topics).to(device)
        thetaWeightedAvg = torch.zeros(1, args.num_topics).to(device)
        cnt = 0
        for idx, ind in enumerate(indices):
            data_batch = data.get_batch(train_tokens, train_counts, ind, args.vocab_size, device)
            sums = data_batch.sum(1).unsqueeze(1)
            cnt += sums.sum(0).squeeze().cpu().numpy()
            if args.bow_norm:
                normalized_data_batch = data_batch / sums
            else:
                normalized_data_batch = data_batch
            theta, _ = model.get_theta(normalized_data_batch)
            thetaAvg += theta.sum(0).unsqueeze(0) / args.num_docs_train
            weighed_theta = sums * theta
            thetaWeightedAvg += weighed_theta.sum(0).unsqueeze(0)
            if idx % 100 == 0 and idx > 0:
                print('batch: {}/{}'.format(idx, len(indices)))
        thetaWeightedAvg = thetaWeightedAvg.squeeze().cpu().numpy() / cnt
        print('\nThe 10 most used topics are {}'.format(thetaWeightedAvg.argsort()[::-1][:10]))

        ## show topics
        beta = model.get_beta()
        topic_indices = list(np.random.choice(args.num_topics, 10))  # 10 random topics
        print('\n')
        for k in range(args.num_topics):  # topic_indices:
            gamma = beta[k]
            top_words = list(gamma.cpu().numpy().argsort()[-args.num_words + 1:][::-1])
            topic_words = [vocab[a] for a in top_words]
            print('Topic {}: {}'.format(k, topic_words))

        if args.train_embeddings:
            ## show etm embeddings
            try:
                rho_etm = model.rho.weight.cpu()
            except:
                rho_etm = model.rho.cpu()
            queries = ['россия', 'человек', 'компьютер', 'спорт', 'религия', 'любовь',
                       'доллар', 'правительство', 'здоровье', 'семья']
            print('\n')
            print('ETM embeddings...')
            for word in queries:
                print('word: {} .. etm neighbors: {}'.format(word, nearest_neighbors(word, rho_etm, vocab)))
            print('\n')
else:
    with open(ckpt, 'rb') as f:
        model = torch.load(f)
    model = model.to(device)
    model.eval()

    with torch.no_grad():
        ## get document completion perplexities

        indices = torch.split(torch.tensor(range(args.num_docs_test)), args.eval_batch_size)
        tokens = test_tokens
        counts = test_counts

        beta = model.get_beta()
        beta[beta < 9e-7] = 0

        ppx = 0
        nd = 0
        indices_1 = torch.split(torch.tensor(range(args.num_docs_test_1)), args.eval_batch_size)
        for idx, ind in enumerate(indices_1):
            data_batch_1 = data.get_batch(test_1_tokens, test_1_counts, ind, args.vocab_size, device)
            theta, _ = model.get_theta(data_batch_1)
            theta[theta < 5e-5] = 0

            np.seterr(divide='ignore')
            p_wd = np.matmul(theta.cpu().numpy(), beta.cpu().numpy())
            p_wd = np.log(p_wd)
            p_wd[np.isneginf(p_wd)] = 0
            recon_loss = -1 * np.multiply(p_wd, data_batch_1.cpu().numpy()).sum()

            sums = data_batch_1.sum()
            ppx += recon_loss.item()
            nd += sums

        avg_ppx = ppx / nd
        res_ppx = math.exp(avg_ppx)
        indices = torch.split(torch.tensor(range(args.num_docs_train)), 1000)
        data_batch = data.get_batch(train_tokens, train_counts, indices[0], args.vocab_size, device)
        show_viz(model, data_batch, vocab)

        dictionary, docs = get_dictionary(vocab, data_batch)
        print(dictionary.token2id['дело'])
        topics = get_topics(vocab, beta, dictionary)

        print(f'result perplexity: {res_ppx}')
        print(f'beta sparsity: {1.0 - torch.count_nonzero(beta) / torch.numel(beta)}')
        print(f'theta sparsity: {1.0 - torch.count_nonzero(theta) / torch.numel(theta)}')
        print(f'Coherence c_v: {get_coherence_gensim(topics, dictionary, "c_v", docs)}')
        print(f'Coherence c_nmpi: {get_coherence_gensim(topics, dictionary, "c_npmi", docs)}')


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

""" Helper functions and class to calculate Average Precisions for 3D object detection.
"""
import os
import sys
import numpy as np
import torch
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
from eval_det import eval_det_cls, eval_det_multiprocessing
from eval_det import get_iou_obb
from nms import nms_2d_faster, nms_3d_faster, nms_3d_faster_samecls
from box_util import get_3d_box
sys.path.append(os.path.join(ROOT_DIR, 'sunrgbd'))
from sunrgbd_utils import extract_pc_in_box3d

def flip_axis_to_camera(pc):
    ''' Flip X-right,Y-forward,Z-up to X-right,Y-down,Z-forward
    Input and output are both (N,3) array
    '''
    pc2 = np.copy(pc)
    pc2[...,[0,1,2]] = pc2[...,[0,2,1]] # cam X,Y,Z = depth X,-Z,Y
    pc2[...,1] *= -1
    return pc2

def flip_axis_to_depth(pc):
    pc2 = np.copy(pc)
    pc2[...,[0,1,2]] = pc2[...,[0,2,1]] # depth X,Y,Z = cam X,Z,-Y
    pc2[...,2] *= -1
    return pc2

def softmax(x):
    ''' Numpy function for softmax'''
    shape = x.shape
    probs = np.exp(x - np.max(x, axis=len(shape)-1, keepdims=True))
    probs /= np.sum(probs, axis=len(shape)-1, keepdims=True)
    return probs

def parse_predictions(end_points, config_dict):
    """ Parse predictions to OBB parameters and suppress overlapping boxes

    Args:
        end_points: dict
            {point_clouds, center, heading_scores, heading_residuals,
            size_scores, size_residuals, sem_cls_scores}
        config_dict: dict
            {dataset_config, remove_empty_box, use_3d_nms, nms_iou,
            use_old_type_nms, conf_thresh, per_class_proposal}

    Returns:
        batch_pred_map_cls: a list of len == batch size (BS)
            [pred_list_i], i = 0, 1, ..., BS-1
            where pred_list_i = [(pred_sem_cls, box_params, box_score)_j]
            where j = 0, ..., num of valid detections - 1 from sample input i
    """
    pred_center = end_points['center'] # B,num_proposal,3
    pred_heading_class = torch.argmax(end_points['heading_scores'], -1) # B,num_proposal
    pred_heading_residual = torch.gather(end_points['heading_residuals'], 2,
        pred_heading_class.unsqueeze(-1)) # B,num_proposal,1
    pred_heading_residual.squeeze_(2)
    pred_size_class = torch.argmax(end_points['size_scores'], -1) # B,num_proposal
    pred_size_residual = torch.gather(end_points['size_residuals'], 2,
        pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(1,1,1,3)) # B,num_proposal,1,3
    pred_size_residual.squeeze_(2)
    pred_sem_cls = torch.argmax(end_points['sem_cls_scores'], -1) # B,num_proposal
    sem_cls_probs = softmax(end_points['sem_cls_scores'].detach().cpu().numpy()) # B,num_proposal,10
    pred_sem_cls_prob = np.max(sem_cls_probs,-1) # B,num_proposal

    num_proposal = pred_center.shape[1]
    # Since we operate in upright_depth coord for points, while util functions
    # assume upright_camera coord.
    bsize = pred_center.shape[0]
    pred_corners_3d_upright_camera = np.zeros((bsize, num_proposal, 8, 3))
    pred_center_upright_camera = flip_axis_to_camera(pred_center.detach().cpu().numpy())
    for i in range(bsize):
        for j in range(num_proposal):
            heading_angle = config_dict['dataset_config'].class2angle(\
                pred_heading_class[i,j].detach().cpu().numpy(), pred_heading_residual[i,j].detach().cpu().numpy())
            box_size = config_dict['dataset_config'].class2size(\
                int(pred_size_class[i,j].detach().cpu().numpy()), pred_size_residual[i,j].detach().cpu().numpy())
            corners_3d_upright_camera = get_3d_box(box_size, heading_angle, pred_center_upright_camera[i,j,:])
            pred_corners_3d_upright_camera[i,j] = corners_3d_upright_camera

    K = pred_center.shape[1] # K==num_proposal
    nonempty_box_mask = np.ones((bsize, K))

    if config_dict['remove_empty_box']:
        # -------------------------------------
        # Remove predicted boxes without any point within them..
        batch_pc = end_points['point_clouds'].cpu().numpy()[:,:,0:3] # B,N,3
        for i in range(bsize):
            pc = batch_pc[i,:,:] # (N,3)
            for j in range(K):
                box3d = pred_corners_3d_upright_camera[i,j,:,:] # (8,3)
                box3d = flip_axis_to_depth(box3d)
                pc_in_box,inds = extract_pc_in_box3d(pc, box3d)
                if len(pc_in_box) < 5:
                    nonempty_box_mask[i,j] = 0
        # -------------------------------------

    obj_logits = end_points['objectness_scores'].detach().cpu().numpy()
    obj_prob = softmax(obj_logits)[:,:,1] # (B,K)
    if not config_dict['use_3d_nms']:
        # ---------- NMS input: pred_with_prob in (B,K,7) -----------
        pred_mask = np.zeros((bsize, K))
        for i in range(bsize):
            boxes_2d_with_prob = np.zeros((K,5))
            for j in range(K):
                boxes_2d_with_prob[j,0] = np.min(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_2d_with_prob[j,2] = np.max(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_2d_with_prob[j,1] = np.min(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_2d_with_prob[j,3] = np.max(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_2d_with_prob[j,4] = obj_prob[i,j]
            nonempty_box_inds = np.where(nonempty_box_mask[i,:]==1)[0]
            pick = nms_2d_faster(boxes_2d_with_prob[nonempty_box_mask[i,:]==1,:],
                config_dict['nms_iou'], config_dict['use_old_type_nms'])
            assert(len(pick)>0)
            pred_mask[i, nonempty_box_inds[pick]] = 1
        end_points['pred_mask'] = pred_mask
        # ---------- NMS output: pred_mask in (B,K) -----------
    elif config_dict['use_3d_nms'] and (not config_dict['cls_nms']):
        # ---------- NMS input: pred_with_prob in (B,K,7) -----------
        pred_mask = np.zeros((bsize, K))
        for i in range(bsize):
            boxes_3d_with_prob = np.zeros((K,7))
            for j in range(K):
                boxes_3d_with_prob[j,0] = np.min(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_3d_with_prob[j,1] = np.min(pred_corners_3d_upright_camera[i,j,:,1])
                boxes_3d_with_prob[j,2] = np.min(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_3d_with_prob[j,3] = np.max(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_3d_with_prob[j,4] = np.max(pred_corners_3d_upright_camera[i,j,:,1])
                boxes_3d_with_prob[j,5] = np.max(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_3d_with_prob[j,6] = obj_prob[i,j]
            nonempty_box_inds = np.where(nonempty_box_mask[i,:]==1)[0]
            pick = nms_3d_faster(boxes_3d_with_prob[nonempty_box_mask[i,:]==1,:],
                config_dict['nms_iou'], config_dict['use_old_type_nms'])
            assert(len(pick)>0)
            pred_mask[i, nonempty_box_inds[pick]] = 1
        end_points['pred_mask'] = pred_mask
        # ---------- NMS output: pred_mask in (B,K) -----------
    elif config_dict['use_3d_nms'] and config_dict['cls_nms']:
        # ---------- NMS input: pred_with_prob in (B,K,8) -----------
        pred_mask = np.zeros((bsize, K))
        for i in range(bsize):
            boxes_3d_with_prob = np.zeros((K,8))
            for j in range(K):
                boxes_3d_with_prob[j,0] = np.min(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_3d_with_prob[j,1] = np.min(pred_corners_3d_upright_camera[i,j,:,1])
                boxes_3d_with_prob[j,2] = np.min(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_3d_with_prob[j,3] = np.max(pred_corners_3d_upright_camera[i,j,:,0])
                boxes_3d_with_prob[j,4] = np.max(pred_corners_3d_upright_camera[i,j,:,1])
                boxes_3d_with_prob[j,5] = np.max(pred_corners_3d_upright_camera[i,j,:,2])
                boxes_3d_with_prob[j,6] = obj_prob[i,j]
                boxes_3d_with_prob[j,7] = pred_sem_cls[i,j] # only suppress if the two boxes are of the same class!!
            nonempty_box_inds = np.where(nonempty_box_mask[i,:]==1)[0]
            pick = nms_3d_faster_samecls(boxes_3d_with_prob[nonempty_box_mask[i,:]==1,:],
                config_dict['nms_iou'], config_dict['use_old_type_nms'])
            assert(len(pick)>0)
            pred_mask[i, nonempty_box_inds[pick]] = 1
        end_points['pred_mask'] = pred_mask
        # ---------- NMS output: pred_mask in (B,K) -----------

    batch_pred_map_cls = [] # a list (len: batch_size) of list (len: num of predictions per sample) of tuples of pred_cls, pred_box and conf (0-1)
    for i in range(bsize):
        if config_dict['per_class_proposal']:
            cur_list = []
            for ii in range(config_dict['dataset_config'].num_class):
                cur_list += [(ii, pred_corners_3d_upright_camera[i,j], sem_cls_probs[i,j,ii]*obj_prob[i,j]) \
                    for j in range(pred_center.shape[1]) if pred_mask[i,j]==1 and obj_prob[i,j]>config_dict['conf_thresh']]
            batch_pred_map_cls.append(cur_list)
        else:
            batch_pred_map_cls.append([(pred_sem_cls[i,j].item(), pred_corners_3d_upright_camera[i,j], obj_prob[i,j]) \
                for j in range(pred_center.shape[1]) if pred_mask[i,j]==1 and obj_prob[i,j]>config_dict['conf_thresh']])
    end_points['batch_pred_map_cls'] = batch_pred_map_cls

    return batch_pred_map_cls

def parse_groundtruths(end_points, config_dict):
    """ Parse groundtruth labels to OBB parameters.

    Args:
        end_points: dict
            {center_label, heading_class_label, heading_residual_label,
            size_class_label, size_residual_label, sem_cls_label,
            box_label_mask}
        config_dict: dict
            {dataset_config}

    Returns:
        batch_gt_map_cls: a list  of len == batch_size (BS)
            [gt_list_i], i = 0, 1, ..., BS-1
            where gt_list_i = [(gt_sem_cls, gt_box_params)_j]
            where j = 0, ..., num of objects - 1 at sample input i
    """
    center_label = end_points['center_label']
    heading_class_label = end_points['heading_class_label']
    heading_residual_label = end_points['heading_residual_label']
    size_class_label = end_points['size_class_label']
    size_residual_label = end_points['size_residual_label']
    box_label_mask = end_points['box_label_mask']
    sem_cls_label = end_points['sem_cls_label']
    bsize = center_label.shape[0]

    K2 = center_label.shape[1] # K2==MAX_NUM_OBJ
    gt_corners_3d_upright_camera = np.zeros((bsize, K2, 8, 3))
    gt_center_upright_camera = flip_axis_to_camera(center_label[:,:,0:3].detach().cpu().numpy())
    for i in range(bsize):
        for j in range(K2):
            if box_label_mask[i,j] == 0: continue
            heading_angle = config_dict['dataset_config'].class2angle(heading_class_label[i,j].detach().cpu().numpy(), heading_residual_label[i,j].detach().cpu().numpy())
            box_size = config_dict['dataset_config'].class2size(int(size_class_label[i,j].detach().cpu().numpy()), size_residual_label[i,j].detach().cpu().numpy())
            corners_3d_upright_camera = get_3d_box(box_size, heading_angle, gt_center_upright_camera[i,j,:])
            gt_corners_3d_upright_camera[i,j] = corners_3d_upright_camera

    batch_gt_map_cls = []
    for i in range(bsize):
        batch_gt_map_cls.append([(sem_cls_label[i,j].item(), gt_corners_3d_upright_camera[i,j]) for j in range(gt_corners_3d_upright_camera.shape[1]) if box_label_mask[i,j]==1])
    end_points['batch_gt_map_cls'] = batch_gt_map_cls

    return batch_gt_map_cls

class APCalculator(object):
    ''' Calculating Average Precision '''
    def __init__(self, ap_iou_thresh=0.25, class2type_map=None):
        """
        Args:
            ap_iou_thresh: float between 0 and 1.0
                IoU threshold to judge whether a prediction is positive.
            class2type_map: [optional] dict {class_int:class_name}
        """
        self.ap_iou_thresh = ap_iou_thresh
        self.class2type_map = class2type_map
        self.reset()

    def step(self, batch_pred_map_cls, batch_gt_map_cls):
        """ Accumulate one batch of prediction and groundtruth.

        Args:
            batch_pred_map_cls: a list of lists [[(pred_cls, pred_box_params, score),...],...]
            batch_gt_map_cls: a list of lists [[(gt_cls, gt_box_params),...],...]
                should have the same length with batch_pred_map_cls (batch_size)
        """

        bsize = len(batch_pred_map_cls)
        assert(bsize == len(batch_gt_map_cls))
        for i in range(bsize):
            self.gt_map_cls[self.scan_cnt] = batch_gt_map_cls[i]
            self.pred_map_cls[self.scan_cnt] = batch_pred_map_cls[i]
            self.scan_cnt += 1

    def compute_metrics(self):
        """ Use accumulated predictions and groundtruths to compute Average Precision.
        """
        rec, prec, ap = eval_det_multiprocessing(self.pred_map_cls, self.gt_map_cls, ovthresh=self.ap_iou_thresh, get_iou_func=get_iou_obb)
        ret_dict = {}
        for key in sorted(ap.keys()):
            clsname = self.class2type_map[key] if self.class2type_map else str(key)
            ret_dict['%s Average Precision'%(clsname)] = ap[key]
        temp = []
        for v in list(ap.values()):
            if np.isnan(v):
                continue
            temp.append(v)
        ret_dict['mAP'] = np.mean(temp)
        rec_list = []
        for key in sorted(ap.keys()):
            clsname = self.class2type_map[key] if self.class2type_map else str(key)
            try:
                ret_dict['%s Recall'%(clsname)] = rec[key][-1]
                rec_list.append(rec[key][-1])
            except:
                ret_dict['%s Recall'%(clsname)] = 0
                rec_list.append(0)
        temp = []
        for v in rec_list:
            if np.isnan(v):
                continue
            temp.append(v)
        ret_dict['AR'] = np.mean(temp)
        return ret_dict

    def reset(self):
        self.gt_map_cls = {} # {scan_id: [(classname, bbox)]}
        self.pred_map_cls = {} # {scan_id: [(classname, bbox, score)]}
        self.scan_cnt = 0


#!/usr/bin/env python
#####################################################
# How to run this script
#
# From Linux Terminal:$./calculon.py Results_File Symbol_Ref1 Symbol_Ref2 Atomic_Number Output_File_Name
#
# 'Symbol_Ref1' and 'Symbol_Ref2' can take the following format --> '40Ca+19F' or '85Rb' (without the quotes)
#
# v1.0 Created  by D. Atanasov @ 16.09.2014
# v2.0 Modified by D. Atanasov @ 14.02.2017
# v3.0 Modified to have OOP by D. Atanaso @ 21.03.2017
# v4.0 Modified to import full AME table by J. Karthein @ 11.04.2017
#####################################################
import os
import sys
import re
import numpy as np
import pandas as pd
import platform

#   ######################################################
#   #              Function's Definition                 #
#   ######################################################
class Ame():
    def __init__(self):
        self.scale = 931494.061                          # [amu]->[keV]  #
        self.me = 548.57990946                           # Electron mass #
        self.me_unc = 0.00000022                         # uncertainty   #
        self.load_table()

    def get_number_symbol(self, item):
        """
        Get the Atomic number and the Element symbol from a specific string ()
        Examples: (40Ca19F), (H1H1:O16), 85Rb, 136Cd etc.
        :param item:
        :return: Atomic number, Element Symbol
        """
        return re.findall('[0-9][0-9][0-9]|[0-9][0-9]|[0-9]', item), re.findall('[A-Z][a-z]|[A-Z]', item)

    def load_table(self, ame='AME16.txt'):
        '''
        The function reads the AME table and returns a DataFrame of str's.

        Input name:     'AME16.txt'
        Further info:   If a new table will be uploaded, please change the file + name.
                        Before starting please remove all '*' by ' ' and all '#' by '.'
        Questions to:   jonas.karthein@cern.ch
        '''
        if platform.system() == 'Windows':
            os.chdir('G:\\Experiments\\ISOLTRAP\\Software\\PI-ICR\\Python-DAQ')
        elif platform.system() == 'Darwin': # MacOS
            # os.chdir('/Volumes/dfs/Software/PI-ICR/Python-DAQ/')
            os.chdir('/Users/jonaskarthein/cernbox/Software/Python/piicr-analysis/')
        ame_import = np.genfromtxt(ame, skip_header=39, dtype=['a1', 'int', 'int', 'int', 'int', 'a4', 'a4', 'float', 'float', 'float', 'float', 'a3', 'float', 'float', 'int', 'float', 'float'], delimiter=[1,3,5,5,5,4,4,16,12,12,6,3,11,8,4,13,11])
        ame_table = [list(elem) for elem in ame_import.tolist()]    # convert list of tuples to list of lists

        for i in range(len(ame_table)):
            ame_table[i][15] = ame_table[i][14] * 1E6 + ame_table[i][15]    # calculate full atomic mass (int*1000000+float)
            for j in range(len(ame_table[i])):
                if type(ame_table[i][j]) == str:
                    ame_table[i][j] = ame_table[i][j].replace(" ", "")      # delete all spaces --> makes searching easier
                ame_table[i][j] = str(ame_table[i][j])                      # convert all entries to int (needed for Dinkos class)
        self.df = pd.DataFrame(ame_table, columns=['cc', 'NZ', 'N', 'Z', 'A', 'EL', 'o', 'mass excess / keV', 'mass excess unc / keV', 'binding energy / keV', 'binding energy unc / keV', 'B', 'beta decay energy / keV', 'beta decay energy unc / keV', 'atomic mass (int) / Dalton', 'atomic mass / micro Dalton', 'atomic mass unc'])

    def get_ame_mass(self, el_expr):
        """
        Checks AME Table for existing of the element with the given atomic number.
        Calculates the mass. If a list is provided to the function (such as a molecule) it calculates
        the summed mass of the constituents.
        :param atomic_number:
        :param symbols:
        :param ame:
        :return: el_expr, atom_mass, np.sqrt(unc)
        """
        ame_mass = 0.0
        ame_unc = 0.0
        self.idx = []
        atomic_number, symbols = self.get_number_symbol(el_expr)
        for i in range(len(symbols)):
            self.idx.append(self.df[(self.df['EL'] == symbols[i]) & (self.df['A'] == atomic_number[i])].index.tolist())
        if len(self.idx[0]) == 0:
            print "Element or Atomic number not found in database"
        else:
            for i in range(len(self.idx)):
                ame_mass += float(self.df.get_value(self.idx[i][0], 'atomic mass / micro Dalton'))
                ame_unc += float(self.df.get_value(self.idx[i][0], 'atomic mass unc'))**2
        return el_expr, ame_mass, np.sqrt(ame_unc)

    def get_ion_mass(self, mass_atom, mass_atom_unc, charge=1):
        """
        Calculate the ion mass
        :param mass_atom:
        :param mass_atom_unc:
        :param charge:
        :return: massIon, massIonUnc
        """
        return (mass_atom - charge*self.me), np.sqrt(mass_atom_unc**2 + charge*self.me_unc**2)

    def get_isobars(self, atomic_number):
        """
        Checks AME Table for existing of the element with the given atomic number.
        Calculates the mass. If a list is provided to the function (such as a molecule) it calculates
        the summed mass of the constituents.
        :param atomic_number:
        :param symbols:
        :param ame:
        :return: el_expr, atom_mass, np.sqrt(unc)
        """
        idx = []
        print atomic_number, type(atomic_number)
        print self.df[self.df['A'] == str(atomic_number)]
        idx = self.df[self.df['A'] == str(atomic_number)].index.tolist()
        print len(idx)
        isobars = np.empty((len(idx), 3,), dtype=object)
        if len(idx) == 0:
            print "No isobars found! Please check the atomic Number"
        else:
            for i in range(len(idx)):
                print idx[i], type(idx[i])
                isobars[i][0] = self.df.get_value(idx[i], 'EL')
                isobars[i][1] = float(self.df.get_value(idx[i], 'atomic mass / micro Dalton'))
                isobars[i][2] = float(self.df.get_value(idx[i], 'atomic mass unc'))
        return isobars

if __name__ == "__main__":
    main(sys.argv)


"""
Create extinguished grid more segmented dealing with large grids with enough
memory

All functions are now transformed into generators. As a result, any function
allows computation of a grid in an arbitrary number of chunks. This offers the
possibility to generate grids that cannot fit in memory.

.. note::

    * dependencies have also been updated accordingly.

    * likelihood computations need to be updated to allow computations even if
      the full grid does not fit in memory
"""
import numpy as np
import copy

from astropy import units
from tqdm import tqdm

from beast.physicsmodel.stars import stellib
from beast.physicsmodel.grid import SpectralGrid, SEDGrid
from beast.physicsmodel.prior_weights_dust import PriorWeightsDust
# from beast.external.eztables import Table
from astropy.table import Table
from beast.tools.helpers import generator
from beast.tools import helpers

from beast.observationmodel.noisemodel import absflux_covmat

__all__ = [
    "gen_spectral_grid_from_stellib_given_points",
    "make_extinguished_grid",
    "add_spectral_properties",
    "calc_absflux_cov_matrices",
]

@generator
def gen_spectral_grid_from_stellib_given_points(
    osl, pts, bounds=dict(dlogT=0.1, dlogg=0.3), chunksize=0
):
    """
    Generator that reinterpolates a given stellar spectral library on to
       an Isochrone grid

    It will iterate over a list of `pts` points and generate
       `chunksize` models until all the list of points is processed

    Parameters
    ----------
    osl: stellib.stellib
        a stellar library

    pts: dict like structure of points
        dictionary like or named data structure of points to interpolate at
        must contain logg, logT, logL, and Z

    bounds:  dict, optional (default={dlogT:0.1, dlogg:0.3})
        sensitivity to extrapolation (see grid.get_stellib_boundaries)

    chunksize: int, optional (default=0)
        number of models to generate at each cycle.
        If default <= 0, all models will be returned at once.

    Returns
    -------
    g: SpectralGrid
        Spectral grid (in memory) containing the requested list of stars
        and associated spectra
    """

    helpers.type_checker("osl", osl, stellib.Stellib)

    if chunksize <= 0:
        yield osl.gen_spectral_grid_from_given_points(pts, bounds=bounds)
    else:
        try:
            # Yield successive n-sized chunks from l, assuming we can take
            # slices of the iterator
            for chunk_slice in helpers.chunks(list(range(len(pts))), chunksize):
                chunk_pts = pts[chunk_slice]
                yield osl.gen_spectral_grid_from_given_points(chunk_pts, bounds=bounds)
        except Exception as e:
            # chunks may not work on this as pts is most likely a Table
            print(e)
            for chunk_pts in helpers.chunks(pts, chunksize):
                yield osl.gen_spectral_grid_from_given_points(chunk_pts, bounds=bounds)

def _make_dust_fA_valid_points_generator(it, min_Rv, max_Rv):
    """
    compute the allowed points based on the R(V) versus f_A plane
    duplicates effort for all A(V) values, but it is quick compared to
    other steps

    .. note::

        on 2.74: SMC extinction implies f_A = 0. and Rv = 2.74

    Parameters
    ----------
    it: an iterable
        an initial sequence of points that will be trimmed to only valid ones

    min_Rv: float
        lower Rv limit

    max_Rv: float
        upper Rv limit

    Returns
    -------
    npts: int
        the actual number of valid points

    pts: generator
        a generator that only produce valid points
    """
    itn = copy.copy(it)
    npts = 0

    def is_valid(ak, rk, fk):
        return (
            fk / max_Rv + (1.0 - fk) / 2.74
            <= 1.0 / rk
            <= fk * 1.0 / min_Rv + (1.0 - fk) / 2.74
        )

    # explore the full list once
    # not very time consuming
    for ak, rk, fk in itn:
        if is_valid(ak, rk, fk):
            npts += 1

    # make the iterator
    pts = (
        (float(ak), float(rk), float(fk)) for ak, rk, fk in it if is_valid(ak, rk, fk)
    )

    return npts, pts

def apply_distance_grid(specgrid, distances, redshift=0):
    """
    Distances are applied to the spectral grid by copying the grid and
    applying a scaling factor.

    Parameters
    ----------

    project: str
        project name

    specgrid: grid.SpectralGrid object
        spectral grid to transform

    distances: list of float
        Distances at which models should be shifted
        0 means absolute magnitude.
        Expecting pc units

    redshift: float
        Redshift to which wavelengths should be shifted
        Default is 0 (rest frame)
    """
    g0 = specgrid

    # Current length of the grid
    N0 = len(g0.grid)
    N = N0 * len(distances)

    # Make singleton list if a single distance is given
    if not hasattr(distances, "__iter__"):
        _distances = [distances]
    else:
        _distances = distances

    # Add distance column if multiple distances are specified
    cols = {}
    cols["distance"] = np.empty(N, dtype=float)

    # Existing columns
    keys0 = list(g0.keys())
    for key in keys0:
        cols[key] = np.empty(N, dtype=float)

    n_sed_points = g0.seds.shape[1]
    new_seds = np.empty((N, n_sed_points), dtype=float)

    for count, distance in enumerate(tqdm(_distances, desc="Distance grid")):

        # The range where the current distance points will live
        distance_slice = slice(N0 * count, N0 * (count + 1))

        # The seds default to 10 pc.
        # Therefore, scale them with (d / (10 pc))**(-2).
        distance_pc = distance.to(units.pc).value
        new_seds[distance_slice, :] = g0.seds / (0.1 * distance_pc) ** 2

        # Fill in the distance in the distance column
        cols["distance"][distance_slice] = distance_pc

        # Copy the old columns
        for key in keys0:
            cols[key][distance_slice] = g0.grid[key]

    # apply redshift
    g0.lamb = g0.lamb * (1.0 + redshift)

    # New object
    g = SpectralGrid(g0.lamb, seds=new_seds, grid=Table(cols), backend="memory")
    return g

@generator
def make_extinguished_grid(
    spec_grid,
    filter_names,
    extLaw,
    avs,
    rvs,
    fAs=None,
    av_prior_model={"name": "flat"},
    rv_prior_model={"name": "flat"},
    fA_prior_model={"name": "flat"},
    chunksize=0,
    add_spectral_properties_kwargs=None,
    absflux_cov=False,
    filterLib=None,
):
    """
    Extinguish spectra and extract an SEDGrid through given series of filters
    (all wavelengths in stellar SEDs and filter response functions are assumed
    to be in Angstroms)

    Parameters
    ----------
    spec_grid: string or grid.SpectralGrid
        if string:
        spec_grid is the filename to the grid file with stellar spectra
        the backend to load this grid will be the minimal invasive: 'HDF'
        if possible, 'cache' otherwise.

        if not a string, expecting the corresponding SpectralGrid instance
        (backend already setup)

    filter_names: list
        list of filter names according to the filter lib

    Avs: sequence
        Av values to iterate over

    av_prior_model: list
        list including prior model name and parameters

    Rvs: sequence
        Rv values to iterate over

    rv_prior_model: list
        list including prior model name and parameters

    fAs: sequence (optional)
        f_A values to iterate over
        f_A can be omitted if the extinction Law does not use it or allow
        fixed values

    fA_prior_model: list
        list including prior model name and parameters

    chunksize: int, optional (default=0)
        number of extinction model variations to generate at each cycle.
        Note that this means len(spec_grid * chunksize)
        If default <= 0, all models will be returned at once.

    filterLib:  str
        full filename to the filter library hd5 file

    add_spectral_properties_kwargs: dict
        keyword arguments to call :func:`add_spectral_properties` at each
        iteration to add model properties from the spectra into the grid
        property table

    asbflux_cov: boolean
        set to calculate the absflux covariance matrices for each model
        (can be very slow!!!  But it is the right thing to do)

    Returns
    -------
    g: grid.SpectralGrid
        final grid of reddened SEDs and models
    """
    # Check inputs
    # ============
    # get the stellar grid (no dust yet)
    # if string is provided try to load the most memory efficient backend
    # otherwise use a cache-type backend (load only when needed)
    if isinstance(spec_grid, str):
        ext = spec_grid.split(".")[-1]
        if ext in ["hdf", "hd5", "hdf5"]:
            g0 = SpectralGrid(spec_grid, backend="disk")
        else:
            g0 = SpectralGrid(spec_grid, backend="cache")
    else:
        helpers.type_checker("spec_grid", spec_grid, SpectralGrid)
        g0 = spec_grid

    # Tag fA usage
    if fAs is None:
        with_fA = False
    else:
        with_fA = True

    # get the min/max R(V) values necessary for the grid point definition
    min_Rv = min(rvs)
    max_Rv = max(rvs)

    # Create the sampling mesh
    # ========================
    # basically the dot product from all input 1d vectors
    # setup interation over the full dust parameter grid
    if with_fA:
        dustpriors = PriorWeightsDust(
            avs, av_prior_model, rvs, rv_prior_model, fAs, fA_prior_model
        )

        it = np.nditer(np.ix_(avs, rvs, fAs))
        niter = np.size(avs) * np.size(rvs) * np.size(fAs)
        npts, pts = _make_dust_fA_valid_points_generator(it, min_Rv, max_Rv)

        # Pet the user
        print(
            """number of initially requested points = {0:d}
              number of valid points = {1:d} (based on restrictions in R(V)
                 versus f_A plane)
              """.format(
                niter, npts
            )
        )

        if npts == 0:
            raise AttributeError("No valid points")
    else:
        dustpriors = PriorWeightsDust(
            avs, av_prior_model, rvs, rv_prior_model, [1.0], fA_prior_model
        )

        it = np.nditer(np.ix_(avs, rvs))
        npts = np.size(avs) * np.size(rvs)
        pts = ((float(ak), float(rk)) for ak, rk in it)

    # Generate the Grid
    # =================
    N0 = len(g0.grid)
    N = N0 * npts

    if chunksize <= 0:
        print("Generating a final grid of {0:d} points".format(N))
    else:
        print(
            "Generating a final grid of {0:d} points in {1:d}"
            + " pieces".format(N, int(float(N0) / chunksize + 1.0))
        )

    if chunksize <= 0:
        chunksize = npts

    if add_spectral_properties_kwargs is not None:
        nameformat = add_spectral_properties_kwargs.pop("nameformat", "{0:s}") + "_wd"

    for chunk_pts in helpers.chunks(pts, chunksize):
        # iter over chunks of models

        # setup chunk outputs
        cols = {"Av": np.empty(N, dtype=float), "Rv": np.empty(N, dtype=float)}

        if with_fA:
            cols["Rv_A"] = np.empty(N, dtype=float)
            cols["f_A"] = np.empty(N, dtype=float)

        keys = list(g0.keys())
        for key in keys:
            cols[key] = np.empty(N, dtype=float)

        n_filters = len(filter_names)
        _seds = np.empty((N, n_filters), dtype=float)
        if absflux_cov:
            n_offdiag = ((n_filters ** 2) - n_filters) / 2
            _cov_diag = np.empty((N, n_filters), dtype=float)
            _cov_offdiag = np.empty((N, n_offdiag), dtype=float)

        for count, pt in enumerate(tqdm(chunk_pts, desc="SED grid")):

            if with_fA:
                Av, Rv, f_A = pt
                dust_prior_weight = dustpriors.get_weight(Av, Rv, f_A)
                Rv_MW = extLaw.get_Rv_A(Rv, f_A)
                r = g0.applyExtinctionLaw(extLaw, Av=Av, Rv=Rv, f_A=f_A, inplace=False)
                # add extra "spectral bands" if requested
                if add_spectral_properties_kwargs is not None:
                    r = add_spectral_properties(
                        r,
                        nameformat=nameformat,
                        filterLib=filterLib,
                        **add_spectral_properties_kwargs
                    )
                temp_results = r.getSEDs(filter_names, filterLib=filterLib)
                # adding the dust parameters to the models
                cols["Av"][N0 * count : N0 * (count + 1)] = Av
                cols["Rv"][N0 * count : N0 * (count + 1)] = Rv
                cols["f_A"][N0 * count : N0 * (count + 1)] = f_A
                cols["Rv_A"][N0 * count : N0 * (count + 1)] = Rv_MW

            else:
                Av, Rv = pt
                dust_prior_weight = dustpriors.get_weight(Av, Rv, 1.0)
                r = g0.applyExtinctionLaw(extLaw, Av=Av, Rv=Rv, inplace=False)

                if add_spectral_properties_kwargs is not None:
                    r = add_spectral_properties(
                        r,
                        nameformat=nameformat,
                        filterLib=filterLib,
                        **add_spectral_properties_kwargs
                    )
                temp_results = r.getSEDs(filter_names, filterLib=filterLib)
                # adding the dust parameters to the models
                cols["Av"][N0 * count : N0 * (count + 1)] = Av
                cols["Rv"][N0 * count : N0 * (count + 1)] = Rv

            # get new attributes if exist
            for key in list(temp_results.grid.keys()):
                if key not in keys:
                    k1 = N0 * count
                    k2 = N0 * (count + 1)
                    cols.setdefault(key, np.empty(N, dtype=float))[
                        k1:k2
                    ] = temp_results.grid[key]

            # compute the fractional absflux covariance matrices
            if absflux_cov:
                absflux_covmats = calc_absflux_cov_matrices(
                    r, temp_results, filter_names
                )
                _cov_diag[N0 * count : N0 * (count + 1)] = absflux_covmats[0]
                _cov_offdiag[N0 * count : N0 * (count + 1)] = absflux_covmats[1]

            # assign the extinguished SEDs to the output object
            _seds[N0 * count : N0 * (count + 1)] = temp_results.seds[:]

            # copy the rest of the parameters
            for key in keys:
                cols[key][N0 * count : N0 * (count + 1)] = g0.grid[key]

            # multiply existing prior weights by the dust prior weight
            cols["weight"][N0 * count : N0 * (count + 1)] *= dust_prior_weight
            cols["prior_weight"][N0 * count : N0 * (count + 1)] *= dust_prior_weight

            if count == 0:
                cols["lamb"] = temp_results.lamb[:]

        _lamb = cols.pop("lamb")

        # free the memory of temp_results
        # del temp_results
        # del tempgrid

        # Ship
        if absflux_cov:
            g = SEDGrid(
                _lamb,
                seds=_seds,
                cov_diag=_cov_diag,
                cov_offdiag=_cov_offdiag,
                grid=Table(cols),
                backend="memory",
            )
        else:
            g = SEDGrid(_lamb, seds=_seds, grid=Table(cols), backend="memory")

        g.header["filters"] = " ".join(filter_names)

        yield g

def add_spectral_properties(
    specgrid,
    filternames=None,
    filters=None,
    callables=None,
    nameformat=None,
    filterLib=None,
):
    """
    Addon spectral calculations to spectral grids to extract in the fitting
    routines

    Parameters
    ----------
    specgrid: SpectralGrid instance
        instance of the spectral grid

    filternames: sequence(str)
        compute the integrated values through given filters in the library

    filters: sequence(Filters)
        sequence of filter instances from which extract integrated values

    callables: sequence(callable)
        sequence of functions to apply onto the spectral grid assuming storing
        results is internally processed by the individual functions

    nameformat: str
        naming format to adopt for filternames and filters
        default value is '{0:s}_0' where the value will be the filter name

    filterLib:  str
        full filename to the filter library hd5 file

    Returns
    -------
    specgrid: SpectralGrid instance
        instance of the input spectral grid which will include more properties
    """
    if nameformat is None:
        nameformat = "{0:s}_0"

    if filternames is not None:
        temp = specgrid.getSEDs(filternames, extLaw=None, filterLib=filterLib)

        logtempseds = np.array(temp.seds)
        indxs = np.where(temp.seds > 0)
        if len(indxs) > 0:
            logtempseds[indxs] = np.log10(temp.seds[indxs])
        indxs = np.where(temp.seds <= 0)
        if len(indxs) > 0:
            logtempseds[indxs] = -100.0

        for i, fk in enumerate(filternames):
            specgrid.grid["log" + nameformat.format(fk)] = logtempseds[:, i]
        del temp

    if filters is not None:
        temp = specgrid.getSEDs(filters, extLaw=None)

        logtempseds = np.array(temp.seds)
        indxs = np.where(temp.seds > 0)
        if len(indxs) > 0:
            logtempseds[indxs] = np.log10(temp.seds[indxs])

        indxs = np.where(temp.seds <= 0)
        if len(indxs) > 0:
            logtempseds[indxs] = -100.0

        for i, fk in enumerate(filters):
            specgrid.grid["log" + nameformat.format(fk.name)] = logtempseds[:, i]
        del temp

    if callables is not None:
        for fn in callables:
            fn(specgrid)

    return specgrid

def calc_absflux_cov_matrices(specgrid, sedgrid, filter_names):
    """ Calculate the absflux covariance matrices for each model
    Must be done on the full spectrum of each model to account for
    the changing combined spectral response due to the model SED and
    the filter response curve.

    Parameters
    ----------
    specgrid: SpectralGrid instance
        instance of the spectral grid containing the full spectrum fluxes

    sedgrid: SpectralGrid instance
        instance of the spectral grid containing the band SED fluxes

    Returns
    -------
    absflux_covmat :
    """

    # get the fractional absflux covariance matrix
    absflux_cov_mats = absflux_covmat.hst_frac_matrix(
        filter_names, spectrum=(specgrid.lamb[:], specgrid.seds)
    )

    # setup the output quantities
    n_models = specgrid.seds.shape[0]
    n_filters = len(filter_names)
    n_offdiag = ((n_filters ** 2) - n_filters) / 2
    cov_diag = np.empty((n_models, n_filters), dtype=np.float64)
    cov_offdiag = np.empty((n_models, n_offdiag), dtype=np.float64)

    # pack the resulting covariance matrices into diganonal and
    # non-diagnonal terms
    #   much more efficient for use later in combining with AST results
    #     and fitting
    #   also convert from fractional to physical flux units
    m = 0
    cov_diag[:, n_filters - 1] = absflux_cov_mats[
        :, n_filters - 1, n_filters - 1
    ] * np.square(sedgrid.seds[:, n_filters - 1])
    for k in range(n_filters - 1):
        cov_diag[:, k] = absflux_cov_mats[:, k, k] * np.square(sedgrid.seds[:, k])
        for l in range(k + 1, n_filters):
            cov_offdiag[:, m] = (
                absflux_cov_mats[:, k, l] * sedgrid.seds[:, k] * sedgrid.seds[:, l]
            )
            m += 1

    return (cov_diag, cov_offdiag)


#!/usr/bin/env python3

import fileinput
import numpy as np
import re
import os
from collections import defaultdict
import argparse
import sys
from pathlib import Path

class Atom:
    S = ''
    x = np.array([])

    def __init__(self, string):
        inp = string.split()
        self.S = inp[0]
        self.x = np.array(list(map(float, inp[1:])))

    def __repr__(self):
        return self.S + " " + " ".join(map(str, self.x))

    def scale(self, a):
        self.x *= a

def getMol(qcfile):
    m = {'molecule': re.compile('^[$]molecule', re.IGNORECASE),
         'end':      re.compile('^[$]end', re.IGNORECASE),
         'bohr':     re.compile('input_bohr.*true', re.IGNORECASE),
         }

    bohr = False
    molecule = []

    # find start of the molecule section
    for line in qcfile:
        if m['bohr'].match(line):
            bohr = True
        if m['molecule'].match(line):
            break

    next(qcfile)  # skip over comment line

    # parse in molecules
    for line in qcfile:
        if m['end'].match(line):
            break

        molecule.append(Atom(line))

    # keep looking for a unit conversion factor
    for line in qcfile:
        if m['bohr'].match(line):
            bohr = True

    scale = 0.10  # \AA -> nm
    if bohr:
        scale *= 0.529177210903  # a_0 -> \AA

    for a in molecule:
        a.scale(scale)

    return molecule

def writeGRO(atoms, name="GIFS", residue="QM", output=sys.stdout):
    """gro file format:

    title string (free format string, optional time in ps after ‘t=’)
    number of atoms (free format integer)
    one line for each atom (fixed format, see below)
    box vectors (free format, space separated reals), values: v1(x) v2(y) v3(z)

    [atoms]
    residue number (5 positions, integer)
    residue name (5 characters)
    atom name (5 characters)
    atom number (5 positions, integer)
    position (in nm, x y z in 3 columns: 8 positions with 3 decimal places)
    velocity (in nm/ps, x y z in 3 columns: 8 positions with 4 decimal places)
    fmt="%5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f"

    N.B.: We do not inclued velocities
    """

    print(name, file=output)
    print(len(atoms), file=output)

    for (i, a) in enumerate(atoms):
        print("%5d%-5s%5s%5d%8.3f%8.3f%8.3f" %
              (1, residue, a.S, i, a.x[0], a.x[1], a.x[2]), file=output)

    print("0 0 0", file=output)

def get_block(block_name, f):
    block = []
    start = re.compile('^\[ *' + block_name + ' *\].*', re.IGNORECASE|re.MULTILINE)
    stop = re.compile('^[\[#].*', re.MULTILINE)
    copy = False
    with open(f) as lines:
        for line in lines:
            line = line.strip()
            if copy and stop.search(line) is not None:
                copy = False
                break
            if start.search(line) is not None:
                copy = True

            if copy:
                block.append(line.strip())

    return block

def get_atomtypes(atoms, f):
    ats = set([ ffDict[s2z[atom.S]] for atom in atoms])
    output = []
    with open(f) as lines:
        for line in lines:
            if ats & set(line.split()):
                output.append(line.strip())
    output.append('\n')
    return output

def writeTOP(atoms, name="GIFS", system="GIFS", residue="QM", forcefield=None, noincludes=False, output=sys.stdout):
    # FIXME: want to be able to choose the water topology

    top = []

    if noincludes:
        top += get_block('defaults', ffDict['path'] / "forcefield.itp")

        top.append('\n[ atomtypes ]')
        top.append('; name bond_type z mass charge ptype sigma espilon')
        top += get_atomtypes(atoms, ffDict['path'] / "ffnonbonded.itp")

    else:
        top.append('; Include forcefield parameters')
        top.append(f'#include "{forcefield}/forcefield.itp"')

    top.append("\n[ moleculetype ]")
    top.append("; Name         nrexcl")
    top.append(f"{name}        3")

    top.append("\n[ atoms ]")
    top.append(";  nr type         resnr residue atom      cgnr   charge       mass")

    for (i, a) in enumerate(atoms, start=1):
        top.append(fmtTopAtoms(i, a, residue))

    top.append("\n[ system ]")
    top.append("; Name")
    top.append(f"QMMM {system}")
    top.append("\n[ molecules ]")
    top.append("; Compound        #mols")
    top.append(f"{name}        1")

    print("\n".join(top), file=output)

def fmtTopAtoms(nr, atom, residue, resnr=1):
    #  nr  type       resnr residue  atom    cgnr     charge   mass
    #  1   opls_287      1    GLY      N      1
    #  2   opls_290      1    GLY     H1      1
    # ...
    # or:
    #  1   opls_287      1    GLY      N      1       -0.3    14.0067
    #  2   opls_290      1    GLY     H1      1       0.33      1.008
    # ...

    # Notes: Charge and mass are optional. We do not perturb mass, but
    # set the charge to 0. The name in the atom field must match that
    # defined in type. Charge groups, cgnr, *were previously*
    # constructed such that there are no more than 6 atoms in a single
    # charge group (Gromacs' max is 32). However, this produced errors
    # (blowing-up) with some solvents. Following qforce [1], we simply
    # place each atom in its own charge group.

    # [1]: https://github.com/selimsami/qforce/blob/master/qforce/forcefield.py#L168

    sym = ffDict[s2z[atom.S]]
    cgnr = nr
    chg = 0.0
    return f"{nr:5} {sym:12} {resnr:5} {residue:7} {atom.S:5} {cgnr:5} {chg:8.3}"

def genffD(forcefield):

    ff = Path(forcefield)
    if not ff.exists():
        ff = Path(os.environ['GMXDATA']) / "top" / (forcefield)

    ffDict['path']=ff

    with open (ff / "ffnonbonded.itp") as f:
        o = 0  # special case offset for OPLS-AA
        if forcefield == "oplsaa.ff":
            o = 1
        for line in f:
            field = line.split()
            if len(field) > 4 + o and field[4+o] == 'A':
                ffDict[int(field[1+o])] = field[0]

def main():
    args = getArgs()
    genffD(args.ff)

    print("reading", args.f)
    with open(args.f) as f:
        molecule = getMol(f)

    with open(args.o + ".gro", 'w') as f:
        writeGRO(molecule, name=args.n, residue=args.r, output=f)
    print("wrote", args.o + ".gro")

    with open(args.o + ".top", 'w') as f:
        writeTOP(molecule, name=args.n, system=args.s, residue=args.r, forcefield=args.ff, noincludes=args.noincludes, output=f)
    print("wrote", args.o + ".top")

    print("Be sure to check your topology file for missing atoms")

def getArgs():
    parser = argparse.ArgumentParser(description='Convert a Q-Chem input file for use with Gromacs under GIFS')
    parser.add_argument("-f", metavar="input", required=True,
                        help="name of Q-Chem input file")
    parser.add_argument("-o", metavar="output", required=True,
                        help="name of output .gro/.top files")
    parser.add_argument("-r", metavar="residue", default="QM",
                        help="name for residue; defaults to QM")
    parser.add_argument("-n", metavar="name", default="GIFS",
                        help="molecule name; defaults to GIFS")
    parser.add_argument("-s", metavar="name", default="GIFS",
                        help="system name; defaults to GIFS")
    parser.add_argument("--ff", metavar="forcefield", default="oplsaa.ff",
                        help="forcefield to use; defaults to oplsaa.ff")
    parser.add_argument("--noincludes", default=False, action='store_true',
                        help="generate self-contained topology with no includes")
    return parser.parse_args()

ffDict = defaultdict(lambda: print("Warning: unknown atom in forcefield", file=sys.stderr) or "ff_???")

s2z = {
    'H':   1, 'He':  2, 'Li':  3, 'Be':  4, 'B':   5, 'C':   6,
    'N':   7, 'O':   8, 'F':   9, 'Ne': 10, 'Na': 11, 'Mg': 12,
    'Al': 13, 'Si': 14, 'P':  15, 'S':  16, 'Cl': 17, 'Ar': 18,
    'K':  19, 'Ca': 20, 'Sc': 21, 'Ti': 22, 'V':  23, 'Cr': 24,
    'Mn': 25, 'Fe': 26, 'Co': 27, 'Ni': 28, 'Cu': 29, 'Zn': 30,
    'Ga': 31, 'Ge': 32, 'As': 33, 'Se': 34, 'Br': 35, 'Kr': 36,
    'Rb': 37, 'Sr': 38, 'Y':  39, 'Zr': 40, 'Nb': 41, 'Mo': 42,
    'Tc': 43, 'Ru': 44, 'Rh': 45, 'Pd': 46, 'Ag': 47, 'Cd': 48,
    'In': 49, 'Sn': 50, 'Sb': 51, 'Te': 52, 'I':  53, 'Xe': 54,
    'Cs': 55, 'Ba': 56, 'La': 57, 'Ce': 58, 'Pr': 59, 'Nd': 60,
    'Pm': 61, 'Sm': 62, 'Eu': 63, 'Gd': 64, 'Tb': 65, 'Dy': 66,
    'Ho': 67, 'Er': 68, 'Tm': 69, 'Yb': 70, 'Lu': 71, 'Hf': 72,
    'Ta': 73, 'W':  74, 'Re': 75, 'Os': 76, 'Ir': 77, 'Pt': 78,
    'Au': 79, 'Hg': 80, 'Tl': 81, 'Pb': 82, 'Bi': 83, 'Po': 84,
    'At': 85, 'Rn': 86, 'Fr': 87, 'Ra': 88, 'Ac': 89, 'Th': 90,
    'Pa': 91, 'U':  92, 'Np': 93, 'Pu': 94, 'Am': 95, 'Cm': 96,
    'Bk': 97, 'Cf': 98, 'Es': 99, 'Fm': 100, 'Md': 101, 'No': 102,
    'Lr': 103, 'Rf': 104, 'Db': 105, 'Sg': 106, 'Bh': 107, 'Hs': 108,
    'Mt': 109, 'Ds': 110, 'Rg': 111, 'Cn': 112, 'Nh': 113, 'Fl': 114,
    'Mc': 115, 'Lv': 116, 'Ts': 117, 'Og': 118,
}

z2s = {v:k for k,v in s2z.items()}

if __name__ == "__main__":
    main()
    # genffD("oplsaa.ff")
    # buildffDict()
    # print(ffDict)