import logging
import time
import json
import unittest
from secrets import token_bytes
from typing import Any, Callable, Dict, List, Optional, Set, Tuple

import asyncio
import aiosqlite
import sqlite3

from blspy import AugSchemeMPL, G1Element, G2Element
from hashlib import sha256

from chives.consensus.constants import ConsensusConstants
from chives.consensus.coinbase import create_puzzlehash_for_pk
from chives.util.bech32m import decode_puzzle_hash, encode_puzzle_hash
from chives.util.config import load_config
from chives.util.default_root import DEFAULT_ROOT_PATH
from chives.util.ints import uint32, uint64
from chives.util.hash import std_hash
from chives.util.db_wrapper import DBWrapper
from chives.util.keychain import Keychain, bytes_from_mnemonic, bytes_to_mnemonic, generate_mnemonic, mnemonic_to_seed
from chives.wallet.derive_keys import master_sk_to_farmer_sk, master_sk_to_pool_sk, master_sk_to_wallet_sk
from chives.wallet.wallet_coin_store import WalletCoinStore
from chives.types.blockchain_format.coin import Coin
from chives.types.blockchain_format.program import Program, SerializedProgram
from chives.types.blockchain_format.sized_bytes import bytes32

from chives.wallet.util.wallet_types import WalletType
from chives.wallet.wallet_coin_record import WalletCoinRecord

from chives.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import (
    DEFAULT_HIDDEN_PUZZLE_HASH,
    calculate_synthetic_secret_key,
    puzzle_for_pk,
    solution_for_conditions,
)
from chives.wallet.puzzles.puzzle_utils import (
    make_assert_coin_announcement,
    make_assert_puzzle_announcement,
    make_assert_my_coin_id_condition,
    make_assert_absolute_seconds_exceeds_condition,
    make_create_coin_announcement,
    make_create_puzzle_announcement,
    make_create_coin_condition,
    make_reserve_fee_condition,
)

import redis
r = redis.Redis(host='localhost', port=6379, decode_responses=True)

class TesKeychain():

    async def puzzle_for_puzzle_hash(puzzle_hash: bytes32) -> Program:
        public_key = await TesKeychain.hack_populate_secret_key_for_puzzle_hash(puzzle_hash)
        return puzzle_for_pk(bytes(public_key))

    def make_solution(
        primaries: Optional[List[Dict[str, Any]]] = None,
        min_time=0,
        me=None,
        coin_announcements: Optional[List[bytes32]] = None,
        coin_announcements_to_assert: Optional[List[bytes32]] = None,
        puzzle_announcements=None,
        puzzle_announcements_to_assert=None,
        fee=0,
    ) -> Program:
        assert fee >= 0
        condition_list = []
        if primaries:
            for primary in primaries:
                condition_list.append(make_create_coin_condition(primary["puzzlehash"], primary["amount"]))
        if min_time > 0:
            condition_list.append(make_assert_absolute_seconds_exceeds_condition(min_time))
        if me:
            condition_list.append(make_assert_my_coin_id_condition(me["id"]))
        if fee:
            condition_list.append(make_reserve_fee_condition(fee))
        if coin_announcements:
            for announcement in coin_announcements:
                condition_list.append(make_create_coin_announcement(announcement))
        if coin_announcements_to_assert:
            for announcement_hash in coin_announcements_to_assert:
                condition_list.append(make_assert_coin_announcement(announcement_hash))
        if puzzle_announcements:
            for announcement in puzzle_announcements:
                condition_list.append(make_create_puzzle_announcement(announcement))
        if puzzle_announcements_to_assert:
            for announcement_hash in puzzle_announcements_to_assert:
                condition_list.append(make_assert_puzzle_announcement(announcement_hash))
        return solution_for_conditions(condition_list)

    async def TestTransaction():
        root_path = DEFAULT_ROOT_PATH
        config = load_config(root_path, "config.yaml")
        selected = config["selected_network"]
        prefix = config["network_overrides"]["config"][selected]["address_prefix"]
        log = logging.Logger
        db_connection = await aiosqlite.connect("/home/wang/.chives/mainnet/db/blockchain_v1_mainnet.sqlite")
        mnemonic = generate_mnemonic()
        mnemonic = "hen battle gauge crouch dose weasel blind noble ugly pull cruel mutual slight tragic bean rule once garage valley ritual still couple charge rich"
        entropy = bytes_from_mnemonic(mnemonic)
        seed = mnemonic_to_seed(mnemonic, "")
        seed_key = AugSchemeMPL.key_gen(seed)
        masterPublicKey = seed_key.get_g1()
        fingerprint = masterPublicKey.get_fingerprint()

        MapKeys = {}
        for i in range(10):
            primary_key = master_sk_to_wallet_sk(seed_key, uint32(i))
            public_key = primary_key.get_g1()
            puzzle_hash = create_puzzlehash_for_pk(public_key)
            address = encode_puzzle_hash(puzzle_hash, prefix)
            MapKeys[puzzle_hash] = public_key
            MapKeys[i] = puzzle_hash
            print(puzzle_hash)
        print(MapKeys)

        # Get coin infor
        coin_name = "9d1cbc9cf8a5ad3883933fd05367562bb771ab5ef4cb6200b6b9acdb4b2c8117";
        newpuzzlehash = MapKeys[2]
        SendAmount = 0.01*100000000
        fee = 0
        cursor = await db_connection.execute("SELECT * from coin_record WHERE coin_name=?", (coin_name,))
        row = await cursor.fetchone()
        await cursor.close()
        if row is None:
            return None
        # parent_coin_info puzzle_hash amount
        coin = Coin(bytes32(bytes.fromhex(row[6])), bytes32(bytes.fromhex(row[5])), uint64.from_bytes(row[7]))
        # print(coin)
        WallTypeValue = 0
        WallTypeId = 1
        WalletCoinRecord(
            coin, uint32(row[1]), uint32(row[2]), bool(row[3]), bool(row[4]), WalletType(WallTypeValue), WallTypeId
        )
        # select_coins
        select_coins: Set = set()
        select_coins.add(coin)

        spends: List[CoinSolution] = []
        primary_announcement_hash: Optional[bytes32] = None

        origin_id = None
        primaries: Optional[List[Dict]] = None
        for coin in select_coins:
            # log.info(f"coin from coins {coin}")
            # print(coin)
            print(coin)
            #puzzle: Program = await TesKeychain.puzzle_for_puzzle_hash(coin.puzzle_hash)
            public_key = MapKeys[puzzle_hash]
            assert public_key is not None
            puzzle: Program =  puzzle_for_pk(bytes(public_key))
            #print(public_key)
            #print(puzzle)

            change = coin.amount - SendAmount
            # Only one coin creates outputs
            if primary_announcement_hash is None and origin_id in (None, coin.name()):
                if primaries is None:
                    primaries = [{"puzzlehash": newpuzzlehash, "amount": SendAmount}]
                else:
                    primaries.append({"puzzlehash": newpuzzlehash, "amount": SendAmount})
                if change > 0:
                    # CHANGE 地址为第二个地址
                    change_puzzle_hash: bytes32 = MapKeys[1]
                    primaries.append({"puzzlehash": change_puzzle_hash, "amount": change})
                message_list: List[bytes32] = [c.name() for c in select_coins]
                print(message_list)
                print('#############################')
                for primary in primaries:
                    print(coin.name())
                    coinNew = Coin(coin.name(), primary["puzzlehash"], uint32(primary["amount"])).name()
                    message_list.append(coinNew)
                print('#############################')

                message: bytes32 = std_hash(b"".join(message_list))
                solution: Program = TesKeychain.make_solution(primaries=primaries, fee=fee, coin_announcements=[message])
                primary_announcement_hash = Announcement(coin.name(), message).name()
            else:
                solution = TesKeychain.make_solution(coin_announcements_to_assert=[primary_announcement_hash])

            spends.append(
                CoinSolution(
                    coin, SerializedProgram.from_bytes(bytes(puzzle)), SerializedProgram.from_bytes(bytes(solution))
                )
            )

        #coin_record: WalletCoinRecord = WalletCoinRecord(
        #    coin, height, uint32(0), False, farm_reward, wallet_type, wallet_id
        #)

# xcc1dr0leqc48k0k3ul7386ulxppf8ru5rmqx6gjffdsdff0tgxj4wqssewhcj
# 68dffc83153d9f68f3fe89f5cf982149c7ca0f60369124a5b06a52f5a0d2ab81
# COIN_NAME 7541233a21d81a443c5809680aca026029547108c091869ee8fb1ad3b09850e5
# COIN_NAME 6a5d959896271bbf01cb29c255cc9dfd33125a940676ec97b2da7decd56f5374
# COIN_NAME 7badb9975ec2b4634093a4e74ecd840c527b0fdc81a42d5758b48c770f428cd9
if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(TesKeychain.TestTransaction())


import numpy as np
from sklearn import preprocessing
from sklearn.feature_extraction.text import TfidfVectorizer
from pandas import read_hdf, concat
from sklearn.metrics import f1_score, accuracy_score
from time import time

