"""
Copyright 2017-2018 Government of Canada - Public Services and Procurement Canada - buyandsell.gc.ca

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import logging
import re
from os import makedirs, sys
from os.path import basename, isdir, join as pjoin

import requests
from requests.exceptions import ConnectionError
from von_anchor.tails import Tails

def usage():
    """
    Print usage message.
    """

    print('\nUsage: sync.py <local-tails-dir> <tails-server-host> <tails-server-port> issuer|prover\n')
    print()
    print('where:')
    print('    <local-tails-dir>:   tails directory on anchor host')
    print('    <tails-server-host>: hostname or IP address of remote tails server')
    print('    <tails-server-port>: port for remote tails server')
    print('    issuer|prover:       issuer to upload local-only tails files, prover to download remote-only')

def survey(dir_tails, host, port):
    """
    Return tuple with paths to local tails symbolic links (revocation registry identifiers) and
    revocation registry identifiers that remote server holds.

    Raise ConnectionError on connection failure.

    :param dir_tails: local root tails directory
    :param host: tails server host
    :param port: tails server port
    :return: pair (remote paths to tails links, remote rev reg ids)
    """

    loc = Tails.links(dir_tails)
    url = 'http://{}:{}/tails/list/all'.format(host, port)
    resp = requests.get(url)
    rem = set(resp.json())

    logging.debug('Survey: local=%s, remote=%s', loc, rem)
    return (loc, rem)

def sync_issuer(dir_tails, host, port, local_only):
    """
    Synchronize for issuer: upload any tails files appearing locally but not remotely.

    :param dir_tails: local root tails directory
    :param host: tails server host
    :param port: tails server port
    :param local_only: paths to local tails symbolic links (rev reg ids) without corresponding remote tails files
    """

    for rr_id in local_only:
        url = 'http://{}:{}/tails/{}'.format(host, port, rr_id)
        with open(Tails.linked(dir_tails, rr_id), 'rb') as fh_tails:
            try:
                resp = requests.post(url, files={'file': fh_tails})
                logging.info('Uploaded: url %s status %s', url, resp.status_code)
            except ConnectionError:
                logging.error('POST Connection refused: %s', url)

def sync_holder_prover(dir_tails, host, port, remote_only):
    """
    Synchronize for issuer: download any tails files appearing remotely but not locally.

    :param dir_tails: local root tails directory
    :param host: tails server host
    :param port: tails server port
    :param remote_only: paths to remote rev reg ids without corresponding local tails files
    """

    for rr_id in remote_only:
        dir_cd_id = Tails.dir(dir_tails, rr_id)
        makedirs(dir_cd_id, exist_ok=True)
        url = 'http://{}:{}/tails/{}'.format(host, port, rr_id)
        try:
            resp = requests.get(url, stream=True)
            if resp.status_code == requests.codes.ok:
                re_tails_hash = re.search('filename="(.+)"', resp.headers['content-disposition'])
                tails_hash = re_tails_hash.group(1) if re_tails_hash.lastindex > 0 else None

                if tails_hash:
                    logging.info('Downloaded: url %s tails-hash %s', url, tails_hash)
                    with open(pjoin(dir_cd_id, tails_hash), 'wb') as fh_tails:
                        for chunk in resp.iter_content(chunk_size=1024):
                            if chunk:
                                fh_tails.write(chunk)
                    Tails.associate(dir_tails, rr_id, tails_hash)
                else:
                    logging.error('Downloading url %s, responded with no tails-hash', url)

            else:
                logging.error('Downloading url %s, responded with status %s', url, resp.status_code)
        except ConnectionError:
            logging.error('GET Connection refused: %s', url)

def main(dir_tails, host, port, role):
    """
    Survey local and remote content, dispatch to synchronize issuer or holder-prover.
    """

    if isdir(dir_tails) and port.isdigit() and role in ('issuer', 'prover'):
        try:
            (paths_local, tails_remote) = survey(dir_tails, host, port)

            if role == 'issuer':
                sync_issuer(dir_tails, host, port, set(basename(p) for p in paths_local) - tails_remote)
            else:
                sync_holder_prover(dir_tails, host, port, tails_remote - set(basename(p) for p in paths_local))
        except ConnectionError:
            logging.error('Could not survey tails server at %s:%s - connection refused', host, port)
    else:
        usage()

if __name__ == '__main__':
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)-15s | %(levelname)-8s | %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')

    if len(sys.argv) != 5:
        usage()
    else:
        ARG_DIR_TAILS = sys.argv[1]
        ARG_HOST = sys.argv[2]
        ARG_PORT = sys.argv[3]
        ARG_ROLE = sys.argv[4].lower()

        main(ARG_DIR_TAILS, ARG_HOST, ARG_PORT, ARG_ROLE)


"""
Code to run each Fibonacci implementation and compare run times.

Runs each file in this directory that start "fibonacci*.py".

Usage: python3 test_fib.py <fib_num>
"""

import sys
import os
import glob
import subprocess

def usage(msg=None):
    if msg:
        print('=' * 60)
        print(msg)
        print('=' * 60)
    print(__doc__)

def test_file(path, num):
    result = subprocess.check_output(['python3', path, num]).decode("utf-8")
    return result

# get the number for which the Fibonacci number is required
if len(sys.argv) != 2:
    usage()
    sys.exit(1)

try:
    num = int(sys.argv[1])
except ValueError:
    usage('You must enter an integer number!')
    sys.exit(1)

# get all executable python files that start "fibonacci".
files = glob.glob('./fibonacci*.py')

# now execute each file passing the test number
# accumulate times in a dictionary
results = {}
for path in files:
    path = os.path.basename(path)
    results[path] = subprocess.check_output(['python3', path, f'{num}']).decode("utf-8")

# parse each result string and replace with the time value
# each value string has the form "fibonacci_memo(40)=102334155     took 0.0000310s\n"
# also gather the fastest file name and the time
max_path_width = 0
for (key, value) in results.items():
    (_, time) = value.split('took ')
    time = float(time[:-2])
    results[key] = time
    if len(key) > max_path_width:
        max_path_width = len(key)

# get file+time into a list of tuples and sort by ascending time
results = sorted(results.items(), key=lambda x: x[1])

# now display the results: file | time | speed
(fast_path, fast_time) = results[0]
print(f" {'FILE':{max_path_width}s} |  RUNTIME    | RESULTS")
print(f"{'-'*max_path_width}--+-------------+--------------------------------")
for (path, time) in results:
    if path == fast_path:
        print(f' {path:{max_path_width}s} | {time:10.7f}s | this is the fastest')
    else:
        print(f' {path:{max_path_width}s} | {time:10.7f}s | {time/fast_time:10.1f} times slower')


 #-*- coding: utf-8 -*-
"""
Created on Mon Sep 11 18:08:21 2017

