import re
from itertools import combinations

import helpers

p = re.compile(r"mem.(\d+). = (\d+)")

def part01(lines):
    collection = dict()
    # current_value = None
    current_mask = None
    for line in lines:
        if line.startswith("mask"):
            current_mask = list(line.split()[2])

        if line.startswith("mem"):
            s = p.search(line)
            mem, val = [int(i) for i in s.groups()]
            val = list(str(bin(val))[2:])
            current_value = collection.get(mem, ["0"] * 36)

            val = ["0"] * (36 - len(val)) + val

            r = []
            for old, msk, new in zip(
                reversed(current_value), reversed(current_mask), reversed(val)
            ):
                if msk == "X":
                    r.append(new)
                elif msk == "1":
                    r.append("1")
                elif msk == "0":
                    r.append("0")
            r.reverse()
            collection[mem] = r

    return sum(int("".join(v), 2) for k, v in collection.items())

def part02(lines):
    collection = dict()
    current_mask = None
    for line in lines:
        if line.startswith("mask"):
            current_mask = list(line.split()[2])

        if line.startswith("mem"):
            s = p.search(line)
            mem, val = [int(i) for i in s.groups()]

            mem = list(str(bin(mem))[2:])
            mem = ["0"] * (36 - len(mem)) + mem

            val = list(str(bin(val))[2:])
            val = ["0"] * (36 - len(val)) + val

            r = []
            for msk, new in zip(reversed(current_mask), reversed(mem)):
                if msk == "X":
                    r.append("X")
                elif msk == "1":
                    r.append("1")
                elif msk == "0":
                    r.append(new)
            r.reverse()

            for perm in get_permutations(r):
                collection[int("".join(perm), 2)] = val

    return sum(int("".join(v), 2) for k, v in collection.items())

def get_permutations(items):
    perms = []

    num = items.count("X")
    com = combinations(["0", "1"] * num, num)
    for c in set(com):

        count = 0
        new = []
        for item in items:
            if item == "X":
                new.append(c[count])
                count += 1
            else:
                new.append(item)
        perms.append(new)

    return perms

def run():
    lines = iter(helpers.get_lines(r"./data/day_14.txt"))
    assert part01(lines) == 12135523360904
    lines = iter(helpers.get_lines(r"./data/day_14.txt"))
    assert part02(lines) == 2741969047858

if __name__ == "__main__":
    run()


from xmlrpc.server import SimpleXMLRPCServer
import sys
import requests

name = sys.argv[1]
port = sys.argv[2]

server = SimpleXMLRPCServer(("localhost", int(port)))
print(name + " starting...")

api_url = "https://en.wikipedia.org/w/api.php"

def find_links(page):
    print(name + " looking for links on page " + page)

    params = {
        "action": "parse",
        "page": page,
        "format": "json",
        "prop": "links",
    }
    try:
        json = requests.get(url=api_url, params=params).json()
        links = json["parse"]["links"]

        pages = []
        for link in links:
            pages.append(link["*"])
        return pages
    except Exception as ex:
        print("Error while trying to get links: " + str(ex))
        return False

print(sys.argv[1] + " started.")
server.register_function(find_links, "find_links")
try:
    server.serve_forever()
except Exception as ex:
    print("Serve_forever exception")


from fastapi import FastAPI, Depends
from svo.config.services import get_comparator
from svo.services.comparator import Comparator
from svo.services.models import ComparatorMaxRet, ComparatorData, GenericParams, MaxRetParams
from svo.config.repos import *
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"],
                   allow_headers=["*"])

@app.get("/companies")
def home():
    #a = company_get_repo().retrieve_companies_data().to_dict()
    return company_get_repo().retrieve_companies_data()

@app.post("/simple-comparator-results", response_model=ComparatorData)
def get_option_stats(params: GenericParams, comparator: Comparator=Depends(get_comparator)):
    res = comparator.naive_comparator_stats(market_cap=params.market_cap, n_cies=params.nb_comparables,
                                            back_date=params.back_test_length, maturity=params.maturity,
                                            tax_rate=params.tax_rate, nom_value=params.nom_value)

    return res

@app.post("/max-ret-comparator-results", response_model=ComparatorMaxRet)
def get_comparator_max_ret(params: MaxRetParams, comparator: Comparator=Depends(get_comparator)) -> ComparatorMaxRet:
    res = comparator.comparator_w_max_ret(market_cap=params.market_cap, n_cies=params.nb_comparables,
                                            back_date=params.back_test_length, maturity=params.maturity,
                                            tax_rate=params.tax_rate, max_ret=params.ret_threshold, nom_value=params.nom_value)
    return res

get_option_stats(GenericParams.parse_obj({
    "market_cap": 378933510144,
    "maturity": 10,
    "nb_comparables": 20,
    "back_test_length": 10,
    "tax_rate": 0,
    "nom_value": 10000
}), get_comparator())


from app.models import User, Quiz
from app import db

FINANCE_ANSWERS = {'Q1':'c', 'Q2':'b', 'Q3':'c'}
MARKETING_ANSWERS = {'Q1':'a', 'Q2':'d', 'Q3':'d'}
CHASSIS_ANSWERS = {'Q1':'b', 'Q2':'a', 'Q3':'a'}
VD_ANSWERS = {'Q1':'d', 'Q2':['a', 'b', 'c', 'd'], 'Q3':['d']}
POWERTRAIN_ANSWERS = {'Q1':['a', 'c'], 'Q2':'a', 'Q3':'a'}

FINANCE_MODULE_ANS = {'Q1':'d', 'Q2':'c'}
MARKETING_MODULE_ANS = {'Q1':'d', 'Q2':'b'}
CHASSIS_MODULE_ANS = {'Q1':'b', 'Q2':'a'}
VD_MODULE_ANS = {'Q1':'a', 'Q2':'d'}
POWERTRAIN_MODULE_ANS = {'Q1':'a', 'Q2':'d'}

#--------------FUNCTIONS FOR ACCESSING QUIZ DATA------------------

def get_section_totals(quiz):
    section_totals = {}
    section_totals['Finance'] = get_total(quiz.finance, FINANCE_ANSWERS)
    section_totals['Marketing'] = get_total(quiz.marketing, MARKETING_ANSWERS)
    section_totals['Chassis'] = get_total(quiz.chassis, CHASSIS_ANSWERS)
    section_totals['Vehicle Dynamics'] = get_total(quiz.vehicle_dynamics, VD_ANSWERS)
    section_totals['Powertrain'] = get_total(quiz.powertrain, POWERTRAIN_ANSWERS)
    return section_totals

def get_proportions(totals):
    num_marks = sum(totals.values()) or 1
    percents = {}
    for k in totals.keys():
        percents[k] = totals[k] / num_marks * 100
    return percents

def get_total(section, answers):
    total = 0
    for key, val in section.items():
        if type(val) == list and set(val) == set(answers[key]):
            total += 1
        elif type(val) == str and val == answers[key]:
            total += 1
    return total

def is_completed(quiz):
    finance = check_completion(quiz.finance)
    marketing = check_completion(quiz.marketing)
    chassis = check_completion(quiz.chassis)
    vehicle_dynamics = check_completion(quiz.vehicle_dynamics)
    powertrain = check_completion(quiz.powertrain)
    return finance and marketing and chassis and vehicle_dynamics and powertrain

def check_completion(section):
    for k, v in section.items():
        if v == '':
            return False
    return True


'''
    Library of helper tools for preprocessing and converting the text data into
    compatible tensors
'''

# import all the required stuff
from __future__ import print_function
import re
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
import numpy as np
import os # for the metadata file creator function

def prepare_input_data(table_file_path):
    '''
        function for obtaining the appropriate lists of words from the text file
        implemented by Anindya
        @param
        table_file_path => the path pointing to the location of the table file
        @return
        field_content_words, field_words, and content_words
    '''
    # obtain all the lines from the text file
    data_lines = open(table_file_path, 'r').readlines()

    # initialize all the lists to empty ones
    field_content_words = []; field_words = []; content_words = []

    # iterate over each line:
    for line in data_lines:
        temp = [] # create an empty array
        element_list = line.rstrip("\n").split("\t")
        for element in element_list:
            pair = element.split(":")

            '''
                small modification from Animesh:
                don't filter out the <none> pairs. Use all the data
            '''

            pair[0] = re.sub(r'[0-9]+', '', pair[0]).replace("_","")
            pair[1] = pair[1].split(" ")
            for word in pair[1]:
                temp_word = (pair[0] + " " + word)
                temp.append(temp_word)
                content_words.append(word)
                field_words.append(pair[0])
        field_content_words.append(temp)

    # remove the repeating field names from the list
    # field_words = list(set(field_words))
    # don't do this

    # return the so constructed lists
    return field_content_words, field_words, content_words

def prepare_input_labels(nb_file_path, sent_file_path):
    '''
        function for concatenating the label sentences in the file according to the numbers.
        implemented by Animesh
        @param
        nb_file_path => the path pointing to the location of the train.nb file
        sent_file_path => the path pointing to the location of the train.sent file
        @return
        aligned_labels_sequences
    '''
    # generate lists of numbers and raw label_sentences
    nums = open(nb_file_path, 'r').readlines(); sents = open(sent_file_path, 'r').readlines()

    # make the nums integers
    nums = map(int, nums);

    # make sure the number of sentences and the nums match
    assert sum(nums) == len(sents), "Length mismatch between the train.nb and train.sent files"

    # run a simple loop to concatenate sentences belonging to one single training example
    label_sents = [] # initialize to empty list
    for num in nums:
        count = 0; sent = '<start>' # initialize counter and sentence
        while count < num:
            sent += ' ' + sents.pop(0).strip(); count += 1
        # add the sentence to the label_sents list
        label_sents.append(sent + ' <eos>')

    # return the aligned_labels_sequences
    return label_sents