from sklearn.ensemble import RandomForestClassifier
from Utils.utils import type2idx

# Load data
TrainServices = read_hdf('D:\python_projects\ServeNet_others\data\\ramdom_categorg_percent\RandomSplittedByCatagories9.h5', key='Train')
TestServices = read_hdf('D:\python_projects\ServeNet_others\data\\ramdom_categorg_percent\RandomSplittedByCatagories9.h5', key='Test')
AllData = concat([TrainServices, TestServices])

data_train=list(TrainServices['Service Desciption'])
target_train=list(TrainServices['Service Classification'])
data_test=list(TestServices['Service Desciption'])
target_test=list(TestServices['Service Classification'])

X_train=data_train
Y_train=target_train
X_test=data_test
Y_test=target_test

Type_c = (list(np.unique(target_train)))

encoder = preprocessing.LabelEncoder()
Y_train = encoder.fit_transform(Y_train)
Y_test = encoder.fit_transform(Y_test)

max_features = 2000

tfidf_vectorizer=TfidfVectorizer(sublinear_tf=True, stop_words='english', max_features=max_features)
tfidf_vectorizer.fit(list(AllData['Service Desciption']))

X_train = tfidf_vectorizer.transform(X_train)
X_test = tfidf_vectorizer.transform(X_test)

# Train processing
clf = RandomForestClassifier(n_estimators=2000, max_depth=40)

t0 = time()
clf.fit(X_train, Y_train)
t1 = time()
print("Train time: ", t1 - t0)

train_top5 = clf.predict_proba(X_train)
train_top1 = clf.predict(X_train)

test_pre_top5 = clf.predict_proba(X_test)
test_pre_top1 = clf.predict(X_test)

test_ret = np.empty((len(Y_test),), dtype=np.int)
train_ret = np.empty((len(Y_train),), dtype=np.int)
for i in range(len(Y_test)):
    Top5_test = sorted(zip(clf.classes_, test_pre_top5[i]), key=lambda x: x[1])[-5:]
    Top5_test=list(map(lambda x: x[0], Top5_test))

    if Y_test[i] in Top5_test:
        test_ret[i] = Y_test[i]
    else:
        test_ret[i] = Top5_test[-1]

for i in range(len(Y_train)):
    Top5_train = sorted(zip(clf.classes_, train_top5[i]), key=lambda x: x[1])[-5:]
    Top5_train = list(map(lambda x: x[0], Top5_train))

    if Y_train[i] in Top5_train:
        train_ret[i] = Y_train[i]
    else:
        train_ret[i] = Top5_train[-1]

f1_s = f1_score(Y_test, test_ret, average='micro')

print("=" * 60)
print("Test top5 acc:%.4f,train top5  acc:%.4f" % (accuracy_score(Y_test, test_ret), accuracy_score(Y_train, train_ret)))
print("Test top1 acc:%.4f,train top1 acc:%.4f" % (
accuracy_score(Y_test, test_pre_top1), accuracy_score(Y_train, train_top1)))
print("F1_score:%.4f" % float(f1_s))
print("=" * 60)
####################################################################
# calculate accuracy of each category.
# type_c_index = type2idx(Type_c, Type_c)
#
# result_dict = {}
# total_dict = {}
# for idx in type_c_index:
#     category = Type_c[idx]
#     total_count = 0
#     account = 0
#     for i in range(len(Y_test)):
#         if Y_test[i] == idx:
#             total_count += 1
#             if Y_test[i] == test_ret[i]:
#                 account += 1
#
#     result_dict[category] = account / total_count * 1.
#     total_dict[category] = total_count
#
# for cate in result_dict.keys():
#     total_account = total_dict[cate]
#     acc = result_dict[cate]
#     print("%s (%d): %.4f" % (cate, total_account, acc))

############################################
# top-1 categories
print("=" * 60)
type_c_index = type2idx(Type_c, Type_c)

result_dict = {}
total_dict = {}
avg = 0.0
correct_num = 0
print(Y_test.shape)
print(test_pre_top1.shape)
for idx in type_c_index:
    category = Type_c[idx]
    total_count = 0
    account = 0
    for i in range(len(Y_test)):
        if Y_test[i] == idx:
            total_count += 1
            if Y_test[i] == test_pre_top1[i]:
                account += 1
                correct_num += 1

    result_dict[category] = format(account / total_count * 100., '.2f')
    total_dict[category] = total_count

labels = ["Tools","Financial","Messaging","eCommerce","Payments","Social","Enterprise","Mapping","Telephony","Science",
          "Government","Email","Security","Reference","Video","Travel","Sports","Search","Advertising","Transportation",
          "Education","Games","Music","Photos","Cloud","Bitcoin","Project Management","Data","Backend","Database",
          "Shipping","Weather","Application Development","Analytics","Internet of Things","Medical","Real Estate",
          "Events","Banking","Stocks","Entertainment","Storage","Marketing","File Sharing","News Services","Domains",
          "Chat","Media","Images","Other"]

for label in labels:
    acc = result_dict[label]
    print(acc)


import itertools
import numpy as np
from lib_naloga1 import sample_distance
from lib_naloga1 import nesttools

########################
# Author: Jernej Vivod #
########################

# average_linkage: return average distance between samples in group c1 and samples in group c2.
def average_linkage(c1, c2, data):
	c1_elements = list(nesttools.un_nest(c1)) # Get elements in groups c1 and c2.
	c2_elements = list(nesttools.un_nest(c2))
	prod = itertools.product(c1_elements, c2_elements) 	# Get cartesian product of elements from the groups.

	# Create accumulator for measuring the sum of distances of pairs in cartesian product.
	total_dist = 0
	for pair in prod:
		pair_fst_data = data[pair[0]] # Get data for countries in pair.
		pair_snd_data = data[pair[1]]
		dist = sample_distance.euclidean_dist(pair_fst_data, pair_snd_data) # Compute distance and add to total.
		total_dist += dist

	# Return average distance between elements of groups.
	return total_dist / (len(c1_elements) * len(c2_elements))

# complete_linkage: return maximal distance between two samples where first sample is in group c1 and second sample in group c2.
def complete_linkage(c1, c2, data):
	c1_elements = list(nesttools.un_nest(c1))	# Get elements in groups c1 and c2.
	c2_elements = list(nesttools.un_nest(c2))

	# Get list of of data for each country in each group.
	c1_data = list(map(lambda x: data[x], c1_elements))
	c2_data = list(map(lambda x: data[x], c2_elements))

	# Initialize max distance to 0.
	max_dist = 0

	# Find max distance between samples in different groups.
	for c1_sample in c1_data:
		for c2_sample in c2_data:
			dist = sample_distance.euclidean_dist(c1_sample, c2_sample)
			if dist > max_dist: 			# If distance is new maximal distance...
				max_dist = dist

	# Return found maximal distance
	return max_dist

# single_linkage: return minimal distance between two samples where first sample is in group c1 and second sample in group c2.
def single_linkage(c1, c2, data):
	c1_elements = list(nesttools.un_nest(c1)) # Get elements in groups c1 and c2.
	c2_elements = list(nesttools.un_nest(c2))

	# Get list of of data for each country in each group.
	c1_data = list(map(lambda x: data[x], c1_elements))
	c2_data = list(map(lambda x: data[x], c2_elements))

	# Initialize min distance to a very large value.
	min_dist = int(1e20)

	# Find max distance between samples in different groups.
	for c1_sample in c1_data:
		for c2_sample in c2_data:
			dist = sample_distance.euclidean_dist(c1_sample, c2_sample)
			if dist < min_dist: 	# If distance is new minimal distance...
				min_dist = dist

	# Return found maximal distance
	return min_dist

# ward_distance: compute ward distance between clusters c1 and c2.
def ward_distance(c1, c2, data):
	c1_elements = list(nesttools.un_nest(c1))	# Get elements in groups c1 and c2.
	c2_elements = list(nesttools.un_nest(c2))

	# Get list of of data for each country in each group.
	c1_data = list(map(lambda x: data[x], c1_elements))
	c2_data = list(map(lambda x: data[x], c2_elements))

	# Find centroids of c1 and c2 (average of samples in groups).
	Rc1 = np.zeros(47, dtype = int)
	for el in c1_data:
		Rc1 = np.add(Rc1, el)
	Rc1 = np.true_divide(Rc1, len(c1_data))

	Rc2 = np.zeros(47, dtype = int)
	for el in c2_data:
		Rc2 = np.add(Rc2, el)
	Rc2 = np.true_divide(Rc2, len(c2_data))

	# Find centroid of union(c1 c2) (average of samples in union).
	Rc1c2 = np.zeros(47, dtype = int)
	for el in np.concatenate([c1_data, c2_data]):
		Rc1c2 = np.add(Rc1c2, el)
	Rc1c2 = np.true_divide(Rc1c2, len(np.concatenate([c1_data, c2_data])))

	# Compute and return ward distance using formula.
	sum_1 = 0
	for el in np.concatenate([c1_data, c2_data]):
		sum_1 += sample_distance.manhattan_dist(el, Rc1c2)**2

	sum_2 = 0
	for el in c1_data:
		sum_2 += sample_distance.manhattan_dist(el, Rc1)**2

	sum_3 = 0
	for el in c2_data:
		sum_3 += sample_distance.manhattan_dist(el, Rc2)**2

	return sum_1 - (sum_2 + sum_3)


