__author__ = 'roman'

from bottle import request, Bottle, abort
from bottle import route, run, template, static_file
import json
import jsonpickle
from client import Client
from serverToClientMessage import ServerToClientMessage
from taskmanager import *
from serverPageToServerMessage import *
from serverToServerPageMessage import *
from clientToServerMessage import *
from seleniumTest import *
from multiprocessing import Process

app = Bottle()
allClients = []
serverSocket = None
taskList = []

def sendAddClientMessageToServerPage(client):
    """
        Send new message to server page
        client: icludes all client info
    """
    newClientMessage = ServerToServerPageMessage(ServerToServerPageMessage.NEW_CLIENT_MSG) # add client to server page
    newClientMessage.addClient(client)
    if (serverSocket):
        serverSocket.send(jsonpickle.encode(newClientMessage))

def sendLeaveClientMessageToServerPage(client):
    """
        Sends leave message to the server page
    """
    deletedClientMessage = ServerToServerPageMessage(ServerToServerPageMessage.CLIENT_LEAVED_MSG) # delete client from server page
    deletedClientMessage.deletedClient(client)
    if (serverSocket):
        serverSocket.send(jsonpickle.encode(deletedClientMessage))

def sendWarningMessageToServerPage(message):
    """ Sends warning message to server page """
    warningMessage = ServerToServerPageMessage(ServerToServerPageMessage.WARNING_MSG) # add warning to server page
    warningMessage.warning(message)
    if (serverSocket):
        serverSocket.send(jsonpickle.encode(warningMessage))

def sendClientStatusMessageToServerPage(clientId, status, substringFound, time=None):
    """ Sends client status to server page """
    message = ServerToServerPageMessage(ServerToServerPageMessage.CLIENT_STATUS_MSG) # add client status to server page
    message.clientStatus(clientId, status)
    message.setSubstringFound(substringFound)
    if time != None:
        message.setTime(time)
    if (serverSocket):
        serverSocket.send(jsonpickle.encode(message))

def sendMsgToClient(msg, socket):
    if socket:
        socket.send(jsonpickle.encode(msg)) #send client info to server page

def giveOutTasks(allClients, taskList):
    """ sends tasks to free clients """
    for task in taskList: # set performers to every task
        if task.getClientPerformer() == None:
            for client in allClients:
                if client.isPerformer == False:
                    client.isPerformer = True
                    task.setClientPerformer(client)
                    print ("performer for task is ", client.getId())
                    break
    for task in taskList: # give out Tasks
        client = task.getClientPerformer()
        if (client != None) and (client.busy == False):
            #print("performer for task is client with id ", client.getId())
            msg = ServerToClientMessage(ServerToClientMessage.TASK_MSG, client.getId())
            msg.setTask(task)
            print("message for perfrmer: ", msg)
            if client.getSocket():
                client.getSocket().send(jsonpickle.encode(msg))
                client.busy = True

def removeTaskFromTaskList(performerId):
    """ Removes task with given performer id """
    for task in taskList:
        client = task.getClientPerformer()
        if client != None:
            if client.getId() == performerId:
                client.busy = False
                client.isPerformer = False
                taskList.remove(task)
                print ("Task done by client", performerId)
                print ("Task removed from taskList")

def removeTaskPerformer(cl):
    """ is called when working client suddenly closes """
    for task in taskList:
        client = task.getClientPerformer()
        if (client != None) and (cl != None):
            if client.getId() == cl.getId():
                task.setClientPerformer(None)

def generateClientId():
    """ Generates unique client id """
    currentClientId = 0
    if allClients.__len__() <= 0: #form new id
        currentClientId = 0
    else:
        for cl in allClients:
            currId =  cl.getId()
            if currentClientId <= currId:
                currentClientId = currId + 1
    return currentClientId

@app.route('/websocketClient')
def handle_websocket_client():
    """
        Is called when client sends socket, handles clients messages
    """
    wsock = request.environ.get('wsgi.websocket')

    _client = Client(generateClientId(), wsock, request.environ.get('HTTP_USER_AGENT'))
    print ("client socket received")

    allClients.append(_client) #append cluster in array to handle them
    serverMsg = ServerToClientMessage(ServerToClientMessage.CONECTION_MSG, _client.getId())
    sendMsgToClient(serverMsg, _client.getSocket()) # send message that client is connected

    print("all clients: ", allClients.__len__())
    print("server socket: ", serverSocket)

    sendAddClientMessageToServerPage(_client)

    if not wsock:
        abort(400, 'Expected WebSocket request.')

    while True:
        try:
            msg = wsock.receive()
            if (msg != None):
                message = ClientToServerMessage(msg)
            else: # client leaved
                for cl in allClients:
                    if cl.getSocket() == wsock:
                        allClients.remove(cl)
                        if cl.busy == True:
                            removeTaskPerformer(cl)
                            print("Working client disconnected")
                        sendLeaveClientMessageToServerPage(cl)
                        giveOutTasks(allClients, taskList)

                break
            if (message.type == ClientToServerMessage.STATUS): # send client status message to server page
                for cl in allClients:
                    if cl.getSocket() == wsock:
                        if message.status == 100: # if client completes task
                            # send status with time argument
                            print("time: ", message.time)
                            sendClientStatusMessageToServerPage(cl.getId(), message.status, message.substringFound, message.time)
                            if cl != None:
                                removeTaskFromTaskList(cl.getId())
                            if taskList.__len__() == 0:
                                print("WORK DONE!!!") # send WORK_DONE_MSG
                                if serverSocket:
                                    serverSocket.send(jsonpickle.encode(ServerToServerPageMessage(ServerToServerPageMessage.WORK_DONE_MSG)))
                            else:
                                giveOutTasks(allClients, taskList)
                        else:
                            sendClientStatusMessageToServerPage(cl.getId(), message.status, message.substringFound)

        except WebSocketError:
            for cl in allClients:
                if cl.getSocket() == wsock:
                    allClients.remove(cl)
                    if cl.busy == True:
                        removeTaskPerformer(cl)
                        print("Working client disconnected")
                    sendLeaveClientMessageToServerPage(cl)
                    giveOutTasks(allClients, taskList)
            break

@app.route('/websocketServer')
def handle_websocket_server():
    """
        Is called when server sends socket, handles server messages
    """
    global serverSocket #use global serverSocket because i want to send some data to server page from another function
    global taskList
    serverSocket = request.environ.get('wsgi.websocket')

    if serverSocket:
        print ("Server socket received")
        for cl in allClients:
            sendAddClientMessageToServerPage(cl)

    if not serverSocket:
        abort(400, 'Expected WebSocket request.')

    while True:
        try:
            msg = serverSocket.receive()
            if (msg != None):
                message = ServerPageToServerMessage(msg)

                if message.type == ServerPageToServerMessage.START_SHARING_TASKS_MSG: # if server page wants to start sharing tasks
                    taskManager = TaskManager(message.data) # create task manager, arg - substring to search
                    taskList = taskManager.getTasks(allClients.__len__()) # generate tasks
                    if taskList != None:
                        print("Task manager generated tasks: ")
                        for task in taskList:
                            print("task -- ", task.string[task.string.__len__()-200:task.string.__len__()])
                        print("clients amount: ", allClients.__len__())
                        print("substring to search: ", task.substringToSearch)

                        giveOutTasks(allClients, taskList)
                    else:
                        sendWarningMessageToServerPage("No clients connected")
        except WebSocketError:
            print ("except WebSocketError")
            break

@app.route('/')
def index():
    return template('client_template')

@app.route('/server')
def server():
    print ("all clients: ", allClients)
    return template('server_template')

# Static Routes
@app.route('/<filename:re:.*\.js>')
def javascripts(filename):
    return static_file(filename, root='static/js')

@app.route('/<filename:re:.*\.css>')
def stylesheets(filename):
    return static_file(filename, root='static/css')

@app.route('/<filename:re:.*\.(jpg|png|gif|ico)>')
def images(filename):
    return static_file(filename, root='static/img')

@app.route('/<filename:re:.*\.(eot|ttf|woff|svg)>')
def fonts(filename):
    return static_file(filename, root='static/fonts')

from gevent.pywsgi import WSGIServer
from geventwebsocket import WebSocketError
from geventwebsocket.handler import WebSocketHandler

def runServer():
    server = WSGIServer(("127.0.0.1", 8080), app,
                        handler_class=WebSocketHandler)
    server.serve_forever()

def runTesting():
    server_thread = Process(target=runServer)
    server_thread.start()
    unittest.main()

#runServer()
runTesting()


# -*- coding: utf-8 -*-
"""
@author: Yifan Cheng, Skylar Du, Yashash Gaurav, Mark He
"""

# import os
import time
import random
import requests
import numpy as np
import pandas as pd
from bs4 import BeautifulSoup
from selenium import webdriver
from urllib.request import urlopen
from pandas.core.frame import DataFrame
from user_agent import DESKTOP_USER_AGENTS
from selenium.common.exceptions import NoSuchElementException
import csv

edge_driver_path = (
    "path/to/webdriver"
)

#######################ROI DATA#######################
# Crawl the ROI data with beautifulsoup
# Time Warning: about 10 minutes
def get_roi():
    school_name = []
    rank = []
    twenty_year_roi = []
    total_4_year_cost = []
    graduation_rate = []
    typical_years_to_graduate = []
    average_loan_amount = []
    for page in range(1, 199, 1):
        html = "https://www.payscale.com/college-roi/page/" + str(page)
        # get HTML
        html = urlopen(html)
        bsyc = BeautifulSoup(html.read(), "html.parser")
        fout = open("payscale_temp.txt", "wt", encoding="utf-8")
        fout.write(str(bsyc))
        fout.close()
        # get lists
        tc_table = list(bsyc.body.div.div)
        table = list(tc_table[1].children)
        body = table[4].tbody
        for i in body.children:
            li = list(
                i.find_all("span", {"class": "roi-grid__schoolname--text"})
            )
            for j in li:
                school_name.append(str(j).split(">")[2][:-3])
        for i in body.children:
            li = list(i.find_all("span", {"class": "roi-grid__rank--text"}))
            for j in li:
                rank.append(str(j).split(">")[1][:-6])
        datas = []
        for i in body.children:
            li = list(i.find_all("span", {"class": "data-table__value"}))
            for j in li:
                datas.append(str(j).split(">")[1][:-6])
        for i in range(len(datas)):
            if i % 7 == 2:
                twenty_year_roi.append(datas[i])
            if i % 7 == 3:
                total_4_year_cost.append(datas[i])
            if i % 7 == 4:
                graduation_rate.append(datas[i])
            if i % 7 == 5:
                typical_years_to_graduate.append(datas[i])
            if i % 7 == 6:
                average_loan_amount.append(datas[i])
    # convert to dictionary
    dataf = {
        "Rank": rank,
        "School Name": school_name,
        "20 Year Net ROI": twenty_year_roi,
        "Total 4 Year Cost": total_4_year_cost,
        "Graduation Rate": graduation_rate,
        "Typical Years to Graduate": typical_years_to_graduate,
        "Average Loan Amount": average_loan_amount,
    }
    # Convert to dataFrame and output as Excel
    dataframe = DataFrame(dataf, index=rank)
    dataframe.to_excel("output.xlsx")

# read ROI data
# since the time to crawl is long, I stored them in an excel, and now I will read from the excel
def clean_roi():
    roi = pd.read_excel("output.xlsx")
    roi = roi.drop_duplicates(
        subset=[
            "School Name",
            "20 Year Net ROI",
            "Total 4 Year Cost",
            "Graduation Rate",
            "Typical Years to Graduate",
        ]
    )
    roi["Rank"] = range(len(roi))
    roi["School Name"] = [
        (s.split("-")[0] + " - " + s.split("-")[1]).strip()
        if s.find("-") != -1
        else s
        for s in roi["School Name"]
    ]
    #os.remove("output.xlsx")
    return roi