def prepare_tokenizer(words, max_word_length = None):
    '''
        funtion to generate vocabulary of the given list of words
        implemented by Anindya
        @param
        words => the list of words to be tokenized
    '''
    # flatten the words list:
    print("flattening the words into a single sequence ... ")
    flat_words = []; # initialize to empty list
    for i in range(len(words)):
        flat_words += words[i]
        if(i % 10000 == 0):
            print("joined", i, "examples")

    # obtain a tokenizer
    print("\nmaximum words to work with: ", max_word_length)
    t = Tokenizer(num_words = max_word_length, filters = '') # don't let keras ignore any words
    print("\nKeras's tokenizer kicks off ... ")
    t.fit_on_texts(flat_words)
    field_dict = dict(); rev_field_dict = dict()

    print("\nbuilding the dict and the rev_dict ... ")
    if(max_word_length is not None):
        vals = t.word_index.items()
        vals = sorted(vals, key=lambda x: x[1])
        for key,value in vals[:max_word_length - 1]:
            field_dict[value] = key
            rev_field_dict[key] = value
    else:
        for key,value in t.word_index.items():
            field_dict[value] = key
            rev_field_dict[key] = value

    ''' Small modification from Animesh
        # also add the '<unk>' token to the dictionary at 0th position
    '''
    field_dict[0] = '<unk>'; rev_field_dict['<unk>'] = 0

    print("\nencoding the words using the dictionary ... ")
    for i in range(len(words)):
        for j in range(len(words[i])):
            if(words[i][j] in rev_field_dict):
                words[i][j] = rev_field_dict[words[i][j]]
            else:
                words[i][j] = rev_field_dict['<unk>']

        if(i % 10000 == 0):
            print("encoded", i, "examples")

    vocab_size = len(field_dict)
    return words, field_dict, rev_field_dict, vocab_size

def group_tokenized_sequences(flat_seq, lengths):
    '''
        funtion to group the seqs together to original form after tokenization
        implemented by Animesh
        @param
        flat_seq => flat list of words (in order)
        lengths => list of lengths of each sequence in the dataset
        @return => grouped_seq
    '''

    # check if the lengths and the field_seq and the content_seq lengths are compatible
    assert sum(lengths) == len(flat_seq), "Lengths are not compatible"

    # perform the grouping:
    grouped_seqs = [] # initialize to empty list
    for length in lengths:
        count = 0; temp_grouped_seq = [] # initialize counter and storer list
        while(count < length):
            temp_grouped_seq.append(flat_seq.pop(0))
            count += 1

        # add the so contructed lists to the main groupings
        grouped_seqs.append(temp_grouped_seq)

    # finally return the so created lists
    return grouped_seqs

# create a function to generate a file for the given dictionary in the .vocab format
def create_dot_vocab(save_dict, path):
    '''
        Function for creating the Metadata file for the given vocab dict at the specified path

        @params:
        save_dict = the dictionary object used to create the metadata file
        path = the path where this dictionary is to be saved

        @return:
        None
    '''
    if(os.path.isfile(path)):
        print("The file already exists: ", path)

    else:
        with open(path, 'w') as metadatat_file:
            for (_, value) in save_dict.items():
                metadatat_file.write(value + "\n")
        # print a message that the file has been generated after completion
        print("The file has been created at: ", path)


"""
Unit tests for various types of solvers for cardiac cell models.
"""
__author__ = "Marie E. Rognes (meg@simula.no), 2013"
__all__ = ["TestBasicSingleCellSolver"]

import pytest
from testutils import medium, assert_almost_equal, parametrize

try:
    from ufl_legacy.log import info_red, info_green
except ImportError:
    from ufl.log import info_red, info_green
from cbcbeat import (
    supported_cell_models,
    BasicSingleCellSolver,
    NoCellModel,
    FitzHughNagumoManual,
    Tentusscher_2004_mcell,
    backend,
)
from dolfin import Expression

class TestBasicSingleCellSolver(object):
    "Test functionality for the basic single cell solver."

    references = {
        NoCellModel: {1.0: (0, 0.3), 0.5: (0, 0.2), 0.0: (0, 0.1)},
        FitzHughNagumoManual: {
            1.0: (0, -84.70013280019053),
            0.5: (0, -84.8000503072239979),
            0.0: (0, -84.9),
        },
        Tentusscher_2004_mcell: {
            1.0: (0, -85.89745525156506),
            0.5: (0, -85.99686000794499),
            0.0: (0, -86.09643254164848),
        },
    }

    def _run_solve(self, model, time, theta):
        "Run two time steps for the given model with the given theta solver."
        dt = 0.01
        T = 2 * dt
        interval = (0.0, T)

        # Initialize solver
        params = BasicSingleCellSolver.default_parameters()
        params["theta"] = theta

        params["enable_adjoint"] = False
        solver = BasicSingleCellSolver(model, time, params=params)

        # Assign initial conditions
        (vs_, vs) = solver.solution_fields()
        vs_.assign(model.initial_conditions())

        # Solve for a couple of steps
        solutions = solver.solve(interval, dt)
        for (t0, t1), vs in solutions:
            pass

        # Check that we are at the end time
        assert_almost_equal(t1, T, 1e-10)
        return vs.vector()

    @medium
    @parametrize(("theta"), [0.0, 0.5, 1.0])
    def test_default_basic_single_cell_solver(self, cell_model, theta):
        "Test basic single cell solver."
        time = backend.Constant(0.0)
        model = cell_model
        Model = cell_model.__class__

        if Model == supported_cell_models[3] and theta > 0:
            pytest.xfail("failing configuration (but should work)")

        model.stimulus = Expression("1000*t", t=time, degree=1)

        info_green("\nTesting %s" % model)
        vec_solve = self._run_solve(model, time, theta)

        if Model == supported_cell_models[3] and theta == 0:
            pytest.xfail("failing configuration (but should work)")

        if Model in self.references and theta in self.references[Model]:
            ind, ref_value = self.references[Model][theta]
            print("vec_solve", vec_solve.get_local())
            print("ind", ind, "ref", ref_value)

            assert_almost_equal(vec_solve[ind], ref_value, 1e-10)
        else:
            info_red("Missing references for %r, %r" % (Model, theta))


from unittest import TestCase

from flask import json
from mock import patch

from broker import broker_app
from broker.broker_app import app
from broker.hcaxlsbroker import SpreadsheetSubmission
from broker.ingestapi import IngestApi

class BrokerAppTest(TestCase):

    def setUp(self):
        self.client = app.test_client()

    def test_authorization_failed(self):
        # when:
        response = self.client.post('/api_upload')

        # then:
        self.assertEqual(401, response.status_code)

    def test_failed_save(self):
        with patch.object(broker_app, '_save_file') as save_file:
            # given:
            assert save_file is broker_app._save_file
            save_file.side_effect = Exception("I/O error")

            # when:
            response = self.client.post('/api_upload', headers={'Authorization': 'auth'})

            # then:
            self.assertEqual(500, response.status_code)

    def test_value_error_on_spreadsheet_submission(self):
        with patch.object(broker_app, '_save_file') as save_file, \
                patch.object(SpreadsheetSubmission, 'submit') as submit_spreadsheet:
            # given:
            save_file.return_value = 'path/to/file.xls'
            submit_spreadsheet.side_effect = ValueError("value error")

            # when:
            response = self.client.post('/api_upload', headers={'Authorization': 'auth'})

            # then:
            self.assertEqual(400, response.status_code)

    def test_key_error_on_spreadsheet_submission(self):
        with patch.object(broker_app, '_save_file') as save_file, \
                patch.object(SpreadsheetSubmission, 'submit') as submit_spreadsheet:
            # given:
            save_file.return_value = 'path/to/file.xls'
            submit_spreadsheet.side_effect = KeyError("key error")

            # when:
            response = self.client.post('/api_upload', headers={'Authorization': 'auth'})

            # then:
            self.assertEqual(400, response.status_code)

    def test_submission_successful(self):
        with patch.object(broker_app, '_save_file') as save_file, \
                patch.object(SpreadsheetSubmission, 'submit') as submit_spreadsheet, \
                patch.object(SpreadsheetSubmission, 'createSubmission') as create_submission, \
                patch.object(IngestApi, '__init__') as init_ingest_api, \
                patch.object(IngestApi, 'getObjectUuid') as get_object_uuid:
            # given:
            save_file.return_value = 'path/to/file.xls'

            # and:
            submission_id = '63fde9'
            submission_url = "https://sample.com/submission/%s" % (submission_id)
            create_submission.return_value = submission_url

            # given:
            init_ingest_api.return_value = None
            uuid = 'dc54c78c-8556-4750-b527-b35b26645e39'
            get_object_uuid.return_value = uuid

            # when:
            response = self.client.post('/api_upload', headers={'Authorization': 'auth'})

            # then:
            self.assertEqual(201, response.status_code)

            # and: assert correct JSON details
            self.assertTrue(response.data)
            response_json = json.loads(response.data)
            response_details = response_json['details']
            self.assertEqual(submission_url, response_details['submission_url'])
            self.assertEqual(uuid, response_details['submission_uuid'])
            self.assertEqual(uuid, response_details['display_uuid'])
            self.assertEqual(submission_id, response_details['submission_id'])


# coding=utf-8

""""HTML Converter"""

import io
import os
import tempfile
import unittest