import os

from kivy.lang import Builder
from kivy.metrics import dp
from kivymd.uix.bottomnavigation import MDBottomNavigationItem
from kivymd.uix.imagelist import SmartTile
from kivy.uix.scrollview import ScrollView
from kivymd.uix.gridlayout import MDGridLayout
from kivymd.uix.button import MDFloatingActionButton
from kivymd.uix.filemanager import MDFileManager
from config import Config

Builder.load_file(f"{Config.TEMPLATES_DIR}/imagecollectiontab.kv")

class ImageCell(SmartTile):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.box_color = (0, 0, 0, 0)

class ImageGrid(MDGridLayout):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.padding = (dp(0), dp(0))
        self.spacing = dp(4)

    def get_free_cell(self):
        for image in self.images:
            if not image.source:
                return image
        return

    def add_image_cells(self):
        for image in self.images:
            self.add_widget(image)

class ThreeVerticalImagesGrid(ImageGrid):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.cols = 1
        self.rows = 3
        self.size_hint = (0.2, 0.67)
        self.images = (ImageCell(), ImageCell(), ImageCell())
        self.add_image_cells()

class BigImageGrid(ImageGrid):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.cols = 1
        self.rows = 1
        self.size_hint = (0.6, 0.67)
        self.images = (ImageCell(),)
        self.add_image_cells()

class BlockOfImages(ImageGrid):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._current_grid = None

        self._first_col_grid = None
        self._middle_block_grid = None
        self._last_col_grid = None

        self.rows = 1
        self.size_hint = (1, 0.5)

        self.images = []
        self.padding = (dp(2), dp(2))
        self._make_new_grid()

    def _to_next_grid(self):
        if self._current_grid == self._first_col_grid:
            self._current_grid = self._middle_block_grid
        elif self._current_grid == self._middle_block_grid:
            self._current_grid = self._last_col_grid
        elif self._current_grid == self._last_col_grid:
            self._make_new_grid()

    def get_free_cell(self):
        if self._last_col_grid.children[0].source:
            return
        image = self._current_grid.get_free_cell()
        if not image:
            self._to_next_grid()
            image = self._current_grid.get_free_cell()
        return image

    def _make_new_grid(self):
        self._first_col_grid = ThreeVerticalImagesGrid()
        self._middle_block_grid = BigImageGrid()
        self._last_col_grid = ThreeVerticalImagesGrid()

        self.add_widget(self._first_col_grid)
        self.add_widget(self._middle_block_grid)
        self.add_widget(self._last_col_grid)

        self._current_grid = self._first_col_grid

class ImageGridBuilder(MDGridLayout):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.blocks = [BlockOfImages(), BlockOfImages(), BlockOfImages()]
        self._idx = 0
        self._current_block = self.blocks[self._idx]
        self.cols = 1
        self.size_hint = (1, 1.5)

        for block in self.blocks:
            self.add_widget(block)

    def _to_next_block(self):
        self._idx += 1
        self._current_block = self.blocks[self._idx]

    def add_image(self, source):
        image = self._current_block.get_free_cell()
        if not image:
            self._to_next_block()
            image = self._current_block.get_free_cell()
        image.source = source

class ImageChooser(MDFileManager):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.exit_manager = self.exit
        self.preview = False
        self.external_storage = os.getenv('EXTERNAL_STORAGE')
        self.images_folder = f"{self.external_storage}/Pictures"

    def select_path(self, path):
        ImageCollectionTab.image_collection.builder.add_image(path)
        self.exit()

    def exit(self, *args):
        self.close()

    def open(self):
        self.show(self.images_folder)

class ImageCollection(MDGridLayout):
    def __init__(self, **kwargs):
        super().__init__(cols=1, **kwargs)

        self.__next_image_index = 0

        self.add_image_button = MDFloatingActionButton(
            icon="plus",
            on_release=self.open_image_chooser
        )

        self.scroll_view = ScrollView(size_hint=(1, 1))

        self.builder = ImageGridBuilder()

        self.scroll_view.add_widget(self.builder)
        self.add_widget(self.scroll_view)
        self.add_widget(self.add_image_button)

    def open_image_chooser(self, touch):
        ImageCollectionTab.image_chooser.open()

class ImageCollectionTab(MDBottomNavigationItem):
    """Tab that contains personal information."""

    image_chooser = None
    image_collection = None
    x_size = None

    def __init__(self, **kwargs):
        super().__init__(name="img_collection", text="Images",
                         icon="image-frame", **kwargs)

        ImageCollectionTab.x_size = self.size[0]
        ImageCollectionTab.image_collection = ImageCollection()
        ImageCollectionTab.image_chooser = ImageChooser()

        self.add_widget(ImageCollectionTab.image_collection)


import json
import os
import pickle as pkl
import sys
from time import time, strftime, gmtime
import numpy as np
import torch
import torch.nn as nn
import xarray as xr
from torch.utils.data import DataLoader

import model.Baseline as Baseline
from utils import ModelUtils
from utils.Logger import Logger
from utils.data import DataUtils, DataLoaders

