import os
import sys
from PyQt5.QtWidgets import *
from PyQt5.QtCore import Qt, pyqtSlot
from PyQt5.QtGui import QIcon

from interface import mysqlHandler
from widget import CommonWidget

sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
import Constants

# def openDBConnection(self):
#     self._db_connection = mysqlHandler.dbConnect()
#     self._cursor = mysqlHandler.dbDictCursor(self._db_connection)
#
# def runSQL_Select(self, sql):
#     self._cursor.execute(sql)
#     # object: dictionary
#     result = self._cursor.fetchall()
#     return result
#
# def closeDBConnection(self):
#     self._db_connection.close()

__NOT_SELECTED__ = -1

class OrgWindow(QWidget):
    def __init__(self, parent=None, company=None):
        super().__init__(parent)
        if company:
            self._company = company
        else:
            self._company = 'DEMO'

        if parent:
            self._parent = parent
            self._company = parent._company
        else:
            self._parent = None

        self.initUI()

    def initUI(self):

        self.setWindowTitle(self._company + ' - ' + '조직관리')
        self.setWindowIcon(QIcon(Constants.BASE_PATH + '/img/daesung.ico'))

        self._storedUserList = []
        self._currentCodeList = []

        self.composeLayout()

        self.getTable_Org()
        self.getTable_User()

        self.composeOrgTree()

        self.setSignals()

    # def setWindowIndex(self, index):
    #     self._index = index
    #
    # def getWindowIndex(self):
    #     return self._index

    def composeLayout(self):

        titleLayout = CommonWidget.TitleLayout('조직관리')

        # mainButtonLayout = QVBoxLayout()
        # self._btnClose = QPushButton('닫기')
        # # self._btnClose.setFixedWidth(100)
        # self._btnClose.resize(self._btnClose.sizeHint())
        # mainButtonLayout.addWidget(self._btnClose)
        # mainButtonLayout.setAlignment(Qt.AlignRight)

        bodyLayout = QGridLayout()

        leftLayout = QHBoxLayout()
        self._orgTreeWidget = QTreeWidget()
        # self._orgTreeWidget.setFixedWidth(400)
        # self._orgTreeWidget.setFixedHeight(500)
        leftLayout.addWidget(self._orgTreeWidget)
        self._orgTreeWidget.header().setVisible(True)
        orgHeaderLabels = ['회사조직', '상위조직', '중간조직', '조직']
        self._orgTreeWidget.setHeaderLabels(orgHeaderLabels)
        self._orgTreeWidget.setColumnCount(len(orgHeaderLabels))
        self._orgTreeWidget.header().setSectionResizeMode(QHeaderView.Stretch)
        # self._orgTreeWidget.setAlternatingRowColors(True)
        rightLayout = QVBoxLayout()

        subtitleLayout = QHBoxLayout()
        self._orgLabel = CommonWidget.SubtitleLabel('조직 구성원')
        self._countLabel = CommonWidget.SubtitleLabel('(인원수)')   # 인원수
        subtitleLayout.addWidget(self._orgLabel)
        subtitleLayout.addWidget(self._countLabel)

        codeTableLayout = QHBoxLayout()
        self._userTableWidget = QTableWidget(self)
        self._userTableWidget.setColumnCount(4)
        self._userTableWidget.setHorizontalHeaderLabels(['Chk', '조직', '사용자', '지역'])
        # self._userTableWidget.setColumnWidth(0, 15)
        self._userTableWidget.setAlternatingRowColors(True)      # 행 색상이 번갈아 구분 표시되게 한다.
        self._userTableWidget.setEditTriggers(QAbstractItemView.NoEditTriggers)      # Edit 금지
        self._userTableWidget.setSelectionBehavior(QAbstractItemView.SelectRows)     # Row 단위로 선택이 표되게 한다.
        self._userTableWidget.setSelectionMode(QAbstractItemView.SingleSelection)    # 하나의 Row만 선택하게 함. (다중
        # self._codeTableView.setFixedWidth(300)
        # self._codeTableView.setFixedHeight(380)
        codeTableLayout.addWidget(self._userTableWidget)

        orgItemLayout = QGridLayout()
        self._textBoxCompany = QLineEdit()
        self._textBoxSuperOrg = QLineEdit()
        self._textBoxMiddleOrg = QLineEdit()
        self._textBoxOrg = QLineEdit()
        self._textBoxCompany.setAlignment(Qt.AlignCenter)
        self._textBoxSuperOrg.setAlignment(Qt.AlignCenter)
        self._textBoxMiddleOrg.setAlignment(Qt.AlignCenter)
        self._textBoxOrg.setAlignment(Qt.AlignCenter)
        # self._textBoxCodeID.setFixedWidth(50)
        # self._textBoxCodeName.setFixedWidth(240)
        orgItemLayout.addWidget(QLabel('회사'), 0, 0, Qt.AlignCenter)
        orgItemLayout.addWidget(QLabel('상위조직'), 0, 1, Qt.AlignCenter)
        orgItemLayout.addWidget(QLabel('중간조직'), 0, 2, Qt.AlignCenter)
        orgItemLayout.addWidget(QLabel('부서(팀)'), 0, 3, Qt.AlignCenter)
        orgItemLayout.addWidget(self._textBoxCompany, 1, 0, Qt.AlignCenter)
        orgItemLayout.addWidget(self._textBoxSuperOrg, 1, 1, Qt.AlignCenter)
        orgItemLayout.addWidget(self._textBoxMiddleOrg, 1, 2, Qt.AlignCenter)
        orgItemLayout.addWidget(self._textBoxOrg, 1, 3, Qt.AlignCenter)

        self._btnAdd = QPushButton('추가')
        self._btnRemove = QPushButton('삭제')
        orgButtonLayout = QHBoxLayout()
        orgButtonLayout.addWidget(self._btnAdd)
        orgButtonLayout.addWidget(self._btnRemove)

        orgItemFrame = QFrame()
        orgItemAllLayout = QVBoxLayout()
        orgItemAllLayout.addWidget(CommonWidget.SubtitleLabel('조직'))
        orgItemAllLayout.addLayout(orgItemLayout)
        orgItemAllLayout.addLayout(orgButtonLayout)
        orgItemFrame.setFrameShape(QFrame.StyledPanel)
        orgItemFrame.setFrameShadow(QFrame.Raised)
        orgItemFrame.setLineWidth(1)
        orgItemFrame.setMidLineWidth(3)
        orgItemFrame.setLayout(orgItemAllLayout)

        moveUserFrame = QFrame()
        moveUserAllLayout = QVBoxLayout()
        moveUserLayout = QGridLayout()
        self._textBoxMoveCompany = QLineEdit()
        self._textBoxMoveOrg = QLineEdit()
        self._btnMove = QPushButton('부서 이동')
        moveUserLayout.addWidget(QLabel('이동할 회사'), 0, 0, Qt.AlignCenter)
        moveUserLayout.addWidget(QLabel('이동할 부서(팀)'), 0, 1, Qt.AlignCenter)
        moveUserLayout.addWidget(self._textBoxMoveCompany, 1, 0, Qt.AlignCenter)
        moveUserLayout.addWidget(self._textBoxMoveOrg, 1, 1, Qt.AlignCenter)
        moveUserLayout.addWidget(self._btnMove, 0, 2, 2, 1, Qt.AlignCenter)
        moveUserFrame.setFrameShape(QFrame.StyledPanel)
        moveUserFrame.setFrameShadow(QFrame.Raised)
        moveUserFrame.setLineWidth(1)
        moveUserFrame.setMidLineWidth(3)
        moveUserAllLayout.addWidget(CommonWidget.SubtitleLabel('조직 이동'))
        moveUserAllLayout.addLayout(moveUserLayout)
        moveUserFrame.setLayout(moveUserAllLayout)

        rightLayout.addWidget(orgItemFrame)
        rightLayout.addLayout(subtitleLayout)
        rightLayout.addLayout(codeTableLayout)
        rightLayout.addWidget(moveUserFrame)

        bodyLayout.addLayout(leftLayout, 0, 0)
        bodyLayout.addLayout(rightLayout, 0, 1)

        winLayout = QVBoxLayout()
        winLayout.addLayout(titleLayout)
        # winLayout.addLayout(mainButtonLayout)
        winLayout.addLayout(bodyLayout)

        self.setLayout(winLayout)

    def getTable_User(self):
        db_connection = mysqlHandler.dbConnect()
        cursor_dict = mysqlHandler.dbDictCursor(db_connection)
        sql = 'select org, user_name, region from myasset.it_user ' \
              ' order by org, user_name'
        cursor_dict.execute(sql)
        db_connection.close()
        self._userAllList = cursor_dict.fetchall()

    def getTable_Org(self):
        db_connection = mysqlHandler.dbConnect()
        cursor_dict = mysqlHandler.dbDictCursor(db_connection)
        sql = 'select company, super_org2, super_org1, org from myasset.it_user_org ' \
              ' order by company, super_org2, super_org1, org'
        cursor_dict.execute(sql)
        db_connection.close()
        self._codeTypeList = cursor_dict.fetchall()

    def composeOrgTree(self):
        self._orgTreeWidget.clear()
        rootItem = QTreeWidget.invisibleRootItem(self._orgTreeWidget)
        lastCompanyText = None
        lastSuperOrgText = None
        lastMiddleOrgText = None
        lastOrgText = None
        for r in self._codeTypeList:
            thisCompanyText = r.get('company')
            thisSuperOrgText = r.get('super_org2')
            thisMiddleOrgText = r.get('super_org1')
            thisOrgText = r.get('org')
            if thisCompanyText != lastCompanyText:
                companyItem = QTreeWidgetItem(rootItem)  # company
                companyItem.setText(Constants.TREE_COL_0, thisCompanyText)
                lastCompanyText = thisCompanyText
                lastSuperOrgText = None
                lastMiddleOrgText = None
                lastOrgText = None
            if thisCompanyText == lastCompanyText \
                and thisSuperOrgText != lastSuperOrgText:
                superOrgItem = QTreeWidgetItem(companyItem)  # super_org2
                superOrgItem.setText(Constants.TREE_COL_0, thisSuperOrgText)
                superOrgItem.setText(Constants.TREE_COL_1, thisSuperOrgText)
                lastSuperOrgText = thisSuperOrgText
                lastMiddleOrgText = None
                lastOrgText = None
            if thisCompanyText == lastCompanyText \
                and thisSuperOrgText == lastSuperOrgText \
                and thisMiddleOrgText != lastMiddleOrgText:
                middleOrgItem = QTreeWidgetItem(superOrgItem)  # super_org1
                middleOrgItem.setText(Constants.TREE_COL_0, thisMiddleOrgText)
                middleOrgItem.setText(Constants.TREE_COL_2, thisMiddleOrgText)
                lastMiddleOrgText = thisMiddleOrgText
                lastOrgText = None
            if thisCompanyText == lastCompanyText \
                and thisSuperOrgText == lastSuperOrgText \
                and thisMiddleOrgText == lastMiddleOrgText \
                and thisOrgText != lastOrgText:
                orgItem = QTreeWidgetItem(middleOrgItem)  # org
                orgItem.setText(Constants.TREE_COL_0, thisOrgText)
                orgItem.setText(Constants.TREE_COL_3, thisOrgText)
                lastOrgText = thisOrgText
        self._orgTreeWidget.expandAll()

    # def changeState(self, state):
    #     """
    #     toggle 상태에 따라 배경색과 상태 텍스트 변환
    #     """
    #     self.runState.setStyleSheet("background-color: %s" % ({True: "green", False: "red"}[state]))
    #     self.runState.setText({True: "ON", False: "OFF"}[state])

    def keyReleaseEvent(self, e):
        if e.key() == Qt.Key_Down or e.key() == Qt.Key_Up:
            self.__treeItemClicked()

    def setSignals(self):
        # self._btnClose.clicked.connect(self.close)  # closeEvent
        # self._orgTreeWidget.itemChanged.connect(self.__treeItemChanged)
        self._orgTreeWidget.itemClicked.connect(self.__treeItemClicked)
        self._userTableWidget.cellClicked.connect(self.__cellClicked)
        self._btnAdd.clicked.connect(self.__btnAddClicked)
        self._btnRemove.clicked.connect(self.__btnRemoveClicked)
        # self._textBoxCodeID.textEdited.connect(self.__textEdited)
        # self._textBoxCodeName.textEdited.connect(self.__textEdited)
        # self._btnAdd.clicked.connect(self.moveItem)
        # self.btnUp.clicked.connect(self.moveItem)
        # self.btnDown.clicked.connect(self.moveItem)

    # def moveItem(self):
    #     # sender = self.sender()
    #     #
    #     # if self.btnUp == sender:
    #     #     source = self.lecture_list
    #     #     target = self.wishlist
    #     # else:
    #     #     source = self.wishlist
    #     #     target = self.lecture_list
    #     source = self._orgTreeWidget
    #     target = self._orgTreeWidget
    #     # idx = self._orgTreeWidget.currentIndex()
    #     # item = self._orgTreeWidget.currentIndex().row()
    #     # self._orgTreeWidget.rowsInserted(item, idx+1)
    #     theItem = self._orgTreeWidget.currentItem()
    #     parentItem = theItem.parent()
    #     QTreeWidgetItem.p
    #     item = QTreeWidget.invisibleRootItem(source).takeChild(source.currentIndex().row())
    #     QTreeWidget.invisibleRootItem(target).addChild(item)
    #     QTreeWidget.invisibleRootItem(target).insert
    #     QTreeWidgetItem.insertChild()

    # 다른 객체간으로 아이템 이동
    # def move_Item(self):
    #     sender = self.sender()
    #
    #     if self.btnUp == sender:
    #         source = self.lecture_list
    #         target = self.wishlist
    #     else:
    #         source = self.wishlist
    #         target = self.lecture_list
    #
    #     item = QTreeWidget.invisibleRootItem(source).takeChild(source.currentIndex().row())
    #     QTreeWidget.invisibleRootItem(target).addChild(item)

    # def closeEvent(self, event):
    #     # if self.parent != None:
    #     #     self.parent.windowChildClosed(self)
    #     event.accept()
    #     print("Closed")

    def __textEdited(self):
        pass

    def showOrgUsers(self):
        self._userTableWidget.clearContents()
        self._userTableWidget.setRowCount(len(self._currentCodeList))
        for i, row in enumerate(self._currentCodeList):
            org = QTableWidgetItem(row[0])
            userName = QTableWidgetItem(row[1])
            region = QTableWidgetItem(row[2])
            org.setTextAlignment(Qt.AlignVCenter)
            userName.setTextAlignment(Qt.AlignVCenter)
            region.setTextAlignment(Qt.AlignVCenter)
            # codeID.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)     # not editable
            # codeName.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)   # not editable
            chkCellWidget = QWidget()
            chkLayout = QHBoxLayout(chkCellWidget)
            chkLayout.addWidget(QCheckBox())
            chkLayout.setAlignment(Qt.AlignCenter)
            chkLayout.setContentsMargins(0,0,0,0)
            chkCellWidget.setLayout(chkLayout)
            self._userTableWidget.setCellWidget(i, 0, chkCellWidget)
            self._userTableWidget.setItem(i, 1, org)
            self._userTableWidget.setItem(i, 2, userName)
            self._userTableWidget.setItem(i, 3, region)

        self._userTableWidget.resizeRowsToContents()
        self._userTableWidget.resizeColumnToContents(0)     # to resize CheckBox column
        # self._userTableWidget.resizeColumnsToContents()

    def showOrgData(self):
        selectedItem = self._orgTreeWidget.currentItem()
        orgText = selectedItem.text(0)
        middleOrgItem = selectedItem.parent()
        middleOrgText = middleOrgItem.text(0)
        superOrgItem = middleOrgItem.parent()
        superOrgText = superOrgItem.text(0)
        companyItem = superOrgItem.parent()
        companyText = companyItem.text(0)
        self._textBoxCompany.setText(companyText)
        self._textBoxSuperOrg.setText(superOrgText)
        self._textBoxMiddleOrg.setText(middleOrgText)
        self._textBoxOrg.setText(orgText)

    def __treeItemClicked(self):
        self._orgLabel.setText('조직 구성원')
        self._countLabel.setText('(인원수)')
        self._textBoxCompany.clear()
        self._textBoxSuperOrg.clear()
        self._textBoxMiddleOrg.clear()
        self._textBoxOrg.clear()
        self._userTableWidget.clearContents()
        self._userTableWidget.setRowCount(0)
        selectedItem = self._orgTreeWidget.currentItem()
        # Check if last node having 'org'
        if selectedItem.text(Constants.TREE_COL_0) != selectedItem.text(Constants.TREE_COL_3):
            return

        middleOrgItem = selectedItem.parent()
        superOrgItem = middleOrgItem.parent()
        companyItem = superOrgItem.parent()
        companyText = companyItem.text(Constants.TREE_COL_0)
        orgText = selectedItem.text(Constants.TREE_COL_0)
        self._storedUserList = []
        for c in self._userAllList:
            # 'company' needs to be considered later
            # if c.get('company') == companyText and c.get('org') == orgText:
            if c.get('org') == orgText:
                item = [c.get('org'), c.get('user_name'), c.get('region')]
                self._storedUserList.append(item)

        self._currentCodeList = self._storedUserList.copy()
        self.showOrgUsers()
        self._orgLabel.setText(companyText + ' ' + orgText)
        self._countLabel.setText('(' + str(len(self._currentCodeList)) + ')명')
        self.showOrgData()

        self._codeItemIndex = __NOT_SELECTED__

    def reInitData(self):
        self.getTable_Org()
        self.composeOrgTree()
        self._textBoxCompany.clear()
        self._textBoxSuperOrg.clear()
        self._textBoxMiddleOrg.clear()
        self._textBoxOrg.clear()
        self._orgLabel.setText('조직 구성원')
        self._countLabel.setText('(인원수)')
        self._userTableWidget.clearContents()
        self._userTableWidget.setRowCount(0)

    # primary key 중복 체크 추가 (company, org)
    # 텍스트 내용이 비어 있는 경우에 동작하지 않도록 함.
    def __btnAddClicked(self):
        orgText = self._textBoxOrg.text()
        middleOrgText = self._textBoxMiddleOrg.text()   # super org 1
        superOrgText = self._textBoxSuperOrg.text()     # super org 2
        companyText = self._textBoxCompany.text()
        dbData = [companyText, orgText, middleOrgText, superOrgText]
        try:
            db_connection = mysqlHandler.dbConnect()
            cursor = mysqlHandler.dbCursor(db_connection)
            sql = 'insert into myasset.it_user_org(company, org, super_org1, super_org2) ' \
                  'values (%s, %s, %s, %s)'
            cursor.execute(sql, dbData)
            db_connection.commit()
            __DB_TRANSACTION_FLAG__ = True
        except Exception as e:
            __DB_TRANSACTION_FLAG__ = False
            print(e)
        finally:
            db_connection.close()
        if __DB_TRANSACTION_FLAG__ == True:
            self.reInitData()

    # 해당 조직에 속한 사용자가 있는지 DB에서 체크 -> 있다면 삭제 못하게 막아야 함.
    # 텍스트 내용이 비어 있는 경우에 동작하지 않도록 함.
    def __btnRemoveClicked(self):
        __DB_TRANSACTION_FLAG__ = None
        orgText = self._textBoxOrg.text()
        middleOrgText = self._textBoxMiddleOrg.text()   # super org 1
        superOrgText = self._textBoxSuperOrg.text()     # super org 2
        companyText = self._textBoxCompany.text()
        dbData = [companyText, orgText, middleOrgText, superOrgText]
        try:
            db_connection = mysqlHandler.dbConnect()
            cursor = mysqlHandler.dbCursor(db_connection)
            sql = 'delete from myasset.it_user_org ' \
                  ' where company = %s and org = %s and super_org1 = %s and super_org2 = %s'
            cursor.execute(sql, dbData)
            db_connection.commit()
            __DB_TRANSACTION_FLAG__ = True
        except Exception as e:
            __DB_TRANSACTION_FLAG__ = False
            print(e)
        finally:
            db_connection.close()
        if __DB_TRANSACTION_FLAG__ == True:
            self.reInitData()

    @pyqtSlot(int, int)
    def __cellClicked(self, row, col):
        pass

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle('Fusion')
    w = OrgWindow(None, '대성에너지')
    w.show()
    sys.exit(app.exec_())