#######################NICHE DATA#######################
# Crawl the NICHE SAT data with beautifulsoup
# Time Warning: about 10 minutes
def send_request(link):
    time.sleep(random.choice(range(2, 10)))
    headers = {"user-agent": random.choice(DESKTOP_USER_AGENTS)}
    res = requests.get(link, headers=headers)
    if res == None:
        print("res is empty")
        return res
    return res

def get_niche():
    # create a list with the url of all 82 pages of the ranking
    allurl = ["https://www.niche.com/colleges/search/best-value-colleges/",]
    nexturl = "https://www.niche.com/colleges/search/best-value-colleges/?page="
    for i in range(2,82):
        myurl = nexturl+str(i)
        allurl.append(myurl)

    school_name = []
    fact = []
    location = []
    # iterate and scrape through each link
    for url in allurl:
        #define the link to be scrape and sent request
        res =send_request(url)
        #parse data from each page
        newsoup = BeautifulSoup(res.text,'html.parser')
        #find the info on each school
        #for each school, create a new dict that contain its name and fact
        for div in newsoup.find_all("div", class_="card"):
            # remove the sponsered colleges
            if len(div.find_all("div", class_="search-result__sponsered-bar"))==0:
                school_name.append(div.find("h2", class_="search-result__title"))
                fact.append(div.find_all("span", class_= "search-result-fact__value"))
                if len(div.find_all("li", class_= "search-result-tagline__item"))!=0:
                    location.append(div.find_all("li", class_= "search-result-tagline__item")[1])
                else:
                    location.append('null')

    cleaned_school_name = []
    for i in school_name:
        new = str(i).split('>')
        if (len(new)<2):
            cleaned_school_name.append("null")
        else:
            cleaned_school_name.append(new[1].split('<')[0])

    cleaned_fact = []
    for list in fact:
        temp = []
        for i in list:
            new = str(i).split('>')
            if (len(new)<2):
                temp.append("null")
            else:
                temp.append(new[1].split('<')[0])
        cleaned_fact.append(temp)

    cleaned_location = []
    for i in location:
        new = str(i).split('>')
        if (len(new)<2):
            cleaned_location.append("null")
        else:
            cleaned_location.append(new[1].split('<')[0])

    city = []
    state = []
    for i in cleaned_location:
        new = i.split(', ')
        if (len(new)<2):
            city.append("null")
            state.append("null")
        else:
            city.append(new[0])
            state.append(new[1])

    Acceptance_Rate = []
    Net_Price = []
    SAT_Range = []

    for eachfact in cleaned_fact:
        if len(eachfact)>0:
            Acceptance_Rate.append(eachfact[0]), Net_Price.append(eachfact[1])
            if len(eachfact)==3:
                if len(eachfact[2])>1:
                    SAT_Range.append(eachfact[2])
                else:
                    SAT_Range.append("null")
            else:
                SAT_Range.append("null")
        else:
            Acceptance_Rate.append("null"), Net_Price.append("null"), SAT_Range.append("null")

    file = open("cleaned_niche.csv", "w", newline = "")
    writer = csv.writer(file)
    # field names
    fields = ['School Name', 'City', 'State', 'Acceptance Rate', 'Net Price', 'SAT Range']
    writer.writerow(fields)
    for i in range(len(cleaned_school_name)):
        if cleaned_school_name[i] != "null":
            writer.writerow([cleaned_school_name[i], city[i], state[i], Acceptance_Rate[i], Net_Price[i], SAT_Range[i]])
    file.close()

#######################Merge the data#######################
def add_calculation_columns(merged_data):
    def convert_currency_to_int(currency):
        try:
            return int(currency.replace(",", "").replace("$", ""))
        except:
            return np.NaN

    def get_sat_range_min(sat_range):
        if type(sat_range) is float and np.isnan(sat_range):
            return np.nan
        else:
            return int(sat_range.split("-")[0])

    def get_sat_range_max(sat_range):
        if type(sat_range) is float and np.isnan(sat_range):
            return np.nan
        else:
            return int(sat_range.split("-")[1])

    merged_data["SAT Min"] = merged_data["SAT Range"].map(
        lambda cell: get_sat_range_min(cell)
    )
    merged_data["SAT Max"] = merged_data["SAT Range"].map(
        lambda cell: get_sat_range_max(cell)
    )
    merged_data["Total 4 Year Cost (Integer)"] = merged_data[
        "Total 4 Year Cost"
    ].map(lambda x: convert_currency_to_int(x))

def merge_data():
    roi = clean_roi()
    roi["Rank"] = roi["Rank"] + 1
    niche = pd.read_csv("cleaned_niche.csv", encoding='ISO-8859-1')
    niche = niche.drop_duplicates()
    merged_data = pd.merge(roi, niche, how="left", on="School Name")
    merged_data = merged_data.drop(labels="Unnamed: 0", axis=1)
    add_calculation_columns(merged_data)
    merged_data.to_excel("merged_data.xlsx")

##################### Combining Fetch ###############
def refresh_all_data():
    get_roi()
    get_niche()
    merge_data()
    # get_careers_data() - uncomment if you have webdriver path set
    print("Success! All data is refreshed.")

#################### Best Colleges Data ####################
# this portion of the script is to scrape bestcollege.com
# for information on career data

def get_careers_data():
    # Call your browser
    driver = webdriver.Edge(executable_path=edge_driver_path)

    # Get all career links
    driver.get("https://www.bestcolleges.com/careers/")
    career_anchors = driver.find_elements_by_css_selector(
        "div.swiper-slide a[data-wpel-link='internal']"
    )
    career_links = [career.get_attribute("href") for career in career_anchors]

    career_data_list = []

    # Loop through urls collected and collect data
    for career in career_links:
        driver.get(career)
        career_name = driver.find_element_by_css_selector(
            "section.hero h1"
        ).text
        career_info = driver.find_element_by_css_selector(
            "section.container.content>p:first-child"
        ).text

        try:
            why_career = driver.find_element_by_css_selector(
                'a[id^="why-pursue"]+h2+p'
            ).text
            why_career += (
                " "
                + driver.find_element_by_css_selector(
                    'a[id^="why-pursue"]+h2+p+div+p'
                ).text
            )
        except NoSuchElementException:
            why_career = ""

        try:
            how_to_start = driver.find_element_by_css_selector(
                "a#advancing-your-career+h2+p"
            ).text
            how_to_start += driver.find_element_by_css_selector(
                "a#advancing-your-career+h2+p+div+p"
            ).text
        except NoSuchElementException:
            how_to_start = ""

        career_data_list.append(
            {
                "career_name": career_name,
                "career_info": career_info,
                "why_career": why_career,
                "how_to_start": how_to_start,
            }
        )

    # Create a data frame of the data collected
    career_data = pd.DataFrame(
        career_data_list,
        columns=["career_name", "career_info", "why_career", "how_to_start"],
    )

    # Save the data frame created for future use.
    career_data.to_csv(
        "bestcolleges_careers.csv", index=False, encoding="utf-8"
    )

    driver.close()

merge_data()


# ==========
# Created by Ivan Gadjev
# 2020.02.01
#
# Library of custom functions that aid in projects using pyRadia, genesis, elegant, and varius other codes and analyses.
#
#
# ==========

# import sys
# import os
import time
import bisect
import json

import scipy.constants as pc
import numpy as np

import scipy.signal as sps

from matplotlib.cm import ScalarMappable
import matplotlib.pyplot as plt

# make a custom JSON encoder
class NdArrEncoder(json.JSONEncoder):
    """
    Class to encode nd.array into JSON. Inhehrits from json.JSONEncoder.
    Useful for storing PSO output to file.
    """
    def default(self, obj):
        """
        default() method of JSONEncoder that is re-implemented to handle np.ndarray
        """

        if isinstance(obj, np.ndarray):
            return obj.tolist()
        # this return statement allows JSONEncoder to handle errors and stuff
        return json.JSONEncoder.default(self, obj)

#
# ===== global defs =====
#

xhat = np.array([1,0,0])
yhat = np.array([0,1,0])
zhat = np.array([0,0,1])
origin = np.array([0,0,0])

# physical constants
mc2 = 1e-6 * (pc.m_e*pc.c**2)/pc.elementary_charge # MeV. electron mass
mpc2 = 1e-6 * (pc.m_p*pc.c**2)/pc.elementary_charge # MeV. proton mass

#
# ===== helper functions =====
#

def rot_mat(angle, dim=3, axis='z'):
    """ Create a rotation matrix for rotation around the specified axis.
    """
    # sine cosine
    c = np.cos(angle)
    s = np.sin(angle)

    if dim == 3:
        if axis == 'x':
            rr = np.array( [[1, 0, 0],
                            [0, c, -s],
                            [0, s, c]] )
        elif axis == 'y':
            rr = np.array( [[c, 0, s],
                            [0, 1, 0],
                            [-s, 0, c]] )
        elif axis == 'z':
            rr = np.array( [[c, -s, 0],
                            [s, c, 0],
                            [0, 0, 1]] )
    elif dim == 2:
        rr = np.array( [[c, -s],
                        [s, c]] )
    else:
        print(' error: `dim` variable must be 2 or 3. Specifies the dimension of the vectors that the rotation is acting on.')
    return rr

def rgb_color(value, cmapname='viridis'):
    """ returns the RGB values of the color from the colormap specified with scaled value.
    value - between 0 and 1
    cmapname='viridis' name of the colormap
    """
    cmap = ScalarMappable(cmap=cmapname)
    return list(cmap.to_rgba(value, bytes=False, norm=False)[:-1])

def contourplot(xx, yy, zz,
                 zlims='auto',
                 arat=1, contours=True,
                 cbarlabel='<color_bar_label>', colormap='viridis',
                 xticks='auto', yticks='auto',
                 imkwargs=None, ckwargs=None):
    """ make a contour plot on top of an imshow figure for the given arrays
    xx, yy, zz are arrays from np.meshgrid() see numpy docs for ref
    s"""

    # compute limits for the plots
    xylims = [xx.min(),xx.max(), yy.min(),yy.max()]

    if (type(xticks) is np.ndarray) and (type(yticks) is np.ndarray):
        try:
            xt = np.unique(xx)
            xstep = np.abs(xt[1] - xt[0])
            yt = np.unique(yy)
            ystep = np.abs(yt[1] - yt[0])
            xylims = [xticks.min() - xstep/2, xticks.max() + xstep/2, yticks.min() - ystep/2, yticks.max() + ystep/2]
        except AttributeError:
            pass

    if zlims == 'auto':
        zmin = zz.min()
        zmax = zz.max()
    else:
        zmin = zlims[0]
        zmax = zlims[1]

    fig, ax = plt.subplots(figsize=(12,12))

    if imkwargs == None:
        imkwargs = dict(cmap=colormap, alpha=0.99, extent=xylims, vmin=zmin,vmax=zmax, interpolation='none', origin='lower')

    im = ax.imshow(zz, **imkwargs)
    ax.set_aspect(arat)

    CS = 0
    if contours:
        if ckwargs == None:
            ckwargs = dict(colors='white', extent=xylims, vmin=zmin,vmax=zmax,levels=19, linewidths=0.5, origin='lower')

        CS = ax.contour(xx,yy,zz,**ckwargs)

        # ax.clabel(CS, fontsize=14, inline=True, fmt='%1.0f')

    fs = dict(fontsize=18)
    if cbarlabel != 'none':
        cbar = fig.colorbar(im, shrink=0.7)

        # cbar.ax.set_ylabel('Tesla', **fs)
        cbar.ax.set_ylabel(cbarlabel, **fs)
        # cbar.add_lines(CS)
        cbar.ax.tick_params(labelcolor='k', labelsize=14, width=1)

