#!/usr/bin/env python3

import argparse
from collections import Counter
from multiprocessing import set_start_method
import pdb
import re
import sys
import time
import copy

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
import torch.multiprocessing as mp

import data_producer

from word2atoms import skipgram_atoms, trivial_atoms, morpheme_split

parser = argparse.ArgumentParser()
parser.add_argument("--train", type=str, default="", help="training file")
parser.add_argument("--output", type=str, default="vectors.txt", help="output word embedding file")
parser.add_argument("--atomoutput", type=str, default="atomvectors.txt", help="output atom embedding file")
parser.add_argument("--ctxoutput", type=str, default="atomvectors.txt", help="context vectors file")
parser.add_argument("--ctxatomoutput", type=str, default="atomvectors.txt", help="context atom vectors file")
parser.add_argument("--losslog", type=str, default="all_losses.txt", help="log of training losses")
parser.add_argument("--opt", type=str, default="SGD", help="optimiser to use")
parser.add_argument("--size", type=int, default=300, help="word embedding dimension")
parser.add_argument("--cbow", type=int, default=1, help="1 for cbow, 0 for skipgram")
parser.add_argument("--window", type=int, default=5, help="context window size")
parser.add_argument("--sample", type=float, default=1e-4, help="subsample threshold")
parser.add_argument("--negative", type=int, default=10, help="number of negative samples")
parser.add_argument("--min_count", type=int, default=5, help="minimum frequency of a word")
parser.add_argument("--processes", type=int, default=4, help="number of processes")
parser.add_argument("--num_workers", type=int, default=6, help="number of workers for data processsing")
parser.add_argument("--iter", type=int, default=5, help="number of iterations")
parser.add_argument("--lr", type=float, default=-1.0, help="initial learning rate")
parser.add_argument("--momentum", type=float, default=0.0, help="momentum")
parser.add_argument("--batch_size", type=int, default=100, help="(max) batch size")
parser.add_argument("--megabatch_size", type=int, default=100000, help="(max) megabatch size")
parser.add_argument("--cuda", action='store_true', default=False, help="enable cuda")
parser.add_argument("--output_ctx", action='store_true', default=False, help="output context embeddings")
parser.add_argument("--anneal", action='store_true', default=False, help="anneal the learning rate linearly to 0")
parser.add_argument("--shuffle", action='store_true', default=False, help="shuffle the training data points")
parser.add_argument("--atomizer", type=str, choices=['fasttext', 'morphoseg', 'word2vec'], default="word2vec", help="atomizer to use (vanilla word2vec by default, can also choose fasttext or morphoseg)")
parser.add_argument("--minL", type=int, default=5, help="minimum possible length of n-grams to take when atomizing")
parser.add_argument("--maxL", type=int, default=5, help="maximum possible length of n-grams to take when atomizing")
parser.add_argument("--halfletters", action='store_true', default=False, help="whether to use half-letters/raw Unicode characters when taking n-grams, or whole Tamil letters")

MAX_SENT_LEN = 1000

# Build the vocabulary.
def file_split(f, delim=' \t\n', bufsize=1024):
    prev = ''
    while True:
        s = f.read(bufsize)
        if not s:
            break
        tokens = re.split('['+delim+']{1,}', s)
        if len(tokens) > 1:
            yield prev + tokens[0]
            prev = tokens[-1]
            for x in tokens[1:-1]:
                yield x
        else:
            prev += s
    if prev:
        yield prev

def build_vocab(args):
    vocab = Counter()
    word_count = 0
    for word in file_split(open(args.train)):
        vocab[word] += 1
        word_count += 1
        if word_count % 10000 == 0:
            sys.stdout.write('%d\r' % len(vocab))
    freq = {k:v for k,v in vocab.items() if v >= args.min_count}
    word_count = sum([freq[k] for k in freq])
    word_list = sorted(freq, key=freq.get, reverse=True)
    word2idx = {}
    for i,w in enumerate(word_list):
        word2idx[w] = i

    print("Vocab size: %ld" % len(word2idx))
    print("Words in train file: %ld" % word_count)
    vars(args)['vocab_size'] = len(word2idx)
    vars(args)['train_words'] = word_count

    """ for i, word in enumerate(word_list):
        if word[-1] == '0':
            print(i, word)
        if i == 1761:
            print(word) """
    #print("Initial mattrum count")
    #print(freq['மற்றும்'])
    return word2idx, word_list, freq

def build_morph(args, word2idx, word_list, freq, atomiser):
    morph_list = []
    morph2idx = {}
    word2morph = []
    #total_morphperword = 0
    max_morphperword = 0
    #total_freq = 0
    for i, word in enumerate(word_list):
        idxs = []
        cnt = 0
        for morph in atomiser(word):
            if morph in morph2idx:
                idx = morph2idx[morph]
            else:
                idx = len(morph_list)
                morph_list.append(morph)
                morph2idx[morph] = idx
            idxs.append(idx)
            cnt += 1
        #total_morphperword += freq[word] * cnt
        #total_freq += freq[word]
        if cnt > max_morphperword:
            max_morphperword = cnt
        word2morph.append(torch.LongTensor(idxs))
        if i % 100 == 0:
            sys.stdout.write('%d\r' % i)

    # actual padding index
    word2morph.append(torch.LongTensor([len(morph2idx)]))

    word2morphfinal = torch.zeros((len(word2morph), max_morphperword), dtype=torch.long)
    word2morphfinal = word2morphfinal.fill_(len(morph2idx) +1)
    for i in range(len(word2morph)):
        row = word2morph[i]
        word2morphfinal[i, :row.shape[0]] = row

    #print(word2morphfinal.shape)
    #print(freq[word_list[0]])
    """ indices = [0]
    for index in indices:
        print(word_list[index])
        print(' '.join([morph_list[j] for j in word2morph[index]]))
    print(word2morph[816])
    print(word2morph[0])
    print(word2morph[10520]) """

    print("Morpheme size: %ld" % len(morph2idx))
    #print("Average morphemes per word: %f" % (float(total_morphperword)/total_freq))
    print("Max morphemes per word: %d" % max_morphperword)
    #vars(args)['morph_size'] = len(morph2idx)
    #vars(args)['word2morph'] = word2morph

    return morph2idx, morph_list, word2morphfinal

class CBOWMean(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, lens):
        ctx.save_for_backward(x)
        x = torch.sum(x, 1, keepdim=True)
        x = x.permute(1,2,0) / lens
        return x.permute(2,0,1)
    @staticmethod
    def backward(ctx, g):
        x, = ctx.saved_variables
        return g.expand_as(x), None

class CBOW(nn.Module):
    def __init__(self, args):
        super(CBOW, self).__init__()
        self.emb0_lookup = nn.Embedding(args.vocab_size+1, args.size, padding_idx=args.vocab_size, sparse=True)
        self.emb1_lookup = nn.Embedding(args.vocab_size, args.size, sparse=True)
        self.emb0_lookup.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        self.emb0_lookup.weight.data[args.vocab_size].fill_(0)
        self.emb1_lookup.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        self.window = args.window
        self.negative = args.negative
        self.pad_idx = args.vocab_size

    def forward(self, data):
        ctx_indices = data[:, 0:2*self.window]
        ctx_lens = data[:, 2*self.window].float()
        word_idx = data[:, 2*self.window+1]
        neg_indices = data[:, 2*self.window+2:2*self.window+2+self.negative]
        neg_mask = data[:, 2*self.window+2+self.negative:].float()

        c_embs = self.emb0_lookup(ctx_indices)
        w_embs = self.emb1_lookup(word_idx)
        n_embs = self.emb1_lookup(neg_indices)

        c_embs = CBOWMean.apply(c_embs, ctx_lens)

        pos_ips = torch.sum(c_embs[:,0,:] * w_embs, 1)
        neg_ips = torch.bmm(n_embs, c_embs.permute(0,2,1))[:,:,0]

        # Neg Log Likelihood
        pos_loss = torch.sum( -F.logsigmoid(torch.clamp(pos_ips,max=10,min=-10)) )
        neg_loss = torch.sum( -F.logsigmoid(torch.clamp(-neg_ips,max=10,min=-10)) * neg_mask )

        return pos_loss + neg_loss

class SG(nn.Module):
    def __init__(self, args):
        super(SG, self).__init__()
        self.emb0morph_lookup = nn.Embedding(args.morph_size+2, args.size, padding_idx=args.morph_size, sparse=True)
        self.emb1morph_lookup = nn.Embedding(args.ctxmorph_size+2, args.size, padding_idx=args.ctxmorph_size, sparse=True)

        #self.emb0_lookup = nn.Embedding(args.vocab_size+1, args.size, padding_idx=args.vocab_size, sparse=True)
        #self.emb1_lookup = nn.Embedding(args.vocab_size, args.size, sparse=True)

        self.emb0morph_lookup.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        self.emb0morph_lookup.weight.data[args.morph_size+1].fill_(0)
        # randomly initialise context vectors
        #self.emb1morph_lookup.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        #self.emb1morph_lookup.weight.data[args.ctxmorph_size+1].fill_(0)

        # OR zero initialise them as usual
        self.emb1morph_lookup.weight.data.zero_()

        #self.emb0_lookup.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        #self.emb1_lookup.weight.data.zero_()
        #self.emb1morph_lookup.weight.data.zero_()

        #self.emb0morph_lookup.weight.data[args.morph_size+1].requires_grad = False

        self.window = args.window
        self.negative = args.negative
        self.pad_idx = args.vocab_size
        self.morph_size = args.morph_size

    def forward(self, data, word2morph, word2morph_mask, ctx2morph, ctx2morph_mask):
        word_idx = data[:, 0]
        ctx_idx = data[:, 1]
        neg_indices = data[:, 2:2+self.negative]
        neg_mask = data[:, 2+self.negative:].float()

        #print(word2morph.shape)
        #print(word2morph_mask.shape)

        #print(torch.max(word2morph))
        #print("MORPHEMES")
        #print(word2morph[3])
        #print(torch.norm(self.emb0morph_lookup.weight.data[-1]))
        #print(torch.norm(torch.squeeze(word2morph_mask), dim=0)) # should start nonzero, then eventually be 0

        #print((self.emb0morph_lookup(word2morph[:, 0]) * word2morph_mask[:, 0]).shape)
        #t = self.emb0morph_lookup(word2morph) * word2morph_mask
        #print(t.shape)
        #print(torch.sum(t, dim=1).shape)
        w_embs = torch.sum(self.emb0morph_lookup(word2morph) * word2morph_mask, dim=1)
        #w_embs = self.emb0_lookup(word_idx)

        #t = self.emb1morph_lookup(ctx2morph[:, 0]) * ctx2morph_mask[:, 0]
        #print(t.shape)
        #print(torch.sum(t, dim=1).shape)
        #print((self.emb0morph_lookup(word2morph[:, 1]) * word2morph_mask[:, 1]).shape)
        c_embs = torch.sum(self.emb1morph_lookup(ctx2morph[:, 0]) * ctx2morph_mask[:, 0], dim=1)
        #c_embs = self.emb1_lookup(ctx_idx)

        #t = self.emb1morph_lookup(ctx2morph[:, 1:1+self.negative]) * ctx2morph_mask[:, 1:1+self.negative]
        #print(t.shape)
        #print(torch.sum(t, dim=2).shape)
        #print((self.emb0morph_lookup(word2morph[:, 2:2+self.negative]) * word2morph_mask[:, 2:2+self.negative]).shape)
        n_embs = torch.sum(self.emb1morph_lookup(ctx2morph[:, 1:1+self.negative]) * ctx2morph_mask[:, 1:1+self.negative], dim=2)
        #n_embs = self.emb1_lookup(neg_indices)

        pos_ips = torch.sum(w_embs * c_embs, 1)
        neg_ips = torch.bmm(n_embs, torch.unsqueeze(w_embs,1).permute(0,2,1))[:,:,0]

        # Neg Log Likelihood
        pos_loss = torch.sum( -F.logsigmoid(torch.clamp(pos_ips,max=10,min=-10)) )
        neg_loss = torch.sum( -F.logsigmoid(torch.clamp(-neg_ips,max=10,min=-10)) * neg_mask )

        return pos_loss + neg_loss