from utils.http_client import simple_client
from bs4 import BeautifulSoup

class Simpleisbetterthancomplex:
    SITE_URL = 'https://simpleisbetterthancomplex.com'

    @classmethod
    def get_all_news(cls, page=1):
        if page == 1:
            URL = cls.SITE_URL + '/feed.xml' # RSS
        else:
            return []

        response = simple_client.get(URL)
        content_bs = BeautifulSoup(response.content, 'html.parser')

        all_news = content_bs.find_all('item')

        news_values = []
        for news_bs in all_news:
            news = {}

            news['name_eng'] = news_bs.find('title').getText()
            news['url'] = news_bs.find('guid').getText()
            news['tags'] = ', '.join([tag.getText() for tag in news_bs.find_all('category')])

            big_description = news_bs.find('description').getText()
            description_bs = BeautifulSoup(big_description, 'html.parser')
            news['description'] = description_bs.find('p').getText()

            # regular_search_date = '\d+\.\d+\.\d+'
            # news['dc'] = datetime.datetime.strptime(date_str, '%d.%m.%Y')

            news_values.append(news)

        return news_values


# Script to do the training distributed across all GPUs using DDP package of Pytorch.
import argparse
import glob
import json
import math
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data

from datetime import datetime
from os.path import join, exists
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

# My own libraries
import models as cm
from dataset import DynamicsDataset
from train_cpn import test, train

def create_roots(train_dir):
    roots = glob.glob('data/*') # Get the roots to get data from - TODO: change this - get all of the roots
    test_root_num = math.floor(len(roots) * args.test_ratio)

    # Find random indices to get from roots for test
    test_rand_indices = np.random.choice(np.arange(len(roots)), test_root_num)
    test_roots = []
    train_roots = []
    for i, root in enumerate(roots):
        if i in test_rand_indices:
            test_roots.append(root)
        else:
            train_roots.append(root)

    print('test_roots: {}'.format(test_roots))
    print('train_roots: {}'.format(train_roots))

    if not os.path.isfile(join(train_dir, 'test_roots.json')): # test roots will be the same for all the roots (when all the data is used)
        # Serializing json
        test_roots_json = {
            'test_roots': test_roots
        }
        json_object = json.dumps(test_roots_json, indent = 4)

        # Writing to sample.json
        with open(join(train_dir, 'test_roots.json'), "w") as outfile:
            outfile.write(json_object)

    return train_roots, test_roots