#     plt.axis([-10,10,-10,10])

    if type(xticks) is np.ndarray:
        print('HELLO-------')
        ax.set_xticks(xticks)
    elif xticks == 'auto':
        xticks = np.unique(np.round(xx, decimals=3))
        # print(xticks)
        plt.xticks(xticks)

    if type(yticks) is np.ndarray:
        ax.set_yticks(yticks)
    elif yticks == 'auto':
        yticks = np.unique(np.round(xx, decimals=3))
        # print(yticks)
        plt.yticks(yticks)

    return fig, ax, CS

def gaussian(xvec, mean, std, norm='max'):
    """ Returns the Gaussian of the xvec normalized so that the peak is unity.
    xvec - ndarray of floats
    mean - float. the center of the Gaussian
    std - float. the standard deviation of the Gaussian
    norm - {'max', 'area'} whether to normalize so that the maximum is unity or the area under the curve is unity

    yvec = normconstant * Exp( - (x - mean)^2 / 2std^2 )

    """

    yvec = np.exp( - (xvec - mean)**2 / (2 * std**2) )

    if norm == 'area':
        normconstant = 1 / (std * np.sqrt(2 * pc.pi))
    else:
        normconstant = 1

    return normconstant * yvec

def gaussianDx(x, mu, sig):
    """ First derivative of a Gaussian. Normalized so that its maximum is 1.
    x - ndarray
    mu - float. mean
    sig - float. std dev

    """

    y = (-(x - mu)/sig**2) * gaussian(x, mu, sig)
    # NOTE: normalization to maximum value = 1
    y = y/np.abs(y).max()

    return y

def gaussianDx2(x, mu, sig):
    """ Second derivative of a Gaussian (inverted sombrero). Normalized so that its maximum is 1.
    x - ndarray
    mu - float. mean
    sig - float. std dev

    """

    y = ( -1/sig**2 + ((x - mu)/sig**2)**2 ) * gaussian(x, mu, sig)
    # NOTE: normalization to maximum value = 1
    y = y/np.abs(y).max()
    return y

def stepfunc(xvec, center, width, norm='max'):
    """
    Returns a step function with center and total width.
    xvec - ndarray of floats
    center - float
    width - float. total width of box or step
    norm - {'max', 'area'} whether to normalize so that the maximum is unity or the area under the curve is unity
    """

    yvec = [ int((xx >= center - width/2) and (xx <= center + width/2)) for xx in xvec]
    if norm == 'area':
        normconstant = 1 / np.trapz(yvec, xvec)
    else:
        normconstant = 1

    return normconstant * yvec

