from matplotlib import pyplot as plt
from matplotlib import mlab as ml
import numpy as np
import csv
import os
import time
import multiprocessing
from multiprocessing import Value
from scipy.io import wavfile
from visualization import visualize
import math
import argparse
from random import shuffle
import random
from functools import partial
import generate_spectrograms

parser = argparse.ArgumentParser()
#parser.add_argument('--data', dest='dataDir', default='../elephant_dataset/New_Data/Truth_Logs',
#     type=str, help='The top level directory with the data (e.g. Truth_Logs)')

parser.add_argument('--data_dirs', dest='data_dirs', nargs='+', type=str,
    help='Provide the data_dirs with the files that you want to be processed')

# For use on quatro
#parser.add_argument('--data', dest='dataDir', default='/home/data/elephants/rawdata/raw_2018',
#   type=str, help='The top level directory with the data (e.g. Truth_Logs)')

parser.add_argument('--out', dest='out_dir', default='../elephant_dataset/Train',
     help='The output directory for the processed files. We need to specify this!')
'''
parser.add_argument('--train_dir', default='../elephant_dataset/Train_New',
     help='Directory for the training chunks for the new data')
parser.add_argument('--test_dir', default='../elephant_dataset/Test_New',
     help='Directory for the test chunks for the new data')
'''
parser.add_argument('--NFFT', type=int, default=4096, help='Window size used for creating spectrograms') # was 3208!!
parser.add_argument('--hop', type=int, default=800, help='Hop size used for creating spectrograms') # was 641!!
parser.add_argument('--window', type=int, default=256,
    help='Deterimes the window size in frames of the resulting spectrogram') # Default corresponds to 21s - should just check this!
parser.add_argument('--max_f', dest='max_freq', type=int, default=150, help='Deterimes the maximum frequency band')
parser.add_argument('--pad', dest='pad_to', type=int, default=4096,
    help='Deterimes the padded window size that we want to give a particular grid spacing (i.e. 1.95hz')
parser.add_argument('--neg_fact', type=int, default=1,
    help="Determines number of negative samples to sample as neg_fact x (pos samples)")
parser.add_argument('--full_24_hr', action='store_true',
    help="Determines whether to create all chunks or just ones corresponding to labels and negative samples")
parser.add_argument('--seed', type=int, default=8,
    help="Set the random seed used for creating the datasets. This is primarily important for determining the negative samples")

'''
parser.add_argument('--val_size', type=float, default=0.1, help='Determines the relative size of the val set if we are creating one')
parser.add_argument('--train_val', action='store_true',
    help = 'Generate a train and validation split based on the training data')
parser.add_argument('--val_dir', dest='val_dir', default='../elephant_dataset/Val/',
     help='The output directory for the validation files. Only need this if creating validation files')
'''

VERBOSE = False

def generate_labels(labels, spectrogram_info, len_wav):
    '''
        Given ground truth label file 'label' create the full
        segmentation labeling for a .wav file. Namely, return
        a vector containing a 0/1 labeling for each time slice
        corresponding to the .wav file transformed into a spectrogram.
        The key challenge here is that we want the labeling to match
        up with the corresponding spectrogram without actually creating
        the spectrogram
    '''
    # Formula is: frames = floor((wav - overlap) / hop)
    len_labels = math.floor((len_wav - (spectrogram_info['NFFT'] - spectrogram_info['hop'])) / spectrogram_info['hop'])
    labelMatrix = np.zeros(shape=(len_labels),dtype=int)

    # In the case where there is actually no GT label file because no elephant
    # calls occred in the recording, we simply output all zeros
    if labels is None:
        return labelMatrix

    labelFile = csv.DictReader(open(labels,'rt'), delimiter='\t')

    samplerate = spectrogram_info['samplerate']
    # Iterates through labels and marks the segments with elephant calls
    for row in labelFile:
        # Use the file offset to determine the start of the call
        start_time = float(row['File Offset (s)'])
        call_length = float(row['End Time (s)']) - float(row['Begin Time (s)'])
        end_time = start_time + call_length

        # Figure out which spectrogram slices we are on
        # to get columns that we want to span with the given
        # slice. This math transforms .wav indeces to spectrogram
        # indices
        start_spec = max(math.ceil((start_time * samplerate - spectrogram_info['NFFT'] / 2.) / spectrogram_info['hop']), 0)
        end_spec = min(math.ceil((end_time * samplerate - spectrogram_info['NFFT'] / 2.) / spectrogram_info['hop']), labelMatrix.shape[0])
        labelMatrix[start_spec : end_spec] = 1

    return labelMatrix

def generate_empty_chunks(n, raw_audio, label_vec, spectrogram_info):
    """
        Generate n empty data chunks by uniformally sampling
        time sections with no elephant calls present
    """
    # Step through the labels vector and collect the indeces from
    # which we can define a window with no elephant call
    # i.e. all start indeces such that the window (start, start + window_sz)
    # does not contain an elephant call
    valid_starts = []
    window_size = spectrogram_info['window']
    # Step backwards and keep track of how far away the
    # last elephant call was
    last_elephant = 0  # For now is the size of the window
    for i in range(label_vec.shape[0] - 1, -1, -1):
        last_elephant += 1

        # Check if we encounter an elephant call
        if (label_vec[i] == 1):
            last_elephant = 0

        # If we haven't seen an elephant call
        # for a chunk size than record this index
        if (last_elephant >= window_size):
            valid_starts.append(i)

    # Generate num_empty uniformally random
    # empty chunks
    empty_features = []
    empty_labels = []
    NFFT = spectrogram_info['NFFT']
    samplerate = spectrogram_info['samplerate']
    hop = spectrogram_info['hop']
    max_freq = spectrogram_info['max_freq']
    pad_to = spectrogram_info['pad_to']

    for i in range(n):
        # Generate a valid empty start chunk
        # index by randomly sampling from our
        # ground truth labels
        start = np.random.choice(valid_starts)

        # Now we have to do a litle back conversion to get
        # the raw audio index in raw audio frames
        # Number of hops in we are marks the first raw audio frame to use
        chunk_start = start * spectrogram_info['hop']
        chunk_size = (spectrogram_info['window'] - 1) * spectrogram_info['hop'] + spectrogram_info['NFFT']
        chunk_end = int(chunk_start + chunk_size)

        # Get the spectrogram chunk
        # Extract the spectogram
        [spectrum, freqs, t] = ml.specgram(raw_audio[chunk_start: chunk_end],
                    NFFT=NFFT, Fs=samplerate, noverlap=(NFFT - hop), window=ml.window_hanning, pad_to=pad_to)

        # Cutout the high frequencies that are not of interest
        spectrum = spectrum[(freqs <= max_freq)]
        assert(spectrum.shape[1] == spectrogram_info['window'])

        data_labels = label_vec[start : start + spectrum.shape[1]]
        # Make sure that no call exists in the chunk
        assert(np.sum(data_labels) == 0)

        if VERBOSE:
            new_features = 10*np.log10(spectrum)
            visualize(new_features.T, labels=data_labels)

        # We want spectrograms to be time x freq
        spectrum = spectrum.T
        empty_features.append(spectrum)
        empty_labels.append(data_labels)

    return empty_features, empty_labels

def generate_chunk(start_time, end_time, raw_audio, truth_labels, spectrogram_info):
    '''
        Generate a data chunk around a given elephant call. The data
        chunk is of size "chunk_length" seconds and has the call
        of interest randomly placed inside the window

        Parameters:
        - start_time and end_time in seconds
    '''
    # Convert the times to .wav frames to help ensure
    # robustness of approach
    start_frame = int(math.floor(start_time * spectrogram_info['samplerate']))
    end_frame = int(math.ceil(end_time * spectrogram_info['samplerate']))
    # Convert from window size in spectrogram frames to raw audio size
    # Note we use the -1 term to force the correct number of frames
    # wav = frames * hop - hop + window ==> wav = frames * hop + overlap
    chunk_size = (spectrogram_info['window'] - 1) * spectrogram_info['hop'] + spectrogram_info['NFFT']

    # Padding to call
    call_length = end_frame - start_frame # In .wav frames
    padding_length = chunk_size - call_length
    # if padding_frame is neg skip call
    # but still want to go to the next!!
    if padding_length < 0:
        print ("skipping too long of call") # Maybe don't need this let us se
        return None, None

    # Randomly split the pad to before and after
    pad_front = np.random.randint(0, padding_length + 1)

    # Do some stuff to avoid the front and end!
    chunk_start = start_frame - pad_front
    chunk_end  = start_frame + call_length + (padding_length - pad_front)

    # Do some quick voodo - assume cant have issue where
    # the window of 64 frames is lareger than the sound file!
    if (chunk_start < 0):
        # Amount to transfer to end
        chunk_start = 0
        chunk_end = chunk_size
    # See if we have passed the end of the sound file.
    # Note divide by sr to get sound file length in seconds
    if (chunk_end >= raw_audio.shape[0]):
        chunk_end = raw_audio.shape[0]
        chunk_start = raw_audio.shape[0] - chunk_size

    assert(chunk_end - chunk_start == chunk_size)
    # Make sure the call is fully in the region
    assert(chunk_start <= start_frame and chunk_end >= end_frame)

    NFFT = spectrogram_info['NFFT']
    samplerate = spectrogram_info['samplerate']
    hop = spectrogram_info['hop']
    max_freq = spectrogram_info['max_freq']
    pad_to = spectrogram_info['pad_to']
    # Extract the spectogram
    [spectrum, freqs, t] = ml.specgram(raw_audio[chunk_start: chunk_end],
                NFFT=NFFT, Fs=samplerate, noverlap=(NFFT - hop), window=ml.window_hanning, pad_to=pad_to)

    # Check our math
    assert(spectrum.shape[1] == spectrogram_info['window'])

    # Cutout the high frequencies that are not of interest
    spectrum = spectrum[(freqs <= max_freq)]
    # Get the corresponding labels
    # Calculate the relative start time w/r
    # to the entire spectogram for the given chunk
    start_spec = max(math.ceil((chunk_start - spectrogram_info['NFFT'] / 2.) / spectrogram_info['hop']), 0)
    end_spec = start_spec + spectrum.shape[1]

    data_labels = truth_labels[start_spec: end_spec]

    if VERBOSE:
        new_features = 10*np.log10(spectrum)
        visualize(new_features.T, labels=data_labels)

    # We want spectrograms to be time x freq
    spectrum = spectrum.T
    return spectrum, data_labels

def generate_elephant_chunks(raw_audio, labels, label_vec, spectrogram_info, num_random_call_positions=10):
    """
        Generate the data chunks for each elephant call in a given
        audio recording, where a data chunk is defined as spectrogram
        of a given window size with the given call randomly placed within

        Parameters
        - raw_audio: Vector of raw 1D audio
        - labels: Dictionary representaiton of the truth labels csv
        - label_vec: Vector with 0/1 label for each spectrogram column
        - spectrogram_info: Spectrogram params
    """
    feature_set = []
    label_set = []
    # 2. Now iterate through label file, when find a call, pass to make chunk,
    for call in labels:
        start_time = float(call['File Offset (s)'])
        call_length = float(call['End Time (s)']) - float(call['Begin Time (s)'])
        end_time = start_time + call_length

        for _ in range(num_random_call_positions):
            feature_chunk, label_chunk = generate_chunk(start_time, end_time, raw_audio, label_vec, spectrogram_info)

            if (feature_chunk is not None):
                feature_set.append(feature_chunk)
                label_set.append(label_chunk)

    return feature_set, label_set

def generate_data_chunks(audio_file, label_file, spectrogram_info, num_neg=0):
    """
        Extract either elephant calls from the audio file or
        negative examples. If num_neg = 0 then extract the
        elephant calls, otherwise extract num_neg negative
        samples from the audio file.
    """
    # Just for the bai elephants
    # Need to loook into this
    try:
        samplerate, raw_audio = wavfile.read(audio_file)
        print ("File size", raw_audio.shape)
    except:
        print("FILE Failed", audio_file)
    # No calls to extract
    if label_file is None and num_neg == 0:
        return [], []
    elif label_file is not None and num_neg == 0: # We have a calls to extract
        labels = csv.DictReader(open(label_file,'rt'), delimiter='\t')

    # Generate the spectrogram index labelings
    spectrogram_info['samplerate'] = samplerate
    label_vec = generate_labels(label_file, spectrogram_info, raw_audio.shape[0])

    if num_neg == 0:
        print ("Extracting elephant calls from", audio_file)
        feature_set, label_set = generate_elephant_chunks(raw_audio, labels, label_vec, spectrogram_info)
    else:
        print ("Extracting negative chuncks from", audio_file)
        feature_set, label_set = generate_empty_chunks(num_neg, raw_audio, label_vec, spectrogram_info)

    return feature_set, label_set