def get_dataloaders(train_roots, test_roots, args): # world_size: Number of possible GPUs

    test_dset = DynamicsDataset(roots=test_roots,
                                sec_interval=args.sec_interval,
                                poly_max_deg=args.poly_max_deg)
    test_sampler = data.DistributedSampler(test_dset)
    test_loader = data.DataLoader(test_dset,
                                  batch_size=args.batch_size,
                                  shuffle=False, # This should be false when DistributedSampler is used, here if true it will shuffle the subsamples
                                  num_workers=4,
                                  sampler=test_sampler)

    train_dset = DynamicsDataset(roots=train_roots,
                                 sec_interval=args.sec_interval,
                                 poly_max_deg=args.poly_max_deg)
    train_sampler = data.DistributedSampler(train_dset)
    train_loader = data.DataLoader(train_dset,
                                  batch_size=args.batch_size,
                                  shuffle=False, # This should be false when DistributedSampler is used, here if true it will shuffle the subsamples
                                  num_workers=4,
                                  sampler=train_sampler)

    return train_loader, test_loader

# General method to start multi GPUd training
def start_process(rank, world_size, args, train_dir, train_roots, test_roots):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    dist.barrier() # Wait for all of the processes to start

    # Get dataloader for this rank
    print('rank: {}'.format(rank))

    if rank == 0:
        writer = SummaryWriter(train_dir)

    # print('len(train_loader): {}'.format(len(train_loader))) # This will be divided by the number of GPUs used - originally number of batches is around 1900, with multi processing it's lowered to 500 for each GPU
    train_loader, test_loader = get_dataloaders(train_roots, test_roots, args)
    print('len(train_loader): {}, len(train_loader.dataset): {}'.format(len(train_loader), len(train_loader.dataset)))
    # Create the models with the given rank and device
    obs_dim = (3, 480, 480)
    # action_dim = 2*(args.poly_max_deg+1)
    action_dim = 2

    # Create CUDA device with the given rank
    device = torch.device(f'cuda:{rank}')

    # Initialize models
    encoder = cm.Encoder(args.z_dim, obs_dim[0]).to(device)
    trans = cm.Transition(args.z_dim, action_dim).to(device)
    print('encoder: {}\ntrans: {}'.format(encoder, trans))
    parameters = list(encoder.parameters()) + list(trans.parameters())
    optimizer = optim.Adam(parameters, lr=args.lr, weight_decay=args.weight_decay)

    # Wrap the models in DDP wrapper
    encoder = DDP(encoder, device_ids=[rank], output_device=rank)
    trans = DDP(trans, device_ids=[rank], output_device=rank)

    # Train losses and model checkpoints are only saved when the rank == 0 since the model parameters
    # and gradients are distributed between GPUs and copy of the model is always preserved in all GPUs
    for epoch in range(args.epochs):
        train_loader.sampler.set_epoch(epoch)
        # Train
        print('rank: {} entered -before training- barier'.format(rank))
        dist.barrier() # All GPUs should start training at the same time
        train_loss = train(train_loader, encoder, trans, optimizer, epoch, device, args.batch_size, rank)

        print('rank: {} entered -after training- barier'.format(rank))
        dist.barrier() # Wait until all the processes reached here

        if rank == 0: # Add the train loss to the tensorboard writer
            for i in range(len(train_loss)-1):
                writer.add_scalar('Train Loss', train_loss[i], epoch * args.batch_size + i)

        if rank == 0 and epoch % args.test_interval == 0:
            print('rank: {}, testing epoch: {}'.format(rank, epoch))
            test_loss, sim = test(test_loader, encoder, trans, epoch, device, rank)
            for i in range(len(test_loss)-1):
                writer.add_scalar('Test Loss', test_loss[i], (epoch / args.test_interval) * args.batch_size + i)
                writer.add_scalar('Euclidean Dist', sim[0], (epoch / args.test_interval) * args.batch_size + i)
                writer.add_scalar('Cosine Similarity', sim[1], (epoch / args.test_interval) * args.batch_size + i)
                writer.add_scalar('Dot Product', sim[2], (epoch / args.test_interval) * args.batch_size + i)

        if rank == 0 and epoch % args.model_save_interval == 0:
            checkpoint = {
                'encoder': encoder,
                'trans': trans,
                'optimizer': optimizer,
            }
            torch.save(checkpoint, join(train_dir, f'checkpoint_{epoch}.pt'), _use_new_zipfile_serialization=False)
            # Save models and optimizers' state_dict so that it would be easier to load in the car
            torch.save(encoder.state_dict(), join(train_dir, f'encoder_{epoch}.pt'), _use_new_zipfile_serialization=False)
            torch.save(trans.state_dict(), join(train_dir, f'trans_{epoch}.pt'), _use_new_zipfile_serialization=False)
            torch.save(optimizer.state_dict(), join(train_dir, f'optimizer_{epoch}.pt'), _use_new_zipfile_serialization=False)

    if rank == 0:
        writer.close()