class signal(object):

    def smooth(xvec, yvec, width, mode='gauss', edgemode='valid'):
        """
        xvec
        yvec

        width - float. if mode == 'gauss' width=std, if 'step' [center-width/2, center+width/2]
        mode - {'gauss','step', 'median'}
        edgemode = {'valid','pad'}
        to be used with gaussian or stepfunc

        When 'gauss' or 'step' filtering modes are used, both xvec and yvec are convolved with a gaussian or step function. The output x and y vecs are of shorter length because the 'valid' values begin at the point when the span of the convolution function (gauss/step) is inside the boundaries of the original xvec.
        If the edgemode='valid' is used, then the output xvec and yvec are shorter than the input vectors, which may lead to problems when trying to compare initial to filtered signals.
        The edgemode='pad' features aims to rectify this by padding the filtered yvec with the edge value on both sides, while keeping the original xvec values there. This method does introduce a discontinuity of the the derivative of the filtered signal.
        """
        center = xvec.min() + (xvec.max() - xvec.min() ) / 2
        xstep = xvec[1]-xvec[0]
        xlen = xvec.shape[0]

        if mode in ['gauss', 'step']:
            wvec = 0
            if mode == 'gauss':
                nsig = 5
                xvectemp = np.arange(center - nsig * width, center + nsig * width + xstep, xstep)
                wvec = gaussian(xvectemp, center, width)
            elif mode == 'step':
                xvectemp = np.arange(center - width/2, center + width/2 + xstep, xstep)
                wvec = stepfunc(xvectemp, center, width)

            cnorm = 1 / (np.sum(wvec))
            # use mode='valid', otherwise edge/boundary values are not intuitive
            yvec = cnorm * np.convolve(yvec, wvec, mode='valid')

            ndiff = xlen - yvec.shape[0]
            # if ndiff < 0:
            #     print('WARNING: `width` is larger than `xvec` span.')
            if edgemode == 'pad':

                yvec = np.pad(yvec, [ndiff // 2, ndiff - ndiff // 2], mode='edge')

            elif edgemode == 'valid':
                xvec = cnorm * np.convolve(xvec, wvec, mode='valid') # done to match length of x and y vecs

        elif mode == 'median':
            yvec = sps.medfilt(yvec, kernel_size=width)
        else:
            print("ERROR: _mode_ must be 'gauss', 'step', or 'median' ")
            return 1

        return xvec, yvec

    def fft_scale_phase(timevec, sigvec, convertfactor,
    mode='weight', power=2, freqlims='auto',
    phase=0.0):
        """ Given a time-series signal
        1. rescale its time axis, based on its main FFT frequency and the given conversion factor
        and 2. phase the signal to the
        This action effectively calibrates the time axis to a known length scale, given by convertfactor.
        "Phasing" the signal is based on maximizing the correlation of the signal with a wavelet of its main frequency for a specified phase.
        This is a form of peak alignment. Returns the time-shift needed to do this.

        inputs:
        timevec - ndarray(n,)
        sigvec - ndarray(n,)
        convertfactor - float

        mode = {'weight', 'max'}, takes a weighted sum of the frequencies to find the main freq content. 'max' takes the freq with maximum amplitude.
        power = 1, the power for the weighting function. not used for 'max' mode.

        freqlims='auto, if 'auto', then use (0,inf) interval. else use (fmin, fmax) interval.

        phase - float. radians. the phase of the main freq oscillation. 0 is cosine-like

        returns:
        timescale - scaling factor to calibrate time axis
        timeshift - shift in time for time-align
        freq_main - main frequency component of the signal

        sigfft_freq - frequency vector from FFT
        sigfft - FFT of signal

        """

        # number of samples
        nsamples = sigvec.shape[0]

        # FFT of signal with appropriate normalization
        sigfft = (1 / nsamples) * np.fft.fft(sigvec)

        # time step
        timestep = (timevec[1:] - timevec[:-1]).mean()

        # frequency vector
        sigfft_freq = np.fft.fftfreq(nsamples, d=timestep)
        if freqlims == 'auto':
            # take only positive frequencies and exclude zero
            fmin = 0.0
            indlogic = (sigfft_freq > fmin)
        else:
            fmin = freqlims[0]
            fmax = freqlims[1]
            indlogic = (sigfft_freq > fmin) & (sigfft_freq < fmax)

        if mode == 'weight':
            # find the main frequency, based on the largest FFT amplitudes averaged with weigths
            # weight of freq is a power of their FFT amplitude

            # inside freqlims
            sigfft_bound = sigfft[indlogic]
            sigfft_freq_bound = sigfft_freq[indlogic]

            # number of frequencies to average over
            nfreq = np.sum(indlogic)

            # normalization constant of frequencies
            normconst =  1 / np.sum(np.abs(sigfft_bound)**power)
            # main frequency
            freq_main = normconst * np.sum(sigfft_freq_bound * np.abs(sigfft_bound)**power)

        elif mode == 'max':
            # inside freqlims
            sigfft_bound = sigfft[indlogic]
            sigfft_freq_bound = sigfft_freq[indlogic]

            #  freq with largest amplitude
            freq_main = sigfft_freq_bound[np.abs(sigfft_bound).argmax()]

        else:
            print('ERROR! Mode not recognized. Please use {"weight", "max"}.')
            return 1

        # scale wavelength to period of undulator
        timescale = convertfactor * freq_main

        # time-align

        # create wavelet with same sampling as signal
        omega = 2*pc.pi*freq_main
        wavelet = np.cos(omega*timevec + phase)
        # correlate (order of input vectors matters for timeshift sign.)
        tcorr = np.correlate(sigvec, wavelet, mode='full')
        # convert to a shift in time
        timeshift = (nsamples - (tcorr.argmax() + 1)) * timestep

        # mod to 2pi and shift to inside [-pi,pi]
        timeshift = (1/omega) * (np.mod(omega*timeshift - pc.pi, 2*pc.pi) - pc.pi)

        return timescale, timeshift, freq_main, sigfft_freq, sigfft

    def align(sig1, sig2):
        """
        ---
        EXPERIMENTAL
        ---

        Aligns the signals to maximize their correlation.
        Uses np.correlate().
        sig1 - ndarray
        sig2 - ndarray

        returns: indexshift - this is the index shift for the time vector for the sig1.
        """
        # correlation
        tcorr = np.correlate(sig1, sig2, mode='full')
        nsamples = sig1.shape[0]
        # shift in index for sig1
        indexshift = (nsamples - (tcorr.argmax() + 1))

        return indexshift, tcorr

#
# === Integrate

def numint(xvec, yvec):

    """ Calculates the integral of the function defined by f(xvec) = yvec
    xvec - ndarray
    yvec - ndarray

    """
    # differential element vector
    dxvec = np.diff( xvec, n=1, prepend=(xvec[1]-xvec[0]) )
    numeric_integral = np.nansum( dxvec * yvec )
    print('WARNING! Please consider using the function `np.trapz()` for improved performance.\n`cl.numint()` will be removed in the future.')
    return numeric_integral

#
# === magnets

def fringe(z, z1, z2, rad, a1):
    """
    Approximation to the longitudinal profile of a multipole from a permanent magnet assembly.

    see Wan et al. 2018 for definition and Enge functions paper (Enge 1964)
    """
    zz1 = (z - z1) / (2 * rad / pc.pi)
    zz2 = (z - z2) / (2 * rad / pc.pi)
    fout = ( (1 / ( 2 *  np.tanh((z2 - z1) / (4 * rad / pc.pi)) ) )
            * (np.tanh(zz1 + a1 * zz1**2 )
               - np.tanh(zz2 - a1 * zz2**2) )
           )

    return fout

#
# === physics specific functions
#
class laser():

    def bwlimit(freq0, timeFWHM, mode='gauss'):
        """ Calculate FWHM BW of a laser pulse, given a central frequency and the pulse length FWHM.

        return freqFWHM, wlFWHM
        """

        bwfactor = {'unit':1.0, 'gauss':0.441, 'sech':0.315}
        freqFWHM = bwfactor[mode] / timeFWHM

        # convert to wavelength
        wlFWHM = (pc.c / freq0**2) * freqFWHM

        return freqFWHM, wlFWHM

    def zrayleigh(w0, wavelength):
        """
            Rayleigh range. For a Gaussian beam, it is the distance from the focus at which the area of the laser spot-size is doubled -> beam radius is increased by sqrt(2).
        """
        zr = pc.pi * w0**2 / wavelength

        return zr

    def waist(w0, z, wavelength, M2):

        """
        Waist of the laser beam with minimum size w0. Distance z away from minimum.
        W(z) = W0 * ( 1 + M2^2  ( z / zR )^2 )
        """

        wz = w0 * np.sqrt( 1 + M2**2 * (z / laser.zrayleigh(w0, wavelength) )**2 )

        return wz

    def fluence(waist, energy):
        """
        The fluence per pulse.
        fluence units are Joules / cm^2

        waist - [m] waist of laser pulse
        energy - [J] energy of laser pulse
        """
        waist = 100 * waist # convert to cm
        flu = energy / ( pc.pi * waist**2)
        return flu

    def efld(waist, power):
        """
        Electric field corresponding to a given power and spot-size.

        """
        E0 = np.sqrt( np.sqrt(2)*power / (pc.epsilon_0*pc.c*pc.pi*waist**2) )
        return E0

    def a0potential(E0, wavelength):
        """
        Normalized vector potential of the laser pulse.
        """
        const = pc.elementary_charge / (2*pc.pi*pc.m_e*pc.c**2)
        a0 = const * E0 * wavelength
        return a0

    def photon_energy(wavelength):
        """
        Returns the energy of a photon with the specified wavelength.
        output is in units of Joules.
        divide by elementary charge ~ 1.6e-19 to convert to eV
        """
        return pc.h * pc.c / wavelength

#
# === beam dynamics

class ebeam():
    """
    class to keep functions organized.
    """

    def gamma_beta_rel(totengMeV, restmassMeV):
        """ returns the relativistic factors gamma and beta for a given particle energy (gmc2) and rest mass (mc2)

        gamma = totengMeV / restmassMeV
        beta = sqrt( 1 - 1/g^2)
        """
        gamma = totengMeV / restmassMeV
        beta = np.sqrt( 1 - 1 / gamma**2)

        return gamma, beta

    def mom_rel(totengMeV, restmassMeV):
        """ returns the relativistic momentum (in units of MeV/c) of the particle with the specified total energy (gmc2) and rest mass (mc2)

        pc = sqrt( E^2 - mc2^2 )

        NB: the output is in units of MeV/c.
        """
        prel = np.sqrt( totengMeV**2 - restmassMeV**2)
        return prel

    def mom_rigid(momMeV, charge=1):
        """ Returns the momentum rigidity for a given momentum in MeV and charge in units of elementary charge units, e.

        B * rho = (1e6 / c) * (1 / charge) * pc
        """
        br = (1e6 / pc.c) * (momMeV / charge)

        return br

    def emit_convert(emit_old, frac_old, frac_new):
        """ Convert to emittance that encompases the given fraction.

        emit_old - float. geometric or normalized emittance that encompases `frac_old` of the beam
        frac_old - float. [0,1] the fraction of the beam encompassed by the emit_old
        frac_new - float. [0,1] new fraction encompased by emittance

        This is just a factor that multiplies the input emittance.
        """

        emit_new = emit_old * (np.log(1 - frac_new) / np.log(1 - frac_old))

        return emit_new

    def beta_drift(s, s0, beta0):
        """
        Beta function in a drift.
        s0 - location of minimum of beta function
        beta0 - beta function minimum
        """
        beta = beta0 * ( 1 + ( (s - s0) / beta0 )**2 )
        return beta

    def phase_drift(s, s0, beta0):
        """
        Phase advance inside a drift section.
        s0 - location of minimum of beta function
        beta0 - beta function minimum

        """
        y = s/beta0
        x = 1 - s0 * (s - s0) / beta0**2
        phi = np.arctan2( y, x )
        return phi

## Transfer matrix formalism

def mat_quad(maglen, gradient, totengMeV, restmassMeV, charge=1, focus='focus'):
    """ the 2x2 1st order matrix for quadrupole focusing.

    1/f = Kquad * L
    f - focal length
    Kquad = gradient / momentumrigidity # this is the S.Y.Lee definition (note this is not squared!)
    L - magnetic length
    """
    mrig = mom_rigid(mom_rel(totengMeV,restmassMeV), charge)
    kquad = gradient / mrig # normalized quadrupole gradient
    if focus == 'focus':
        fsign = -1
    else:
        fsign = 1

    fl = fsign / (kquad * maglen)

    mat = np.array( [[ 1, 0 ],
                     [ 1/fl, 1 ]
                     ])
    return mat

def mat_drift(dlen):
    """ return the 2x2 matrix for a drift
    """
    mat = np.array( [[ 1, dlen ],
                     [ 0, 1 ]
                     ])
    return mat

def make_M_mat(mat):
    """ returns the 3x3 matrix that acts on [beta,alpha,gamma] from the 2x2 matrix for [x,x']
    mat - ndarray with shape [2,2]
    """
    m11 = mat[1,1]
    m12 = mat[1,2]
    m21 = mat[2,1]
    m22 = mat[2,2]

    mat3 = np.array( [[ m11**2, -2*m11*m12, m12**2 ],
                      [ -m11*m21, m11*m22 + m12*m21, -m12*m22 ],
                      [ m21**2, -2*m21*m22, m22**2 ]
                     ])
    return mat3

#
##
### === Transfer matrix beamlines
##
#

class BLElement(object):
    """
    Optical element represented by an ABCD matrix.

    drift:
        properties = {'eletype':'drift',
                      'position':0,
                      'length':0}
    lens:
        properties = {'eletype':'lens',
                      'position':0,
                      'focal_len':0}
    quad:
        properties = {'eletype':'quad',
                      'position':0.0,
                      'maglen':0.0,
                      'gradient':0.0,
                      'totengMeV':0.0,
                      'restmassMeV':0.0,
                      'charge':1,
                      'focus':{'focus','defocus'}}
    """
    def __init__(self, name, eleprops={'eletype':'drift','position':0,'length':0}):

        self.name = name
        self.properties = eleprops

        if self.properties['eletype'] == 'drift':
            self.mat2x2 = np.array([[1, self.properties['length']],
                                      [0, 1]])

        elif self.properties['eletype'] == 'lens':
            self.mat2x2 = np.array([[1, 0],
                                      [-1 / self.properties['focal_len'], 1]])
        elif self.properties['eletype'] == 'quad':
            matargs = [  self.properties['maglen']
                        ,self.properties['gradient']
                        ,self.properties['totengMeV']
                        ,self.properties['restmassMeV']
                        ]
            matkwargs ={  'charge':self.properties['charge']
                         ,'focus':self.properties['focus']
                         }
            self.mat2x2 = mat_quad(*matargs,**matkwargs)

        else:
            print('WARNING: Element type not supported. Created drift with zero length.')

class BeamLine(object):
    """
    Optical beamline class.
    Store elements as matricies (ABCD or electron transfer matricies).

    """

    def __init__(self):

        self.element_list = []
        self.element_names = []
        self.element_position = []

    def add_element(self, element):
        """ Adds an element class to the beamline.

        element is a class BLElement() or similar.
        The element is inserted into the list based on its position, using `bisect`.
        """
        try:
            # insert the element based on its positionin along the beamline
            posi = element.properties['position']
            bisect.insort(self.element_position, posi)
            addindex = self.element_position.index(posi)
            self.element_list.insert(addindex, element)
            self.element_names.insert(addindex, element.name)

        except AttributeError:
            print('ERROR! element does not possess the required attributes. Need class with .name and .properties()')
            return 1

        return 0

    def del_element(self, elementname):
        """ Deletes the specified element from the beamline
        """
        try:
            delindex = self.element_names.index(elementname)
            del self.element_list[delindex]
            del self.element_names[delindex]
            del self.element_position[delindex]

        except ValueError:
            print('WARNING: Beamline does not contain element with that name.')
            return 1

        return 0

    def make_mat(self, senter, sexit, ytransport=False):
        """ Based on the given location in the beamline, find the total transfer matrix.
        Assumes that the elements are sorted by their position. This is the case when the .add_element() method was used to add the element to the beamline.
        """
        # init transport matrix
        transportmatrix = np.eye(2)
        # print(transportmatrix)

        # position temporary
        si = senter

        # include elements up to sexit point
        s0index = bisect.bisect(self.element_position, senter)
        s1index = bisect.bisect(self.element_position, sexit)
        # print(s0index)
        # print(s1index)
        for i,ele in enumerate(self.element_list[s0index:s1index],start=s0index):
            # print(ele.name)
            # drift to element from previous position
            driftlen = self.element_position[i] - si
            # print(driftlen)
            # print(transportmatrix)
            matdrift = BLElement('tempdrift', eleprops={'eletype':'drift','position':0, 'length':driftlen}).mat2x2
            # print(matdrift)
            transportmatrix = np.matmul(matdrift, transportmatrix)
            # print(transportmatrix)
            # element
            if ytransport and (ele.properties['eletype'] == 'quad'):
                elemat2x2 = ele.mat2x2 * np.array([[1,1],[-1,1]])
            else:
                elemat2x2 = ele.mat2x2
            transportmatrix = np.matmul(elemat2x2, transportmatrix)
            # print(transportmatrix)
            try:
                # sold = si
                si = self.element_position[i]
            except IndexError:

                continue

        # drift to final position from last element before final position
        driftlen = sexit - si

        matdrift = BLElement('tempdrift', eleprops={'eletype':'drift', 'position':0, 'length':driftlen}).mat2x2

        transportmatrix = np.matmul(matdrift, transportmatrix)
        # print(transportmatrix)

        return transportmatrix

    def ray_trace(self, invec, inpos, outpos, ytransport=False):
        """ Given a set of input vectors with initial transverse postion and angle, calculate their transport through the beamline at the specified outpos.

        invec - ndarray.shape = (2,n)
        inpos - float position at which invec is specified
        outpos = ndarray.shape = (m,)

        """
        # inital transport to first point
        transmat = self.make_mat(inpos, outpos[0])
        outvec = np.matmul(transmat, invec) # include start position in output
        outvec = np.stack((invec, outvec)) # joins arrays along new axis

        for i,ss in enumerate(outpos[1:], start=1):

            transmat = self.make_mat(outpos[i-1], ss, ytransport=ytransport)
            outvectemp = np.matmul(transmat, outvec[-1])
            outvec = np.concatenate((outvec, [outvectemp]), axis=0)

        outpos = np.insert(outpos, 0, inpos)

        # calculate centroid position
        outvec_posmean = np.array( [ outvec[i,0,:].mean() for i in range(outvec.shape[0]) ] )
        outvec_angmean = np.array( [ outvec[i,1,:].mean() for i in range(outvec.shape[0]) ] )
        outvecMEAN = np.c_[outvec_posmean, outvec_angmean]
        # calculate RMS size and angle
        outvec_posstd = np.array( [ outvec[i,0,:].std() for i in range(outvec.shape[0]) ] )
        outvec_angstd = np.array( [ outvec[i,1,:].std() for i in range(outvec.shape[0]) ] )
        outvecRMS = np.c_[outvec_posstd, outvec_angstd]

        return outpos, outvec, outvecMEAN, outvecRMS

#
# === ICS analytic formulas

def eng_photon(wavelen):
    """ Returns the energy of the photon with the given wavelength.

    wavelen - float. meter.

    returns the energy in MeV
    """
    hbark = 1e-6 * pc.hbar * pc.c * 2 * pc.pi / wavelen / pc.e

    return hbark

def ics_gammaCM(engele, englas, angin):
    """ The relativistic gamma factor of the center of momentum frame.
    engele - float. MeV energy of incoming electron
    englas - float MeV energy of incoming laser photon
    angin - float. radians angle between the incoming electron and photon in the lab frame

    """
    x =  (4 * engele * englas) / mc2**2

    gcm = ( engele + englas ) / ( mc2 * np.sqrt(1 + (x/2) * (1 - np.cos(angin)) ) )

    return gcm

def ics_betaCM(engele, englas, angin):
    """ The beta factor of the center of momentum frame.
    engele - float. MeV energy of incoming electron
    englas - float MeV energy of incoming laser photon
    angin - float. radians angle between the incoming electron and photon in the lab frame

    """
    bcm = np.sqrt( 1 - ( 2*engele*englas*(1 - np.cos(angin)) ) / ( engele + englas )**2 )

    return bcm

def ics_eng_ph(engele, englas, angin, angout):

    """ Calculate the energy of the scattered photon in the outgoing angle.
    This is the full formula when eqs 6 and 11 of Curatolo et al. (2017) are combined.

    engele - float. MeV energy of incoming electron
    englas - float MeV energy of incoming laser photon
    angin - float. radians angle between the incoming electron and photon in the lab frame

    returns engph in MeV
    """
    gcm = ics_gammaCM(engele, englas, angin)
    bcm = ics_betaCM(engele, englas, angin)

    angoutdep = ( 1 + bcm*np.cos(angout) ) / ( np.cos(angout)**2 + gcm**2 * np.sin(angout)**2 )

    engph = englas * gcm**2 * (1 - np.cos(angin)) * angoutdep

    return engph

def ics_sig(eng_ele, eng_las):
    """ Calculates the Compton scattering cross-section as a proportion of the Thomson cross section.
    sigma_T = 8*pi*re^2/3

    eng_ele - float. MeV. total energy of the incoming electron
    eng_las - float. MeV. total energy of the incoming laser photons

    x - unitless. defined as: x = 4*E_e*E_L / (m_e*c^2)^2. It is a measure of the recoil effect due to the high electron energy or photon energy when compared to the rest mass of the electron.

    returns sigma/sigma_T

    !!! NB: if x is too small, numerical erros due to division by small numbers and subtraction of small numbers begin to take over. usually around x<1e-4. In small x limit, sigma/sigma_T is well approximated by simply (1-x).
    """
    x = 4 * eng_ele * eng_las / mc2**2

    sigbysigT = ((3/4)
              *(1/x)
              *(1/2 + 8/x - 1/(2*(1+x)**2) + (1 - 4/x - 8/x**2)*np.log(1 + x) )
              )
    return sigbysigT

def ics_lum_headon(eng_ele, eng_las, num_ele, num_las, sigma_ele_x, sigma_ele_y, sigma_las, reprate=1.0):
    """ Calculate the luminosity for a HEAD-ON collision between Gaussian beams.
    eng_ele - float. MeV. total energy of the incoming electron
    eng_las - float. MeV. total energy of the incoming laser photons
    num_ele - float. number of electrons
    num_las - float. number of laser photons
    sigma_ele_x, sigma_ele_y - float. meter. x and y RMS of electron beam
    sigma_las - float. meter. sigma_las = waist/2
    reprate - float. number of beam-beam collisions per second.

    luminosity =  N_e * N_l * reprate * cross_section / overlap
    """

    overlap = 2*pc.pi*np.sqrt(sigma_ele_x**2 + sigma_las**2)*np.sqrt(sigma_ele_y**2 + sigma_las**2)

    sigThomson = pc.physical_constants['Thomson cross section'][0]
    sigTotal = sigThomson * ics_sig(eng_ele, eng_las)
    lum = num_ele * num_las * reprate  * sigTotal/ overlap

    return lum

#
##
### === Particle Swarm Optimization
##
#

class PSO(object):
    """ Class handles Particle Swarm Optimization (PSO) of the magnet dimensions with a target gradient, strenght, magnetic length, etc.

    Particle Swarm Optimization pseudo code:
    - define seach space, SS
    - init particles inside SS
    - init particle velocity
    - init global best solution, gbest
        - compute cgbest = cost(f(gbest))
    while (termination condition not met):
        - init particle best solution, pbest
            - compute cpbest = cost(f(pbest))
        for particle in particles:

            - compute cost
                cparticle = cost(f(particle))

            - update gbest and pbest:
                if cpbest > cparticle:
                    pbest = particle
                    cpbest = cost(f(pbest))
                    if cgbest > cpbest:
                        gbest = pbest
                        cgbest = cpbest

            - update velocity
                v = v + alpha1*rand(0,1)*(pbest - particle) + alpha2*rand(0,1)*(gbest - particle)
            - update position
                particle = particle + v

    """
    def __init__(self):

        self.phi1 = 2.05
        self.phi2 = 2.05

        self.searchspace = np.array([])
        self.domain = np.array([])
        self.target = np.array([])
        self.nparticles = 0
        self.maxiter = 0
        self.precision = 0

        self.output = {} # dictionary contains PSO output arrays from run_pso()

    def cost(self, current, target):
        """ Calculates the total square difference between the current parameter values and the target values.

        current - np.array(n,)
        target - np.array(n,)
        """

        value = np.sqrt( np.sum( (current - target)**2 ))

        return value

    def velocity(self, vin, xin, pbest, gbest):
        """

        Updates the input velocity.
        vin - np.array(nparticles, ) input velocity
        xin - np.array(nparticles, ) input position
        pbest - np.array(nparticles, ) the best position for the particle so far
        gbest - np.float the best position for any particle
        phi1=2.05, phi2=2.05 - regulate the 'randomness' in the velocity update as described below

        Clerc and Kennedy (2002) noted that there can be many ways to implement the constriction coefficient. One of the simplest methods of incorporating it is the following :
        v_(i+1) = chi * ( v_i + U(0,phi1) * (p_i - x_i) + U(0,phi2) * (pg - x_i) )
        x_(i+1) = x_i + v_i
        where,
        phi = phi1 + phi2 > 4
        chi = 2 / ( phi - 2 + sqrt(phi^2 - 4*phi) )

        When Clerc's constriction method is used, phi is commonly set to 4.1, phi1=phi2 and the constant multiplier chi is approximately 0.7298. This results in the previous velocity being multiploied by 0.7298 and each of the two (p - x) terms being multiplied by a random number limited by 0.7398*2.05 = 1.49618.
        """
        phi = self.phi1 + self.phi2
        chi = 2 / ( phi - 2 + np.sqrt(phi**2 - 4 * phi) )

        vout = chi * ( vin + self.phi1 * np.random.random(xin.shape) * (pbest - xin) + self.phi2 * np.random.random(xin.shape) * (gbest - xin) )

        return vout

    def run_pso(self, function, searchspace, target, nparticles, maxiter, precision, domain, rngseed=42, verbose=True):

        """ Performs a PSO for the given function in the searchspace, looking for the target, which is in the output space.

        function - the function to be optimized. Its domain must include the seachspace and its output must be in the space of target.
        searchspace - np.array((ssdim, 2))
        target - np.array((tdim, ))
        nparticles - number of particles to use in the optimization
        maxiter - maximum number of iterations to the optimization routine
        precision - how close to the target to attemp to get
        domain - absolute boundaries on the trial solutions/particles
        rngseed=42 int. the seed for the numpy RNG.
        outputs:
        xarr - particle positions for all iterations
        varr - particle velocities for all iterations
        parr - best particle position for all iterations
        cparr - cost of best particle position for all iterations
        garr - best global position for all iterations
        cgarr - cost of best global positions for all iterations

        """
        # update attributes

        self.searchspace = searchspace
        self.domain = domain
        self.maxiter = maxiter
        self.precision = precision
        self.nparticles = nparticles

        # search space dimensionality
        if searchspace.shape[1] != 2:
            print('WARNING! searchspace does not have dimenstions (N,2).')
        ssdim = searchspace.shape[0]

        # init particle positions and velocities
        rng = np.random.default_rng(rngseed)
        xpart = rng.random((nparticles, ssdim))
        for ii in range(ssdim):
            xpart[:,ii] = (searchspace[ii,1] - searchspace[ii,0]) * xpart[:,ii] + searchspace[ii,0] # scale the uniform radnom dist

        vpart = np.zeros(xpart.shape)

        # init particle best solution
        pbest = 1.0 * xpart
        cpbest = np.array([ self.cost(function(*xp), target) for xp in pbest ])
        # init global best solutions
        im = np.argmin(cpbest)
        gbest = pbest[im]
        cgbest = cpbest[im]

        # intermediate arrays
        # multiply by 1.0 to make copies not bind references
        xarr = 1.0 * xpart[:,:,None]
        varr = 1.0 * vpart[:,:,None]
        parr = 1.0 * pbest[:,:,None]
        cparr = 1.0 * cpbest[:,None]
        garr = 1.0 * gbest[:,None]
        cgarr = 1.0 * np.array([cgbest])

        iternum = 0

        t1 = time.time()
        tprint = time.time()
        while (iternum <= maxiter) and ( cgbest > precision):
        # while (iternum <= maxiter):

            for pp in range(nparticles):

                # update velocity
                vpart[pp] = self.velocity(vpart[pp], xpart[pp], pbest[pp], gbest)
                # update position
                xpart[pp] = xpart[pp] + vpart[pp]

                # keeps particles inside the absolute boundaries given by `domain`
                xpart[pp] = np.maximum(xpart[pp], domain[:,0])
                xpart[pp] = np.minimum(xpart[pp], domain[:,1])

                # compute cost of new position
                cpp = self.cost(function(*xpart[pp]) , target )

                # update best position
                if cpp < cpbest[pp]:
                    pbest[pp] = xpart[pp]
                    cpbest[pp] = cpp
                if cpp < cgbest:
                    gbest = xpart[pp]
                    cgbest = cpp

            xarr = np.concatenate((xarr, xpart[:,:,None]),axis=2)
            varr = np.concatenate((varr, vpart[:,:,None]), axis=2)
            parr = np.concatenate((parr, pbest[:,:,None]), axis=2)
            cparr = np.concatenate((cparr, cpbest[:,None]), axis=1)
            garr = np.concatenate((garr, gbest[:,None]), axis=1)
            cgarr = np.append(cgarr, cgbest)

            if verbose:
                maxiter10 = int(maxiter / 10)
                if iternum%maxiter10 == 0:

                    print( ( 'finished iterations: {:d}'
                            +'\nelapsed time: {:5.2f} seconds'
                            +'\n---'
                            ).format(*[iternum, time.time() - tprint ])
                        )
                    tprint = time.time()

            iternum += 1

        t2 = time.time()
        if verbose:
            print('optimization took {:5.2f} seconds'.format(*[t2-t1]))

        arrkeys = ['xarr', 'varr', 'parr', 'cparr', 'garr', 'cgarr']
        arrlist = [xarr, varr, parr, cparr, garr, cgarr]

        # update the output dictionary
        self.format_out_dict(arrkeys, arrlist)

        return 0

    def format_out_dict(self, outarrkeys=None, arrlist=None, notes='none'):
        """
        Helper method to format the output dictionary for JSONEncoder.
        outarrkeys - list of strings corresponding the array names of PSO output
        outarrlist - list of output np.ndarrays created in the
        notes = 'none' optionally add a string of notes e.g. 'varied rin13, rin2, and zlen. optimized for a target integrated gradient.'
        """
        # update the output
        infodict = { 'searchspace' : self.searchspace
                    ,'domain' : self.domain
                    ,'target' : self.target
                    ,'nparticles' : self.nparticles
                    ,'maxiter' : self.maxiter
                    ,'precision' : self.precision
                    ,'notes' : notes
                    }
        self.output['info'] = infodict

        if (outarrkeys != None) and (arrlist != None):
            # add arrays to output dict
            for k,a in zip(outarrkeys, arrlist): self.output[k] = a

        return 0

    def save_pso_output(self, outfilepath):
        """
        Saves the output np.ndarrays to a JSON file in the specified output path.
        Requires NdArrEncoder class.time
        """

        with open(outfilepath, 'w') as outfile:
            json.dump(self.output, outfile, cls=NdArrEncoder)

        return 0

    def load_pso_output(self, outfilepath):
        """
        Load a json file created from previous PSO run.

        Expected arrays:
        xarr - particle positions for all iterations
        varr - particle velocities for all iterations
        parr - best particle position for all iterations
        cparr - cost of best particle position for all iterations
        garr - best global position for all iterations
        cgarr - cost of best global positions for all iterations
        """

        with open(outfilepath, 'r') as readfile:
            pso_output = json.load(readfile)

        # convert back to np.ndarray from nested lists
        for k in pso_output.keys():
            if isinstance(pso_output[k], list):
                pso_output[k] = np.array(pso_output[k])

        return pso_output


from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import collections
import enum

from googlecloudsdk.api_lib.storage import cloud_api
from googlecloudsdk.command_lib.storage import errors
from googlecloudsdk.command_lib.storage import plurality_checkable_iterator
from googlecloudsdk.command_lib.storage import storage_url
from googlecloudsdk.command_lib.storage import wildcard_iterator
from googlecloudsdk.command_lib.storage.resources import resource_reference
from googlecloudsdk.core import log
from googlecloudsdk.core.util import debug_output

class BucketSetting(enum.Enum):
  """An enum saying whether or not to include matched buckets."""

  YES = '_yes'
  NO = '_no'
  NO_WITH_ERROR = '_no_with_error'

class RecursionSetting(enum.Enum):
  """An enum saying whether or not recursion is requested."""
  YES = '_yes'
  NO = '_no'
  NO_WITH_WARNING = '_no_with_warning'

class NameExpansionIterator:
  """Expand all urls passed as arguments, and yield NameExpansionResult.

  For each url, expands wildcards, object-less bucket names,
  subdir bucket names, and directory names, and generates a flat listing of
  all the matching objects/files.
  The resulting objects/files are wrapped within a NameExpansionResult instance.
  See NameExpansionResult docstring for more info.
  """

  def __init__(
      self,
      urls_iterable,
      all_versions=False,
      fields_scope=cloud_api.FieldsScope.NO_ACL,
      ignore_symlinks=False,
      include_buckets=BucketSetting.NO,
      preserve_symlinks=False,
      recursion_requested=RecursionSetting.NO_WITH_WARNING,
  ):
    """Instantiates NameExpansionIterator.

    Args:
      urls_iterable (Iterable[str]): The URLs to expand.
      all_versions (bool): True if all versions of objects should be fetched,
        else False.
      fields_scope (cloud_api.FieldsScope): Determines amount of metadata
        returned by API.
      ignore_symlinks (bool): Skip over symlinks instead of following them.
      include_buckets (BucketSetting): Whether to fetch matched buckets and
        whether to raise an error.
      preserve_symlinks (bool): Preserve symlinks instead of following them.
      recursion_requested (RecursionSetting): Says whether or not recursion is
        requested.
    """
    self.all_versions = all_versions

    self._urls_iterator = (
        plurality_checkable_iterator.PluralityCheckableIterator(urls_iterable))
    self._fields_scope = fields_scope
    self._ignore_symlinks = ignore_symlinks
    self._include_buckets = include_buckets
    self._preserve_symlinks = preserve_symlinks
    self._recursion_requested = recursion_requested

    self._top_level_iterator = (
        plurality_checkable_iterator.PluralityCheckableIterator(
            self._get_top_level_iterator()))
    self._has_multiple_top_level_resources = None
    self._url_found_match_tracker = collections.OrderedDict()

  def _get_wildcard_iterator(self, url):
    """Returns get_wildcard_iterator with instance variables as args."""
    return wildcard_iterator.get_wildcard_iterator(
        url,
        all_versions=self.all_versions,
        fetch_encrypted_object_hashes=True,
        fields_scope=self._fields_scope,
        ignore_symlinks=self._ignore_symlinks,
        preserve_symlinks=self._preserve_symlinks,
    )

  @property
  def has_multiple_top_level_resources(self):
    """Returns if the iterator yields plural items without recursing.

    Also returns True if the iterator was created with multiple URLs.
    This may not be true if one URL doesn't return anything, but it's
    consistent with gsutil and the user's probable intentions.

    Returns:
      Boolean indicating if iterator contains multiple top-level sources.
    """
    if self._has_multiple_top_level_resources is None:
      self._has_multiple_top_level_resources = (
          self._urls_iterator.is_plural() or
          self._top_level_iterator.is_plural())
    return self._has_multiple_top_level_resources

  def _get_top_level_iterator(self):
    """Iterates over user-entered URLs and does initial processing."""
    for url in self._urls_iterator:
      original_storage_url = storage_url.storage_url_from_string(url)
      if (
          isinstance(original_storage_url, storage_url.CloudUrl)
          and original_storage_url.is_bucket()
          and self._recursion_requested is not RecursionSetting.YES
          and self._include_buckets is BucketSetting.NO_WITH_ERROR
      ):
        raise errors.InvalidUrlError(
            'Expected object URL. Received: {}'.format(url)
        )
      # Set to True if associated Cloud resource found in __iter__.
      self._url_found_match_tracker[url] = False
      for resource in self._get_wildcard_iterator(url):
        yield url, self._get_name_expansion_result(resource,
                                                   resource.storage_url,
                                                   original_storage_url)

  def _get_nested_objects_iterator(self, parent_name_expansion_result):
    new_storage_url = parent_name_expansion_result.resource.storage_url.join(
        '**')
    child_resources = self._get_wildcard_iterator(new_storage_url.url_string)
    for child_resource in child_resources:
      yield self._get_name_expansion_result(
          child_resource, parent_name_expansion_result.resource.storage_url,
          parent_name_expansion_result.original_url)

  def _get_name_expansion_result(self, resource, expanded_url, original_url):
    """Returns a NameExpansionResult, removing generations when appropriate."""
    keep_generation_in_url = (
        self.all_versions or
        original_url.generation  # User requested a specific generation.
    )
    if not keep_generation_in_url:
      new_storage_url = storage_url.storage_url_from_string(
          resource.storage_url.versionless_url_string)
      resource.storage_url = new_storage_url
    return NameExpansionResult(resource, expanded_url, original_url)

  def _raise_no_url_match_error_if_necessary(self):
    non_matching_urls = [
        url for url, found_match in self._url_found_match_tracker.items()
        if not found_match
    ]
    if non_matching_urls:
      raise errors.InvalidUrlError(
          'The following URLs matched no objects or files:\n-{}'.format(
              '\n-'.join(non_matching_urls)))

  def _raise_error_if_multiple_sources_include_stdin(self, expanded_url):
    if (self._has_multiple_top_level_resources and isinstance(
        expanded_url, storage_url.FileUrl) and expanded_url.is_stdio):
      raise errors.Error(
          'Multiple URL strings are not supported when transferring from'
          ' stdin.')

  def __iter__(self):
    """Iterates over each URL in self._urls_iterator and yield the expanded result.

    Yields:
      NameExpansionResult instance.

    Raises:
      InvalidUrlError: No matching objects found.
    """
    self._has_multiple_top_level_resources = self._top_level_iterator.is_plural(
    )
    for input_url, name_expansion_result in self._top_level_iterator:

      self._raise_error_if_multiple_sources_include_stdin(
          name_expansion_result.expanded_url)

      should_return_bucket = self._include_buckets is BucketSetting.YES and (
          name_expansion_result.resource.storage_url.is_bucket()
      )
      if not resource_reference.is_container_or_has_container_url(
          name_expansion_result.resource) or should_return_bucket:
        self._url_found_match_tracker[input_url] = True
        yield name_expansion_result

      if resource_reference.is_container_or_has_container_url(
          name_expansion_result.resource):
        if self._recursion_requested is RecursionSetting.YES:
          for nested_name_expansion_result in self._get_nested_objects_iterator(
              name_expansion_result):
            self._url_found_match_tracker[input_url] = True
            yield nested_name_expansion_result

        elif not should_return_bucket:
          if self._recursion_requested is RecursionSetting.NO_WITH_WARNING:
            # Does not warn about buckets processed above because it's confusing
            # to warn about something that was successfully processed.
            log.warning('Omitting {} because it is a container, and recursion'
                        ' is not enabled.'.format(
                            name_expansion_result.resource))

    self._raise_no_url_match_error_if_necessary()

class NameExpansionResult:
  """Holds one fully expanded result from iterating over NameExpansionIterator.

  This class is required to pass the expanded_url information to the caller.
  This information is required for cp and rsync command, where the destination
  name is determined based on the expanded source url.
  For example, let's say we have the following objects:
  gs://bucket/dir1/a.txt
  gs://bucket/dir1/b/c.txt

  If we run the following command:
  cp -r gs://bucket/dir* foo

  We would need to know that gs://bucket/dir* was expanded to gs://bucket/dir1
  so that we can determine destination paths (foo/a.txt, foo/b/c.txt) assuming
  that foo does not exist.

  Attributes:
    resource (Resource): Yielded by the WildcardIterator.
    expanded_url (StorageUrl): The expanded wildcard url.
    original_url (StorageUrl): Pre-expanded URL.
  """

  def __init__(self, resource, expanded_url, original_url):
    """Initialize NameExpansionResult.

    Args:
      resource (resource_reference.Resource): Yielded by the WildcardIterator.
      expanded_url (StorageUrl): The expanded url string without any wildcard.
        This value should preserve generation even if not available in
        resource.storage_url. The versionless version of this should be same
        as resource.storage_url if recursion was not requested. This field is
        intended for only the cp and rsync commands.
      original_url (StorageUrl): Pre-expanded URL. Useful for knowing intention.
    """
    self.resource = resource
    self.expanded_url = expanded_url
    self.original_url = original_url

  def __repr__(self):
    return debug_output.generic_repr(self)

  def __str__(self):
    return self.resource.storage_url.url_string

  def __eq__(self, other):
    if not isinstance(other, type(self)):
      return NotImplemented
    # Leave out original_url because two different URLs can expand to the same
    # result. This is a "results" class.
    return (self.resource == other.resource
            and self.expanded_url == other.expanded_url)


# from pudb import set_trace; set_trace()
from typing import List
import math

class Solution1:
    def countQuadruplets(self, nums: List[int]) -> int:
        """TLE, because num_smaller and num_larger takes too long to build.
        """
        N = len(nums)
        # num_smaller[i][j] is the number of vals smaller than nums[i] to the
        # left of nums[j]
        num_smaller = [[0] * N for _ in range(N)]
        # num_larger[i][j] is the number of vals bigger than nums[i] to the
        # right of nums[j]
        num_larger = [[0] * N for _ in range(N)]
        for i in range(N):
            pre = 0
            for j in range(i + 1):
                num_smaller[i][j] = pre + int(nums[j] < nums[i])
                pre = num_smaller[i][j]
        for i in range(N - 1, -1, -1):
            pre = 0
            for j in range(N - 1, i - 1, -1):
                num_larger[i][j] = pre + int(nums[j] > nums[i])
                pre = num_larger[i][j]
        # find j, k of the triplet i, j, k, l, such that nums[j] > nums[k]

        # for j in range(1, N - 2):
        #     for k in range(j + 1, N - 1):
        #         if nums[j] > nums[k]:
        #             print(nums[j], nums[k], num_smaller[k], num_larger[j])
        return sum(num_smaller[k][j] * num_larger[j][k] for j in range(1, N - 2) for k in range(j + 1, N - 1) if nums[j] > nums[k])

class Solution2:
    def countQuadruplets(self, nums: List[int]) -> int:
        """Use a 2D array to handle num_smaller matrix. The observation is that
        if I am handling value 3. I check the idx of value 2. If value 2 is to
        my right, that means we don't have to do anything about value 3. The
        num_smaller of value 3 is the same as num_smaller of value 2.

        If value 2 is to my left, then I only need to fill the values from 2 to
        3, by adding 1 to the last num_smaller of value 2.

        num_larger can be handled in the same way. Thus, we can complete
        num_smaller and num_larger in O(N^2) time with some saveings and
        O(N^2) space.

        Once the prefix sum matrices for values smaller and values bigger are
        created, we pick any pair of nums[j] and nums[k] where j < k and
        nums[j] > nums[k]. Then we just need to find the number of values to
        the left of nums[j] that is smaller than nums[k] (all of these values
        can serve as nums[i]) and the number of values to the right of nums[k]
        that is larger than nums[j] (all of these values can serve as nums[l]).

        The product of these two values are the total number of quadruplets
        that can be formed with the current pair of nums[j] and nums[k].

        O(N^2), 6504 ms, faster than 21.21%
        """
        N = len(nums)
        num_idx_map = {n: i for i, n in enumerate(nums)}
        # num_smaller[i][j] is the number of vals smaller than nums[i] to the
        # left of nums[j]
        num_smaller = {num_idx_map[1]: [0] * N}
        for n in range(2, N + 1):
            i, j = num_idx_map[n], num_idx_map[n - 1]
            if j > i:
                num_smaller[i] = num_smaller[j]
            else:
                num_smaller[i] = num_smaller[j][:]
                pre = num_smaller[j][j]
                for k in range(j, i + 1):
                    num_smaller[i][k] = pre + int(nums[k] < nums[i])
                    pre = num_smaller[i][k]
        # num_larger[i][j] is the number of vals bigger than nums[i] to the
        # right of nums[j]
        num_larger = {num_idx_map[n]: [0] * N}
        for n in range(N - 1, 0, -1):
            i, j = num_idx_map[n], num_idx_map[n + 1]
            if j < i:
                num_larger[i] = num_larger[j]
            else:
                num_larger[i] = num_larger[j][:]
                pre = num_larger[j][j]
                for k in range(j, i - 1, -1):
                    num_larger[i][k] = pre + int(nums[k] > nums[i])
                    pre = num_larger[i][k]
        return sum(num_smaller[k][j] * num_larger[j][k] for j in range(1, N - 2) for k in range(j + 1, N - 1) if nums[j] > nums[k])

class Solution3:
    def countQuadruplets(self, nums: List[int]) -> int:
        """Solution from: https://leetcode.com/problems/count-increasing-quadruplets/discuss/3111697/C%2B%2BorJavaorPython3-Cleanest-DP-with-Clarification-O(n2)

        This is just too rich, too good.

        Let dp[j] be the number of valid triplets i, j, k, such that i < j < k
        and nums[i] < nums[k] < nums[j], given an l > j.

        Thus, as we iterate j from 0 to l - 1, each time nums[j] < nums[l], we
        can add dp[j] to the answer.

        Also, we keep track of the number of vals smaller than nums[l] along the
        way. We use pre_small to keep track of this number.
        Thus, if nums[j] > nums[l], then the number of valid triplets centering
        in nums[j] increase by pre_small because we can include nums[l] into the
        triplet. This is to prepare the DP triplet count for the next round of
        l.

        O(N^2), 3045 ms, faster than 68.70%
        """
        N = len(nums)
        dp = [0] * N
        res = 0
        for l in range(N):
            pre_small = 0
            for j in range(l):
                if nums[j] < nums[l]:
                    res += dp[j]
                    pre_small += 1
                else:
                    dp[j] += pre_small
        return res

sol = Solution2()
tests = [
    ([1,3,2,4,5], 2),
    ([1,2,3,4], 0),
    ([1,3,5,2,4], 1),
]

for i, (nums, ans) in enumerate(tests):
    res = sol.countQuadruplets(nums)
    if res == ans:
        print(f'Test {i}: PASS')
    else:
        print(f'Test {i}; Fail. Ans: {ans}, Res: {res}')


import codecs
import datetime
import logging
import os
import re
import sys
import urllib2
import hashlib

from cghub_python_api import WSAPIRequest, SOLRRequest
from cghub_python_api.utils import urlopen

from django.conf import settings
from django.template.loader import render_to_string
from django.utils import timezone

from .attributes import ATTRIBUTES, SORT_BY_ATTRIBUTES
from .templatetags.search_tags import file_size as file_size_to_str
from .utils import makedirs_group_write

if settings.API_TYPE == 'WSAPI':
    REQUEST_CLASS = WSAPIRequest
elif settings.API_TYPE == 'SOLR':
    REQUEST_CLASS = SOLRRequest

api_logger = logging.getLogger('api.request')

def get_from_test_cache(url, format='xml'):
    """
    Used while testing.
    Trying to get response from cache, if it fails - get response from server and save it to cache.

    :param url: url that passed to urlopen
    :param format: 'xml' or 'json'

    :return: file object
    """
    FORMAT_CHOICES = {
        'xml': 'text/xml',
        'json': 'application/json'
    }
    CACHE_DIR = settings.TEST_CACHE_DIR
    if not os.path.exists(CACHE_DIR) or not os.path.isdir(CACHE_DIR):
        makedirs_group_write(CACHE_DIR)
    md5 = hashlib.md5(url)
    path = os.path.join(CACHE_DIR, '%s_%s.%s.cache' % (
            settings.API_TYPE.lower(), md5.hexdigest(), format))
    if os.path.exists(path):
        if format == 'json':
            return open(path, 'r')
        else:
            return codecs.open(path, 'r', encoding='utf-8')
    headers = {'Accept': FORMAT_CHOICES.get(format, FORMAT_CHOICES['xml'])}
    req = urllib2.Request(url, headers=headers)
    content = urllib2.urlopen(req).read()
    if format == 'json':
        # JSON is always ASCII and uses \uXXXX to represent non-ASCII symbols.
        with open(path, 'w') as f:
            f.write(content)
        return open(path, 'r')
    with codecs.open(path, 'w', encoding='utf-8') as f:
        f.write(content)
    return codecs.open(path, 'r', encoding='utf-8')

def build_wsapi_xml(result):
    """
    Makes solr xml response the same as wsapi response.
    <str name="analysis_id"> -> <analysis_id>, etc.
    """
    files_count = 0
    files_size = 0
    files = []
    for i in range(10):
        if result['filename.%s' % i].text == None:
            break
        files.append({
            'filename': result['filename.%s' % i].text,
            'filesize': result['filesize.%s' % i].text,
            'checksum': result['checksum.%s' % i].text,
            'checksum_method': result['checksum_method.%s' % i].text,
            })
        files_count += 1
        files_size += int(result['filesize.%s' % i].text)
    size_str = file_size_to_str(files_size)
    try:
        files_size, files_units = size_str.split(' ')
        files_size = files_size.replace(',', '.')
    except ValueError:
        files_size = '0'
        files_units = 'KB'
    xml = render_to_string('xml/wsapi_response.xml', {
            'timestamp': datetime.datetime.strftime(
                    timezone.now(), '%Y-%m-%d %H:%M:%S'),
            'result': result,
            'files': files,
            'files_count': files_count,
            'files_size': files_size,
            'files_units': files_units,
            'server_url': settings.CGHUB_DOWNLOAD_SERVER})
    return xml

class RequestBase(REQUEST_CLASS):

    def patch_input_data(self):
        server_url = getattr(settings, 'CGHUB_SERVER', None)
        if server_url:
            self.server_url = server_url
        if settings.API_TYPE == 'SOLR':
            # set SOLR uri if specified
            self.uri = getattr(settings, 'CGHUB_SERVER_SOLR_URI', self.uri)
        # filter sort_by
        if self.sort_by:
            if self.sort_by.find('-') == 0:
                if self.sort_by[1:] not in SORT_BY_ATTRIBUTES:
                    self.sort_by = None
            else:
                if self.sort_by not in SORT_BY_ATTRIBUTES:
                    self.sort_by = None
        # patch for preservation_method filter
        # ffpe: xml_text:("EXPERIMENT_ATTRIBUTE TAG SAMPLE_PRESERVATION TAG VALUE FFPE")
        # frozen: -xml_text:("EXPERIMENT_ATTRIBUTE TAG SAMPLE_PRESERVATION TAG VALUE FFPE")
        if 'preservation_method' in self.query:
            if 'ffpe' in self.query['preservation_method']:
                if 'xml_text' in self.query:
                    if isinstance(self.query['xml_text'], list):
                        pass
                    elif isinstance(self.query['xml_text'], tuple):
                        self.query['xml_text'] = list(self.query['xml_text'])
                    else:
                        self.query['xml_text'] = [str(self.query['xml_text'])]
                else:
                    self.query['xml_text'] = []
                self.query['xml_text'].append('EXPERIMENT_ATTRIBUTE TAG SAMPLE_PRESERVATION TAG VALUE FFPE')
            elif 'frozen' in self.query['preservation_method']:
                if '-xml_text' in self.query:
                    if isinstance(self.query['-xml_text'], list):
                        pass
                    elif isinstance(self.query['-xml_text'], tuple):
                        self.query['-xml_text'] = list(self.query['-xml_text'])
                    else:
                        self.query['-xml_text'] = [str(self.query['-xml_text'])]
                else:
                    self.query['-xml_text'] = []
                self.query['-xml_text'].append('EXPERIMENT_ATTRIBUTE TAG SAMPLE_PRESERVATION TAG VALUE FFPE')
            del self.query['preservation_method']

    def get_source_file(self, url):
        if 'test' in sys.argv:
            return get_from_test_cache(url=url, format=self.format)
        api_logger.debug(urllib2.unquote(url))
        return urlopen(
                url=url,
                format=self.format,
                max_attempts=getattr(settings, 'API_HTTP_ERROR_ATTEMPTS', 5),
                sleep_time=getattr(settings, 'API_HTTP_ERROR_SLEEP_AFTER', 1))

    def patch_result(self, result, result_xml):
        new_result = {}
        for attr in ATTRIBUTES:
            if result[attr].exist:
                new_result[attr] = result[attr].text
        # files
        new_result['files'] = []
        for i in range(10):
            if not result['filename.%d' % i].exist:
                break
            new_result['files'].append({
                'filename': result['filename.%d' % i].text,
                'filesize': int(result['filesize.%d' % i].text),
                'checksum': result['checksum.%d' % i].text,
            })
        try:
            new_result['filename'] = new_result['files'][0]['filename']
            new_result['checksum'] = new_result['files'][0]['checksum']
            new_result['files_size'] = new_result['files'][0]['filesize']
        except KeyError:
            new_result['filename'] = ''
            new_result['checksum'] = ''
            new_result['files_size'] = 0
        return new_result

class RequestID(RequestBase):
    """
    URI: analysisID uri.
    Fields: analysis_id.
    """

    def patch_input_data(self):
        super(RequestID, self).patch_input_data()
        self.uri = getattr(self, 'CGHUB_ANALYSIS_ID_URI', self.uri)
        self.fields = ['analysis_id']

    def patch_result(self, result, result_xml):
        return {'analysis_id': result['analysis_id'].text}

class RequestMinimal(RequestBase):
    """
    URI: analysisDetail uri.
    Fields: analysis_id, last_modified.
    """

    def patch_input_data(self):
        super(RequestMinimal, self).patch_input_data()
        self.uri = getattr(self, 'CGHUB_ANALYSIS_DETAIL_URI', self.uri)
        self.fields = ['analysis_id', 'last_modified']

    def patch_result(self, result, result_xml):
        return {
                'analysis_id': result['analysis_id'].text,
                'last_modified': result['last_modified'].text}

class RequestDetail(RequestBase):
    """
    URI: analysisDetail uri.
    Fields: analysis_id, state, reason, last_modified,
            upload_date, published_date, center_name, study,
            aliquot_id, filename, filesize, checksum,
            sample_accession, legacy_sample_id, disease_abbr,
            tss_id, participant_id, sample_id, analyte_code,
            sample_type, library_strategy, platform,
            refassem_short_name
    """

    def patch_input_data(self):
        super(RequestDetail, self).patch_input_data()
        self.uri = getattr(self, 'CGHUB_ANALYSIS_DETAIL_URI', self.uri)
        self.fields = [
            'analysis_id', 'state', 'reason', 'last_modified',
            'upload_date', 'published_date', 'center_name', 'study',
            'aliquot_id', 'filename', 'filesize', 'checksum',
            'sample_accession', 'legacy_sample_id', 'disease_abbr',
            'tss_id', 'participant_id', 'sample_id', 'analyte_code',
            'sample_type', 'library_strategy', 'platform',
            'refassem_short_name']

class RequestDetailJSON(RequestDetail):
    """
    URI: analysisDetail uri.
    Fields: analysis_id, refassem_short_name,
        legacy_sample_id, center_name, checksum, disease_abbr,
        analyte_code, filename, filesize, library_strategy,
        last_modified, platform, sample_accession, sample_type,
        state, study, upload_date

    Used json response from wsapi/solr.
    """

    def patch_input_data(self):
        super(RequestDetailJSON, self).patch_input_data()
        self.format = self.FORMAT_JSON

    def patch_result(self, result, result_xml):
        try:
            if result.get('files'):
                for f in result['files']:
                    f['checksum'] = f['checksum']['#text']
            elif isinstance(result['filename'], list):
                result['files'] = []
                for i in range(len(result['filename'])):
                    result['files'].append({
                        'filename': result['filename'][i],
                        'checksum': result['checksum'][i],
                        'filesize': result['filesize'][i],
                    })
            result['filename'] = result['files'][0]['filename']
            result['checksum'] = result['files'][0]['checksum']
            result['files_size'] = result['files'][0]['filesize']
        except KeyError:
            result['filename'] = ''
            result['checksum'] = ''
            result['files_size'] = 0
        return result

class RequestFull(RequestBase):
    """
    URI: analysisFull uri.
    Fields: all.
    Raw xml added to results.
    """

    def patch_input_data(self):
        super(RequestFull, self).patch_input_data()
        self.uri = getattr(self, 'CGHUB_ANALYSIS_FULL_URI', self.uri)
        self.fields = None

    def patch_result(self, result, result_xml):
        if settings.API_TYPE == 'WSAPI':
            xml = result_xml
        elif settings.API_TYPE == 'SOLR':
            # create the same xml as WSAPI returns
            xml = build_wsapi_xml(result)
        new_result = super(RequestFull, self).patch_result(result, result_xml)
        new_result['reason'] = result['reason'].text or ''
        new_result['xml'] = xml
        return new_result

class ResultFromSOLRFile(SOLRRequest):
    """
    Used by RequestsTestCase.test_build_wsapi_xml
    """

    def get_source_file(self, url):
        filename = self.query['filename']
        return codecs.open(filename, 'r')

def get_results_for_ids(ids, sort_by=None):
    """
    Obtain all necessary attributes for specified analysis_ids
    """
    if not ids:
        return []
    api_request = RequestDetailJSON(query={'analysis_id': ids}, sort_by=sort_by)
    results = []
    for result in api_request.call():
        results.append(result)
    return results

class SearchByIDs(object):
    """
    Allows to search by multiple ids.
    """

    ID_ATTRS = ('analysis_id', 'aliquot_id', 'participant_id', 'sample_id')

    def __init__(self, ids, request_cls=RequestID, filters=None):
        """
        :param ids: list of ids to search by
        :param request_cls: Request class will be used to obtain results (RequestID, RequestDetail, ...)
        :param filters: filters dict to use in query
        """
        self.request_cls = request_cls
        self.filters = filters or {}
        self.legacy_sample_ids = []
        self.other_ids = []
        id_regexp = re.compile(settings.ID_PATTERN)
        for i in ids:
            if id_regexp.match(i.lower()):
                # other ids: aliquot_id, analysis_id, participant_id, sample_id
                self.other_ids.append(i.lower())
            else:
                self.legacy_sample_ids.append(i.upper())
        self.search()

    def search(self):
        self.results = {}
        if self.other_ids:
            for attr in self.ID_ATTRS:
                self.results[attr] = []
                query = dict(self.filters)
                query[attr] = self.other_ids
                api_request = self.request_cls(query=query)
                found = False
                for result in api_request.call():
                    self.results[attr].append(result)
                    found = True
                # stop if single id was found
                if found and len(self.other_ids) == 1:
                    break

        if self.legacy_sample_ids:
            self.results['legacy_sample_id'] = []
            query = dict(self.filters)
            query['legacy_sample_id'] = self.legacy_sample_ids
            api_request = self.request_cls(query=query)
            for result in api_request.call():
                self.results['legacy_sample_id'].append(result)

    def get_ids(self):
        ids = []
        for attr in self.results:
            for result in self.results[attr]:
                if result['analysis_id'] not in ids:
                    ids.append(result['analysis_id'])
        return ids

    def get_results(self):
        results = []
        ids = []
        for attr in self.results:
            for result in self.results[attr]:
                if result['analysis_id'] not in ids:
                    ids.append(result['analysis_id'])
                    results.append(result)
        return results

    def is_empty(self):
        return not any([val for key, val in self.results.iteritems()])


from skimage import io
import numpy as np

im1 = io.imread('1.pgm')
im2 = io.imread('2.pgm')

x = np.array(im1)
y = np.array(im2)
z = abs(x/100 - y/100)
z = z*100
SAD = np.sum(z)
SSD = np.sum(z**2)
print(SAD, SSD)


__author__ = 'brian'
import atom
import numpy as np

class MStructureObj:
    def __init__(self, data, num_species,num_beg_rules, num_cluster_rules, num_j_rules):
        data = data.split()
        self.num_beg_rules = num_beg_rules
        self.num_cluster_rules = num_cluster_rules
        self.num_j_rules = num_j_rules
        itter = 0
        self.composition = [0] * num_species
        self.num_Atoms = 0
        for i in range(num_species):
            self.num_Atoms += int(data[itter + 1])
            self.composition[i] = int(data[itter + 1])
            itter += 1
        self.phase_name = data[itter + 1]
        self.mag_phase = data[itter + 2]
        self.name = data[itter + 3]
        self.u = float(data[itter + 4])
        self.enrg = float(data[itter + 5])
        a = float(data[itter + 6])
        b = float(data[itter + 7])
        c = float(data[itter + 8])
        self.LCs = [a, b, c]
        find_c = [abs(b - c), abs(a - c), abs(a - b)]
        self.Cindex = find_c.index(min(find_c))
        self.C = self.LCs[self.Cindex]
        self.LCs[self.Cindex] = self.LCs[2]
        self.LCs[2] = self.C
        self.weight = 1.0
        self.BEG_sums = [0.0] * num_beg_rules
        self.Cluster_sums = [0.0] * num_cluster_rules
        self.J_sums = [0.0] * num_j_rules
        self.basis = []
        self.distances = np.ones([self.num_Atoms, self.num_Atoms * 27]) * 100
        self.mins = np.ones([self.num_Atoms, 10]) * 100

    def set_atom_properties(self, index, atom_data):
        atom_data = atom_data.split()
        mag = float(atom_data[1])
        pos = [round(float(atom_data[2]), 5), round(float(atom_data[3]), 5), round(float(atom_data[4]), 5)]
        self.basis.append(atom.AtomObj(index, self.composition, mag, pos, self.Cindex))

    def create_super_cell(self):
        for i in range(-1, 2):
            for j in range(-1, 2):
                for k in range(-1, 2):
                    if abs(i) + abs(j) + abs(k) != 0:
                        for l in range(self.num_Atoms):
                            atom_copy = self.basis[l]
                            mag = atom_copy.mag
                            pos = [atom_copy.a_pos + i, atom_copy.b_pos + j, atom_copy.c_pos + k]
                            self.basis.append(atom.AtomObj(l, self.composition, mag, pos))

    def calculate_distances(self):
        for i in range(self.num_Atoms):
            for j in range(len(self.basis)):
                dist_a = round(self.basis[i].a_pos - self.basis[j].a_pos, 5)
                dist_b = round(self.basis[i].b_pos - self.basis[j].b_pos, 5)
                dist_c = round(self.basis[i].c_pos - self.basis[j].c_pos, 5)
                self.distances[i, j] = round((dist_a ** 2 + dist_b ** 2 + dist_c ** 2) ** (0.5), 5)
                if self.distances[i, j] == 0:
                    self.distances[i, j] = 100

    def calculate_minimums(self):
        dists = [0] * len(self.basis)
        for i in range(self.num_Atoms):
            old_min = 0.0
            for j in range(len(self.basis)):
                dists[j] = self.distances[i, j]
            for j in range(10):
                self.mins[i, j] = self.next_min(dists, old_min)
                old_min = self.mins[i, j]

    def next_min(self, dists, old_min):
        for i in range(len(dists)):
            if round(dists[i], 5) <= old_min:
                dists[i] = 100
        new_min = np.min(dists)
        return new_min

    def check_plane(self, home_atom_index, neighbor_atom_index):
        if round(self.basis[home_atom_index].c_pos, 5) == round(self.basis[neighbor_atom_index].c_pos, 5):
            plane = 'IN'
        else:
            plane = 'OUT'
        if self.phase_name == 'aust':
            plane = 'ALL'
        return plane


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import math
import torch
from torch import nn
import torch.nn.functional as F
import torchvision.models as models

from .network_utils import *

class StateEncoderShallowCNN(torch.nn.Module):
    """
    state encoder (shallow CNN)
    """
    def __init__(self, output_dim, drop):
        super(StateEncoderShallowCNN, self).__init__()
        self.output_dim = output_dim

        self.conv1 = nn.Conv2d(3, 32, 8, stride=4)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 4, stride=2)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 64, 3, stride=1)
        self.bn3 = nn.BatchNorm2d(64)
        self.fc = nn.Linear(64 * 7 * 7, output_dim)
        self.drop = drop
        self.apply(weights_init)

    def forward(self, state):
        x = F.relu(self.bn1(self.conv1(state)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = F.relu(F.dropout(self.fc(x), p=self.drop))
        return x

class StateEncoderResNet(torch.nn.Module):
    def __init__(self, output_dim, drop):
        super(StateEncoderResNet, self).__init__()
        self.output_dim = output_dim
        self.backbone = models.resnet50(pretrained=True)
        for param in self.backbone.parameters():
            param.requires_grad = True
        # self.conv1 = nn.Conv2d(3, 32, 8, stride=4)
        # self.bn1 = nn.BatchNorm2d(32)
        # self.conv2 = nn.Conv2d(32, 64, 4, stride=2)
        # self.bn2 = nn.BatchNorm2d(64)
        # self.conv3 = nn.Conv2d(64, 64, 3, stride=1)
        # self.bn3 = nn.BatchNorm2d(64)
        self.fc = nn.Linear(1000, output_dim)
        self.drop = drop
        self.apply(weights_init)

    def forward(self, x):
        x = self.backbone(x)
        # x = F.relu(self.bn1(self.conv1(x)))
        # x = F.relu(self.bn2(self.conv2(x)))
        # x = F.relu(self.bn3(self.conv3(x)))
        # x = x.view(-1, 64 * 7 * 7)
        x = F.dropout(F.relu(self.fc(x)), p=self.drop)
        return x