@author: home
"""

#csv ausgabe
from collections import OrderedDict
import csv
import os
import time
def main():
    '''---nur zu Testzwecken---'''
    name="test"
    datain={}
#    gsr=[7,8,9]
#    datain["gsr"]=gsr
    for i in range(2):
        datain.clear()
        hr=[14,15,16]
        rr=[38,39]
        datain["hr"]=hr
        datain["rr"]=rr
#    duration=len(data["hr"])
#    c_rr=len(data["rr"])
#    if c_rr<duration:
#        while c_rr<duration:
#            data["rr"].append(0)
#            c_rr+=1

        writedatacsv(name,datain,1)
        datain.clear()
        gsr=[27,28,29]
        datain["gsr"]=gsr
        writedatacsv(name,datain,0)
        datain.clear()
        gsr=[21,22,23]
        hr=[17,18,19]
        rr=[30,31]
        datain["hr"]=hr
        datain["rr"]=rr
        datain["gsr"]=gsr
        writedatacsv(name,datain,1)
    datain.clear()
    datain["end"]=[time.time()]
    writedatacsv(name,datain,0)
    datain.clear()
    gsr=[20]
    hr=[10]
    rr=[30]
    datain["hr"]=hr
    datain["rr"]=rr
    datain["gsr"]=gsr
    writedatacsv(name,datain,1)

def csvreader(name,append):
    '''---liest CSV file mit Name(name) ein (append==0) oder nur die Keys(append==1)---'''

    reader=csv.DictReader(open(name+'.csv', 'r'))
    dread=OrderedDict()

    #get vorhandene Spalten(not empty)
    if reader.fieldnames:
        try:
            for key in reader.fieldnames:
                dread[key]=[]

        except:append=0
    else:append=0
    #vorhandene Spalten mit Daten befüllen
    if append==0:
        for row in reader:
            for key in row:
                try:
                 if not row[key]=="":
                     dread[key].append(row[key])
                except:next

    return dread

def csvwriter(dread=dict(),f=str(),append=bool()):
    '''überschreibt Daten(dread=dict()) des Files f (!langsamer!) mit (append==0) oder (schneller) hängt diese nur an (append==1)---'''
#            print(for key in dread: dread[key][0])
    line=1
    count=0
    if append==0:
        w = csv.writer(open(f+'.csv', 'w'))
        w.writerow(dread.keys())
    else:
        w = csv.writer(open(f+'.csv', 'a'))

    #über datenlänge iterieren ´= Zeilenanzahl
    while count<line:
        c=[]
        #über Spalten Zeilen schreiben
        for item in dread:

            if len(dread[item])>line:
                line =len(dread[item])
            try:
                c.append(dread[item][count])
            except:
                c.append("")
        count+=1
        w.writerow(c)

def writedatacsv(name=str(),data=dict(),append=bool()):
    '''---liest CSV file mit Name(name) ein und schreibt Daten(data=dict()) komplett neu (!langsamer!) mit (append==0) oder (schneller) hängt diese nur an (append==1)---'''
    exist=os.path.isfile(name +'.csv')

    if exist == 0:

        writer = csv.writer(open(name+'.csv', 'w'))
        writer.writerow(data.keys())

   #---------------------------------------------------------------------------------------
    dread=csvreader(name,append)
    if not dread.items():
        append=0
   #überprüfen ob input Dict(data) schon in csv vorhanden falls nicht erstellen, falls ja Daten anhängen
    for key in data:
        if key in dread:
            for d in range(len(data[key])):
                dread[key].append(data[key][d])

        else:
            dread[key]=[]
            dread[key]=data[key]

    #daten schreiben
    csvwriter(dread,name,append)

if __name__ == '__main__':
    main()


# =====================================================
# CODE AUTHOR: RAUL ESPINOSA
# The views for the Cart. I used the Django
# documentation to learn about views:

# https://docs.djangoproject.com/en/2.1/topics/http/views/
# https://docs.djangoproject.com/en/2.1/ref/views/

# EDIT: Don't know why, but my last push from when I was working in my branch
# didn't push either this or my forms.py file to the Development branch.
# =====================================================

from django.shortcuts import render, redirect, get_object_or_404
# We will need this for views that intend to change data; those
# require HTTP POST requests.
from django.views.decorators.http import require_POST

# This is the Book model from the bookDetails package I made.
from bookDetails.models import Book, Purchase
# These are the cart and cart forms.
from .cart import Cart
from .forms import AddToCartForm
from users.models import Profile

# This is the view that will handle adding/updating items

cartBooks = []  #Store ids of multiples books that are in the shopping cart

def addToCart(request, book_id):
    userCart = Cart(request)
    cartBooks.append(book_id)
    # Attempt to get the Book that has the
    # given id
    book = get_object_or_404(Book, id=book_id)
    # Validate the form for adding the item to the cart
    form = AddToCartForm(request.POST)

    # If the form is successfully validated, we
    # proceed with adding to/updating the book's
    # amount in the cart
    if form.is_valid():
        data = form.cleaned_data
        userCart.add(book=book,
                     amount=data['amount'],
                     change_amount=data['change_amount'])

    # Once finished, the function redirects the user to the page
    # that shows them the contents of their cart
    return redirect('cart:cart_info')

# This view will handle removing items.

def removeFromCart(request, book_id):
    userCart = Cart(request)
    cartBooks.remove(book_id)
    # Same as addToCart function
    book = get_object_or_404(Book, id=book_id)
    cartBooks.remove(book_id)
    # Simply remove the specified Book
    # from the cart
    userCart.remove(book)

    # Again, redirect to cart contents page
    return redirect('cart:cart_info')

# This view will handle adding items to
# the Saved For Later (SFL) list

def addToSFL(request, book_id):
    userCart = Cart(request)

    book = get_object_or_404(Book, id=book_id)

    # Add the specified book to the SFL list
    userCart.addSFL(book)

    return redirect('cart:cart_info')

# This view will handle adding items back
# the cart from the SFL List. Note that removing
# an item from the SFL list is equivalent to removing
# an item from the cart, so we can just use the regular
# remove function for that

def removeFromSFL(request, book_id):
    userCart = Cart(request)

    book = get_object_or_404(Book, id=book_id)

    # Remove the specified book from the SFL
    # list by putting it back in the cart
    userCart.removeSFL(book)

    return redirect('cart:cart_info')

# This view displays the cart and its contents

def cart_info(request):
    userCart = Cart(request)

    # Iterates through the cart and creates a form allowing for
    # the modification of the amount of books in the cart
    for current in userCart:
        current['update_amount_form'] = AddToCartForm(
            initial={'amount': current['amount'],
                     'change_amount': True}
        )

    return render(request, 'cart/info.html', {'userCart': userCart})

# This view displays the checkout page and lets users purchase books (no credit card validation)
def checkout(request):
    userCart = Cart(request)
    User = get_object_or_404(Profile, user=request.user)
    for i in range(len(cartBooks)):
        book = get_object_or_404(Book, id=cartBooks[i])     #Get book using ids
        total_copies = userCart.get_total_copies(book)      #Total copies of each book in the cart
        if Purchase.objects.filter(book=book, User=User):   #Check if user has already bought book
            if total_copies > 0:
                purchase = get_object_or_404(Purchase, book=book, User=User)
                purchase.amount = purchase.amount + total_copies    #Add total copies to amount user has bought
                purchase.save()
        else:
            if total_copies > 1:                            #Check if user wants to buy more than one copy
                Purchase.objects.create(book=book, User=User, amount = (total_copies))  #Create purchase with total copies
            else:
                Purchase.objects.create(book=book, User=User)                           #Create the one copy of a book

    # Remove all books from the cart
    userCart.clear()
    cartBooks.clear()   #Clear list

    return render(request, 'cart/checkout.html', {'userCart': userCart})


# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2012 NEC Corporation
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import logging

from django.core.urlresolvers import reverse_lazy
from django.utils.translation import ugettext_lazy as _

from horizon import exceptions
from horizon import forms
from horizon import tables
from horizon.utils import memoized

from openstack_dashboard import api

from openstack_dashboard.dashboards.cluster_admin.qos_policy \
    import forms as project_forms
from openstack_dashboard.dashboards.cluster_admin.qos_policy. \
    bandwidth_limit_rules import tables as bw_rules_tables
from openstack_dashboard.dashboards.cluster_admin.qos_policy \
    import tables as project_tables

LOG = logging.getLogger(__name__)

class IndexView(tables.DataTableView):
    table_class = project_tables.QosPolicyTable
    template_name = 'cluster_admin/qos_policy/index.html'

    def get_tenants(self):
        try:
            tenants, has_more = api.keystone.tenant_list(
                self.request,
                user=self.request.user.id)
        except Exception as e:
            LOG.error(e.message)
            tenants = []
            exceptions.handle(self.request,
                              _("Unable to retrieve project list."))

        tenant_ids = [tenant.id for tenant in tenants]
        return tenant_ids

    def get_data(self):
        try:
            policies = api.neutron.list_qos_policies(self.request)
        except Exception as e:
            LOG.error(e.message)
            policies = []
            msg = _('policies list can not be retrieved.')
            exceptions.handle(self.request, msg)

        if policies:
            tenants = self.get_tenants()
            policies = [policy for policy in policies if
                policy.tenant_id in tenants]
        return policies

class CreateView(forms.ModalFormView):
    form_class = project_forms.CreateQosPolicy
    template_name = 'cluster_admin/qos_policy/create.html'
    success_url = reverse_lazy('horizon:cluster_admin:qos_policy:index')

class DetailView(tables.MultiTableView):
    table_classes = (bw_rules_tables.BandwidthLimitRulesTable,)
    template_name = 'cluster_admin/qos_policy/detail.html'
    failure_url = reverse_lazy('horizon:cluster_admin:qos_policy:index')

    def get_bw_rules_data(self):
        try:
            policy_id = self.kwargs['policy_id']
            bw_rules = api.neutron.list_bandwidth_limit_rules(
                self.request, policy_id)
        except Exception:
            bw_rules = []
            msg = _('bandwidth limit rule list can not be retrieved.')
            exceptions.handle(self.request, msg)
        return bw_rules

    @memoized.memoized_method
    def _get_data(self):
        try:
            policy_id = self.kwargs['policy_id']
            policy = api.neutron.qos_policy_get(self.request, policy_id)
        except Exception:
            redirect = self.failure_url
            exceptions.handle(self.request,
                              _('Unable to retrieve details for '
                                'qos policy "%s".') % policy_id,
                                redirect=redirect)
        return policy

    def get_context_data(self, **kwargs):
        context = super(DetailView, self).get_context_data(**kwargs)
        context["policy"] = self._get_data()
        return context

class UpdateView(forms.ModalFormView):
    form_class = project_forms.UpdateQosPolicy
    template_name = 'cluster_admin/qos_policy/update.html'
    success_url = reverse_lazy('horizon:cluster_admin:qos_policy:index')

    def get_initial(self):
        policy = self._get_object()
        return {'id': policy['id'],
                'name': policy['name'],
                'description': policy['description']}


# encoding: utf-8

"""
date: 2020/03/01/17/36