# Initialize model.
def init_net(args):
    if args.cbow == 1:
        if args.lr == -1.0:
            vars(args)['lr'] = 0.05
        return CBOW(args)
    elif args.cbow == 0:
        if args.lr == -1.0:
            vars(args)['lr'] = 0.025
        return SG(args)

# Training
def train_process_sent_producer(p_id, data_queue, word_count_actual, word2idx, word_list, freq, args):
    if args.negative > 0:
        table_ptr_val = data_producer.init_unigram_table(word_list, freq, args.train_words)

    train_file = open(args.train)
    file_pos = args.file_size * p_id // args.processes
    train_file.seek(file_pos, 0)
    while True:
        try:
            train_file.read(1)
        except UnicodeDecodeError:
            file_pos -= 1
            train_file.seek(file_pos, 0)
        else:
            train_file.seek(file_pos, 0)
            break

    batch_count = 0
    if args.cbow == 1:
        batch_placeholder = np.zeros((args.megabatch_size, 2*args.window+2+2*args.negative), 'int64')
    else:
        batch_placeholder = np.zeros((args.megabatch_size, 2+2*args.negative), 'int64')
    #mattrum_cnt = 0
    for it in range(args.iter):
        train_file.seek(file_pos, 0)

        last_word_cnt = 0
        word_cnt = 0
        sentence = []
        prev = ''
        eof = False
        while True:
            if eof or train_file.tell() > file_pos + args.file_size / args.processes:
                break

            while True:
                s = train_file.read(1)
                if not s:
                    eof = True
                    break
                elif s == ' ' or s == '\t':
                    if prev in word2idx:
                        sentence.append(prev)
                    prev = ''
                    if len(sentence) >= MAX_SENT_LEN:
                        break
                elif s == '\n':
                    if prev in word2idx:
                        sentence.append(prev)
                    prev = ''
                    break
                else:
                    prev += s

            if len(sentence) > 0:
                #print("Full sentence")
                #print(' '.join(sentence))
                # subsampling
                sent_id = []
                trimmed = []
                if args.sample != 0:
                    sent_len = len(sentence)
                    i = 0
                    while i < sent_len:
                        word = sentence[i]
                        f = freq[word] / args.train_words
                        pb = (np.sqrt(f / args.sample) + 1) * args.sample / f

                        if pb > np.random.random_sample():
                            sent_id.append( word2idx[word] )
                            """ if word2idx[word] == 'மற்றும்' and mattrum_cnt % 1000 == 0:
                                print("Hit another 1000 mattrums")
                                mattrum_cnt += 1
                        else:
                            trimmed.append(word) """
                        i += 1

                if len(sent_id) < 2:
                    word_cnt += len(sentence)
                    sentence.clear()
                    continue

                #print("Killed words")
                #print(' '.join(trimmed))
                #print("Trimmed sentence")
                #print(' '.join([word_list[index] for index in sent_id]))

                next_random = (2**24) * np.random.randint(0, 2**24) + np.random.randint(0, 2**24)
                if args.cbow == 1: # train CBOW
                    chunk = data_producer.cbow_producer(sent_id, len(sent_id), table_ptr_val,
                                args.window, args.negative, args.vocab_size, args.batch_size, next_random)
                elif args.cbow == 0: # train skipgram
                    chunk = data_producer.sg_producer(sent_id, len(sent_id), table_ptr_val,
                                args.window, args.negative, args.vocab_size, args.batch_size, next_random)

                #print("Data points")
                #print(chunk)

                chunk_pos = 0
                while chunk_pos < chunk.shape[0]:
                    remain_space = args.megabatch_size - batch_count
                    remain_chunk = chunk.shape[0] - chunk_pos

                    if remain_chunk < remain_space:
                        take_from_chunk = remain_chunk
                    else:
                        take_from_chunk = remain_space

                    batch_placeholder[batch_count:batch_count+take_from_chunk, :] = chunk[chunk_pos:chunk_pos+take_from_chunk, :]
                    batch_count += take_from_chunk

                    if batch_count == args.megabatch_size:
                        if args.shuffle:
                            p = torch.randperm(batch_count)
                            batch_placeholder = batch_placeholder[p]

                        start = 0
                        while start < batch_count:
                            data_queue.put(batch_placeholder[start : min(start + args.batch_size, batch_count)])
                            start += args.batch_size
                        #print("Batch placeholder")
                        #print(batch_placeholder)
                        batch_count = 0

                    chunk_pos += take_from_chunk

                word_cnt += len(sentence)
                if word_cnt - last_word_cnt > 10000:
                    with word_count_actual.get_lock():
                        word_count_actual.value += word_cnt - last_word_cnt
                    last_word_cnt = word_cnt
                sentence.clear()

        with word_count_actual.get_lock():
            word_count_actual.value += word_cnt - last_word_cnt

    #print("Total occurrences of mattrum: " + str(mattrum_cnt))
    #print("Total non-occurrences of mattrum: " + str(non_mattrum_cnt))
    if batch_count > 0:
        if args.shuffle:
            p = torch.randperm(batch_count)
            batch_placeholder[:batch_count] = batch_placeholder[p]

        start = 0
        while start < batch_count:
            data_queue.put(batch_placeholder[start : min(start + args.batch_size, batch_count)])
            start += args.batch_size
        #print("Batch placeholder")
        #print(batch_placeholder)
        batch_count = 0
    data_queue.put(None)

def train_process(p_id, word_count_actual, word2idx, word_list, freq, args, model, word2morph, word2morph_mask, ctx2morph, ctx2morph_mask):
    data_queue = mp.SimpleQueue()

    if args.opt == "Adagrad":
        optimizer = optim.Adagrad(model.parameters(), lr=args.lr)
    elif args.opt == "SGD":
        optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    elif args.opt == 'SparseAdam':
        optimizer = optim.SparseAdam(model.parameters(), lr=args.lr)

    t = mp.Process(target=train_process_sent_producer, args=(p_id, data_queue, word_count_actual, word2idx, word_list, freq, args))
    t.start()

    # get from data_queue and feed to model
    prev_word_cnt = 0
    losses_cnt = 0
    total_loss = 0.0
    losses_file = open(args.losslog, 'w')
    lr = args.lr
    #mattrum_cnt = 0
    #non_mattrum_cnt = 0
    while True:
        d = data_queue.get()
        if d is None:
            break
        else:
            # lr anneal
            if args.anneal:
                if word_count_actual.value - prev_word_cnt > 10000:
                    lr = args.lr * (1 - word_count_actual.value / (args.iter * args.train_words))
                    if lr < 0.0001 * args.lr:
                        lr = 0.0001 * args.lr
                    for param_group in optimizer.param_groups:
                        param_group['lr'] = lr
            else:
                lr = args.lr

            if args.cuda:
                data = Variable(torch.LongTensor(d).cuda(), requires_grad=False)
            else:
                data = Variable(torch.LongTensor(d), requires_grad=False)

            if args.cbow == 1:
                optimizer.zero_grad()
                loss = model(data)
                loss.backward()
                optimizer.step()
                model.emb0_lookup.weight.data[args.vocab_size].fill_(0)
            elif args.cbow == 0:
                optimizer.zero_grad()
                #print("WORD")
                #print(data[3][0])
                loss = model(data, word2morph[data[:, 0]], word2morph_mask[data[:, 0]], ctx2morph[data[:, 1:2+args.negative]], ctx2morph_mask[data[:, 1:2+args.negative]])
                loss.backward()
                #model.emb0morph_lookup.weight.data.grad[args.morph_size+1].fill_(0)
                optimizer.step()
                #model.emb0morph_lookup.weight.data[args.morph_size+1].zero_()

            losses_cnt += data.shape[0]
            total_loss += loss

            # output
            if word_count_actual.value - prev_word_cnt > 10000:
                avg_loss = total_loss/losses_cnt
                sys.stdout.write("\rAlpha: %0.8f, Loss: %0.8f, Progress: %0.2f, Words/sec: %f" % (lr, avg_loss, word_count_actual.value / (args.iter * args.train_words) * 100, word_count_actual.value / (time.monotonic() - args.t_start)))
                sys.stdout.flush()
                prev_word_cnt = word_count_actual.value
                losses_cnt = 0
                total_loss = 0.0
                losses_file.write(str(avg_loss.item()) + '\n')

    losses_file.close()
    t.join()