def extract_data_chunks(audio_file, label_file, spectrogram_info):
    """
        Extract an equal number of empty and elephant call
        data chunks from a raw audio file. Each data chunk
        is of a fixed size, and for data chunks with an
        elephant call we make a chunk around a given reference
        call and then randomly place that call within the fram
    """
    print ("Processing", audio_file)
    samplerate, raw_audio = wavfile.read(audio_file)
    labels = csv.DictReader(open(label_file,'rt'), delimiter='\t')
    # Generate the spectrogram index labelings
    spectrogram_info['samplerate'] = samplerate
    label_vec = generate_labels(label_file, spectrogram_info, raw_audio.shape[0])

    feature_set, label_set = generate_elephant_chunks(raw_audio, labels, label_vec, spectrogram_info)

    empty_features, empty_labels = generate_empty_chunks(len(feature_set) * spectrogram_info['neg_fact'],
            raw_audio, label_vec, spectrogram_info)
    feature_set.extend(empty_features)
    label_set.extend(empty_labels)

    return feature_set, label_set

if __name__ == '__main__':
    args = parser.parse_args()
    data_dirs = args.data_dirs
    out_dir = args.out_dir
    spectrogram_info = {'NFFT': args.NFFT,
                        'hop': args.hop,
                        'max_freq': args.max_freq,
                        'window': args.window,
                        'pad_to': args.pad_to,
                        'neg_fact': args.neg_fact}

    print(args)

    np.random.seed(args.seed)
    random.seed(args.seed)

    # Collect the wav/txt file pairs
    # Iterate through all files with in data directories
    data_pairs = {}
    for currentDir in data_dirs:
        #Iterate through each dir and get files within
        for(dirpath, dirnames, filenames) in os.walk(currentDir):
            # Create data/label
            # pairs (i.e. (.wav, .txt))
            for eachFile in filenames:
                # Strip off the location and time tags
                tags = eachFile.split('_')
                data_id = tags[0] + '_' + tags[1]
                file_type = eachFile.split('.')[1]

                if (file_type not in ['wav', 'txt']):
                    continue

                # Insert the file name into the dictionary
                # with the file type tag for a given id
                if not data_id in data_pairs:
                    data_pairs[data_id] = {}
                    data_pairs[data_id]['id'] = data_id
                    data_pairs[data_id]['dir'] = currentDir

                data_pairs[data_id][file_type] = eachFile

    # data_pairs now contains all of the wav/txt data file pairs.
    # Let us create a list of these pairs to randomly split into
    # train/val if the flag is set

    # Some wav files do not have a gt because there are no calls!
    file_pairs = []
    for _, pair in data_pairs.items():
        # We must be careful in the case where there is no
        # label file basically there are no calls!
        if 'wav' in pair:
            gt_labels = None if 'txt' not in pair else pair['txt']
            file_pairs.append(((pair['wav'], gt_labels, pair['id'], pair['dir'])))

    if args.full_24_hr:
        print("Processing all 24 hours and saving chunks from them")

        spectrogram_info['samplerate'] = 8000

        def wrapper_processFull24Hours(directory, data_pair):
            """
            This worker function is called on every data sample
            """
            audio_file = data_pair[0]
            label_file = data_pair[1]
            data_id = data_pair[2]
            curren_dir = data_pair[3]

            # Catch case where no calls exist so the gt file does not
            label_path = curren_dir + '/' + label_file if label_file is not None else None
            full_24_hr_spectogram = generate_spectrograms.generate_whole_spectogram(curren_dir + '/' + audio_file, spectrogram_info, "-1")
            labels = generate_spectrograms.generate_labels(label_path, spectrogram_info, full_24_hr_spectogram.shape[0])
            print("Shapes of full spectrograms and labels")
            print(full_24_hr_spectogram.shape)
            print(labels.shape)

            # Save the individual files seperately for each location!
            num_chunks = 0
            for i in range(math.floor(full_24_hr_spectogram.shape[0] / 256)):
                feature = full_24_hr_spectogram[i * 256:(i + 1) * 256, :]
                label = labels[i * 256:(i + 1) * 256]

                if feature.shape[0] != 256:
                    print("MAJOR PROBLEMSSSS WHY DOESNT MULTIPROCESSING SURFACE ERRORS")
                assert feature.shape[0] == 256
                assert label.shape[0] == 256

                np.save(directory + '/' + data_id + "_features_" + str(i), feature)
                np.save(directory + '/' + data_id + "_labels_" + str(i), label)
                call_counter.value += 1
                num_chunks += 1

            print("Saved successfully {} chunks.".format(num_chunks))

        out_dir += '/Full_24_hrs'
        if not os.path.isdir(out_dir):
            os.mkdir(out_dir)

        # Process Given Data pairs.
        print ("Processing full 24 hours examples")
        print ("Num Files: ", len(file_pairs))

        call_counter = Value("i", 0) # Shared thread variable to count the number of postive call examples
        pool = multiprocessing.Pool(20)
        print('Multiprocessing')
        start_time = time.time()
        pool.map(partial(wrapper_processFull24Hours, out_dir), file_pairs)
        print('Multiprocessed took {}'.format(time.time()-start_time))
        pool.close()
        print('Multiprocessed took {}'.format(time.time()-start_time))

        quit()

    def wrapper_processPos(directory, data_pair):
        """
        This worker function is called on every data sample
        """
        audio_file = data_pair[0]
        label_file = data_pair[1]
        data_id = data_pair[2]
        curren_dir = data_pair[3]

        # Catch case where no calls exist so the gt file does not
        label_path = curren_dir + '/' + label_file if label_file is not None else None
        feature_set, label_set = generate_data_chunks(curren_dir + '/' + audio_file,
                                        label_path, spectrogram_info)
        call_counter.value += len(feature_set)

        # Save the individual files seperately for each location!
        for i in range(len(feature_set)):
            np.save(directory + '/' + data_id + "_features_" + str(i), feature_set[i])
            np.save(directory + '/' + data_id + "_labels_" + str(i), label_set[i])

    # Add now the random seed used to create it
    out_dir += '/Neg_Samples_x' + str(args.neg_fact) + '_Seed_' + str(args.seed)
    if not os.path.isdir(out_dir):
        os.mkdir(out_dir)

    # Process Given Data pairs.
    print ("Processing Positive examples")
    print ("Num Files: ", len(file_pairs))

    call_counter = Value("i", 0) # Shared thread variable to count the number of postive call examples
    pool = multiprocessing.Pool()
    print('Multiprocessing on {} CPU cores'.format(os.cpu_count()))
    start_time = time.time()
    pool.map(partial(wrapper_processPos, out_dir), file_pairs)
    print('Multiprocessed took {}'.format(time.time()-start_time))
    pool.close()
    print('Multiprocessed took {}'.format(time.time()-start_time))

    num_calls = call_counter.value
    num_sound_files = len(file_pairs)

    num_neg_samples = spectrogram_info['neg_fact'] * num_calls
    samples_per_file = int(math.ceil(num_neg_samples / float(num_sound_files)))
    print ("Num calls:", num_calls)
    print("Num neg:", num_neg_samples)
    print("samples per:", samples_per_file)

    def wrapper_processNeg(directory, num_negative, data_pair):
        """
        This worker function is called on every data sample
        """
        audio_file = data_pair[0]
        label_file = data_pair[1]
        data_id = data_pair[2]
        curren_dir = data_pair[3]

        # Catch case where no calls exist so the gt file does not
        label_path = curren_dir + '/' + label_file if label_file is not None else None
        feature_set, label_set = generate_data_chunks(curren_dir + '/' + audio_file,
                                        label_path, spectrogram_info, num_negative)

        # Save the individual files seperately for each location!
        for i in range(len(feature_set)):
            np.save(directory + '/' + data_id + "_neg-features_" + str(i), feature_set[i])
            np.save(directory + '/' + data_id + "_neg-labels_" + str(i), label_set[i])

    # Generate num_neg_samples negative examples where we randomly
    # sample "samples_per_file" examples from each file
    print ("Processing Positive examples")
    print ("Size: ", len(file_pairs))
    pool = multiprocessing.Pool()
    print('Multiprocessing on {} CPU cores'.format(os.cpu_count()))
    start_time = time.time()
    pool.map(partial(wrapper_processNeg, out_dir, samples_per_file), file_pairs)
    print('Multiprocessed took {}'.format(time.time()-start_time))
    pool.close()
    print('Multiprocessed took {}'.format(time.time()-start_time))

    # Save which files were used for main data files
    with open(out_dir + '/files.txt', 'w') as f:
        for i in range(len(file_pairs)):
            file = file_pairs[i]
            # write the id of each file pair (wav/txt)
            # along with which time period it came from
            # e.g. jan/id
            dirs = file[3].split('/')
            time_tag = dirs[-1]

            path = time_tag + '/' + file[2]
            if i != len(file_pairs) - 1:
                path += '\n'

            f.write(path)


import time
from model.product import Product

class ProductHelper:

    def __init__(self, app):
        self.app = app

    def get_product_list(self, pattert=''):
        wd = self.app.wd
        self.find_products(pattert)

        product_list = wd.find_element_by_css_selector("[class='table table-striped table-bordered table-condensed']").text.split('\n')
        product_count = len(product_list) // 3
        list = []
        for i in range(product_count):
            product_kwargs = dict(zip(['name', 'description', 'price'], product_list[3*i+1:3*i+4]))
            list.append(Product(**product_kwargs))
        return list

    def find_products(self, pattert=''):
        wd = self.app.wd
        self.app.open_home_page()
        wd.find_element_by_xpath("//html//body//nav//div//ul//li[4]//form//input").click()
        wd.find_element_by_xpath("//html//body//nav//div//ul//li[4]//form//input").clear()
        wd.find_element_by_xpath("//html//body//nav//div//ul//li[4]//form//input").send_keys(pattert)
        time.sleep(1)
        wd.find_element_by_id("searchButton").click()
        time.sleep(2)


import os
import unittest

import angr

try:
    import pysoot
except ModuleNotFoundError:
    pysoot = None

test_location = os.path.join(os.path.dirname(os.path.realpath(str(__file__))), "..", "..", "binaries", "tests")

# pylint: disable=missing-class-docstring
# pylint: disable=no-self-use
@unittest.skipUnless(pysoot, "pysoot not available")
class TestCfgfastSoot(unittest.TestCase):
    def test_simple1(self):
        binary_path = os.path.join(test_location, "java", "simple1.jar")
        p = angr.Project(binary_path, main_opts={"entry_point": "simple1.Class1.main"}, auto_load_libs=False)
        cfg = p.analyses.CFGFastSoot()
        assert cfg.graph.nodes()

    def test_simple2(self):
        binary_path = os.path.join(test_location, "java", "simple2.jar")
        p = angr.Project(binary_path, main_opts={"entry_point": "simple2.Class1.main"}, auto_load_libs=False)
        cfg = p.analyses.CFGFastSoot()
        assert cfg.graph.nodes()

if __name__ == "__main__":
    unittest.main()


# -*-coding:utf-8-*-
# @作者：haiyu.ma
# @创建日期：2020-11-03 23:10
# @Software：PyCharm
# ----------------------------------------------------------------------------
# 一、int类型
# 作用：记录年龄、身份证号、数量、个数等等
# 定义：
age = 10  # age = int(10)

# 类型转换:（了解）
res = int('102033')  # 纯数字的字符串转换成int
print(res, type(res))
# 进制转换
# 十进制转二进制  （除二取余，然后倒序排列，高位补零）
# 11 -> 1011
# print(bin(11))  # 0b1011  0b代表二进制
# 十进制转八进制
# print(oct(11))  # 0o13  0o代表二进制    # 1*8^1 + 3*8^0
# 十进制转十六进制
# print(hex(11))  # 0xb  0x代表二进制
# 二进制转十进制
# 1011 -> 1*2^3 + 0*2^2 + 1*2^1 + 1*2^0 = 8+0+2+1
# print(int(0b1011))
# 二进制转八进制
# print(int(0o13))
# 二进制转十六进制
# print(int(0xb))