"""

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def reverseList(self, head):
        cur, pre = head, None
        while (cur):
            cur.next, pre, cur = pre, cur, cur.next
        return pre


try:
    # for Python2
    from Tkinter import Tk, messagebox, Label, Button, Entry, FileDialog, RIGHT, YES, X, DISABLED, StringVar, Frame, TOP, LEFT
except ImportError:
    # for Python3
    from tkinter import Tk, messagebox, Label, Button, Entry, filedialog, RIGHT, YES, X, DISABLED, StringVar, Frame, TOP, LEFT

import numpy as np
import os
from itertools import groupby
from operator import itemgetter

#Input Form
fields = ('Vehicles Ignored', 'Saturation Flow Headway', 'Min Vehs per Measurement', 'Minimum Valid Cycles')

class SatFlowsCalc:
    def check_params(self, *args):
        for p in self.params:
            try:
                t = int(p.get())
                if t>0:
                    self.run_button.config(state='normal')
                else:
                    self.run_button.config(state='disabled')
                    break
            except ValueError:
                break

    def __init__(self, master):
        self.master = master
        master.title("Saturation Flows Calculator")

        vcmd    = master.register(self.validate) # we have to wrap the command
        vcmd_fp = master.register(self.validate_inputFolder) # we have to wrap the command

        self.fp_dis = StringVar()
        e = Entry(master, textvariable=self.fp_dis,validate="key", validatecommand=(vcmd_fp,'%P'))
        b = Button(master, text="Browse",
                   command=lambda:self.fp_dis.set(filedialog.askdirectory()))
        e.pack()
        b.pack()

        self.vehsIgnored = StringVar()
        self.satFlowHead = StringVar()
        self.minVehsPerMeas = StringVar()
        self.minValidCycles = StringVar()

        self.params = [self.vehsIgnored, self.satFlowHead, self.minVehsPerMeas, self.minValidCycles]

        i=0
        for p in self.params:
            row = Frame(master)
            lab = Label(row, width=22, text=fields[i]+": ", anchor='w')
            self.ent = Entry(row, textvariable=p, validate="key", validatecommand=(vcmd, '%P'))
            self.ent.insert(0,'0')
            row.pack(side=TOP, fill=X, padx=5, pady=5)
            lab.pack(side=LEFT)
            self.ent.pack(side=RIGHT, expand=YES, fill=X)
            p.trace("w", self.check_params)
            i+=1

        self.run_button = Button(master, text='Calculate Saturation Flow',
                                 state=DISABLED, command=lambda: self.Main(self.params))
        self.run_button.pack(padx=5, pady=5,expand=YES, fill=X)

    def validate_inputFolder(self, fp):
        print('in')
        if os.listdir(fp):
            return False
        else:
            return True

    def validate(self, new_text):
        if not new_text: # the field is being cleared
            self.entered_number = 0
            return True
        try:
            self.entered_number = int(new_text)
            return True
        except ValueError:
            return False

    def Main(self,params):
        #Close the window
        self.master.withdraw()
        vehsIgnored = self.vehsIgnored.get()
        satFlowHead = self.satFlowHead.get()
        minValidCycles = self.minValidCycles.get()
        minVehsPerMeas = self.minVehsPerMeas.get()

        csv_output_fp = r'Output_SatFlows.csv'
        dis_input_fp = self.fp_dis.get()
        #--------------------------------------------------------------------------------------------------------------#
        satFlow_data = self.GetDisData(dis_input_fp)

        #Keep the original Data in memory
        wt = satFlow_data
        #Remove the rows of the data where crossVehIdx <= VehsIgnored
        wt = self.RemoveIgnoredVehs(wt,vehsIgnored)
        #Return a dictionary with the total valid crossing time of vehicles for each Signal Head and Simulation
        satFlows = self.CalculateSatFlows(wt,satFlowHead,minValidCycles,minVehsPerMeas)
        #For the valid Signal Heads we can calculate the average of Saturation Flows
        #For the rest we print the list
        avgSatFlows = self.CalculateAvgSatFlows(satFlows)
        self.Output(csv_output_fp,avgSatFlows)
        completion_msg = 'End of execution - Results in {}'.format(csv_output_fp)
        print(completion_msg )
        messagebox.showinfo('Result', completion_msg )

        self.master.destroy()

    def GetDisData(self,dis_input_fp):
        #Function to read the dis files generated by Vissim
        #The output is a table with columns['sigHead','SimRun','cycleStart','timeGap']
        import glob

        sigHeads = []
        Simulation_Runs = []
        cycleStarts = []
        crossVehIdxs = []
        timeGaps = []

        path = dis_input_fp + r'\*.dis'
        for dis_file in glob.glob(path):
            with open(dis_file) as f:
                simRun = int(dis_file.rsplit('_',1)[1].replace('.dis',''))
                lines = f.readlines()
                for line in lines:
                    line = line.strip()
                    if line.startswith('Discharge at'):
                        sigHead = int(line.rsplit(' ', 1)[1])
                    else:
                        first_val = line.split(' ',1)[0]
                        #Try to convert value in float
                        try:
                            cycleStart = float(first_val)

                            #remove the consecutive spaces from the string
                            line_tmp = " ".join(line.split())

                            #First remove everything after the parenthesis
                            timeGaps_inLine = line_tmp.split('(')[0].strip()
                            timeGaps_inLine = timeGaps_inLine.split(' ')
                            del timeGaps_inLine[0]

                            crossVehIdx = 1
                            for timeGap in timeGaps_inLine:
                                sigHeads.append(sigHead)
                                Simulation_Runs.append(simRun)
                                cycleStarts.append(cycleStart)
                                crossVehIdxs.append(crossVehIdx)
                                timeGaps.append(timeGap)
                                crossVehIdx += 1
                        except:
                            continue
        #Create the numpy array
        #np_array = np.array(sigHeads,Simulation_Runs,cycleStarts,crossVehIdxs,timeGaps)
        wtype=np.dtype([('sigHead','i4'),('SimRun','i4'),
                        ('cycleStart','f4'),('crossVehIdx','f4'),
                        ('timeGap','f4')])

        wt=np.empty(len(sigHeads),dtype=wtype)
        wt['sigHead'] = sigHeads
        wt['SimRun'] = Simulation_Runs
        wt['cycleStart'] = cycleStarts
        wt['crossVehIdx'] = crossVehIdxs
        wt['timeGap'] = timeGaps

        wt.sort(order='sigHead')
        return wt

    def RemoveIgnoredVehs(self,inp_tbl,vehsIgnored):
        inds_to_del = np.where(inp_tbl['crossVehIdx'] <= vehsIgnored)[0]
        out_tbl = np.delete(inp_tbl,inds_to_del,axis=0)
        return out_tbl

    def CalculateSatFlows(self,wt,satFlowHead,minValidCycles,minVehsPerMeas):
        """ Return a dictionary with keys [Signal_Head_Simulation_Run]
            and values the total time of vehicles crossing (number of subsequent vehicles with
            acceptable timeGap, greater than minVehsPerMeas) and their count
        """
        satFlows = {}
        for signalHead in np.unique(wt['sigHead']):
            #dictionary to store the sat flows of the signal head for each iteration
            d = {}
            for simRun in np.unique(wt['SimRun']):
                #Find the unique values of cycleStarts for this signalHead and simRun
                sliced_tbl = wt[(wt['sigHead'] == signalHead) & (wt['SimRun'] == simRun)]

                unique_cycleStarts = np.unique(sliced_tbl['cycleStart'])

                if len(unique_cycleStarts) >= minValidCycles:
                    for cycleStart in unique_cycleStarts:
                        #Indices of the vehicles with low (acceptable) relative gap
                        idx_validGaps = np.where((wt['sigHead'] == signalHead) &
                                                 (wt['SimRun'] == simRun) &
                                                 (wt['cycleStart'] == cycleStart) &
                                                 (wt['timeGap'] <= satFlowHead))[0]

                        #Check if there is a sequence of acceptable indices with more element than the limit (minVehsPerMeas)

                        #Group the array in groups of consecutive integers
                        for k, g in groupby(enumerate(idx_validGaps), lambda ix: ix[0]-ix[1]):
                            idxs = map(itemgetter(1), g)

                            if len(idxs) >= minVehsPerMeas:
                                #We have found a sequence of vehicles with more elements than minVehsPerMeas so we can move on
                                crossTime = 0
                                for idx in idxs:
                                    #Add the sum to a dictionary with the number of vehicles that were counted
                                    crossTime += wt['timeGap'][idx]

                                sf = round((3600 * len(idxs)/crossTime),2)
                                d.setdefault(simRun,[]).append(sf)
                                break

                satFlows[signalHead] = d
        return satFlows

    def CalculateAvgSatFlows(self,satFlows):
        #Create a dictionary with the average sat flows for each Signal Head and each iteration
        avgSatFlows = {}

        for signalHead, simRun_SatFlows in satFlows.iteritems():
            #dictionary with the avg sat flows for each iteration
            d = {}
            for simRun, satFlows in simRun_SatFlows.iteritems():
                if len(satFlows) > 0:
                    avg = sum(satFlows)/len(satFlows)
                else:
                    avg=-1

                d[simRun] = round(avg,2)

            avgSatFlows[signalHead] = d

            #Calculate also the avg of all simRuns
            tot = 0
            for k, v in d.iteritems():
                tot += v
            if len(d) > 0:
                d['AllSimRuns'] = round(tot/len(d),2)
            else:
                d['AllSimRuns'] = -1

        return avgSatFlows

    def Output(self,filepath_out,results):
        import csv
        import os
        import collections as cl

        directory = os.path.dirname(filepath_out)

        if not os.path.exists(directory):
            os.makedirs(directory)

        with open(filepath_out,'wb') as f:
            w = csv.writer(f)
            w.writerow(['Signal_Head','Avg_Saturation_Flow'])
            results_ordered = cl.OrderedDict(sorted(results.items()))
            for k, v in results_ordered.iteritems():
                dout = v
                avgAllSimRuns = dout.pop('AllSimRuns',-1)
                l = [[k],[avgAllSimRuns],dout.values()]
                #flatten the list
                row = [i for slst in l for i in slst]
                w.writerow(row)

if __name__ == "__main__":
    root = Tk()
    my_gui = SatFlowsCalc(root)
root.mainloop()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import sys
import re
from string import uppercase, lowercase
from urllib import quote
sys.path.append('../decipherclips')
import decipherclips

pyFlakes = (os, uppercase, lowercase, quote, decipherclips)

class ClipParserException(Exception):
    pass

class ClipFunctionParser(object):
    """Parses out functions from a .vim plugin placing all
    parsed functions in global namespace"""
    FUNC_RGX = re.compile(r'\s*def ([A-Z].*?)\(.*')  # testable funcs are capitalized FUPEP8

    def __init__(self, clipPath):
        try:
            self.document = open(clipPath).readlines()
        except IOError:
            msg = "Could not open plugin: {path}"
            raise ClipParserException(msg.format(path=clipPath))

    def get_indent(self, string):
        return len(string) - len(string.lstrip())

    def parse(self):
        parsing = False
        capture = []
        indent = 0

        for line in self.document:
            if parsing:
                if line.strip() and self.get_indent(line) <= indent:
                    self._exec(''.join(capture))
                    capture = []
                    parsing = False
                    indent = 0
                else:
                    capture.append(line[indent:])
            elif self.FUNC_RGX.match(line):
                parsing = True
                indent = self.get_indent(line)
                capture.append(line.lstrip())

    def _exec(self, code):
        exec(code, globals(), globals())

ClipFunctionParser('../decipher_clips.vim').parse()


from http.client import HTTPConnection
from http.client import HTTPSConnection
from time import strptime
from urllib.parse import urlparse
from urllib.parse import urlencode
from xml.dom.minidom import parseString
import json

class API:
    # Initialize the API class
    def __init__(self, key, secret, base_url = 'users.licensario.com', use_ssl = True, debug = False):
        self.key, self.secret, self.base_url, self.use_ssl, self.debug = key, secret, base_url, use_ssl, debug

    # Build the request URL
    def request_url(self, url):
        if self.use_ssl:
            full_url = "https://"
        else:
            full_url = "http://"
        full_url += self.base_url
        full_url += url
        return full_url

    # Execute the Request
    def do_request(self, method, url, params = {}):
        final_url = url
        if method == 'GET' and params:
            final_url += '?' + urlencode(params)
            request_body = None
        elif (method == 'POST' or method == 'PUT') and params:
            request_body = urlencode(params)
        else:
            request_body = None
        print("Starting {method} {final_url}".format(method = method, final_url = final_url))
        if self.use_ssl:
            con = HTTPSConnection(self.base_url)
        else:
            con = HTTPConnection(self.base_url)
        headers = {}
        headers["Content-Type"] = "application/x-www-form-urlencoded"
        headers["ISV_API_KEY"] = self.key
        headers["ISV_API_SECRET"] = self.secret
        headers["LCNS_DISABLE_SIGN"] = "true"
        con.request(method, final_url, request_body, headers)
        res = con.getresponse()
        body = res.read()
        status = res.status
        if self.debug:
            debug_message = ""
            debug_message += "################################\n"
            debug_message += "REQUEST: {method} {uri}\n".format(method = method.upper(), uri = self.request_url(final_url) )
            if params and ( method == 'POST' or method == 'GET' ):
                debug_message += "PARAMETERS: {params}\n".format(params = params)
            debug_message += "RESPONSE CODE: {status}\n".format(status = status)
            debug_message += "BODY:\n{body}\n".format(body = body)
            debug_message += "================================\n"
            print(debug_message)
        return {'body': body, 'status': status}

    # Ensure that an external user exists
    def ensure_external_user_exists(self, external_user_id, email):
        self.do_request("PUT", "/api/v1/users/external/{eid}".format(eid = external_user_id), { 'email': email })


"""
This file produces the data for Task 1. Running it prompts the user to input initial values for the system,
after which the necessary data is produced and output to a file in the parent directory.