from .html_pages import HtmlPagesConverter, FileAccessWrapper

class HtmlPagesTest(unittest.TestCase):

    def test_inserts_br_tags_for_linebreaks(self):
        filename = os.path.join(tempfile.gettempdir(), 'file.txt')
        f = open(filename, 'w', encoding='UTF-8')
        f.write('plain text\n')
        f.close()
        converter = HtmlPagesConverter(FileAccessWrapper(filename))
        new_text = converter.get_html_page(0)
        self.assertEqual(new_text, 'plain text
')

    def test_quotes_escaped(self):
        converter = HtmlPagesConverter(FakeFileWrapper("text with 'quotes'"))
        new_text = converter.get_html_page(0)
        self.assertEqual(new_text, 'text with 'quotes'
')

    def test_random_access_pages(self):
        text = 'Page one\nPAGE_BREAK\npage two\nPAGE_BREAK\n'
        converter = HtmlPagesConverter(FakeFileWrapper(text))
        page_two = converter.get_html_page(1)
        self.assertEqual('page two
', page_two)

    def test_nonexistent_file(self):
        converter = HtmlPagesConverter(FileAccessWrapper('missing'))
        self.assertEqual('', converter.get_html_page(0))

class FakeFileWrapper(object):

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

    def open(self):
        return io.StringIO(self.text)


# Import socket module
from socket import *
import sys # In order to terminate the program

serverName = '10.84.97.64'

# Assign a port number
serverPort = 12000

# Bind the socket to server address and server port
clientSocket = socket(AF_INET, SOCK_DGRAM)
sentence = 'This is a message from tobi ' #this is our data

#clientSocket.connect((serverName, serverPort))
clientSocket.sendto( sentence.encode(), (serverName, serverPort))

#receiving back from server
modifiedSentence, serverAddress = clientSocket.recvfrom(2048)

#print server address backd
print('From server: ', modifiedSentence.decode())
clientSocket.close()


import json
import os
import shutil
import uuid
import warnings
from collections import OrderedDict

import numpy as np
import pandas as pd
from mleap.bundle.serialize import MLeapSerializer, MLeapDeserializer, Vector
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, MinMaxScaler, Binarizer, PolynomialFeatures, OneHotEncoder, LabelEncoder
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils import column_or_1d
from sklearn.utils.validation import check_is_fitted

class ops(object):
    def __init__(self):
        self.STANDARD_SCALER = 'standard_scaler'
        self.MIN_MAX_SCALER = 'min_max_scaler'
        self.LABEL_ENCODER = 'string_indexer'
        self.ONE_HOT_ENCODER = 'one_hot_encoder'
        self.IMPUTER = 'imputer'
        self.NDARRAYTODATAFRAME = 'one_dim_array_to_dataframe'
        self.TODENSE = 'dense_transformer'
        self.BINARIZER = 'sklearn_binarizer'
        self.POLYNOMIALEXPANSION = 'sklearn_polynomial_expansion'

ops = ops()

def serialize_to_bundle(self, path, model_name):
    serializer = SimpleSerializer()
    return serializer.serialize_to_bundle(self, path, model_name)

def deserialize_from_bundle(self, path, node_name):
    serializer = SimpleSerializer()
    return serializer.deserialize_from_bundle(self, path, node_name)

def mleap_init(self, prior_tf, output_features=None):

    self.name = "{}_{}".format(self.op, uuid.uuid1())

    if prior_tf.op == 'vector_assembler':
        self.input_features = prior_tf.output_vector
        output_shape = 'tensor'
        output_size = len(prior_tf.input_features)
        output_feature_name = prior_tf.output_vector

    else:
        self.input_features = prior_tf.output_features
        output_shape = 'scalar'
        output_size = 1
        output_feature_name = prior_tf.output_features

    class_name = self.__class__.__name__

    if output_features is not None:
        self.output_features = output_features
    else:
        self.output_features = "{}_{}".format(output_feature_name, class_name.lower())

    if class_name == 'PolynomialFeatures':
        self.input_size = len(prior_tf.input_features)

    elif output_shape == 'tensor':
        self.input_shapes = {'data_shape': [{'shape':'tensor',
                                                 "tensor_shape": {"dimensions": [{"size": output_size}]}}]}
    else:
        self.input_shapes = {'data_shape': [{'shape': output_shape}]}

setattr(StandardScaler, 'op', ops.STANDARD_SCALER)
setattr(StandardScaler, 'serialize_to_bundle', serialize_to_bundle)
setattr(StandardScaler, 'deserialize_from_bundle', deserialize_from_bundle)
setattr(StandardScaler, 'mlinit', mleap_init)
setattr(StandardScaler, 'serializable', True)

setattr(MinMaxScaler, 'op', ops.MIN_MAX_SCALER)
setattr(MinMaxScaler, 'mlinit', mleap_init)
setattr(MinMaxScaler, 'serialize_to_bundle', serialize_to_bundle)
setattr(MinMaxScaler, 'deserialize_from_bundle', deserialize_from_bundle)
setattr(MinMaxScaler, 'serializable', True)

setattr(SimpleImputer, 'op', ops.IMPUTER)
setattr(SimpleImputer, 'mlinit', mleap_init)
setattr(SimpleImputer, 'serialize_to_bundle', serialize_to_bundle)
setattr(SimpleImputer, 'serializable', True)

setattr(OneHotEncoder, 'op', ops.ONE_HOT_ENCODER)
setattr(OneHotEncoder, 'mlinit', mleap_init)
setattr(OneHotEncoder, 'serialize_to_bundle', serialize_to_bundle)
setattr(OneHotEncoder, 'deserialize_from_bundle', deserialize_from_bundle)
setattr(OneHotEncoder, 'serializable', True)

setattr(Binarizer, 'op', ops.BINARIZER)
setattr(Binarizer, 'mlinit', mleap_init)
setattr(Binarizer, 'serialize_to_bundle', serialize_to_bundle)
setattr(Binarizer, 'deserialize_from_bundle', deserialize_from_bundle)
setattr(Binarizer, 'serializable', True)

setattr(PolynomialFeatures, 'op', ops.POLYNOMIALEXPANSION)
setattr(PolynomialFeatures, 'mlinit', mleap_init)
setattr(PolynomialFeatures, 'serialize_to_bundle', serialize_to_bundle)
setattr(PolynomialFeatures, 'deserialize_from_bundle', deserialize_from_bundle)
setattr(PolynomialFeatures, 'serializable', True)

class SimpleSerializer(object):
    """
    Extends base scikit-learn transformers to include:
        - set_input_features
        - set_output_features
        - serialize_to_bundle
    methods. The set input/output features are required to maintain parity with the Spark/MLeap execution engine.
    """
    def __init__(self):
        super(SimpleSerializer, self).__init__()

    @staticmethod
    def _choose_serializer(transformer):
        serializer = None
        if transformer.op == ops.STANDARD_SCALER:
            serializer = StandardScalerSerializer()
        elif transformer.op == ops.MIN_MAX_SCALER:
            serializer = MinMaxScalerSerializer()
        elif transformer.op == ops.ONE_HOT_ENCODER:
            serializer = OneHotEncoderSerializer()
        elif transformer.op == ops.IMPUTER:
            serializer = ImputerSerializer()
        elif transformer.op == ops.BINARIZER:
            serializer = BinarizerSerializer()
        elif transformer.op == ops.POLYNOMIALEXPANSION:
            serializer = PolynomialExpansionSerializer()
        return serializer

    @staticmethod
    def set_input_features(transformer, input_features):
        transformer.input_features = input_features

    @staticmethod
    def set_output_features(transformer, output_features):
        transformer.output_features = output_features

    def serialize_to_bundle(self, transformer, path, model_name):
        serializer = self._choose_serializer(transformer)
        serializer.serialize_to_bundle(transformer, path, model_name)

    def deserialize_from_bundle(self, transformer, path, node_name):
        deserializer = self._choose_serializer(transformer)
        transformer = deserializer.deserialize_from_bundle(transformer, path, node_name)
        return transformer

def _to_list(x):
    if isinstance(x, list):
        return x
    return list(x)

class FeatureExtractor(BaseEstimator, TransformerMixin, MLeapSerializer):
    def __init__(self, input_scalars=None, input_vectors=None, output_vector=None, output_vector_items=None):
        """
        Selects a subset of features from a pandas dataframe that are then passed into a subsequent transformer.
        MLeap treats this transformer like a VectorAssembler equivalent in spark.
        >>> data = pd.DataFrame([['a', 0, 1], ['b', 1, 2], ['c', 3, 4]], columns=['col_a', 'col_b', 'col_c'])
        >>> vector_items = ['col_b', 'col_c']
        >>> input_shapes = {'data_shape': [{'shape':'scalar'}, {'shape':'scalar'}]}
        >>> input_shapes = {'data_shape': [{'shape':'tensor', "tensor_shape": {"dimensions": [{"size": 23}]}}]}
        >>> feature_extractor2_tf = FeatureExtractor(vector_items, 'continuous_features', vector_items, input_shapes)
        >>> feature_extractor2_tf.fit_transform(data).head(1).values
        >>> array([[0, 1]])
        :param input_vectors: List of scalar feature names that are being extracted from a DataFrame
        :param input_vectors: List of FeatureExtractors that were used to generate the input vectors
        :param input_features: List of features to extracts from a pandas data frame
        :param output_vector: Name of the output vector, only used for serialization
        :param output_vector_items: List of output feature names
        :param input_shapes: the shape of each input feature, whether it is scalar or a vector
        if it's a vector, then we include size information of the vector
        :return:
        """
        self.input_scalars = input_scalars
        self.input_vectors = input_vectors
        self.input_features = self.get_input_features(input_scalars, input_vectors)
        self.output_vector_items = output_vector_items
        self.output_vector = output_vector
        self.op = 'vector_assembler'
        self.name = "{}_{}".format(self.op, uuid.uuid1())
        self.dtypes = None
        self.serializable = True
        self.skip_fit_transform = False
        self.input_size = None
        self.input_shapes = None

    def get_input_features(self, input_scalars, input_vectors):
        if input_scalars is not None:
            return input_scalars
        elif input_vectors is not None and isinstance(input_vectors, list) and len(input_vectors)>0:
            features = list()
            for x in input_vectors:
                if 'output_vector' in x.__dict__:
                    features.append(x.output_vector)
                elif 'output_features' in x.__dict__:
                    features.append(x.output_features)
            return features
        else:
            raise BaseException("Unable To Define Input Features")

    def transform(self, df, **params):
        if not self.skip_fit_transform:
            return df[self.input_features]
        return df

    def assign_input_shapes(self, X):
        # Figure out the data shape
        if self.input_scalars is not None and isinstance(X, pd.DataFrame):
            self.input_shapes = {'data_shape': [{'shape': 'scalar'}]*len(self.input_features)}

        elif self.input_vectors is not None:
            self.input_shapes = {'data_shape': []}
            for vector in self.input_vectors:
                if vector.op not in [ops.ONE_HOT_ENCODER, ops.POLYNOMIALEXPANSION]:
                    shape = {'shape': 'tensor', "tensor_shape": {"dimensions": [{"size": len(vector.input_features)}]}}
                    self.input_shapes['data_shape'].append(shape)
                elif vector.op == ops.ONE_HOT_ENCODER:
                    shape = {'shape': 'tensor', "tensor_shape": {"dimensions": [{"size": int(vector.n_values_[0] - 1)}]}}
                    self.input_shapes['data_shape'].append(shape)
        return self

    def fit(self, X, y=None, **fit_params):
        self.assign_input_shapes(X)
        if not self.skip_fit_transform:
            self.dtypes = X[self.input_features].dtypes.to_dict()
            if len([x for x in self.dtypes.values() if x.type == np.object_]) != 0:
                self.serializable = False
        return self

    def fit_transform(self, X, y=None, **fit_params):
        if not self.skip_fit_transform:
            self.fit(X)
        else:
            self.assign_input_shapes(X)

        df_subset = self.transform(X)
        return df_subset

    def serialize_to_bundle(self, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(("input_shapes", self.input_shapes))

        # define node inputs and outputs
        inputs = [{'name': x, 'port': 'input{}'.format(self.input_features.index(x))} for x in self.input_features]

        outputs = [{
                  "name": self.output_vector,
                  "port": "output"
                }]

        self.serialize(self, path, model_name, attributes, inputs, outputs)

class LabelEncoder(BaseEstimator, TransformerMixin, MLeapSerializer, MLeapDeserializer):
    """
    Copied from sklearn, but enables passing X and Y features, which allows this transformer
    to be used in Pipelines.

    Converts categorical values of a single column into categorical indices. This transformer should be followed by a
    NDArrayToDataFrame transformer to maintain a data structure required by scikit pipelines.

    NOTE: You can only LabelEncode/String Index one feature at a time!!!

    >>> data = pd.DataFrame([['a', 0], ['b', 1], ['b', 3], ['c', 1]], columns=['col_a', 'col_b'])
    >>> # Label Encoder for x1 Label
    >>> label_encoder_tf = LabelEncoder(input_features = ['col_a'] , output_features='col_a_label_le')
    >>> # Convert output of Label Encoder to Data Frame instead of 1d-array
    >>> n_dim_array_to_df_tf = NDArrayToDataFrame('col_a_label_le')
    >>> n_dim_array_to_df_tf.fit_transform(label_encoder_tf.fit_transform(data['col_a']))

    Encode labels with value between 0 and n_classes-1.

    Read more in the :ref:`User Guide <preprocessing_targets>`.

    Attributes
    ----------
    classes_ : array of shape (n_class,)
        Holds the label for each class.

    Examples
    --------
    `LabelEncoder` can be used to normalize labels.

    >>> from sklearn import preprocessing
    >>> le = preprocessing.LabelEncoder()
    >>> le.fit([1, 2, 2, 6])
    LabelEncoder()
    >>> le.classes_
    array([1, 2, 6])
    >>> le.transform([1, 1, 2, 6]) #doctest: +ELLIPSIS
    array([0, 0, 1, 2]...)
    >>> le.inverse_transform([0, 0, 1, 2])
    array([1, 1, 2, 6])

    It can also be used to transform non-numerical labels (as long as they are
    hashable and comparable) to numerical labels.

    >>> le = preprocessing.LabelEncoder()
    >>> le.fit(["paris", "paris", "tokyo", "amsterdam"])
    LabelEncoder()
    >>> list(le.classes_)
    ['amsterdam', 'paris', 'tokyo']
    >>> le.transform(["tokyo", "tokyo", "paris"]) #doctest: +ELLIPSIS
    array([2, 2, 1]...)
    >>> list(le.inverse_transform([2, 2, 1]))
    ['tokyo', 'tokyo', 'paris']

    """
    def __init__(self, input_features=None, output_features=None):
        self.input_features = input_features
        self.output_features = output_features
        self.op = 'string_indexer'
        self.name = "{}_{}".format(self.op, uuid.uuid1())
        self.serializable = True

    def fit(self, X):
        """Fit label encoder

        Parameters
        ----------
        y : array-like of shape (n_samples,)
            Target values.

        Returns
        -------
        self : returns an instance of self.
        """
        X = column_or_1d(X, warn=True)
        self.classes_ = np.unique(X)
        return self

    def fit_transform(self, X, y=None, **fit_params):
        """Fit label encoder and return encoded labels

        Parameters
        ----------
        y : array-like of shape [n_samples]
            Target values.

        Returns
        -------
        y : array-like of shape [n_samples]
        """
        y = column_or_1d(X, warn=True)
        self.classes_, X = np.unique(X, return_inverse=True)
        return X

    def transform(self, y):
        """Transform labels to normalized encoding.

        Parameters
        ----------
        y : array-like of shape [n_samples]
            Target values.

        Returns
        -------
        y : array-like of shape [n_samples]
        """
        check_is_fitted(self, 'classes_')
        y = column_or_1d(y, warn=True)

        classes = np.unique(y)
        if len(np.intersect1d(classes, self.classes_)) < len(classes):
            diff = np.setdiff1d(classes, self.classes_)
            raise ValueError("y contains new labels: %s" % str(diff))
        return np.searchsorted(self.classes_, y)

    def inverse_transform(self, y):
        """Transform labels back to original encoding.

        Parameters
        ----------
        y : numpy array of shape [n_samples]
            Target values.

        Returns
        -------
        y : numpy array of shape [n_samples]
        """
        check_is_fitted(self, 'classes_')

        diff = np.setdiff1d(y, np.arange(len(self.classes_)))
        if diff:
            raise ValueError("y contains new labels: %s" % str(diff))
        y = np.asarray(y)
        return self.classes_[y]

    def serialize_to_bundle(self, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('labels', self.classes_.tolist()))
        attributes.append(('nullable_input', False))

        # define node inputs and outputs
        inputs = [{
                  "name": self.input_features[0],
                  "port": "input"
                }]

        outputs = [{
                  "name": self.output_features,
                  "port": "output"
                }]

        self.serialize(self, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, node_path, node_name):

        attributes_map = {
            'labels': 'classes_'
        }

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(self, full_node_path, attributes_map)

        return transformer

class MinMaxScalerSerializer(MLeapSerializer, MLeapDeserializer):
    """
    Scales features by the range of values using calculated min and max from training data.

    >>> data = pd.DataFrame([[1], [5], [6], [1]], columns=['col_a'])
    >>> minmax_scaler_tf = MinMaxScaler()
    >>> minmax_scaler_tf.mlinit(input_features='col_a', output_features='scaled_cont_features')

    >>> minmax_scaler_tf.fit_transform(data)
    >>> array([[ 0.],
    >>>      [ 0.8],
    >>>      [ 1.],
    >>>      [ 0.]])
    """
    def __init__(self):
        super(MinMaxScalerSerializer, self).__init__()

    def serialize_to_bundle(self, transformer, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('min', _to_list(transformer.data_min_)))
        attributes.append(('max', _to_list(transformer.data_max_)))

        # define node inputs and outputs
        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, transformer, node_path, node_name):

        # Default (and only range supported)
        feature_range = (0, 1)

        attributes_map = {
            'min': 'data_min_',
            'max': 'data_max_'
        }

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(transformer, full_node_path, attributes_map)
        transformer.data_range_ = np.array(transformer.data_max_) - np.array(transformer.data_min_)

        transformer.scale_ = ((feature_range[1] - feature_range[0]) / transformer.data_range_)

        transformer.min_ = feature_range[0] - transformer.data_min_ * transformer.scale_

        return transformer

class ImputerSerializer(MLeapSerializer):
    def __init__(self):
        super(ImputerSerializer, self).__init__()
        self.serializable = False

    def serialize_to_bundle(self, transformer, path, model_name):
        if transformer.strategy == 'most_frequent' or transformer.strategy == 'constant':
            raise NotImplementedError(f"Scikit-learn's Imputer strategy `{transformer.strategy}` is not supported by MLeap")
        if transformer.add_indicator:
            raise NotImplementedError("Scikit-learn's Imputer parameter `add_indicator` is not supported by MLeap")
        if len(transformer.statistics_.tolist()) != 1:
            raise NotImplementedError("MLeap's Imputer only supports imputing a single feature at a time")

        attributes = list()
        attributes.append(('strategy', transformer.strategy))
        attributes.append(('surrogate_value', transformer.statistics_.tolist()[0]))
        if not np.isnan(transformer.missing_values):
            attributes.append(('missing_value', transformer.missing_values))

        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

class StandardScalerSerializer(MLeapSerializer, MLeapDeserializer):
    """
    Standardizes features by removing the mean and scaling to unit variance using mean and standard deviation from
    training data.
    >>> data = pd.DataFrame([[1], [5], [6], [1]], columns=['col_a'])
    >>> standard_scaler_tf = StandardScaler()
    >>> standard_scaler_tf.mlinit(input_features='col_a', output_features='scaled_cont_features')
    >>> standard_scaler_tf.fit_transform(data)
    >>> array([[-0.98787834],
    >>>         [ 0.76834982],
    >>>         [ 1.20740686],
    >>>         [-0.98787834]])
    """
    def __init__(self):
        super(StandardScalerSerializer, self).__init__()

    def serialize_to_bundle(self, transformer, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('mean', transformer.mean_.tolist()))
        attributes.append(('std', [np.sqrt(x) for x in transformer.var_]))

        # define node inputs and outputs
        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, transformer, node_path, node_name):

        attributes_map = {
            'mean': 'mean_',
            'std': 'scale_'
        }

        # Set serialized attributes
        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(transformer, full_node_path, attributes_map)

        # Set Additional Attributes
        if 'mean_' in transformer.__dict__:
            transformer.with_mean = True
        else:
            transformer.with_mean = False

        if 'scale_' in transformer.__dict__:
            transformer.with_std = True
            transformer.var = np.square(transformer.scale_)
        else:
            transformer.with_std = False

        return transformer

class OneHotEncoderSerializer(MLeapSerializer, MLeapDeserializer):
    """
    A one-hot encoder maps a single column of categorical indices to a
    column of binary vectors, which can be re-assamble back to a DataFrame using a ToDense transformer.
    """

    def __init__(self):
        super(OneHotEncoderSerializer, self).__init__()

    def serialize_to_bundle(self, transformer, path, model_name):
        if len(transformer.categories_) != 1:
            raise NotImplementedError("MLeap can only one-hot encode a single feature at a time")
        single_feature_categories = transformer.categories_[0]
        if not np.array_equal(single_feature_categories, np.arange(single_feature_categories.size)):
            raise ValueError(f"Categories {single_feature_categories} do not form a valid index range")
        if transformer.drop is not None:
            raise NotImplementedError("Scikit-learn's OneHotEncoder `drop` parameter is not supported by MLeap")
        if transformer.dtype != np.float64:
            raise NotImplementedError("Scikit-learn's OneHotEncoder `dtype` parameter is not supported by MLeap")

        attributes = list()
        attributes.append(('size', single_feature_categories.size))
        if transformer.handle_unknown == 'ignore':
            # MLeap's OneHotEncoderModel adds an extra column when keeping invalid data
            # so we drop that extra column to match sklearn's ignore behavior
            attributes.append(('handle_invalid', 'keep'))
            attributes.append(('drop_last', True))
        else:
            attributes.append(('handle_invalid', 'error'))
            attributes.append(('drop_last', False))

        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, transformer, node_path, node_name):

        attributes_map = {
            'size': 'categories_',
            'handle_invalid': 'handle_unknown',
        }

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(transformer, full_node_path, attributes_map)

        transformer.categories_ = np.asarray([range(transformer.categories_)])
        if transformer.handle_unknown == 'keep':
            transformer.handle_unknown = 'ignore'
        transformer.drop_idx_ = None

        transformer.sparse = False

        return transformer

class BinarizerSerializer(MLeapSerializer, MLeapDeserializer):
    def __init__(self):
        super(BinarizerSerializer, self).__init__()

    def serialize_to_bundle(self, transformer, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('threshold', float(transformer.threshold)))
        attributes.append(("input_shapes", transformer.input_shapes))

        # define node inputs and outputs
        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, transformer, node_path, node_name):

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(transformer, full_node_path)

        return transformer