# 使用

# 二、float类型
# 作用：记录薪资、身高、体重、小数等等
# 定义:
salary = 15.5  # salary = float(15.5)
# 类型转换:
res = float("15.5")
print(res, type(res))
# 使用:
# int 与 float没有需要掌握的内置方法
# 他们的使用就是数学运算+比较运算


from urllib.parse import urlparse, urlencode, urlunparse
from contextlib import contextmanager

def build_url(baseurl, path, args_dict=None):
    # Returns a list in the structure of urlparse.ParseResult
    url_parts = list(urlparse(baseurl))
    url_parts[2] = path
    if args_dict:
        url_parts[4] = urlencode(args_dict)
    return urlunparse(url_parts)

def slack_link(text, url):
    return '<{}|{}>'.format(url, text)

@contextmanager
def db_session_scope(session_maker):
    if not session_maker:
        raise RuntimeError('no db_session_maker specified for this command')

    session = session_maker()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


#!/usr/bin/env python3

import argparse as arg
import os.path as path
import os
import shutil
import subprocess

import scipy.io as sio
import numpy as np
from numpy import linalg as LA
import itertools
import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
import pandas as pd

def covariance_spread(x, y, ref=None, normalized=False):
    if ref is None:
        ref = np.identity(2)
    cov = np.cov(x, y)
    if normalized:
        cov = cov / LA.norm(cov, 2)
    return LA.norm(normalized - ref)

def eeg_covariance_spreads(eeg, window=10, frames=60, normalized=False):
    N = 400 * window
    start = 0
    stop = N - 1
    window_starts = (np.linspace(0, 10*60 - 2 * window, frames) * 400)\
        .astype(int)

    # Covariance over the whole data
    ref = {(i, j): np.cov(eeg.ix[start:window_starts[-1] + stop, i],
                          eeg.ix[start:window_starts[-1] + stop, j])
           for (i, j) in itertools.combinations(range(16), 2)}

    # Not sure if this have a meaning
    if normalized:
        ref = {k: c / LA.norm(c, 2) for k, c in ref.items()}

    data = {(i, j): [covariance_spread(eeg.ix[ws + start:ws + stop, i],
                                       eeg.ix[ws + start:ws + stop, j],
                                       ref[(i, j)])
                     for ws in window_starts]
            for (i, j) in itertools.combinations(range(16), 2)}
    return pd.DataFrame(data)

if __name__ == "__main__":
    parser = arg.ArgumentParser()
    parser.add_argument('matfile')
    parser.add_argument('-w', '--window', type=float,
                        default=10, help='window width')
    parser.add_argument('-f', '--frames', type=int,
                        default=60, help='number of frames')
    parser.add_argument('-n', '--normalized', action='store_true',
                        default=False,
                        help='whether matrices should be normalized to 1')
    opts = parser.parse_args()

    basename, _ = path.splitext(path.basename(opts.matfile))

    data = sio.loadmat(opts.matfile,
                       struct_as_record=False, squeeze_me=True)['dataStruct']
    eeg = pd.DataFrame(data.data)

    cv = eeg_covariance_spreads(eeg, window=opts.window,
                                frames=opts.frames,
                                normalized=opts.normalized)
    cv.to_csv("cor-spread-{0:s}.csv".format(basename), header=False, index=False)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 20 11:20:47 2020

@author: amelielaurens

Source : A simple model for nanofiber formation by rotary jet-spinning
by mellado et al.

Goal : Verify the model by comparing the results obtained with the model and
       the ones obtained with the experiment

Inputs : Concerning the polymer : the surface tension, the density, the viscosity
         Concerning the machine : the collector radius, the orifice radius,
         the angular viscosity of the spinneret
         The initial velocity and the discretisation

All data are in SI units.
"""

from models_RJS import *
import numpy
import matplotlib.pyplot as plt

discretisation = 20
# The higher the discretisation number is, the finer the discretisation will be,
# there will be more points on the graphic.

# Values used by Mellado
Rc = 0.135
rho = 1540
surface_tension = 0.027
initial_velocity = 0.1
orifice_radius = 0.000170
mu = 0.113
omega = 200

# X Axis : the scaling factor
scaling_factor = numpy.linspace(1.0, 2.6, discretisation)

# Y Axis : radius ratio
radius_ratio = []
for k in range(discretisation):
    Sigma = sigma(surface_tension, Rc, orifice_radius, initial_velocity)
    # Numerator
    r_num = Radius(orifice_radius, rho, initial_velocity, Rc, mu, Sigma,
                   scaling_factor[k]*omega)
    # Denominator
    r_deno = Radius(orifice_radius, rho, initial_velocity, Rc, mu, Sigma,
                    omega)
    radius_ratio.append(r_num/r_deno)
radius_ratio = numpy.array(radius_ratio)

# Plot the figure
fig = plt.figure()
axes = fig.add_subplot(1, 1, 1)

# Plot the points of the model
axes.plot(scaling_factor, radius_ratio, 'ro', label='model')

# Plot the points of the experiment
x = [1.0, 1.1, 1.3, 1.5, 1.8, 2.1, 2.4]
y = [1.0, 0.88, 0.79, 0.76, 0.69, 0.68, 0.64]

axes.plot(x, y, 'bx', label='experiment')

axes.grid()
axes.set_xlabel("Scaling factor", fontsize=16)
axes.set_ylabel("Radius ratio", fontsize=16)
axes.set_title("Comparison of the model with experiment", fontsize=16, y=1.)

axes.legend()
plt.show()


from flask import Flask, render_template, request, redirect, url_for
from flask.ext.script import Manager
from flask.ext.bootstrap import Bootstrap
from flask.ext.moment import Moment
from flask.ext.wtf import Form
from wtforms import StringField, SubmitField
from wtforms.validators import Required
from wtforms import StringField, TextAreaField, SubmitField, validators, ValidationError

app = Flask(__name__)
app.config['SECRET_KEY'] = 'hard to guess string'

manager = Manager(app)
bootstrap = Bootstrap(app)
moment = Moment(app)

class tuharChacha(Form):
    name = StringField("Name",  [Required()])
    email = StringField("Email",  [Required()])
    subject = StringField("Subject",  [Required()])
    message = TextAreaField("Message",  [Required()])
    submit = SubmitField("Send")
    #print "---------------NameForm--------------------------"
    #print "|\n\n|{}\n||{}\n\n|".format(name, submit)
    #print "|\n\n|{}\n||{}\n\n|".format(type(name), type(submit))
    #print "---------------NameForm--------------------------"

@app.errorhandler(404)
def page_not_found(e):
    return render_template('404.html'), 404

@app.errorhandler(500)
def internal_server_error(e):
    return render_template('500.html'), 500

@app.route('/', methods=['GET', 'POST'])
def index():
    name = None
    form = tuharChacha()
    if request.method == 'GET':
        return render_template('index.html', form=form, name=name)
    if request.method == 'POST':
        name = form.name.data
        form.name.data = ''
        return render_template('index.html', form=form, name=name)
    '''if form.validate_on_submit():
        name = form.name.data
        form.name.data = ''
        return render_template('index.html', form=form, name=name)'''

@app.route('/form/', methods=['POST'])
def rdct():
    return redirect(url_for('index'))

if __name__ == '__main__':
   app.run(debug=True, host='0.0.0.0', port=6665)


k#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun 24 19:43:17 2019

@author: vidhi

"""

import pymc3 as pm
import pandas as pd
#import numpy as np
import autograd.numpy as np
from autograd import elementwise_grad, jacobian, grad
import taylor_vi as tv
import theano.tensor as tt
import matplotlib.pylab as plt
import warnings
warnings.filterwarnings("ignore")

varnames = ['s_1', 'ls_2','s_3', 'ls_4','ls_5','s_6','ls_7','alpha_8','s_9','ls_10','n_11']

# Analytical variational inference for CO_2 data

def kernel(theta, X1, X2):

     # se +  sexper + rq + se

     s_1 = theta[0]
     ls_2 = theta[1]
     s_3 = theta[2]
     ls_4 = theta[3]
     ls_5 = theta[4]
     s_6 = theta[5]
     ls_7 = theta[6]
     alpha_8 = theta[7]

     sqdist = np.sum(X1**2, 1).reshape(-1, 1) + np.sum(X2**2, 1) - 2 * np.dot(X1, X2.T)
     dist = np.abs(np.sum(X1,1).reshape(-1,1) - np.sum(X2,1))

     sk1 = s_1**2 * np.exp(-0.5 / ls_2**2 * sqdist)
     sk2 = s_3**2 * np.exp(-0.5 / ls_4**2 * sqdist) * np.exp(-2*(np.sin(np.pi*dist)/ls_5)**2)
     sk3 = s_6**2 * np.power(1.0 + 0.5*(sqdist / (alpha_8 * ls_7**2)), -1*alpha_8)

     return sk1 + sk2 + sk3

def get_kernel_matrix_blocks(X, X_star, n_train, theta):

      s_1 = theta[0]
      ls_2 = theta[1]
      s_3 = theta[2]
      ls_4 = theta[3]
      ls_5 = theta[4]
      s_6 = theta[5]
      ls_7 = theta[6]
      alpha_8 = theta[7]
      s_9 = theta[8]
      ls_10 = theta[9]
      n_11 = theta[10]

      k1 = pm.gp.cov.Constant(s_1**2)*pm.gp.cov.ExpQuad(1, ls_2)
      k2 = pm.gp.cov.Constant(s_3**2)*pm.gp.cov.ExpQuad(1, ls_4)*pm.gp.cov.Periodic(1, period=1, ls=ls_5)
      k3 = pm.gp.cov.Constant(s_6**2)*pm.gp.cov.RatQuad(1, alpha=alpha_8, ls=ls_7)
      k4 = pm.gp.cov.Constant(s_9**2)*pm.gp.cov.ExpQuad(1, ls_10) +  pm.gp.cov.WhiteNoise(n_11)

      cov_sig =  k1 + k2 + k3
      cov_noise = k4

      K = cov_sig(X)
      K_s = cov_sig(X, X_star)
      K_ss = cov_sig(X_star, X_star)
      K_noise = K + cov_noise(X)*tt.eye(n_train)
      K_inv = np.linalg.inv(K_noise.eval())
      return K.eval(), K_s.eval(), K_ss.eval(), K_noise.eval(), K_inv

def get_empirical_covariance(trace_df, varnames):

      #df = pm.trace_to_dataframe(trace)
      return pd.DataFrame(np.cov(trace_df[varnames], rowvar=False), index=varnames, columns=varnames)

def gp_mean(theta, X, y, X_star):

     s_9 = theta[8]
     ls_10 = theta[9]
     n_11 = theta[10]

     #sqdist = np.sum(X**2, 1).reshape(-1, 1) + np.sum(X_star**2, 1) - 2 * np.dot(X, X_star.T)
     sqdist_X =  np.sum(X**2, 1).reshape(-1, 1) + np.sum(X**2, 1) - 2 * np.dot(X, X.T)
     sk4 = s_9**2 * np.exp(-0.5 / ls_10**2 * sqdist_X) + n_11**2*np.eye(len(X))

     K = kernel(theta, X, X)
     K_noise = K + sk4*np.eye(len(X))
     K_s = kernel(theta, X, X_star)
     return np.matmul(np.matmul(K_s.T, np.linalg.inv(K_noise)), y)

def gp_cov(theta, X, y, X_star):

     s_9 = theta[8]
     ls_10 = theta[9]
     n_11 = theta[10]

     #sqdist = np.sum(X**2, 1).reshape(-1, 1) + np.sum(X_star**2, 1) - 2 * np.dot(X, X_star.T)
     sqdist_X =  np.sum(X**2, 1).reshape(-1, 1) + np.sum(X**2, 1) - 2 * np.dot(X, X.T)
     sk4 = s_9**2 * np.exp(-0.5 / ls_10**2 * sqdist_X) + n_11**2*np.eye(len(X))

     K = kernel(theta, X, X)
     K_noise = K + sk4*np.eye(len(X))
     K_s = kernel(theta, X, X_star)
     K_ss = kernel(theta, X_star, X_star)
     return K_ss - np.matmul(np.matmul(K_s.T, np.linalg.inv(K_noise)), K_s)

