from uuid import uuid4

import factory

from rest_framework import status
from rest_framework.reverse import reverse

from datahub.company.test.factories import AdviserFactory
from datahub.core.test_utils import (
    APITestMixin,
)

from datahub.task.test.factories import TaskFactory

class BaseTaskTests(APITestMixin):
    def adviser_api_client(self, adviser):
        """Create an api client where the adviser is the authenticated user"""
        return self.create_api_client(user=adviser)

class BaseListTaskTests(BaseTaskTests):
    reverse_url = None
    task_type_factory = None

    def test_get_all_tasks_returns_empty_when_no_tasks_exist(self):
        url = reverse(self.reverse_url)
        response = self.api_client.get(url).json()
        assert response == {
            'count': 0,
            'next': None,
            'previous': None,
            'results': [],
        }

    def test_get_all_tasks_returns_not_archived_tasks(self):
        TaskFactory(archived=True)
        not_archived_task = TaskFactory(archived=False)
        id = not_archived_task.id

        if self.task_type_factory:
            task_type = self.task_type_factory(task=not_archived_task)
            id = task_type.id

        url = f'{reverse(self.reverse_url)}?archived=false'

        response = self.api_client.get(url).json()
        assert response['count'] == 1
        assert response['results'][0]['id'] == str(id)

    def test_get_all_tasks_returns_archived_tasks(self):
        archived_task = TaskFactory(archived=True)
        TaskFactory(archived=False)
        id = archived_task.id

        if self.task_type_factory:
            task_type = self.task_type_factory(task=archived_task)
            id = task_type.id

        url = f'{reverse(self.reverse_url)}?archived=true'

        response = self.api_client.get(url).json()
        assert response['count'] == 1
        assert response['results'][0]['id'] == str(id)

    def test_get_all_tasks_returns_tasks_belonging_to_an_adviser(self):
        adviser = AdviserFactory()
        adviser_tasks = TaskFactory.create_batch(2, advisers=[adviser])
        TaskFactory.create_batch(3)
        task_ids = [str(x.id) for x in adviser_tasks]

        if self.task_type_factory:
            task_types = self.task_type_factory.create_batch(
                2,
                task=factory.Iterator(adviser_tasks),
            )
            task_ids = [str(x.id) for x in task_types]

        url = f'{reverse(self.reverse_url)}?advisers={adviser.id}'

        response = self.api_client.get(url).json()
        assert response['count'] == len(adviser_tasks)

        for result in response['results']:
            assert result['id'] in task_ids

    def test_get_all_tasks_returns_no_tasks_when_adviser_has_no_assigned_tasks(self):
        adviser = AdviserFactory()
        TaskFactory.create_batch(3)

        url = f'{reverse(self.reverse_url)}?advisers={adviser.id}'

        response = self.api_client.get(url).json()
        assert response == {
            'count': 0,
            'next': None,
            'previous': None,
            'results': [],
        }

class BaseEditTaskTests(BaseTaskTests):
    reverse_url = None
    task_type_factory = None
    nested_task_field_name = None

    def test_edit_task_with_unknown_advisor_id_returns_bad_request(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)

        data = {'advisers': [uuid4()]}

        response = self._call_task_endpoint_assert_response(
            adviser,
            task,
            data,
            status.HTTP_400_BAD_REQUEST,
        )

        json = response.json()

        if self.nested_task_field_name:
            json = json.get(self.nested_task_field_name)

        assert list(json.keys()) == ['advisers']

    def test_edit_task_returns_ok_when_user_not_creator_or_assigned_to_task(self):
        task = TaskFactory()
        adviser = AdviserFactory()

        data = {'advisers': [adviser.id]}
        self._call_task_endpoint_assert_response(adviser, task, data, status.HTTP_200_OK)

    def test_edit_task_returns_success_when_user_is_creator_but_not_assigned_to_task(self):
        adviser = AdviserFactory()
        task = TaskFactory(created_by=adviser)

        data = {'advisers': [adviser.id]}

        self._call_task_endpoint_assert_response(adviser, task, data, status.HTTP_200_OK)

    def test_edit_task_returns_success_when_user_is_not_creator_but_is_assigned_to_task(self):
        adviser = AdviserFactory()
        task = TaskFactory(advisers=[adviser])

        data = {'advisers': [adviser.id]}

        self._call_task_endpoint_assert_response(adviser, task, data, status.HTTP_200_OK)

    def test_edit_task_returns_success_when_user_is_creator_and_is_assigned_to_task(self):
        adviser = AdviserFactory()
        task = TaskFactory(advisers=[adviser], created_by=adviser)

        data = {'advisers': [adviser.id], 'title': 'New title'}

        self._call_task_endpoint_assert_response(adviser, task, data, status.HTTP_200_OK)

    def _call_task_endpoint_assert_response(self, adviser, task, data, status_code):
        """
        Call the task endpoint and check the response is expected
        """
        id = task.id

        if self.task_type_factory:
            task_type = self.task_type_factory(task=task, created_by=task.created_by)
            id = task_type.id
            data = {'task': data}

        url = reverse(self.reverse_url, kwargs={'pk': id})

        response = self.adviser_api_client(adviser).patch(
            url,
            data,
        )

        assert response.status_code == status_code
        return response


"""
如何实时测速？
    水缸：

    没有手动触发
        1，启动即执行
        2，循环执行

    'en0'
如何
    def fun1():
    print('hello')
    app.after(1000, fun1)

app = Tk()
app.after(1000, fun1)
app.mainloop() # 两个循环会导致卡死
"""

from tkinter import *
import psutil
import time

# while True:
#     s1 = psutil.net_io_counters(pernic=True)['en0']
#     time.sleep(1)
#     s2 = psutil.net_io_counters(pernic=True)['en0']
#     res = s2.bytes_recv - s1.bytes_recv
#     print(res/1024)

def make_app():
    app = Tk()
    app.geometry('300x150')
    app.config(bg='#303030')
    Label(text='Speed Monitor',
          font=('Hack', 25, 'bold'),
          bg='#303030',
          fg='white').pack()
    Label(
          name='lb2',
          text='-',
          font=('Hack', 20, 'bold'),
          bg='#303030',
          fg='white'
          ).pack()
    return app

def speed_test():
    s1 = psutil.net_io_counters(pernic=True)['en0']
    time.sleep(1)
    s2 = psutil.net_io_counters(pernic=True)['en0']
    res = s2.bytes_recv - s1.bytes_recv
    # print(res/1024)
    return str(res/1024) + 'kb/s'

def ui_update(do):
    data = do()
    lb2 = app.children['lb2']
    lb2.config(text=data)
    lb2.after(1000, lambda: ui_update(do))

app = make_app()
app.after(1000, lambda: ui_update(speed_test))
app.mainloop()


import os
import matplotlib.pyplot as plt
import math
from matplotlib.lines import Line2D
import seaborn as sns
"""
Program: hist.py
Author: Prabu Gugagantha
Date: Aug/5/2020
Details: Trying to seperate the historgram processes from the video ones.
         Aiming to only create graphs that describe the whole system in this program.
"""

"""
Details: The fileReader object initialization begins here.
@param fileName: The name of the file we want to read.
@param directory: The path directory for that file.

Fields:
name: The name of the file we will be reading.
dir: The directory of the file we are reading.
file: A file object that opens the file using self.name and self.dir.
generator: A iterator object that will be used to parse through a file.
#frame: A frame object that will store the file data and produce a resulting graph.
"""
class FileReader:

    def __init__(self,fileName, directory):
        self.name = fileName
        self.dir = directory
        self.file = open(self.dir + self.name, "r")
        self.generator = self.generator_read()
        self.title = self.fileTitle()

    """
    """
    def fileTitle(self):
        splitName = self.name.split("_")

        filetype = splitName[0]
        experimentNum = splitName[3]
        stress = splitName[7]
        if("." in stress):
            stress = stress[6:7] + "_" + stress[8:9]
        else:
            stress = stress[5:].split("c")[0]
        #filetype + "_" + "experiment" + experimentNum + "_" + "stress" + stress
        return "stress" + stress
    """
    Details: The method readLine reads a single line from a given file.
    Not used currently but is useful for testing the file reader.
    @return self.file.readline(): Uses our file object's readline method to return a line of text.
    """
    def readLine(self):
        return self.file.readline()

    """
    Details: The method readLines reads every line from a given file.
    Not used currently but is useful for testing the file reader.
    @return self.file.readlines(): Uses our file object's readlines method to return a list of rows.
    """
    def readLines(self):
        return self.file.readlines()

    """
    Details: The method generator_read creates an iterator that is stored in self.generator_read
    @yield row: Returns the next line of text in our file
    """
    def generator_read(self):
        for row in (self.file):
            yield row

    """
    Details: The method next simply calls next on our iterator object self.generator
    @return self.generator.__next__(): Returns the next line after calling __next__().
    """
    def next(self):
        try:
            return self.generator.__next__()
        except StopIteration:
            return None

    """
    Details: The method skip will iterate through the lines of a file a set number of times.
    @param num: The number of times you want to iterate.
    """
    def skip(self, num): # num = 17 for our files.
        for i in range(num):
            self.next()

    """
    Details: The method frameDetails will use the file name and extract the graph title.
    @return packingFrac + stress: The title of the frame which will be created.
    """
    def frameDetails(self):
        splitName = self.name.split("_")
        packingFrac = splitName[1][7:].split("B")[0]
        stress = splitName[7].split("cl")[0]
        return packingFrac + stress

"""
Method: parser
Details: Used to count how many contact/non-contact forces there are.
@param intReader: A file reader that will read different interaction files.
@return numBlue : The number of non-contact Forces.
@return numRed : The number of contact forces.
@return line : The current line we are on in the int file.
"""
def parser(intReader):
    numBlue = 0;
    numRed = 0;
    condition = False;
    intReader.skip(6);
    line = intReader.next()
    while(condition ==  False):
        line = line.split(" ")
        if(line[2] == "0"): #no-contact
            numBlue = numBlue + 1
        elif(line[2] != "0"): #contact
            numRed = numRed + 1
        line = intReader.next()
        if(not line == None):
            condition = '#' in line #error when hits final line
        else:
            condition = True
    return numBlue,numRed,line

"""
Method: color
Details: Provides the color of a specific graph using matplotlibs basic colors.
@param i : What color we want from our list.
"""
def color(i):
    #color = ['b','g','r','c','m','y','k'] #7
    color = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
    return color[int(i % len(color))]