class PolynomialExpansionSerializer(MLeapSerializer, MLeapDeserializer):
    def __init__(self):
        super(PolynomialExpansionSerializer, self).__init__()

    def serialize_to_bundle(self, transformer, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('combinations', str(transformer.get_feature_names()).replace("'", "").replace(", ", ",")))

        # define node inputs and outputs
        inputs = [{
                  "name": transformer.input_features,
                  "port": "input"
                }]

        outputs = [{
                  "name": transformer.output_features,
                  "port": "output"
                }]

        self.serialize(transformer, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, transformer, node_path, node_name):

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(transformer, full_node_path)
        transformer.include_bias = False

        # Get number of input features
        with open("{}/node.json".format(full_node_path)) as json_data:
            node_j = json.load(json_data)

        # Set powers
        transformer.n_input_features_ = len(node_j['shape']['inputs'][0]['name']) # TODO: Make this dynamic
        transformer.interaction_only = False

        transformer.n_output_features_ = sum(1 for _ in transformer._combinations(transformer.n_input_features_,
                                                                                transformer.degree,
                                                                                transformer.interaction_only,
                                                                                transformer.include_bias))
        return transformer

class ReshapeArrayToN1(BaseEstimator, TransformerMixin):
    def __init__(self):
        self.serializable=False
        self.op='reshape_array'
        self.name = "{}_{}".format(self.op, uuid.uuid1())

    def transform(self, X, **params):
        """
        :type X: np.ndarray
        :param X:
        :param params:
        :return:
        """
        return X.reshape(X.size, 1)

    def fit(self, df, y=None, **fit_params):
        return self

    def fit_transform(self, X, y=None, **fit_params):
        return self.transform(X)