if __name__ == "__main__":

      # Loading data

      uni_path = '/home/vidhi/Desktop/Workspace/CoreML/GP/Hierarchical GP/'
      mac_path = '/Users/vidhi.lalchand/Desktop/Workspace/CoreML/GP/Hierarchical GP/'
      desk_home_path = '/home/vr308/Desktop/Workspace/CoreML/GP/Hierarchical GP/'

      path = mac_path

      data_path = path + 'Data/Co2/'
      results_path = path + 'Results/Co2/'

      # Load Co2 data

      df = pd.read_table(data_path + 'mauna.txt', names=['year', 'co2'], infer_datetime_format=True, na_values=-99.99, delim_whitespace=True, keep_default_na=False)

      # creat a date index for the data - convert properly from the decimal year

      #df.index = pd.date_range(start='1958-01-15', periods=len(df), freq='M')

      df.dropna(inplace=True)

      mean_co2 = df['co2'][0]
      std_co2 = np.std(df['co2'])

      # normalize co2 levels

      #y = normalize(df['co2'])
      y = df['co2']
      t = df['year'] - df['year'][0]

      sep_idx = 545

      y_train = y[0:sep_idx].values
      y_test = y[sep_idx:].values

      t_train = t[0:sep_idx].values[:,None]
      t_test = t[sep_idx:].values[:,None]

      # Read in trace_fr

      #trace_fr_df = pm.trace_to_dataframe(trace_fr)

      trace_fr_df = pd.read_csv(results_path + '/trace_fr_df.csv', sep=',')

      # Read in fr_param

      fr_df_raw = pd.read_csv(results_path + 'VI/fr_df_raw.csv', sep=',', index_col=0)

      mu_theta = fr_df_raw['mu_implicit'][varnames] #maybe update
      cov_theta = get_empirical_covariance(trace_fr_df, varnames)

      # Analytical mean and var

      theta = np.array(mu_theta)

      dh = elementwise_grad(gp_mean)
      d2h = jacobian(dh)
      dg = grad(gp_cov)
      d2g = jacobian(dg)

      mu_taylor, std_taylor = tv.get_vi_analytical(t_train, y_train, t_test, dh, d2h, d2g, theta, mu_theta, cov_theta)

      #mu_taylor = pred_ng_mean
      #std_taylor = [np.sqrt(x) for x in pred_ng_var]

      rmse_taylor = pa.rmse(mu_taylor, y_test)
      se_rmse_taylor = pa.se_of_rmse(mu_taylor, y_test)
      lppd_mf, lpd_mf = pa.log_predictive_density(y_test, mu_taylor, std_taylor)

      print('rmse_mf:' + str(rmse_mf))
      print('se_rmse_mf:' + str(se_rmse_mf))
      print('lpd_mf:' + str(lpd_mf))

      #sample_mcvi_means = pd.read_csv(results_path + 'pred_dist/means_fr.csv', sep=',')
      #sample_mcvi_stds = pd.read_csv(results_path + 'pred_dist/std_fr.csv', sep=',')
      #lower_fr, upper_fr = pa.get_posterior_predictive_uncertainty_intervals(sample_mcvi_means, sample_mcvi_stds)

      # Plotting

#      plt.figure()
#      plt.plot(df['year'][sep_idx:], df['co2'][sep_idx:], 'ko', markersize=1)
#      plt.plot(df['year'][sep_idx:], np.mean(sample_mcvi_means), alpha=1, label='MCVI', color='g')
#      plt.plot(df['year'][sep_idx:], pred_ng_mean, alpha=1, label='Analytical VI', color='b')
#      plt.fill_between(df['year'][sep_idx:], lower_fr, upper_fr, color='green', alpha=0.5)
#      plt.fill_between(df['year'][sep_idx:], (pred_ng_mean - 1.96*np.sqrt(pred_ng_var)), (pred_ng_mean + 1.96*np.sqrt(pred_ng_var)), color='b', alpha=0.3)
#      plt.plot(df['year'][sep_idx:], mu_test, alpha=1, label='Type II ML', color='r')
#      plt.fill_between(df['year'][sep_idx:], (mu_test - 1.96*std_test), (mu_test + 1.96*std_test), color='red', alpha=0.3)
#      plt.legend(fontsize='x-small')
#      plt.title('Co2 - Test Predictions')


grades_tuple = ('sum_t','abc',80,47,85)
grades_list = ['sum_l','abc',80,47,85] # lists are ordered

grades_list.append('hello') #adds an item at the end
grades_list.reverse()

#grades_tuple.append(12) tuples are immutable
grades_tuple=grades_tuple+(100,109)

grades_set = {70,80,90,101,101,70} # sets are unordered and unique
grades_list[0]='hell'
grades_set.add('xyz')
print((grades_list[0],grades_tuple[0],grades_set))

#set operations

set_one = {1,2,3,4,5}
set_two = {5,6,7,8,1}
print(set_one.intersection(set_two))
print(set_one.issuperset(set_two))
print(set_one.difference(set_two))
print(set_one.union(set_two).difference(set_one.intersection(set_two)))


# def message(name, age):
#     if name == 'hayley':
#             print('hayley is great')
#     else:
#         print(name, 'is', age, 'years old')

# message('hayley', 22)
# message('lj', 1000)
# message ('marc', 119)
# message('james', 12)

def subtract (num1, num2):
    print(num1, num2)
    return num1 - num2

num1 = 10
num2 = 2
print (num1,num2)
subtract(3,6)

# print(subtract(10, 2))

# num3 = subtract(34, 23)

# print(num3)


from data_utils import load_PROTEST
from data_utils import load_TARGET
from data_utils import save_data
import scipy.misc
import numpy as np
import glob
from six.moves import cPickle as pickle
import config

def resize_arr(data_arr,H,W):
    N,_,_,D = data_arr.shape

    new_arr = list()
    for i in range(N):
        new_arr.append(scipy.misc.imresize(data_arr[i], (H,W), interp='bilinear'))
    arr = np.stack(new_arr,axis=0)
    return arr

def resize_data(data, size_tuple):
    H, W = size_tuple
    new_data = dict()
    #data is a hash map
    for key in data:
        if key[0]=='y':
            new_data[key] = data[key]
        else:
            new_data[key] = resize_arr(data[key],H,W)

    return new_data

#loop through batches to resize
def convert_all_data():
    batch_files = glob.glob(config.ALL_DATA+"*.save")
    for batch_file in batch_files:
        raw_batch = load_PROTEST(filename=batch_file)
        resized_batch = resize_data(raw_batch, (224,224))
        suffix = batch_file.split("/")[-1]
        new_filename = "%s%s" % (config.ALL_DATA_224,suffix)
        f = open(new_filename, 'wb')
        pickle.dump(resized_batch, f, protocol=pickle.HIGHEST_PROTOCOL)
        f.close()
        exit()

#loop through batches to resize
def convert_all_data_by_filenames(count=33):
    prefix = "/DATA/SHARE/datasets/all_data/"
    suffix = "_batch.save"

    batchno = 11
    num_converted = 0

    while num_converted < count:
        batch_file = "%s%d%s" % (prefix,batchno,suffix)
        #skip number 9
        if batchno == 9:
            batchno += 1
            continue
        raw_batch = load_PROTEST(filename=batch_file)
        resized_batch = resize_data(raw_batch, (224,224))
        new_filename = "%s%d%s" % (config.ALL_DATA_224,batchno,suffix)
        f = open(new_filename, 'wb')
        pickle.dump(resized_batch, f, protocol=pickle.HIGHEST_PROTOCOL)
        f.close()

        print("Converted batch number: %d" % batchno)
        num_converted += 1
        batchno += 1

convert_all_data_by_filenames()

#raw_data=load_PROTEST(20)
#raw_data = load_TARGET()
#new_data = resize_data(raw_data, (224, 224))
#save_data(new_data,"20_224")


#Program plots the 2-meter surface temperature

from pylab import *
import numpy
import datetime
import time
import glob, os
import math
import netCDF4 as netcdf
import matplotlib.colors as colors
from scipy import stats
from mpl_toolkits.basemap import Basemap

#Making pathway to folder with all data
directory 		= '../../../Data/UH-CESM'

#-----------------------------------------------------------------------------------------
#--------------------------------MAIN SCRIPT STARTS HERE----------------------------------
#-----------------------------------------------------------------------------------------

month_start	= 6 	#1 = January, 2 = February, 3 = March, ..., 13 = January (+ 1), ...
month_end	= 11	#12 = December, 13 = January (+ 1), 14 = February (+ 1), ...

temp_present		= ma.masked_all((5, 280, 354))
temp_future			= ma.masked_all((5, 280, 354))
temp_global_present	= ma.masked_all(5)
temp_global_future	= ma.masked_all(5)
#-----------------------------------------------------------------------------------------

for period in ['PRESENT', 'FUTURE']:
	#Loop over both periods
	if period == 'PRESENT':
		ensemble_number	= [1, 2, 3, 4, 5]

	if period == 'FUTURE':
		ensemble_number	= [6, 7, 8, 9, 10]

	for ensemble_i in range(len(ensemble_number)):
		#Loop over each ensemble
		print(ensemble_i)
		fh 			= netcdf.Dataset(directory+'_'+str(ensemble_number[ensemble_i]).zfill(3)+'/Atmosphere/TEMP_2m_NA_month_'+str(month_start)+'-'+str(month_end)+'.nc', 'r')

		lon			= fh.variables['lon'][:]
		lat			= fh.variables['lat'][:]
		temp		= fh.variables['TEMP'][:]
		temp_global	= fh.variables['TEMP_global'][:]

		if ensemble_number[ensemble_i] <= 5:
			#Present-day ensemble
			temp_present[ensemble_i]		= temp
			temp_global_present[ensemble_i]	= temp_global

		else:
			#Future ensemble
			temp_future[ensemble_i]			= temp
			temp_global_future[ensemble_i]	= temp_global

#Take ensemble mean
temp_present_plot	= np.mean(temp_present, axis = 0)
temp_future_plot	= np.mean(temp_future, axis = 0)
temp_global_present	= np.mean(temp_global_present, axis = 0)
temp_global_future	= np.mean(temp_global_future, axis = 0)
temp_future_plot	= temp_future_plot - temp_present_plot - (temp_global_future - temp_global_present)
print('Global mean 2-meter surface temperature increase:', temp_global_future - temp_global_present)

#-----------------------------------------------------------------------------------------

fig, ax	= subplots()

m = Basemap(projection = 'merc', llcrnrlat=-3, urcrnrlat=61, llcrnrlon=250, urcrnrlon=360.0001, resolution='i')
m.drawcoastlines()

par = m.drawparallels(np.arange(-80,80,10),labels=[1,0,0,0])
mer = m.drawmeridians(np.arange(0,360.1,20),labels=[0,0,0,1])

x, y	= np.meshgrid(lon, lat)
x, y	= m(x, y)
CS	= contourf(x, y, temp_future_plot, levels = np.arange(-3, 3.1, 0.1), extend = 'both', cmap = 'RdBu_r')
cbar	= m.colorbar(CS, ticks = np.arange(-3, 3.1, 1))
cbar.set_label('2-meter surface temperature difference ($^{\circ}$C)')

ax.set_title('c) 2-meter surface temperature, $\Delta$UH-CESM')

show()


import numpy as np
from glob import glob
import matplotlib.pyplot as plt
from sklearn import linear_model, datasets,svm
from sklearn.decomposition import PCA