#K = 1 + 3. 322 logx:round(1 + 3.322*(math.log(len(fractionList))))
#sqrt(x):round(math.sqrt(len(fractionList)))
"""
Method: histograph
Details: Takes the fractionList and plots the histogram with plt.bar.
@param fractionList : The resulting list fraction of contact forces in a specific experiment.
@param intReader: A file reader that will read different interaction files.
@param colorNum : The color of the histogram
@return fractionList : An updated fractionList with begining values trimmed off.
"""
def histograph(fractionList, intReader, colorNum):
    fractionList.sort()
    if("_" not in intReader.title):
        fractionList = fractionList[10:]
    binNum = 100 #round(1 + 3.322*(math.log(len(fractionList))))
    listMin = round(min(fractionList),7)
    listMax = round(max(fractionList),7)
    binWidth = (listMax-listMin)/binNum
    sumVal = sum(fractionList)

    #print(newfraclist)
    # you are here !!!!!!!!!!!!!!!!!!!!!!!!!!!!!11
    newfraclist = []#  a list showing the bin boundary a.k.a x coordinates for the line graph
    for i in range(binNum + 1):
        newfraclist.append(listMin + binWidth/2 + (binWidth*i))
    #print(newfraclist,listMax,listMin,len(newfraclist))
    #plt.ylim(0,400)

    binList = [0]* (len(newfraclist)) # a list holding the counts in each bin    a.k.a y coordinates for line graph
    for i in fractionList:
        newval = math.floor((i - listMin)/binWidth) #---> which b3in does a number fall in.
        binList[newval] = binList[newval] + 1
    totalCount = sum(binList)
    #print(binList)

    divisor = [] # a list of the new count values normalized
    for count in binList:
        divisor.append((count/totalCount)/binWidth)

    #print(newfraclist,listMax)
    #print(color(colorNum),colorNum)
    #print(len(newfraclist),len(divisor),len(fractionList))
    #print(plt.hist(fractionList, bins = binNum, label = intReader.title + " #: " + str(len(fractionList)))[1])
    #print(newfraclist)

    #.366788         .463898
    barBool = 3
    if(barBool == 1):
        plt.bar(newfraclist,divisor,binWidth,edgecolor = color(colorNum),linewidth = 3,color = 'white',label = intReader.title + " #: " + str(len(fractionList)))
        print(sum(divisor)*binWidth)
        plt.plot(newfraclist,divisor,color = color(colorNum + 1),linewidth = 3,label = intReader.title + " #: " + str(len(fractionList)))
    elif(barBool == 2):
        print(sum(divisor)*binWidth)
        newx = [newfraclist[0]-(binWidth/2),newfraclist[0]-(binWidth/2)]
        newy = [0,divisor[0]]
        for i in range(1,len(newfraclist)):
            newx.append(newfraclist[i]-(binWidth/2))
            newy.append(divisor[i-1])

            newx.append(newfraclist[i]-(binWidth/2))
            newy.append(divisor[i])
        newx.append(newfraclist[len(newfraclist) - 1]+(binWidth/2))
        newy.append(divisor[len(divisor) - 1])

        newx.append(newfraclist[len(newfraclist) - 1]+(binWidth/2))
        newy.append(0)

        #print(" ")
        #print(newx,newy)
        plt.plot(newx,newy,color = color(colorNum),linewidth = 3,label = intReader.title + " #: " + str(len(fractionList)))
        plt.plot(newfraclist,divisor,color = color(colorNum + 1),linewidth = 3,label = intReader.title + " #: " + str(len(fractionList)))
        #newfraclist,divisor,binWidth,edgecolor = color(colorNum),linewidth = 3,color = 'white',label = intReader.title + " #: " + str(len(fractionList))
    elif(barBool == 3):
        print(sum(divisor)*binWidth)
        plt.plot(newfraclist,divisor,color = color(colorNum),linewidth = 3,label = intReader.title + " #: " + str(len(fractionList)))
    return fractionList, intReader.title + " #: " + str(len(fractionList))

"""
Method: graph
Details: Edits the graph with labels and shows the resulting graph.
@param intReader: A file reader that will read different interaction files.
@param fractionListLength : The length of the fractionList.
@param saveTitle : The resulting file name that the graph will be saved as.

"""
def graph(intReader,fractionListLength, saveTitle):
    """
    legendElements = [Line2D([0],[0], marker = 'o',color = 'w', label = '#ofPoints: ' + str(len(fractionList)), markerfacecolor = 'black', markersize = 5),
                      Line2D([0],[0], marker = 'o',color = 'w', label = '#ofBins: ' + str(binNum), markerfacecolor = 'black', markersize = 5),
                      Line2D([0],[0], marker = 'o',color = 'w', label = 'BinWidth: ' + str(binWidth)[:7], markerfacecolor = 'black', markersize = 5)]
    plt.legend(handles = legendElements,loc = 'upper right')#bbox_to_anchor=(1, 1));
    """
    #print(fractionListLength)
    legendElements = [Line2D([0],[0], marker = 'o',color = 'w', label = '#ofPoints: ' + str(fractionListLength), markerfacecolor = 'black', markersize = 5)]
    plt.legend(loc = 'upper right')
    plt.xlabel("Fraction of Contact Forces")
    plt.ylabel("Probability Density")
    plt.title("Distribution: " + intReader.frameDetails().split("s")[0])

    temp = "histtest"
    #temp = "histSquareRoot"
    parent = "E:/Chakraborty/frames/"
    path = os.path.join(parent,temp)
    if not os.path.exists(path):
        os.makedirs(parent + temp)
    print(saveTitle.split(" ")[0])
    #plt.savefig("D:\\Chakraborty\\frames\\histtest\\" + saveTitle.split(" ")[0]) # D:\Chakraborty\frames\histtest
    plt.show()
    plt.close()

"""
Method: forceCount
Details: Runs through an int file and counts\calculates the percent of contact forces.
@param folderString : A string representation of the folder we want to access.
@param fileNum : The specific file id we are going to access data from.
@return fractionList : The resulting list fraction of contact forces in a specific experiment.
@return intReader: A file reader that will read different interaction files.
"""
def forceCount(folderString, fileNum):

    fractionList = []

    for i in range(1,11): #10 experiments

        directionInt = folderString + str(i) +"/"
        listDirInt = sorted(os.listdir(directionInt))# returns a list of files within this directory
        listDirIntLength = len(sorted(os.listdir(directionInt)))

        intReader = FileReader(listDirInt[fileNum], directionInt)
        intReader.skip(20)
        print(listDirInt[fileNum])

        line = ""
        while(line != None):
            b,r,line = parser(intReader)
            fractionList.append(round(r/(r+b),7))

    return fractionList, intReader

def main():
    folderString = "E:\\Chakraborty\\int\\"
    directionInt = folderString + "1/"
    listDirInt = sorted(os.listdir(directionInt))# returns a list of files within this directory
    listDirIntLength = len(listDirInt)
    print(listDirInt)
    key = [9,11,13,14,16,18,7]  #[0,1,2,3,4,5] # [9,11,13,14,16,18,7] # [8,10,12,15,17,6] #all [1:18]
    intReader = None
    fractionList = []
    for i in range(len(listDirInt)): # len(listDirInt)   ... listDirIntLength    #g1: 0,6   #g2:6,12?
        if(i in key):
            print("File Number: " + str(i))
            fractionList,intReader = forceCount(folderString,i)
            fractionList,saveTitle  = histograph(fractionList,intReader,i)

    saveTitle = "CrossBreed"
    graph(intReader,len(fractionList), saveTitle)
main()

"""
.1  -->0
.2  -->1
.3  -->2
.5  -->3
.7  -->4
.8  -->5

1  -->9
2  -->11
3  -->13
4  -->14
5  -->16
7  -->18
10  -->7

15  -->8
20  -->10
30  -->12
50  -->15
75  -->17
100  -->6
"""


import datetime
import json
import re
from copy import deepcopy
from PyQt5 import QtWidgets, uic, QtCore, QtGui
from QuantLib import Date, China, Period
import traceback

from ui import *
from utils import *
import trade