class NDArrayToDataFrame(BaseEstimator, TransformerMixin):
    def __init__(self, input_features):
        self.input_features = input_features
        self.output_features = input_features
        self.op = 'one_dim_array_to_dataframe'
        self.name = "{}_{}".format(self.op, uuid.uuid1())

    def transform(self, X, **params):
        if isinstance(X, np.ndarray):
            return pd.DataFrame(X, columns=[self.input_features])
        return pd.DataFrame(X.todense(), columns=[self.input_features])

    def fit(self, df, y=None, **fit_params):
        return self

    def fit_transform(self, X, y=None, **fit_params):
        return self.transform(X)

    def get_mleap_model(self):
        js = {
          "op": self.op
        }
        return js

    def get_mleap_node(self):
        js = {
          "name": self.name,
          "shape": {
            "inputs": [{'name': x, 'port': 'input{}'.format(self.input_features.index(x))} for x in self.input_features],
            "outputs": [{
              "name": self.output_features,
              "port": "output"
            }]
          }
        }
        return js

    def _serialize_to_bundle(self, path, model_name):

        # If bundle path already exists, delte it and create a clean directory
        if os.path.exists("{}/{}".format(path, model_name)):
            shutil.rmtree("{}/{}".format(path, model_name))

        model_dir = "{}/{}".format(path, model_name)
        os.mkdir(model_dir)

        # Write bundle file
        with open("{}/{}".format(model_dir, 'model.json'), 'w') as outfile:
            json.dump(self.get_mleap_model(), outfile, indent=3)

        # Write node file
        with open("{}/{}".format(model_dir, 'node.json'), 'w') as outfile:
            json.dump(self.get_mleap_node(), outfile, indent=3)

class ToDense(BaseEstimator, TransformerMixin):
    def __init__(self, input_features):
        self.op = 'dense_transformer'
        self.name = "{}_{}".format(self.op, uuid.uuid1())
        self.input_features = input_features
        self.output_features = input_features

    def transform(self, X, **params):
        return X.todense()

    def fit(self, df, y=None, **fit_params):
        return self

    def fit_transform(self, X, y=None, **fit_params):
        return self.transform(X)

    def get_mleap_model(self):
        js = {
          "op": self.op
        }
        return js

    def get_mleap_node(self):
        js = {
          "name": self.name,
          "shape": {
            "inputs": {'name': self.input_features, 'port': 'input0'},
            "outputs": [{
              "name": self.output_features,
              "port": "output"
            }]
          }
        }
        return js

    def _serialize_to_bundle(self, path, model_name):
        # If bundle path already exists, delte it and create a clean directory
        if os.path.exists("{}/{}".format(path, model_name)):
            shutil.rmtree("{}/{}".format(path, model_name))

        model_dir = "{}/{}".format(path, model_name)
        os.mkdir(model_dir)

        # Write bundle file
        with open("{}/{}".format(model_dir, 'model.json'), 'w') as outfile:
            json.dump(self.get_mleap_model(), outfile, indent=3)

        # Write node file
        with open("{}/{}".format(model_dir, 'node.json'), 'w') as outfile:
            json.dump(self.get_mleap_node(), outfile, indent=3)