data_path = '/Users/Naji/Box Sync/CMU/Masters Thesis/Data/'
dirs_list = glob(data_path+'*')
realisations_dict = dict()
is_labeled_dict = dict()
num_of_realisations = 5
database_names = []
for dir in dirs_list:
    files_list = glob(dir+'/*')
    # clearing up the data to only cinldude it's full name
    database_name  = dir.split('/')[-1][:-5]
    database_names.append(database_name)
    realisations_dict[database_name] = dict()
    realisations_dict[database_name]['train'] = []
    realisations_dict[database_name]['test'] = []

    is_labeled_dict[database_name] = dict()
    is_labeled_dict[database_name]['new_train'] = []
    is_labeled_dict[database_name]['test'] = []
    is_labeled_dict[database_name]['original_train'] = []

    for i in range(num_of_realisations):
        temp_test_data = np.loadtxt(dir+'/'+database_name+'_test_data_'+str(i+1)+'.asc' ,skiprows=6)
        temp_test_labels = np.loadtxt(dir+'/'+database_name+'_test_labels_'+str(i+1)+'.asc' ,skiprows=6)
        temp_train_data = np.loadtxt(dir+'/'+database_name+'_train_data_'+str(i+1)+'.asc' ,skiprows=6)
        temp_train_labels = np.loadtxt(dir+'/'+database_name+'_train_labels_'+str(i+1)+'.asc' ,skiprows=6)
        if i==0:
            train_size = temp_train_data.shape[0]
            test_size = temp_test_data.shape[0]
            train_test_ratio = train_size  / float(test_size)
        realisations_dict[database_name]['test'].append([temp_test_data, temp_test_labels])
        realisations_dict[database_name]['train'].append([temp_train_data, temp_train_labels])
        temp_positive_train_labels_idx = temp_train_labels == 1
        temp_new_labeled_data = np.vstack((temp_train_data[temp_positive_train_labels_idx],temp_test_data[:int(temp_test_data.shape[0]*train_test_ratio)]))
        temp_new_labeled_labels= np.hstack((temp_train_labels[temp_positive_train_labels_idx],-np.ones(min(temp_test_data.shape[0],int(temp_test_data.shape[0]*train_test_ratio)))))
        is_labeled_dict[database_name]['new_train'].append([temp_new_labeled_data,temp_new_labeled_labels])
        is_labeled_dict[database_name]['original_train'].append([temp_train_data, temp_train_labels])
        is_labeled_dict[database_name]['test'].append([temp_test_data[int(temp_test_data.shape[0]*train_test_ratio):],temp_test_labels[int(temp_test_data.shape[0]*train_test_ratio):]])
plt_dict = dict()
def log_reg_plots():
    for database_name in database_names:
        logreg = linear_model.LogisticRegression(C=1e5)
        h = .02  # step size in the mesh
        X = is_labeled_dict[database_name]['new_train'][1][0]
        two_dim_X = PCA(n_components=2).fit_transform(X)
        Y = is_labeled_dict[database_name]['new_train'][1][1]

        # we create an instance of Neighbours Classifier and fit the data.
        logreg.fit(two_dim_X, Y)

        # Plot the decision boundary. For that, we will assign a color to each
        # point in the mesh [x_min, x_max]x[y_min, y_max].
        x_min, x_max = X[:, 0].min() - 2., X[:, 0].max() + 2.
        y_min, y_max = X[:, 1].min() - 2., X[:, 1].max() + 2.
        xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
        Z = logreg.predict(np.c_[xx.ravel(), yy.ravel()])

        # Put the result into a color plot
        Z = Z.reshape(xx.shape)
        plt.figure(1, figsize=(4, 3))
        plt.pcolormesh(xx, yy, Z, cmap=plt.cm.Paired)

        # Plot also the training points
        positive_idx = Y==1
        color_list=np.asarray(['b,r'])
        plt_dict[database_name] = [plt.scatter(two_dim_X[positive_idx, 0], two_dim_X[positive_idx, 1], c='r', marker='+', edgecolors='k',cmap=plt.cm.coolwarm),
                                    plt.scatter(two_dim_X[~positive_idx, 0], two_dim_X[~positive_idx, 1], c='b', marker='+', edgecolors='k',cmap=plt.cm.coolwarm)]
        plt.xlabel('X_0')
        plt.ylabel('X_1')
        plt.legend((plt_dict[database_name][0], plt_dict[database_name][1]), ('Positive','Negative'))
        plt.xlim(xx.min(), xx.max())
        plt.ylim(yy.min(), yy.max())
    # plt.xticks(())
    # plt.yticks(())
        plt.savefig('/Users/Naji/Box Sync/CMU/Masters Thesis/Plots/'+database_name+'.png')
    # plt.show()
def SVM_plots(kernel='rbf'):
    for database_name in database_names:
        h = .02  # step size in the mesh
        X = is_labeled_dict[database_name]['new_train'][1][0]
        if X.shape[1]>2:
            two_dim_X = PCA(n_components=2).fit_transform(X)
        else:
            two_dim_X = X
        Y = is_labeled_dict[database_name]['new_train'][1][1]
        clf = svm.SVC(kernel=kernel)
        clf.fit(two_dim_X, Y)

        plt.axis('tight')
        x_min = two_dim_X[:, 0].min()
        x_max = two_dim_X[:, 0].max()
        y_min = two_dim_X[:, 1].min()
        y_max = two_dim_X[:, 1].max()

        XX, YY = np.mgrid[x_min:x_max:200j, y_min:y_max:200j]
        Z = clf.decision_function(np.c_[XX.ravel(), YY.ravel()])

        # Put the result into a color plot
        Z = Z.reshape(XX.shape)
        plt.pcolormesh(XX, YY, Z > 0, cmap=plt.cm.Paired)
        # plt.contour(XX, YY, Z, colors=['k', 'k', 'k'], linestyles=['--', '-', '--'],
        #             levels=[-.5, 0, .5])

        plt.title(database_name)

        # Plot also the training points
        positive_idx = Y==1
        color_list=np.asarray(['b,r'])
        plt_dict[database_name] = [plt.scatter(two_dim_X[positive_idx, 0], two_dim_X[positive_idx, 1], c='r', marker='+', edgecolors='k',cmap=plt.cm.coolwarm),
                                    plt.scatter(two_dim_X[~positive_idx, 0], two_dim_X[~positive_idx, 1], c='b', marker='+', edgecolors='k',cmap=plt.cm.coolwarm)]
        plt.xlabel('X_0')
        plt.ylabel('X_1')
        plt.legend((plt_dict[database_name][0], plt_dict[database_name][1]), ('Positive','Negative'))
        plt.xlim(x_min, x_max)
        plt.ylim(y_min, y_max)
        plt.savefig('/Users/Naji/Box Sync/CMU/Masters Thesis/Plots/'+database_name+'.png')
SVM_plots()


import numpy as np
import pandas as pd
import statsmodels.api as sm
import scipy as sp
import scipy.stats as st
import sys
import warnings
import inspect
import patsy
from decimal import Decimal
from IPython.core.display import display, HTML
from scipy.stats import poisson
from scipy.stats import nbinom

FLOAT_EPS = np.finfo(float).eps
pd.options.display.float_format = '{:,.12f}'.format

__all__ = ['ZeroInflated']

class LinkClass(object):
    def __init__(self):
        return NotImplementedError
    def link(self, mu):
        return NotImplementedError
    def link_inv(self, eta):
        return NotImplementedError
    def link_inv_deriv(self, eta):
        return NotImplementedError

class Logit(LinkClass):
    def __init__(self):
        self.linkclass = sm.genmod.families.links.logit
    def link(self, p):
        return np.log(p/(1.0-p))
        #return sp.special.logit(p)
    def link_inv(self, eta):
        thresh = 30.0
        eta = np.minimum(np.maximum(eta,-thresh), thresh)
        exp_eta = np.exp(-eta)
        return 1.0/(1.0+exp_eta)
        #return sp.special.expit(eta)
    def link_inv_deriv(self, eta):
        thresh = 30.0
        eta[abs(eta) > thresh] = FLOAT_EPS
        return np.exp(eta)/(1+np.exp(eta))**2
    def __repr__(self):
        display_string = f"\n    linkstr: logit"
        display_string += '\n    link: log(p/(1-p))'
        display_string += '\n    linkinv: exp(eta)/(1+exp(eta))'
        return display_string

class Probit(LinkClass):
    def __init__(self):
        self.linkclass = sm.genmod.families.links.probit
    def link(self, mu):
        return st.norm.ppf(mu)
    def link_inv(self, eta):
        thresh = -st.norm.ppf(FLOAT_EPS)
        eta = np.minimum(np.maximum(eta,-thresh),thresh)
        return st.norm.cdf(eta)
    def link_inv_deriv(self, eta):
        return np.maximum(st.norm.pdf(eta),FLOAT_EPS)
    def __repr__(self):
        display_string = f"\n    linkstr: probit"
        display_string += '\n    link: norm.ppf(mu)'
        display_string += '\n    linkinv: norm.cdf(eta)'
        return display_string

class CLogLog(LinkClass):
    def __init__(self):
        self.linkclass = sm.genmod.families.links.cloglog
    def link(self, mu):
        return np.log(-np.log(1 - mu))
    def link_inv(self, eta):
        return np.maximum(np.minimum(-np.expm1(-np.exp(eta)),1-FLOAT_EPS),FLOAT_EPS)
    def link_inv_deriv(self, eta):
        eta = np.minimum(eta,700)
        return np.maximum(np.exp(eta)*np.exp(-np.exp(eta)),FLOAT_EPS)
    def __repr__(self):
        display_string = f"\n    linkstr: cloglog"
        display_string += '\n    link: log(-log(1 - mu))'
        display_string += '\n    linkinv: 1-exp(-exp(eta))'
        return display_string

class Cauchit(LinkClass):
    def __init__(self):
        self.linkclass = sm.genmod.families.links.cauchy
    def link(self, mu):
        return st.cauchy.ppf(mu)
    def link_inv(self, eta):
        thresh = -st.cauchy.ppf(FLOAT_EPS)
        eta = np.minimum(np.maximum(eta,-thresh),thresh)
        return st.cauchy.cdf(eta)
    def link_inv_deriv(self, eta):
        return nnp.maximum(st.cauchy.pdf(eta),FLOAT_EPS)
    def __repr__(self):
        display_string = f"\n    linkstr: cauchit"
        display_string += '\n    link: cauchy.ppf(mu)'
        display_string += '\n    linkinv: cauchy.cdf(eta)'
        return display_string

class Log(LinkClass):
    def __init__(self):
        self.linkclass = sm.genmod.families.links.log
    def link(self, mu):
        return np.log(mu)
    def link_inv(self, eta):
        return np.maximum(np.exp(eta), FLOAT_EPS)
    def link_inv_deriv(self, eta):
        return np.maximum(np.exp(eta), FLOAT_EPS)
    def __repr__(self):
        display_string = f"\n    linkstr: log"
        display_string += '\n    link: log(mu)'
        display_string += '\n    linkinv: exp(eta)'
        return display_string