class Ui(QtWidgets.QMainWindow):
    def __init__(self):
        super(Ui, self).__init__()
        self.resize(1500, 1000)
        layout_1 = QtWidgets.QHBoxLayout()
        layout_2 = QtWidgets.QHBoxLayout()

        self.trader_ui = TraderUi()
        self.counterparty_ui = CounterpartyUi()
        self.bond_info_ui = BondInfoUi()
        self.transfer_ui = TransferUi()

        layout_1.addWidget(self.trader_ui)
        layout_1.addWidget(self.counterparty_ui)
        layout_2.addWidget(self.bond_info_ui)
        layout_2.addWidget(self.transfer_ui)

        mainlayout = QtWidgets.QVBoxLayout()
        mainlayout.addLayout(layout_1, 10)
        mainlayout.addLayout(layout_2, 60)
        widget = QtWidgets.QWidget()
        widget.setLayout(mainlayout)
        self.setCentralWidget(widget)

        self.bond_info_ui.get_position.clicked.connect(self.getPosition)
        self.bond_info_ui.send_order.clicked.connect(self.sendOrder)
        self.bond_info_ui.is_last_trade.clicked.connect(self.updateTplus1)
        self.bond_info_ui.send_settlement.clicked.connect(self.sendSettlement)
        self.bond_info_ui.save_log.clicked.connect(self.saveLog)
        self.bond_info_ui.parent = self

        self.transfer_ui.send_order.clicked.connect(self.sendTransferOrder)
        self.transfer_ui.is_last_trade.clicked.connect(self.updateTransfer)

        self.trader_ui.now_time_button.clicked.connect(self.changeNowtime)
        self.trader_ui.check_info.clicked.connect(self.checkTraderInfo)

        # create Portfolio instances for all accounts.
        self.portfolios = {}
        self.bond_info_ui.initial_portfolios = {}
        portflios = json.load(open('trader.json', encoding='utf-8'))
        for key_1 in portflios.keys():
            for key_2, item_2 in portflios[key_1].items():
                account = key_1 + key_2
                self.portfolios[key_1 + key_2] =\
                    portfolio_utils.create_portfolio(account)
                # Makes it easier to update json file
                self.portfolios[key_1 + key_2].key_1 = key_1
                self.portfolios[key_1 + key_2].key_2 = key_2

                self.bond_info_ui.initial_portfolios[key_1 + key_2] =\
                    portfolio_utils.create_portfolio(account)
                x = self.bond_info_ui.initial_portfolios[key_1 + key_2].bonds
                self.bond_info_ui.initial_portfolios[key_1 + key_2].init = True
                self.bond_info_ui.initial_portfolios[key_1 +
                                                     key_2].initial_value = x
                self.bond_info_ui.initial_portfolios[key_1 +
                                                     key_2].key_1 = key_1
                self.bond_info_ui.initial_portfolios[key_1 +
                                                     key_2].key_2 = key_2

    def getPosition(self):
        trader_position = json.load(open('trader.json', encoding='utf-8'))[
            self.trader_ui.list_type.currentText()]

        trader_id = self.trader_ui.account_list.currentText()
        code = self.bond_info_ui.code.text()

        bond_position = trader_position[trader_id]['position'].get(code, 0)
        cash_position = trader_position[trader_id]['cash']

        self.bond_info_ui.bond_position.setText(str(bond_position))
        self.bond_info_ui.cash_position.setText(str(cash_position))

    def changeNowtime(self):
        list_type = self.trader_ui.list_type.currentText()
        trader_id = self.trader_ui.account_list.currentText()
        key = list_type + trader_id

        date = self.trader_ui.now_time.text()
        self.portfolios[key].now_time = date
        QtWidgets.QMessageBox().about(self, '提示', '账户时间被设置为"{}"'.format(date))

    def checkTraderInfo(self):
        list_type = self.trader_ui.list_type.currentText()
        trader_id = self.trader_ui.account_list.currentText()
        key = list_type + trader_id
        portfolio = self.portfolios[key]

        popup = QtWidgets.QMainWindow(parent=self)
        mainlayout = QtWidgets.QVBoxLayout()
        widget = QtWidgets.QWidget()
        widget.setLayout(mainlayout)
        popup.setCentralWidget(widget)

        trader_info = '''
        账户名称:{}
        内置时间:{}
        全部资金:{}
        可用资金:{}
        '''.format(portfolio.account, portfolio.now_time, portfolio.cash, portfolio.free_cash)
        msg = QtWidgets.QTextBrowser()
        msg.setText(trader_info)
        mainlayout.addWidget(msg)

        df_ = displayDataFrame(portfolio.bonds, key, self)
        mainlayout.addWidget(df_)

        df_2 = displayDataFrame(portfolio.all_trade, key, self)
        mainlayout.addWidget(df_)

        msg_2 = QtWidgets.QMessageBox()
        msg_2.setStandardButtons(QtWidgets.QMessageBox.Ok)
        mainlayout.addWidget(msg_2)
        popup.resize(600, 600)
        popup.show()
        msg_2.exec_()
        popup.close()

    def _export_trader_info(self):
        excel_utils._export_trader_info(self)

    def sendOrder(self):
        self.bond_info_ui._export_info()
        self._export_trader_info()
        x = compliance.check_spot_order()
        if ('失败' in x) or ('不合规' in x) or ('不达标' in x):
            error_message = '报单失败，失败原因：\n'+'!\n'.join(x.split('！'))
            QtWidgets.QMessageBox().about(self, '错误信息', error_message)
            return False

        elif '预警' in x:
            error_message = '交易预警：\n'+'!\n'.join(x.split('！'))
            QtWidgets.QMessageBox().about(self, '预警', error_message)

        try:
            trade = trade_utils.create_spot_trade()
            print(trade.trade_time, type(trade.trade_time))
            print(trade.settlement_date, type(trade.settlement_date))

        except:
            QtWidgets.QMessageBox().about(self, '错误信息', traceback.format_exc())
            return False

        # 更新Portfolio对象
        try:
            account = trade.inside_id
            self.portfolios[account].append_waiting_trade(trade)
            self.portfolios[account].portfolio_update_t0(trade)
            print(self.portfolios[account].all_trade)
            print(self.portfolios[account].now_time)
            # print(self.portfolios[account].now_time, type(self.portfolios[account].now_time))
            # print(type(self.portfolios[account].all_trade.iloc[0, 1]))

            if trade.is_inside_trade:
                account = trade.other_inside_id
                self.portfolios[account].append_waiting_trade(
                    trade.reversed())
                self.portfolios[account].portfolio_update_t0(
                    trade.reversed())

        except:
            QtWidgets.QMessageBox().about(
                self, '错误信息', traceback.format_exc())
            return False

        # 写入Excel和json文件
        if not trade.is_inside_trade:
            self.portfolios[trade.inside_id].to_excel()
            self.portfolios[trade.inside_id].to_json()
        else:
            self.portfolios[trade.inside_id].to_excel()
            self.portfolios[trade.other_inside_id].to_excel()

            self.portfolios[trade.inside_id].to_json()
            portfolio_utils.to_json(self.portfolios[trade.other_inside_id])

        QtWidgets.QMessageBox().about(self, '', '报单完成')

    def sendTransferOrder(self):
        # self.transfer_ui._export_transfer_info()
        x = compliance.check_transfer_order()
        if x != '交易成功':
            error_message = '转托管报单失败，失败原因：\n'+'!\n'.join(x.split('！'))
            QtWidgets.QMessageBox().about(self, '错误信息', error_message)
            return False

        try:
            trade = trade_utils.create_transfer_trade()
            trade_ = deepcopy(trade)
            print(trade.settlement_date)
            trade.settlement_date = trade.trade_time
            print(trade_.settlement_date)
        except:
            QtWidgets.QMessageBox().about(self, '错误信息', traceback.format_exc())
            return False

        # 更新Portfolio对象
        try:
            trade.settlement_days == 'T+1'
            account = trade.inside_id
            self.portfolios[account].append_waiting_trade(trade)
            self.portfolios[account].portfolio_update_t0(trade)

            account = trade.other_inside_id
            self.portfolios[account].append_waiting_trade(
                trade_.reversed())
            self.portfolios[account].portfolio_update_t0(trade_.reversed())

        except:
            QtWidgets.QMessageBox().about(
                self, '错误信息', traceback.format_exc())
            return False

        # 写入Excel和json文件
        if not trade.is_inside_trade:
            portfolio_utils.to_excel(self.portfolios[trade.inside_id])
            portfolio_utils.to_json(self.portfolios[trade.inside_id])
        else:
            portfolio_utils.to_excel(self.portfolios[trade.inside_id])
            portfolio_utils.to_excel(self.portfolios[trade.other_inside_id])

            portfolio_utils.to_json(self.portfolios[trade.inside_id])
            portfolio_utils.to_json(self.portfolios[trade.other_inside_id])

        QtWidgets.QMessageBox().about(self, '', '报单完成')

    def sendSettlement(self):
        # 当前账户提交清算申请
        list_type = self.trader_ui.list_type.currentText()
        trader_id = self.trader_ui.account_list.currentText()
        key = list_type + trader_id

        self.portfolios[key].settle()
        self.portfolios[key].to_json()
        self.portfolios[key].to_excel()

        self.hintFailures()

    def checkSufficiency(self):
        prompt_msg = []
        popup = QtWidgets.QMainWindow(parent=self)
        mainlayout = QtWidgets.QVBoxLayout()
        widget = QtWidgets.QWidget()
        widget.setLayout(mainlayout)
        popup.setCentralWidget(widget)

        for key in self.portfolios:
            temp_portfolio = deepcopy(self.portfolios[key])
            temp_portfolio.settle()
            bonds_not_enough = temp_portfolio.bonds[temp_portfolio.bonds['par_amount'] < 0]
            if not bonds_not_enough.empty:
                x = bonds_not_enough['bond_code'].to_list()
                msg = ['{}-{}现券持仓不足'.format(key, i) for i in x]
                prompt_msg.extend(msg)
            # display current positions
            df_ = displayDataFrame(temp_portfolio.bonds, key, self)
            mainlayout.addWidget(df_)

        msg = QtWidgets.QMessageBox()
        text = "若不进行转托管操作，账户持仓如下，请选择是否进行接下去的操作"
        msg.setText(text)
        msg.setWindowTitle("请选择")
        msg.setStandardButtons(
            QtWidgets.QMessageBox.No | QtWidgets.QMessageBox.Yes)
        mainlayout.addWidget(msg)
        popup.resize(600, 600)
        popup.show()
        retval = msg.exec_()
        popup.close()
        return prompt_msg, retval

    def updateTplus1(self):
        # 更新明天到账的T+1交易，调用后now_time会+1
        # 首先检查是否存在需要弥补交易
        prompt_msg, retval = self.checkSufficiency()
        if retval == QtWidgets.QMessageBox.No:
            # 中止本函数运行
            return

        if prompt_msg:
            msg = QtWidgets.QMessageBox()
            text = "更新持仓中止，以下现券不足，请选择是否进行自动内部转托管\n" + '\n'.join(prompt_msg)
            msg.setText(text)
            msg.setWindowTitle("持仓不足警示")
            msg.setStandardButtons(
                QtWidgets.QMessageBox.No | QtWidgets.QMessageBox.Yes)
            retval = msg.exec_()

            # 若选择Yes，则进入此分支，进行自动转托管
            if retval == QtWidgets.QMessageBox.Yes:
                self.autoTransfer()
                QtWidgets.QMessageBox().about(self, '提示信息', '自动转托管完成，请再次点击“完成今日交易”按钮')
                return
            # 若不选择，则继续更新T+1持仓

        # 更新T+1
        cal = China(China.IB)
        for key in self.portfolios:
            date = self.portfolios[key].now_time
            if isinstance(date, datetime.datetime):
                date = date.date().isoformat().replace('-', '/')
            x = date.split('/')
            ql_date = Date(int(x[2]), int(x[1]), int(x[0]))
            next_trading_day = cal.advance(
                ql_date, Period('1D')).ISO().split('-')
            next_trading_day = ql_date.ISO().replace('-', '/')
            self.portfolios[key].now_time = next_trading_day
            self.portfolios[key].portfolio_update_t1()
            self.portfolios[key].now_time = date

            self.portfolios[key].to_excel()
            self.portfolios[key].to_json()
            self.portfolios[key].save_position()
            self.portfolios[key].log()

        QtWidgets.QMessageBox().about(self, '', '更新完成')

    def updateTransfer(self):
        # 更新明天到账的转托管交易，调用后now_time会+1
        cal = China(China.IB)
        for key in self.portfolios:
            date = self.portfolios[key].now_time
            if isinstance(date, datetime.datetime):
                date = date.date().isoformat().replace('-', '/')
            x = date.split('/')
            ql_date = Date(int(x[2]), int(x[1]), int(x[0]))
            next_trading_day = cal.advance(
                ql_date, Period('1D')).ISO().replace('-', '/')
            self.portfolios[key].now_time = next_trading_day
            other_portfolios = [self.portfolios[i]
                                for i in self.portfolios if i != key]
            self.portfolios[key].portfolio_update_transfer(direction='in',
                                                           other_portfolios=other_portfolios)
            self.portfolios[key].now_time = date

            self.portfolios[key].to_excel()
            self.portfolios[key].to_json()
            self.portfolios[key].save_position()
            self.portfolios[key].log()

        QtWidgets.QMessageBox().about(self, '', '更新完成')

    def autoTransfer(self):
        # 判断各内部账户中是否存在互补，并进行提示
        # 转托管仍需要手动输入
        temp_portfolios = {}
        for key in self.portfolios:
            temp_portfolios[key] = deepcopy(self.portfolios[key])
            temp_portfolios[key].settle()
            bonds_not_enough = temp_portfolios[key].bonds[temp_portfolios.bonds['par_amount'] < 0]
            temp_portfolios[key].unfilled = bonds_not_enough

    def hintFailures(self):
        popup = QtWidgets.QMainWindow(parent=self)
        mainlayout = QtWidgets.QVBoxLayout()
        widget = QtWidgets.QWidget()
        widget.setLayout(mainlayout)
        popup.setCentralWidget(widget)

        msg = QtWidgets.QMessageBox()
        msg.setText("结算报单完成，该账户以下Trade属于Failed Trade")
        mainlayout.addWidget(msg)

        list_type = self.trader_ui.list_type.currentText()
        trader_id = self.trader_ui.account_list.currentText()
        key = list_type + trader_id
        df = self.portfolios[key].failed_trade
        df = displayDataFrame(df, key, self)

        mainlayout.addWidget(df)
        popup.resize(600, 600)
        popup.show()
        msg.exec_()
        popup.close()

    def saveLog(self):
        for x in self.portfolios.values():
            x.log()