class MathUnary(BaseEstimator, TransformerMixin, MLeapSerializer, MLeapDeserializer):
    """
    Performs basic math operations on a single feature (column of a DataFrame). Supported operations include:
        - log
        - exp
        - sqrt
        - sin
        - cos
        - tan
    Note, currently we only support 1d-arrays.
    Inputs need to be floats.
    """
    def __init__(self, input_features=None, output_features=None, transform_type=None):
        self.valid_transforms = ['log', 'exp', 'sqrt', 'sin', 'cos', 'tan', 'abs']
        self.op = 'math_unary'
        self.name = "{}_{}".format(self.op, uuid.uuid4())
        self.input_features = input_features
        self.output_features = output_features
        self.transform_type = transform_type
        self.serializable = True

    def fit(self, X, y=None, **fit_params):
        """
        Fit Unary Math Operator
        :param X:
        :return:
        """
        X = column_or_1d(X, warn=True)
        if self.transform_type not in self.valid_transforms:
                warnings.warn("Invalid transform type.", stacklevel=2)
        return self

    def transform(self, y):
        """
        Transform features per specified math function.
        :param y:
        :return:
        """
        if self.transform_type == 'log':
            return np.log(y)
        elif self.transform_type == 'exp':
            return np.exp(y)
        elif self.transform_type == 'sqrt':
            return np.sqrt(y)
        elif self.transform_type == 'sin':
            return np.sin(y)
        elif self.transform_type == 'cos':
            return np.cos(y)
        elif self.transform_type == 'tan':
            return np.tan(y)
        elif self.transform_type == 'abs':
             return np.abs(y)

    def fit_transform(self, X, y=None, **fit_params):
        """
        Transform data per specified math function.
        :param X:
        :param y:
        :param fit_params:
        :return:
        """
        self.fit(X)

        return self.transform(X)

    def serialize_to_bundle(self, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('operation', self.transform_type))

        # define node inputs and outputs
        inputs = [{
                  "name": self.input_features[0],
                  "port": "input"
                }]

        outputs = [{
                  "name": self.output_features,
                  "port": "output"
                }]

        self.serialize(self, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, node_path, node_name):
        attributes_map = {
            'operation': 'transform_type'
        }
        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(self, full_node_path, attributes_map)
        return transformer

class MathBinary(BaseEstimator, TransformerMixin, MLeapSerializer, MLeapDeserializer):
    """
    Performs basic math operations on two features (columns of a DataFrame). Supported operations include:
        - add: Add x + y
        - sub: Subtract x - y
        - mul: Multiply x * y
        - div: Divide x / y
        - rem: Remainder x % y
        - logn: LogN log(x) / log(y)
        - pow: Power x^y
    These transforms work on 2-dimensional arrays/vectors, where the the first column is x and second column is y.
    Inputs need to be floats.
    """
    def __init__(self, input_features=None, output_features=None, transform_type=None):
        self.valid_transforms = ['add', 'sub', 'mul', 'div', 'rem', 'logn', 'pow']
        self.op = 'math_binary'
        self.name = "{}_{}".format(self.op, uuid.uuid4())
        self.input_features = input_features
        self.output_features = output_features
        self.transform_type = transform_type
        self.serializable = True

    def _return(self, y):
        if type(y, np.ndarray):
            return

    def fit(self, X, y=None, **fit_params):
        """
        Fit Unary Math Operator
        :param y:
        :return:
        """
        if self.transform_type not in self.valid_transforms:
                warnings.warn("Invalid transform type.", stacklevel=2)
        return self

    def transform(self, y):
        """
        Transform features per specified math function.
        :param y:
        :return:
        """
        if isinstance(y, pd.DataFrame):
            x = y.iloc[:,0]
            y = y.iloc[:,1]
        else:
            x = y[:,0]
            y = y[:,1]
        if self.transform_type == 'add':
            return pd.DataFrame(np.add(x, y), columns=[self.output_features])
        elif self.transform_type == 'sub':
            return pd.DataFrame(np.subtract(x, y), columns=[self.output_features])
        elif self.transform_type == 'mul':
            return pd.DataFrame(np.multiply(x, y), columns=[self.output_features])
        elif self.transform_type == 'div':
            return pd.DataFrame(np.divide(x, y), columns=[self.output_features])
        elif self.transform_type == 'rem':
            return pd.DataFrame(np.remainder(x, y), columns=[self.output_features])
        elif self.transform_type == 'pow':
            return pd.DataFrame(x**y, columns=[self.output_features])

    def fit_transform(self, X, y=None, **fit_params):
        """
        Transform data per specified math function.
        :param X:
        :param y:
        :param fit_params:
        :return:
        """
        self.fit(X)

        return self.transform(X)

    def serialize_to_bundle(self, path, model_name):

        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(('operation', self.transform_type))

        # define node inputs and outputs
        inputs = [{
                  "name": self.input_features[0],
                  "port": "input_a"
                },
                {
                    "name": self.input_features[1],
                    "port": "input_b"
                }]

        outputs = [{
                  "name": self.output_features,
                  "port": "output"
                }]

        self.serialize(self, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, node_path, node_name):
        attributes_map = {
            'operation': 'transform_type'
        }
        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(self, full_node_path, attributes_map)
        return transformer

class StringMap(BaseEstimator, TransformerMixin, MLeapSerializer, MLeapDeserializer):

    def __init__(self, input_features=None, output_features=None, labels=None):
        self.op = 'string_map'
        self.name = "{}_{}".format(self.op, uuid.uuid4())
        self.input_features = input_features
        self.output_features = output_features
        self.serializable = True
        self.labels = labels
        if labels is not None:
            if not isinstance(self.labels, OrderedDict):
                self.labels = OrderedDict(
                    sorted(self.labels.items(), key=lambda x: x[0])
                )
            self.label_keys = self.labels.keys
            self.label_values = self.labels.values

    def fit(self, X, y=None, **fit_params):
        if self.labels is None:
            self.labels = dict(zip(self.label_keys, self.label_values))
        return self

    def transform(self, y):
       return y.applymap(lambda input : self.labels[input]).values

    def fit_transform(self, X, y=None, **fit_params):
        self.fit(X)
        return self.transform(X)

    def serialize_to_bundle(self, path, model_name):
        # compile tuples of model attributes to serialize
        attributes = list()
        attributes.append(("labels", list(self.labels.keys())))
        attributes.append(("values", Vector(list(self.labels.values()))))

        # define node inputs and outputs
        inputs = [{
            "name": self.input_features[0],
            "port": "input"
        }]

        outputs = [{
            "name": self.output_features,
            "port": "output"
        }]

        self.serialize(self, path, model_name, attributes, inputs, outputs)

    def deserialize_from_bundle(self, node_path, node_name):
        attributes_map = {
            'labels': 'label_keys',
            'values': 'label_values'
        }

        full_node_path = os.path.join(node_path, node_name)
        transformer = self.deserialize_single_input_output(self, full_node_path, attributes_map)
        return transformer


import speech_recognition as sr
from pydub import AudioSegment
import os
import math
import sys

# directory where the audio files are
root = "C:/Users/awalker/Downloads/reduced"
if sys.argv.count > 1 :
   root = sys.argv[1]
if not os.path.exists(root) or not os.path.isdir(root):
   print('Could not open containing folder')
   exit()
files = os.listdir(root)

# open output file
audio_text = open('audio_text.txt', 'w+')

# write header
audio_text.write('Speech to text:\n\n')

file_index = 1

# create speech recognizer
recognizer = sr.Recognizer()

# go through each file in audio dir
for file in files :
   # write which file it is
   audio_text.write("File " + str(file_index) + ":\n")

   # get file name
   file_name = root + '/' + file

   audio = AudioSegment.from_file(file_name, format='wav')

   # length of the audio in milliseconds
   audio_length = len(audio)
   print(f'Audio Length: {audio_length}')

   split_size = 6000

   # get number of 8 second chunks
   chunk_count = len(audio) / split_size
   chunk_count = math.ceil(chunk_count)

   for i in range(0, chunk_count) :
      # get start and end, use that to get current chunk
      start = split_size*i
      end = split_size + split_size * i if i != chunk_count - 1 else len(audio)
      chunk = audio[start:end]

      # chunk name
      chunk_name = f'chunk_{i}'
      # storing the chunk to local storage
      chunk.export(chunk_name, format = 'wav')
      # printing the chunk
      print(f'{chunk_name} start: {start} end: {end}')

      # creating a listened audio
      with sr.AudioFile(chunk_name) as chunk_audio:
         chunk_listened = recognizer.listen(chunk_audio)

      # recognizing content from the audio
      try:
         # getting content from the chunk
         content = recognizer.recognize_vosk(chunk_listened)

         # writing to the file, removes the premade vosk format
         audio_text.write(content[14:-3] + ' ')
      # if not recognized
      except sr.UnknownValueError:
         print('Audio is not recognized')
      # request error if incorrectly setup
      except sr.RequestError as Error:
         print('Request Error')

   # write next line for next file
   audio_text.write('\n')
   file_index += 1
audio_text.close()


from airflow                   import configuration, DAG
from airflow.operators.python  import BranchPythonOperator
from airflow.operators.dummy   import DummyOperator
from airflow.models            import Variable
from airflow.utils.task_group  import TaskGroup
from utils.dag_notification    import *

from airflow.providers.ssh.hooks.ssh    import *

import datetime as dt
import logging

######### VARIABLES ###########

log       = logging.getLogger(__name__)
path      = configuration.get('core','dags_folder')
MAIN_PATH = path + "/../data"

MAIN_FOLDER = ["ODP", "B2S"]
SUB_FOLDER  = ["POS"]

###############################

def _list_file(ti, hookname, mainfolder, subfolder, tablename):
    SSH_HOOK  = SSHHook(
        ssh_conn_id = hookname,
        banner_timeout = 120,
        conn_timeout   = 120,
        keepalive_interval = 15
    )

    with SSH_HOOK.get_conn() as ssh_client:
        SFTP_HOOK = ssh_client.open_sftp()
        file_list = SFTP_HOOK.listdir(f"/{subfolder}/outbound/{tablename}/")
        SFTP_HOOK.close()

    if file_list == [ 'archive' ]:
        return f"no_alert_{mainfolder}_{subfolder}_{tablename}"
    else:
        file_list.remove('archive')
        ti.xcom_push(key='stuck_files', value=file_list)
        ti.xcom_push(key='sftp_path', value=f"/{subfolder}/outbound/{tablename}/")
        return f"send_alert_{mainfolder}_{subfolder}_{tablename}"