class ZeroInflated(object):
    __doc__ = """
    Zero Inflated model for count data
    %(params)s
    %(extra_params)s
    Attributes
    -----------
    formula_str : string
        A reference to the endogenous response variable.
    data : pandas dataframe
        A reference to the exogenous design.
    dist: string
        A reference to the zero-inflated exogenous design.
    link: string
        A reference to
    """
    def __init__(self, formula_str, data, dist = 'poisson', offsetx = None, offsetz = None,
                 link = 'logit', weights = None, missing='none', **kwargs):
        self._set_data(formula_str, data, missing)
        self.terms = {'Y':self.endog.columns.values[0],'X':self.X.columns.values,\
                      'Z':self.Z.columns.values}
        self.formula = formula_str
        self.dist = self._dist_processing(dist)
        self.link = self._link_processing(link)
        self.n = len(self.endog)
        self._set_wt_offset(weights, offsetx, offsetz)
        self.linkobj = self._LinkClass_processing(self.link)
        self._set_loglik(self.dist)
        self.call = f"ZeroInflated(formula_str='{formula_str}', data={self._retrieve_name(data)}, dist='{dist}', offsetx={offsetx}, offsetz={offsetz},"
        self.call = self.call + f" link='{link}', weights={weights}, missing='{missing}')"

        # Convenience variables
        self.kx = self.X.shape[1]
        self.kz = self.Z.shape[1]
        self.Y = np.squeeze(self.endog)
        self.Y0 = self.Y <= 0
        self.Y1 = self.Y > 0
        self.EM = True
        self.reltol = (np.finfo(float).eps)**(1/1.6)

    def print_obj(self):
        print(self)

    def _retrieve_name(self, var):
        """
        Gets the name of var. Does it from the out most frame inner-wards.
        :param var: variable to get name from.
        :return: string
        """
        for fi in reversed(inspect.stack()):
            names = [var_name for var_name, var_val in fi.frame.f_locals.items() if var_val is var]
            if len(names) > 0:
                return names[-1]

    def _set_data(self, formula_str, data, missing):
        self.endog, self.X, self.Z = self.formula_processing(formula_str, data, missing=missing)

    def _set_wt_offset(self, weights, offsetx, offsetz):
        ## weights and offset

        if weights is None:
            weights = 1.0
        weights = np.ndarray.flatten(np.array(weights))
        if weights.size == 1:
            weights = np.repeat(weights,self.n)
        weights = pd.Series(data = weights, index = self.X.index)

        if offsetx is None:
            offsetx = 0.0
        offsetx = np.ndarray.flatten(np.array(offsetx))
        if offsetx.size == 1:
            offsetx = np.repeat(offsetx,self.n)

        if offsetz is None:
            offsetz = 0.0
        offsetz = np.ndarray.flatten(np.array(offsetz))
        if offsetz.size == 1:
            offsetz = np.repeat(offsetz,self.n)

        self.weights = weights
        self.offsetx = offsetx
        self.offsetz = offsetz

    def _set_loglik(self, dist):
        if dist is 'poisson':
            self.loglikfun = self.ziPoisson
            self.gradfun = self.gradPoisson
        elif dist is 'negbin':
            self.loglikfun = self.ziNegBin
            self.gradfun = self.gradNegBin
        else:
            self.loglikfun = self.ziGeom
            self.gradfun = self.gradGeom

    def ziPoisson(self, parms, sign = 1.0):
        """
        Log-likelihood of Zero Inflated Poisson.
        """

        ## count mean
        mu = np.exp(np.dot(self.X,parms[np.arange(self.kx)]) + self.offsetx)
        ## binary mean
        phi = self.linkobj.link_inv(np.dot(self.Z, parms[np.arange(self.kx,self.kx+self.kz)]) +\
                                    self.offsetz)
        ## log-likelihood for y = 0 and y >= 1
        loglik0 = np.log(phi + np.exp(np.log(1-phi) - mu))
        loglik1 = np.log(1-phi) + sp.stats.poisson.logpmf(self.Y, mu)
        ## collect and return
        loglik = np.dot(self.weights[self.Y0],loglik0[self.Y0])+np.dot(self.weights[self.Y1],loglik1[self.Y1])
        return sign*loglik

    def gradPoisson(self, parms, sign = 1.0):
        """
        Gradient of Zero Inflated Poisson Log-likelihood.
        """

        ## count mean
        eta = np.dot(self.X,parms[np.arange(self.kx)]) + self.offsetx
        mu = np.exp(eta)
        ## binary mean
        etaz = np.dot(self.Z, parms[np.arange(self.kx,self.kx+self.kz)]) + self.offsetz
        muz = self.linkobj.link_inv(etaz)
        ## densities at 0
        clogdens0 = -mu
        dens0 = muz*(1-self.Y1.astype(float)) + np.exp(np.log(1 - muz) + clogdens0)
        ## working residuals
        wres_count = np.where(self.Y1,self.Y-mu,-np.exp(-np.log(dens0) +
                                          np.log(1 - muz) + clogdens0 + np.log(mu)))
        link_etaz = self.linkobj.link_inv_deriv(etaz)
        wres_zero  = np.where(self.Y1,-1/(1-muz) * link_etaz, \
                          (link_etaz - np.exp(clogdens0) * link_etaz)/dens0)

        return sign*(np.hstack((np.expand_dims(wres_count*self.weights,axis=1)*self.X, \
                np.expand_dims(wres_zero*self.weights,axis=1)*self.Z))).sum(axis=0)

    def ziNegBin(self, parms, sign = 1.0):
        """
        Log-Likelihood of Zero Inflated Negative Binomial.
        """
        ## count mean
        mu = np.exp(np.dot(self.X,parms[np.arange(self.kx)]) + self.offsetx)
        ## binary mean
        phi = self.linkobj.link_inv(np.dot(self.Z, parms[np.arange(self.kx,self.kx+self.kz)]) + self.offsetz)
        ## negbin size
        theta = np.exp(parms[self.kx+self.kz])

        ## log-likelihood for y = 0 and y >= 1 sp.stats.poisson.logpmf(Y, mu)
        loglik0 = np.log(phi + np.exp(np.log(1-phi) + \
                                   st.nbinom.logpmf(0,*self.convert_params(theta = theta, mu = mu)) ) )
        loglik1 = np.log(1-phi) + st.nbinom.logpmf(self.Y,*self.convert_params(theta = theta, mu = mu))

        ## collect and return
        loglik = np.dot(self.weights[self.Y0],loglik0[self.Y0])+np.dot(self.weights[self.Y1],loglik1[self.Y1])
        return sign*loglik

    def ziGeom(self, parms, sign = 1.0):
        return self.ziNegBin(np.hstack((parms, 0)), sign)

    def gradGeom(self, parms, sign = 1.0):
        """
        Gradient of Zero Inflated Geometric Log-likelihood.

        """
        ## count mean
        eta = np.dot(self.X,parms[np.arange(self.kx)]) + self.offsetx
        mu = np.exp(eta)
        ## binary mean
        etaz = np.dot(self.Z, parms[np.arange(self.kx,self.kx+self.kz)]) + self.offsetz
        muz = self.linkobj.link_inv(etaz)

        ## densities at 0
        clogdens0 = st.nbinom.logpmf(0,*self.convert_params(theta = 1, mu = mu))
        dens0 = muz*(1-self.Y1.astype(float)) + np.exp(np.log(1 - muz) + clogdens0)

        ## working residuals
        wres_count = np.where(self.Y1,self.Y - mu*(self.Y + 1)/(mu + 1), \
                              -np.exp(-np.log(dens0) + np.log(1 - muz) + clogdens0 +\
                                      -np.log(mu+1) + np.log(mu)))
        link_etaz = self.linkobj.link_inv_deriv(etaz)
        wres_zero  = np.where(self.Y1,-1/(1-muz) * link_etaz, \
                          (link_etaz - np.exp(clogdens0) * link_etaz)/dens0)

        return sign*(np.hstack((np.expand_dims(wres_count*self.weights,axis=1)*self.X, \
                np.expand_dims(wres_zero*self.weights,axis=1)*self.Z))).sum(axis=0)

    def gradNegBin(self, parms, sign = 1.0):
        """
        Gradient of Zero Inflated Negative Binomial Log-likelihood.
        (Negetive Binomial2 to be specific.)

        """
        ## count mean
        eta = np.dot(self.X,parms[np.arange(self.kx)]) + self.offsetx
        mu = np.exp(eta)
        ## binary mean
        etaz = np.dot(self.Z, parms[np.arange(self.kx,self.kx+self.kz)]) + self.offsetz
        muz = self.linkobj.link_inv(etaz)
        ## negbin size
        theta = np.exp(parms[self.kx+self.kz])

        ## densities at 0
        clogdens0 = st.nbinom.logpmf(0,*self.convert_params(theta = theta, mu = mu))
        dens0 = np.array(muz*(1-self.Y1.astype(float)) + np.exp(np.log(1 - muz) + clogdens0))

        ## working residuals
        wres_count = np.where(self.Y1,self.Y - mu*(self.Y + theta)/(mu + theta), \
                              -np.exp(-np.log(dens0) + np.log(1 - muz) + clogdens0 + np.log(theta) \
                                      -np.log(mu+theta) + np.log(mu)))
        link_etaz = self.linkobj.link_inv_deriv(etaz)
        wres_zero  = np.where(self.Y1,-1/(1-muz) * link_etaz, \
                          (link_etaz - np.exp(clogdens0) * link_etaz)/dens0)

        wres_theta = theta*np.where(self.Y1, sp.special.digamma(self.Y + theta) - sp.special.digamma(theta) +\
                                   np.log(theta) - np.log(mu + theta) + 1 - (self.Y + theta)/(mu + theta),\
                                   np.exp(-np.log(dens0) + np.log(1 - muz) + clogdens0)*\
                                   (np.log(theta) - np.log(mu + theta) + 1 - theta/(mu+theta) ) )

        return sign*(np.hstack((np.expand_dims(wres_count*self.weights,axis=1)*self.X, \
                np.expand_dims(wres_zero*self.weights,axis=1)*self.Z, \
                               np.expand_dims(wres_theta,axis=1)))).sum(axis=0)

    def EM_estimate(self):
        ## EM estimation of starting values

        model_count = sm.GLM(endog = self.Y, exog = self.X, family = sm.families.Poisson(),\
                                  offset = self.offsetx , freq_weights = self.weights).fit()
        model_zero = sm.GLM(self.Y0.astype(int), exog = self.Z, family=sm.families.Binomial(link = self.linkobj.linkclass), \
                   offset = self.offsetz , freq_weights = self.weights).fit()
        self.start = {'zero':model_zero.params, 'count':model_count.params}

        if self.dist is 'negbin':
            self.start['theta'] = 1.0

        if (self.EM is True) and (self.dist is 'poisson'):
            mui = model_count.predict()
            probi = model_zero.predict()
            probi = probi/(probi + (1-probi)*sp.stats.poisson.pmf(0, mui))
            probi[self.Y1] = 0
            probi
            ll_new = self.loglikfun(np.hstack((self.start['count'].values,self.start['zero'].values)))
            ll_old = 2 * ll_new

            while np.absolute((ll_old - ll_new)/ll_old) > self.reltol :
                ll_old = ll_new
                model_count = sm.GLM(endog = self.Y, exog = self.X, family = sm.families.Poisson(),\
                                  offset = self.offsetx , freq_weights = self.weights*(1-probi) \
                                              ).fit(start_params = self.start['count'].values)
                model_zero = sm.GLM(probi, exog = self.Z, family=sm.families.Binomial(link = self.linkobj.linkclass),\
                        offset = self.offsetz, freq_weights = self.weights \
                               ).fit(start_params = self.start['zero'].values)
                self.start = {'zero':model_zero.params, 'count':model_count.params}

                mui = model_count.predict()
                probi = model_zero.predict()
                probi = probi/(probi + (1-probi)*sp.stats.poisson.pmf(0, mui))
                probi[self.Y1] = 0

                ll_new = self.loglikfun(np.hstack((self.start['count'].values,self.start['zero'].values)))

        if (self.EM is True) and (self.dist is 'geom'):
            mui = model_count.predict()
            probi = model_zero.predict()
            probi = probi/(probi + (1-probi)*st.nbinom.pmf(0,*self.convert_params(theta = 1, mu = mui)))
            probi[self.Y1] = 0

            ll_new = self.loglikfun(np.hstack((self.start['count'].values,self.start['zero'].values)))
            ll_old = 2 * ll_new

            while np.absolute((ll_old - ll_new)/ll_old) > self.reltol :
                ll_old = ll_new
                model_count = sm.GLM(endog = self.Y, exog = self.X, family = sm.families.NegativeBinomial(alpha = 1.0),\
                                  offset = self.offsetx , freq_weights = self.weights*(1-probi)).fit(\
                                        start_params = self.start['count'].values)
                                    #sm.families.NegativeBinomial(alpha = 1.0\
                                    #                            ).starting_mu(y=self.start['count'].values))
                model_zero = sm.GLM(probi, exog = self.Z, family=sm.families.Binomial(link = self.linkobj.linkclass),\
                        offset = self.offsetz, freq_weights = self.weights).fit(start_params = self.start['zero'].values)
                self.start = {'zero':model_zero.params, 'count':model_count.params}

                mui = model_count.predict()
                probi = model_zero.predict()
                probi = probi/(probi + (1-probi)*st.nbinom.pmf(0,*self.convert_params(theta = 1, mu = mui)))
                probi[self.Y1] = 0

                ll_new = self.loglikfun(np.hstack((self.start['count'].values,self.start['zero'].values)))

        if (self.EM is True) and (self.dist is 'negbin'):
            warnings.warn('EM estimation of starting values not optimal for Negetive Binomial.')
            mui = model_count.predict() # or model_count.mu
            probi = model_zero.predict()
            probi = probi/(probi + (1-probi)*st.nbinom.pmf(0,*self.convert_params(theta = self.start['theta'], mu = mui)))
            probi[self.Y1] = 0

            ll_new = self.loglikfun(np.hstack((self.start['count'].values,self.start['zero'].values,np.log(self.start['theta']))))
            ll_old = 2 * ll_new

            while np.absolute((ll_old - ll_new)/ll_old) > self.reltol :
                ll_old = ll_new
                model_count = sm.GLM(endog = self.Y, exog = self.X, family = \
                                sm.families.NegativeBinomial(alpha = 1/self.start['theta']),method = 'newton',\
                                  offset = self.offsetx , freq_weights = self.weights*(1-probi) \
                                      ).fit(start_params = self.start['count'])
                model_zero = sm.GLM(probi, exog = self.Z, family=sm.families.Binomial(link = self.linkobj.linkclass),\
                        offset = self.offsetz, freq_weights = self.weights, \
                        start_params = self.start['zero']).fit()

                mui = model_count.predict()
                theta = model_count.scale

                probi = model_zero.predict()
                probi = probi/(probi + (1-probi)*st.nbinom.pmf(0,*self.convert_params(theta = theta, mu = mui)))

                probi[self.Y1] = 0
                self.start = {'zero':model_zero.params, 'count':model_count.params, 'theta':theta}

                ll_new = self.loglikfun(np.hstack((self.start['count'].values,\
                                self.start['zero'].values,np.log(self.start['theta']))))

        return self.start

    def fit(self, method = 'BFGS', EM = True, start = None, reltol = None,\
            options = {'disp': False, 'maxiter': 10000}, factr = 1.0):
        self.set_tolerance(factr, reltol)
        self.optim_options = options
        self.method = method
        self.EM = EM
        self.set_start(start)

        ## ML Estimation
        if (self.dist is 'negbin'):
            x0 = np.hstack((self.start['count'].values,self.start['zero'].values,\
                                         np.log(self.start['theta'])))
        else:
            x0 = np.hstack((self.start['count'].values,self.start['zero'].values))

        fitResult = sp.optimize.minimize(self.loglikfun, args=(-1.0,), x0 = x0, \
                                        method=self.method, jac=self.gradfun, options=self.optim_options)

        ## coefficients and covariances
        coefc = pd.Series(data = fitResult.x[0:self.kx], index = self.X.columns.values)
        coefz = pd.Series(data = fitResult.x[self.kx:self.kx+self.kz], index = self.Z.columns.values)

        if self.method == 'L-BFGS-B':
            vc_data = fitResult.hess_inv.todense()
        elif self.method == 'BFGS':
            vc_data = fitResult.hess_inv
        else:
            warnings.warn('Not tested for methods other than BFGS and L-BFGS-B.')

        vc = pd.DataFrame(data = vc_data[np.arange(self.kx+self.kz)[:,None],np.arange(self.kx+self.kz)], \
                      index = np.append(self.X.columns.values, self.Z.columns.values),\
                 columns = np.append(self.X.columns.values, self.Z.columns.values))
        if self.dist == 'negbin':
            ntheta = self.kx + self.kz
            theta = np.exp(fitResult.x[ntheta])
            SE_logtheta = np.sqrt(np.diagonal(vc_data)[ntheta])
        else:
            theta = None
            SE_logtheta = None

        ## fitted and residuals
        mu = np.exp(np.dot(self.X,coefc)+self.offsetx)
        phi = self.linkobj.link_inv(np.dot(self.Z,coefz)+self.offsetz)
        Yhat = (1-phi) * mu
        res = np.sqrt(self.weights) * (self.Y - Yhat)

        ## effective observations
        nobs = np.sum(self.weights > 0)

        Result = ZeroInflatedResults(self.call, self.formula, self.terms, self.kx, self.kz, \
                                     self.dist, self.link, self.linkobj, self.optim_options, self.method, self.start,\
                                     self.reltol, self.weights, self.offsetx, self.offsetz,\
                                     fitResult, coefc, coefz, theta, SE_logtheta, nobs, res, Yhat, vc, self.endog,\
                                    self.X, self.Z)
        return Result

    def set_start(self, start):
        if start is not None:
            valid = True
            if ('count' in start) is False:
                valid = False
                warnings.warn("invalid starting values, count model coefficients not specified")
                start['count'] = pd.Series(np.repeat(0,kx), index = X.columns.values)
            if ('zero' in start) is False:
                valid = False
                warnings.warn("invalid starting values, zero model coefficients not specified")
                start['zero'] = pd.Series(np.repeat(0,kz), index = Z.columns.values)
            if len(start['count']) != kx:
                valid = False
                warnings.warn("invalid starting values, wrong number of count model coefficients")
            if len(start['zero']) != kz:
                valid = False
                warnings.warn("invalid starting values, wrong number of zero model coefficients")
            if dist is 'negbin':
                if ('theta' in start) is False:
                    start['theta'] = 1.0
                start = {'zero':start['zero'], 'count':start['count'], 'theta' : (start['theta'][0]).astype(float)}
            else:
                start = {'zero':start['zero'], 'count':start['count']}

            if valid is False:
                start = None

        if start is None:
            self.EM_estimate()
        else:
            self.start = start

    def set_tolerance(self, factr=1.0, reltol = FLOAT_EPS**(1/1.6)):
        if factr < 1.0:
            warnings.warn('Minimum value of factr is 1.0.')
            factr = 1.0
        if reltol is None:
            self.reltol = factr*((np.finfo(float).eps)**(1/1.6))

    @staticmethod
    def formula_processing(formula_str, data, missing='none'):
        # ToDo: Add 'missing' operations on df
        X_formula,Z_formula = formula_str.split("|")
        Z_formula = X_formula.split("~")[0]+" ~ "+ Z_formula
        y, X = patsy.dmatrices(X_formula, data, return_type='dataframe')
        Z = patsy.dmatrices(Z_formula, data, return_type='dataframe')[1]

        Y = np.squeeze(y)
        ## sanity checks
        if len(Y) < 1:
            sys.exit("empty model")
        if np.all(Y > 0):
            sys.exit("invalid dependent variable, minimum count is not zero")
        if np.array_equal(np.asarray(Y), (np.round(Y + 0.001)).astype(int)) is False:
            sys.exit("invalid dependent variable, non-integer values")
        Y = (np.round(y + 0.001)).astype(int)
        if np.any(Y < 0):
            sys.exit("invalid dependent variable, negative counts")

        return y,X,Z

    @staticmethod
    def convert_params(mu, theta):
        """
        Convert mean/dispersion parameterization of a negative binomial to the ones scipy supports

        """
        n = theta
        p = theta/(theta+mu)
        return n, p

    @staticmethod
    def _link_processing(link):
        ## binary link processing
        linkstr = link
        linkList = ['logit','probit','cauchit','cloglog','log']
        if linkstr not in linkList:
            warnings.warn(linkstr +" link not valid. Available links are: " + str(linkList))
            linkstr = 'logit'
        return(linkstr)

    @staticmethod
    def _LinkClass_processing(linkstr):
        Link = {
            'logit': Logit(),
            'probit': Probit(),
            'cloglog': CLogLog(),
            'cauchit': Cauchit(),
            'log': Log(),
        }
        return Link.get(linkstr, Logit())

    @staticmethod
    def _dist_processing(dist):
        if dist not in ['poisson','negbin','geom']:
            sys.exit(dist+" method not yet implemented")
        return dist