class PandasModel(QtCore.QAbstractTableModel):

    def __init__(self, data):
        QtCore.QAbstractTableModel.__init__(self)
        self._data = data

    def rowCount(self, parent=None):
        return self._data.shape[0]

    def columnCount(self, parnet=None):
        return self._data.shape[1]

    def data(self, index, role=QtCore.Qt.DisplayRole):
        if index.isValid():
            if role == QtCore.Qt.DisplayRole:
                return str(self._data.iloc[index.row(), index.column()])
        return None

    def headerData(self, col, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return self._data.columns[col]
        return None

def displayDataFrame(df, title='', parent=None):
    popup = QtWidgets.QMainWindow(parent=parent)
    popup.setWindowTitle(title)
    popup.view = QtWidgets.QTableView()
    model = PandasModel(df)
    popup.view.setModel(model)
    popup.view.resize(800, 400)
    popup.setCentralWidget(popup.view)
    return popup


"""
  Name     : c12_31_scatter_random_numbers.py
  Book     : Python for Finance (2nd ed.)
  Publisher: Packt Publishing Ltd.
  Author   : Yuxing Yan
  Date     : 6/6/2017
  email    : yany@canisius.edu
             paulyxy@hotmail.com
"""

import numpy as np
import matplotlib.pyplot as plt
np.random.seed(12345)
n=200
a = np.random.uniform(size=(n*2))
plt.scatter(a[:n], a[n:])
plt.show()

"""
import cython
#%load_ext cythonmagic
def bessel_J0(double x):
    return gsl_sf_bessel_J0(x)

%%cython -l gsl
# Include the right header file and declare the function
cdef extern from "gsl/gsl_sf_bessel.h":
    double gsl_sf_bessel_J0(double x)

# small python wrapper to give it a nicer name
def bessel_J0(double x):
    return gsl_sf_bessel_J0(x)

"""


import logging

from django.db import transaction
from rest_framework import serializers

from datahub.email_ingestion.email_processor import EmailProcessor
from datahub.interaction.email_processors.exceptions import (
    BadCalendarInviteError,
    InvalidInviteError,
    NoContactsError,
)
from datahub.interaction.email_processors.notify import (
    notify_meeting_ingest_failure,
    notify_meeting_ingest_success,
)
from datahub.interaction.email_processors.parsers import (
    CalendarInteractionEmailParser,
    InteractionEmailParser,
)
from datahub.interaction.email_processors.utils import (
    get_all_recipients,
    get_best_match_adviser_by_email,
)
from datahub.interaction.models import Interaction
from datahub.interaction.serializers import InteractionSerializer

logger = logging.getLogger(__name__)

# Invalid email error codes that we can notify users about - we should not
# notify users if the email is malformed or if the sender was unverified
# as these situations *could* point to malicious activity and we should offer
# no hints in these situations.
EXCEPTION_NOTIFY_MESSAGES = {
    BadCalendarInviteError: (
        'It looks like you sent something to Data Hub that was not a calendar invitation. '
        'Please send to Data Hub a calendar invitation (not an email or another type of '
        'message). If you still see this error message after sending a calendar email, '
        'please contact the Data Hub support team.'
    ),
    NoContactsError: (
        'The calendar invitation you sent to Data Hub did not include an email address '
        'that is recognised as a contact on Data Hub. The invitation must also be sent '
        'to Data Hub at the same time (and not forwarded to Data Hub later). Please '
        'check both of these things when resending your invitation and if you still '
        'see this message after resending the invitation, contact the Data Hub support '
        'team.'
    ),
}

def _flatten_serializer_errors_to_list(serializer_errors):
    """
    Flatten DRF Serializer validation errors to a list with one field per item.
    """
    field_errors = []
    for field_name, details in serializer_errors.items():
        details_string = ','.join([str(detail) for detail in details])
        field_errors.append(f'{field_name}: {details_string}')
    return field_errors

def _filter_contacts_to_single_company(contacts, company):
    """
    Given a list of contacts and a company, return all of the contacts who are
    attributed to that company.
    """
    return [contact for contact in contacts if contact.company == company]

def _get_meeting_subject(sender, contacts, secondary_advisers):
    """
    Construct and return a meeting subject given a sender, contacts and secondary
    advisers (if present).
    """
    adviser_names = [
        adviser.name
        for adviser in (sender, *secondary_advisers)
        if adviser.name
    ]
    if not adviser_names:
        adviser_names = ['DIT']
    contact_names = [contact.name for contact in contacts if contact.name]
    if not contact_names:
        try:
            contact_names = [contacts[0].company.name]
        except IndexError:
            # This is not going to be a valid interaction, but this helper function
            # should not raise an exception
            contact_names = ['a company']
    all_names = (*adviser_names, *contact_names)
    comma_names = ', '.join(all_names[:-1])
    return f'Meeting between {comma_names} and {all_names[-1]}'

class BaseInteractionEmailProcessor:

    def _notify_meeting_ingest_failure(self, message, errors):
        try:
            sender_email = message.from_[0][1]
        except IndexError:
            logger.info(
                'Cannot extract email of the sender from message. '
                'Failure notification will not be sent.',
            )
            return
        sender_adviser = get_best_match_adviser_by_email(sender_email)
        if not sender_adviser:
            logger.info(
                'Cannot find adviser matching the sender email in message. '
                'Failure notification will not be sent.',
            )
            return
        recipient_emails = get_all_recipients(message)
        notify_meeting_ingest_failure(sender_adviser, errors, recipient_emails)

    def validate_with_serializer(self, data):
        """
        Transforms extracted data into a dict suitable for use with InteractionSerializer
        and then runs data through this serializer for validation.

        Returns the instantiated serializer.
        """
        transformed_data = self._to_serializer_format(data)
        serializer = InteractionSerializer(
            context={'check_association_permissions': False},
            data=transformed_data,
        )
        serializer.is_valid(raise_exception=True)
        return serializer

    def _to_serializer_format(self, data):
        dit_participants = [
            {
                'adviser': {'id': adviser.id},
                'team': adviser.dit_team,
            }
            for adviser in (data['sender'], *data['secondary_advisers'])
        ]

        data_for_serializer = {
            'contacts': [{'id': contact.id} for contact in data['contacts']],
            'company': {'id': data['top_company'].id},
            'date': data['date'],
            'dit_participants': dit_participants,
            'kind': Interaction.Kind.INTERACTION,
            'status': Interaction.Status.DRAFT,
            'subject': data['subject'],
            'was_policy_feedback_provided': False,
        }
        if 'body' in data:
            data_for_serializer['notes'] = data['body']

        return data_for_serializer

class CalendarInteractionEmailProcessor(BaseInteractionEmailProcessor, EmailProcessor):
    """
    An EmailProcessor which checks whether incoming email is a valid DBT/company
    meeting, parses meeting information and creates a draft Interaction model
    instance for it if the information is valid.
    """

    def _handle_invalid_invite(self, exception, message):
        """
        Given an InvalidInviteError and an email message, log the error that
        the user triggered and notify them with an explanation (if applicable).
        """
        error_str = repr(exception)
        error_with_email_info = (
            f'Ingested email with ID "{message.message_id}" (received '
            f'{message.received[0]["date_utc"]}) was not valid: {error_str}'
        )
        logger.warning(error_with_email_info)
        exc_class = exception.__class__
        # Only notify users if the error code is one that we can notify users about
        if exc_class in EXCEPTION_NOTIFY_MESSAGES:
            readable_error = EXCEPTION_NOTIFY_MESSAGES[exc_class]
            self._notify_meeting_ingest_failure(message, [readable_error])

    @transaction.atomic
    def save_serializer_as_interaction(self, serializer, interaction_data):
        """
        Create the interaction model instance from the validated serializer.
        """
        # Provide an overridden value for source - so that we save the meeting
        # data properly
        interaction = serializer.save(
            source={
                'meeting': {'id': interaction_data['meeting_details']['uid']},
            },
        )
        return interaction

    def process_email(self, message):
        """
        Review the metadata and calendar attachment (if present) of an email
        message to see if it fits the our criteria of a valid Data Hub meeting
        request.  If it does, create a draft Interaction for it.

        :param message: mailparser.MailParser object - the message to process
        """
        # Parse the email for interaction data
        email_parser = CalendarInteractionEmailParser(message)
        try:
            interaction_data = email_parser.extract_interaction_data_from_email()
        except InvalidInviteError as exc:
            self._handle_invalid_invite(exc, message)
            return (False, repr(exc), None)

        # Make the same-company check easy to remove later if we allow Interactions
        # to have contacts from more than one company
        sanitised_contacts = _filter_contacts_to_single_company(
            interaction_data['contacts'],
            interaction_data['top_company'],
        )
        interaction_data['contacts'] = sanitised_contacts

        # Replace the meeting invite subject with one which details the people attending
        interaction_data['subject'] = _get_meeting_subject(
            interaction_data['sender'],
            interaction_data['contacts'],
            interaction_data['secondary_advisers'],
        )

        # Get a serializer for the interaction data
        try:
            serializer = self.validate_with_serializer(interaction_data)
        except serializers.ValidationError as exc:
            errors = _flatten_serializer_errors_to_list(exc.detail)
            self._notify_meeting_ingest_failure(message, errors)
            return (False, ', '.join(errors), None)

        # For our initial iteration of this feature, we are ignoring meeting updates
        matching_interactions = Interaction.objects.filter(
            source__contains={'meeting': {'id': interaction_data['meeting_details']['uid']}},
        )
        if matching_interactions.exists():
            return (False, 'Meeting already exists as an interaction', None)

        interaction = self.save_serializer_as_interaction(serializer, interaction_data)
        notify_meeting_ingest_success(
            interaction_data['sender'],
            interaction,
            get_all_recipients(message),
        )
        return (True, f'Successfully created interaction #{interaction.id}', interaction.id)

class InteractionPlainEmailProcessor(BaseInteractionEmailProcessor, EmailProcessor):
    """
    An EmailProcessor which checks whether incoming email is a valid DBT/company
    interaction, parses meeting information and creates a draft Interaction model
    instance for it if the information is valid.
    """

    @transaction.atomic
    def save_serializer_as_interaction(self, serializer, interaction_data):
        """
        Create the interaction model instance from the validated serializer.
        """
        # Provide an overridden value for source - so that we save the meeting
        # data properly
        interaction = serializer.save(
            source={
                'email': {'id': interaction_data['id']},
            },
        )
        return interaction

    def process_email(self, message):
        """
        Review the metadata of an email message to see if it fits the our criteria
        of a valid Data Hub interaction. If it does, create a draft Interaction for it.

        :param message: mailparser.MailParser object - the message to process
        """
        # Parse the email for interaction data
        email_parser = InteractionEmailParser(message)
        interaction_data = email_parser.extract_interaction_data_from_email()

        # Make the same-company check easy to remove later if we allow Interactions
        # to have contacts from more than one company
        sanitised_contacts = _filter_contacts_to_single_company(
            interaction_data['contacts'],
            interaction_data['top_company'],
        )
        interaction_data['contacts'] = sanitised_contacts

        try:
            serializer = self.validate_with_serializer(interaction_data)
        except serializers.ValidationError as exc:
            errors = _flatten_serializer_errors_to_list(exc.detail)
            self._notify_meeting_ingest_failure(message, errors)
            return (False, ', '.join(errors), None)

        matching_interactions = Interaction.objects.filter(
            source__contains={'email': {'id': interaction_data['id']}},
        )
        if matching_interactions.exists():
            return (False, 'Email already exists as an interaction', None)

        interaction = self.save_serializer_as_interaction(serializer, interaction_data)
        notify_meeting_ingest_success(
            interaction_data['sender'],
            interaction,
            get_all_recipients(message),
        )

        return (True, f'Successfully created interaction #{interaction.id}', interaction.id)


# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/data_module.py
import time

from pytorch_lightning import LightningDataModule
from soundstorm.s1.AR.data.bucket_sampler import DistributedBucketSampler
from soundstorm.s1.AR.data.dataset_librilight_6k import Text2SemanticDataset
from torch.utils.data import DataLoader

class Text2SemanticDataModule(LightningDataModule):
    def __init__(self,
                 config,
                 train_semantic_dirs,
                 train_phoneme_dirs,
                 dev_semantic_dirs,
                 dev_phoneme_dirs,
                 train_non_speech_dirs=None,
                 dev_non_speech_dirs=None):
        super().__init__()
        self.config = config
        self.train_semantic_dirs = train_semantic_dirs
        self.train_phoneme_dirs = train_phoneme_dirs
        self.dev_semantic_dirs = dev_semantic_dirs
        self.dev_phoneme_dirs = dev_phoneme_dirs
        self.train_non_speech_dirs = train_non_speech_dirs
        self.dev_non_speech_dirs = dev_non_speech_dirs
        self.num_workers = self.config['data']['num_workers']
        print("self.num_workers:", self.num_workers)
        self.persistent_workers = True if self.num_workers > 0 else False
        self.prefetch_factor = 2

    def prepare_data(self):
        pass

    def setup(self, stage=None, output_logs=False):
        start_build_time = time.time()
        self._train_dataset = Text2SemanticDataset(
            phoneme_dirs=self.train_phoneme_dirs,
            semantic_dirs=self.train_semantic_dirs,
            non_speech_dirs=self.train_non_speech_dirs,
            max_sec=self.config['data']['max_sec'],
            pad_val=self.config['data']['pad_val'],
            min_ps_ratio=self.config['data'].get('min_ps_ratio', 6),
            max_ps_ratio=self.config['data'].get('max_ps_ratio', 22), )
        self._dev_dataset = Text2SemanticDataset(
            phoneme_dirs=self.dev_phoneme_dirs,
            semantic_dirs=self.dev_semantic_dirs,
            non_speech_dirs=self.dev_non_speech_dirs,
            max_sample=self.config['data']['max_eval_sample'],
            max_sec=self.config['data']['max_sec'],
            pad_val=self.config['data']['pad_val'],
            min_ps_ratio=self.config['data'].get('min_ps_ratio', 6),
            max_ps_ratio=self.config['data'].get('max_ps_ratio', 22), )
        print(
            f"time of build dataloader: {round(time.time() - start_build_time, 2)}s"
        )

    def train_dataloader(self):
        batch_size = self.config['train']['batch_size']
        sampler = DistributedBucketSampler(
            self._train_dataset, batch_size=batch_size)
        return DataLoader(
            self._train_dataset,
            batch_size=batch_size,
            sampler=sampler,
            collate_fn=self._train_dataset.collate,
            num_workers=self.num_workers,
            pin_memory=False,
            persistent_workers=self.persistent_workers,
            prefetch_factor=self.prefetch_factor, )

    def val_dataloader(self):
        return DataLoader(
            self._dev_dataset,
            batch_size=1,
            shuffle=False,
            collate_fn=self._train_dataset.collate,
            num_workers=self.num_workers,
            pin_memory=False,
            persistent_workers=self.persistent_workers,
            prefetch_factor=self.prefetch_factor, )

    # 这个会使用到嘛？
    def test_dataloader(self):
        return DataLoader(
            self._dev_dataset,
            batch_size=1,
            shuffle=False,
            collate_fn=self._train_dataset.collate)


import os
#from this import d

import numpy as np
import pandas as pd
from scipy.optimize import minimize, basinhopping

import matplotlib.pyplot as pl
from matplotlib import rcParams

from .. import logging as logg
from .. import settings
from ..core import get_n_jobs, parallelize, SplicingDynamics
from ..preprocessing.moments import get_connectivities
from .dynamical_model_utils import BaseDynamics, convolve, linreg, get_norm_std, use_model
from .utils import make_unique_list, test_bimodality

MAX_LOSS = 1e6

class DynamicsRecovery(BaseDynamics):
    def __init__(self, adata, gene, load_pars=None, **kwargs):
        super().__init__(adata, gene, **kwargs)
        if load_pars and "fit_alpha" in adata.var.keys():
            try:
                self.load_pars(adata, gene)
            except Exception:
                self.recoverable = False
                logg.warn(f"Model for {self.gene} could not be instantiated in load_pars.")

        elif self.recoverable:
            self.initialize()

    def initialize_WX(self, method):
        dim = self.n_TFs
        #tmp = self.gamma - np.sqrt(self.gamma*self.gamma+8*self.beta)
        #max_WX = (self.gamma-tmp/2)*self.alpha*np.exp(-tmp*tmp/(16*self.beta)) + self.gamma*self.rho
        #max_x = self.x_f.sum(1).max()
        #norm_term = max_WX/max_x
        if method == 'rand':
            weight = np.random.rand(dim)*2/dim  #* (2*norm_term)
        elif method == 'ones':
            #weight = np.ones(dim) * norm_term
            #weight = np.ones(dim)/dim
            #weight = (self.y_f.mean()/self.x_f.mean(0))/dim
            weight = np.ones(dim)/self.x_f.sum(1).std()
        elif method == 'knockTF_Log2FC':
            weight = self.knockTF_Log2FC / np.dot(self.x_f, self.knockTF_Log2FC.reshape(-1,1)).std()
        elif method == 'correlation':
            weight = self.TFs_correlation / np.dot(self.x_f, self.TFs_correlation.reshape(-1,1)).std()
        elif method == 'PCA':
            from sklearn.decomposition import PCA
            pca = PCA(n_components=1)
            pca.fit_transform(self.x_f)
            weight = pca.components_[0]
        elif method == 'gamma':
            from scipy.optimize import lsq_linear
            tmp1 = self.max_yf_x.copy()
            tmp2 = self.gamma * self.max_yf
            lb = -self.WX_thres
            hb = self.WX_thres
            weight = lsq_linear(tmp1, tmp2, bounds=(lb, hb), method='bvls', lsmr_tol='auto', verbose=0)
            weight = weight.x
        weight = weight.reshape(-1,1)
        WX = np.ravel(np.dot(self.x, weight))
        return WX, weight

    def init_par(self, alpha=None, beta=None, omega=None, theta=None, gamma=None):
        self.alpha = float(self.max_yf.mean()-self.min_yf.mean())/2 if alpha is None else alpha
        self.beta = float(self.max_yf.mean()+self.min_yf.mean())/2 if beta is None else beta
        self.omega = 2*np.pi if omega is None else omega
        self.theta = 0 if theta is None else theta
        self.gamma = 1 if gamma is None else gamma

        if not self.init_weight_method is None:
            self.WX, self.weights_init = self.initialize_WX(method=self.init_weight_method)
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)

        # update object with initialized vars
        self.likelihood, self.varx = 0, 0

        # initialize time point assignment
        if not self.init_weight_method is None:
            self.WX_model, self.delta = self.update_WX(self.alpha, self.beta, self.omega, self.theta,
                self.gamma, method=self.WX_method, thres=self.WX_thres)
        self.WX = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta)
        max_yf_WX = use_model(self.WX_model, self.max_yf_x, method=self.WX_method, bias=self.delta)
        self.gamma = float(np.mean(max_yf_WX/self.max_yf))
        #if self.gamma < self.gamma_thres:
        #    self.gamma = self.gamma_thres
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)
        self.pars = np.array([self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta,
            self.likelihood, self.varx, self.scaling_y] + list(self.scaling))[:, None]

        self.loss = [self.get_loss(refit_time=False, norm_std=True, reg=self.reg)]
        return

    def initialize(self):
        TFs_x, target_y, f = self.TFs_x, self.target_y, self.filtered
        TFs_x_f = TFs_x[f]
        target_y_f = target_y[f]

        # initialize scaling
        #for x_i in self.TFs_expression
        self.std_TFs_x, self.std_target_y = np.std(TFs_x_f, 0), np.std(target_y_f)
        #self.std_TFs_x[self.std_TFs_x==0] = 1e6 #TFs_x_f[:,self.std_TFs_x==0].mean(0)
        if self.std_target_y == 0:
            self.std_TFs_x = self.std_target_y = 1
        _scaling = self.fit_scaling
        if isinstance(_scaling, bool):
            if _scaling:
                scaling = 1/self.std_TFs_x
                scaling = np.nan_to_num(scaling, nan=0, posinf=0, neginf=0)
                scaling_y = 1/self.std_target_y
            else:
                scaling = np.ones_like(self.std_TFs_x)
                scaling_y = 1
        else:
            scaling = _scaling

        x, x_f = TFs_x * scaling, TFs_x_f * scaling
        y, y_f = target_y * scaling_y, target_y_f * scaling_y
        self.x, self.x_f, self.y, self.y_f = x, x_f, y, y_f
        self.scaling, self.scaling_y = scaling, scaling_y

        self.max_yf_idx = np.ravel(y_f >= np.percentile(y_f, 99, axis=0))
        self.max_yf_x = x_f[self.max_yf_idx]
        self.max_yf = np.ravel(y_f[self.max_yf_idx])
        self.min_yf_idx = np.ravel(y_f <= np.percentile(y_f, 1, axis=0))
        self.min_yf = np.ravel(y_f[self.min_yf_idx])

        self.init_par()
        return

    def fit(self):
        if self.fix_theta:
            self.search_and_fit_no_theta(reg=self.reg)
        else:
            self.search_and_fit(reg=self.reg)
        self.WX = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta)
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)
        self.varx = self.get_variance(to_filter=True, refit_time=False, norm_std=True)
        self.likelihood = self.get_likelihood(to_filter=True, varx=self.varx, refit_time=False, norm_std=True)
        self.weights = np.ravel(self.WX_model)
        self.weights_final = self.weights * self.x_f.sum(0)
        self.weights_final = self.weights_final / abs(self.weights_final).max()
        #if :
        #    self.WX_model = - self.WX_model
        #    self.fit_paras_(reg=self.reg)
        self.get_raw_velocity()

    def save_par(self):
        self.best_loss = self.loss[-1]
        self.alpha_best, self.beta_best, self.omega_best, self.theta_best, self.gamma_best, self.delta_best = \
            self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta
        self.loss_best = self.loss
        self.t_best = self.t
        self.WX_model_best = self.WX_model
        self.WX_best = self.WX
        return
    def load_par(self):
        self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta =\
            self.alpha_best, self.beta_best, self.omega_best, self.theta_best, self.gamma_best, self.delta_best
        self.loss = self.loss_best
        self.t = self.t_best
        self.WX_model = self.WX_model_best
        self.WX = self.WX_best
        return
    def search_and_fit(self, **kwargs):
        omega_vales = [2*np.pi] #np.linspace(1/4*np.pi, 2*np.pi, num=8)
        self.save_par()
        self.best_loss = MAX_LOSS
        for omega in omega_vales:
            theta_vals = np.linspace(-0.5, 0.5, num=1+int(omega/(np.pi/4))) * omega
            for theta in theta_vals:
                self.init_par(omega=omega, theta=theta)
                self.fit_paras_(**kwargs)
                if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                    continue
                if self.loss[-1] < self.best_loss:
                    self.save_par()
        self.load_par()
        self.tmp_best_loss = self.best_loss

        gamma_vals = np.linspace(2.5, 7.5, num=3)
        for omega in omega_vales:
            theta_vals = np.linspace(-0.5, 0.5, num=1+int(omega/(np.pi/4))) * omega
            for theta in theta_vals:
                for gamma in gamma_vals:
                    if gamma < self.gamma_thres:
                        continue
                    self.init_par(omega=omega, theta=theta, gamma=gamma)
                    self.fit_paras_(**kwargs)
                    if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                        continue
                    if (self.loss[-1] < self.best_loss) and (self.loss[-1] < 0.9*self.tmp_best_loss):
                        self.save_par()
        self.load_par()
        return

    def fit_paras_(self, **kwargs):
        def mse(data):
            #return self.get_mse(alpha=data[0], beta=data[1], omega=data[2], theta=data[3], gamma=data[4], **kwargs)
            return self.get_mse(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], **kwargs)
        #data_0, cb = np.array([self.alpha, self.beta, self.omega, self.theta, self.gamma]), self.cb_fit_paras_
        data_0, cb = np.array([self.alpha, self.beta, self.theta, self.gamma]), self.cb_fit_paras_
        #bounds=((0, None), (None, None), (0, 2*np.pi), (-np.pi, np.pi), (0, None))
        bounds=((0, None), (None, None), (-np.pi, np.pi), (self.gamma_thres, None))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_fit_paras_(res.x)
    def cb_fit_paras_(self, data):
        #self.update(alpha=data[0], beta=data[1], omega=data[2], theta=data[3], gamma=data[4], refit_time=True, update_weight=True)
        self.update(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], refit_time=True, update_weight=True)

    def search_and_fit_no_theta(self, **kwargs):
        omega = 2*np.pi
        theta = 0
        self.save_par()
        self.best_loss = MAX_LOSS
        self.init_par(omega=omega, theta=theta)
        self.fit_paras_no_theta_(**kwargs)
        if self.loss[-1] < self.best_loss:
            self.save_par()
        self.load_par()
        self.tmp_best_loss = self.best_loss

        gamma_vals = np.linspace(2.5, 7.5, num=3)
        for gamma in gamma_vals:
            if gamma < self.gamma_thres:
                continue
            self.init_par(omega=omega, theta=theta, gamma=gamma)
            self.fit_paras_no_theta_(**kwargs)
            if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                continue
            if (self.loss[-1] < self.best_loss) and (self.loss[-1] < 0.9*self.tmp_best_loss):
                self.save_par()
        self.load_par()
        return

    def fit_paras_no_theta_(self, **kwargs):
        def mse(data):
            return self.get_mse(alpha=data[0], beta=data[1], gamma=data[2], **kwargs)
        data_0, cb = np.array([self.alpha, self.beta, self.gamma]), self.cb_fit_paras_no_theta_
        bounds=((0, None), (None, None), (self.gamma_thres, None))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_fit_paras_no_theta_(res.x)
    def cb_fit_paras_no_theta_(self, data):
        self.update(alpha=data[0], beta=data[1], gamma=data[2], refit_time=True, update_weight=True)

    def refit_paras_(self, **kwargs):
        def mse(data):
            return self.get_mse(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], delta=data[4], **kwargs)
        data_0, cb = np.array([self.alpha, self.beta, self.theta, self.gamma, self.delta]), self.cb_refit_paras_
        bounds=((0, None), (None, None), (-np.pi, np.pi), (0, None), (-10, 10))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_refit_paras_(res.x)
    def cb_refit_paras_(self, data):
        self.update(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], delta=data[4], refit_time=True, update_weight=False)

    def update(
        self,
        t=None,
        alpha=None,
        beta=None,
        omega=None,
        theta=None,
        gamma=None,
        delta=None,
        refit_time=False,
        update_weight=True,
    ):
        loss_prev = self.loss[-1] if len(self.loss) > 0 else MAX_LOSS

        alpha, beta, omega, theta, gamma, delta = self.get_vars(alpha, beta, omega, theta, gamma, delta)

        if update_weight:
            WX_model, delta = self.update_WX(alpha, beta, omega, theta, gamma, method=self.WX_method, thres=self.WX_thres)
        else:
            WX_model = self.WX_model

        WX = use_model(WX_model, self.x, method=self.WX_method, bias=delta)
        if refit_time:
            t = self.assign_time(WX, self.y, alpha, beta, omega, theta, gamma, self.filtered, norm_std=True)
        else:
            t = self.t

        loss = self.get_loss(t, alpha, beta, omega, theta, gamma, delta, WX_model, refit_time=False, norm_std=True, reg=self.reg)
        perform_update = loss < loss_prev

        if perform_update:
            self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta = alpha, beta, omega, theta, gamma, delta
            self.t = t
            self.WX = WX
            self.WX_model = WX_model
            new_pars = np.array([alpha, beta, omega, theta, gamma, delta, self.likelihood, self.varx, self.scaling_y] + self.scaling.reshape(-1).tolist())[:, None]
            self.pars = np.c_[self.pars, new_pars]
            self.loss.append(loss)
            #print('loss:',loss)

        return perform_update

    def assign_time(
        self, WX, y, alpha, beta, omega, theta, gamma, filtered, norm_std=False
    ):
        num = np.clip(len(y), self.n_time_points, self.n_time_points*2)
        tpoints = np.linspace(0, 1, num=num)
        data_model = SplicingDynamics(alpha=alpha, beta=beta, omega=omega, theta=theta,
            gamma=gamma).get_solution(tpoints)

        if norm_std:
            std_WX = get_norm_std(data_model[:,0], WX[filtered]) # WX[filtered], data_model[:,0]
            std_y = get_norm_std(data_model[:,1], y[filtered], par=1.5) # y[filtered], data_model[:,1]
        else:
            std_WX, std_y = 1, 1
        data_model = data_model/np.array([std_WX, std_y])
        data_obs = np.vstack([WX/std_WX, np.ravel(y)/std_y]).T
        t = tpoints[
            ((data_model[None, :, :] - data_obs[:, None, :]) ** 2).sum(axis=2).argmin(axis=1)
        ]
        return t

    def get_raw_velocity(self):
        self.velo_hat = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta) - np.ravel(self.gamma * self.y)
        tmp1 = self.omega * self.t + self.theta
        self.y_t = self.alpha * np.sin(tmp1) + self.beta
        self.velo_t = self.alpha * self.omega * np.cos(tmp1)
        self.velo_normed = (self.velo_hat-self.velo_hat.min()) / (self.velo_hat.max()-self.velo_hat.min()+1e-6)
        return