with DAG(
    dag_id="check_sftp_files",
    # schedule_interval=None,
    schedule_interval="59 23 * * *",
    start_date=dt.datetime(2022, 7, 7),
    catchup=True,
    max_active_runs=1,
    tags=['convz', 'production', 'mario', 'alert', 'sftp'],
    description='Alert for stuck files on SFTP server',
    render_template_as_native_obj=True,
    default_args={
        'on_failure_callback': ofm_task_fail_slack_alert,
        'retries': 0
    }
) as dag:

    start_task = DummyOperator(task_id = "start_task")
    end_task   = DummyOperator(task_id = "end_task", trigger_rule='none_failed')

    iterable_sources_list = Variable.get(
        key=f'sftp_alert',
        deserialize_json=True
    )
    # iterable_sources_list = {
    #     "ODP_JDA": ["BCH_JDA_DataPlatform_APADDR"],
    #     "B2S_JDA": ["BCH_JDA_DataPlatform_APADDR"],
    #     "ODP_POS": ["POS_DataPlatform_Txn_Translator"],
    #     "B2S_POS": ["POS_DataPlatform_Txn_Translator"]
    # }

    with TaskGroup(
        f'load_{SUB_FOLDER[0]}_tasks_group',
        prefix_group_id=False,
    ) as load_source_tasks_group:

        for source in MAIN_FOLDER:

            if source == "ODP": source = "OFM"
            HOOK_NAME   = f"sftp-{source.lower()}-connection-id"
            if source == "OFM": source = "ODP"

            for subsource in SUB_FOLDER:

                start_source = DummyOperator(task_id = f"start_{source}_{subsource}")

                with TaskGroup(
                    f'load_{source}_{subsource}_tasks_group',
                    prefix_group_id=False,
                ) as load_tables_tasks_group:

                    for table in iterable_sources_list.get(f"{source}_{subsource}"):

                        TABLE_ID = f'{table}'

                        list_file = BranchPythonOperator(
                            task_id=f'list_file_{source}_{subsource}_{table}',
                            python_callable=_list_file,
                            pool='sftp_connect_pool',
                            op_kwargs = {
                                'hookname'  : HOOK_NAME,
                                'mainfolder': source,
                                'subfolder' : subsource,
                                'tablename' : table
                            }
                        )

                        send_alert = DummyOperator(
                            task_id =f"send_alert_{source}_{subsource}_{table}",
                            on_success_callback = ofm_stuck_sftp_file_slack_alert
                        )

                        no_alert = DummyOperator(task_id =f"no_alert_{source}_{subsource}_{table}")

                        list_file >> [ send_alert, no_alert ]

                start_source >> load_tables_tasks_group

    start_task >> load_source_tasks_group >> end_task


"""Function for solving nonlinear least squares problems."""

from scipy.optimize import OptimizeResult

from desc.backend import jnp
from desc.utils import errorif, setdefault

from .bound_utils import (
    cl_scaling_vector,
    find_active_constraints,
    in_bounds,
    make_strictly_feasible,
    select_step,
)
from .tr_subproblems import (
    trust_region_step_exact_cho,
    trust_region_step_exact_svd,
    update_tr_radius,
)
from .utils import (
    STATUS_MESSAGES,
    check_termination,
    compute_jac_scale,
    print_header_nonlinear,
    print_iteration_nonlinear,
)