if __name__ == '__main__':
    set_start_method('forkserver')

    args = parser.parse_args()
    print("Starting training using file %s" % args.train)
    train_file = open(args.train)
    train_file.seek(0, 2)
    vars(args)['file_size'] = train_file.tell()

    word2idx, word_list, freq = build_vocab(args)
    # constructing and applying atomizer to all words
    minL = args.minL
    maxL = args.maxL
    use_listify = not args.halfletters
    if args.atomizer == 'word2vec':
        atomizer = lambda w: trivial_atoms(w)
    elif args.atomizer == 'fasttext':
        atomizer = lambda w: skipgram_atoms(w, minL=minL, maxL=maxL)
    elif args.atomizer == 'morphoseg':
        atomizer = lambda w: morpheme_split(w, minL=minL, maxL=maxL, use_listify=use_listify, to_stem=True)

    morph2idx, morph_list, word2morph = build_morph(args, word2idx, word_list, freq, atomizer)
    vars(args)['morph_size'] = len(morph2idx)

    ctxmorph2idx, ctxmorph_list, ctx2morph = build_morph(args, word2idx, word_list, freq, trivial_atoms)
    vars(args)['ctxmorph_size'] = len(ctxmorph2idx)

    #print(word2morph.shape)
    #print(ctx2morph.shape)

    """ idx = 200
    word = word_list[idx]
    print("Word: " + word)
    print("Index should be: " + str(word2idx[word]))
    all_morphs = word2morph[idx]
    for midx in all_morphs:
        print("Morpheme: " + morph_list[midx])
        print("Index used for lookup: " + str(midx))
        print("Index retrieved: " + str(morph2idx[morph_list[midx]])) """

    word_count_actual = mp.Value('L', 0)

    model = init_net(args)
    model.share_memory()
    word2morph_mask = torch.unsqueeze((word2morph <= args.morph_size), dim=2).type(model.emb0morph_lookup.weight.data.dtype)
    ctx2morph_mask = torch.unsqueeze((ctx2morph <= args.ctxmorph_size), dim=2).type(model.emb1morph_lookup.weight.data.dtype)

    if args.cuda:
        model.cuda()
        word2morph = word2morph.cuda()
        word2morph_mask = word2morph_mask.cuda()
        ctx2morph = ctx2morph.cuda()
        ctx2morph_mask = ctx2morph_mask.cuda()

    #print(word2morph_mask.shape)
    #print(ctx2morph_mask.shape)

    vars(args)['t_start'] = time.monotonic()
    processes = []
    for p_id in range(args.processes):
        p = mp.Process(target=train_process, args=(p_id, word_count_actual, word2idx, word_list, freq, args, model, word2morph, word2morph_mask, ctx2morph, ctx2morph_mask))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    torch.cuda.empty_cache()

    # print out the atom input vectors
    if args.cuda:
        embs = model.emb0morph_lookup.weight.data.cpu().numpy()
    else:
        embs = model.emb0morph_lookup.weight.data.numpy()

    #print(embs.shape)
    data_producer.write_embs(args.atomoutput, morph_list, embs, args.morph_size, args.size)
    del embs
    torch.cuda.empty_cache()

    wordembs = np.zeros((word2morph.shape[0], args.size), dtype=np.float32)
    print(wordembs.shape)
    # word input vectors
    if args.cuda:
        numpieces = 5
        cnt = word2morph.shape[0] // numpieces
        print("Need to handle " + str(word2morph.shape[0]) + " words")
        for j in range(numpieces):
            start = j * cnt
            end = start + cnt
            if j == numpieces - 1:
                end = word2morph.shape[0]
            print("Handling all words from " + str(start) + " to " + str(end))
            wordembs[start:end] = torch.sum(model.emb0morph_lookup(word2morph[start:end]) * word2morph_mask[start:end], dim=1).detach().cpu().numpy()
    else:
        wordembs = torch.sum(model.emb0morph_lookup(word2morph) * word2morph_mask, dim=1).detach().numpy()

    #print(wordembs.shape)
    data_producer.write_embs(args.output, word_list, wordembs, args.vocab_size, args.size)
    del wordembs
    torch.cuda.empty_cache()

    # atom context vectors
    if args.cuda:
        atembs = model.emb1morph_lookup.weight.data.cpu().numpy()
    else:
        atembs = model.emb1morph_lookup.weight.data.numpy()

    #print(atembs.shape)
    data_producer.write_embs(args.ctxatomoutput, ctxmorph_list, atembs, args.ctxmorph_size, args.size)
    del atembs
    torch.cuda.empty_cache()

    # word context vectors
    ctxembs = np.zeros((ctx2morph.shape[0], args.size), dtype=np.float32)
    print(ctxembs.shape)
    if args.cuda:
        numpieces = 5
        cnt = ctx2morph.shape[0] // numpieces
        print("Need to handle " + str(ctx2morph.shape[0]) + " words")
        for j in range(numpieces):
            start = j * cnt
            end = start + cnt
            if j == numpieces - 1:
                end = ctx2morph.shape[0]
            print("Handling all words from " + str(start) + " to " + str(end))
            ctxembs[start:end] = torch.sum(model.emb1morph_lookup(ctx2morph[start:end]) * ctx2morph_mask[start:end], dim=1).detach().cpu().numpy()
    else:
        ctxembs = torch.sum(model.emb1morph_lookup(ctx2morph) * ctx2morph_mask, dim=1).detach().numpy()

    #print(ctxembs.shape)
    data_producer.write_embs(args.ctxoutput, word_list, ctxembs, args.vocab_size, args.size)
    del ctxembs
    torch.cuda.empty_cache()

    print("")


###################################
# Script :
# 1) Contains class to validate models
# built using SAS datasets
#
# ganesans - Salilab - UCSF
# ganesans@salilab.org
###################################
import pandas as pd
import sys,os,math
import numpy as np
import pandas as pd
import re,pickle,requests,json
from sklearn.linear_model import LinearRegression
from decimal import Decimal
from validation import get_input_information
from subprocess import run, call, PIPE
from decouple import config