# main method for training a neural network to directly predict the 2m-temperature prediction
def runModel(config, data_dictionary, data_statistics, train_test_folds):
    program_start_time = time()

    # assign all program arguments to local variables
    with open(config['model']['path']) as handle:
        ModelDict = json.loads(handle.read())

    # check if station and grid time invariant features should be used and set the list of desired parameters
    if not ('grid_time_invariant' in ModelDict and ModelDict['grid_time_invariant']): config['grid_time_invariant_parameters'] =[]
    if not ('station_time_invariant' in ModelDict and ModelDict['station_time_invariant']): config['station_parameters'] = []

    # update general static model information
    experiment_info = config
    experiment_info['model'] = ModelDict
    experiment_info['code_commit'] = ModelUtils.get_git_revision_short_hash()

    # if needed, load time invariant features
    with open("%s/%s/grid_size_%s/time_invariant_data_per_station.pkl" % (config['input_source'], config['preprocessing'], config['original_grid_size']), "rb") as input_file:
        time_invarian_data = pkl.load(input_file)

    # initialize feature scaling function for each feature
    featureScaleFunctions = DataUtils.getFeatureScaleFunctions(ModelUtils.ParamNormalizationDict, data_statistics)

    # get optimizer config
    optimizer_config = config['optimizer']

    # generate output path for experiment information
    setting_string = '%s_grid_%s_bs_%s_tf_%s_optim_%s_lr_%s_sl_%s' % (
        config['model']['name'], config['grid_size'], config['batch_size'], config['test_fraction'], optimizer_config['algorithm'], optimizer_config['learning_rate'], config['slice_size'])
    output_path = '%s/%s' % (config['experiment_path'], setting_string)
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # time for the set up until first run
    experiment_info['set_up_time'] = time() - program_start_time
    print('[Time]: Set-up %s' % strftime("%H:%M:%S", gmtime(experiment_info['set_up_time'])))
    sys.stdout.flush()

    # initialize statistics
    error_statistics = None
    run_times = None
    skip_statistics = None
    if 'per_station_rmse' in config:
        error_per_station_statistics = None

    # keep used learning rates
    experiment_info['scheduled_learning_rates'] = []

    # cross validation
    for run in range(config['runs']):
        # logger  for tensorboardX
        train_logger = Logger(output_path + '/logs/run_%s/train' % run)
        test_logger = Logger(output_path + '/logs/run_%s/test' % run)

        print('[Run %s] Cross-validation test fold %s' % (str(run + 1), str(run + 1)))

        # take the right preprocessed train/test data set for the current run
        train_fold, test_fold = train_test_folds[run]

        # initialize best epoch test error
        best_epoch_test_rmse = float("inf")

        # use different data loader if we want to train a 3nn model approach
        if "knn" in ModelDict:
            # initialize train and test dataloaders
            trainset = DataLoaders.CosmoData3NNData(
                config=config,
                station_data_dict=data_dictionary,
                files=train_fold,
                featureScaling=featureScaleFunctions,
                time_invariant_data=time_invarian_data)
            trainloader = DataLoader(trainset, batch_size=config['batch_size'], shuffle=True,
                                     num_workers=config['n_loaders'], collate_fn=DataLoaders.collate_fn)

            testset = DataLoaders.CosmoData3NNData(
                config=config,
                station_data_dict=data_dictionary,
                files=test_fold,
                featureScaling=featureScaleFunctions,
                time_invariant_data=time_invarian_data)
            testloader = DataLoader(testset, batch_size=config['batch_size'], shuffle=True,
                                    num_workers=config['n_loaders'], collate_fn=DataLoaders.collate_fn)
        else:
            # initialize train and test dataloaders
            trainset = DataLoaders.CosmoDataGridData(
                config=config,
                station_data_dict=data_dictionary,
                files=train_fold,
                featureScaling=featureScaleFunctions,
                time_invariant_data=time_invarian_data)
            trainloader = DataLoader(trainset, batch_size=config['batch_size'], shuffle=True,
                                     num_workers=config['n_loaders'], collate_fn=DataLoaders.collate_fn)

            testset = DataLoaders.CosmoDataGridData(
                config=config,
                station_data_dict=data_dictionary,
                files=test_fold,
                featureScaling=featureScaleFunctions,
                time_invariant_data=time_invarian_data)
            testloader = DataLoader(testset, batch_size=config['batch_size'], shuffle=True,
                                    num_workers=config['n_loaders'], collate_fn=DataLoaders.collate_fn)

        # initialize network, optimizer and loss function
        net = Baseline.model_factory(ModelDict, trainset.n_parameters, trainset.n_grid_time_invariant_parameters,
                                     config['grid_size'], config['prediction_times'])
        # store class name
        experiment_info['model_class'] = net.__class__.__name__

        if torch.cuda.device_count() > 1:
            net = nn.DataParallel(net)

        if torch.cuda.is_available():
            net.cuda()

        # load number of train and test samples
        n_train_samples, n_test_samples = len(train_fold), len(test_fold)

        optimizer, scheduler = ModelUtils.initializeOptimizer(optimizer_config, net)
        criterion = nn.MSELoss()

        # keep number of processed smaples over all epochs for tensorboard
        processed_train_samples_global = 0
        processed_test_samples_global = 0

        # start learning
        for epoch in range(config['epochs']):
            epoch_train_time = np.zeros((5,))
            epoch_start_time = time()
            print('Epoch: ' + str(epoch + 1) + '\n------------------------------------------------------------')

            # adapt learning rate and store information in experiment attributes
            if scheduler is not None:
                scheduler.step()
                if run == 0: experiment_info['scheduled_learning_rates'] += scheduler.get_lr()
                print('Using learning rate %s' % str(scheduler.get_lr()))

            # TRAINING
            # initialize variables for epoch statistics
            LABELS, MODELoutputs, COSMOoutputs = None, None, None
            processed_train_samples = 0
            net.train(True)

            train_start_time = time()
            # loop over complete train set
            for i, data in enumerate(trainloader, 0):
                time_start = time()
                try:
                    # get training batch, e.g. label, cosmo-1 output and time inv. features for station
                    DATA = data
                    # DATA has only length 4 if we do not use the station time invariant features
                    if len(DATA) == 4:
                        Blabel, Bip2d, BTimeData, init_station_temp = DATA
                        station_time_inv_input = None
                    elif len(DATA) == 5:
                        Blabel, Bip2d, BTimeData, StationTimeInv, init_station_temp = DATA
                        station_time_inv_input = ModelUtils.getVariable(StationTimeInv).float()
                    else:
                        raise Exception('Unknown data format for training...')
                    input = ModelUtils.getVariable(Bip2d).float()
                    time_data = ModelUtils.getVariable(BTimeData).float()
                    target = ModelUtils.getVariable(Blabel).float()

                except TypeError:
                    # when the batch size is small, it could happen, that all labels have been corrupted and therefore
                    # collate_fn would return an empty list
                    print('Value error...')
                    continue
                time_after_data_preparation = time()

                processed_train_samples += len(Blabel)

                optimizer.zero_grad()
                out = net(input, time_data, station_time_inv_input)
                time_after_forward_pass = time()
                loss = criterion(out, target)
                loss.backward()
                optimizer.step()
                time_after_backward_pass = time()

                if LABELS is None:
                    LABELS = Blabel.data
                    MODELoutputs = out.data
                    COSMOoutputs = init_station_temp[2].data
                else:
                    LABELS = np.vstack((LABELS, Blabel.data))
                    MODELoutputs = np.vstack((MODELoutputs, out.data))
                    COSMOoutputs = np.vstack((COSMOoutputs, init_station_temp[2].data))

                time_after_label_stack = time()

                if (i + 1) % 64 == 0:

                    print('Sample: %s \t Loss: %s' % (processed_train_samples, float(np.sqrt(loss.data))))

                    # ============ TensorBoard logging ============#
                    # (1) Log the scalar values
                    info = {
                        setting_string: np.sqrt(loss.item()),
                    }

                    for tag, value in info.items():
                        train_logger.scalar_summary(tag, value, processed_train_samples_global + processed_train_samples)

                    # (2) Log values and gradients of the parameters (histogram)
                    for tag, value in net.named_parameters():
                        tag = tag.replace('.', '/')
                        train_logger.histo_summary(tag, ModelUtils.to_np(value), i + 1)
                        train_logger.histo_summary(tag + '/grad', ModelUtils.to_np(value.grad), i + 1)

                    epoch_train_time += np.array((time_start - time_end,
                                                  time_after_data_preparation - time_start,
                                                  time_after_forward_pass - time_after_data_preparation,
                                                  time_after_backward_pass - time_after_forward_pass,
                                                  time_after_label_stack - time_after_backward_pass))

                time_end = time()

            # calculate error statistic of current epoch
            diff_model = MODELoutputs - LABELS
            diff_cosmo = COSMOoutputs - LABELS
            epoch_train_rmse_model = np.apply_along_axis(func1d=ModelUtils.rmse, arr=diff_model, axis=0)
            epoch_train_rmse_cosmo = np.apply_along_axis(func1d=ModelUtils.rmse, arr=diff_cosmo, axis=0)

            # update global processed samples
            processed_train_samples_global += processed_train_samples

            if np.isnan(epoch_train_rmse_model).any():
                print("Learning rate too large resulted in NaN-error while training. Stopped training...")
                return
            # print epoch training times
            print('Timing: Waiting on data=%s, Data Preparation=%s,'
                  'Forward Pass=%s, Backward Pass=%s, Data Stacking=%s' % tuple(list(epoch_train_time / len(epoch_train_time))))

            # RMSE of epoch
            print('Train/test statistic for epoch: %s' % str(epoch + 1))
            print('Train RMSE COSMO: ' , ", ".join(["T=%s: %s" % (idx, epoch_train_rmse_cosmo[idx]) for idx in range(len(epoch_train_rmse_cosmo))]))
            print('Train RMSE Model: ' , ", ".join(["T=%s: %s" % (idx, epoch_train_rmse_model[idx]) for idx in range(len(epoch_train_rmse_model))]))
            sys.stdout.flush()

            train_time = time() - train_start_time

            # TESTING
            test_start_time = time()

            LABELS, MODELoutputs, COSMOoutputs, STATION = None, None, None, None
            processed_test_samples = 0
            net.eval()
            for i, data in enumerate(testloader, 0):
                try:
                    # get training batch, e.g. label, cosmo-1 output and time inv. features for station
                    DATA = data
                    # DATA has only length 4 if we do not use the station time invariant features
                    if len(DATA) == 4:
                        Blabel, Bip2d, BTimeData, init_station_temp = DATA
                        station_time_inv_input = None
                    elif len(DATA) == 5:
                        Blabel, Bip2d, BTimeData, StationTimeInv, init_station_temp = DATA
                        station_time_inv_input = ModelUtils.getVariable(StationTimeInv).float()
                    else:
                        raise Exception('Unknown data format for training...')
                    input = ModelUtils.getVariable(Bip2d).float()
                    time_data = ModelUtils.getVariable(BTimeData).float()
                    target = ModelUtils.getVariable(Blabel).float()

                except TypeError:
                    # when the batch size is small, it could happen, that all labels have been corrupted and therefore
                    # collate_fn would return an empty list
                    print('Value error...')
                    continue

                processed_test_samples += len(Blabel)

                out = net(input, time_data, station_time_inv_input)
                loss = criterion(out, target)

                if LABELS is None:
                    LABELS = Blabel.data
                    MODELoutputs = out.data
                    COSMOoutputs = init_station_temp[2].data
                    STATION = init_station_temp[1].data
                else:
                    LABELS = np.vstack((LABELS, Blabel.data))
                    MODELoutputs = np.vstack((MODELoutputs, out.data))
                    COSMOoutputs = np.vstack((COSMOoutputs, init_station_temp[2].data))
                    STATION = np.hstack((STATION, init_station_temp[1].data))

                if i % 16:
                    # ============ TensorBoard logging ============#
                    # (1) Log the scalar values
                    info = {
                        setting_string: np.sqrt(loss.item()),
                    }

                    for tag, value in info.items():
                        test_logger.scalar_summary(tag, value, processed_test_samples_global + processed_test_samples)

            # calculate error statistic of current epoch
            diff_model = MODELoutputs - LABELS
            diff_cosmo = COSMOoutputs - LABELS

            # rmse
            epoch_test_rmse_model = np.apply_along_axis(func1d=ModelUtils.rmse, arr=diff_model, axis=0)
            epoch_test_rmse_cosmo = np.apply_along_axis(func1d=ModelUtils.rmse, arr=diff_cosmo, axis=0)
            overall_test_rmse_model = ModelUtils.rmse(diff_model)
            overall_test_rmse_cosmo = ModelUtils.rmse(diff_cosmo)

            # mae
            epoch_test_mae_model = np.apply_along_axis(func1d=ModelUtils.mae, arr=diff_model, axis=0)
            epoch_test_mae_cosmo = np.apply_along_axis(func1d=ModelUtils.mae, arr=diff_cosmo, axis=0)
            overall_test_mae_model = ModelUtils.mae(diff_model)
            overall_test_mae_cosmo = ModelUtils.mae(diff_cosmo)

            # calculate per station rmse if desired (especially for K-fold station generalization experiment
            if "per_station_rmse" in config:
                max_station_id = 1435

                squared_errors_per_epoch = np.array((np.square(diff_model), np.square(diff_cosmo))).squeeze()

                # the highest index of data is 1435, thus we expect at least 1435 entries, which we can access by
                # station id
                test_samples_per_station = np.bincount(STATION, minlength=max_station_id+1)
                model_squared_error_per_station = np.bincount(STATION, weights=squared_errors_per_epoch[0], minlength=max_station_id+1)
                cosmo_squared_error_per_station = np.bincount(STATION, weights=squared_errors_per_epoch[1], minlength=max_station_id+1)

                # set division by zero/NaN warning to 'ignore'
                np.seterr(divide='ignore', invalid='ignore')

                # calculate rmse per station
                rmse_per_station = np.vstack((np.sqrt(np.divide(model_squared_error_per_station, test_samples_per_station)),
                                              np.sqrt(np.divide(cosmo_squared_error_per_station, test_samples_per_station)))).T

                # set division by zero/NaN warning to 'warn'
                np.seterr(divide='warn', invalid='warn')

            # update global processed samples
            processed_test_samples_global += processed_test_samples

            # RMSE of epoch
            print('Test RMSE COSMO: ', ", ".join(
                ["T=%s: %s" % (idx, epoch_test_rmse_cosmo[idx]) for idx in range(len(epoch_test_rmse_cosmo))]),
                  " (Overall: %s" % overall_test_rmse_cosmo)
            print('Test RMSE Model: ' , ", ".join(["T=%s: %s" % (idx, epoch_test_rmse_model[idx]) for idx in range(len(epoch_test_rmse_model))]),
                  " (Overall: %s" % overall_test_rmse_model)
            # mae of epoch
            print('Test MAE COSMO: ', ", ".join(
                ["T=%s: %s" % (idx, epoch_test_mae_cosmo[idx]) for idx in range(len(epoch_test_mae_cosmo))]),
                  " (Overall: %s" % overall_test_mae_cosmo)
            print('Test MAE Model: ' , ", ".join(["T=%s: %s" % (idx, epoch_test_mae_model[idx]) for idx in range(len(epoch_test_mae_model))]),
                  " (Overall: %s" % overall_test_mae_model)

            sys.stdout.flush()

            test_time = time() - test_start_time

            # time for epoch
            epoch_time = time() - epoch_start_time

            # update error statistics
            error_statistics = ModelUtils.updateErrorStatistic(error_statistics,
                                                               np.array([epoch_train_rmse_model, epoch_test_rmse_model])[None, None, ...],
                                                               run, epoch, config['prediction_times'])
            # update run times statistic
            run_times = ModelUtils.updateRuntimeStatistic(run_times, np.array([epoch_time, train_time, test_time])[None, None, ...],
                                                          run, epoch)
            # update skip statistic
            skip_statistics = ModelUtils.updateSkipStatistic(skip_statistics,
                                                             np.array([n_train_samples, processed_train_samples,
                                                                       n_test_samples, processed_test_samples])[None, None, ...],
                                                             run, epoch)

            # update per station rmse data array over runs if desired (especially for K-fold station generalization experiment
            if "per_station_rmse" in config:
                error_per_station_statistics = ModelUtils.updatePerStationErrorStatistic(error_per_station_statistics, rmse_per_station, run, epoch, np.arange(max_station_id+1))

            # store model if it was the best yes
            is_best = overall_test_rmse_model <= best_epoch_test_rmse
            best_epoch_test_rmse = min(overall_test_rmse_model, best_epoch_test_rmse)
            ModelUtils.save_checkpoint({
                'epoch': epoch,
                'run': run,
                'arch': net.__class__.__name__,
                'state_dict': net.state_dict(),
                'overall_test_rmse': overall_test_rmse_model,
                'lead_test_rmse' : overall_test_rmse_model,
                'best_epoch_test_rmse': best_epoch_test_rmse,
                'optimizer': optimizer.state_dict(),
            }, is_best, output_path + '/stored_models/run_%s' % run)

            # flush output to see progress
            sys.stdout.flush()

    # update statistics dict
    ModelUtils.get_model_details(experiment_info, net, optimizer, criterion)

    # complete program runtime
    experiment_info['program_runtime'] = time() - program_start_time

    # generate data set of all experiment statistics and additional information
    experiment_statistic = xr.Dataset({
        'error_statistic' : error_statistics,
        'run_time_statistic': run_times,
        'samples_statistic' : skip_statistics}).assign_attrs(experiment_info)

    # dump experiment statistic
    with open(output_path + '/experiment_statistic.pkl', 'wb') as handle:
        pkl.dump(experiment_statistic, handle, protocol=pkl.HIGHEST_PROTOCOL)

    if 'per_station_rmse' in config:
        # dump experiment statistic
        with open(output_path + '/rmse_per_station.pkl', 'wb') as handle:
            pkl.dump(error_per_station_statistics, handle, protocol=pkl.HIGHEST_PROTOCOL)

    # print program execution time
    m, s = divmod(experiment_info['program_runtime'], 60)
    h, m = divmod(m, 60)
    print('Experiment has successfully finished in %dh %02dmin %02ds' % (h, m, s))