default_pars_names = ["alpha", "beta", "omega", "theta", "gamma", "delta", "likelihood", "varx", "scaling_y"]
default_pars_names_vec = ["scaling", "weights", "weights_init", "weights_final"]
default_pars_names += default_pars_names_vec

def read_pars(adata, pars_names=None, key="fit"):
    pars = []
    for name in default_pars_names if pars_names is None else pars_names:
        pkey = f"{key}_{name}"
        if name in default_pars_names_vec:
            par = np.zeros([adata.n_vars, adata.var['n_TFs'].max()])
        else:
            par = np.zeros(adata.n_vars) * np.nan
        if pkey in adata.var.keys():
            par = adata.var[pkey].values
        if pkey in adata.varm.keys():
            par = adata.varm[pkey]
        pars.append(par)

    return pars

default_results_names = ["velo_hat", "velo_t", "y_t", "velo_normed", "filtered", "WX", "y"]
def write_result(adata, results):
    for i, name in enumerate(default_results_names):
        adata.layers[name] = results[i]

def write_pars(adata, pars, pars_names=None, add_key="fit"):
    for i, name in enumerate(default_pars_names if pars_names is None else pars_names):
        if name in default_pars_names_vec:
            adata.varm[f"{add_key}_{name}"] = pars[i]
        else:
            adata.var[f"{add_key}_{name}"] = pars[i]