This file must be in the same folder as sources.py to be run successfully.
"""
from sources import *

# ==================================================
# #### Defining constants #####
# ==================================================
G = 4*np.pi**2
AU = 1.495978707E8  # (km)
M_TOT = 1
END_TIME = 100

# ====================================================
# #### Defining derivatives of the object's trajectory#####
# ====================================================
def derivatives_twobody(variables, *_, g=G, m_tot=M_TOT):
    """Calculates the derivatives with respect to time of the variables in the two body system.

    We make use of the EOM for relative motion between the two masses. Note that these EOM have no explicit time
    depencency; the second positional argument of this function is in place for compatibility with runge_kutta(),
    but we don't need any actual time input.

    Args:
        variables: NumPy array containing the relative variables x, y, vx, and vy
        g (optional): Gravitational constant G in Newton's law of gravitation.
        m_tot (optional): Total mass M1 + M2 contained in the system.

    Returns:
        A NumPy array containing the derivatives dx/dt, dy/dt, dvx/dt, and dvy/dt.
    """
    x, y, vx, vy = variables
    r = np.sqrt(x**2 + y**2)  # For use in dvx and dvy
    const = -g * m_tot / r**3   # Ditto

    # We calculate the actual derivatives
    dx = vx
    dy = vy
    dvx = const * x
    dvy = const * y

    return np.array([dx, dy, dvx, dvy])

# ==================================================
# #### Producing data file and running simulation #####
# ==================================================
def main():
    m_1 = set_value("Please enter the mass of the first primary in kg: ")
    m_2 = set_value("Please enter the mass of the second primary in kg: ")
    m1 = m_1 / (m_1 + m_2)
    m2 = m_2 / (m_1 + m_2)
    x = set_value("Please enter the initial x separation between the masses in km: ")/AU
    y = set_value("Please enter the initial y separation between the masses in km: ")/AU
    vx = set_value("Please enter the initial relative x velocity in km/s: ")*3600*24*365.2422/AU
    vy = set_value("Please enter the initial relative y velocity in km/s: ")*3600*24*365.2422/AU
    dt = set_value("Please enter the integration step dt in years (e.g. 0.01): ")
    variables = np.array([x, y, vx, vy])
    produce_data(1, variables, derivatives_twobody, END_TIME, dt, m1, m2)

main()


import xml.dom.minidom
import xml.etree.ElementTree as ET
import numpy as np
import pickle
import pdb
import os

xodr_root = 'xodr_files/'
available_maps = ['Town01', 'Town02', 'Town03', 'Town04', 'Town05']

link_types = ['predecessor', 'successor']

for map_name in available_maps:

    xodr_fname = xodr_root + map_name + ".xodr"
    print(xodr_fname, ' has loaded.')

    with open(xodr_fname) as fi:
        xml_tree = ET.parse(fi)

    parent_map = {c: p for p in xml_tree.iter() for c in p}
    root_f = xml_tree.getroot()

    root_t = ET.Element(map_name)

    for road_f in root_f.iter('road'):
    #     pdb.set_trace()

        road_t = ET.SubElement(root_t, "road")
        road_t.set('id', road_f.attrib['id'])
        road_t.set('junction', road_f.attrib['junction'])

        lane_t = ET.SubElement(road_t, "lane")
        lane_ids = []
        lane_widths =[]

        for lane_f in road_f.iter('lane'):
            if lane_f.attrib['id'] == '0':
                continue

            lane_ids.append(lane_f.attrib['id'])
            for width_f in lane_f.iter('width'):
                lane_widths.append(float(width_f.attrib['a']))
                if (float(width_f.attrib['b']) !=0
                    or float(width_f.attrib['c']) !=0
                    or float(width_f.attrib['d']) !=0):
                    print('lane marking is not constant function') # means width of road changes over d
                    print(road_f.attrib)
#                     print(lane_f.attrib)
#                     print(width_f.attrib)
                break

        lane_t.set('ids', lane_ids)
        lane_t.set('widths', lane_widths)

        for link_type in link_types:
            lane_id_from=[]
            lane_id_to=[]
            for pre_f in road_f.iter(link_type):
                if 'elementType' in pre_f.attrib.keys():
                    pre_t = ET.SubElement(road_t, link_type)
                    pre_t.set('elementType', pre_f.attrib['elementType'])
                    pre_t.set('elementId', pre_f.attrib['elementId'])
                else:
                    lane_id_to.append(pre_f.attrib['id'])
                    lane_id_from.append(parent_map[parent_map[pre_f]].attrib['id'])
            pre_t.set('lane_id_from', lane_id_from)
            pre_t.set('lane_id_to', lane_id_to)

    for junction_f in root_f.iter('junction'):
        junction_t = ET.SubElement(root_t, "junction")
        junction_t.set('id', junction_f.attrib['id'])

        for connection_f in junction_f.iter('connection'):
            connection_t = ET.SubElement(junction_t, "connection")
            connection_t.set('id', connection_f.attrib['id'])
            connection_t.set('incomingRoad', connection_f.attrib['incomingRoad'])
            connection_t.set('connectingRoad', connection_f.attrib['connectingRoad'])
            connection_t.set('contactPoint', connection_f.attrib['contactPoint'])

            for laneLink_f in connection_f.iter('laneLink'):
                laneLink_t = ET.SubElement(connection_t, "laneLink")
                laneLink_t.set('from', laneLink_f.attrib['from'])
                laneLink_t.set('to', laneLink_f.attrib['to'])

    v_tree = ET.ElementTree(root_t)

    fname_light=  xodr_root + map_name+"_light.xodr"
    v_tree.write(fname_light, encoding="utf-8", xml_declaration=True)
    dom = xml.dom.minidom.parse(fname_light)

    pretty_xml_as_string = dom.toprettyxml()

    fname_pretty = xodr_root + map_name+"_light_pretty.xodr"
    with open(fname_pretty, "w") as files:
        files.write(pretty_xml_as_string)
    print(fname_pretty, ' has saved.')
    os.remove(fname_light)


from pygame import *

screen = display.set_mode((800,600))

class BoyAI:

    def __init__(self, x_pos, y_pos, pics, startX, startY, endX, endY):
        self.x_pos = x_pos
        self.y_pos = y_pos
        self.pics = pics #2d list
        self.hitbox = Rect(x_pos - 35, y_pos - 35, 70, 70)
        self.startX = startX
        self.startY = startY
        self.endX = endX
        self.endY = endY
        self.targetX = endX
        self.targetY = endY

        right = []
        down = []
        up = []
        left = []

        for i in range(3):
            right.append(image.load("Boy_AI/right/right" + str(i) + ".png"))
        for i in range(3):
            down.append(image.load("Boy_AI/down/down" + str(i) + ".png"))
        for i in range(3):
            up.append(image.load("Boy_AI/up/up" + str(i) + ".png"))
        for i in range(3):
            left.append(image.load("Boy_AI/left/left" + str(i) + ".png"))

        pics[0] = right
        pics[1] = down
        pics[2] = up
        pics[3] = left

        self.move = 0
        self.frame = 0

    def moveBoy(self, newX, newY):

        newMove = -1

        if self.x_pos < newX:
            newMove = 0
            self.x_pos += 4
        elif self.y_pos < newY:
            newMove = 1
            self.y_pos += 4
        elif self.y_pos > newY:
            newMove = 2
            self.y_pos -= 4
        elif self.x_pos > newX:
            newMove = 3
            self.x_pos -= 4
        else:
            self.frame = 0

        move = self.move

        if move == newMove:  # 0 is a standing pose, so we want to skip over it when we are moving
            self.frame = self.frame + 0.27
            if self.frame >= len(self.pics[self.move]):
                self.frame = 1
        elif newMove != -1:  # a move was selected
            self.move = newMove  # make that our current move
            self.frame = 1

        if self.frame == 0:
            return "Turn"

    def makeMove(self, name, start,end):
        move = []
        for i in range(start,end+1):
            move.append(image.load("%s/%s%d.png" % (name,name,i)))

        return move

    def drawBoyAI(self):

        move = self.move
        frame = self.frame
        pics = self.pics

        pic = pics[move][int(frame)]
        screen.blit(pic,(self.x_pos-pic.get_width()//2,self.y_pos-pic.get_height()//2))

    def drawHitbox(self):
        draw.rect(screen, (255,0,0), self.hitbox)

    def update(self):
        self.hitbox = Rect(self.x_pos - 35, self.y_pos - 35, 70, 70)


import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler

import numpy as np

import matplotlib.pyplot as plt

import time
import os
import copy

from tqdm import tqdm

def train_model(model, dataloaders, transform, criterion, optimizer, scheduler, num_epochs, device, is_inception=False):
    """
    Return:Best Model,val_acc_history
    model:model for training
    dataloaders:dictionary of dataloaders
    critetion:loss function
    optimizer:optimizer
    num_epochs:epochs
    is_inception:True if model is Inception
    """
    since = time.time()

    val_acc_history = []
    val_loss_history = []
    train_acc_history = []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            print("model mode to ",phase)
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            print("init loss")
            running_loss = 0.0
            running_corrects = 0

            print("init transforms")
            dataloaders[phase].dataset.transform = transform[phase]

            # print("Phase:", phase)
            # print("transform:", dataloaders[phase].dataset.transform)
            # Iterate over data.
            for inputs, labels in tqdm(dataloaders[phase]):
                inputs = inputs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # Get model outputs and calculate loss
                    # Special case for inception because in training it has an auxiliary output. In train
                    #   mode we calculate the loss by summing the final output and the auxiliary output
                    #   but in testing we only consider the final output.
                    if is_inception and phase == 'train':
                        # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958
                        outputs, aux_outputs = model(inputs)
                        loss1 = criterion(outputs, labels)
                        loss2 = criterion(aux_outputs, labels)
                        loss = loss1 + 0.4*loss2
                    else:
                        outputs = model(inputs)
                        loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if hasattr(dataloaders[phase],"sampler"):
                epoch_loss = running_loss / len(dataloaders[phase].sampler)
                epoch_acc = running_corrects.double() / len(dataloaders[phase].sampler)
            else:
                epoch_loss = running_loss / len(dataloaders[phase].dataset)
                epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'val':
                val_acc_history.append(epoch_acc)
                val_loss_history.append(epoch_loss)

                # scheduler
                scheduler.step(epoch_loss)
            if phase == 'train':
                train_acc_history.append(epoch_acc)

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model, (val_acc_history, val_loss_history, train_acc_history)


#!/usr/bin/env python
import random
from inspect import signature
from functools import partial, reduce

import torch
import numpy as np
import tensornetwork as tn
import torch.optim as optim

"""
### NOTES ON TENSOR NETWORK FORMAT ###

A tensor network is specified as a list of tensors, each of whose shape is
formatted in a specific way. In particular, for a network with n tensor
cores, the shape will be:

tensor_i.shape = (r_1, r_2, ..., r_i, ..., r_n),

where r_j gives the tensor network rank (TN-rank) connecting cores i and j
when j != i, while r_i := d_i gives the dimension of the input to core i.
This implies that tensor_i.shape[j] == tensor_j.shape[i], and stacking the
shapes of all the tensors in order gives the adjacency matrix of the
network as follows:

     [[d_1, r_{1,2}, ..., r_{1,n}],
      [r_{1,2}, d_2, ..., r_{2,n}],
                     ...
      [r_{1,n}, r_{2,n}, ..., d_n]]

The diagonal entries here are precisely the input dimensions for each of
the core tensors, while the symmetric nature of the matrix is necessary
for different cores to contract with each other.

Besides the above matrix format, the ranks are also sometimes specified
in the following upper-triangular format:

     [[r_{1,2}, r_{1,3}, ..., r_{1,n}], [r_{2,3}, ..., r_{2,n}], ...
      ..., [r_{n-2,n-1}, r_{n-2,n}], [r_{n-1,n}]]
"""

# Set global defaults
tn.set_default_backend("pytorch")
torch.set_default_tensor_type(torch.DoubleTensor)

def continuous_optim(tensor_list, train_data, loss_fun, epochs=10,
                     val_data=None, other_args=dict()):
    """
    Train a tensor network using gradient descent on input dataset

    Args:
        tensor_list: List of tensors encoding the network being trained
        train_data:  The data used to train the network
        loss_fun:    Scalar-valued loss function of the type
                        tens_list, data -> scalar_loss
                     (This depends on the task being learned)
        epochs:      Number of epochs to train for. When val_data is given,
                     setting epochs=None implements early stopping
        val_data:    The data used for validation
        other_args:  Dictionary of other arguments for the optimization,
                     with some options below (feel free to add more)

                        optim: Choice of Pytorch optimizer (default='SGD')
                        lr:    Learning rate for optimizer (default=1e-3)
                        bsize: Minibatch size for training (default=100)
                        reps:  Number of times to repeat
                               training data per epoch     (default=1)
                        print: Whether to print info       (default=True)
                        dyn_print: use dynamic printing    (default=False)
                        hist:  Whether to return losses
                               from train and val sets     (default=False)
                        momentum: Momentum value for
                                  continuous optimization  (default=0)
                        cvg_threshold: threshold to test convergence of
                            optimization (optimization is stopped if
                            |(prev_loss - cur_loss)/prev_loss| <  cvg_threshold
                            If None, convergence is not checked. If epochs is
                            set as well, then optimziation is stopped either when
                            convergence criteria is met or when epochs is reached
                                                            (default:None)
                        lr_scheduler: a function taking an optimizer as input
                        and returning a learning rate scheduler for this optimizer
                                                            (default:None)
                        save_optimizer_state: if True, other_args should have an empty
                            dict for the key optimizer_state. This dict will contain
                              {optimizer_state: optimizer state_dict,
                              lr_scheduler_state: scheduler state_dict (if any)}
                            after the function returns.     (default:False)
                        load_optimzer_state: a dictionnary that will be used to
                            initialize the optimizer (and scheduler if any) from a
                            previously saved optimizer state.
                                                            (default: None)
                        grad_masking_function: a function taking the list of tensor
                            parameters between the backward pass and the optimizer step
                            (can be used to e.g. zero out parts of the gradient)
                                                            (default: None)
                        stop_condition: a function taking the training and validation loss
                            as input after each epoch and returning True if optimization
                            should be stopped               (default: None)

    Returns:
        better_list: List of tensors with same shape as tensor_list, but
                     having been optimized using the appropriate optimizer.
                     When validation data is given, the model with the
                     lowest validation loss is output, otherwise the model
                     with lowest training loss
        first_loss:  Initial loss of the model on the validation set,
                     before any training. If no val set is provided, the
                     first training loss is instead returned
        best_loss:   The value of the validation/training loss for the
                     model output as better_list
        best_epoch:  epoch at which best_model was found
        loss_record: If hist=True in other_args, history of all validation
                     and training losses is returned as a tuple of Pytorch
                     vectors (train_loss, val_loss), with each vector
                     having length equal to number of epochs of training.
                     When no validation loss is provided, the second item
                     (val_loss) is an empty tensor.
    """
    # Check input and initialize local record variables
    early_stop = epochs is None
    has_val = val_data is not None
    optim = other_args['optim'] if 'optim' in other_args else 'SGD'
    lr    = other_args['lr']    if 'lr'    in other_args else 1e-3
    bsize = other_args['bsize'] if 'bsize' in other_args else 100
    reps  = other_args['reps']  if 'reps'  in other_args else 1
    prnt  = other_args['print'] if 'print' in other_args else True
    hist  = other_args['hist']  if 'hist'  in other_args else False
    dyn_print  = other_args['dyn_print']  if 'dyn_print'  in other_args else False
    lr_scheduler  = other_args['lr_scheduler']  if 'lr_scheduler'  in other_args else None
    cvg_threshold  = other_args['cvg_threshold']  if 'cvg_threshold'  in other_args else None
    save_optimizer_state  = other_args['save_optimizer_state']  if 'save_optimizer_state'  in other_args else None
    load_optimizer_state  = other_args['load_optimizer_state']  if 'load_optimizer_state'  in other_args else None
    grad_masking_function  = other_args['grad_masking_function']  if 'grad_masking_function'  in other_args else None
    momentum = other_args['momentum'] if 'momentum' in other_args else 0

    stop_condition  = other_args['stop_condition']  if 'stop_condition'  in other_args else None

    if save_optimizer_state and (not 'optimizer_state' in other_args):
        raise ValueError("an empty dictionnary should be passed as the optimizer_state argument to store the"
            " optimizer state.")
    if early_stop and not has_val:
         raise ValueError("Early stopping (epochs=None) requires val_data "
                          "to be input")
    loss_rec, first_loss, best_loss, best_network, best_epoch = [], None, np.infty, tensor_list, 0
    if hist: loss_record = ([], [])    # (train_record, val_record)

    # Function to maybe print, conditioned on `prnt`
    m_print = lambda s: print(s, end='\r' if dyn_print else '\n') if prnt else None

    # Function to record loss information and return whether to stop
    def record_loss(new_loss, new_network, epoch_num):
        # Load record variables from outer scope
        nonlocal loss_rec, first_loss, best_loss, best_network, best_epoch

        # Check for first and best loss
        if best_loss is None or new_loss < best_loss:
            best_loss, best_network, best_epoch = new_loss, new_network, epoch_num
        if first_loss is None:
            first_loss = new_loss

        # Update loss record and check for early stopping. If you want to
        # change early stopping criteria, this is the place to do it.
        window = 2    # Number of epochs kept for checking early stopping
        warmup = 1    # Number of epochs before early stopping is checked
        if len(loss_rec) < window:
            stop, loss_rec = False, loss_rec + [new_loss]
        else:
            # stop = new_loss > sum(loss_rec)/len(loss_rec)
            stop = (new_loss > max(loss_rec)) and (epoch_num >= warmup)
            loss_rec = loss_rec[1:] + [new_loss]

        return stop

    # Another loss logging function, but for recording *all* loss history
    @torch.no_grad()
    def loss_history(new_loss, is_val):
        if not hist: return
        nonlocal loss_record
        loss_record[int(is_val)].append(new_loss)

    # Function to run TN on validation data
    @torch.no_grad()
    def run_val(t_list):
        val_loss = []

        # Note that `batchify` uses different logic for different types
        # of input, so update batchify when you work on tensor completion
        for batch in batchify(val_data):
            val_loss.append(loss_fun(t_list, batch))
        if has_val:
            val_loss = torch.mean(torch.tensor(val_loss))

        return val_loss

    # Copy tensor_list so the original is unchanged
    tensor_list = copy_network(tensor_list)

    # Record the initial validation loss (if we validation dataset)
    if has_val: record_loss(run_val(tensor_list), tensor_list, 0)

    # Initialize optimizer, using only the keyword args in the
    optim = getattr(torch.optim, optim)
    opt_args = signature(optim).parameters.keys()
    kwargs = {'lr':lr, 'momentum':momentum}    # <- Add new options here
    kwargs = {k: v for (k, v) in kwargs.items() if k in opt_args}
    optim = optim(tensor_list, **kwargs)    # Initialize the optimizer
    if lr_scheduler: # instantiate learning rate scheduler
        scheduler = lr_scheduler(optim)

    if load_optimizer_state:
        optim.load_state_dict(other_args["load_optimizer_state"]["optimizer_state"])
        if lr_scheduler:
            scheduler.load_state_dict(other_args["load_optimizer_state"]["lr_scheduler_state"])

    # Loop over validation and training for given number of epochs
    ep = 1
    prev_loss = np.infty

    while epochs is None or ep <= epochs:

        # Train network on all the training data
        #from copy import deepcopy
        prev_tensor_list = copy_network(tensor_list)
        #prev_tensor_list = tensor_list
        train_loss, num_train = 0., 0
        for batch in batchify(train_data, batch_size=bsize, reps=reps):
            loss = loss_fun(tensor_list, batch)
            optim.zero_grad()
            loss.backward()
            if grad_masking_function:
                grad_masking_function(tensor_list)
            optim.step()

            with torch.no_grad():
                num_train += 1
                train_loss += loss

        train_loss /= num_train

        if lr_scheduler:
            scheduler.step(train_loss)

        loss_history(train_loss, is_val=False)

        val_loss = run_val(tensor_list) if has_val else None

        val_loss_str = f"Val. loss:  {val_loss.data:.10f}" if has_val else ""
        m_print(f"EPOCH {ep} {'('+str(reps)+' reps)' if reps > 1 else ''}\t\t{val_loss_str}\t\t Train loss: {train_loss.data:.10f}\t\t Convergence: {np.abs(train_loss-prev_loss)/prev_loss:.10f}")

        # Get validation loss if we have it, otherwise record training loss
        if has_val:
            # Get and record validation loss, check early stopping condition
            loss_history(val_loss, is_val=True)
            if record_loss(val_loss, copy_network(tensor_list) if has_val else prev_tensor_list, ep) and early_stop:
                print(f"\nEarly stopping condition reached")
                break
        else:
            record_loss(train_loss, copy_network(tensor_list) if has_val else prev_tensor_list, ep)

        if cvg_threshold and np.abs(train_loss-prev_loss)/prev_loss < cvg_threshold:
            print(f"\nConvergence criteria reached")
            break
        if stop_condition and stop_condition(train_loss=train_loss,val_loss=val_loss):
            print(f"\nStopping condition reached")
            break

        prev_loss = train_loss

        ep += 1
    m_print("")

    # Save the optimizer state if needed
    if save_optimizer_state:
        other_args["optimizer_state"]["optimizer_state"] = optim.state_dict()
        if lr_scheduler:
            other_args["optimizer_state"]["lr_scheduler_state"] = scheduler.state_dict()

    if hist:
        loss_record = tuple(torch.tensor(fr) for fr in loss_record)
        return best_network, first_loss, best_loss, best_epoch, loss_record
    else:
        return best_network, first_loss, best_loss

def contract_network(nodes, contractor='auto', edge_order=None):
    """
    Contract a tensor network that has already been 'wired' together

    Args:
        nodes:      One or more nodes in the network of interest. All
                    nodes connected to this one will get contracted
        contractor: Name of the TensorNetwork contractor used to contract
                    the network. Options include 'greedy', 'optimal',
                    'bucket', 'branch', and 'auto' (default)
        edge_order: When expanding to a dense tensor, giving a list of the
                    dangling edges of the tensor is required to set the
                    order of the indices of the (large) output tensor

    Returns:
        output:     Pytorch tensor containing the contracted network,
                    which here will always be a scalar or batch vector
    """
    contractor = getattr(tn.contractors, contractor)
    output = contractor(tn.reachable(nodes), output_edge_order=edge_order)
    return output.tensor

def evaluate_input(tensor_rep, input_list):
    """
    Contract input vectors with large tensor to get scalar output

    Args:
        tensor_rep:  The tensor network whose input indices are being
                     contracted. This currently only supports TN input
                     (i.e. list of properly formatted tensor cores), but
                     please ask if you want support for inputs that are
                     specified as one big dense tensor
        input_list:  Batch of inputs to feed to the cores in our tensor.
                     This can be either a list of matrices with shapes
                     (batch_dim, input_dim_i) or a single Pytorch tensor
                     with shape (num_cores, batch_dim, input_dim)

    Returns:
        closed_list: Scalar or batch tensor giving output of contraction
                     between tensor network and input data
    """
    num_cores = len(tensor_rep)
    assert len(input_list) == num_cores
    assert isinstance(tensor_rep, (tuple, list, torch.Tensor))
    assert len(set(len(inp.shape) for inp in input_list)) == 1

    # Get batch information about our input
    input_shape = input_list[0].shape
    has_batch = len(input_shape) == 2
    assert len(input_shape) in (1, 2)
    if has_batch:
        batch_dim = input_shape[0]
        assert all(i.shape[0] == batch_dim for i in input_list)

        # Generate batch node for dealing with batch dims
        batch_edges = batch_node(num_cores, batch_dim)
        assert len(batch_edges) == num_cores + 1

    # Convert all tensor cores to list of wired nodes
    node_list = wire_network(tensor_rep)

    # Go through and contract all inputs with corresponding cores
    for i, node, inp in zip(range(num_cores), node_list, input_list):
        inp_node = tn.Node(inp)
        node[i] ^ inp_node[int(has_batch)]

        # Explicitly contract batch indices together if we need that
        if has_batch:
            inp_node[0] ^ batch_edges[i]

    return contract_network(node_list)

def tn_inner_prod(tensor_list1, tensor_list2):
    """
    Get inner product of two tensor networks with identical input dims

    Args:
        tensor_list1: List of (properly formatted) tensors that encodes
                      the first (open) tensor network
        tensor_list2: List of (properly formatted) tensors that encodes
                      the second (open) tensor network

    Returns:
        inner_prod:   Scalar giving inner product between input networks
    """
    num_cores = len(tensor_list1)
    assert len(tensor_list1) == len(tensor_list2)
    assert all(tensor_list1[i].shape[i] == tensor_list2[i].shape[i]
                for i in range(num_cores))

    # Wire up each of the networks
    network1 = wire_network(tensor_list1)
    network2 = wire_network(tensor_list2)

    # Contract all input indices together
    for i in range(num_cores):
        network1[i][i] ^ network2[i][i]

    return contract_network(network1 + network2)

def l2_norm(tensor_list):
    """Compute the Frobenius norm of tensor network"""
    return torch.sqrt(tn_inner_prod(tensor_list, tensor_list))

def l2_distance(tensor_list1, tensor_list2):
    """
    Compute L2 distance between two tensor networks with same input dims
    """
    norm1, norm2 = l2_norm(tensor_list1), l2_norm(tensor_list2)
    inner_prod = tn_inner_prod(tensor_list1, tensor_list2)

    return torch.sqrt(norm1**2 + norm2**2 - 2*inner_prod)

def wire_network(tensor_list, give_dense=False):
    """
    Convert list of tensor cores into fully wired network of TN Nodes

    If give_dense=True, the wired network is contracted together and a
    single (large) tensor is returned
    """
    num_cores = len(tensor_list)
    assert valid_formatting(tensor_list)

    # Wire together all internal edges connecting cores
    node_list = [tn.Node(core) for core in tensor_list]
    for i in range(num_cores):
        for j in range(i+1, num_cores):
            node_list[i][j] ^ node_list[j][i]

    if give_dense:
        edge_order = [node[i] for i, node in enumerate(node_list)]
        return contract_network(node_list, edge_order=edge_order)
    else:
        return node_list

def expand_network(tensor_list):
    """
    Contract all cores of a tensor network together, yielding dense tensor
    """
    return wire_network(tensor_list, give_dense=True)

def random_tn(input_dims=None, rank=1):
    """
    Initialize a tensor network with random (normally distributed) cores

    Args:
        input_dims:  List of input dimensions for each core in the network
        rank:        Scalar or list of rank connecting different cores.
                     For scalar inputs, all ranks will be initialized at
                     the specified number, whereas more fine-grained ranks
                     are specified via a square matrix, or in the
                     following triangular format:
                     [[r_{1,2}, r_{1,3}, ..., r_{1,n}], [r_{2,3}, ..., r_{2,n}],
                      ..., [r_{n-2,n-1}, r_{n-2,n}], [r_{n-1,n}]]

    Returns:
        tensor_list: List of randomly initialized and properly formatted
                     tensors encoding our tensor network
    """
    # Convert rank object to list of shapes
    if hasattr(rank, 'shape'):
        # Matrix format
        shape_mat = rank.shape
        assert len(shape_mat) == 2
        assert shape_mat[0] == shape_mat[1]
        shape_list = [tuple(int(r) for r in row) for row in shape_mat]
    elif hasattr(rank, '__len__'):
        if len(set(len(row) for row in rank)) == 1:
            # Matrix-type format
            shape_list = [tuple(int(r) for r in row) for row in rank]
        else:
            # Upper-triangular format
            shape_list = unpack_ranks(input_dims, rank)
    else:
        # Scalar format
        assert hasattr(input_dims, '__len__')
        r, n_c = rank, len(input_dims)
        shape_list = [(r,)*i + (d,) + (r,)*(n_c-1-i)
                      for i, d in enumerate(input_dims)]
    num_cores = len(shape_list)

    # Check that diagonals match input_dims
    if input_dims is not None:
        assert len(input_dims) == num_cores
        assert all(shape_list[i][i] == d for i, d in enumerate(input_dims))

    # Heuristic function used to set stdev of normally-distributed tensor
    # elements, feel free to replace as desired
    def stdev_fun(shape):
        # Square root of num_core'th root of number of elements in shape
        num_el = torch.prod(torch.tensor(shape)).to(dtype=torch.float)
        return torch.exp(torch.log(num_el) / (2 * num_cores))

    # Use shapes to instantiate random core tensors
    # The variance of the normal distributions is chosen to make the tensor norm 1 in expectation.
    tensor_list = []
    for i,shape in enumerate(shape_list):
        std = 1/np.power(torch.prod(torch.tensor(shape)) * shape[i],0.25).to(dtype=torch.float)
        tensor_list.append(std * torch.randn(shape))

    return tensor_list

make_trainable = lambda tensor_list: [t.requires_grad_()
                                      for t in tensor_list]
make_trainable.__doc__ = "Returns trainable version of tensor list"

def copy_network(tensor_list):
    """Returns detached copy of tensor list"""
    my_copy = [t.clone().detach() for t in tensor_list]
    my_copy = [ct.requires_grad_() if t.requires_grad else ct
               for t, ct in zip(tensor_list, my_copy)]
    return my_copy

def generate_regression_data(target_tensor, num_data, noise=1e-5):
    """
    Use a target tensor network to get pair of batch (input, output) data

    Args:
        target_tensor: List of tensors encoding a target tensor network,
                       which is used to generate the data
        num_data:     The number of inputs and outputs to generate
        noise:         Stdev of Guassian noise to add to real output value

    Return:
        rand_ins:      List of matrices, with i'th entry having shape
                       (num_data, input_dim_i), with input_dim_i being
                       the i'th visible dimension of target_tensor
        rand_out:      Vector of length num_data holding noisy outputs
    """
    num_cores = len(target_tensor)
    in_dims = torch.tensor(get_indims(target_tensor))
    rand_ins = [torch.randn((num_data, d)) / torch.sqrt(d.double())
                    for d in in_dims]

    # Convert into tensors when possible
    if len(set(in_dims)) == 1:
        rand_ins = torch.tensor(rand_ins)

    # Produce outputs in small batches
    eval_fun = partial(evaluate_input, target_tensor)
    num, mini_size, rand_out = 0, 50, []
    while num < num_data:
        this_in = [r_in[num:num+mini_size] for r_in in rand_ins]
        rand_out.append(evaluate_input(target_tensor, this_in))
        num += mini_size
    rand_out = torch.cat(rand_out)

    # Add noise to output data
    rand_out += noise * torch.randn_like(rand_out)

    return rand_ins, rand_out

def generate_completion_data(target_tensor, num_data, noise=1e-5):
    """
    Use a target tensor network to get dataset of tensor elements

    Args:
        target_tensor: List of tensors encoding a target tensor network,
                       which is used to generate the data
        num_data:      The number of elements and values to generate
        noise:         Stdev of Guassian noise to add to tensor elements

    Return:
        rand_elms:     Random integer-valued Pytorch matrix of shape
                       (num_cores, num_data) containing random elements
                       that target_tensor is evaluated on
        rand_vals:     Vector of length num_data holding noisy values of
                       target_tensor at elements in rand_elms
    """
    num_cores = len(target_tensor)
    in_dims = torch.tensor(get_indims(target_tensor))

    # Generate rand_elms, then convert into one-hot format
    rand_elms = torch.stack([torch.randint(d, (num_data,)) for d in in_dims])
    one_hot = torch.functional.F.one_hot
    rand_ins = [one_hot(vec, d).double()
                    for vec, d in zip(rand_elms, in_dims)]

    # Produce outputs in small batches
    eval_fun = partial(evaluate_input, target_tensor)
    num, mini_size, rand_vals = 0, 50, []
    while num < num_data:
        this_in = [r_in[num:num+mini_size] for r_in in rand_ins]
        rand_vals.append(evaluate_input(target_tensor, this_in))
        num += mini_size
    rand_vals = torch.cat(rand_vals)

    # Add noise to output data
    rand_vals += noise * torch.randn_like(rand_vals)

    return rand_elms, rand_vals

def unpack_ranks(in_dims, ranks):
    """Converts triangular `ranks` structure to list of tensor shapes"""
    num_cores = len(in_dims)
    assert [len(rl) for rl in ranks] == list(range(num_cores - 1, 0, -1))

    shape_list = []
    for i in range(num_cores):
        shape = [ranks[j][i-j-1] for j in range(i)]               # Lower triangular
        shape += [ranks[i][j-i-1] for j in range(i+1, num_cores)] # Upper triangular
        shape.insert(i, in_dims[i])                               # Diagonals
        shape_list.append(tuple(shape))

    return shape_list

def print_ranks(tensor_list):
    """Print out the ranks of edges in tensor network"""
    all_shapes = torch.tensor([t.shape for t in tensor_list])
    print(torch.triu(all_shapes).numpy())

def get_indims(tensor_list):
    """Return the input dimensions of nodes in tensor network"""
    all_shapes = torch.tensor([t.shape for t in tensor_list])
    return tuple(int(s[i]) for i, s in enumerate(all_shapes))

def valid_formatting(tensor_list):
    """Check if a tensor list is correctly formatted"""
    num_cores = len(tensor_list)
    try:
        shape_mat = torch.tensor([core.shape for core in tensor_list])
        assert torch.all(shape_mat == shape_mat.T)
        return True
    except:
        return False

def batch_node(num_inputs, batch_dim):
    """
    Return a network of small CopyNodes which emulates a large CopyNode

    This network is used for reproducing the standard batch functionality
    available in Pytorch, and requires connecting the `num_inputs` edges
    returned by batch_node to the respective batch indices of our inputs.
    The sole remaining dangling edge will then give the batch index of
    whatever contraction occurs later with the input.

    Args:
        num_inputs: The number of batch indices to contract together
        batch_dim:  The batch dimension we intend to reproduce

    Returns:
        edge_list:  List of edges of our composite CopyNode object
    """
    # For small numbers of edges, just use a single CopyNode
    num_edges = num_inputs + 1
    if num_edges < 4:
        node = tn.CopyNode(rank=num_edges, dimension=batch_dim)
        return node.get_all_edges()

    # Initialize list of free edges with output of trivial identity mats
    input_node = tn.Node(torch.eye(batch_dim))
    edge_list, dummy_list = zip(*[input_node.copy().get_all_edges()
                                  for _ in range(num_edges)])

    # Contract dummy edges as a binary tree via third-order tensors
    dummy_len = len(dummy_list)
    while dummy_len > 4:
        odd = dummy_len % 2 == 1
        half_len = dummy_len // 2

        # Apply third order tensor to contract two dummy indices together
        temp_list = []
        for i in range(half_len):
            temp_node = tn.CopyNode(rank=3, dimension=batch_dim)
            temp_node[1] ^ dummy_list[2 * i]
            temp_node[2] ^ dummy_list[2 * i + 1]
            temp_list.append(temp_node[0])
        if odd:
            temp_list.append(dummy_list[-1])

        dummy_list = temp_list
        dummy_len = len(dummy_list)

    # Contract the 3 or less dummy indices together
    last_node = tn.CopyNode(rank=dummy_len, dimension=batch_dim)
    [last_node[i] ^ dummy_list[i] for i in range(dummy_len)]

    return edge_list

def increase_rank(slim_list, vertex1, vertex2, rank_inc=1, pad_noise=1e-6):
    """
    Increase the rank of one bond in a tensor network

    Args:
        slim_list: List of tensors encoding a tensor network
        vertex1:   Node number for one end of the edge being increased
        vertex2:   Node number for the other end of the edge being
                   increased, which can't equal vertex1
        rank_inc:  Amount to increase the rank by (default 1)
        pad_noise: Increasing the rank involves embedding the original
                   TN in a larger parameter space, and adding a bit of
                   noise (set by pad_noise) helps later in training

    Returns:
        fat_list:  List of tensors encoding the same network, but with the
                   rank of the edge connecting nodes vertex1 and vertex2
                   increased by rank_inc
    """
    num_tensors = len(slim_list)
    assert 0 <= vertex1 < num_tensors
    assert 0 <= vertex2 < num_tensors
    assert rank_inc >= 0
    assert pad_noise >= 0

    # Function for increasing one index of one tensor
    def pad_tensor(tensor, ind):
        shape = list(tensor.shape)
        shape[ind] = rank_inc
        pad_mat = torch.randn(shape) * pad_noise
        return torch.cat([tensor, pad_mat], dim=ind)

    # Pad both of the tensors along the index of the other
    fat_list = slim_list
    fat_list[vertex1] = pad_tensor(fat_list[vertex1], vertex2)
    fat_list[vertex2] = pad_tensor(fat_list[vertex2], vertex1)

    return fat_list

def num_params(tensor_list):
    """
    Get number of parameters associated with a tensor network

    Args:
        tensor_list: List of tensors encoding a tensor network

    Return:
        param_count: The number of parameters in the tensor network
    """
    return sum(t.numel() for t in tensor_list)

def tensor_recovery_loss(tensor_list, target_tensor):
    """
    Compute the L2 distance between our tensor network and a target network

    Args:
        tensor_list:   List of tensors encoding a tensor network
        target_tensor: The tensor network we are trying to approximate

    Returns:
        loss:          The L2 distance

    """
    return l2_distance(tensor_list, target_tensor)

def regression_loss(tensor_list, dataset, p=2):
    """
    Compute the regression loss of TN when used as scalar-valued fun,
    relative to dataset of inputs and outputs

    Args:
        tensor_list: List of tensors encoding a tensor network
        dataset:     Tuple of the form (input_list, targets), with targets
                     a Pytorch vector containing target values, and
                     input_list either a list of matrices with shapes
                     (batch_dim, input_dim_i), or else a single Pytorch
                     tensor with shape (num_cores, batch_dim, input_dim)
        p:           Sets which p-norm is used for the loss (default: 2)

    Returns:
        loss:        The sum of distance between targets and the output of
                     tensor_list when fed the contents of input_list
    """
    # Unpack dataset and evaluate inputs using tensor network
    input_list, targets = dataset
    num_ins = len(targets)
    outputs = evaluate_input(tensor_list, input_list)

    return torch.dist(targets, outputs, p=p)

def completion_loss(tensor_list, dataset, p=2, use_full_tensor=True):
    """
    Compute sum of distances between actual and target tensor elements,
    for all elements with known values in the dataset

    Args:
        tensor_list: List of tensors encoding a tensor network
        dataset:     Tuple of the form (input_elms, targets), with targets
                     a Pytorch vector of target values, and input_elms an
                     integer Pytorch matrix of shape (num_cores, batch_dim)
        p:           Sets which p-norm is used for the loss (default: 2)
        use_full_tensor: if True the tensor_list is contracted to a full tensor
            to extract observed entries (does not scale to very large tensors!),
            if False, efficient contractions with one hots is done.

    Returns:
        loss:        The sum of distance between targets and the output of
                     tensor_list when fed the contents of input_list
    """
    # Convert input_elms to one-hot vectors, then call regression loss
    if use_full_tensor:
        full_tensor = wire_network(tensor_list,give_dense=True)
        idx, vals = dataset
        return torch.dist(full_tensor[idx.chunk(chunks=idx.shape[-1])],vals, p=p)
    else:
        input_elms, targets = dataset
        in_dims = get_indims(tensor_list)
        one_hot = torch.functional.F.one_hot
        input_list = [one_hot(vec, d).double()
                          for vec, d in zip(input_elms, in_dims)]
        return regression_loss(tensor_list, (input_list, targets), p=p)

@torch.no_grad()
def batchify(dataset, batch_size=100, reps=1):
    """
    Convert dataset into iterator over minibatches of data

    The nature of this iteration depends on the task at hand, so batchify
    should be modified when new tasks are proposed. Currently supported
    options are tensor recover and regression
    """
    # Figure out which task we're dealing with
    if dataset is None: return  # Trivial input data with no iteration

    elif isinstance(dataset, list) and valid_formatting(dataset):
        task = 'recovery'
    elif len(dataset) == 2 and isinstance(dataset[0][0], torch.Tensor):
        num_ax = len(dataset[0][0].shape)
        assert num_ax in [1, 2]
        task = 'regression' if num_ax == 2 else 'completion'
        inputs, targets = dataset
        tensor_input = isinstance(inputs, torch.Tensor)
    else:
        raise NotImplementedError

    # Loop until reps is 0
    while reps > 0:
        # For tensor recovery, just give the target tensor
        if task == 'recovery' or batch_size < 0:
            yield dataset

        # For regression, return minibatches of (input, target) data
        elif task in ['regression', 'completion']:
            ind = 0
            while ind < len(dataset[1]):
                inp = [ip[ind: ind+batch_size] for ip in inputs]
                if tensor_input or task == 'completion':
                    inp = torch.stack(inp)
                tar = targets[ind: ind+batch_size]
                ind += batch_size
                yield inp, tar
        reps = reps - 1
    return

### OLDER CODE ###

def get_repeated_Indices(list_of_indices):
    #input: string of indices for the tensors.
    #output: string of repeated indices. code takes in indices and output only the repeated indices
    #Ex: Suppose tensor A has index ijkp, and tensor B has index klpm.
    #then get_repeated_Indices('ijkp', 'klpm') will return the following string: 'kp'
    myList = list_of_indices
    #convert List to string
    myString =''.join(myList)
    #break the string into indivual list of characters ex.  'abc' ->['a','b', 'c']
    myList = list(myString)
    #get the repeated frequencies of each indices
    my_dict = {i:myList.count(i) for i in myList}

    repeatedList = []
    for item in my_dict:
        if my_dict[item] > 1:
            repeatedList.append(item)
    return repeatedList

def remove_Repeated_indices(List_of_indices):
    #inputs: tensor indices in the form of string
    #output: string of non repeated indicies
    #Ex: remove_Repeated_indices('abc', 'cde')
    #output of the example would be: 'abde'

    myList = List_of_indices
    #turn myList into String: Ex: ['abc','cde'] -> 'abccde'
    myString = ''.join(myList)
    #turn back into lists again: Exp: from 'abccde' -> ['a','b','c','c','d','e']
    myList = list(myString)
    repeated_indices = get_repeated_Indices(List_of_indices)
    #print('the repeated list of indices are:', repeated_indices)
    unique_indices = []
    #now we remove repeated indices from myList
    for item in myList:
        if item not in repeated_indices:
            unique_indices.append(item)
    uniqueString = ''.join(unique_indices)
    return uniqueString

# TODO: This is a fantastic use of einsum, but should probably be rewritten
#       for two reasons (even though it breaks my heart):
#   (1) If the entire contraction of the tensor network is specified through
#       a single einsum string, we can only handle 52 (=2*26) different
#       edges, which means at most 9 nodes in our network. This might be
#       sufficient, but situations with 10 nodes could definitely still pop
#       up, and we don't want to be limited by this.
#   (2) We still need to generate indxList in order to use this function,
#       and indxList is currently written by hand. This isn't possible when
#       we have a variable number of nodes in our network.
def einSum_Contraction(tensorList, indxList):  #<----should rename this to einSum_Contraction to replace old code
    #Purpose: this function takes a list of tensors, and list of indices, and indix to contract and uses einstien summation to perform contraction
    #ex: tensorList = [tensor1, tensor2, tensor3]
    #indxList   = [indx1, indx2, indx3]
    myList = []
    uniqueIndices = remove_Repeated_indices(indxList)
    inputIndices = [indxList]
    N = len(indxList)
    #myList = [indx1, ',',indx2,',',indx3,'->', uniqueIndices]
    for i in range(N - 1):
        myList.append(indxList[i])
        myList.append(',')
    myList.append(indxList[N-1])
    myList.append('->')
    myList.append(uniqueIndices)
    #convert myList to a string: i.e.  [indx1, ',',indx2,',',indx3,'->', uniqueIndices]  - >'ijk,klm,mjp->ilp'
    myString = ''.join(myList)
    #print('myString = ', myString)
    C = torch.einsum(myString, tensorList)
    return C

def padTensor(tensor, pad_axis):
    #this is for the discrete optimization
    #this function takes a tensor and append an extra dimension of ~ zeros along the specified axis (we call the pad axis)
    if pad_axis == -1:
        return tensor #don't pad anything
    tensorShape = list(tensor.shape)
    tensorShape[pad_axis] = 1  #increase the dimension up by 1
    zerosPad = torch.rand(tensorShape) *1e-6  #pad with values approx. equal to zero
    padded_tensor = torch.cat([tensor, zerosPad], pad_axis)
    #print('padded_tensor.shape = ', padded_tensor.shape)
    #print('padded_tensor function output = ', padded_tensor)
    return padded_tensor

def increaseRank(Tensor1, Tensor2, indx1, indx2):
    # The indx 1 and index2 represents the indices for tensor 1 and 2 respectively.
    #There is only one repeated index in the list (indx1, indx2). The repeated index represents the shared edge between
    #the two tensors. For ex: ijkl, lmno
    alpha = get_repeated_Indices([indx1, indx2])
    if len(alpha) != 0 :
        #convert alpha to string
        alpha = ''.join(alpha)
        # find the position of the repeated index alpha in indx1 and indx2
        padAxes1 = indx1.index(alpha)
        padAxes2 = indx2.index(alpha)
        Tensor1 = padTensor(Tensor1, padAxes1)
        Tensor2 = padTensor(Tensor2, padAxes2)
    return  Tensor1, Tensor2

def Tensor_Generator(TensorDimension):
#input: desired target tensor dimension in the form of a list. Ex: input d1xd2xd3 as [d1, d2, d3]
#output: target tensor with random entries drawn from a normal distribution of mean=0 and variance=1
    Tensor = torch.randn(TensorDimension)
    return Tensor

def getOneData_point(W):
# W is input tensor
#X are drawn from a normal distribution of mean=0 and variance=1
#output: Xi, yi  = targetTensor * X
    indxList = []
    alphabet = 'abcdefghijklmnopqrstuvwxyz'
    Xi = Tensor_Generator(W.shape) #generate Tensor Xi with same shape as input tensor W
    for j in range(len(Xi.shape)):
        indxList.append(alphabet[j])
        indx = ''.join(indxList)
    tensorList = [W, Xi]
    #W and X have same index so that yi is a scaler
    yi = einSum_Contraction([W,Xi], [indx, indx])
    return Xi, yi, indx  #W and Xi shares the same indx

def data_Set_Generator(tensor, N):
    #N = number of training data you want
    #tensor is usually the target tensor
    Xi_set = []
    yi_set = []
    for i in range(N):
        Xi, yi, indx = getOneData_point(tensor)
        Xi_set.append(Xi)
        yi_set.append(yi)
    return Xi_set, yi_set,indx  #tensor and Xi shares the same indx

def getYi_set(Xi_set, tensor, indxList):
    #this function generates a N number of yi's using yi = tensor*Xi_set[i].
    #The input tensor is usually either the approx. tensor or the target tensor.
    #indxList = [indx_Xi, indx_tensor]
    yi_set = []
    N = len(Xi_set)      #is there are N elements of Xi in Xi_set, there will be N elements of yi in yi_set
    for i in range(N):
        yi = einSum_Contraction([tensor,Xi_set[i]], indxList)
        yi_set.append(yi)
    return yi_set

def innerProduct_Tensor(T,A):
    #input: inner product of two tensors T and A of same dimension equals the sum of the product of their entries
    #covert T, and A tensors in to 1 D tensors
    T = T.view(1,-1)   #convert tensor  to 1D
    T = T.squeeze()    #squeeze out any extra dimension
    A = A.view(1,-1)
    A = A.squeeze()
    #perform inner product of two 1D tensors
    yi = sum(torch.mul(T,A))
    return yi

# TODO: Code should work for an arbitrary number of cores and edges in the
#       tensor network, not just for 4
def printRank(TensorList):
    [r5,r6,r7,r8] = (TensorList[4].shape) #tensor G dimension
    [r4,d4,r3,r5] = (TensorList[3].shape) #tensor D dimension
    [r2,d1,r1,r7] = (TensorList[0].shape) #tensor A dimension
    #print(' [r1, r2, r3,r4,r5,r6,r7,r8] = ', '[', r1, ',', r2, ',', r3, ',', r4, ',', r5,',', r6, ',', r7, ',', r8, ']')
    r = [r1,r2,r3,r4,r5,r6,r7,r8]
    return r

# TODO: Code should work for an arbitrary number of cores and edges in the
#       tensor network, not just for 4
def getNumParams(r):
    #r = list of ranks of a tensor
    numParam = d1*r[0]*r[6]*r[1] + d2*r[0]*r[7]*r[3]+ d3*r[1]*r[5]*r[2] + d4*r[3]*r[4]*r[2] + r[4]*r[5]*r[6]*r[7]
    return numParam

# TODO: Code should work for an arbitrary number of cores and edges in the
#       tensor network, not just for 4
def getNumParams_stoch(r):
    #r = [r0,r1,r2,r3,...,r7]
    #we want to compuate the dimension of the block we start with. For exmape
    #we started with 4 order tensor, which can be approximated by 5 tensors
    #below computes the number of params for each of the tensor
    numParam1 = d1*r[0]*r[6]*r[1]
    numParam2 = d2*r[0]*r[7]*r[3]
    numParam3 = d3*r[1]*r[5]*r[2]
    numParam4 = d4*r[3]*r[4]*r[2]
    numParamCore = r[4]*r[5]*r[6]*r[7]
    numParamList = [numParam1, numParam2, numParam3, numParam4, numParamCore]
    return numParamList

def computeLoss_Regression(targetData, approxTensor):
    # Unpack targetData
    yi_set, Xi_set = targetData

    #N is total number of training/test data
    sum = 0
    N = len(yi_set)
    for i in range(N):
        # y_approx = innerProduct(ApproxTensor, X_i)
        y_approx = innerProduct_Tensor(approxTensor, Xi_set[i])
        loss = (yi_set[i] - y_approx)**2
        sum = sum + loss
    #divide by N to average the squared error of the cost function  so that the cost function doesn't depend on the number
    #of elements in the training set.
    total_loss = 1/(2*N)*sum
    return(total_loss)

def computeLoss_Factorization(targetTensor, approxTensor):
    cost = torch.norm(approxTensor - targetTensor, 'fro')
    return(cost)

def get_RandomSeqence(seqLength,d):
    #the function generates a random sequence of length 'seqLength' with range from 1 to range d.
    #For examplae, if we have 8 ranks, then seqLength. Code will generate 8 random numbers in the range between
    #1 to d inclusive.
    r = []
    for i in range(seqLength):
        r.append(random.randint(1,d))
    return r

def get_Next_randomEdge(r_list):
    #this function is for random walk
    #input: r_list = [r0,r1,...]
    nextDirection = random.randint(0, len(r_list)-1)
    print('nextDirection = ', nextDirection)
    r_list[nextDirection] = 1 + r_list[nextDirection]
    return r_list

def solve_Continuous(targetData, tensorList, indxList, iterNum, lossFun,
                     hyperparams):
#input: list of tensors and their corresponding indices
#Goal: The purpose of this function is to solve the innerloop of the optimization for the problem
    len_Tensor = len(tensorList)
    len_Indx   = len(indxList)
    for i in range(len_Tensor):
        tensorList[i] = tensorList[i].detach()
        tensorList[i].requires_grad = True

    # Unpack hyperparameters
    lr = hyperparams['lr']
    momentum = hyperparams['momentum'] if 'momentum' in hyperparams else None

    #defines a SGD optimizer to update the parameters
    #optimizer = optim.SGD(tensorList lr = 0.001, momentum=0.2)
    optimizer = optim.Adam(tensorList, lr=0.009)
    #initialize parameters
    LostList = []              #use this to plot the lost function
    for i in range(iterNum):
        optimizer.zero_grad()
        tensor_approx = einSum_Contraction(tensorList, indxList)
        #loss_fn = computeLoss(tensor_approx, target_Tensor)   # this is for tensor decomp: ||W_target - W_approx||_Fˆ2
        loss_fn = lossFun(targetData, tensor_approx)
        loss_fn.backward()
        optimizer.step()                # the new A,B,C will be A_k+1,B_k+1, C_k+1 after optimizer.step
        LostList.append(float(loss_fn))
    return tensorList, indxList, LostList


"""Exercise 4: Use the Call Stack"""

def fibonacci(n):
    """Calculates the fibonacci number of n."""
    # Base case
    if n <= 0:
        return 0
    elif n == 1:
        return 1

    # Recursive case
    return fibonacci(n-1) + fibonacci(n-2)

if __name__ == '__main__':
    result = fibonacci(6)
    print(result)