from datetime import datetime
import os.path
import time
import sys
import tensorflow as tf
import numpy as np
import importlib
import argparse
import facenet
import random
import itertools
from sklearn.metrics.pairwise import euclidean_distances
from sklearn.metrics.pairwise import pairwise_distances
import recordRetriever as rr
from copy import deepcopy
from collections import Counter

def main(args):
    tf.reset_default_graph()
    network = importlib.import_module(args.model_def)

#    subdir = datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S')
#    log_dir = os.path.join(os.path.expanduser(args.logs_base_dir), subdir)
#    if not os.path.isdir(log_dir):  # Create the log directory if it doesn't exist
#        os.makedirs(log_dir)
#    model_dir = os.path.join(os.path.expanduser(args.models_base_dir), subdir)
#    if not os.path.isdir(model_dir):  # Create the model directory if it doesn't exist
#        os.makedirs(model_dir)

    # Write arguments to a text file
#    facenet.write_arguments_to_file(args, os.path.join(log_dir, 'arguments.txt'))

    # Store some git revision info in a text file in the log directory
#    src_path = ''
#    facenet.store_revision_info(src_path, log_dir, ' '.join(sys.argv))

    np.random.seed(int(time.time()))

    DB_TF_RECORD_PATH_LIST = rr.selectedDiseasePath(args.select_disease, args.select_anatomy,
                                            args.with_biopsy, args.with_mask, args.root)
    """Get shuffled train and validation tf record path"""
#    TRAIN_TF_RECORD_PATH_LIST, VALID_TF_RECORD_PATH_LIST, overall_dict, training_dict, test_dict = \
#                                        rr.split_training_validation_sets(TF_RECORD_PATH_LIST, args.valid_fraction)

    """ensure that the dataset split is correct"""
#    for key, value in overall_dict.items():
#        assert overall_dict[key] == training_dict[key] + test_dict[key], 'Train data + Test data != Overall data'

    """ CHECKING VALIDITY OF FILE PATH """
#    rr.check_files(TRAIN_TF_RECORD_PATH_LIST, VALID_TF_RECORD_PATH_LIST)

    """print out the stats of my dataset"""
#    TOTAL_TRAIN_DATA = len(TRAIN_TF_RECORD_PATH_LIST)
#    TOTAL_VALID_DATA = len(VALID_TF_RECORD_PATH_LIST)
#    TOTAL_DATA = len(TF_RECORD_PATH_LIST)

#    print ('total train data: ', TOTAL_TRAIN_DATA)
#    print ('total valid data: ', TOTAL_VALID_DATA)
#    print ('total data available: ', TOTAL_DATA)

    """organize into list of disease classes for sample_disease"""
#    train_set_dict = {}
#    for path in TRAIN_TF_RECORD_PATH_LIST:
#        disease_name = path.split('/')[-1].split('_')[0]
#        if disease_name not in train_set_dict:
#            train_set_dict[disease_name] = [path]
#        else:
#            train_set_dict[disease_name].append(path)
#    train_set = facenet.get_dataset(train_set_dict)

    """to get the number of brains in the most common disease for augmentation"""