def recover_dynamics(
    data,
    var_names="all",#"velocity_genes",
    n_top_genes=None,
    max_iter=10,
    WX_method='lsq_linear',
    WX_thres=20,
    max_n_TF=999,
    t_max=None,
    fit_time=True,
    fit_scaling=True,
    fit_steady_states=True,
    fit_connected_states=None,
    fit_basal_transcription=None,
    use_raw=False,
    load_pars=None,
    return_model=None,
    plot_results=False,
    steady_state_prior=None,
    add_key="fit",
    copy=False,
    n_jobs=None,
    backend="loky",
    reg=None,
    init_weight_method=None,
    fix_theta=0,
    gamma_thres=0,
    n_time_points=1000,
    **kwargs,
):
    """Recovers the full splicing kinetics of specified genes.

    The model infers transcription rates, splicing rates, degradation rates,
    as well as cell-specific latent time and transcriptional states,
    estimated iteratively by expectation-maximization.

    .. image:: https://user-images.githubusercontent.com/31883718/69636459-ef862800-1056-11ea-8803-0a787ede5ce9.png

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix.
    var_names: `str`,  list of `str` (default: `'velocity_genes'`)
        Names of variables/genes to use for the fitting. If `var_names='velocity_genes'`
        but there is no column `'velocity_genes'` in `adata.var`, velocity genes are
        estimated using the steady state model.
    n_top_genes: `int` or `None` (default: `None`)
        Number of top velocity genes to use for the dynamical model.
    max_iter:`int` (default: `10`)
        Maximal iterations in the EM-Algorithm.
    t_max: `float`, `False` or `None` (default: `None`)
        Total range for time assignments.
    fit_scaling: `bool` or `float` or `None` (default: `True`)
        Whether to fit scaling between unspliced and spliced.
    fit_time: `bool` or `float` or `None` (default: `True`)
        Whether to fit time or keep initially given time fixed.
    fit_steady_states: `bool` or `None` (default: `True`)
        Whether to explicitly model and fit steady states (next to induction/repression)
    fit_connected_states: `bool` or `None` (default: `None`)
        Restricts fitting to neighbors given by connectivities.
    fit_basal_transcription: `bool` or `None` (default: `None`)
        Enables model to incorporate basal transcriptions.
    use_raw: `bool` or `None` (default: `None`)
        if True, use .layers['sliced'], else use moments from .layers['Ms']
    load_pars: `bool` or `None` (default: `None`)
        Load parameters from past fits.
    return_model: `bool` or `None` (default: `None`)
        Whether to return the model as :DynamicsRecovery: object.
    plot_results: `bool` or `None` (default: `False`)
        Plot results after parameter inference.
    steady_state_prior: list of `bool` or `None` (default: `None`)
        Mask for indices used for steady state regression.
    add_key: `str` (default: `'fit'`)
        Key to add to parameter names, e.g. 'fit_t' for fitted time.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to `adata`.
    n_jobs: `int` or `None` (default: `None`)
        Number of parallel jobs.
    backend: `str` (default: "loky")
        Backend used for multiprocessing. See :class:`joblib.Parallel` for valid
        options.

    Returns
    -------
    fit_alpha: `.var`
        inferred transcription rates
    fit_beta: `.var`
        inferred splicing rates
    fit_gamma: `.var`
        inferred degradation rates
    fit_weight: `.var`
        inferred weight of a TF
    fit_t_: `.var`
        inferred switching time points
    fit_scaling: `.var`
        internal variance scaling factor for un/spliced counts
    fit_likelihood: `.var`
        likelihood of model fit
    fit_alignment_scaling: `.var`
        scaling factor to align gene-wise latent times to a universal latent time
    """  # noqa E501

    adata = data.copy() if copy else data

    n_jobs = get_n_jobs(n_jobs=n_jobs)
    logg.info(f"recovering dynamics (using {n_jobs}/{os.cpu_count()} cores)", r=True)

    if len(set(adata.var_names)) != len(adata.var_names):
        logg.warn("Duplicate var_names found. Making them unique.")
        adata.var_names_make_unique()

    if "M_total" not in adata.layers.keys():
        use_raw = True
    if fit_connected_states is None:
        fit_connected_states = not use_raw

    adata.uns["recover_dynamics"] = {
        "fit_connected_states": fit_connected_states,
        "fit_basal_transcription": fit_basal_transcription,
        "use_raw": use_raw,
    }

    if isinstance(var_names, str) and var_names not in adata.var_names:
        if var_names in adata.var.keys():
            if adata.var[var_names][0] in ['True', 'False']:
                adata.var[var_names] = adata.var[var_names].map({'True': True, 'False': False})
            var_names = adata.var_names[adata.var[var_names].values]
        elif use_raw or var_names == "all":
            var_names = adata.var_names
        elif "_genes" in var_names:
            from .velocity import Velocity

            velo = Velocity(adata, use_raw=use_raw)
            #velo.compute_deterministic(perc=[5, 95])
            #var_names = adata.var_names[velo._velocity_genes]
            #adata.var["fit_r2"] = velo._r2
        else:
            # raise ValueError("Variable name not found in var keys.")
            return False
    if not isinstance(var_names, str):
        var_names = list(np.ravel(var_names))

    var_names = make_unique_list(var_names, allow_array=True)
    var_names = np.array([name for name in var_names if name in adata.var_names])
    if len(var_names) == 0:
        # raise ValueError("Variable name not found in var keys.")
        return False
    if n_top_genes is not None and len(var_names) > n_top_genes:
        X = adata[:, var_names].layers[("total" if use_raw else "M_total")]
        var_names = var_names[np.argsort(np.sum(X, 0))[::-1][:n_top_genes]]
    if return_model is None:
        return_model = len(var_names) < 5

    pars = read_pars(adata)
    alpha, beta, omega, theta, gamma, delta, likelihood, varx, scaling_y = pars[:-4]
    scaling, weights, weights_init, weights_final = pars[-4], pars[-3], pars[-2], pars[-1]
    # likelihood[np.isnan(likelihood)] = 0
    idx, L = [], []
    velo_hat, velo_t, velo_normed = np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan
    y_t, y, filtered = np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan
    WX = np.zeros(adata.shape) * np.nan

    T = np.zeros(adata.shape) * np.nan
    if "fit_t" in adata.layers.keys():
        T = adata.layers["fit_t"]

    conn = get_connectivities(adata) if fit_connected_states else None

    res = parallelize(
        _fit_recovery,
        var_names,
        n_jobs,
        unit="gene",
        as_array=False,
        backend=backend,
        show_progress_bar=len(var_names) > 9,
    )(
        adata=adata,
        use_raw=use_raw,
        load_pars=load_pars,
        max_iter=max_iter,
        fit_time=fit_time,
        fit_steady_states=fit_steady_states,
        fit_scaling=fit_scaling,
        fit_basal_transcription=fit_basal_transcription,
        conn=conn,
        WX_method=WX_method,
        WX_thres=WX_thres,
        max_n_TF=max_n_TF,
        reg=reg,
        init_weight_method=init_weight_method,
        fix_theta=fix_theta,
        gamma_thres=gamma_thres,
        n_time_points=n_time_points,
        **kwargs,
    )
    idx, dms = map(_flatten, zip(*res)) # gene_id_list, model_list

    for ix, dm in zip(idx, dms):
        T[:, ix] = dm.t
        if WX_method in ['lsq_linear']:
            weight = np.ravel(dm.WX_model)
            weights[ix, :len(weight)] = weight
            weights_init[ix, :len(weight)] = np.ravel(dm.weights_init)
            weights_final[ix, :len(weight)] = np.ravel(dm.weights_final)
        alpha[ix], beta[ix], omega[ix], theta[ix], gamma[ix], delta[ix] = dm.alpha, dm.beta, dm.omega, dm.theta, dm.gamma, dm.delta #dm.pars[:5, -1]
        scaling[ix][:dm.n_TFs], scaling_y[ix] = dm.scaling, dm.scaling_y # dm.pars[10, -1],  dm.pars[7, -1]
        likelihood[ix], varx[ix] = dm.likelihood, dm.varx # dm.pars[5, -1],  dm.pars[6, -1]
        velo_hat[:,ix], velo_t[:,ix], velo_normed[:,ix] = dm.velo_hat.reshape(-1), dm.velo_t.reshape(-1), dm.velo_normed.reshape(-1)
        y_t[:,ix], y[:,ix], filtered[:,ix] =  dm.y_t.reshape(-1), dm.y.reshape(-1), dm.filtered.reshape(-1)
        WX[:,ix] = dm.WX.reshape(-1)
        L.append(dm.loss)

    _pars = [
        alpha, beta, omega, theta, gamma, delta, likelihood, varx, scaling_y, \
        scaling, weights, weights_init, weights_final
    ]
    write_pars(adata, _pars)
    write_result(adata, [velo_hat, velo_t, y_t, velo_normed, filtered, WX, y])

    adata.layers["fit_t_raw"] = T.copy()
    #if "fit_t" in adata.layers.keys():
    #    adata.layers["fit_t"][:, idx] = (
    #        T[:, idx] if conn is None else conn.dot(T[:, idx])
    #    )
    #else:
    adata.layers["fit_t"] = T if conn is None else conn.dot(T)

    if L:  # is False if only one invalid / irrecoverable gene was given in var_names
        cur_len = adata.varm["loss"].shape[1] if "loss" in adata.varm.keys() else 2
        max_len = max(np.max([len(loss) for loss in L]), cur_len) if L else cur_len
        loss = np.ones((adata.n_vars, max_len)) * np.nan

        if "loss" in adata.varm.keys():
            loss[:, :cur_len] = adata.varm["loss"]

        loss[idx] = np.vstack(
            [
                np.concatenate([loss, np.ones(max_len - len(loss)) * np.nan])
                for loss in L
            ]
        )
        adata.varm["loss"] = loss

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint(
        "added \n"
        f"    '{add_key}_pars', "
        f"fitted parameters for splicing dynamics (adata.var)"
    )

    '''
    P = []
    if plot_results:  # Plot Parameter Stats
        n_rows, n_cols = len(var_names[:4]), 6
        figsize = [2 * n_cols, 1.5 * n_rows]  # rcParams['figure.figsize']
        fontsize = rcParams["font.size"]
        fig, axes = pl.subplots(nrows=n_rows, ncols=6, figsize=figsize)
        pl.subplots_adjust(wspace=0.7, hspace=0.5)
        for i, gene in enumerate(var_names[:4]):
            #if t_max is not False:
            #    mi = dm.m[i]
            #    P[i] *= np.array([1 / mi, 1 / mi, 1 / mi, mi, 1])[:, None]
            ax = axes[i] if n_rows > 1 else axes
            for j, pij in enumerate(P[i]):
                ax[j].plot(pij)
            ax[len(P[i])].plot(L[i])
            if i == 0:
                pars_names = ["alpha", "beta", "gamma", "weight", "t_", "scaling", "loss"]
                for j, name in enumerate(pars_names):
                    ax[j].set_title(name, fontsize=fontsize)
    '''
    if return_model:
        logg.info("\noutputs model fit of gene:", dm.gene)

    #return dm if return_model else adata if copy else None
    return True