class ZeroInflatedResults(object):
    def __init__(self, call, formula, terms, kx, kz, dist, link, linkobj, options, method, start, reltol, \
                 weights, offsetx, offsetz, fitResult, coefc, coefz, theta, SE_logtheta,\
                nobs, res, Yhat, vc, y, X, Z):

        # Need to change final results objects names to standard names used in python
        self.call = call
        self.formula = formula
        self.terms = terms
        self.kx = kx
        self.kz = kz
        self.n = self.nobs = nobs

        # fit paramters
        self.dist = dist
        self.linkstr = link
        self.link = linkobj.link
        self.linkinv = linkobj.link_inv
        self.optim_options = options
        self.method = method
        self.start = start
        self.reltol = reltol

        self.weights = weights
        self.offsetx = offsetx
        self.offsetz = offsetz
        self.linkobj = linkobj

        # Optimization Results
        self.fit = fitResult
        self.loglik = fitResult.fun*(-1)    # log-likelihood
        self.converged = fitResult.success
        self.iters = fitResult.nit;          # number of iterations for convergence
        self.coefficients = {'count':coefc ,'zero': coefz}
        self.theta = theta if (dist is 'negbin') else None
        self.SE_logtheta = SE_logtheta
        self.df_null = nobs - 2
        self.df_resid = nobs - (kx + kz + (dist == "negbin"))
        self.df_model = (kx + kz + (dist == "negbin"))
        self.residuals = res
        self.fitted_values = Yhat
        self.vcov = vc
        self.y = y
        self._X = X
        self._Z = Z

        self.deviance = 0
        self.pearson_chi2 = 0
        self.cov_type = None
        self.use_t = False

    def __repr__(self):
        display_string = "Call:\n    "+self.call
        display_string += "\n\nformula:\n    "+self.formula
        #display_string += f"\nterms:\n    Y: {self.terms['Y']}\n    X: {self.terms['X']}\n    Z: {self.terms['Z']}"
        display_string += "\ndist: "+self.dist
        display_string += "\nlink: "+self.linkstr
        display_string += "\nlinkobj:"+ self.linkobj.__repr__()
        display_string += f"\nMessage: {self.fit.message}"
        display_string += f"\nResult: \n    Count:\n"
        display_string += f"{self.coefficients['count']}\n    Zero:\n{self.coefficients['zero']}"
        display_string += f"\n    theta: {self.theta:0.12f}" if self.dist is 'negbin' else " "
        display_string += f"\ndf_null: {self.df_null} \ndf_resid: {self.df_resid}"
        return display_string

    # Merged function definitions from zipModel start here:
    def printModel(self):
        MODEL1_HEADER = f"Count model cefficients ({self.dist} log link): "
        MODEL2_HEADER = f"Zero-inflation model coefficients (binomial with {self.linkstr} link): "

        p_count = self.kx;   # num of preds (+ intercept) for poisson count
        p_zero  = self.kz;    # num of preds (+ intercept) for zero-inflated model

        # print string representation of model
        print("\nCall:\n" + self.call + "\n")

        # part 1: poisson count model
        print(MODEL1_HEADER )
        #count_coeff = list(zip(self.terms['X'], np.round(self.coefficients['count'], 4)))
        #c_df = pd.DataFrame(np.array(count_coeff).reshape(len(count_coeff),2), columns = list("  "))
        c_df = pd.DataFrame(np.array(self.coefficients['count']),index = self.terms['X'], columns = list(" ")\
                           ).applymap("{0:.5f}".format)
        #print (z_df.to_string())
        print(c_df.T)
        #display(HTML(pd.DataFrame(c_df).to_html))
        #print('\n')
        if self.dist == 'negbin':
            print(f'Theta = {self.theta:0.4f}')
        # for line in [self.terms['X'], np.round(self.coefficients['count'],4)]:
        #     print(('{:>12}' * p_count).format(*line))
        print("\n")

        # part 2: logit model for predicting excess zeros
        print(MODEL2_HEADER)

        #zero_coeff = list(zip(self.terms['Z'], np.round(self.coefficients['zero'], 4)))
        #z_df = pd.DataFrame(np.array(zero_coeff).reshape(len(zero_coeff),2), columns = list("  "))
        #z_df = pd.DataFrame(np.array(zero_coeff).reshape(len(zero_coeff),2), columns = list("  "))
        z_df = pd.DataFrame(np.array(self.coefficients['zero']),index = self.terms['Z'], columns = list(" ")\
                           ).applymap("{0:.5f}".format)
        #print (z_df.to_string())
        print(z_df.T)

    # return variance-covariance matrix for calculation purposes
    def covar(self):
        return(self.vcov)

    # formatted version of the variance-covariance matrix
    #def printCov(self):
    #    pd.set_option('display.float_format', '{:.3E}'.format)
    #    cov_mat = pd.DataFrame(self.vcov)
    #    display(HTML(cov_mat.to_html(index=False)))

    def summary(self):
        RESIDUAL_OUTPUT = "Pearson residuals:"
        MODEL1_HEADER = f"Count model cefficients ({self.dist} log link): "
        MODEL2_HEADER = f"Zero-inflation model coefficients (binomial with {self.linkstr} link): "

        ## chunk 1: output call, formula
        print("\nCall:\n" + self.call +'\n')

        ## chunk 2: output pearson residuals -- residuals function still needs to be implemented
        # object$residuals = residuals(object, type = "pearson")
        resid_summ = np.round(st.mstats.mquantiles(self.residuals, prob = [0, 0.25, 0.5, 0.75, 1.0]), 5)
        resid_str  = ['Min', '1Q', 'Median', '3Q','Max'];
        print(RESIDUAL_OUTPUT)
        for line in [resid_str, resid_summ]:
            print(('{:>10}' * len(resid_summ)).format(*line))
        #print("\n")

        # compute standard error for all coefficients (both models)
        se = np.sqrt(np.diagonal(self.vcov))

        # compute z statistics for both models
        z_count = np.array(self.coefficients['count']) / se[0:self.kx]
        z_zip = np.array(self.coefficients['zero']) / se[self.kx:self.kx+self.kz]

        # compute p-values
        pval_count = 2 * st.norm.cdf(-np.abs(z_count));
        pval_zip = 2 * st.norm.cdf(-np.abs(z_zip));

        # format p-values for output
        pc_format = [0] * len(pval_count)
        pz_format = [0] * len(pval_zip)

        for i in np.arange(len(pval_count)):
            if pval_count[i] < 2e-16:
                pc_format[i] = str("<2e-16")
            else:
                pc_format[i] = str("{:.4e}".format(Decimal(pval_count[i])))

        for i in np.arange(len(pval_zip)):
            if pval_zip[i] < 2e-16:
                pz_format[i] = str("<2e-16")
            else:
                pz_format[i] = str("{:.4e}".format(Decimal(pval_zip[i])))

        ## chunk 3: output count model coefficients
        print('\n')
        print(MODEL1_HEADER)
        coeff_label = ['', 'Estimate', 'Std. Error', 'z value', 'Pr(>|z|)'];
        data_count = [coeff_label] + list(zip(self.terms['X'], np.round(self.coefficients['count'],4), \
                                        np.round(se[0:self.kx], 4), np.round(z_count, 3), pc_format))
        coeff_maxlen = np.max([len(x) for x in self.terms['X']])
        SPACE = 20 if coeff_maxlen > 10 else 12

        for i, d in enumerate(data_count):
            line = '|'.join(str(x).rjust(SPACE if x==d[0] else 10) for x in d)
            print(line)
            if i == 0:
                print('-' * len(line))

        print('\n')

        ## chunk 4: output zero-inflation model coefficients
        print(MODEL2_HEADER)
        data_zero = [coeff_label] + list(zip(self.terms['Z'], np.round(self.coefficients['zero'],4), \
                                        np.round(se[self.kx:], 4), np.round(z_zip, 3), pz_format))
        coeff_maxlen = np.max([len(x) for x in self.terms['Z']])
        SPACE = 20 if coeff_maxlen > 10 else 12
        for i, d in enumerate(data_zero):
            line = '|'.join(str(x).rjust(SPACE if x==d[0] else 10) for x in d)
            print(line)
            if i == 0:
                print('-' * len(line))
        #print('\n')
        print('---')

        ## chunk 5: Number of iterations, log-likelihood
        if self.dist == 'negbin':
            print(f'Theta = {self.theta:0.4f}')
        print(f"Number of iterations in {self.method} optimization: " + str(self.iters));
        if self.converged is False:
            print("Failed to converge.")
        print("Log-likelihood: "+str("{:.3e}".format(Decimal(self.loglik)))+f" on {self.df_model} Df.")

    # m: fitted zip model
    # pred_type: 'response', 'count', 'zero', or 'prob'
    # newdata_X: design matrix for count model
    # newdata_Z: design matrix for zero model
    def predict(m, pred_type, newdata_X = None, newdata_Z = None):

        if pred_type not in ["response", "count", "zero", "prob"]:
            print("Error: Prediction of type \'" + pred_type + "\' is not supported")
            return

        if newdata_X is None and newdata_Z is None:
            rval = m.fitted_values
            if pred_type != 'response':
                X = m._X                   # training matrix for count model with intercept in first column
                Z = m._Z                   # training matrix for zero model with intercept in first column

                ## offset
                # offsetx = 0.0
                #offsetx = np.repeat(0.0,self.n)
                offsetx = np.repeat(0.0, X.shape[0])

                # offsetz = 0.0
                # offsetz = np.repeat(0.0,self.n)
                offsetz = np.repeat(0.0, Z.shape[0])

                ## count mean
                # mu = np.exp(np.dot(X,m.coefficients['count']) + offsetx)
                ## binary mean
                # phi = self.linkobj.link_inv(np.dot(Z, m.coefficients['zero']) + offsetz)
                # Not sure if it's what's needed here but consider code commented above. was compatible
                # with offsetx and offsetz in earlier functions.

                mu = np.exp(np.matmul(X, m.coefficients['count']) + offsetx)
                phi = m.linkinv(np.matmul(Z, m.coefficients['zero']) + offsetz)
        elif newdata_X is None or newdata_Z is None:
            print("Design Matrix for both count model and zero-inflated model must be specified")
            return
        else:
            X = newdata_X
            Z = newdata_Z
            offsetx = np.repeat(0.0, X.shape[0])
            offsetz = np.repeat(0.0, Z.shape[0])
            mu = np.exp(np.matmul(X, m.coefficients['count']) + offsetx)
            phi = m.linkinv(np.matmul(Z, m.coefficients['zero']) + offsetz)
            rval = (1 - phi) * mu

        # predicted means for count/zero component
        if pred_type == 'count':
            rval = mu
        if pred_type == 'zero':
            rval = phi

        if pred_type == 'prob':
            if(y is not none):
                y = np.squeeze(m.y)
            else:
                print("Predicted Probabilities require non-null values for y")

            yUnique = np.arange(0, np.max(m0.y) + 1)
            nUnique = len(yUnique)
            rval = np.zeros(shape = (len(rval), nUnique))

            if m.dist == 'poisson':
                rval[:,0] = phi + (1 - phi) * np.exp(-mu) # first column of r_val
                for i in np.arange(1,nUnique):
                    # can this be vectorized instead of for loop?
                    rval[:,i] = (1 - phi) * poisson.pmf(yUnique[i], mu)
            elif m.dist == 'negbin':
                theta = self.theta
                rval[:,0] = phi + (1 - phi) * nbinom.pmf(yUnique[i],*self.convert_params(theta = theta, mu = mu))
                for i in np.arange(1,nUnique):
                    # can this be vectorized instead of for loop?
                    # should size be 1 in negbin?
                    rval[:,i] = (1 - phi) * nbinom.pmf(yUnique[i],*self.convert_params(theta = 1, mu = mu))
            elif m.dist == 'geom':
                rval[:,0] = phi + (1 - phi) * nbinom.pmf(yUnique[i],*self.convert_params(theta = 1, mu = mu))
                for i in np.arange(1,nUnique):
                    # can this be vectorized instead of for loop?
                    rval[:,i] = (1 - phi) * nbinom.pmf(yUnique[i],*self.convert_params(theta = 1, mu = mu))
            else:
                print("unsupported distribution")

        return(rval)