#    max_brains_in_single_disease = max([value for key, value in training_dict.items()])
#    print('max brains in single disease :', max_brains_in_single_disease)

    """to multiply more cases for the rare cases"""
#    for imageclass in train_set:
#        nbr_brains_in_imageclass = len(imageclass.image_paths)
#        multiplier = int(min(4, max_brains_in_single_disease/nbr_brains_in_imageclass))
#        imageclass.image_paths *= multiplier
#

#        curr_brain_records = imageclass.image_paths.copy()
#        for augmentation_step in range(multiplier-1):
#            imageclass.image_paths += list(map(lambda x: x+'aug'+str(augmentation_step+1), curr_brain_records))
#        print('mutliplier: ', multiplier)
#        print('number of brains in {} changed from {} to {}'.format(imageclass.name, nbr_brains_in_imageclass,
#              nbr_brains_in_imageclass * multiplier))
#    print('disease in train set after augmenting, if there is')
#    for disease in train_set:
#        print(disease.name,': ',len(disease.image_paths))

    if args.pretrained_model:
        print('Pre-trained model: %s' % os.path.expanduser(args.pretrained_model))

    with tf.Graph().as_default():

#        last_run_step = 0

#        if args.pretrained_model:
#            last_run_step += int(args.pretrained_model.split("-")[-1])

#        global_step = tf.Variable(max(0, last_run_step), trainable=False)

#        Placeholder for TL_scalar
#        TL_scalar = tf.constant(args.TL_scalar, tf.float32, name = 'TL_scalar')

        # Placeholder for the learning rate
#        with tf.name_scope('learning_rate'):
#            learning_rate_placeholder = tf.placeholder(tf.float32, name='learning_rate')

#        with tf.name_scope('labels'):
#            labels_placeholder = tf.placeholder(tf.int32, shape=[None], name = 'labels')

        with tf.name_scope('to_train_BN'):
            phase_train_BN_placeholder = tf.placeholder(tf.bool, name='to_train_BN')

        with tf.name_scope('to_train_DO'):
            phase_train_DO_placeholder = tf.placeholder(tf.bool, name='to_train_DO')

        with tf.name_scope('to_train'):
            phase_train_placeholder = tf.placeholder(tf.bool, name='to_train')

        with tf.name_scope('filenames'):
            filenames_placeholder = tf.placeholder(tf.string, shape=[None], name = 'filenames')

        iterator, next_batch = rr.dataPreparation(filenames_placeholder, batch_size = args.batch_size,
                                            image_shape=[args.slices, args.image_size, args.image_size, 1],
                                            new_image_shape = [args.slices, args.new_image_size, args.new_image_size, 1],
                                            training = phase_train_placeholder)

        t1_op, t2_op, _, _, _, _, _, mask_op, _, _, _, \
        _, label_op, _, filename_op, _, _, _, _, _ = next_batch

        t1t2mask_images = tf.concat([t1_op, t2_op, mask_op], axis = 4)

        # Build the inference graph
        prelogits, end_points = network.inference(t1t2mask_images, args.keep_probability, args,
                                    phase_train_BN = phase_train_BN_placeholder,
                                    phase_train_DO = phase_train_DO_placeholder,
                                    bottleneck_layer_size=args.embedding_size, weight_decay=args.weight_decay)

        tensor_lst = []
        for key, value in end_points.items():
            tensor_lst.append(value)

        print('1')

        embeddings = tf.nn.l2_normalize(prelogits, 1, 1e-10, name='embeddings')

#        triplet_loss, fraction_positive_triplets = facenet.batch_all_triplet_loss(labels_placeholder, embeddings, args, squared=True)

#        if args.lossless == True:
#            TL_scalar = 1.0

#        triplet_loss = tf.scalar_mul(TL_scalar, triplet_loss)

#        learning_rate = tf.train.exponential_decay(learning_rate_placeholder, global_step,
#            args.learning_rate_decay_epochs*args.epoch_size, args.learning_rate_decay_factor, staircase=True)

        # Calculate the total losses
#        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
#        total_loss = tf.add_n([triplet_loss] + regularization_losses, name='total_loss')

        print('2')

#        if args.optimizer =='ADAM':
#            opt = tf.train.AdamOptimizer(learning_rate, beta1=0.9, beta2=0.999, epsilon=0.1).minimize(total_loss, global_step = global_step)
#        else:
#            raise ValueError('Invalid optimization algorithm')

        # Create a saver
        saver = tf.train.Saver(tf.trainable_variables(), max_to_keep=1)

        # Build the summary operation based on the TF collection of Summaries.
        summary_op = tf.summary.merge_all()

        print('3')

        # Start running operations on the Graph.
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

        print('4')

        # Initialize variables
        sess.run(tf.global_variables_initializer())

        sess.run(tf.local_variables_initializer())

        print('5')

#        summary_writer = tf.summary.FileWriter(log_dir, sess.graph)
        print('6')

        best_top1_accu = 0

        with sess.as_default():

            if args.pretrained_model:
                print('Restoring pretrained model: %s' % args.pretrained_model)
                saver = tf.train.import_meta_graph(os.path.expanduser(args.metadata))
                saver.restore(sess, os.path.expanduser(args.pretrained_model))

            print('7')

            # Training and validation loop
            epoch = 0
#            while epoch < args.max_nrof_epochs:
#            step = sess.run(global_step, feed_dict=None)
#            epoch = step // args.epoch_size
            # Train for one epoch
            generate_emb(args, DB_TF_RECORD_PATH_LIST, sess, phase_train_BN_placeholder, phase_train_DO_placeholder, embeddings, summary_op,
                  iterator, t1_op, t2_op, mask_op, label_op, filename_op, filenames_placeholder, t1t2mask_images, prelogits, tensor_lst, phase_train_placeholder)

#                #select data from train set to compute training accuracy
#                selected_set = train_data_for_accuracy(train_set, len(VALID_TF_RECORD_PATH_LIST))
#                print('selected train eval set...')
#                print(Counter(list(map(lambda x: get_disease_from_path(x), selected_set))))
#
#                # Evaluate on test set
#                top1_accu = evaluate(args, sess, selected_set, VALID_TF_RECORD_PATH_LIST, embeddings, learning_rate_placeholder,
#                                     phase_train_BN_placeholder, phase_train_DO_placeholder, log_dir, step, summary_writer, iterator,
#                                     t1_op, t2_op, mask_op, label_op, filename_op, filenames_placeholder, t1t2mask_images, prelogits, tensor_lst, phase_train_placeholder)
#
#                # if accuracy is best so far, save variables and the metagraph if it doesn't exist already
#                if top1_accu > best_top1_accu:
#                    best_top1_accu = top1_accu
#                    save_variables_and_metagraph(sess, saver, summary_writer, model_dir, subdir, step)
#
#                print('best top1 accuracy: {}'.format(best_top1_accu))

    return None

def generate_emb(args, DB_TF_RECORD_PATH_LIST, sess, phase_train_BN_placeholder, phase_train_DO_placeholder,
          embeddings, summary_op,
          iterator, t1_op, t2_op, mask_op, label_op, filename_op, filenames_placeholder, t1t2mask_images, prelogits, tensor_lst, phase_train_placeholder):

    batch_size = args.batch_size
    sess.run(iterator.initializer, feed_dict={filenames_placeholder: DB_TF_RECORD_PATH_LIST,
                                              phase_train_placeholder: False})
    nrof_images = len(DB_TF_RECORD_PATH_LIST)
    emb_array = np.zeros((nrof_images, args.embedding_size))
    nrof_batches = int(np.ceil(nrof_images / batch_size))
    label_check_array = np.zeros((nrof_images,))

    for i in range(nrof_batches):
        batch_size = min(nrof_images-i*batch_size, batch_size)
        labels = list(range(i * args.batch_size, i * args.batch_size + batch_size))
        print('labels: ', labels)
        emb = sess.run([embeddings], feed_dict={phase_train_BN_placeholder: False,
                                                phase_train_DO_placeholder: False})
        emb_array[labels,:] = emb
        label_check_array[labels] = 1

    print('asserting')
    assert(np.all(label_check_array==1))
    print('preparing data')
    database_disease = list(map(lambda x: get_disease_from_path(x), DB_TF_RECORD_PATH_LIST))
    database_disease_labels = list(map(lambda x: rr.ACRONYM_LABEL_MAPPER[x], database_disease))
    database_disease_labels_np = np.asarray(database_disease_labels)
    print('saving...')
    np.savez('20180817-231301-'+'5400-'+'database.npz', disease_emb=emb_array, disease_labels=database_disease_labels_np)
    print('saved')

    return None