def latent_time(
    data,
    vkey="velocity",
    t_key='normed_t',
    min_likelihood=0.1,
    min_confidence=0.75,
    min_corr_diffusion=None,
    weight_diffusion=None,
    root_key=None,
    end_key=None,
    t_max=None,
    copy=False,
):
    """Computes a gene-shared latent time.

    Gene-specific latent timepoints obtained from the dynamical model are coupled to a
    universal gene-shared latent time, which represents the cell’s internal clock and
    is based only on its transcriptional dynamics.

    .. image:: https://user-images.githubusercontent.com/31883718/69636500-03318e80-1057-11ea-9e14-ae9f907711cc.png

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix
    vkey: `str` (default: `'velocity'`)
        Name of velocity estimates to be used.
    min_likelihood: `float` between `0` and `1` or `None` (default: `.1`)
        Minimal likelihood fitness for genes to be included to the weighting.
    min_confidence: `float` between `0` and `1` (default: `.75`)
        Parameter for local coherence selection.
    min_corr_diffusion: `float` between `0` and `1` or `None` (default: `None`)
        Only select genes that correlate with velocity pseudotime obtained
        from diffusion random walk on velocity graph.
    weight_diffusion: `float` or `None` (default: `None`)
        Weight applied to couple latent time with diffusion-based velocity pseudotime.
    root_key: `str` or `None` (default: `'root_cells'`)
        Key (.uns, .obs) of root cell to be used.
        If not set, it obtains root cells from velocity-inferred transition matrix.
    end_key: `str` or `None` (default: `None`)
        Key (.obs) of end points to be used.
    t_max: `float` or `None` (default: `None`)
        Overall duration of differentiation process.
        If not set, a overall transcriptional timescale of 20 hours is used as prior.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to `adata`.

    Returns
    -------
    latent_time: `.obs`
        latent time from learned dynamics for each cell
    """  # noqa E501

    adata = data.copy() if copy else data

    from .dynamical_model_utils import compute_shared_time, root_time
    from .terminal_states import terminal_states
    from .utils import scale, vcorrcoef
    from .velocity_graph import velocity_graph
    from .velocity_pseudotime import velocity_pseudotime

    if t_key not in adata.layers.keys():
        raise ValueError("you need to run `tl.recover_dynamics` first.")

    if f"{vkey}_graph" not in adata.uns.keys():
        velocity_graph(adata, approx=True)

    if root_key is None:
        terminal_keys = ["root_cells", "starting_cells", "root_states_probs"]
        keys = [key for key in terminal_keys if key in adata.obs.keys()]
        if len(keys) > 0:
            root_key = keys[0]
    if root_key not in adata.uns.keys() and root_key not in adata.obs.keys():
        root_key = "root_cells"
    if root_key not in adata.obs.keys():
        terminal_states(adata, vkey=vkey)

    t = np.array(adata.layers[t_key])
    idx_valid = ~np.isnan(t.sum(0))
    if min_likelihood is not None:
        likelihood = adata.var["fit_likelihood"].values
        idx_valid &= np.array(likelihood >= min_likelihood, dtype=bool)
    t = t[:, idx_valid]
    t_sum = np.sum(t, 1)
    conn = get_connectivities(adata)

    if root_key not in adata.uns.keys():
        roots = np.argsort(t_sum)
        idx_roots = np.array(adata.obs[root_key][roots])
        idx_roots[pd.isnull(idx_roots)] = 0
        if np.any([isinstance(ix, str) for ix in idx_roots]):
            idx_roots = np.array([isinstance(ix, str) for ix in idx_roots], dtype=int)
        idx_roots = idx_roots.astype(float) > 1 - 1e-3
        if np.sum(idx_roots) > 0:
            roots = roots[idx_roots]
        else:
            logg.warn(
                "No root cells detected. Consider specifying "
                "root cells to improve latent time prediction."
            )
    else:
        roots = [adata.uns[root_key]]
        root_key = f"root cell {adata.uns[root_key]}"

    if end_key in adata.obs.keys():
        fates = np.argsort(t_sum)[::-1]
        idx_fates = np.array(adata.obs[end_key][fates])
        idx_fates[pd.isnull(idx_fates)] = 0
        if np.any([isinstance(ix, str) for ix in idx_fates]):
            idx_fates = np.array([isinstance(ix, str) for ix in idx_fates], dtype=int)
        idx_fates = idx_fates.astype(float) > 1 - 1e-3
        if np.sum(idx_fates) > 0:
            fates = fates[idx_fates]
    else:
        fates = [None]

    logg.info(
        f"computing latent time using {root_key}"
        f"{', ' + end_key if end_key in adata.obs.keys() else ''} as prior",
        r=True,
    )

    VPT = velocity_pseudotime(
        adata, vkey, root_key=roots[0], end_key=fates[0], return_model=True
    )
    vpt = VPT.pseudotime

    if min_corr_diffusion is not None:
        corr = vcorrcoef(t.T, vpt)
        t = t[:, np.array(corr >= min_corr_diffusion, dtype=bool)]

    if root_key in adata.uns.keys():
        root = adata.uns[root_key]
        t, t_ = root_time(t, root=root)
        latent_time = compute_shared_time(t)
    else:
        roots = roots[:4]
        latent_time = np.ones(shape=(len(roots), adata.n_obs))
        for i, root in enumerate(roots):
            t, t_ = root_time(t, root=root)
            latent_time[i] = compute_shared_time(t)
        latent_time = scale(np.mean(latent_time, axis=0))

    if fates[0] is not None:
        fates = fates[:4]
        latent_time_ = np.ones(shape=(len(fates), adata.n_obs))
        for i, fate in enumerate(fates):
            t, t_ = root_time(t, root=fate)
            latent_time_[i] = 1 - compute_shared_time(t)
        latent_time = scale(latent_time + 0.2 * scale(np.mean(latent_time_, axis=0)))

    tl = latent_time
    tc = conn.dot(latent_time)

    z = tl.dot(tc) / tc.dot(tc)
    tl_conf = (1 - np.abs(tl / np.max(tl) - tc * z / np.max(tl))) ** 2
    idx_low_confidence = tl_conf < min_confidence

    if weight_diffusion is not None:
        w = weight_diffusion
        latent_time = (1 - w) * latent_time + w * vpt
        latent_time[idx_low_confidence] = vpt[idx_low_confidence]
    else:
        conn_new = conn.copy()
        conn_new[:, idx_low_confidence] = 0
        conn_new.eliminate_zeros()
        latent_time = conn_new.dot(latent_time)

    latent_time = scale(latent_time)
    if t_max is not None:
        latent_time *= t_max

    adata.obs["latent_time"] = latent_time

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint("added \n" "    'latent_time', shared time (adata.obs)")
    return adata if copy else None

recover_latent_time = latent_time