#Multiples of 3 and 5
#If we list all the natural numbers below 10 that are multiples of 3 or 5, we get 3, 5, 6 and 9. The sum of these multiples is 23.
#Find the sum of all the multiples of 3 or 5 below 1000.

#intuition would have me believe that simply iterating through multiples and summing them would be an easy brute force method that would
#prove effective up to the problem statement limit of 1000. However, this method would likely scale poorly to extremely large inputs. Perhaps
#there are shortcuts in the math to be found...

#5 and 10
import timeit
def Problem1():
    x = 5
    fives = 0
    while x < 1000:
        fives += x
        x += 5
        print(x)
        print(fives)

    y=3
    threes = 0
    while y <= 999:
        if y % 15 == 0:
            y+=3
        else:
            threes += y
            y += 3

    total = threes + fives
    print(total)

#alternatively one can use a single iteration through the digits by using an OR and iterating through each digit. This mehtod, although also
#brute force, affords a massive run time efficiency improvement.

def Problem1_v2():
    sum1=0
    for i in range (999):
        if i%3 ==0 or i % 5 == 0:
            sum1 += i

#Finally the use of the formula ∑k=1/2n(n+1). will allow for a few simple math calculations rather than an iterative process and thus
#not scaling at runtime according to n other than byte size whereas this  previous strategy certainly has nth order scaling. However, it must be
#remembered that multiples of 15 are counted twice as they are multiples of both 3 and 5 and that is accounted for in the third term of the following

def Dividesum(n,p):
    return n*(p/n)*((p/n)+1)/2

def Problem1_v3():
    result = Dividesum(3,999)+Dividesum(5,999)-Dividesum(15,999)
    print(result)

elapsed_time = timeit.timeit(Problem1, number=10)/10
elapsed_time_v2 = timeit.timeit(Problem1_v2, number = 10)/10
elapsed_time_v3 = timeit.timeit(Problem1_v3, number = 10)/10
print(elapsed_time)
print(elapsed_time_v2)
print(elapsed_time_v3)


#!/usr/bin/python3

from itemsqueue import ItemsQueue
from allegro import Allegro
from concurrent.futures import ThreadPoolExecutor
import json
import datetime
import time
import io
import sys

queue = ItemsQueue(sys.argv[1])
filebase = sys.argv[2]

def today_date():
  return datetime.datetime.now().strftime("%Y-%m-%d")

def dumpitems(items):
  with io.open("%s.%s.txt" % (filebase, today_date()), "a+", encoding='utf8') as itemsfile:
    for item in items:
      mark_status([item['itemInfo']['itId']], ItemsQueue.DONE)
      j = json.dumps(item, ensure_ascii=False)
      itemsfile.write(j + "\n")

def mark_status(ids, status):
  for id in ids:
    queue.markStatus(id, status)

allegro = Allegro()
allegro.load_credentials('.credentials')
concurrency = int(sys.argv[3])

while True:
  sample = queue.getWaiting(50000)
  print("Will download %d items" % len(sample))
  if len(sample) > 0:
    executor = ThreadPoolExecutor(max_workers=concurrency)
    futures = []
    for chunk in [sample[x:x+25] for x in range(0, len(sample), 25)]:
      futures.append(executor.submit(allegro.get_items_info, chunk))

    for idx, future in enumerate(futures):
      try:
        done, not_found, killed = future.result()
        mark_status(not_found, ItemsQueue.NOT_FOUND)
        mark_status(killed, ItemsQueue.KILLED)
        dumpitems(done)
        queue.commit()
        futures[idx] = None
      except:
        print("Chunk failed", sample[idx * 25:idx *25 + 25], sys.exc_info()[1])

      sys.stdout.write('.' if len(done) == 25 else '+')
      sys.stdout.flush()

    sys.stdout.write("\n")
    sys.stdout.flush()
    executor.shutdown()
  else:
    time.sleep(60)


#!/usr/bin/env python
# D. Jones, M. Dai - 9/16/19

import sys, argparse, configparser, os

usagestring = """Lightcurve fitting
"""

parser = argparse.ArgumentParser(usage=usagestring, conflict_handler="resolve")
parser.add_argument('config', default='pipeline_training.txt', type=str,
                    help='configuration file')
options = parser.parse_args()
from saltshaker.pipeline.pipeline import *

def runpipe():
    pipe = SALT3pipe(finput=options.config)
    pipe.build(data=False,mode='customize',onlyrun=['lcfit'])
    pipe.configure()

    pipe.run(onlyrun=['lcfit'])

if __name__ == "__main__":
	runpipe()


from django.conf.urls import url, include
from api import views
from rest_framework import routers

router = routers.DefaultRouter()
router.register(r'users', views.UserViewSet)
router.register(r'todo', views.ToDoViewSet)

urlpatterns = [
    url(r'^', include(router.urls)),
    url(r'^api-auth/', include('rest_framework.urls')),
    url(r'^authenticate/', views.CustomObtainAuthToken.as_view()),
]


'''
102. Binary Tree Level Order Traversal
Given a binary tree, return the level order traversal of its nodes' values. (ie, from left to right, level by level).

For example:
Given binary tree {3,9,20,#,#,15,7},
    3
   / \
  9  20
    /  \
   15   7
return its level order traversal as:
[
  [3],
  [9,20],
  [15,7]
]
'''
# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def levelOrder(self, root):
        """
        :type root: TreeNode
        :rtype: List[List[int]]
        """
        if root is None:
            return []

        queue = []
        level = -1
        result = []
        queue.append((root,0))
        while len(queue) != 0:
            node, plevel = queue.pop(0)
            if plevel == level:
                result[-1].append(node.val)
            else:
                result.append([node.val])
                level = plevel
            if node.left:
                queue.append((node.left, plevel + 1))
            if node.right:
                queue.append((node.right, plevel + 1))
        return result