def main():
    # Create the directory to save the outputs
    if args.train_dir == 'train': # We should create a train_dir with timestamp if there is no additional value given
        now = datetime.now()
        time_str = now.strftime('%d%m%Y_%H%M%S')
        train_dir = join(args.train_out, 'train_{}'.format(time_str))

    else:
        train_dir = join(args.train_out, args.train_dir)
    os.mkdir(train_dir)
    train_roots, test_roots = create_roots(train_dir)

    n_gpus = torch.cuda.device_count()
    print('args: {}'.format(args))
    print('n_gpus = world_size: {}'.format(n_gpus))
    world_size = 4
    mp.spawn(start_process,
        args=(world_size, args, train_dir, train_roots, test_roots),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29503"

    # Parse arguments
    parser = argparse.ArgumentParser()

    # Dataset Parameters
    parser.add_argument('--root', type=str, default='data/28012018_111425')
    parser.add_argument('--train_out', type=str, default='out')
    parser.add_argument('--train_dir', type=str, default='train')

    # Learning Parameters
    parser.add_argument('--lr', type=float, default=5e-4)
    parser.add_argument('--weight_decay', type=float, default=1e-5)
    parser.add_argument('--epochs', type=int, default=51)
    parser.add_argument('--test_interval', type=int, default=2)
    parser.add_argument('--test_ratio', type=float, default=0.25)
    parser.add_argument('--model_save_interval', type=int, default=5)
    parser.add_argument('--model_load', type=bool, default=False)
    parser.add_argument('--model_load_file', type=str, default='checkpoint_100.pt')

    # InfoNCE Parameters
    # Negative Samples = Batch Size
    parser.add_argument('--batch_size', type=int, default=256) # TODO: Change this
    # parser.add_argument('--action_dim', type=int, default=12)
    parser.add_argument('--poly_max_deg', type=int, default=10)
    parser.add_argument('--sec_interval', type=int, default=0.5)
    parser.add_argument('--z_dim', type=int, default=64)
    parser.add_argument('--name', type=str, default='arya')
    parser.add_argument('--seed', type=int, default=17)

    args = parser.parse_args()

    main()


import pandas as pd
import nni
import pandas_ta as pd_ta
from hep_ml import reweight
import matplotlib.pyplot as plt
from hep_ml.metrics_utils import ks_2samp_weighted
from fetch_alpaca_data import fetch_alpaca_csv
from pandas import DataFrame
from sklearn.metrics import balanced_accuracy_score, accuracy_score, roc_auc_score, f1_score
import warnings
from sklearn.ensemble import RandomForestClassifier
import seaborn as sns
from tqdm import tqdm
from sklearn.model_selection import train_test_split, cross_val_score
from xgboost import XGBClassifier
import pywt
import pandas_market_calendars as mcal
from collections import Counter
import numpy as np
from time import perf_counter
import yfinance as yf
from scipy.stats import ttest_ind
from sklearn.preprocessing import quantile_transform

warnings.filterwarnings('ignore')

# set random seed on numpy
np.random.RandomState(42)
np.random.seed(42)

# performance csv headers
header = ['test_acc', 'live_acc']

params: dict = {
    'years': 3,
    'wavelet': 'bior1.1',
    'window': 70
}
optimized_params = nni.get_next_parameter()
params.update(optimized_params)

# create script parameters (dynamic)
years = params['years']
trend = params['window']
wavelet = params['wavelet']
shuffle_split: bool = True

# find yesterday closing price
ticker = yf.Ticker('SPY')
yesterday_price = ticker.history(interval='1d')['Close'][-2]

# mean absolute error evaluator for wavelet transform
def madev(d, axis=None):
    """ Mean absolute deviation of a signal """
    return np.mean(np.absolute(d - np.mean(d, axis)), axis)

# de-noise original signal with selected wavelet function
def wavelet_denoising(x, wavelet='db4', level=1):
    coeff = pywt.wavedec(x, wavelet, mode="per")
    sigma = (1 / 0.6745) * madev(coeff[-level])
    uthresh = sigma * np.sqrt(2 * np.log(len(x)))
    coeff[1:] = (pywt.threshold(i, value=uthresh, mode='hard') for i in coeff[1:])
    return pywt.waverec(coeff, wavelet, mode='per')

# smooth initial model features in order to extract longer trend pattern
def _wavelet_smooth(data, wavelet):
    """wavelet transformation applied to curve"""
    processed_arr = []
    for col in ['open', 'high', 'low', 'close', 'volume']:
        processed_arr.append(wavelet_denoising(data[col], wavelet=wavelet, level=1))
    return processed_arr

def _get_indicator_data(data_raw):
    # technical analysis indicators
    data_raw['rsi'] = pd_ta.rsi(data_raw['close'])
    data_raw['macd'] = pd_ta.macd(data_raw['close']).iloc[:, -1]
    data_raw['willr'] = pd_ta.willr(data_raw['high'], data_raw['low'], data_raw['close'])
    data_raw['obv'] = pd_ta.obv(data_raw['close'], data_raw['volume'])
    data_raw['proc'] = pd_ta.roc(data_raw['close'])
    data_raw['stoch_k'] = pd_ta.stoch(data_raw['high'], data_raw['low'], data_raw['close']).iloc[:, 0]

    # Remove columns that won't be used as features
    del (data_raw['open'])
    del (data_raw['high'])
    del (data_raw['low'])
    del (data_raw['volume'])
    try:
        del (data_raw['Adj Close'])
    except KeyError:
        pass
    return data_raw

def _produce_prediction_gaps(data, transformed_data, window, from_hardcoded=False, year=None) -> DataFrame:
    """create predicted target values"""
    try:
        if from_hardcoded:
            return pd.read_csv(f'targets/{year}y_{window}w.csv')
    except FileNotFoundError:
        pass
    # create containers and calendars
    nyse = mcal.get_calendar('NYSE')
    calendar = nyse.schedule(start_date=data['day'].iloc[0], end_date=data['day'].iloc[-1])
    pred_days = data['day'].unique()[-(window // 14):]
    dropped_idx: list = []
    missing_days = []
    missing_counter = 0
    value_counts = data['day'].value_counts()
    for key, value in value_counts.items():
        if value != 14:
            missing_days.append(key)
            missing_counter += 1
            print(value, key)
    data = data.loc[~data['day'].isin(missing_days)]
    transformed_data = transformed_data.loc[~transformed_data['day'].isin(missing_days)]

    # custom mapping for trading window at given intraday time
    map: dict = {
        '09:30:00': 13, '10:00:00': 12, '10:30:00': 11, '11:00:00': 10, '11:30:00': 9, "12:00:00": 8,
        '12:30:00': 7, '13:00:00': 6, '13:30:00': 5, '14:00:00': 4, '14:30:00': 3, '15:00:00': 2,
        '15:30:00': 1, '16:00:00': 0
    }
    # create containers
    times, closes, days = data['time'], data['close'], data['day']
    targets = []
    end_dates = []
    compare_closes = []

    for start_idx in tqdm(range(data.shape[0])):
        # get prices to compare for target variable
        current_close = closes.iloc[start_idx]

        if days.iloc[start_idx] in pred_days:
            end_dates.append(np.nan)
            compare_closes.append(np.nan)
            targets.append(np.nan)
            continue

        # get price comparisons
        try:
            compare_close = closes.iloc[start_idx + window + map[times.iloc[start_idx]]]
            end_dates.append(days.iloc[start_idx + window + map[times.iloc[start_idx]]])
            compare_closes.append(compare_close)
        except IndexError:
            compare_close = np.nan
            end_dates.append(np.nan)
            compare_closes.append(np.nan)  # only at the -window span at the end of the dataframe

        # check market day difference
        if not np.isnan(compare_close):
            diff = calendar.loc[days.iloc[start_idx]:days.iloc[start_idx + window + map[times.iloc[start_idx]]]].shape[
                       0] - 1
            if diff != window // 14 or '16:00:00' not in times.iloc[
                start_idx + window + map[times.iloc[start_idx]]] or compare_close == current_close:
                dropped_idx.append(start_idx)
                targets.append('GAP')
                continue
            else:
                if current_close < compare_close:
                    targets.append(1)
                else:
                    targets.append(0)
        else:
            # check if larger or smaller
            targets.append('GAP')
    print(data.shape[0], len(end_dates))
    data['end_timestamps'] = end_dates
    data['compare_closes'] = compare_closes
    data['pred'] = targets

    # remove dropped indices from the dropped_idx array
    data = data.loc[~data.index.isin(dropped_idx)]
    transformed_data = transformed_data.loc[~transformed_data.index.isin(dropped_idx)]
    # reset indices
    data.reset_index(inplace=True, drop=True)
    transformed_data.reset_index(inplace=True, drop=True)
    # count number of GAP target variables
    same_price_idx = data[data['pred'] == 'GAP'].index.values
    data = data.loc[~data.index.isin(same_price_idx)]
    transformed_data = transformed_data.loc[~transformed_data.index.isin(same_price_idx)]
    # insert original target variables into transformed dataframe
    transformed_data['pred'] = data['pred']
    transformed_data['end_timestamps'] = data['end_timestamps']
    transformed_data['compare_closes'] = data['compare_closes']
    return transformed_data

def _produce_prediction(data, window):
    """create predicted target values"""

    # custom mapping for trading window at given intraday time
    map: dict = {
        '09:30:00': 13, '10:00:00': 12, '10:30:00': 11, '11:00:00': 10, '11:30:00': 9, "12:00:00": 8,
        '12:30:00': 7, '13:00:00': 6, '13:30:00': 5, '14:00:00': 4, '14:30:00': 3, '15:00:00': 2,
        '15:30:00': 1, '16:00:00': 0
    }
    # create
    times, closes = data['time'], data['close']
    targets = np.array([])

    for start_idx in range(data.shape[0]):
        # get prices to compare for target variable
        current_close = closes.iloc[start_idx]
        try:
            compare_close = closes.iloc[start_idx + window + map[times.iloc[start_idx]]]
        except IndexError:
            compare_close = np.nan

        # check if larger or smaller
        if np.isnan(compare_close):
            targets = np.append(targets, np.nan)
        else:
            if current_close < compare_close:
                targets = np.append(targets, 1)
            else:
                targets = np.append(targets, 0)
    data['pred'] = targets
    return data['pred']

def test_boundaries_whole(df, ranges, training_data_boundaries):
    boundary_flag: bool = True
    for var, range in zip(['rsi', 'macd', 'willr', 'obv', 'proc', 'stoch_k'], ranges):
        if training_data_boundaries(range[0], range[1], df[var]):
            boundary_flag = False
    return boundary_flag

def live_accuracy(trends, strikes, comp_price):
    """check the live trades accuracy"""

    # create output container
    container: list = []

    for i, (trend, strike) in enumerate(zip(trends, strikes)):
        if i <= 13:
            if trend == 1:
                if strike < comp_price:
                    container.append(True)
                else:
                    container.append(False)
            else:
                if strike > comp_price:
                    container.append(True)
                else:
                    container.append(False)
        else:
            container.append('-')
    return container

features = ['rsi', 'macd', 'willr', 'obv', 'proc', 'stoch_k']
def draw_distributions(original, target, new_original_weights, plot=False):
    ks_vals: list = []
    for id, column in enumerate(features, 1):
        if plot:
            plt.figure(figsize=[15, 7])
            plt.subplot(2, 3, id)
            sns.kdeplot(data=original, x=column, color='red', weights=new_original_weights)
            sns.kdeplot(data=target, x=column, color='blue')
            plt.title(column)
            plt.show()
        ks_vals.append(ks_2samp_weighted(original[column], target[column],
                                         weights1=new_original_weights, weights2=np.ones(len(target), dtype=float)))
    return sum(ks_vals)/len(ks_vals)

ticker = yf.Ticker('SPY')
yesterday_price_backtest = ticker.history(interval='1d', start="2020-01-01", end="2022-01-21")['Close']
nyse = mcal.get_calendar('NYSE')

def evaluate_preds(strike, close, trend):
  """get actual price movements to live data"""
  if strike < close and trend == 1:
    return True
  elif strike > close and trend == 0:
    return True
  else:
    return False

from datetime import timedelta
from datetime import datetime
def get_n_market_days_ahead(n, start_date):
  """get the next market day in an n-day interval"""
  start_date = str(start_date)
  start_date = datetime.strptime(start_date, "%Y-%m-%d")
  end_date = start_date + timedelta(days=15)

  early = nyse.schedule(start_date=start_date, end_date=end_date)
  actual_end = str(early.iloc[:n+1].iloc[-1]['market_open']).split(' ')[0]
  return yesterday_price_backtest[yesterday_price_backtest.index == str(actual_end)].iloc[0]

def backtest_func(live_data, model, threshold):
  """backtesting function"""
  live_data['trend'] = model.predict(live_data[features])
  live_data['day'] = live_data['start_date'].apply(lambda x: str(x).split(' ')[0])
  live_data['close_price'] = live_data['day'].apply(lambda x: get_n_market_days_ahead(params['window']//14, str(x)))

  live_data['actual'] = live_data.apply(lambda x: evaluate_preds(x.strikes, x.close_price, x.trend), axis=1).astype(int)
  live_data.drop('day', axis=1, inplace=True)
  print(live_data[live_data['prob'] > threshold]['actual'].value_counts())

# backtesting mode turn on/off
backtest: bool = True

def create_model(years=params['years'], trend=params['window'], wavelet=params['wavelet'], facade=False):
    # get data from backend
    data = fetch_alpaca_csv(years)
    data.rename(columns={col: col.lower() for col in data.columns}, inplace=True)
    copy_data = data.copy()

    print(data.shape)  # replace with wavelet transformed version
    data['open'], data['high'], data['low'], data['close'], data['volume'] = _wavelet_smooth(data, wavelet)

    # create indicator features
    data = _get_indicator_data(data)
    copy_data = _get_indicator_data(copy_data)
    data.drop('unnamed: 0', axis=1, inplace=True)
    copy_data.drop('unnamed: 0', axis=1, inplace=True)
    if data.isna().sum().sum() != copy_data.isna().sum().sum():
        print('NaN Error')
        return
    data.dropna(inplace=True)
    copy_data.dropna(inplace=True)

    # remove nan values and reset indices on both dataframes
    data.reset_index(inplace=True)
    copy_data.reset_index(inplace=True)
    data.drop(['index'], axis=1, inplace=True)
    copy_data.drop(['index'], axis=1, inplace=True)
    start = perf_counter()
    data['pred'] = _produce_prediction(copy_data, window=trend)
    print("Target Variable Creation Runtime", perf_counter() - start)

    # create live data samples from equivalent training distribution
    live_data = data[-trend:].drop('pred', axis=1)[['rsi', 'macd', 'willr', 'obv', 'proc', 'stoch_k']]
    strike_prices = copy_data[-trend:]['close']
    live_dates = data[-trend:]['timestamp']

    # resolve formatting and nan issues
    del (data['close'])
    data.dropna(inplace=True)
    print(data.shape, 'tracking data shape 4')
    #data = data.sample(frac=1).reset_index(drop=True)
    features = ['rsi', 'macd', 'willr', 'obv', 'proc', 'stoch_k']
    target = 'pred'

    # train-test data splits
    test_ratio = 0.30
    #X_train, X_test, y_train, y_test = train_test_split(data[features], data[target],
    #                                                    test_size=test_ratio, random_state=42)
    X_train, X_test = data[features][:int(0.8*data.shape[0])], data[features][int(0.8*data.shape[0]):]
    y_train, y_test = data[target][:int(0.8*data.shape[0])], data[target][int(0.8*data.shape[0]):]

    # transform distributions into the normal distribution to standardize training-testing-live samples
    live_data = live_data[features]
    transform: bool = False
    if transform:
        output_dist = 'uniform'
        train_trans = quantile_transform(X_train, output_distribution=output_dist, n_quantiles=len(X_train))
        test_trans = quantile_transform(X_test, output_distribution=output_dist, n_quantiles=len(X_train))
        live_trans = quantile_transform(live_data, output_distribution=output_dist, n_quantiles=len(X_train))
        for i, col in enumerate(X_train.columns):
            X_train[col] = train_trans[:, i]
            X_test[col] = test_trans[:, i]
            live_data[col] = live_trans[:, i]

    reweight_check = False
    opt_weights = np.ones(len(X_train))
    if reweight_check and not transform:
        opt_weights, distance = None, 10000
        for estimator in tqdm([150, 200, 250]):
            for lr in [0.05, 0.1]:
                for depth in [3, 4]:
                    for leaf in [1000, 2000, 3000]:
                        reweighter = reweight.GBReweighter(n_estimators=estimator, learning_rate=lr, max_depth=depth,
                                                           min_samples_leaf=leaf)
                        reweighter.fit(X_train, live_data[features])
                        gb_weights_test = reweighter.predict_weights(X_train)
                        dist_distance = draw_distributions(X_train, live_data, gb_weights_test)
                        if dist_distance < distance:
                            print(dist_distance, distance)
                            opt_weights = gb_weights_test
                            distance = dist_distance

    print(draw_distributions(X_train, live_data, opt_weights, plot=False)) ### DISABLE PLOTTING WHEN TUNING
    ratio = Counter(y_train)
    scale_pos_weight = max(ratio.values()) / min(ratio.values())
    model = XGBClassifier(scale_pos_weight=scale_pos_weight).fit(X_train, y_train, sample_weight=opt_weights)
    #model = RandomForestClassifier(class_weight='balanced').fit(X_train, y_train)

    stat, p_value = ttest_ind(X_train, live_data)
    print("T-test P-value: ", p_value)

    # determine accuracy and append to results
    rf_prediction = model.predict(X_test)
    print("Raw Accuracy: ", accuracy_score(list(y_test), list(rf_prediction)))
    print("ROC AUC: ", roc_auc_score(y_test, rf_prediction))
    print('F1 Score: ', f1_score(y_test, rf_prediction, average='macro'))
    rf_accuracy = balanced_accuracy_score(list(y_test), list(rf_prediction))
    print('Test Accuracy: ', rf_accuracy)
    print(years, trend, wavelet)

    # make live predictions
    live_predictions = model.predict(live_data)
    live_prob_predictions = model.predict_proba(live_data)
    live_data['pred'] = live_predictions
    live_data['strikes'] = strike_prices
    live_data['prob'] = [i[np.argmax(i)] for i in live_prob_predictions]
    live_data['date'] = live_dates
    live_data['window'] = [trend // 14 for _ in range(live_data.shape[0])]
    live_data['day'] = live_data['date'].apply(lambda x: str(x).split(' ')[0])
    live_data = live_data[live_data['day'].isin(live_data['day'].unique()[-3:])].drop('day', axis=1)
    live_data.rename(columns={'pred': 'trend', 'date': 'start_date'}, inplace=True)

    # measure accuracy across different confidence filters
    correct, incorrect, skipped = 0, 0, 0
    threshold = 0.9
    try:
        for pred, prob, actual in zip(rf_prediction, model.predict_proba(X_test), y_test):
            if prob[int(pred)] >= threshold:
                if pred == actual:
                    correct += 1
                else:
                    incorrect += 1
            else:
                skipped += 1
        print((f"Acc: {correct / (correct + incorrect)}"), f'Skip Ratio: {skipped / len(y_test)}')
    except ZeroDivisionError:
        pass

    # create a backtest
    if backtest:
        backtest_func(live_data, model, 0.9)

    if not facade:
        live_data.to_csv('live_data_backtest.csv')
        print(live_data[['window', 'trend', 'start_date', 'prob', 'strikes']][live_data['prob'] > 0.9])

    # report the NNI console
    nni.report_final_result(rf_accuracy)
    #nni.report_intermediate_result(cross_val_acc)

    if not facade:
        return model, rf_accuracy, live_data[['trend', 'prob', 'strikes', 'start_date']], live_data[features], \
               X_train, X_test, y_train, y_test, data
    else:
        return model, accuracy_score(list(y_test), list(rf_prediction)), \
               live_data[['window', 'trend', 'start_date', 'prob', 'strikes']][live_data['prob'] > 0.9]

create_model()
#create_model(years=3, trend=42, wavelet='coif2')


__version__ = '0.5.3'
__author__ = 'Qing Yu <yuq@sustech.edu.cn>'

# module level doc-string
__doc__ = """
TransBigData - A Python package develop for transportation spatio-temporal big
data processing, analysis and visualization.
=====================================================================

`TransBigData` is a Python package developed for transportation spatio-temporal
big data processing, analysis and visualization. `TransBigData` provides fast
and concise methods for processing common transportation spatio-temporal big
data such as Taxi GPS data, bicycle sharing data and bus GPS data.
`TransBigData` provides a variety of processing methods for each stage of
transportation spatio-temporal big data analysis. The code with `TransBigData`
is clean, efficient, flexible, and easy to use, allowing complex data tasks to
be achieved withconcise code.

For some specific types of data, `TransBigData` also provides targeted tools
for specific needs, such as extraction of Origin and Destination(OD) of taxi
trips from taxi GPS data and identification of arrival and departure
information from bus GPS data. The latest stable release of the software can
be installed via pip and full documentation can be found at
https://transbigdata.readthedocs.io/en/latest/.

Main Features
-------------

Currently, `TransBigData` mainly provides the following methods:

  - Data Quality: Provides methods to quickly obtain the general information of
  the dataset, including the data amount the time period and the sampling
  interval.
  - Data Preprocess: Provides methods to clean multiple types of data error.
  - Data Gridding: Provides methods to generate multiple types of geographic
  grids (Rectangular grids, Hexagonal grids) in the research area. Provides
  fast algorithms to map GPS data to the generated grids.
  - Data Aggregating: Provides methods to aggregate GPS data and OD data into
  geographic polygon.
  - Data Visualization: Built-in visualization capabilities leverage the
  visualization package keplergl to interactively visualize data on Jupyter
  notebook with simple code.
  - Trajectory Processing: Provides methods to process trajectory data,
  including generating trajectory linestring from GPS points, and trajectory
  densification, etc.
  - Basemap Loading: Provides methods to display Mapbox basemap on matplotlib
  figures
"""

from transbigdata.plotmap import (
    plot_map,
    plotscale,
    set_mapboxtoken,
    set_imgsavepath,
    read_imgsavepath,
    read_mapboxtoken
)
from transbigdata.coordinates import (
    gcj02tobd09,
    bd09togcj02,
    wgs84togcj02,
    gcj02towgs84,
    wgs84tobd09,
    bd09towgs84,
    bd09mctobd09,
    getdistance,
    transform_shape
)
from transbigdata.grids import (
    area_to_grid,
    area_to_params,
    GPS_to_grid,
    grid_to_centre,
    grid_to_polygon,
    grid_to_area,
    grid_to_params,
    grid_params_optimize,
    geohash_encode,
    geohash_decode,
    geohash_togrid
)
from transbigdata.gisprocess import (
    ckdnearest,
    ckdnearest_point,
    ckdnearest_line,
    splitline_with_length,
    merge_polygon,
    polyon_exterior
)
from transbigdata.odprocess import (
    odagg_grid,
    odagg_shape,
    tolinewitharrow
)
from transbigdata.preprocess import (
    clean_outofbounds,
    clean_outofshape,
    dataagg,
    id_reindex_disgap,
    id_reindex
)
from transbigdata.bikedata import (
    bikedata_to_od
)
from transbigdata.taxigps import (
    clean_taxi_status,
    taxigps_to_od,
    taxigps_traj_point
)
from transbigdata.mobilephonedata import (
    mobile_stay_duration,
    mobile_identify_home,
    mobile_identify_work,
    )
from transbigdata.traj import (
    traj_mapmatch,
    traj_clean_drift,
    traj_clean_redundant,
    traj_slice,
    traj_smooth,
    traj_segment,
    traj_densify,
    traj_sparsify,
    traj_stay_move,
    traj_to_linestring,
    traj_length
)

from transbigdata.utils import (
    dumpjson
)

from transbigdata.quality import (
    sample_duration,
    data_summary
)
from transbigdata.busgps import (
    busgps_arriveinfo,
    busgps_onewaytime
)
from transbigdata.crawler import (
    getadmin,
    getbusdata,
    get_isochrone_amap,
    get_isochrone_mapbox,
)
from transbigdata.metro import (
    split_subwayline,
    metro_network,
    get_path_traveltime,
    get_shortest_path,
    get_k_shortest_paths
)
from transbigdata.visualization import (
    visualization_trip,
    visualization_od,
    visualization_data
)

from transbigdata.activity import (
    entropy,
    entropy_rate,
    ellipse_params,
    ellipse_plot,
    plot_activity
)


import asyncio
import os
import warnings
import json
from types import SimpleNamespace
import jwt
from base64 import b64decode

from bson import ObjectId
from django.utils.decorators import sync_and_async_middleware
from django import http
from django.utils.cache import patch_vary_headers
from django.utils.deprecation import RemovedInDjango40Warning

from urllib.parse import urlparse

from base.aes import AES
from base.rf.response import Response
from base.rf import status

from .cors_deafults import *

AES_KEY = bytes(os.getenv('AES_KEY').encode('utf-8'))
JWT_PUBLIC = os.getenv('JWT_PUBLIC').replace('\\n', '\n').strip().replace('"', '')
JWT_ALGORITHM = os.getenv('JWT_ALGORITHM')

ACCESS_CONTROL_ALLOW_ORIGIN = 'Access-Control-Allow-Origin'
ACCESS_CONTROL_EXPOSE_HEADERS = 'Access-Control-Expose-Headers'
ACCESS_CONTROL_ALLOW_CREDENTIALS = 'Access-Control-Allow-Credentials'
ACCESS_CONTROL_ALLOW_HEADERS = 'Access-Control-Allow-Headers'
ACCESS_CONTROL_ALLOW_METHODS = 'Access-Control-Allow-Methods'
ACCESS_CONTROL_MAX_AGE = 'Access-Control-Max-Age'

def get_data(auth):
    if auth:
        auth = auth.strip()
        auth = auth.split(' ')[-1]
        jwt_token = auth
        payload = jwt.decode(jwt_token, JWT_PUBLIC,
                             algorithms=[JWT_ALGORITHM])
        data = b64decode(payload['data'])
        iv = b64decode(payload['identifier'])

        dt = AES(AES_KEY).decrypt_ctr(data, iv)
        dt = dt.replace("'", '"').replace("None", "null")
        data1 = json.loads(dt, object_hook=lambda e: SimpleNamespace(**e))
        return data1
    return None

def authorize(request):
    access_token = request.headers.get('Authorization', None)
    if access_token:
        try:
            user_data = get_data(access_token)
        except jwt.exceptions.InvalidSignatureError:
            return Response(messages=[{'message': "Invalid Signature Error"}],
                            status_code=401
                            , status=status.HTTP_401_UNAUTHORIZED)
        except jwt.exceptions.ExpiredSignatureError:
            return Response(messages=[{'message': 'Signature is Expired'}], status_code=401,
                            status=status.HTTP_401_UNAUTHORIZED)
        except Exception:
            return Response(messages=[{'message': 'Invalid Token'}], status_code=401,
                            status=status.HTTP_401_UNAUTHORIZED)
        request._user = user_data
        request._user.is_authenticated = True
        request._user.id = ObjectId(user_data._id)
    else:
        request._user = None

@sync_and_async_middleware
def auth_middleware(get_response):
    """
        check token exists and decode token and get data for authorization
    """
    if asyncio.iscoroutinefunction(get_response):
        async def middleware(request):
            resp = authorize(request)
            if isinstance(resp, Response):
                return resp
            response = await get_response(request)
            return response

    else:

        def middleware(request):
            resp = authorize(request)
            if isinstance(resp, Response):
                return resp
            response = get_response(request)
            return response

    return middleware

AuthMiddleware = auth_middleware

class CORSMiddleWare:
    sync_capable = True
    async_capable = True

    # RemovedInDjango40Warning: when the deprecation ends, replace with:
    #   def __init__(self, get_response):
    def __init__(self, get_response=None):
        self._get_response_none_deprecation(get_response)
        self.get_response = get_response
        self._async_check()
        super().__init__()

    def _async_check(self):
        """
        If get_response is a coroutine function, turns us into async mode so
        a thread is not consumed during a whole request.
        """
        if asyncio.iscoroutinefunction(self.get_response):
            # Mark the class as async-capable, but do the actual switch
            # inside __call__ to avoid swapping out dunder methods
            self._is_coroutine = asyncio.coroutines._is_coroutine

    def __call__(self, request):
        # Exit out to async mode, if needed
        if asyncio.iscoroutinefunction(self.get_response):
            return self.__acall__(request)
        response = None
        if hasattr(self, 'process_request'):
            response = self.process_request(request)
        response = response or self.get_response(request)
        if hasattr(self, 'process_response'):
            response = self.process_response(request, response)
        return response

    async def __acall__(self, request):
        """
        Async version of __call__ that is swapped in when an async request
        is running.
        """
        response = None
        if hasattr(self, 'async_process_request'):
            response = await self.async_process_request(request)
        response = response or await self.get_response(request)
        if hasattr(self, 'async_process_response'):
            response = await self.async_process_response(request, response)
        return response

    def _get_response_none_deprecation(self, get_response):
        if get_response is None:
            warnings.warn(
                'Passing None for the middleware get_response argument is '
                'deprecated.',
                RemovedInDjango40Warning, stacklevel=3,
            )

    def _https_referer_replace(self, request):
        """
        When https is enabled, django CSRF checking includes referer checking
        which breaks when using CORS. This function updates the HTTP_REFERER
        header to make sure it matches HTTP_HOST, provided that our cors logic
        succeeds
        """
        origin = request.META.get('HTTP_ORIGIN')

        if (request.is_secure() and origin and
                'ORIGINAL_HTTP_REFERER' not in request.META):
            url = urlparse(origin)
            if (not CORS_ORIGIN_ALLOW_ALL and
                    self.origin_not_found_in_white_lists(origin, url)):
                return

            try:
                http_referer = request.META['HTTP_REFERER']
                http_host = "https://%s/" % request.META['HTTP_HOST']
                request.META = request.META.copy()
                request.META['ORIGINAL_HTTP_REFERER'] = http_referer
                request.META['HTTP_REFERER'] = http_host
            except KeyError:
                pass

    def process_request(self, request):
        """
        If CORS preflight header, then create an
        empty body response (200 OK) and return it
        Django won't bother calling any other request
        view/exception middleware along with the requested view;
        it will call any response middlewares
        """
        if CORS_REPLACE_HTTPS_REFERER:
            self._https_referer_replace(request)

        if (request.method == 'OPTIONS' and
                "HTTP_ACCESS_CONTROL_REQUEST_METHOD" in request.META):
            response = http.HttpResponse()
            return response
        return None

    # def process_view(self, request, callback, callback_args, callback_kwargs):
    #     """
    #     Do the referer replacement here as well
    #     """
    #     if CORS_REPLACE_HTTPS_REFERER:
    #         self._https_referer_replace(request)
    #     return None

    def process_response(self, request, response):
        """
        Add the respective CORS headers
        """
        origin = request.META.get('HTTP_ORIGIN')
        if origin:
            # todo: check hostname from db instead
            url = urlparse(origin)

            if CORS_MODEL is not None:
                response[ACCESS_CONTROL_ALLOW_ORIGIN] = origin

            if (not CORS_ORIGIN_ALLOW_ALL and
                    self.origin_not_found_in_white_lists(origin, url)):
                return response

            if CORS_ORIGIN_ALLOW_ALL and not CORS_ALLOW_CREDENTIALS:
                response[ACCESS_CONTROL_ALLOW_ORIGIN] = "*"
            else:
                response[ACCESS_CONTROL_ALLOW_ORIGIN] = origin
                patch_vary_headers(response, ['Origin'])

            if len(CORS_EXPOSE_HEADERS):
                response[ACCESS_CONTROL_EXPOSE_HEADERS] = ', '.join(
                    CORS_EXPOSE_HEADERS)

            if CORS_ALLOW_CREDENTIALS:
                response[ACCESS_CONTROL_ALLOW_CREDENTIALS] = 'true'

            if request.method == 'OPTIONS':
                response[ACCESS_CONTROL_ALLOW_HEADERS] = ', '.join(
                    CORS_ALLOW_HEADERS)
                response[ACCESS_CONTROL_ALLOW_METHODS] = ', '.join(
                    CORS_ALLOW_METHODS)
                if CORS_PREFLIGHT_MAX_AGE:
                    response[ACCESS_CONTROL_MAX_AGE] = \
                        CORS_PREFLIGHT_MAX_AGE
        return response

    async def async_process_response(self, request, response):
        """
        Add the respective CORS headers
        """
        origin = request.META.get('HTTP_ORIGIN')
        if origin:
            # todo: check hostname from db instead
            response[ACCESS_CONTROL_ALLOW_ORIGIN] = origin

            if (not CORS_ORIGIN_ALLOW_ALL and
                    self.origin_not_found_in_white_lists(origin, origin)):
                return response

            if CORS_ORIGIN_ALLOW_ALL and not CORS_ALLOW_CREDENTIALS:
                response[ACCESS_CONTROL_ALLOW_ORIGIN] = "*"
            else:
                response[ACCESS_CONTROL_ALLOW_ORIGIN] = origin
                patch_vary_headers(response, ['Origin'])

            if len(CORS_EXPOSE_HEADERS):
                response[ACCESS_CONTROL_EXPOSE_HEADERS] = ', '.join(
                    CORS_EXPOSE_HEADERS)

            if CORS_ALLOW_CREDENTIALS:
                response[ACCESS_CONTROL_ALLOW_CREDENTIALS] = 'true'

            if request.method == 'OPTIONS':
                response[ACCESS_CONTROL_ALLOW_HEADERS] = ', '.join(
                    CORS_ALLOW_HEADERS)
                response[ACCESS_CONTROL_ALLOW_METHODS] = ', '.join(
                    CORS_ALLOW_METHODS)
                if CORS_PREFLIGHT_MAX_AGE:
                    response[ACCESS_CONTROL_MAX_AGE] = \
                        CORS_PREFLIGHT_MAX_AGE
        return response

    def origin_not_found_in_white_lists(self, origin, url):
        return url.netloc not in CORS_ORIGIN_WHITELIST

    async def async_process_request(self, request):
        """
                If CORS preflight header, then create an
                empty body response (200 OK) and return it
                Django won't bother calling any other request
                view/exception middleware along with the requested view;
                it will call any response middlewares
                """
        if CORS_REPLACE_HTTPS_REFERER:
            self._https_referer_replace(request)

        if (request.method == 'OPTIONS' and
                "HTTP_ACCESS_CONTROL_REQUEST_METHOD" in request.META):
            response = http.HttpResponse()
            return response
        return None


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def isCousins(self, root: TreeNode, x: int, y: int) -> bool:
        if root.val==x or root.val==y:
            return False
        x_parent, x_depth = self.search(root,x,0)
        y_parent, y_depth = self.search(root,y,0)
        print(x_parent,y_parent,x_depth,y_depth)
        if ((x_depth==y_depth) and (x_parent!=y_parent)):
            return True
        else:
            return False

    def search(self, root, x, d):
        p1,p2,dep1,dep2 = (None,None,None,None)

    # if left and right are empty
        if (root.left==None) and (root.right == None):
            return (None, None)

    # else if left or right has expected value
        if (root.left != None):
            if (root.left.val == x):
                return (root.val, d+1)
        if (root.right != None):
            if (root.right.val == x):
                return (root.val, d+1)

    # else search left and right if possible
        if root.left!=None:
            p1,dep1 = self.search(root.left, x, d+1)
            if p1 != None:
                return (p1,dep1)
        if root.right!=None:
            p2,dep2 = self.search(root.right, x, d+1)
            if p2!= None:
                return (p2,dep2)

    # if nothing found in left and right sub tree, return None
        return (None, None)


from optuna.samplers import RandomSampler
import optuna
import random
import torch
import numpy as np
import utils.utils as utils
from data.data_generation import get_datasets
from models.gmsm_instantiations import GMSM_binary

def run_hyper_tuning(config_hyper, config_data):
    # Load data
    seed = config_hyper["seed"]
    _ = set_seeds(seed)
    datasets = get_datasets(config_data)
    models = config_hyper["models"]
    tuning_ranges = config_hyper["tuning_ranges"]
    hyper_path = "/hyperparams/" + config_hyper["name"] + "/params/"

    # Tune
    for model in models:
        tune_sampler = set_seeds(seed)
        obj = get_objective(model, datasets, tuning_ranges)
        study = tune_objective(obj, study_name=model, num_samples=config_hyper["num_samples"], sampler=tune_sampler)
        best_params = study.best_trial.params
        # Save params
        path = hyper_path + model
        utils.save_yaml(path, best_params)
    print("Done tuning")

def set_seeds(seed):
    # Set seeds
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    tune_sampler = RandomSampler(seed=seed)
    return tune_sampler

def get_objective(key, datasets, tuning_ranges):
    def obj(trial):
        gmsm = GMSM_binary(datasets["causal_graph"], y_type=datasets["d_train"].datatypes["y_type"])
        #Sample params
        if key == "y" and datasets["d_train"].datatypes["y_type"] == "continuous":
            config = sample_params(trial, tuning_ranges["nf"])
        else:
            config = sample_params(trial, tuning_ranges["mlp"])
        # Fit
        if key in list(datasets["d_train"].data["m"].keys()):
            train_results = gmsm.fit_distribution_by_keys([key], config, d_train=datasets["d_train"], d_val=datasets["d_val"],
                                                          datatype="discrete")
        elif key == "y":
            train_results = gmsm.fit_distribution_by_keys([key], config, d_train=datasets["d_train"], d_val=datasets["d_val"],
                                                          datatype=datasets["d_train"].datatypes["y_type"])
        elif key == "y_regression":
            train_results = gmsm.fit_y_regression(config, d_train=datasets["d_train"], d_val=datasets["d_val"])
        elif key == "propensity":
            train_results = gmsm.fit_distribution_by_keys("a", config, d_train=datasets["d_train"], d_val=datasets["d_val"],
                                                          datatype=datasets["d_train"].datatypes["a_type"])
        else:
            raise ValueError("Invalid key in config: " + key)
        return train_results["val_results"][0]["val_obj"]

    return obj

def tune_objective(objective, study_name, num_samples=10, sampler=None):
    if sampler is not None:
        study = optuna.create_study(direction="minimize", study_name=study_name, sampler=sampler)
    else:
        study = optuna.create_study(direction="minimize", study_name=study_name)
    study.optimize(objective, n_trials=num_samples)

    print("Finished. Best trial:")
    trial_best = study.best_trial

    print("  Value: ", trial_best.value)

    print("  Params: ")
    for key, value in trial_best.params.items():
        print("    {}: {}".format(key, value))
    return study

def sample_params(trial, tuning_ranges):
    params = {}
    for param in tuning_ranges.keys():
        params[param] = trial.suggest_categorical(param, tuning_ranges[param])
    params["neptune"] = False
    return params


from __future__ import print_function
from builtins import str
import os
import unittest
import shutil
import unittest
import yaml
from fundamentals.utKit import utKit
from fundamentals import tools
from os.path import expanduser
home = expanduser("~")

packageDirectory = utKit("").get_project_root()
settingsFile = packageDirectory + "/test_settings.yaml"
# settingsFile = home + "/.config/soxspipe.recipes/soxspipe.recipes.yaml"
su = tools(
    arguments={"settingsFile": settingsFile},
    docString=__doc__,
    logLevel="DEBUG",
    options_first=False,
    projectName=None,
    defaultSettingsFile=False
)
arguments, settings, log, dbConn = su.setup()

# SETUP AND TEARDOWN FIXTURE FUNCTIONS FOR THE ENTIRE MODULE
moduleDirectory = os.path.dirname(__file__)
utKit = utKit(moduleDirectory)
log, dbConn, pathToInputDir, pathToOutputDir = utKit.setupModule()
utKit.tearDownModule()

try:
    shutil.rmtree(pathToOutputDir)
except:
    pass
# COPY INPUT TO OUTPUT DIR
shutil.copytree(pathToInputDir, pathToOutputDir)

# Recursively create missing directories
if not os.path.exists(pathToOutputDir):
    os.makedirs(pathToOutputDir)

# xt-setup-unit-testing-files-and-folders

pathToSqlite = pathToInputDir + "/test.sqlite"

class test_sqlite2mysql(unittest.TestCase):

    def test_sqlite2mysql_function(self):

        from fundamentals.mysql import sqlite2mysql
        this = sqlite2mysql(
            log=log,
            settings=settings,
            pathToSqlite=pathToSqlite,
            tablePrefix="Marvin",
            dbConn=dbConn
        )
        this.convert_sqlite_to_mysql()

    def test_sqlite2mysql_function_exception(self):

        from fundamentals.mysql import sqlite2mysql
        try:
            this = sqlite2mysql(
                log=log,
                settings=settings,
                pathToSqlite=pathToSqlite,
                fakeKey="break the code"
            )
            this.get()
            assert False
        except Exception as e:
            assert True
            # print(str(e))

        # x-print-testpage-for-pessto-marshall-web-object

    # x-class-to-test-named-worker-function


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Dec 28 16:49:18 2018

@author: weanl
"""

'''
    0.

    1. showScoreData()

    2. show_score_forecast()

    3. show_kpi_forecast()

    3.4 cmpMF(DBID='210100063')

'''

'''
'''
FileDir = '../given_data_1128/'
ForecastStep = 10

'''
'''
from BuildSeq import conSeqs

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

'''
    1. showScoreData()
'''
def showScoreData(DBID='210100063'):
    print('--->call--->showScoreData')

    # 1 load data
    agg_scorex_seqs, agg_scorey_seqs, agg_score_inputs, agg_score_outputs = conSeqs(DBID=DBID)

    slct = 210
    x_seq = agg_scorex_seqs[slct]
    y_seq = agg_scorey_seqs[slct]
    input = agg_score_inputs[:, slct]
    output = agg_score_outputs[:, slct]
    '''
    '''
    input[0] = input[0] + input[1]
    output[0] = output[0] + output[1]

    # 2 plot
    look_back = x_seq.shape[0]
    forecast_step = 10
    dec_lv = input.shape[0] - 1
    t1 = np.arange(0, look_back)
    t2 = np.arange(look_back, look_back+forecast_step)

    plt.figure(1, figsize=(16, 24))
    ROW = dec_lv + 3
    ## 2.1 plot x_seq and y_seq
    plt.subplot(ROW, 1, 1)
    plt.title('DBID = '+DBID)
    plt.plot(t1, x_seq, 'b.-')
    plt.plot(t2, y_seq, 'r.-')
    ## 2.1 plot sum(input) and sum(output)
    plt.subplot(ROW, 1, 2)
    x_sum = np.sum(input, axis=0)
    y_sum = np.sum(output, axis=0)
    plt.plot(t1, x_sum, 'b.-')
    plt.plot(t2, y_sum, 'r.-')
    ## 2.2 plot input and output
    for i in range(dec_lv+1):
        plt.subplot(ROW, 1, i+3)
        plt.plot(t1, input[i], 'b.-')
        plt.plot(t2, output[i], 'r.-')

    plt.show()

    print('--->return from--->showScoreData')
    return 0

'''
    2. show_score_forecast()
'''
def show_score_forecast(DBID='210100063'):
    print('--->call--->show_score_forecast')
    # 1 load data
    file_path = FileDir + DBID + '/' + 'score_forecast.npz'
    data = np.load(file_path)

    # 2
    y_test_pred = data['y_test_pred'].reshape(-1, )
    y_test_truth = data['y_test_truth'].reshape(-1, )
    y_train_pred = data['y_train_pred'].reshape(-1, )
    y_train_truth = data['y_train_truth'].reshape(-1, )

    # 3 compute mase and mse for test
    ae_test = abs(y_test_pred-y_test_truth)
    mae_test = np.mean(ae_test)
    ae_var_test = np.var(ae_test)
    print('mae_test: \t\t', mae_test)
    #print('mae_var_test: \t', ae_var_test)
    square_error_test = np.square(y_test_pred-y_test_truth)
    mse_test = np.mean(square_error_test)
    print('mse_test: \t\t', mse_test)

    # 4 compute mase and mse for train
    ae_train = abs(y_train_pred-y_train_truth)
    mae_train = np.mean(ae_train)
    ae_var_train = np.var(ae_train)
    print('mae_train: \t\t', mae_train)
    #print('mae_var_train: \t', ae_var_train)
    square_error_train = np.square(y_train_pred-y_train_truth)
    mse_train = np.mean(square_error_train)
    print('mse_train: \t\t', mse_train)
    '''
    # 4
    t = np.arange(testTruth.shape[0])
    #plt.plot(t, testTrue, label='testTrue', color='b')
    plt.plot(t, testTruth, label='testTruth', color='g')

    plt.plot(t, testPredict, label='prediction', color='r')

    # set interval size of grid along x
    plt.xticks(np.arange(1, testTruth.shape[0]+1, ForecastStep))

    plt.grid(linestyle='-.')
    plt.legend(loc='best')

    SavePath = FileDir + DBID + '/' + 'show_score_forecast.png'
    #plt.savefig(SavePath, dpi=256)
    plt.show()
    '''

    print('--->return from--->show_score_forecast')
    return mae_test, mse_test, mae_train, mse_train

'''
    # 2.2 cmpScoreBaseline(DBID='210100063')
'''
def cmpScoreBaseline(DBID='210100063'):
    print('--->call--->cmpScoreBaseline')

    # 0.1 load data of averModelScore
    print('# 0.1 load data of averModelScore')
    file_path_averModelScore = FileDir + DBID + '/' + 'score_forecast_averModelScore.npz'
    averModelScore = np.load(file_path_averModelScore)
    aver_y_test_pred = averModelScore['y_test_pred']
    aver_y_test = averModelScore['y_test']
    aver_y_train_pred = averModelScore['y_train_pred']
    aver_y_train = averModelScore['y_train']
    #print('aver_y_test_pred:\n', aver_y_test_pred)
    #print('aver_y_test:\n', aver_y_test)
    aver_y_test_pred = np.concatenate([aver_y_test_pred, aver_y_train_pred], axis=0)
    aver_y_test = np.concatenate([aver_y_test, aver_y_train], axis=0)
    print('np.std(aver_y_test) = ', np.std(aver_y_test))

    # 0.2 compute MAE and RMSE for averModelScore
    print('# 0.2 compute MAE and RMSE')
    aver_AE_test = abs(aver_y_test_pred-aver_y_test)
    aver_MAE_test = np.mean(aver_AE_test, axis=0)
    aver_SE_test = np.square(aver_y_test_pred-aver_y_test)
    aver_MSE_test = np.mean(aver_AE_test, axis=0)
    aver_RMSE_test= np.sqrt(aver_MSE_test)
    #print('aver_MAE_test: \n', aver_MAE_test)
    #print('aver_MAE_test mean = ', np.mean(aver_MAE_test))
    print('aver_RMSE_test: \n', aver_RMSE_test)
    print('aver_RMSE_test mean = ', np.mean(aver_RMSE_test))

    # 1.1 load data of arModelScore
    print('# 1.1 load data of arModelScore')
    file_path_arModelScore = FileDir + DBID + '/' + 'score_forecast_arModelScore.npz'
    arModelScore = np.load(file_path_arModelScore)
    ar_y_test_pred = arModelScore['y_test_pred']
    ar_y_test = arModelScore['y_test']
    ar_y_train_pred = arModelScore['y_train_pred']
    ar_y_train = arModelScore['y_train']
    #print('ar_y_test_pred:\n', ar_y_test_pred)
    #print('ar_y_test:\n', ar_y_test)
    ar_y_test_pred = np.concatenate([ar_y_test_pred, ar_y_train_pred], axis=0)
    ar_y_test = np.concatenate([ar_y_test, ar_y_train], axis=0)
    print('np.std(ar_y_test) = ', np.std(ar_y_test))

    # 1.2 compute MAE and RMSE for arModelScore
    print('# 1.2 compute MAE and RMSE')
    ar_AE_test = abs(ar_y_test_pred-ar_y_test)
    ar_MAE_test = np.mean(ar_AE_test, axis=0)
    ar_SE_test = np.square(ar_y_test_pred-ar_y_test)
    ar_MSE_test = np.mean(ar_AE_test, axis=0)
    ar_RMSE_test= np.sqrt(ar_MSE_test)
    #print('ar_MAE_test: \n', ar_MAE_test)
    #print('ar_MAE_test mean = ', np.mean(ar_MAE_test))
    print('ar_RMSE_test: \n', ar_RMSE_test)
    print('ar_RMSE_test mean = ', np.mean(ar_RMSE_test))

    # ardwtModelScore (recursive)
    # rnndwtModelScore (recursive)

    # 2.1 load data of ourModelScore
    print('# 2.1 load data of ourModelScore')
    file_path_ourModelScore = FileDir + DBID + '/' + 'score_forecast.npz'
    ourModelScore = np.load(file_path_ourModelScore)
    our_y_test_pred = ourModelScore['y_test_pred'].reshape(-1, ForecastStep)
    our_y_test = ourModelScore['y_test_truth']
    our_y_train_pred = ourModelScore['y_train_pred'].reshape(-1, ForecastStep)
    our_y_train = ourModelScore['y_train_truth']
    #print('our_y_test_pred:\n', our_y_test_pred)
    #print('our_y_test:\n', our_y_test)
    our_y_test_pred = np.concatenate([our_y_test_pred, our_y_train_pred], axis=0)
    our_y_test = np.concatenate([our_y_test, our_y_train], axis=0)
    print('np.std(our_y_test) = ', np.std(our_y_test))

    # 2.2 compute MAE and RMSE
    print('# 2.2 compute MAE and RMSE')
    our_AE_test = abs(our_y_test_pred-our_y_test)
    our_MAE_test = np.mean(our_AE_test, axis=0)
    our_SE_test = np.square(our_y_test_pred-our_y_test)
    our_MSE_test = np.mean(our_AE_test, axis=0)
    our_RMSE_test= np.sqrt(our_MSE_test)
    #print('our_MAE_test: \n', our_MAE_test)
    #print('our_MAE_test mean = ', np.mean(our_MAE_test))
    print('our_RMSE_test: \n', our_RMSE_test)
    print('our_RMSE_test mean = ', np.mean(our_RMSE_test))

    # -1 show RMSE for different score forecast method
    plt.figure(1, figsize=(16, 12))
    t = np.arange(ForecastStep)+1
    plt.plot(t, aver_RMSE_test, 'k.-', label='average:'
    +str(np.mean(aver_RMSE_test)))
    plt.plot(t, ar_RMSE_test, 'g.-', label='auto-regression:'
    +str(np.mean(ar_RMSE_test)))
    plt.xticks(t, fontsize=16)

    plt.plot(t, our_RMSE_test, 'r.-', label='MDWT-Seq2Seq:'
    +str(np.mean(our_RMSE_test)))
    plt.ylim(0, 1.6)
    plt.xlabel('ForecastStep', fontsize=20)
    plt.ylabel('RMSE', fontsize=20)

    plt.legend()
    plt.grid(True, linestyle='-.')

    plt.show()

    print('--->return from--->cmpScoreBaseline')
    return 0

'''
    3.1 show_kpi_forecast()
'''
def show_kpi_forecast(DBID='210100063'):
    print('--->call--->show_kpi_forecast')
    # 1 load data
    file_path = FileDir + DBID + '/' + 'kpi_forecast.npz'
    data = np.load(file_path)

    # 2
    y_test_pred = data['y_test_pred']
    y_test_truth = data['y_test_truth']
    y_train_pred = data['y_train_pred']
    y_train_truth = data['y_train_truth']
    KPIsNum = y_test_pred.shape[-1]

    # 3 compute mase and mse for test (all kpi)
    ae_test = abs(y_test_pred-y_test_truth)
    mae_test = np.mean(ae_test)
    ae_var_test = np.var(ae_test)
    print('mae_test: \t\t', mae_test)
    #print('mae_var_test: \t', ae_var_test)
    square_error_test = np.square(y_test_pred-y_test_truth)
    mse_test = np.mean(square_error_test)
    print('mse_test: \t\t', mse_test)

    # 4 compute mase and mse for train (all kpi)
    ae_train = abs(y_train_pred-y_train_truth)
    mae_train = np.mean(ae_train)
    ae_var_train = np.var(ae_train)
    print('mae_train: \t\t', mae_train)
    #print('mae_var_train: \t', ae_var_train)
    square_error_train = np.square(y_train_pred-y_train_truth)
    mse_train = np.mean(square_error_train)
    print('mse_train: \t\t', mse_train)

    '''
    # 4 plot
    pick = 1
    t = np.arange(y_pred.shape[0])
    plt.figure(1)
    plt.plot(t, y_pred[:, pick], 'g.-', label='y_pred')
    plt.plot(t, y_test[:, pick], 'r.-', label='y_test')
    plt.legend()
    plt.show()
    '''
    print('--->return from--->show_kpi_forecast')
    return mae_test, mse_test, mae_train, mse_train

'''
    3.2 show_joint_forecast()
'''
def show_joint_forecast(DBID='210100063'):
    print('--->call--->show_joint_forecast')
    # 1 load data
    file_path = FileDir + DBID + '/' + 'joint_forecast.npz'
    data = np.load(file_path)

    # 2
    scoreTestPrediction = data['scoreTestPrediction']
    scoreTestTruth = data['scoreTestTruth']
    scoreTrainPredict = data['scoreTrainPredict']
    scoreTrainTruth = data['scoreTrainTruth']
    kpicorTestPredict = data['kpicorTestPredict'] # ()
    kpicorTestTruth = data['kpicorTestTruth']
    kpicorTrainPredict = data['kpicorTrainPredict']
    kpicorTrainTruth = data['kpicorTrainTruth']
    '''
    print('kpicorTestPredict.shape = ', kpicorTestPredict.shape)
    print('kpicorTestTruth.shape = ', kpicorTestTruth.shape)
    print('kpicorTrainPredict.shape = ', kpicorTrainPredict.shape)
    print('kpicorTrainTruth.shape = ', kpicorTrainTruth.shape)
    '''
    KPIsNum = kpicorTestPredict.shape[-1]
    test_num = kpicorTestPredict.shape[0]
    train_num = kpicorTrainPredict.shape[0]
    kpicorTestTruth = kpicorTestTruth.reshape(test_num, KPIsNum)
    kpicorTrainTruth = kpicorTrainTruth.reshape(train_num, KPIsNum)

    #
    # 3 compute mae and mse for score test
    ae_score_test = abs(scoreTestPrediction - scoreTestTruth)
    mae_score_test = np.mean(ae_score_test)
    ae_var_score_test = np.var(ae_score_test)
    print('mae_score_test: \t\t', mae_score_test)
    #print('mae_var: \t', ae_var)
    square_error_score_test = np.square(scoreTestPrediction - scoreTestTruth)
    mse_score_test = np.mean(square_error_score_test)
    print('mse_score_test: \t\t', mse_score_test)

    # 4 compute mae and mse for score train
    ae_score_train = abs(scoreTrainPredict - scoreTrainTruth)
    mae_score_train = np.mean(ae_score_train)
    ae_var_score_train = np.var(ae_score_train)
    print('mae_score_train: \t\t', mae_score_train)
    #print('mae_var: \t', ae_var)
    square_error_score_train = np.square(scoreTrainPredict - scoreTrainTruth)
    mse_score_train = np.mean(square_error_score_train)
    print('mse_score_train: \t\t', mse_score_train)

    # 5 compute mae and mse for kpicor test
    ae_kpicor_test = abs(kpicorTestPredict - kpicorTestTruth)
    mae_kpicor_test = np.mean(ae_kpicor_test)
    ae_var_kpicor_test = np.var(ae_kpicor_test)
    print('mae_kpicor_test: \t\t', mae_kpicor_test)
    #print('mae_var: \t', ae_var)
    square_error_kpicor_test = np.square(kpicorTestPredict - kpicorTestTruth)
    mse_kpicor_test = np.mean(square_error_kpicor_test)
    print('mse_kpicor_test: \t\t', mse_kpicor_test)

    # 6 compute mae and mse for kpicor train
    ae_kpicor_train = abs(kpicorTrainPredict - kpicorTrainTruth)
    mae_kpicor_train = np.mean(ae_kpicor_train)
    ae_var_kpicor_train = np.var(ae_kpicor_train)
    print('mae_kpicor_train: \t\t', mae_kpicor_train)
    #print('mae_var: \t', ae_var)
    square_error_kpicor_train = np.square(kpicorTrainPredict - kpicorTrainTruth)
    mse_kpicor_train = np.mean(square_error_kpicor_train)
    print('mse_kpicor_train: \t\t', mse_kpicor_train)

    '''
    # 4 plot
    pick = 1
    t = np.arange(kpicorPredict.shape[0])
    plt.figure(1)
    plt.plot(t, kpicorPredict[:, pick], 'g.-', label='kpicorPredict')
    plt.plot(t, kpicorTruth[:, pick], 'r.-', label='kpicorTruth')
    plt.legend()
    plt.show()
    '''
    print('--->return from--->show_joint_forecast')
    return mae_score_test, mse_score_test, mae_score_train, mse_score_train, mae_kpicor_test, mse_kpicor_test, mae_kpicor_train, mse_kpicor_train

'''
    3.3 show_forecast_compared()
'''
def show_forecast_compared(DBID='210100063'):

    mae_sep, mse_sep = show_kpi_forecast(DBID)
    mae_joint, mse_joint = show_joint_forecast(DBID)

    mae_sep_mean = np.mean(mae_sep)
    mse_sep_mean = np.mean(mse_sep)
    print('mae_sep_mean = ', mae_sep_mean)
    print('mse_sep_mean = ', mse_sep_mean)

    mae_joint_mean = np.mean(mae_joint)
    mse_joint_mean = np.mean(mse_joint)
    print('mae_joint_mean = ', mae_joint_mean)
    print('mse_joint_mean = ', mse_joint_mean)

    return 0

'''
    3.4 cmpMF(DBID='210100063')
'''
def cmpMF(DBID='210100063'):
    print('---> call ---> cmpMF')

    # 1 load data
    file_path1 = FileDir + DBID + '/' + 'score_forecast.npz'
    file_path2 = FileDir + DBID + '/' + 'joint_forecast.npz'
    data_sep = np.load(file_path1)
    data_joint = np.load(file_path2)

    scorePredSep = data_sep['y_test_pred'].reshape(-1, ForecastStep)
    scoreTruthSep = data_sep['y_test_truth']
    scorePredJoint = data_joint['scoreTestPrediction']
    scoreTruthJoint = data_joint['scoreTestTruth']

    # 2
    SE_sep = np.square(scorePredSep-scoreTruthSep)
    MSE_sep = np.mean(SE_sep, axis=0)
    RMSE_sep = np.sqrt(MSE_sep)

    SE_joint = np.square(scorePredJoint-scoreTruthJoint)
    MSE_joint = np.mean(SE_joint, axis=0)
    RMSE_joint = np.sqrt(MSE_joint)

    # 3
    plt.figure(1, figsize=(16, 12))
    t = np.arange(1, ForecastStep+1)
    plt.subplot(1, 1, 1)
    plt.plot(t, RMSE_sep, 'g')
    plt.plot(t, RMSE_joint, 'r')

    plt.show()

    print('---> return from --> cmpMF')
    return 0

'''
    3.5 cmpCor(DBID='210100063')
'''
def cmpCor(DBID='210100063'):    # 1 load data
    file_path1 = FileDir + DBID + '/' + 'score_forecast.npz'
    file_path2 = FileDir + DBID + '/' + 'joint_forecast.npz'
    data_sep = np.load(file_path1)
    data_joint = np.load(file_path2)
    print('---> call ---> cmpCor')

    # 1 load data
    file_path1 = FileDir + DBID + '/' + 'kpi_forecast.npz'
    file_path2 = FileDir + DBID + '/' + 'joint_forecast.npz'
    data_sep = np.load(file_path1)
    data_joint = np.load(file_path2)

    kpiPredSep = data_sep['y_test_pred']
    KPIsNum = kpiPredSep.shape[-1]
    kpiTruthSep = data_sep['y_test_truth']
    kpiPredJoint = data_joint['kpicorTestPredict']
    kpiTruthJoint = data_joint['kpicorTestTruth'].reshape(-1, KPIsNum)
    '''
    print('kpiPredSep.shape = ', kpiPredSep.shape)
    print('kpiTruthSep.shape = ', kpiTruthSep.shape)
    print('kpiPredJoint.shape = ', kpiPredJoint.shape)
    print('kpiTruthJoint.shape = ', kpiTruthJoint.shape)
    '''

    # 2
    SE_sep = np.square(kpiPredSep-kpiTruthSep)
    MSE_sep = np.mean(SE_sep, axis=0)
    RMSE_sep = np.sqrt(MSE_sep)
    print('RMSE_sep.shape = ', RMSE_sep.shape)

    SE_joint = np.square(kpiPredJoint-kpiTruthJoint)
    MSE_joint = np.mean(SE_joint, axis=0)
    RMSE_joint = np.sqrt(MSE_joint)
    print('RMSE_joint.shape = ', RMSE_joint.shape)

    # 3
    plt.figure(1, figsize=(16, 12))
    t = np.arange(1,KPIsNum+1)
    plt.subplot(1, 1, 1)
    mean_sep = np.mean(RMSE_sep)
    plt.bar(t, RMSE_sep, width=0.4, label='Seperate: '+str(mean_sep))
    mean_joint = np.mean(RMSE_joint)
    plt.bar(t+0.4, RMSE_joint, width=0.4, label='Joint: '+str(mean_joint))

    plt.legend()
    plt.show()

    print('---> return from ---> cmpCor')
    return 0

'''
    3.6 cmpCorTop(DBID='210100063')
'''
def cmpCorTop(DBID='210100063'):
    file_path1 = FileDir + DBID + '/' + 'score_forecast.npz'
    file_path2 = FileDir + DBID + '/' + 'joint_forecast.npz'
    data_sep = np.load(file_path1)
    data_joint = np.load(file_path2)
    print('---> call ---> cmpCor')

    # 1 load data
    file_path1 = FileDir + DBID + '/' + 'kpi_forecast.npz'
    file_path2 = FileDir + DBID + '/' + 'joint_forecast.npz'
    data_sep = np.load(file_path1)
    data_joint = np.load(file_path2)

    kpiPredSep = data_sep['y_test_pred']
    KPIsNum = kpiPredSep.shape[-1]
    kpiTruthSep = data_sep['y_test_truth']
    kpiPredJoint = data_joint['kpicorTestPredict']
    kpiTruthJoint = data_joint['kpicorTestTruth'].reshape(-1, KPIsNum)

    # 2 compute mean recall fo Top M with predict
    kpiGroundTruth = abs(kpiTruthSep)
    kpiPredSep = abs(kpiPredSep)
    kpiPredJoint = abs(kpiPredJoint)
    print('kpiGroundTruth.shape = ', kpiGroundTruth.shape)

    M = 8
    TruthIndex = np.argsort(kpiGroundTruth, axis=1)[:, -M:]
    PredSepIndex = np.argsort(kpiPredSep, axis=1)[:, -M:]
    PredJointIndex = np.argsort(kpiPredJoint, axis=1)[:, -M:]
    print('TruthIndex: \n', TruthIndex)
    print('PredSepIndex: \n', PredSepIndex)
    print('PredJointIndex: \n', PredJointIndex)

    return 0

if __name__ == '__main__':
    print('------------ ShowData ------------')

    #showScoreData(DBID='210100063')
    # 1
    '''
    show_score_forecast(DBID='210100063')
    show_kpi_forecast()
    show_joint_forecast()
    '''

    # 2
    '''
    show_kpi_forecast()
    show_joint_forecast()
    '''
    #show_forecast_compared()

    # 3
    #cmpMF(DBID='210100063')

    # 4
    #cmpCor(DBID='210100063')

    # 5
    #cmpCorTop(DBID='210100063')

    # 6
    cmpScoreBaseline(DBID='210100063')

#   END OF FILE