def rank_dynamical_genes(data, n_genes=100, groupby=None, copy=False):
    """Rank genes by likelihoods per cluster/regime.

    This ranks genes by their likelihood obtained from the
    dynamical model grouped by clusters specified in groupby.

    Arguments
    ----------
    data : :class:`~anndata.AnnData`
        Annotated data matrix.
    n_genes : `int`, optional (default: 100)
        The number of genes that appear in the returned tables.
    groupby: `str`, `list` or `np.ndarray` (default: `None`)
        Key of observations grouping to consider.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to data.

    Returns
    -------
    rank_dynamical_genes : `.uns`
        Structured array to be indexed by group id storing the gene
        names. Ordered according to scores.
    """

    from .dynamical_model_utils import get_divergence

    adata = data.copy() if copy else data

    logg.info("ranking genes by cluster-specific likelihoods", r=True)

    groupby = (
        groupby
        if isinstance(groupby, str) and groupby in adata.obs.keys()
        else "clusters"
        if "clusters" in adata.obs.keys()
        else "louvain"
        if "louvain" in adata.obs.keys()
        else "velocity_clusters"
        if "velocity_clusters" in adata.obs.keys()
        else None
    )

    vdata = adata[:, ~np.isnan(adata.var["fit_alpha"])]
    groups = vdata.obs[groupby].cat.categories

    ll = get_divergence(
        vdata,
        mode="gene_likelihood",
        use_connectivities=True,
        clusters=adata.obs[groupby],
    )
    idx_sorted = np.argsort(np.nan_to_num(ll), 1)[:, ::-1][:, :n_genes]
    rankings_gene_names = vdata.var_names[idx_sorted]
    rankings_gene_scores = np.sort(np.nan_to_num(ll), 1)[:, ::-1][:, :n_genes]

    key = "rank_dynamical_genes"
    if key not in adata.uns.keys():
        adata.uns[key] = {}

    adata.uns[key] = {
        "names": np.rec.fromarrays(
            [n for n in rankings_gene_names],
            dtype=[(f"{rn}", "U50") for rn in groups],
        ),
        "scores": np.rec.fromarrays(
            [n.round(2) for n in rankings_gene_scores],
            dtype=[(f"{rn}", "float32") for rn in groups],
        ),
    }

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint("added \n" f"    '{key}', sorted scores by group ids (adata.uns)")

    return adata if copy else None

def _fit_recovery(
    var_names,
    adata,
    use_raw,
    load_pars,
    max_iter,
    fit_time,
    fit_steady_states,
    conn,
    fit_scaling,
    fit_basal_transcription,
    max_n_TF,
    queue,
    WX_method,
    reg,
    WX_thres,
    init_weight_method,
    fix_theta,
    gamma_thres,
    n_time_points,
    **kwargs,
):

    idx, dms = [], []
    for gene in var_names:#################
        dm = DynamicsRecovery(
            adata,
            gene,
            use_raw=use_raw,
            load_pars=load_pars,
            max_iter=max_iter,
            fit_time=fit_time,
            fit_steady_states=fit_steady_states,
            fit_connected_states=conn,
            fit_scaling=fit_scaling,
            fit_basal_transcription=fit_basal_transcription,
            WX_method=WX_method,
            WX_thres=WX_thres,
            max_n_TF=max_n_TF,
            reg=reg,
            init_weight_method=init_weight_method,
            fix_theta=fix_theta,
            gamma_thres=gamma_thres,
            n_time_points=n_time_points,
            **kwargs,
        )
        tmp = str(list(var_names).index(gene)) + '/' + str(len(var_names))
        if dm.recoverable:
            #try:
            print('Processing', tmp, dm.gene)
            dm.fit()

            ix = np.where(adata.var_names == gene)[0][0]
            idx.append(ix)
            dms.append(dm)
            print(tmp, dm.gene, 'FINISHED with n_TFs:', dm.n_TFs)
            #except:
            #    dm.recoverable = False
            #    logg.warn(dm.gene, "not recoverable during optimization.")

        else:
            logg.warn(tmp, dm.gene, "not recoverable due to insufficient samples.")

        if queue is not None:
            queue.put(1)

    if queue is not None:
        queue.put(None)

    return idx, dms # gene_id_list, model_list

def _flatten(iterable):
    return [i for it in iterable for i in it]

{'PCA': 'sklearn.decomposition.PCA', 'lsq_linear': 'scipy.optimize.lsq_linear'}


#!/usr/bin/env python
# coding: utf-8
#
# Author: Kazuto Nakashima
# URL:    https://kazuto1011.github.io
# Date:   04 March 2019

import hashlib
import os
import shutil
import tarfile
import tempfile
from collections import OrderedDict
from glob import glob
from urllib.request import urlopen

import click
import h5py
import numpy as np
import tensorflow
import torch
import torch.nn as nn
from tensorflow import keras
from tqdm import tqdm

from .modules import _ConvBnReLU, _Flatten, _SeparableConv2d, _SepConvBnReLU
from .sync_batchnorm.batchnorm import SynchronizedBatchNorm2d

def _parse_tensorflow_ckpt(model_path):
    reader = tensorflow.train.NewCheckpointReader(model_path)
    var_to_shape_map = reader.get_variable_to_shape_map()
    data = dict()
    for name in var_to_shape_map:
        if "global_step" not in name.lower():
            tensor = reader.get_tensor(name)
            data[name] = tensor
    return data

def load_tensorflow_resnet(tar_url, ckpt_relpath, model_torch, model_name):

    torch_home = os.path.expanduser(os.getenv("TORCH_HOME", "~/.torch"))
    model_dir = os.getenv("TORCH_MODEL_ZOO", os.path.join(torch_home, "models"))
    os.makedirs(model_dir, exist_ok=True)

    save_path = os.path.join(model_dir, model_name)
    if os.path.exists(save_path):
        return torch.load(save_path)

    # Download TensorFlow model
    u = urlopen(tar_url)
    meta = u.info()
    if hasattr(meta, "getheaders"):
        content_length = meta.getheaders("Content-Length")
    else:
        content_length = meta.get_all("Content-Length")
    if content_length is not None and len(content_length) > 0:
        file_size = int(content_length[0])
    else:
        file_size = None
    with tempfile.NamedTemporaryFile() as f:
        print(f"Downloading: {tar_url}")
        with tqdm(total=file_size) as pbar:
            while True:
                buffer = u.read(8192)
                if len(buffer) == 0:
                    break
                f.write(buffer)
                pbar.update(len(buffer))
        with tempfile.TemporaryDirectory() as tar_dest:
            with tarfile.open(f.name, "r:gz") as tar:
                tar.extractall(tar_dest)
            state_dict_tf = _parse_tensorflow_ckpt(os.path.join(tar_dest, ckpt_relpath))

    # Parse tensorflow model
    layers_tf = {"conv": set(), "bn": set(), "fc": set()}
    for name in state_dict_tf:
        name = name.split("/")
        name, param = "/".join(name[:-1]), name[-1]
        if "conv" in name:
            layers_tf["conv"].add(name)
        elif "batch_normalization" in name:
            layers_tf["bn"].add(name)
        elif "dense" in name:
            layers_tf["fc"].add(name)

    # Parse PyTorch model
    ref_state_dict = model_torch.state_dict()
    layers_torch = {"conv": [], "bn": [], "fc": []}
    for name, module in model_torch.named_modules():
        if isinstance(module, nn.Conv2d):
            layers_torch["conv"].append(name)
        elif isinstance(module, nn.BatchNorm2d):
            layers_torch["bn"].append(name)
        elif isinstance(module, SynchronizedBatchNorm2d):
            layers_torch["bn"].append(name)
        elif isinstance(module, nn.Linear):
            layers_torch["fc"].append(name)

    assert len(layers_tf["conv"]) == len(layers_torch["conv"])
    assert len(layers_tf["bn"]) == len(layers_torch["bn"])
    assert len(layers_tf["fc"]) == len(layers_torch["fc"])

    # Generate
    key_pairs = []
    add_suffix = lambda idx: f"_{idx}" if idx != 0 else ""
    for idx, key_torch in enumerate(layers_torch["conv"]):
        key_tf = "resnet_model/conv2d" + add_suffix(idx)
        key_pairs.append((key_torch + ".weight", key_tf + "/kernel"))
    for idx, key_torch in enumerate(layers_torch["bn"]):
        key_tf = "resnet_model/batch_normalization" + add_suffix(idx)
        key_pairs.append((key_torch + ".weight", key_tf + "/gamma"))
        key_pairs.append((key_torch + ".running_mean", key_tf + "/moving_mean"))
        key_pairs.append((key_torch + ".bias", key_tf + "/beta"))
        key_pairs.append((key_torch + ".running_var", key_tf + "/moving_variance"))
    for idx, key_torch in enumerate(layers_torch["fc"]):
        key_tf = "resnet_model/dense" + add_suffix(idx)
        key_pairs.append((key_torch + ".weight", key_tf + "/kernel"))
        key_pairs.append((key_torch + ".bias", key_tf + "/bias"))

    # Convert weights
    new_state_dict = OrderedDict()
    for key_torch, key_tf in key_pairs:
        param_pth = ref_state_dict[key_torch]
        param_tf = state_dict_tf[key_tf]
        if param_tf.ndim == 2:
            param_tf = param_tf.transpose(1, 0)
        elif param_tf.ndim == 4:
            param_tf = param_tf.transpose(3, 2, 0, 1)
        assert (
            tuple(param_pth.shape) == param_tf.shape
        ), "Inconsistent shape: {}, {}".format(tuple(param_pth.shape), param_tf.shape)
        new_state_dict[key_torch] = torch.FloatTensor(param_tf)

    model_torch.load_state_dict(new_state_dict)

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    torch.save(new_state_dict, save_path)

    return new_state_dict

def load_keras_xceptionv1(model_torch):

    assert model_torch.n_classes == 1000

    weights_path = keras.utils.get_file(
        "xception_weights_tf_dim_ordering_tf_kernels.h5",
        (
            "https://github.com/fchollet/deep-learning-models/"
            "releases/download/v0.4/"
            "xception_weights_tf_dim_ordering_tf_kernels.h5"
        ),
        cache_subdir="models",
        file_hash="0a58e3b7378bc2990ea3b43d5981f1f6",
    )

    with h5py.File(weights_path, mode="r") as f:
        state_dict = OrderedDict()

        s_idx = 1
        b_idx = 1
        for pth_layer, pth_module in model_torch.named_modules():
            if isinstance(pth_module, _SeparableConv2d):
                h5_layer = f"separableconvolution2d_{s_idx}"
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_depthwise_kernel:0"])
                h5_value = torch.tensor(h5_value.transpose(2, 3, 0, 1))
                state_dict[pth_layer + ".depthwise.weight"] = h5_value
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_pointwise_kernel:0"])
                h5_value = torch.tensor(h5_value.transpose(3, 2, 0, 1))
                state_dict[pth_layer + ".pointwise.weight"] = h5_value
                s_idx += 1

            if isinstance(pth_module, nn.BatchNorm2d):
                h5_layer = f"batchnormalization_{b_idx}"
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_gamma:0"])
                state_dict[pth_layer + ".weight"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_beta:0"])
                state_dict[pth_layer + ".bias"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_running_mean:0"])
                state_dict[pth_layer + ".running_mean"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_running_std:0"])
                state_dict[pth_layer + ".running_var"] = torch.tensor(h5_value)
                b_idx += 1
            elif isinstance(pth_module, SynchronizedBatchNorm2d):
                h5_layer = f"batchnormalization_{b_idx}"
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_gamma:0"])
                state_dict[pth_layer + ".weight"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_beta:0"])
                state_dict[pth_layer + ".bias"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_running_mean:0"])
                state_dict[pth_layer + ".running_mean"] = torch.tensor(h5_value)
                h5_value = np.asarray(f[h5_layer][f"{h5_layer}_running_std:0"])
                state_dict[pth_layer + ".running_var"] = torch.tensor(h5_value)
                b_idx += 1

        idx = 1
        for pth_layer, pth_module in model_torch.named_modules():
            if isinstance(pth_module, nn.Conv2d):
                if pth_layer + ".weight" not in state_dict:
                    h5_layer = f"convolution2d_{idx}"
                    h5_value = np.asarray(f[h5_layer][f"{h5_layer}_W:0"])
                    h5_value = torch.tensor(h5_value.transpose(3, 2, 0, 1))
                    state_dict[pth_layer + ".weight"] = h5_value
                    idx += 1

        h5_value = torch.tensor(np.asarray(f["dense_2"]["dense_2_W:0"]))
        state_dict["exit_flow.fc.weight"] = h5_value.permute(1, 0)
        h5_value = torch.tensor(np.asarray(f["dense_2"]["dense_2_b:0"]))
        state_dict["exit_flow.fc.bias"] = h5_value

        return state_dict

def load_keras_model():
    pass