def get_disease_from_path(path):
    """pass in tfrecord path to get the disease
    args: tfrecord path
    Returns: disease name"""
    disease_name = path.split('/')[-1].split('_')[0]

    return disease_name

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--logs_base_dir', type=str,
        help='Directory where to write event logs.', default='logs/facenet')
    parser.add_argument('--models_base_dir', type=str,
        help='Directory where to write trained models and checkpoints.', default='models/facenet')
    parser.add_argument('--gpu_memory_fraction', type=float,
        help='Upper bound on the amount of GPU memory that will be used by the process.', default=1.0)
    parser.add_argument('--pretrained_model', type=str,
#                                help='Load a pretrained model before training starts.')
        help='Load a pretrained model before training starts.', default = 'models/facenet/20180817-231301/model-20180817-231301.ckpt-5400')
    parser.add_argument('--metadata', type=str,
#                                help ='Load pretrained model metadata before training starts.')
        help ='Load pretrained model metadata before training starts.', default = 'models/facenet/20180817-231301/model-20180817-231301.meta')
    parser.add_argument('--model_def', type=str,
        help='Model definition. Points to a module containing the definition of the inference graph.', default='nn3')
    parser.add_argument('--max_nrof_epochs', type=int,
        help='Number of epochs to run.', default=2000)
    parser.add_argument('--batch_size', type=int,
        help='Number of images to process in a batch.', default=9)
    parser.add_argument('--image_size', type=int,
        help='Image size (height, width) in pixels.', default=320)
    parser.add_argument('--slices', type=int,
        help='number of slices in patients brain.', default=24)
    parser.add_argument('--brains_per_disease', type=int,
        help='max number of brains per disease.', default=3)
    parser.add_argument('--epoch_size', type=int,
        help='Number of batches per epoch.', default=200)
    parser.add_argument('--alpha', type=float,
        help='Positive to negative triplet distance margin.', default=0.3)
    parser.add_argument('--embedding_size', type=int,
        help='Dimensionality of the embedding.', default=1024)
    parser.add_argument('--keep_probability', type=float,
        help='Keep probability of dropout for the fully connected layer(s).', default=1.0)
    parser.add_argument('--weight_decay', type=float,
        help='L2 weight regularization.', default=1e-4)
    parser.add_argument('--optimizer', type=str, choices=['ADAGRAD', 'ADADELTA', 'ADAM', 'RMSPROP', 'MOM'],
        help='The optimization algorithm to use', default='ADAM')
    parser.add_argument('--learning_rate', type=float,
        help='Initial learning rate.', default=0.01)
    parser.add_argument('--learning_rate_decay_epochs', type=int,
        help='Number of epochs between learning rate decay.', default=100)
    parser.add_argument('--learning_rate_decay_factor', type=float,
        help='Learning rate decay factor.', default= 0.8)
    parser.add_argument('--evaluation_metric', type=str,
        help ='The type of evaluation metric (cosine, euclidean or etc)', default = 'euclidean')
    parser.add_argument('--TL_scalar', type=float,
        help ='The amount to scale the triplet loss to match regularization loss', default = 10.0)
    parser.add_argument('--select_disease', type=None,
        help ='Which disease to select? Empty list will select all disease', default = [])
    parser.add_argument('--select_anatomy', type=None,
        help ='Which anatomy to select? Empty list will select all anatomies', default = [])
    parser.add_argument('--with_biopsy', type=bool,
        help ='select only patients with biopsy?', default = False)
    parser.add_argument('--with_mask', type=bool,
        help ='select only patients with mask?', default = True)
    parser.add_argument('--valid_fraction', type=float,
        help ='percentage of validation data', default = 0.05)
    parser.add_argument('--root', type=str,
        help ='The root directory where all TFrecords are stored',
        default = '/data/tumor/dicoms/TFRECORD_GRAND_5')
    parser.add_argument('--new_image_size', type=int,
        help ='resize t1, t2, mask to this size', default = 160)
    parser.add_argument('--lossless', type=bool,
        help ='use lossless triplet loss?', default = False)
    return parser.parse_args(argv)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))
#    print(emb_array)
#    assert(np.all(label_check_array==1))
#    database_disease = list(map(lambda x: get_disease_from_path(x), DB_TF_RECORD_PATH_LIST))
#    database_disease_labels = list(map(lambda x: rr.ACRONYM_LABEL_MAPPER[x], database_disease))
#    database_disease_labels_np = np.asarray(database_disease_labels)
#    np.savez('database.npz', disease_emb=emb_array, disease_labels=database_disease_labels_np)


from models import *
from database import db_session
from pathlib import Path

class GroupInfo:
    def __init__(self, read_source):
        self.read_source = read_source

    def run(self):
        print("Agrupando objetos")
        devices = db_session.query(Device).all()
        for device in devices:
            path = Path(device.folder)
            device.group = path.parent
            db_session.add(device)
        db_session.commit()


# -*- coding: utf-8 -*-
import logging
import sys

from PyQt5.QtCore import (Qt, QModelIndex)
from PyQt5.QtGui import QKeySequence
from PyQt5.QtWidgets import (QMainWindow, QAction, QApplication,
                             QSplitter, QGroupBox, QFormLayout,
                             QLabel, QBoxLayout, QWidget, QGridLayout,
                             QStyle, QTreeView, QFileDialog,
                             QMessageBox, QTabWidget)

from mhw_armor_edit import AmDat
from mhw_armor_edit.assets import Definitions
from mhw_armor_edit.tree import ArmorSetTreeModel, ArmorSetNode, ArmorListModel
from mhw_armor_edit.view_ctrl import (ComboBoxWidgetCtrl, SpinBoxWidgetCtrl,
                                      LabelWidgetCtrl,
                                      PieceViewCtrl)

log = logging.getLogger()
logging.basicConfig(level=logging.DEBUG)

def groupbox(layout, title=None):
    box = QGroupBox()
    box.setStyleSheet("QGroupBox {font-weight:bold}")
    if title:
        box.setTitle(title)
        box.setFlat(True)
    box.setLayout(layout)
    return box, layout

def tree_index_is_root(index: QModelIndex):
    return not index.isValid()

def create_action(icon, title, handler, shortcut=None):
    action = QAction(icon, title)
    if shortcut is not None:
        action.setShortcut(shortcut)
    action.triggered.connect(handler)
    return action

class FileModel:
    def __init__(self, path, data):
        self.path = path
        self.data = data

    def save(self):
        with open(self.path, "wb") as fp:
            fp.write(self.data.data)

    @classmethod
    def load(cls, path):
        with open(path, "rb") as fp:
            data = AmDat.make(fp)
        return cls(path, data)