class sas_validation(get_input_information):
    def __init__(self,mmcif_file):
        super().__init__(mmcif_file)
        self.ID=str(get_input_information.get_id(self))
        self.nos=get_input_information.get_number_of_models(self)
        self.dataset=get_input_information.get_dataset_comp(self)
        self.imagepath='../static/images/'
        self.saslink='https://www.sasbdb.org/media/sascif/sascif_files/'
        self.sasentry='https://www.sasbdb.org/rest-api/entry/summary/'

    def get_SASBDB_code(self)->list:
        '''
        function to get all SASBDB codes used in the model,
        returns a list of SASBDB codes
        '''
        SAS_db_codes=[]
        for indx,datatype in enumerate(self.dataset['Dataset type']):
            if 'SAS' in str(datatype):
                SAS_db_codes.append(self.dataset['Data access code'][indx])
        return SAS_db_codes

    def clean_SASBDB_code(self)->list:
        '''
        function to clean SASBDB list of codes
        as some might have 'None' or can be repetitive
        '''
        codes=list(set(self.get_SASBDB_code()))
        cleaned_code=[i for i in codes if i != 'None']
        return cleaned_code

    def get_data_from_SASBDB(self)->dict:
        '''
        get data from JSON
        '''
        data_dic={}
        for code in self.get_SASBDB_code():
            if 'None' not in str(code):
                url_f=self.sasentry+code+'.json'
                response=requests.get(url_f, data={'key':'value'});
                if response.status_code!=200:
                    print ("Error....unable to fetch data from SASBDB, please check the entry ID")
                data_dic[code]=response.json();
                with open (code+'.json', 'w') as f:
                    formatted_data=json.dumps(response.json(), indent = 4, sort_keys=True);
                    f.write(formatted_data);
        return data_dic

    def get_sascif_file(self):
        '''
        get data from SASCIF files
        '''
        for code in self.get_SASBDB_code():
            if 'None' not in str(code):
                url_f=self.saslink+code+'.sascif'
                response=requests.get(url_f);
                if response.status_code!=200:
                    print ("Error....unable to fetch data from SASBDB, please check the entry ID")
                with open (code+'.sascif', 'w') as f:
                    f.write(response.text);

    def get_all_sascif(self,sasbdb)->list:
        '''
        get a list of all lines in a SASCIF file
        '''
        if 'None' not in str(sasbdb):
            file=open(sasbdb+'.sascif','r')
            all_lines=[data.strip().split() for indx,data in enumerate(file.readlines())]
        return all_lines

    def get_intensities(self)->dict:
        '''
        get intensity data from SASCIF file
        if SASCIF file is not present, this information will not be present/used in the report
        JSON file typically has raw data
        '''
        self.get_sascif_file()
        Int_dict={}
        for code in self.clean_SASBDB_code():
            #Int_dict={}
            all_lines=self.get_all_sascif(code)
            data={}
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<2 and len(sascifline)>0 and 'scan_intensity' in sascifline[0]:
                    data[(sascifline[0].split('.')[1])]=[]
                if len(sascifline)>2 and len(all_lines[indx-1])>0 and'scan_intensity' in all_lines[indx-1][0]:
                    for indx_sub,sascifline_sub in enumerate(all_lines[indx:]):
                        if len(sascifline_sub)>2 and '#' not in sascifline_sub and 'sas' not in sascifline_sub[0] :
                            for num,key in enumerate(list(data.keys())):
                                data[key].append(sascifline_sub[num])
                        else:
                            break
            I_df=pd.DataFrame(list(data.values()),index=list(data.keys())).T
            I_df_re=I_df[['momentum_transfer','intensity','intensity_su_counting']]
            I_df_re.rename(columns={'momentum_transfer':'Q','intensity':'I','intensity_su_counting':'E'},inplace=True)
            Int_dict[code]=I_df_re
        return Int_dict

    def modify_intensity(self)->dict:
        '''
        modify intensity data to calcualte errors and log values
        '''
        Int_dict=self.get_intensities()
        Int_dict_modify={}
        rg_and_io=self.get_rg_and_io()
        for key,val in Int_dict.items():
            Rg=rg_and_io[key][0]
            IO=rg_and_io[key][1]
            dim_num=Rg*Rg/IO
            I_df=val.astype({'Q':float,'I':float,'E':float})
            I_df=I_df[I_df['I']-I_df['E']>0]
            I_df['Q']=I_df['Q']*10
            I_df['err_x']=I_df.apply(lambda row: (row['Q'],row['Q']), axis=1)
            I_df['err_y']=I_df.apply(lambda row: (np.log(row['I']-row['E']),np.log(row['I']+row['E'])),axis=1)
            I_df['logI']=np.log(I_df['I'])
            I_df['logQ']=np.log(I_df['Q'])
            I_df['logX']=I_df.apply(lambda row: (row['logQ'],row['logQ']), axis=1)
            I_df['Ky']=I_df['Q']*I_df['Q']*I_df['I']*dim_num
            I_df['Kx']=I_df['Q']*Rg
            I_df['Px']=I_df['Q']**4
            I_df['Px'].round(3)
            I_df['Py']=I_df['Px']*I_df['I']
            Int_dict_modify[key]=I_df
        return Int_dict_modify

    def modify_intensity_dep(self)->dict:
        '''
        depreciated function to get intensities from JSON/raw data
        '''
        Int_dict=self.get_intensities()
        Int_dict_modify={}
        for key,val in Int_dict.items():
            I_df=val.astype({'Q':float,'I':float,'E':float})
            I_df.head()
            I_df=I_df[I_df['I']-I_df['E']>0]
            I_df['Q']=I_df['Q']*10
            I_df['err_x']=I_df.apply(lambda row: (row['Q'],row['Q']), axis=1)
            I_df['err_y']=I_df.apply(lambda row: (np.log(row['I']-row['E']),np.log(row['I']+row['E'])),axis=1)
            I_df['logI']=np.log(I_df['I'])
            I_df['logQ']=np.log(I_df['Q'])
            I_df['logX']=I_df.apply(lambda row: (row['logQ'],row['logQ']), axis=1)
            I_df['Ky']=I_df['Q']*I_df['Q']*I_df['I']
            I_df['Px']=I_df['Q']**4
            I_df['Px'].round(3)
            I_df['Py']=I_df['Px']*I_df['I']
            Int_dict_modify[key]=I_df
        return Int_dict_modify

    def get_rg_for_plot(self)->dict:
        '''
        get Rg values from SASCIF file, if unavailabel, get it from JSON
        '''
        self.get_sascif_file()
        Rg_dict={}
        for code in self.clean_SASBDB_code():
            rg={};
            all_lines=self.get_all_sascif(code)
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<3 and len(sascifline)>0 and 'sas_result.Rg_from_PR' in sascifline[0] and'sas_result.Rg_from_PR_' not in sascifline[0] :
                    rg[sascifline[0].split('.')[1]]=float(sascifline[1])
                if len(sascifline)<3 and len(sascifline)>0 and 'sas_result.Rg_from_Guinier' in sascifline[0] and'sas_result.Rg_from_Guinier_' not in sascifline[0] :
                    rg[sascifline[0].split('.')[1]]=float(sascifline[1])
            Rg_dict[code]=list(rg.values())

        if len(list(rg.values()))<1:
            data_dic=self.get_data_from_SASBDB()
            for key,val in data_dic.items():
                Rg_dict[key]=[]
                Rg_dict[key].append(round(float(val['guinier_rg']),2))
                Rg_dict[key].append(round(float(val['pddf_rg']),2))
        return Rg_dict

    def get_rg_and_io(self)->dict:
        '''
        get rg information from SASCIF file
        '''
        self.get_sascif_file()
        rg_and_io={}
        for code in self.clean_SASBDB_code():
            all_lines=self.get_all_sascif(code)
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<3 and len(sascifline)>0 and 'sas_result.Rg_from_PR' in sascifline[0] and'sas_result.Rg_from_PR_' not in sascifline[0] :
                    rg=float(sascifline[1])
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.I0_from_PR' in sascifline[0] and'_sas_result.I0_from_PR_' not in sascifline[0] :
                    io=float(sascifline[1])
            rg_and_io[code]=(rg,io)
        return rg_and_io

    def get_rg_table_many(self)->dict:
        '''
        get rg information from multiple SASCIF files
        '''
        data_dic=self.get_data_from_SASBDB()
        rg_table={'SASDB ID':[],'Rg':[],'Rg error':[],'MW':[],'MW error':[]}
        for key,val in data_dic.items():
            rg_table['Rg'].append(str(round(float(val['guinier_rg']),2))+ ' nm')
            try:
                rg_table['Rg error'].append(val['guinier_rg_error']+ ' nm')
            except:
                rg_table['Rg error'].append('N/A')
            try:
                rg_table['MW'].append(val['guinier_i0_mw']+ ' nm')
            except:
                rg_table['MW'].append('N/A')
            try:
                rg_table['MW error'].append(val['guinier_i0_mw_error']+ ' nm')
            except:
                rg_table['MW error'].append('N/A')
            rg_table['SASDB ID'].append(key)
        return rg_table

    def get_fits_for_plot(self)->dict:
        '''
        get chi-squared values from SASCIF files
        '''
        self.get_sascif_file()
        fit_dict={}
        for code in self.clean_SASBDB_code():
            fits=[];
            all_lines=self.get_all_sascif(code)
            for indx,sascifline in enumerate(all_lines):
                if (len(sascifline)<3) and (len(sascifline)>0) and ('sas_model_fitting_details.chi_square' in sascifline[0]) and (float(sascifline[1])>0.00000):
                    fits.append(round(float(sascifline[1]),2))
            if len(fits)>0:
                fit_dict[code]=fits
        return fit_dict

    def get_pofr(self)->dict:
        '''
        get pair-dist distribution from SASCIF files
        '''
        pofr_dict={}
        for code in self.clean_SASBDB_code():
            all_lines=self.get_all_sascif(code)
            data={}
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<2 and len(sascifline)>0 and 'sas_p_of_R.' in sascifline[0]:
                    data[(sascifline[0].split('.')[1])]=[]
                if len(sascifline)>2 and len(all_lines[indx-1])>0 and 'sas_p_of_R.' in all_lines[indx-1][0]:
                    for subindx,subval in enumerate(all_lines[indx:]):
                        if len(subval)>2 and '#' not in subval and 'sas' not in subval[0] :
                            for num,key in enumerate(list(data.keys())):
                                data[key].append(subval[num])
                        else:
                            break
            pdf=pd.DataFrame(list(data.values()),index=list(data.keys())).T
            pdf_re=pdf[['r','P','P_error']]
            pdf_re.rename(columns={'r':'R','P':'P','P_error':'E'},inplace=True)
            pofr_dict[code]=pdf_re
        return pofr_dict

    def get_pvals(self)->dict:
        '''
        get p-values from ATSAS
        '''
        data_dic=self.get_data_from_SASBDB()
        num_of_fits=self.get_number_of_fits()
        pval_table={'SASDB ID':[],'Model':[],'p-value':[]}
        for key,val in data_dic.items():
            num=num_of_fits[key]
            if num>0:
                for fitnum in range(0,num):
                    pval_table['SASDB ID'].append(key)
                    pval_table['Model'].append(fitnum+1)
                    target_url=val['fits'][fitnum]['fit_data']
                    fit = requests.get(target_url);
                    if fit.status_code !=200:
                        print ("Error....unable to fetch data from SASBDB, please check the entry ID")
                    fname=key+str(fitnum)+'fit.csv'
                    with open (fname,'w') as f:
                        f.write(fit.text)
                    f_df=pd.read_csv(fname, skiprows=3,delim_whitespace=True, names=['Q','Ie','Ib','E'])
                    if abs(f_df.iloc[22,2]-f_df.iloc[22,1])>abs(f_df.iloc[22,3]-f_df.iloc[22,1]):
                        f_df.rename(columns={'Q':'Q','Ie':'Ie','Ib':'E','E':'Ib'},inplace=True)
                    fit_1=f_df[['Q','Ie']]
                    fit_1.to_csv('fit1.csv',header=False,index=False)
                    fit_2=f_df[['Q','Ib']]
                    fit_2.to_csv('fit2.csv',header=False,index=False)
                    f1=open('pval.txt','w+')
                    with f1 as outfile:
                        run([config('ATSAS'),'fit1.csv','fit2.csv'],stdout=outfile)
                    f2=open('pval.txt','r')
                    all_lines=[j.strip().split() for i,j in enumerate(f2.readlines())]
                    p_val=[all_lines[i+1][4] for i,j in enumerate(all_lines) if 'adj' in j][0]
                    pval_table['p-value'].append('%.2E' % Decimal(p_val))
            else:
                pval_table['SASDB ID'].append(key)
                pval_table['Model'].append('N/A')
                pval_table['p-value'].append('N/A')
        return pval_table

    def get_pofr_ext(self)->dict:
        '''
        get pair-distance details from SASCIF files
        '''
        pofr_dict={}
        for code in self.clean_SASBDB_code():
            all_lines=self.get_all_sascif(code)
            data={}
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<2 and len(sascifline)>0 and '_sas_p_of_R_extrapolated_intensity.' in sascifline[0]:
                    data[(sascifline[0].split('.')[1])]=[]
                if len(sascifline)>2 and len(all_lines[indx-1])>0 and '_sas_p_of_R_extrapolated_intensity.' in all_lines[indx-1][0]:
                    for subindx,subval in enumerate(all_lines[indx:]):
                        if len(subval)>2 and '#' not in subval and 'sas' not in subval[0] :
                            for num,key in enumerate(list(data.keys())):
                                data[key].append(subval[num])
                        else:
                            break
            pdf=pd.DataFrame(list(data.values()),index=list(data.keys())).T
            pdf_re=pdf[['momentum_transfer','intensity_reg']]
            pdf_re.rename(columns={'momentum_transfer':'Q','intensity_reg':'I'},inplace=True)
            pdf_re=pdf_re.astype({'Q':float,'I':float})
            pdf_re['Q']=pdf_re['Q']*10
            pdf_re['logI']=np.log(pdf_re['I'])
            pofr_dict[code]=pdf_re
        return pofr_dict

    def get_pofr_errors(self)->dict:
        '''
        get pair-distance details and errors from JSON files
        '''
        pofr_dict=self.get_pofr_ext()
        Int_dict=self.modify_intensity()
        compiled_dict={}
        for code in self.clean_SASBDB_code():
            I_df=Int_dict[code]
            I_df_dict= dict(zip(I_df.Q, I_df.I))
            I_df_err_dict= dict(zip(I_df.Q, I_df.E))
            p_df=pofr_dict[code]
            p_df_dict=dict(zip(p_df.Q, p_df.I))
            errors=[]
            for Q,I in p_df_dict.items():
                data_Q=self.findMinDiff(list(I_df_dict.keys()),Q)
                if data_Q != 9999:
                    data_I=I_df_dict[data_Q]
                    delta_I=(I-data_I)
                    if I_df_err_dict[data_Q] != 0:
                        wt_delta_I=delta_I/I_df_err_dict[data_Q]
                    else:
                        wt_delta_I=0
                    errors.append([Q,delta_I,wt_delta_I])
            errors_df=pd.DataFrame(errors,columns=['Q','R','WR'])
            compiled_dict[code]=errors_df
        return compiled_dict

    def findMinDiff(self,listn: list, num: int)-> int:
        '''
        quick min diff operation for calculating errors
        '''
        list_sub=[(i,abs(j-num)) for i,j in enumerate(listn)]
        list_sort=sorted(list_sub, key=lambda x: x[1])
        if list_sort[0][1]<0.00001:
            return listn[list_sort[0][0]]
        else:
            return 9999

    def get_Guinier_data(self)->(dict,dict):
        '''
        get Guinier plot data from JSON files
        '''
        Int_dict=self.get_intensities()
        data_dic=self.get_data_from_SASBDB()
        Guinier_dict={};Guinier_score={}
        for key,val in Int_dict.items():
            G_df=val.astype({'Q':float,'I':float,'E':float})
            G_df['logI']=np.log(G_df['I'])
            rg=float(data_dic[key]['pddf_rg'])
            dmax=float(data_dic[key]['pddf_dmax'])
            index_low=int(data_dic[key]['guinier_point_first'])
            index_high=int(data_dic[key]['guinier_point_last'])
            q_min=math.pi/(dmax*10)
            q_max=1.3/(rg*10)
            G_df_range=G_df[G_df['Q']<q_max].copy()
            G_df_range['Q']=G_df['Q']*10
            G_df_range['Q2']=G_df_range['Q']**2
            X=G_df_range[['Q2']].values
            y=G_df_range['logI'].values
            regression=LinearRegression(fit_intercept=True)
            regression.fit(X,y)
            G_df_range['y_pred']=regression.predict(X)
            G_df_range['res']=y-regression.predict(X)
            G_df_range['Q2A']=G_df_range['Q2']*100
            score='%.2f' %regression.score(X,y)
            Guinier_score[key]=score
            Guinier_dict[key]=G_df_range
        return Guinier_score,Guinier_dict

    def get_parameters_vol_many_dep(self)->dict:
        '''
        get volume parameters from JSON files
        '''
        data_dic=self.get_data_from_SASBDB()
        parameter_table={'SASDB ID':[],'Estimated volume':[],'Estimated volume method':[],'Porod volume':[]}
        for key,val in data_dic.items():
            try:
                if parameter_table['Estimated volume'] is None:
                    parameter_table['Estimated volume'].append('N/A')
                else:
                    parameter_table['Estimated volume'].append(val['estimated_volume'])
                if parameter_table['Estimated volume method'] is None:
                    parameter_table['Estimated volume method'].append('N/A')
                else:
                    parameter_table['Estimated volume method'].append(val['estimated_volume_method'])
            except:
                parameter_table['Estimated volume'].append('N/A')
                parameter_table['Estimated volume method'].append('N/A')
            try:
                parameter_table['Porod volume'].append(val['porod_volume']+' nm\u00b3')
            except:
                parameter_table['Porod volume'].append('N/A')
            parameter_table['SASDB ID'].append(key)
        return parameter_table

    def get_parameters_vol_many(self)->dict:
        '''
        get volume details from SASCIF files
        '''
        self.get_sascif_file()
        parameter_table={'SASDB ID':[],'Estimated Volume':[],'Porod Volume':[],'Specific Volume':[],
        'Sample Contrast':[],'Sample Concentration':[]}
        for code in self.clean_SASBDB_code():
            parameter_table['SASDB ID'].append(code)
            all_lines=self.get_all_sascif(code)
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_sample.specimen_concentration' in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Sample Concentration'].append(str(round(float(sascifline[1]),2))+' mg/ml')
                    else:
                        parameter_table['Sample Concentration'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_sample.contrast' in sascifline[0] :
                    if len(sascifline[1])>1:
                        parameter_table['Sample Contrast'].append(str(round(float(sascifline[1]),2)))
                    else:
                        parameter_table['Sample Contrast'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_sample.specific_vol' in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Specific Volume'].append(str(round(float(sascifline[1]),2))+' nm\u00b3')
                    else:
                        parameter_table['Specific Volume'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.Porod_volume' in sascifline[0] and '_sas_result.Porod_volume_error' not in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Porod Volume'].append(str(round(float(sascifline[1]),2))+' nm\u00b3')
                    else:
                        parameter_table['Porod Volume'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.estimated_volume' in sascifline[0] and '_sas_result.estimated_volume_error' not in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Estimated Volume'].append(str(round(float(sascifline[1]),2))+' nm\u00b3')
                    else:
                        parameter_table['Estimated Volume'].append('N/A')
        return parameter_table

    def get_parameters_mw_many(self)->dict:
        '''
        get MW details from SASCIF files
        '''
        self.get_sascif_file()
        parameter_table={'SASDB ID':[],'Chemical composition MW':[],'Standard MW':[],'Porod Volume/MW':[]}
        for code in self.clean_SASBDB_code():
            parameter_table['SASDB ID'].append(code)
            all_lines=self.get_all_sascif(code)
            for indx,sascifline in enumerate(all_lines):
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.experimental_MW' in sascifline[0] and '_sas_result.experimental_MW_error' not in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Chemical composition MW'].append(str(round(float(sascifline[1]),2))+' kDa')
                    else:
                        parameter_table['Chemical composition MW'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.MW_standard' in sascifline[0] and '_sas_result.MW_standard_error' not in sascifline[0]:
                    if len(sascifline[1])>1:
                        parameter_table['Standard MW'].append(str(round(float(sascifline[1]),2))+' kDa')
                    else:
                        parameter_table['Standard MW'].append('N/A')
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.MW_Porod' in sascifline[0] and '_sas_result.MW_Porod_error' not in sascifline[0]:
                    if len(sascifline[1])>1:
                        Porod_MW=round(float(sascifline[1]),2)
                    else:
                        Porod_MW=0
                if len(sascifline)<3 and len(sascifline)>0 and '_sas_result.Porod_volume' in sascifline[0] and '_sas_result.Porod_volume_error' not in sascifline[0]:
                    if len(sascifline[1])>1 and Porod_MW>0:
                        Porod_V=round(float(sascifline[1]),2)/Porod_MW
                        parameter_table['Porod Volume/MW'].append(str(round(Porod_V,2))+' nm \u00b3/kDa')
                    else:
                        parameter_table['Porod Volume/MW'].append('N/A')

        return parameter_table

    def get_parameters_mw_many_dep(self)->dict:
        '''
        depreciated function on getting MW from JSON
        '''
        data_dic=self.get_data_from_SASBDB()
        #MW based on chemical composition
        parameter_table={'SASDB ID':[],'Sequence MW':[],'Experimental MW':[],'Porod MW':[]}
        for key,val in data_dic.items():
            try:
                parameter_table['Experimental MW'].append(list(data_dic.values())[0]['experimental_mw']+' kDa')
            except:
                parameter_tabel['Experimental MW'].append('N/A')
            try:
                parameter_table['Porod MW'].append(list(data_dic.values())[0]['porod_mw']+' kDa')
            except:
                parameter_table['Porod MW'].append('N/A')
            try:
                parameter_table['Sequence MW'].append(list(data_dic.values())[0]['experiment']['sample']['molecule'][0]['total_mw']+' kDa')
            except:
                parameter_table['Sequence MW'].append('N/A')
            parameter_table['SASDB ID'].append(key)
        return parameter_table

    def get_pddf(self)->dict:
        '''
        get p(r) data from JSON
        '''
        data_dic=self.get_data_from_SASBDB()
        pofr_dic=self.get_pofr()
        pddf_dic={}
        for key,val in pofr_dic.items():
            pd_df=val.astype({'P':float,'R':float,'E':float})
            pd_df['R']=pd_df['R']/10;
            pd_df['err_x']=pd_df.apply(lambda row: (row['R'],row['R']), axis=1)
            pd_df['err_y']=pd_df.apply(lambda row: (row['P']-row['E'],row['P']+row['E']),axis=1)
            pddf_dic[key]=pd_df
        return pddf_dic

    def get_pddf_info(self)->dict:
        '''
        get p(r) related info from JSON
        '''
        data_dic=self.get_data_from_SASBDB()
        pddf_info={'SASDB ID':[],'Software used':[],'Dmax':[],'Dmax error':[],'Rg':[],'Rg error':[]}
        for key,val in data_dic.items():
            pddf_info['Software used'].append(str(val['pddf_software']))
            try:
                pddf_info['Dmax'].append(str(val['pddf_dmax'])+' nm')
            except:
                pddf_info['Dmax'].append('N/A')
            try:
                pddf_info['Rg'].append(str(val['pddf_rg'])+' nm')
            except:
                pddf_info['Rg'].append('N/A')
            try:
                if val['pddf_dmax_error'] is None:
                    pddf_info['Dmax error'].append('N/A')
                else:
                    pddf_info['Dmax error'].append(str(val['pddf_dmax_error'])+' nm')
            except:
                pddf_info['Dmax error'].append('N/A')
            try:
                pddf_info['Rg error'].append(str(val['pddf_rg_error'])+' nm')
            except:
                pddf_info['Rg error'].append('N/A')
            pddf_info['SASDB ID'].append(key)
        return pddf_info

    def get_number_of_fits(self)->dict:
        '''
        get number of fits from JSON, deprecated
        '''
        data_dic=self.get_data_from_SASBDB()
        num_of_fits={}
        for key,val in data_dic.items():
            num_of_fits[key]=len(val['fits'])
        return num_of_fits

    def get_chi_table(self)->dict:
        '''
        get chi value from JSON, deprecated
        '''
        data_dic=self.get_data_from_SASBDB()
        chi_table={'SASDB ID':[],'Model':[],'\u03C7\u00b2':[]}
        for key,val in data_dic.items():
            numoffits=self.get_number_of_fits()[key]
            if numoffits>0:
                for fitnum in range(0,numoffits):
                    count=fitnum+1
                    chi_table['SASDB ID'].append(key)
                    chi_table['Model'].append(str(count))
                    chi_value=val['fits'][fitnum]['chi_square_value']
                    chi_value_round=round(chi_value,2)
                    chi_table['\u03C7'+'\u00b2'].append(chi_value_round)
            else:
                chi_table['SASDB ID'].append(key)
                chi_table['Model'].append('N/A')
                chi_table['\u03C7'+'\u00b2'].append('N/A')
        return chi_table

    def get_sasdb_code_fits(self)->list:
        '''
        get number of fits per SASBDB ID
        '''
        fit_dict=self.get_number_of_fits()
        return list(fit_dict.values())

    def get_fit_data(self)->dict:
        '''
        get fit information to make plots, from JSON
        '''
        data_dic=self.get_data_from_SASBDB()
        num_of_fits=self.get_number_of_fits()
        data_fit={}
        for key,val in data_dic.items():
            num=num_of_fits[key]
            fits={}
            if num>0:
                for fitnum in range(0,num):
                    target_url=val['fits'][fitnum]['fit_data']
                    fit = requests.get(target_url);
                    if fit.status_code !=200:
                        print ("Error....unable to fetch data from SASBDB, please check the entry ID")

                    fname=key+str(fitnum)+'fit.csv'
                    with open (fname,'w') as f:
                        f.write(fit.text)

                    f_df=pd.read_csv(fname, skiprows=3,delim_whitespace=True, names=['Q','Ie','Ib','E'])
                    if abs(f_df.iloc[22,2]-f_df.iloc[22,1])>abs(f_df.iloc[22,3]-f_df.iloc[22,1]):
                        f_df.rename(columns={'Q':'Q','Ie':'Ie','Ib':'E','E':'Ib'},inplace=True)
                    f_df['logIe']=np.log(f_df['Ie'])
                    f_df['logIb']=np.log(f_df['Ib'])
                    f_df['r']=f_df['Ie']-f_df['Ib']

                    if f_df['E'].isnull().values.any():
                        f_df['rsigma']=0
                    else:
                        f_df['rsigma']=f_df['r']/f_df['E']
                    #f_df['rsigma']=f_df['r']/f_df['E']
                    f_df['logr']=f_df['logIe']-f_df['logIb']
                    f_df['r2a']=(f_df['Ib']-f_df['Ie'].mean())**2
                    f_df['r2b']=(f_df['Ie']-f_df['Ie'].mean())**2
                    fits[fitnum]=(self.get_fit_r2(f_df),f_df)
            else:
                fits[0]=(0,pd.DataFrame())
                #data_fit=None
            data_fit[key]=fits
        return data_fit

    def get_fit_r2(self,df:pd.DataFrame)->int:
        rsquared=df['r2a'].sum()/df['r2b'].sum()
        return round(rsquared,2)

    def get_total_fits(self)->int:
        '''
        get number of fits
        '''
        data_dic=self.get_data_from_SASBDB()
        num_of_fits=0
        for key,val in data_dic.items():
            num_of_fits += len(val['fits'])
        return num_of_fits

    def get_fit_image(self):
        '''
        get fit image from fit, deprecated
        '''
        data_dic=self.get_data_from_SASBDB()
        num_of_fits=self.get_number_of_fits()
        data_fit={}
        for key,val in data_dic.items():
            num=num_of_fits[key]
            if num>0:
                for fitnum in range(0,num):
                    target_url=val['fits'][fitnum]['models'][0]['model_plot']
                    fitdata = requests.get(target_url);
                    if fitdata.status_code !=200:
                        print ("Error....unable to fetch data from SASBDB, please check the entry ID")
                    dirname=os.path.dirname(os.path.abspath(__file__))
                    filename = os.path.abspath(os.path.join(os.getcwd(),self.imagepath,self.ID+key+str(fitnum)+'fit.png'))
                    with open (filename,'wb') as f:
                        f.write(fitdata.content)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Package: mesxr.calibration
Module: trimscan
Author: Patrick VanMeter
Affiliation: Department of Physics, University of Wisconsin-Madison
Last Updated: November 2018

Description:
    This module is a significant overhaul to the overiginal ME-SXR calibration software. The
    goals of the re-write are to improve usability and transperency, and to more naturally
    facilitate features which were added on to the original code.
Acknowledgements:
    - Novimir Pablant and Jacob Maddox, for the original calibration code.
    - Luis Felipe Delgado-Aparicio, for heading the PPPL/MST collaboration.
    - Daniel Den Hartog and Lisa Reusch, for advising me.
"""
import copy
import os
import multiprocessing as mp
import numpy as np
import scipy as sp
from scipy.optimize import curve_fit
from scipy.special import erf
import mesxr.calibration.utilities as utilities

def calibrate_mp(pixel):
    """
    This is a helper function in order to use multiprocessing to invoke the calibrate() method for an array
    of Pilatus_Pixel objects.
    """
    pixel.calibrate()
    return copy.copy(pixel)

def calibrate_module_mp(chip):
    """
    This helper function is used to calibrate an entire module using multiprocessing. This ensures that each individual chip does
    not itself attempt to spawn additional processes.
    """
    chip.calibrate(num_cores=1)
    return copy.copy(chip)

class Pilatus_Pixel(object):
    """
    Description:
        TBD
    Arguments:
        - calibration_data = (list of np.array) Contains the calibration data for this pixel. The list should have
                a number of entries equal to the length of self.energies, and each entry should be an array with a
                length equal to the length of the corresponding entry in self.trimbits.
        - calibration_energies = (np.array) The photon energy for each calibration source element.
        - elements = (list) Each entry should be a string which identifies the element corresponding to the
                same index in calibration_energies.
        - trimbits = (list of np.array) Identifies the trimbit values which correspond to the arrays in calibration_data.
                This allows the user to remove certain trimbit settings from the calibration, if necessary.
        - coords = (tuple of ints) The global (x,y) coordinates of this pixel.
        - size = (tuple of float) The physical dimensions (x,y) of the pixel, in micrometers.
    """
    def __init__(self, calibration_data, calibration_energies, elements, trimbits, coords, chip_number, module_number, size=(172.0, 172.0)):
        # Physical parameters
        self.data = calibration_data
        self.energies = calibration_energies
        self.elements = elements
        self.trimbits = trimbits
        self.coords = coords
        self.chip = chip_number
        self.module = module_number
        self.size = size
        self.edge_pixel = False

        # Calibration parameters to be fit to the data
        self.num_elements = len(self.elements)
        self.trimfit_params = np.zeros([self.num_elements, 6])
        self.trimfit_cov    = np.zeros([self.num_elements, 6, 6])

        self.enfit_params = np.zeros(3)
        self.enfit_cov    = np.zeros([3, 3])

        self.trimfit_chi2 = np.zeros(self.num_elements)
        self.enfit_chi2 = 0.0

        # Keep track of failed fits
        self.good_trimfits = np.array([False for x in range(self.num_elements)])
        self.good_enfit = False

    def __str__(self):
        return str(self.coords)

    def s_curve(self, trim, A0, A1, A2, A3, A4, A5):
        """
        This function parameterizes the detector response to a given trimbit setting. This returns the predicted
        photon counts given the calibration parameters and supplied trimbit value. This function permits non-integer
        trimbit values, which in reality must be rounded to the nearest integer.
        """
        return 0.5*(erf( -1*(trim - A0)/(np.sqrt(2)*A1) ) + 1.)*(A2 + A3*(trim - A0)) + A4 + A5*(trim - A0)

    def s_curve_model(self, element, trim):
        """
        Use the fit results to return values from the analytic model for the S-curve fit.
        """
        elem_index = self.elements.index(element)
        return self.s_curve(trim, *self.trimfit_params[elem_index, :])

    def en_curve(self, energy, C0, C1, C2):
        """
        This function parameterizes the mapping between the trimbit of the S-curve inflection point and the calibration
        line energies.
        """
        return C0*energy**2 + C1*energy + C2

    def en_curve_model(self, energy):
        """
        Use the fit results to return values from the analytic model for the energy fit.
        """
        return self.en_curve(energy, *self.enfit_params)

    def en_curve_uncertainty(self, energy):
        """
        Returns the uncertainty in the trimbit required to set the threshold to a given energy, based on the
        energy fit results.
        """
        var_2 = (energy**4)*self.enfit_cov[0,0] + (energy**2)*self.enfit_cov[1,1] + self.enfit_cov[2,2]
        cov_2 = 2*energy*( (energy**2)*self.enfit_cov[0,1] + energy*self.enfit_cov[0,2] + self.enfit_cov[1,2] )
        return np.sqrt(var_2 + cov_2)

    def trimbit_from_threshold(self, energy):
        """
        This is just a wrapper for the en_curve_model function which provides a more consistent notation for some
        use cases.
        """
        return self.en_curve_model(energy)

    def trimbit_uncertainty(self, energy):
        """
        This is just a wrapper for the en_curve_model function which provides a more consistent notation for some
        use cases.
        """
        return self.en_curve_uncertainty(energy)

    def threshold_from_trimbit(self, trimbit):
        """
        Invert the threshold fit. That is, get the lower threshold energy corresponding to a specific trimbit value. This
        is useful when considering the impact of trimbit rounding.
        """
        c0, c1, c2 = self.enfit_params
        return (-c1 + np.sqrt(c1**2 - 4*c0*(c2 - trimbit))) / (2*c0)

    def threshold_uncertainty(self, trimbit):
        """
        """
        return 0

    def exclude_trimbits(self, remove_bits):
        """
        Remove the given trimbits and corresponding data so that they will not be used in the calibration.
        """
        for elem_index in range(self.num_elements):
            if set(remove_bits).issubset(set(self.trimbits[elem_index])):
                exclude_indices = [np.where(tbit == self.trimbits[elem_index])[0][0] for tbit  in remove_bits]
                slice_indices = [i for i in range(len(self.trimbits[elem_index])) if i not in exclude_indices]
                self.trimbits[elem_index] = self.trimbits[elem_index][slice_indices]
                self.data[elem_index] = self.data[elem_index][slice_indices]

    def exclude_elements(self, remove_elem):
        """
         Remove the given elements and corresponding data so that they will not be used in the calibration.
        """
        # Determine the indices of the entries to keep
        try:
            remove_indices = [self.elements.index(elem) for elem in remove_elem]
        except:
            print('ERROR: Supplied element(s) not in the elements array.')
            remove_indices = []

        slice_indices = [i for i in range(len(self.elements)) if i not in remove_indices]

        # Remove the unwanted elements
        self.elements = [self.elements[i] for i in slice_indices]
        self.data = [self.data[i] for i in slice_indices]
        self.trimbits = [self.trimbits[i] for i in slice_indices]
        self.energies = self.energies[slice_indices]
        self.num_elements = len(self.elements)

        # Remake the calibration result arrays which depend on the number of elements
        self.trimfit_params = np.zeros([self.num_elements, 6])
        self.trimfit_cov    = np.zeros([self.num_elements, 6, 6])
        self.trimfit_chi2   = np.zeros(self.num_elements)
        self.good_trimfits  = np.array([False for x in range(self.num_elements)])

    # Calibration functions
    def trimbit_fit(self):
        """
        Use the calibration data to determine the best-fit parameters and covariance matrix for this pixel.
        """
        for elem_index in range(self.num_elements):
            # Use a gradient method to guess the trimbit of the inflection point
            data_slope = np.gradient(self.data[elem_index], 1)
            index = np.argmin(data_slope) + self.trimbits[elem_index][0]

            # Initial guesses
            p0 = [float(index),                     # A0 - Inflection trimbit
                1.0,                              # A1 - S-curve width
                np.amax(self.data[elem_index]),   # A2 - Response amplitude
                0.0,                              # A3 - CX slope
                np.amin(self.data[elem_index]),   # A4 - BG amplitude
                0.0]                              # A5 - BG CX slope

            # Do the fit
            bounds = (np.array([-np.inf, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf]),   # Lower bounds
                      np.array([np.inf, np.inf, np.inf, 0.0, np.inf, 0.0]))               # Upper bounds

            sigma = np.sqrt(self.data[elem_index])

            try:
                self.trimfit_params[elem_index, :], self.trimfit_cov[elem_index, :, :] = curve_fit(self.s_curve, self.trimbits[elem_index], self.data[elem_index],
                                                                                p0=p0, bounds=bounds, sigma=sigma, absolute_sigma=True)

                # Determine the reduced Chi^2 of the fit
                model_data = self.s_curve(self.trimbits[elem_index], *self.trimfit_params[elem_index, :])
                self.trimfit_dof = len(self.data[elem_index]) - 6.0
                self.trimfit_chi2[elem_index] = np.sum( (self.data[elem_index] - model_data)**2/(sigma**2) )

                self.good_trimfits[elem_index] = True
            except:
                # The fit did not work - consider including fallback options here
                self.good_trimfits[elem_index] = False
                self.trimfit_params[elem_index, :] = np.nan
                self.trimfit_cov[elem_index, :, :] = np.nan
                self.trimfit_chi2[elem_index] = np.nan
                self.trimfit_dof = np.nan

    def energy_fit(self):
        """
        Use the results of the trimbit_fit to fit the en_curve.
        """
        # Only use the successful fits
        if np.sum(self.good_trimfits) >= 4:
            trim_data = self.trimfit_params[self.good_trimfits, 0]
            trim_sigma = np.sqrt(self.trimfit_cov[self.good_trimfits, 0, 0])
            trim_energies = self.energies[self.good_trimfits]

            #self.enfit_params, self.enfit_cov = np.polyfit(self.energies, trim_data, 2, w=1/trim_sigma, cov=True)
            p0 = [np.mean(trim_data), 0.0, 0.0]

            try:
                self.enfit_params[:], self.enfit_cov[:,:] = curve_fit(self.en_curve, trim_energies, trim_data, p0=p0,
                                                                      sigma=trim_sigma, absolute_sigma=True)

                # Determine the Chi^2 of the fit
                trim_model = self.en_curve(trim_energies, *self.enfit_params)
                self.enfit_dof = len(trim_data) - 3.0
                self.enfit_chi2 = np.sum( (trim_data - trim_model)**2/(trim_sigma**2) )
                self.good_enfit = True
            except:
                # The fit did not work - consider including fallback options here
                self.good_enfit = False
                self.enfit_params[:] = np.nan
                self.enfit_cov[:] = np.nan
                self.enfit_chi2 = np.nan
                self.enfit_dof = np.nan
        else:
            # There is too little good data for a reliable fit
            self.good_enfit = False
            self.enfit_params[:] = np.nan
            self.enfit_cov[:,:] = np.nan
            self.enfit_chi2 = np.nan
            self.enfit_dof = np.nan

    def calibrate(self):
        """
        Performs both the trimbit fit and the energy fit in a single command.
        """
        self.trimbit_fit()
        self.energy_fit()

class Pilatus_Null_Pixel(Pilatus_Pixel):
    """
    This class exists to allow the creation of placeholder Pixel objects which do nothing and hold no data, but share the
    same methods. This is useful as a placeholder and for representing the empty row between pixels.
    """
    def __init__(self):
        self.coords = (-1,-1)
        self.chip = -1
        self.module = -1
        self.edge_pixel = False

        # Calibration parameters to be fit to the data
        self.elements       = ['none']
        self.num_elements   = len(self.elements)
        self.trimfit_params = np.array([[np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]])
        self.trimfit_cov    = np.zeros([1,6,6])
        self.trimfit_cov[:,:,:] = np.nan

        self.enfit_params   = np.array([np.nan, np.nan, np.nan])
        self.enfit_cov      = np.nan

        self.trimfit_chi2   = np.array([np.nan])
        self.enfit_chi2     = np.nan

        self.good_trimfits = np.array([False])
        self.good_enfit = False

    def s_curve_model(self, element, trim):
        return -1.0

    def en_curve_model(self, energy):
        return -1.0

    def trimbit_fit(self):
        pass

    def energy_fit(self):
        pass

class Pilatus_Chip(object):
    """
    Description:
        The chip is the basic organizational unit for pixels in the PILATUS series of X-ray detectors. Chips essentially
        consist of an array of pixels and some additional properties.
    Inputs:
        - pixel_dims = (tuple of int) The size n x m of the array of pixels on the chip.
        - chip_number = (int) ID number for the chip.
        - calibration_data = (list of np.array) Each array in the list corresponds to the calibration data for the
                corresponding entry in elements. The arrays are 3D with indices corresponding to [x, y, trimbit].
        - calibration_energies = (np.array) The photon energy for each calibration source element.
        - elements = (list) Each entry should be a string which identifies the element corresponding to the
                same index in calibration_energies.
        - trimbits = (list of np.array) The trimbits to include for the calibration of this chip. This is expected to vary
                on a chip-by-chip basis. This should be the same length as the third axis of each element in the
                calibration_data list.
        - pixel_coords = (np.array) This array assigns the global coordinates for each pixel on the chip. The array
                is indexed by [x, y, coord] where coord=0 gives the global X coordinate and coord=1 gives the global Y.
                The indices x,y refer to the local (chip-level) coordinates.
        - vcmp = (float) The Vcmp value for this chip. This is the global value with the offset applied.
    """
    def __init__(self, pixel_dims, chip_number, module_number, calibration_data, calibration_energies, elements, trimbits, pixel_coords, vcmp):
        self.pixel_dims = pixel_dims
        self.number = chip_number
        self.module = module_number
        self.trimbits = trimbits

        # Create the array of pixels
        self.pixels = np.empty(self.pixel_dims, dtype=object)

        for y in range(self.pixel_dims[1]):
            for x in range(self.pixel_dims[0]):
                self.pixels[x, y] = Pilatus_Pixel([data[x, y, :] for data in calibration_data], calibration_energies,
                                                  elements, copy.copy(self.trimbits), tuple(pixel_coords[x,y,:]), self.number, self.module)

                if x == 0 or y == 0 or x == self.pixel_dims[0]-1 or y == self.pixel_dims[1]-1:
                    self.pixels[x, y].edge_pixel = True

        # Store chip-level detector properties
        self.vcmp = vcmp

    def __str__(self):
        return str(self.number)

    def exclude_trimbits(self, remove_bits):
        """
        Remove the specified trimbits and associated data for all pixels on the chip.
        """
        for pixel in self.pixels.ravel():
            pixel.exclude_trimbits(remove_bits)

    def remove_elements(self, remove_elem):
        """
        Remove the specified elements and associated data for all pixels on the chip.
        """
        for pixel in self.pixels.ravel():
            pixel.exclude_elements(remove_elem)

    def calibrate(self, num_cores=16):
        """
        Calibrate all pixels on a chip. Returns the reduced chi^2 metric for the energy fits. This is mostly of interest when
        investigating individual modules. Otherwise, see the detector-level calibrate() function.
        """
        if num_cores == 1:
            [pixel.calibrate() for pixel in self.pixels.ravel()]
        else:
            pool = mp.Pool(num_cores)
            pixels = pool.map(calibrate_mp, self.pixels.ravel())
            self.pixels = np.array(pixels).reshape(self.pixel_dims)

class Pilatus_Module(object):
    """
    Description:
        TBD
    Inputs:
        - chip_dims = (tuple of int) The global (X, Y) layout of chips on the module.
        - pixel_dims = (tuple of int) The local layout (x,y) of pixels on a chip.
        - module_number = (int) Identifier for the module. This is important for multi-module detectors.
        - calibration_data = (list of np.array) Each array in the list corresponds to the calibration data for the
                corresponding entry in elements. The arrays are 3D with indices corresponding to [X, Y, trimbit], where
                (X,Y) are the pixel coordinates in the module frame. It is assumed that no trimbit data has been dropped
        - calibration_energies = (np.array) The photon energy for each calibration source element.
        - elements = (list) Each entry should be a string which identifies the element corresponding to the
                same index in calibration_energies.
        - trimbit_start = (np.array) The trimbit to start the scan for each element for each module. This array is indexed
                by [elem, chip_num]. Any value above zero will cause some data to be omitted from the calibration procedure.
        - chip_coords = (np.array) Entries in this array define the starting pixel (upper-left) for each chip. This is indexed
                by [chip_num, dim] where dim=0 is x and dim=1 is y. For example, on the Pilatus3 chip 0 start with pixel (0,0)
                and chip 1 starts with (61, 0) so chip_coords[0,:] = [0,0] and chip_coords[1,:] = [61,0].
        - num_trimbits = [keyword](int) Specify the number of trimbits available to each pixel.
    """
    def __init__(self, chip_dims, pixel_dims, module_number, calibration_data, calibration_energies, elements,
                 trimbit_start, chip_coords, num_trimbits=64):
        self.number = module_number
        self.chip_dims = chip_dims
        self.pixel_dims = pixel_dims
        self.num_trimbits = num_trimbits

        # Initialize the chips on the module
        self.chips = np.empty(self.chip_dims, dtype=object)
        chip_num = 0

        for y in range(self.chip_dims[1]):
            for x in range(self.chip_dims[0]):
                vcmp = 0

                # Determine the pixel coordinates based on the supplied dimensions and the start coordinates
                pixel_coords = np.zeros([self.pixel_dims[0], self.pixel_dims[1], 2], dtype=int)
                for pix_x in range(self.pixel_dims[0]):
                    for pix_y in range(self.pixel_dims[1]):
                        pixel_coords[pix_x, pix_y, :] = chip_coords[chip_num, :] + [pix_x, pix_y]

                # Trim the calibration data to the appropriate size
                trimbits = [np.arange(trimbit_start[elem, chip_num], self.num_trimbits) for elem in range(len(elements))]

                start_x = chip_coords[chip_num, 0]
                end_x   = chip_coords[chip_num, 0]+self.pixel_dims[0]
                start_y = chip_coords[chip_num, 1]
                end_y   = chip_coords[chip_num, 1]+self.pixel_dims[1]
                pixel_data = [calibration_data[elem][start_x:end_x, start_y:end_y, trimbits[elem][0]:] for elem in range(len(elements))]

                self.chips[x, y] = Pilatus_Chip(self.pixel_dims, chip_num, self.number, pixel_data, calibration_energies, elements,
                                                trimbits, pixel_coords, vcmp)
                chip_num += 1

    def __str__(self):
        return str(self.number)

    def get_pixels(self):
        """
        Return an array with all pixels organized by module coordinates.
        """
        x_dim = 487
        y_dim = 195
        pixels_mod = np.array([[Pilatus_Null_Pixel() for y in range(y_dim)] for x in range(x_dim)])

        for chip in self.chips.ravel():
            for pixel in chip.pixels.ravel():
                pixels_mod[pixel.coords] = pixel

        return pixels_mod

    def calibrate(self, num_cores=16):
        """
        Use multiprocessing to calibrate all pixels. If possible it is advisable to set num_cores equal to the
        number of chips on the module.
        """
        if num_cores == 1:
            chips = [chip.calibrate() for chip in self.chips.ravel()]
        else:
            pool = mp.Pool(num_cores)
            chips = pool.map(calibrate_module_mp, self.chips.ravel())

        self.chips = np.array(chips).reshape(self.chip_dims)

class Pilatus_Detector(object):
    """
    Description:
        A detector is an object which holds an array of Modules and the associated global settings. Objects
        of this class also provide methods for initiating the calibration process.
    Usage:
        This class is intended to be inhereted by classes representing the implementation of specific Pilatus
        detectors. As such, direct instances of this class will be of limited use. For an example implementation,
        see the Pilatus3_100k class below.
    Parameters:
        - TBD
    """
    def __init__(self, *args, **kwargs):
        # Detector properties
        self.global_pixel_dims = (0, 0)     # Dimensions of all pixels on the detector face, including fake pixels
        self.num_chips    = 0               # The number of chips which compose the module
        self.module_dims  = (0, 0)          # The layout of modules on the detector
        self.chip_dims    = (0, 0)          # The layout of chips on the module
        self.pixel_dims   = (0, 0)          # The layout of pixels on a chip
        self.num_trimbits = 0               # The number of trimbits available to each pixel
        self.trimbit_start = [[0]]          # Value to start trimbit data at, indexed by [elem, chip_num]
        self.calibrated = False             # Flag to tell whether the detector has been calibrated already
        self.bad_pixels = []                # Keep track of which pixels are problematic

        # Initialize the modules and global settings
        self.load_trimscan_data([], [])
        self.load_global_settings(0)
        self.init_modules()

        # Make pixels easily accessible by global coordinates
        self.pixels = self.get_pixels()

    def init_modules(self):
        """
        Run this to initialize the modules which compose the detector. This is separated out from the __init__ method
        since it does generally not need to be overwritten for inhereted classes.
        """
        self.modules = np.empty(self.module_dims, dtype=object)
        module_number = 0

        for y in range(self.module_dims[1]):
            for x in range(self.module_dims[0]):
                self.modules[x,y] = Pilatus_Module(self.chip_dims, self.pixel_dims, module_number, self.trimscan_data, self.trimscan_energy,
                                                   self.elements, self.trimbit_start, self.get_chip_coords(module_number),
                                                   num_trimbits=self.num_trimbits)
                module_number += 1

    def load_global_settings(self, settings_path):
        """
        Load the detector settings in from the data file. This may require slight modification for multi-module
        detectors since I do not know which settings are shared across modules. I have separated out the VCMP settings
        due to its unique chip-level variation.
        """
        self.vcmp = np.zeros(self.num_chips)
        self.settings = {}

        # Load the settings data from the file
        if settings_path != 0:
            try:
                fname = os.path.join(settings_path, 'setdacs_b01_m01.dat')
                settings_file = open(fname, 'r')

                for line in settings_file.readlines():
                    line_elems = line.split()
                    if line_elems[0] == 'set':
                        key = line_elems[1].split('_')[2]
                        self.settings[key] = float(line_elems[2])

                settings_file.close()

                # Pull out the VCMP settings
                for key in self.settings.keys():
                    if 'VCMP' in key:
                        index = int(key.split('VCMP')[1])
                        self.vcmp[index] = self.settings[key]

            except:
                print('ERROR: Failed to load supplied settings file. Check path name.')
        else:
            pass

    def load_trimscan_data(self, elements, trimscan_paths):
        """
        """
        self.elements = elements

        self.trimscan_data = [np.zeros([self.global_pixel_dims[0], self.global_pixel_dims[1], self.num_trimbits]) for i in self.elements]
        self.trimscan_energy = utilities.get_line_energy(self.elements)
        self.trimbits = [np.arange(self.num_trimbits) for elem in self.elements]

        for elem, fname in enumerate(trimscan_paths):
            self.trimscan_data[elem] = utilities.load_calibration_data(fname)

    def generate_settings_file(self, output_path):
        """
        Used to generate the settings file for use in ME-SXR operation. This removes the requirement of keeping
        up with the original file used in the calibration procedure.
        """
        pass

    def pixel_mapping(self, x, y):
        """
        Returns the mapping between between global (x,y) coordinates and local (mod_num, chip_num, (pix_x, pix_y))
        coordinates. This is important for mapping the trimbit scan data to the Detector data structure.
        """
        return (0, 0, (x, y))

    def get_chip_coords(self, mod_number):
        """
        Override to set the global coordinates for the first (upper-left) pixel on each chip.
        """
        return np.array([[0,0]])

    def get_pixels(self):
        """
        Return an array of all pixel objects following the global coordinates.
        """
        pass

    def calibrate(self, num_cores=16):
        """
        Use multiprocessing to calibrate all pixels.
        """
        pass

    def determine_bad_pixels(self, include=[]):
        """
        Update the list of bad pixels. This should generally be called immediately after calibration. Use
        the 'include' keyword to manually add points, even if the fit did not fail. This keyword accepts a
        list of tuples.
        """
        # Ensure that the specified fits are marked as bad
        for coords in include:
            if coords not in self.bad_pixels:
                self.pixels[coords].good_enfit = False

        # Now add all remaining pixels for which a fit was not found
        for pixel in self.pixels.ravel():
            if not pixel.good_enfit and type(pixel) != Pilatus_Null_Pixel:
                if pixel.coords not in self.bad_pixels:
                    self.bad_pixels.append(pixel.coords)

    def get_trimbit_map(self, threshold_map):
        """
        Returns the trimbit values for the supplied threshold map (in keV).
        """
        if self.calibrated:
            trimbit_map = np.zeros(self.global_pixel_dims, dtype=int)

            for x in range(self.global_pixel_dims[0]):
                for y in range(self.global_pixel_dims[1]):
                    # Ensure that the trimbit is within the valid range
                    tbit = self.pixels[x,y].trimbit_from_threshold(threshold_map[x,y])
                    if tbit < 0 or np.isnan(tbit):
                        tbit = 0
                    elif tbit > self.num_trimbits - 1:
                        tbit = self.num_trimbits - 1
                    else:
                        tbit = int(round(tbit))

                    trimbit_map[x,y] = tbit
        else:
            print('Detector is not yet calibrated.')
            trimbit_map = -1*np.ones(self.global_pixel_dims, dtype=int)
        return trimbit_map

# ---------------------------------------- Specific Detector Implementations ----------------------------------------

class Pilatus3_100k(Pilatus_Detector):
    """
    Description:
        This is a specific implementation of the Pilatus_Detector which integrates the appropraite pixel geometry
        and device settings. The PILATUS3 100K detector, produced by DECTRIS Ltd., is composed of a single module
        with approximately 100,000 total pixels.
    """
    def __init__(self, elements, trimscan_paths, settings_path, trimbit_start=0, omit_trimbits=False, calib_name='MST'):
        # Pilatus3 pixel properties
        self.global_pixel_dims = (487, 195)
        self.num_chips         = 16
        self.module_dims       = (  1,   1)
        self.chip_dims         = (  8,   2)
        self.pixel_dims        = ( 60,  97)
        self.num_trimbits      = 64
        self.calibrated        = False
        self.bad_pixels        = []
        self.name              = calib_name

        if not omit_trimbits:
            self.trimbit_start = np.zeros([len(elements), self.num_chips], dtype=int)
        else:
            self.trimbit_start = trimbit_start

        # Pilatus3 physical characteristics

        # Load modules and settings
        self.load_trimscan_data(elements, trimscan_paths)
        self.load_global_settings(settings_path)
        self.init_modules()

        # Make pixels easily accessible by global coordinates
        self.pixels = self.get_pixels()

    def pixel_mapping(self, x, y):
        """
        Implement the proper mapping for the PILATUS3 100K detector. This is most easily done using the pre-existing
        mapping function from Novi.
        """
        chip_num, chip_x, chip_y = utilities.get_chip_coords(x, y)
        return (0, chip_num, (chip_x, chip_y))

    def get_chip_coords(self, mod_number):
        """
        Returns the global coordinates for the first (upper-left) pixel on each chip for the PILATUS3 100K.
        """
        dx = 61
        dy = 98
        chip_coords = np.array([[0*dx, 0], [1*dx, 0], [2*dx, 0], [3*dx, 0], [4*dx, 0], [5*dx, 0], [6*dx, 0], [7*dx, 0],
                                [0*dx,dy], [1*dx,dy], [2*dx,dy], [3*dx,dy], [4*dx,dy], [5*dx,dy], [6*dx,dy], [7*dx,dy]])
        return chip_coords

    def write_settings_file(self, filepath):
        """
        Writes the settings used to generate the trimscan back out to a file.
        """
        # Format the settings into a string with linebreaks
        settings = '# /dev/shm/setdacs_b01_m01.dat\n'
        settings += 'set B01_M01_VTRM {0: 2.4f}\n'.format(self.settings['VTRM'])

        for vcmp in range(self.num_chips):
            vcmp_str = 'VCMP{0:}'.format(vcmp)
            settings += 'set B01_M01_{0:} {1: 2.4f}\n'.format(vcmp_str, self.settings[vcmp_str])

        for key in ['VCCA', 'VRF', 'VRFS', 'VCAL', 'VDEL', 'VADJ']:
            settings += 'set B01_M01_{0:} {1: 2.4f}\n'.format(key, self.settings[key])

        # Remove the final unwanted newline
        settings = settings[:-1]

        with open(os.path.join(filepath, 'setdacs_b01_m01.dat'), 'wb') as f:
            f.write(settings)

        print('Settings file written to {0:}.'.format(filepath))

    def write_autotrim_files(self, threshold_map, filepath, filename='autotrim'):
        """
        Generate autotrim files given a specified threshold map using the calibration.
        """
        trimbit_map = self.get_trimbit_map(threshold_map)
        trimbit_sorted = np.zeros([self.num_chips, self.pixel_dims[0], self.pixel_dims[1]])

        for g_x in range(self.global_pixel_dims[0]):
            for g_y in range(self.global_pixel_dims[1]):
                cn, cx, cy = utilities.get_chip_coords(g_x, g_y)
                if cn != -1:
                    trimbit_sorted[cn, cx, cy] = trimbit_map[g_x, g_y]

        for cn in range(self.num_chips):
            trim_string = '# AUTOTRIM files for PILATUS3 100K by UW Madison code\n'
            trim_string += 'set B01_M01_CHSEL{0:} 1\n'.format(cn)

            for cx in range(self.pixel_dims[0]):
                for cy in range(self.pixel_dims[1]):
                    trim_string += 'trim {0:} {1:} {2:}\n'.format(cx, cy, hex(int(trimbit_sorted[cn, cx, cy])))

            trim_string += 'settrims'

            # Write the output file
            fname = '{0:}_b01_m01_c{1:02d}.dat'.format(filename, cn)
            with open(os.path.join(filepath, fname), 'wb') as f:
                f.write(trim_string)

            print('Saved {0:} to file.'.format(fname))

    def write_bad_pixels_file(self, filepath):
        """
        Outputs a CSV file containing all of the bad pixels. This is useful in later analysis.
        """
        fname = 'bad_pixels.csv'
        np.savetxt(os.path.join(filepath, fname), self.bad_pixels, delimiter=',', fmt='%d')

    def get_pixels(self):
        """
        Since this detector has only one module, just call that module's get_pixels() method.
        """
        return self.modules[0,0].get_pixels()

    def calibrate(self, num_cores=16, bad_pix=[]):
        """
        Since this detector has only one module, just call that module's calibrate() method. Use the 'bad_pix'
        to manually ensure that the specified pixels get marked as bad, even if the fit does not fail.
        """
        self.modules[0,0].calibrate(num_cores=num_cores)
        self.calibrated = True
        self.pixels = self.get_pixels()
        self.determine_bad_pixels(include=bad_pix)