def lsqtr(  # noqa: C901 - FIXME: simplify this
    fun,
    x0,
    jac,
    bounds=(-jnp.inf, jnp.inf),
    args=(),
    x_scale="jac",
    ftol=1e-6,
    xtol=1e-6,
    gtol=1e-6,
    verbose=1,
    maxiter=None,
    callback=None,
    options=None,
):
    """Solve a least squares problem using a (quasi)-Newton trust region method.

    Parameters
    ----------
    fun : callable
        objective to be minimized. Should have a signature like fun(x,*args)-> 1d array
    x0 : array-like
        initial guess
    jac : callable:
        function to compute Jacobian matrix of fun
    bounds : tuple of array-like
        Lower and upper bounds on independent variables. Defaults to no bounds.
        Each array must match the size of x0 or be a scalar, in the latter case a
        bound will be the same for all variables. Use np.inf with an appropriate sign
        to disable bounds on all or some variables.
    args : tuple
        additional arguments passed to fun, grad, and jac
    x_scale : array_like or ``'jac'``, optional
        Characteristic scale of each variable. Setting ``x_scale`` is equivalent
        to reformulating the problem in scaled variables ``xs = x / x_scale``.
        An alternative view is that the size of a trust region along jth
        dimension is proportional to ``x_scale[j]``. Improved convergence may
        be achieved by setting ``x_scale`` such that a step of a given size
        along any of the scaled variables has a similar effect on the cost
        function. If set to ``'jac'``, the scale is iteratively updated using the
        inverse norms of the columns of the Jacobian matrix.
    ftol : float or None, optional
        Tolerance for termination by the change of the cost function.
        The optimization process is stopped when ``dF < ftol * F``,
        and there was an adequate agreement between a local quadratic model and
        the true model in the last step. If None, the termination by this
        condition is disabled.
    xtol : float or None, optional
        Tolerance for termination by the change of the independent variables.
        Optimization is stopped when ``norm(dx) < xtol * (xtol + norm(x))``.
        If None, the termination by this condition is disabled.
    gtol : float or None, optional
        Absolute tolerance for termination by the norm of the gradient.
        Optimizer terminates when ``max(abs(g)) < gtol``.
        If None, the termination by this condition is disabled.
    verbose : {0, 1, 2}, optional
        * 0 : work silently.
        * 1 (default) : display a termination report.
        * 2 : display progress during iterations
    maxiter : int, optional
        maximum number of iterations. Defaults to size(x)*100
    callback : callable, optional
        Called after each iteration. Should be a callable with
        the signature:

            ``callback(xk, *args) -> bool``

        where ``xk`` is the current parameter vector, and ``args``
        are the same arguments passed to fun and jac. If callback returns True
        the algorithm execution is terminated.
    options : dict, optional
        dictionary of optional keyword arguments to override default solver settings.

        - ``"max_nfev"`` : (int > 0) Maximum number of function evaluations (each
          iteration may take more than one function evaluation). Default is
          ``5*maxiter+1``
        - ``"max_dx"`` : (float > 0) Maximum allowed change in the norm of x from its
          starting point. Default np.inf.
        - ``"initial_trust_radius"`` : (``"scipy"``, ``"conngould"``, ``"mix"`` or
          float > 0) Initial trust region radius. ``"scipy"`` uses the scaled norm of
          x0, which is the default behavior in ``scipy.optimize.least_squares``.
          ``"conngould"`` uses the norm of the Cauchy point, as recommended in ch17
          of [1]_. ``"mix"`` uses the geometric mean of the previous two options. A
          float can also be passed to specify the trust radius directly.
          Default is ``"scipy"``.
        - ``"initial_trust_ratio"`` : (float > 0) A extra scaling factor that is
          applied after one of the previous heuristics to determine the initial trust
          radius. Default 1.
        - ``"max_trust_radius"`` : (float > 0) Maximum allowable trust region radius.
          Default ``np.inf``.
        - ``"min_trust_radius"`` : (float >= 0) Minimum allowable trust region radius.
          Optimization is terminated if the trust region falls below this value.
          Default ``np.finfo(x0.dtype).eps``.
        - ``"tr_increase_threshold"`` : (0 < float < 1) Increase the trust region
          radius when the ratio of actual to predicted reduction exceeds this threshold.
          Default 0.75.
        - ``"tr_decrease_threshold"`` : (0 < float < 1) Decrease the trust region
          radius when the ratio of actual to predicted reduction is less than this
          threshold. Default 0.25.
        - ``"tr_increase_ratio"`` : (float > 1) Factor to increase the trust region
          radius by when  the ratio of actual to predicted reduction exceeds threshold.
          Default 2.
        - ``"tr_decrease_ratio"`` : (0 < float < 1) Factor to decrease the trust region
          radius by when  the ratio of actual to predicted reduction falls below
          threshold. Default 0.25.
        - ``"tr_method"`` : ``"svd"``, ``"cho"``) Method to use for solving the trust
          region subproblem. ``"cho"`` uses a sequence of cholesky factorizations
          (generally 2-3), while ``"svd"`` uses one singular value decomposition.
          ``"cho"`` is generally faster for large systems, especially on GPU, but may
          be less accurate for badly scaled systems. Default ``"svd"``

    Returns
    -------
    res : OptimizeResult
        The optimization result represented as a ``OptimizeResult`` object.
        Important attributes are: ``x`` the solution array, ``success`` a
        Boolean flag indicating if the optimizer exited successfully and
        ``message`` which describes the cause of the termination. See
        ``OptimizeResult`` for a description of other attributes.

    References
    ----------
    .. [1] Conn, Andrew, and Gould, Nicholas, and Toint, Philippe. "Trust-region
           methods" (2000).

    """
    options = {} if options is None else options
    errorif(
        isinstance(x_scale, str) and x_scale not in ["jac", "auto"],
        ValueError,
        "x_scale should be one of 'jac', 'auto' or array-like, got {}".format(x_scale),
    )

    nfev = 0
    njev = 0
    iteration = 0

    n = x0.size
    x = x0.copy()
    lb, ub = jnp.broadcast_to(bounds[0], x.size), jnp.broadcast_to(bounds[1], x.size)
    bounded = jnp.any(lb != -jnp.inf) | jnp.any(ub != jnp.inf)
    assert in_bounds(x, lb, ub), "x0 is infeasible"
    x = make_strictly_feasible(x, lb, ub)

    f = fun(x, *args)
    nfev += 1
    cost = 0.5 * jnp.dot(f, f)
    J = jac(x, *args)
    njev += 1
    g = jnp.dot(J.T, f)

    maxiter = setdefault(maxiter, n * 100)
    max_nfev = options.pop("max_nfev", 5 * maxiter + 1)
    max_dx = options.pop("max_dx", jnp.inf)

    jac_scale = isinstance(x_scale, str) and x_scale in ["jac", "auto"]
    if jac_scale:
        scale, scale_inv = compute_jac_scale(J)
    else:
        x_scale = jnp.broadcast_to(x_scale, x.shape)
        scale, scale_inv = x_scale, 1 / x_scale

    v, dv = cl_scaling_vector(x, g, lb, ub)
    v = jnp.where(dv != 0, v * scale_inv, v)
    d = v**0.5 * scale
    diag_h = g * dv * scale

    g_h = g * d
    J_h = J * d
    g_norm = jnp.linalg.norm(g * v, ord=jnp.inf)

    # conngould : norm of the cauchy point, as recommended in ch17 of Conn & Gould
    # scipy : norm of the scaled x, as used in scipy
    # mix : geometric mean of conngould and scipy
    init_tr = {
        "scipy": jnp.linalg.norm(x * scale_inv / v**0.5),
        "conngould": jnp.sum(g_h**2) / jnp.sum((J_h @ g_h) ** 2),
        "mix": jnp.sqrt(
            jnp.sum(g_h**2)
            / jnp.sum((J_h @ g_h) ** 2)
            * jnp.linalg.norm(x * scale_inv / v**0.5)
        ),
    }
    trust_radius = options.pop("initial_trust_radius", "scipy")
    tr_ratio = options.pop("initial_trust_ratio", 1.0)
    trust_radius = init_tr.get(trust_radius, trust_radius)
    trust_radius *= tr_ratio
    trust_radius = trust_radius if (trust_radius > 0) else 1.0

    max_trust_radius = options.pop("max_trust_radius", jnp.inf)
    min_trust_radius = options.pop("min_trust_radius", jnp.finfo(x0.dtype).eps)
    tr_increase_threshold = options.pop("tr_increase_threshold", 0.75)
    tr_decrease_threshold = options.pop("tr_decrease_threshold", 0.25)
    tr_increase_ratio = options.pop("tr_increase_ratio", 2)
    tr_decrease_ratio = options.pop("tr_decrease_ratio", 0.25)
    tr_method = options.pop("tr_method", "svd")

    errorif(
        len(options) > 0,
        ValueError,
        "Unknown options: {}".format([key for key in options]),
    )
    errorif(
        tr_method not in ["cho", "svd"],
        ValueError,
        "tr_method should be one of 'cho', 'svd', got {}".format(tr_method),
    )

    callback = setdefault(callback, lambda *args: False)

    x_norm = jnp.linalg.norm(x, ord=2)
    success = None
    message = None
    step_norm = jnp.inf
    actual_reduction = jnp.inf
    reduction_ratio = 0

    if verbose > 1:
        print_header_nonlinear()
        print_iteration_nonlinear(
            iteration, nfev, cost, actual_reduction, step_norm, g_norm
        )

    allx = [x]
    alltr = [trust_radius]
    if g_norm < gtol:
        success, message = True, STATUS_MESSAGES["gtol"]

    alpha = 0  # "Levenberg-Marquardt" parameter

    while iteration < maxiter and success is None:

        # we don't want to factorize the extra stuff if we don't need to
        J_a = jnp.vstack([J_h, jnp.diag(diag_h**0.5)]) if bounded else J_h
        f_a = jnp.concatenate([f, jnp.zeros(diag_h.size)]) if bounded else f

        if tr_method == "svd":
            U, s, Vt = jnp.linalg.svd(J_a, full_matrices=False)
        elif tr_method == "cho":
            B_h = jnp.dot(J_a.T, J_a)

        actual_reduction = -1

        # theta controls step back step ratio from the bounds.
        theta = max(0.995, 1 - g_norm)

        while actual_reduction <= 0 and nfev <= max_nfev:
            # Solve the sub-problem.
            # This gives us the proposed step relative to the current position
            # and it tells us whether the proposed step
            # has reached the trust region boundary or not.
            if tr_method == "svd":
                step_h, hits_boundary, alpha = trust_region_step_exact_svd(
                    f_a, U, s, Vt.T, trust_radius, alpha
                )
            elif tr_method == "cho":
                step_h, hits_boundary, alpha = trust_region_step_exact_cho(
                    g_h, B_h, trust_radius, alpha
                )
            step = d * step_h  # Trust-region solution in the original space.

            step, step_h, predicted_reduction = select_step(
                x,
                J_h,
                diag_h,
                g_h,
                step,
                step_h,
                d,
                trust_radius,
                lb,
                ub,
                theta,
                mode="jac",
            )

            step_h_norm = jnp.linalg.norm(step_h, ord=2)
            step_norm = jnp.linalg.norm(step, ord=2)

            x_new = make_strictly_feasible(x + step, lb, ub, rstep=0)
            f_new = fun(x_new, *args)
            nfev += 1

            cost_new = 0.5 * jnp.dot(f_new, f_new)
            actual_reduction = cost - cost_new

            # update the trust radius according to the actual/predicted ratio
            tr_old = trust_radius
            trust_radius, reduction_ratio = update_tr_radius(
                trust_radius,
                actual_reduction,
                predicted_reduction,
                step_h_norm,
                hits_boundary,
                max_trust_radius,
                tr_increase_threshold,
                tr_increase_ratio,
                tr_decrease_threshold,
                tr_decrease_ratio,
            )
            alltr.append(trust_radius)
            alpha *= tr_old / trust_radius
            # TODO: does this need to move to the outer loop?
            success, message = check_termination(
                actual_reduction,
                cost,
                step_norm,
                x_norm,
                g_norm,
                reduction_ratio,
                ftol,
                xtol,
                gtol,
                iteration,
                maxiter,
                nfev,
                max_nfev,
                min_trust_radius=min_trust_radius,
                dx_total=jnp.linalg.norm(x - x0),
                max_dx=max_dx,
            )
            if success is not None:
                break

        # if reduction was enough, accept the step
        if actual_reduction > 0:
            x = x_new
            allx.append(x)
            f = f_new
            cost = cost_new
            J = jac(x, *args)
            njev += 1
            g = jnp.dot(J.T, f)

            if jac_scale:
                scale, scale_inv = compute_jac_scale(J, scale_inv)

            v, dv = cl_scaling_vector(x, g, lb, ub)
            v = jnp.where(dv != 0, v * scale_inv, v)
            d = v**0.5 * scale
            diag_h = g * dv * scale

            g_h = g * d
            J_h = J * d
            x_norm = jnp.linalg.norm(x, ord=2)
            g_norm = jnp.linalg.norm(g * v, ord=jnp.inf)

            if g_norm < gtol:
                success, message = True, STATUS_MESSAGES["gtol"]

            if callback(jnp.copy(x), *args):
                success, message = False, STATUS_MESSAGES["callback"]

        else:
            step_norm = actual_reduction = 0

        iteration += 1
        if verbose > 1:
            print_iteration_nonlinear(
                iteration, nfev, cost, actual_reduction, step_norm, g_norm
            )

    if g_norm < gtol:
        success, message = True, STATUS_MESSAGES["gtol"]
    if (iteration == maxiter) and success is None:
        success, message = False, STATUS_MESSAGES["maxiter"]
    active_mask = find_active_constraints(x, lb, ub, rtol=xtol)
    result = OptimizeResult(
        x=x,
        success=success,
        cost=cost,
        fun=f,
        grad=g,
        v=v,
        jac=J,
        optimality=g_norm,
        nfev=nfev,
        njev=njev,
        nit=iteration,
        message=message,
        active_mask=active_mask,
        allx=allx,
        alltr=alltr,
    )
    if verbose > 0:
        if result["success"]:
            print(result["message"])
        else:
            print("Warning: " + result["message"])
        print("         Current function value: {:.3e}".format(result["cost"]))
        print(
            "         Total delta_x: {:.3e}".format(jnp.linalg.norm(x0 - result["x"]))
        )
        print("         Iterations: {:d}".format(result["nit"]))
        print("         Function evaluations: {:d}".format(result["nfev"]))
        print("         Jacobian evaluations: {:d}".format(result["njev"]))

    return result


from .hkReferencedObject import hkReferencedObject
from enum import Enum
from .hkaSkeletonMapperData import hkaSkeletonMapperData

class ConstraintSource(Enum):
    NO_CONSTRAINTS = 0
    REFERENCE_POSE = 1
    CURRENT_POSE = 2

class hkaSkeletonMapper(hkReferencedObject):
    mapping: hkaSkeletonMapperData

    def __init__(self, infile):
        self.mapping = hkaSkeletonMapperData(infile)  # TYPE_STRUCT:TYPE_VOID

    def __repr__(self):
        return "<{class_name} mapping={mapping}>".format(**{
            "class_name": self.__class__.__name__,
            "mapping": self.mapping,
        })