class ArmorPieceWidget(QWidget):
    def __init__(self, view, *args, **kwargs):
        super().__init__(*args, *kwargs)
        self._init(view)

    def _init(self, view):
        layout = QBoxLayout(QBoxLayout.TopToBottom)
        self.setLayout(layout)
        self._init_basic(layout, view)
        self._init_resistance(layout, view)
        self._init_gem_slots(layout, view)
        self._init_set_skills(layout, view)
        self._init_piece_skills(layout, view)

    def _init_piece_skills(self, layout, view):
        box, box_layout = groupbox(QFormLayout(), "Piece Skills")
        layout.addWidget(box, 0)
        view.skill1.ctrl = ComboBoxWidgetCtrl(Definitions.skill, completer=True)
        view.skill1_lvl.ctrl = SpinBoxWidgetCtrl(0, 10)
        view.skill2.ctrl = ComboBoxWidgetCtrl(Definitions.skill, completer=True)
        view.skill2_lvl.ctrl = SpinBoxWidgetCtrl(0, 10)
        view.skill3.ctrl = ComboBoxWidgetCtrl(Definitions.skill, completer=True)
        view.skill3_lvl.ctrl = SpinBoxWidgetCtrl(0, 10)
        box_layout.addRow(QLabel("Skill 1"), view.skill1.ctrl.widget)
        box_layout.addRow(QLabel("Level"), view.skill1_lvl.ctrl.widget)
        box_layout.addRow(QLabel("Skill 2"), view.skill2.ctrl.widget)
        box_layout.addRow(QLabel("Level"), view.skill2_lvl.ctrl.widget)
        box_layout.addRow(QLabel("Skill 3"), view.skill3.ctrl.widget)
        box_layout.addRow(QLabel("Level"), view.skill3_lvl.ctrl.widget)

    def _init_set_skills(self, layout, view):
        box, box_layout = groupbox(QFormLayout(), "Set Skills")
        layout.addWidget(box, 0)
        view.set_skill1.ctrl = ComboBoxWidgetCtrl(Definitions.skill, completer=True)
        view.set_skill1_lvl.ctrl = SpinBoxWidgetCtrl(0, 10)
        view.set_skill2.ctrl = ComboBoxWidgetCtrl(Definitions.skill, completer=True)
        view.set_skill2_lvl.ctrl = SpinBoxWidgetCtrl(0, 10)
        box_layout.addRow(QLabel("Skill 1"), view.set_skill1.ctrl.widget)
        box_layout.addRow(QLabel("Level"), view.set_skill1_lvl.ctrl.widget)
        box_layout.addRow(QLabel("Skill 2"), view.set_skill2.ctrl.widget)
        box_layout.addRow(QLabel("Level"), view.set_skill2_lvl.ctrl.widget)

    def _init_gem_slots(self, layout, view):
        box, box_layout = groupbox(QFormLayout(), "Gem Slots")
        layout.addWidget(box, 0)
        view.num_gem_slots.ctrl = ComboBoxWidgetCtrl(Definitions.gem_slot)
        view.gem_slot1_lvl.ctrl = ComboBoxWidgetCtrl(Definitions.gem_slot)
        view.gem_slot2_lvl.ctrl = ComboBoxWidgetCtrl(Definitions.gem_slot)
        view.gem_slot3_lvl.ctrl = ComboBoxWidgetCtrl(Definitions.gem_slot)
        box_layout.addRow(QLabel("Active slots"), view.num_gem_slots.ctrl.widget)
        box_layout.addRow(QLabel("Slot 1 Level"), view.gem_slot1_lvl.ctrl.widget)
        box_layout.addRow(QLabel("Slot 2 Level"), view.gem_slot2_lvl.ctrl.widget)
        box_layout.addRow(QLabel("Slot 3 Level"), view.gem_slot3_lvl.ctrl.widget)

    def _init_resistance(self, layout, view):
        box, box_layout = groupbox(QFormLayout(), "Resistance")
        layout.addWidget(box, 0)
        view.fire_res.ctrl = SpinBoxWidgetCtrl(-127, 127)
        view.water_res.ctrl = SpinBoxWidgetCtrl(-127, 127)
        view.thunder_res.ctrl = SpinBoxWidgetCtrl(-127, 127)
        view.ice_res.ctrl = SpinBoxWidgetCtrl(-127, 127)
        view.dragon_res.ctrl = SpinBoxWidgetCtrl(-127, 127)
        box_layout.addRow(QLabel("Fire"), view.fire_res.ctrl.widget)
        box_layout.addRow(QLabel("Water"), view.water_res.ctrl.widget)
        box_layout.addRow(QLabel("Thunder"), view.thunder_res.ctrl.widget)
        box_layout.addRow(QLabel("Ice"), view.ice_res.ctrl.widget)
        box_layout.addRow(QLabel("Dragon"), view.dragon_res.ctrl.widget)

    def _init_basic(self, layout, view):
        section_box, section_layout = groupbox(QGridLayout())
        layout.addWidget(section_box)
        section_layout.setColumnStretch(0, 0)
        section_layout.setColumnStretch(1, 1)
        section_layout.setColumnStretch(2, 0)
        section_layout.setColumnStretch(3, 1)

        view.set_name.ctrl = LabelWidgetCtrl(Definitions.set)
        section_layout.addWidget(QLabel("Set:"), 0, 0, Qt.AlignLeft)
        section_layout.addWidget(view.set_name.ctrl.widget, 0, 1, Qt.AlignLeft)

        section_layout.addWidget(QLabel("Index:"), 0, 2, Qt.AlignLeft)
        view.index.ctrl = LabelWidgetCtrl([])
        section_layout.addWidget(view.index.ctrl.widget, 0, 3, Qt.AlignLeft)

        section_layout.addWidget(QLabel("Variant:"), 2, 0, Qt.AlignLeft)
        view.variant.ctrl = LabelWidgetCtrl(Definitions.variant)
        section_layout.addWidget(view.variant.ctrl.widget, 2, 1, Qt.AlignLeft)
        section_layout.addWidget(QLabel("Equip Slot:"), 2, 2, Qt.AlignLeft)
        view.equip_slot.ctrl = LabelWidgetCtrl(Definitions.equip_slot)
        section_layout.addWidget(view.equip_slot.ctrl.widget, 2, 3, Qt.AlignLeft)

        section_box, section_layout = groupbox(QFormLayout(), "Basic")
        layout.addWidget(section_box, 0)
        view.defense.ctrl = SpinBoxWidgetCtrl(0, 0xffff)
        section_layout.addRow(QLabel("Defense"), view.defense.ctrl.widget)
        view.rarity.ctrl = ComboBoxWidgetCtrl(Definitions.rarity)
        section_layout.addRow(QLabel("Rarity"), view.rarity.ctrl.widget)
        view.cost.ctrl = SpinBoxWidgetCtrl(0, 0xffff)
        section_layout.addRow(QLabel("Cost"), view.cost.ctrl.widget)

class StructuredEditorWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.file_model = None
        self.current_piece_view_ctrl = PieceViewCtrl()
        self.init_actions()
        self.init_toolbar()
        self.init_menubar()
        self.init_ui()
        self.current_piece_view_ctrl.update(None)

    def get_icon(self, name):
        return self.style().standardIcon(name)

    def init_actions(self):
        self.open_file_action = create_action(
            self.get_icon(QStyle.SP_DialogOpenButton),
            "Open file ...",
            self.handle_open_file_action,
            QKeySequence.Open)
        self.save_file_action = create_action(
            self.get_icon(QStyle.SP_DialogSaveButton),
            "Save ...",
            self.handle_save_file_action,
            QKeySequence.Save)
        self.save_file_as_action = create_action(
            self.get_icon(QStyle.SP_DialogSaveButton),
            "Save as ...",
            self.handle_save_file_as_action,
            QKeySequence.SaveAs)
        self.close_file_action = create_action(
            self.get_icon(QStyle.SP_DialogCloseButton),
            "Close file",
            self.handle_close_file_action,
            QKeySequence(Qt.CTRL + Qt.Key_W))

    def init_menubar(self):
        menubar = self.menuBar()
        file_menu = menubar.addMenu("File")
        file_menu.insertAction(None, self.open_file_action)
        file_menu.insertAction(None, self.save_file_action)
        file_menu.insertAction(None, self.save_file_as_action)
        file_menu.insertAction(None, self.close_file_action)

    def init_ui(self):
        split = QSplitter(Qt.Horizontal, self)
        split.setChildrenCollapsible(False)
        tab_widget = QTabWidget(split)
        tab_widget.addTab(self.init_parts_tree(), "Sets")
        tab_widget.addTab(self.init_parts_list(), "List")
        split.addWidget(tab_widget)
        split.addWidget(ArmorPieceWidget(self.current_piece_view_ctrl))
        self.setCentralWidget(split)
        self.setGeometry(300, 300, 600, 400)
        self.setWindowTitle('Armor Editor')
        self.show()

    def init_toolbar(self):
        toolbar = self.addToolBar("Main")
        toolbar.insertAction(None, self.open_file_action)
        toolbar.insertAction(None, self.save_file_action)
        toolbar.insertAction(None, self.close_file_action)

    def init_parts_list(self):
        self.parts_list_view = QTreeView()
        self.parts_list_view.activated.connect(self.handle_parts_list_activated)
        return self.parts_list_view

    def init_parts_tree(self):
        self.parts_tree_view = QTreeView()
        self.parts_tree_view.activated.connect(self.handle_parts_tree_activated)
        return self.parts_tree_view

    def handle_open_file_action(self):
        file_path, _ = QFileDialog.getOpenFileName(parent=self)
        if file_path:
            self.handle_file_selected(file_path)

    def handle_save_file_action(self):
        if self.file_model is None:
            return
        try:
            self.file_model.save()
        except Exception as e:
            QMessageBox.warning(self,
                                "Error writing file", str(e),
                                QMessageBox.Ok, QMessageBox.Ok)

    def handle_save_file_as_action(self):
        if self.file_model is None:
            return
        file_path, _ = QFileDialog.getSaveFileName(self)
        if file_path:
            self.file_model.path = file_path
            self.handle_save_file_action()

    def handle_close_file_action(self):
        self.file_model = None
        self.parts_tree_view.setModel(None)
        self.parts_list_view.setModel(None)
        self.current_piece_view_ctrl.update(None)

    def handle_file_selected(self, file_path):
        try:
            self.file_model = FileModel.load(file_path)
        except Exception as e:
            self.file_model = None
            QMessageBox.warning(self,
                                "Error opening file", str(e),
                                QMessageBox.Ok, QMessageBox.Ok)
            return
        self.parts_tree_view.setModel(
            ArmorSetTreeModel(self.file_model.data.entries))
        self.parts_list_view.setModel(
            ArmorListModel(self.file_model.data.entries))

    def handle_parts_tree_activated(self, qindex):
        if isinstance(qindex.internalPointer(), ArmorSetNode):
            return
        index = qindex.internalPointer().ref.index
        model = self.file_model.data.find_first(index=index)
        self.current_piece_view_ctrl.update(model)

    def handle_parts_list_activated(self, qindex):
        index = qindex.row()
        model = self.file_model.data.find_first(index=index)
        self.current_piece_view_ctrl.update(model)

if __name__ == '__main__':
    Definitions.load()
    app = QApplication(sys.argv)
    ex = StructuredEditorWindow()
    sys.exit(app.exec_())