import sys
import glob
import scipy.io as sio
import numpy as np

import torch
from torch.utils.data.dataloader import DataLoader

from framework_pytorch.core import Trainer, BasicDataset
from framework_pytorch.modelnet.model_vnet_nointer import VNet
from framework_pytorch.modelnet.model import SegModel
from framework_pytorch.utils import loss_fuctions as LF, process_methods as P, utils as U
from framework_pytorch.utils.data_augmentation import RandomAugmentation

# training parameters
epochs = 500
learning_rate = 1e-4
train_batch_size = 1
eval_batch_size = 1
loss_function = {LF.CrossEntropy(): 1.0, LF.SoftDice(): 1.0} # loss functions {method: weight}
aug_rate = 0.8
output_path = './results/cmr_unet'

# training setttings
test_only = False

log_train_image = False
log_validation_image = True

data_path = './CMRxMotion/training/'

data_list = np.load(data_path + '/train_list.npy', allow_pickle=True).item()

train = data_list['d1_train'] + data_list['d2_train']
valid = data_list['d1_valid'] + data_list['d2_valid']

train = [data_path + '/data/' + x.split('-E')[0] + '/' + x + '.nii.gz' for x in train]
valid = [data_path + '/data/' + x.split('-E')[0] + '/' + x + '.nii.gz' for x in valid]

# set suffix for image and label (the difference between image path and label path) for data loading
img_suffix = '.nii.gz'
lab_suffix = '-label.nii.gz'

# set pre-process functions for image and label
pre = {img_suffix: [np.squeeze,
                    P.CenterCrop([256,256,10]),
                    P.min_max,
                    P.ExpandDim(0),
                    ],
        lab_suffix: [P.CenterCrop([256,256,10]),
                    #  lambda x: np.array(x>0.5, np.uint8),
                     P.OneHot([0,1,2,3])]
        }

aug = RandomAugmentation(img_suffix, lab_suffix, aug_rate=aug_rate)

# set device to gpu if gpu is available, otherwise use cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# build pytorch dataset, see core/basic_dataset
train_set = BasicDataset(train, [img_suffix, lab_suffix], pre, aug)
valid_set = BasicDataset(valid, [img_suffix, lab_suffix], pre)

# build pytorch data loader, shuffle train set while training
trainloader = DataLoader(train_set, batch_size=train_batch_size, shuffle=True)
validloader = DataLoader(valid_set, batch_size=eval_batch_size)

# get a random image for graph draw
random_img = torch.tensor([train_set[0][img_suffix]]).to(device)

# build 3D unet for 5 classes segmentation, the input channel is 1
net = VNet(input_channel=1, class_number=4)
# init optimizer, adam is used here
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
# init the model, #TODO: more details see modelnet/model
model = SegModel(net, optimizer, device, img_suffix, lab_suffix, dropout_rate=0, loss_functions=loss_function)

# init train and start train
trainer = Trainer(model)
if not test_only:
    trainer.train(trainloader, validloader,
        epochs=epochs,
        output_path=output_path,
        log_train_image=log_train_image,
        log_validation_image=log_validation_image,
        log_graph_input=random_img)

# test on test set
trainer.restore(output_path+'/ckpt/model_final.pt')

# test set 1
test1 = data_list['d1_valid'] + data_list['d2_valid']
test1 = [data_path + '/data/' + x.split('-E')[0] + '/' + x + '.nii.gz' for x in test1]
test1_set = BasicDataset(test1, [img_suffix, lab_suffix], pre)
test1loader = DataLoader(test1_set, batch_size=eval_batch_size)

test_results = trainer.test(test1loader, output_path, log_image=True, logsuffix='1')
sio.savemat(f'{output_path}/test1_results.mat', test_results)
with open(f'{output_path}/test1_results.txt', 'a+') as f:
    f.write(U.dict_to_str(test_results) + '\n')


class Stories:
	def __init__(self):
		self.memberStory = {"validator" : {"Faults":0}}

	def newStory(self, member):
		if member in self.memberStory:
			aux = self.memberStory[member]
			aux["Faults"] += 1
			self.memberStory[member] = aux
		else:
			self.memberStory[member] = {"Faults" : 1}

	def getStories(self, member= None):
		if member is not None:
			if self.memberStory.get(member) is not None:
				target = self.memberStory[member]
				faults = target["Faults"]
				requested = "Found, {0} has {1} stories in this server".format(member, str(faults))
			else:
				requested = "Error, User: {} not found!".format(member)
		else:
			requested = ""
			for (i, j) in zip(self.memberStory.keys(), self.memberStory.values()):
				requested += "{0} has {1} stories in this server\n".format(i,j)
		return requested

	def manageStories(self, member, command):
		if command == "delete":
			ifIn = self.tool_CheckIfIn(member, self.memberStory, [
				"Stories of {} succesfully cleared.".format(member),
				"{} doesn't have an story yet".format(member)])
			if ifIn[0]:
				del self.memberStory[member]
			requested = ifIn[1]
		else:
			ifIn = self.tool_CheckIfIn(member, self.memberStory, [
				"Stories of {} succesfully undone.".format(member),
				"{} doesn't have an story yet".format(member)])
			if ifIn[0]:
				aux = self.memberStory[member]
				aux["Faults"] -= 1
				self.memberStory[member] = aux
			requested = ifIn[1]
		return requested

	@staticmethod
	def tool_CheckIfIn(what, where, msg):
		if what in where:
			result = [True, msg[0]]
		else:
			result = [False, msg[1]]
		return result

if __name__ == '__main__':
	h = Stories()
	h.newStory("fabrizio")
	h.newStory("fabrizio")
	print(h.getStories())
	h.manageStories("fabrizio", "undo")
	print(h.getStories("fabrizio"))
	h.manageStories("fabrizio", "penalty")
	print(h.getStories("fabrizio"))
	h.manageStories("fabrizio", "delete")
	print(h.getStories("fabrizio"))


# python3

from abc import ABC
from collections import namedtuple
from sys import setrecursionlimit, stdin
from threading import stack_size, Thread
from typing import AnyStr, IO, List
from unittest import TestCase

setrecursionlimit(10 ** 6)
stack_size(2 ** 27)

test = namedtuple('test', 'input expected')

class TreeOrders:
    def __init__(self):
        self.n = 0
        self.key = self.left = self.right = None

    def read(self, src: IO):
        self.n = int(src.readline())
        self.key = [0 for _ in range(self.n)]
        self.left = [0 for _ in range(self.n)]
        self.right = [0 for _ in range(self.n)]
        for i in range(self.n):
            [self.key[i], self.left[i], self.right[i]] = map(int, src.readline().split())
        return self

    def walk(self) -> List[List[int]]:
        result, buf = [], [0 for _ in range(self.n)]

        for walk in [self.in_order, self.pre_order, self.post_order]:
            walk(buf)
            result.append(buf[:])

        return result

    def in_order(self, buf: List[int], node: int = 0, position: int = 0) -> int:
        if self.left[node] != -1:
            position = self.in_order(buf, self.left[node], position)

        buf[position] = self.key[node]
        position += 1

        if self.right[node] != -1:
            position = self.in_order(buf, self.right[node], position)

        return position

    def pre_order(self, buf: List[int], node: int = 0, position: int = 0) -> int:
        buf[position] = self.key[node]
        position += 1

        if self.left[node] != -1:
            position = self.pre_order(buf, self.left[node], position)

        if self.right[node] != -1:
            position = self.pre_order(buf, self.right[node], position)

        return position

    def post_order(self, buf: List[int], node: int = 0, position: int = 0) -> int:
        if self.left[node] != -1:
            position = self.post_order(buf, self.left[node], position)

        if self.right[node] != -1:
            position = self.post_order(buf, self.right[node], position)

        buf[position] = self.key[node]
        position += 1

        return position

class Fake(IO, ABC):
    def __init__(self, rows: List[str]):
        self.__i = -1
        self.__rows = [str(len(rows))] + rows

    def readline(self, limit: int = -1) -> AnyStr:
        self.__i += 1
        return self.__rows[self.__i]

class Test(TestCase):
    def test_tree_orders(self):
        tests = [
            # samples
            test([
                '4 1 2',
                '2 3 4',
                '5 -1 -1',
                '1 -1 -1',
                '3 -1 -1',
            ], [
                '1 2 3 4 5',
                '4 2 1 3 5',
                '1 3 2 5 4',
            ]),
            test([
                '0 7 2',
                '10 -1 -1',
                '20 -1 6',
                '30 8 9',
                '40 3 -1',
                '50 -1 -1',
                '60 1 -1',
                '70 5 4',
                '80 -1 -1',
                '90 -1 -1',
            ], [
                '50 70 80 30 90 40 0 20 10 60',
                '0 70 50 40 30 80 90 20 60 10',
                '50 80 90 30 40 70 10 60 20 0',
            ]),
        ]
        for i, t in enumerate(tests):
            src = Fake(t.input)
            self.assertEqual(t.expected,
                             list(map(lambda way: ' '.join(str(step) for step in way), TreeOrders().read(src).walk())),
                             msg='at {} position'.format(i))

def main():
    for way in TreeOrders().read(stdin).walk():
        print(' '.join(str(step) for step in way))

if __name__ == '__main__':
    Thread(target=main).start()


#!/usr/bin/env python3

import sys
sys.path.append('../Lib/')

import torch
import torch.nn as nn

from torch.utils.data import DataLoader

import random, argparse
import wikidata

from transformers import (
    AdamW,
    T5ForConditionalGeneration,
    T5Tokenizer,
    get_linear_schedule_with_warmup)

# deterministic determinism
torch.manual_seed(2020)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False
random.seed(2020)

class T5FineTuner(nn.Module):
  """A transformative experience"""

  def __init__(self):
    """Some of the best constructors in the world"""

    super(T5FineTuner, self).__init__()
    self.model = T5ForConditionalGeneration.from_pretrained(
      hparams.model_name)

  def forward(
   self,
   input_ids,
   attention_mask,
   decoder_input_ids,
   decoder_attention_mask,
   labels):
    """Forwarding"""

    output = self.model(
      input_ids=input_ids,
      attention_mask=attention_mask,
      decoder_input_ids=decoder_input_ids,
      decoder_attention_mask=decoder_attention_mask,
      labels=labels)

    return output

def fit(model, train_loader, val_loader, tokenizer):
  """Training routine"""

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model.to(device)

  optimizer = AdamW(model.parameters())

  for epoch in range(1, hparams.n_epochs + 1):
    train_loss, num_train_steps = 0, 0
    model.train()

    for batch in train_loader:
      optimizer.zero_grad()

      batch = tuple(t.to(device) for t in batch)
      source_ids, source_mask, target_ids, target_mask = batch

      labels = target_ids
      labels[labels[:, :] == tokenizer.pad_token_id] = -100

      outputs = model(
        input_ids=source_ids,
        attention_mask=source_mask,
        decoder_input_ids=None,
        decoder_attention_mask=target_mask,
        labels=labels)
      loss = outputs[0]

      loss.backward()

      torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
      optimizer.step()

      train_loss += loss.item()
      num_train_steps += 1

    av_loss = train_loss / num_train_steps
    val_loss = evaluate(model, val_loader, tokenizer)
    print('ep: %d, steps: %d, tr loss: %.3f, val loss: %.3f' % \
          (epoch, num_train_steps, av_loss, val_loss))

def evaluate(model, data_loader, tokenizer):
  """Evaluation routine"""

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model.to(device)

  total_loss, num_steps = 0, 0
  model.eval()

  for batch in data_loader:
    batch = tuple(t.to(device) for t in batch)
    source_ids, source_mask, target_ids, target_mask = batch

    labels = target_ids
    labels[labels[:, :] == tokenizer.pad_token_id] = -100

    with torch.no_grad():
      outputs = model(
        input_ids=source_ids,
        attention_mask=source_mask,
        decoder_input_ids=None,
        decoder_attention_mask=target_mask,
        labels=labels)
      loss = outputs[0]

    total_loss += loss.item()
    num_steps += 1

  average_loss = total_loss / num_steps
  return average_loss

def generate(model, data_loader, tokenizer):
  """Need to add 'summarize' if run before training"""

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model.to(device)
  model.eval()

  for batch in data_loader:
    batch = tuple(t.to(device) for t in batch)
    source_ids, source_mask, target_ids, target_mask = batch

    inputs = tokenizer.batch_decode(source_ids)
    targets = tokenizer.batch_decode(target_ids)

    predictions = model.model.generate(
      input_ids=source_ids,
      max_length=hparams.max_output_length,
      early_stopping=True,
      num_beams=2,
      attention_mask=source_mask)
    predictions = tokenizer.batch_decode(
      predictions,
      skip_special_tokens=True,
      clean_up_tokenization_spaces=True)

    print('input:', inputs[0])
    print('target:', targets[0])
    print('prediction:', predictions[0])
    print()

def main():
  """Fine-tune on summarization data"""

  tokenizer = T5Tokenizer.from_pretrained(hparams.model_name)

  train_dataset = wikidata.WikiHow(
    tokenizer=tokenizer,
    split='train',
    input_length=hparams.max_input_length,
    output_length=hparams.max_output_length)
  train_data_loader = DataLoader(
    train_dataset,
    batch_size=hparams.batch_size)

  val_dataset = wikidata.WikiHow(
    tokenizer=tokenizer,
    split='validation',
    input_length=hparams.max_input_length,
    output_length=hparams.max_output_length)
  val_data_loader = DataLoader(
    val_dataset,
    batch_size=hparams.batch_size)

  model = T5FineTuner()

  fit(model, train_data_loader, val_data_loader, tokenizer)
  generate(model, val_data_loader, tokenizer)

if __name__ == "__main__":
  "My kind of street"

  hparam_dict = dict(
    model_name='t5-small',
    max_input_length=512,
    max_output_length=150,
    batch_size=32,
    n_epochs=5)
  hparams = argparse.Namespace(**hparam_dict)
  print('hyper-parameters:', hparams)

  main()


from copy import copy

from dateutil.parser import parse
from django.test import TestCase

from ..models import Report

class ReportTestCase(TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.standard_data = {
            "subscriber": "11981094379",
            "destination": "1198109437",
            "call_started_at": parse("2016-01-01T15:00:01Z"),
            "call_ended_at": parse("2016-01-01T15:50:37Z"),
        }
        cls.reduced_data = {
            "subscriber": "11981094379",
            "destination": "1198109437",
            "call_started_at": parse("2016-01-01T22:00:00Z"),
            "call_ended_at": parse("2016-01-02T05:59:59Z"),
        }
        cls.standard_and_reduced_data = {
            "subscriber": "11981094379",
            "destination": "1198109437",
            "call_started_at": parse("2016-01-01T21:57:13Z"),
            "call_ended_at": parse("2016-01-01T22:17:53Z"),
        }
        cls.reduced_and_standard_data = {
            "subscriber": "11981094379",
            "destination": "1198109437",
            "call_started_at": parse("2017-12-12T04:57:13Z"),
            "call_ended_at": parse("2017-12-12T06:10:56Z"),
        }

        cls.report_standard_only = Report.objects.create(**cls.standard_data)
        cls.report_reduced_only = Report.objects.create(**cls.reduced_data)
        cls.report_standard_and_reduced = Report.objects.create(
            **cls.standard_and_reduced_data
        )
        cls.report_reduced_and_standard = Report.objects.create(
            **cls.reduced_and_standard_data
        )

    def test_call_price_standard_only(self):
        self.assertEqual(self.report_standard_only.price_label, "R$ 4,86")
        self.assertEqual(
            self.report_standard_only.minutes_in_each_tarrif(),
            {"standard": 50, "reduced": 0},
        )

    def test_call_price_reduced_only(self):
        self.assertEqual(self.report_reduced_only.price_label, "R$ 0,36")
        self.assertEqual(
            self.report_reduced_only.minutes_in_each_tarrif(),
            {"standard": 0, "reduced": 479},
        )

    def test_call_price_standard_and_reduced(self):
        self.assertEqual(self.report_standard_and_reduced.price_label, "R$ 0,54")
        self.assertEqual(
            self.report_standard_and_reduced.minutes_in_each_tarrif(),
            {"standard": 2, "reduced": 18},
        )

    def test_call_price_reduced_and_standard(self):
        self.assertEqual(self.report_reduced_and_standard.price_label, "R$ 1,26")
        self.assertEqual(
            self.report_reduced_and_standard.minutes_in_each_tarrif(),
            {"standard": 10, "reduced": 63},
        )

    def test_call_price_over_day(self):
        data = copy(self.standard_data)
        data["call_started_at"] = parse("2017-12-13T21:57:13Z")
        data["call_ended_at"] = parse("2017-12-14T22:10:56Z")
        report = Report.objects.create(**data)
        self.assertEqual(report.price_label, "R$ 86,85")

    def test_call_price_one_minute(self):
        data = copy(self.standard_data)
        data["call_started_at"] = parse("2017-12-13T21:57:13Z")
        data["call_ended_at"] = parse("2017-12-13T21:58:13Z")
        report = Report.objects.create(**data)
        self.assertEqual(report.price_label, "R$ 0,45")

    def test_call_price_less_than_one_minute(self):
        data = copy(self.standard_data)
        data["call_ended_at"] = parse("2016-01-01T15:00:59Z")
        report = Report.objects.create(**data)
        self.assertEqual(report.price_label, "R$ 0,36")

    def test_call_price_two_hours(self):
        data = copy(self.standard_data)
        data["call_ended_at"] = parse("2016-01-01T17:00:01Z")
        report = Report.objects.create(**data)
        self.assertEqual(report.price_label, "R$ 11,16")

    def test_call_start_date(self):
        self.assertEqual(
            self.report_standard_only.call_start_date,
            self.standard_data["call_started_at"].date(),
        )

    def test_call_start_time(self):
        self.assertEqual(
            self.report_standard_only.call_start_time,
            self.standard_data["call_started_at"].time(),
        )

    def test_call_duration(self):
        self.assertEqual(self.report_standard_only.call_duration, "0h50m36s")
        self.assertEqual(self.report_reduced_only.call_duration, "7h59m59s")
        self.assertEqual(self.report_standard_and_reduced.call_duration, "0h20m40s")


import os
import pickle

import yaml

__all__ = ['save_yaml',
           'load_yaml',
           'save_obj',
           'load_obj',
           'get_dataset_tags',
           'save_experiment',
           'get_corr_pred'
           ]

def save_yaml(obj, path):
    """Save yaml object to specified path.
    """
    if not os.path.exists(os.path.dirname(path)):
        os.makedirs(os.path.dirname(path))
    with open(path, 'w') as outfile:
        yaml.dump(obj, outfile, default_flow_style=False)

def load_yaml(filepath):
    """Load yaml object from specified path.
    """
    return yaml.load(open(filepath), Loader=yaml.FullLoader)

def save_obj(obj, path):
    """Save pickle object to specified path.
    """
    if not os.path.exists(os.path.dirname(path)):
        os.makedirs(os.path.dirname(path))
    with open(path, 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(path):
    """Load pickle object from specified path.
    """
    with open(path, 'rb') as f:
        return pickle.load(f)

def save_experiment(d_spec, target_folder):
    """Save dataset specifications for experiment to yaml file.
    """
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)
    for t in d_spec:
        save_yaml(d_spec[t], os.path.join(target_folder, str(t) + '.yml'))

def get_dataset_tags(spec):
    """Given a spec dict return a list of its tags.
    """
    tags = [mode['specification']['tag'] for mode in spec['modes']]
    return list(set(tags))

def get_corr_pred(study_name):
    """
    A study name is of the form CORR-factor1-factor2-factor3_PRED-factor1.
    If no factors are correlated, then CORR_PRED-factor1.
    This function returns lists of correlated and predicted factors from this string.
    """
    if 'CORR-' in study_name:
        # There are correlations
        corrs, preds = study_name.split('CORR-')[1].split('_PRED-')
        corrs = corrs.split('-')
        preds = preds.split('-')
    else:
        # There are no correlations
        corrs = []
        preds = study_name.split('PRED-')[1].split('-')
    return corrs, preds


#! /usr/bin/env python

import rospy
from move_base_msgs.msg import MoveBaseActionResult

def debug():
    pub = rospy.Publisher('move_base/result',MoveBaseActionResult,queue_size=10)
    rospy.init_node('debug',anonymous=True)

    result = MoveBaseActionResult()
    result.status.status = 3
    pub.publish(result)
    rospy.loginfo('succeed sending ')

if __name__ == '__main__':
    debug()


import torch.nn as nn
import torch.optim as optim
from Models.Resnet_Model.BasicBlock import BasicBlock
from Models.Resnet_Model.Bottle_neck import Bottle_neck
from Utils import conv_1x1

class ResNet(nn.Module):
    def __init__(self, input_channel, layers, num_classes=10, zero_init_residual=True):
        super(ResNet, self).__init__()
        self.in_planes = 64

        self.inputLayer = nn.Sequential(nn.Conv2d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False),
                                        nn.BatchNorm2d(64),
                                        nn.ReLU(inplace=True),
                                        nn.MaxPool2d(kernel_size=3, stride=2, padding=1))

        self.inputLayer2 = nn.Sequential(self.make_layer(Bottle_neck, 64, layers[0]),
                                         self.make_layer(Bottle_neck, 128, layers[1], stride=2),
                                         self.make_layer(Bottle_neck, 256, layers[2], stride=2),
                                         self.make_layer(Bottle_neck, 512, layers[3], stride=2))

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * Bottle_neck.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottle_neck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

        self.optimizer = optim.Adam(self.parameters(), lr=0.001)
        self.cross_entropy = nn.CrossEntropyLoss()

    def make_layer(self, block, planes, blocks, stride=1):
        down_sample = None

        if stride != 1 or self.in_planes != planes * block.expansion:
            down_sample = nn.Sequential(
                conv_1x1(self.in_planes, planes * block.expansion, stride),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.in_planes, planes, stride, down_sample))

        self.in_planes = planes * block.expansion

        for _ in range(1, blocks):
            layers.append(block(self.in_planes, planes))

        final_layers = nn.Sequential(*layers)
        return final_layers

    def forward(self, x):
        x = self.inputLayer(x)
        x = self.inputLayer2(x)
        x = self.avgpool(x)

        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x

    def Learn(self, inputs, labels):
        predict = self.forward(inputs)
        loss = self.cross_entropy(predict, labels)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss.item()


from typing import Union

from pipenv.vendor.plette.models import Package, PackageCollection
from pipenv.vendor.tomlkit.container import Container
from pipenv.vendor.tomlkit.items import AoT, Array, Bool, InlineTable, Item, String, Table

try:
    import tomllib as toml
except ImportError:
    from pipenv.vendor import tomli as toml

from pipenv.vendor import tomlkit

TOML_DICT_TYPES = Union[Container, Package, PackageCollection, Table, InlineTable]
TOML_DICT_OBJECTS = (Container, Package, Table, InlineTable, PackageCollection)
TOML_DICT_NAMES = [o.__class__.__name__ for o in TOML_DICT_OBJECTS]

def cleanup_toml(tml):
    # Remove all empty lines from TOML.
    toml = "\n".join(line for line in tml.split("\n") if line.strip())
    new_toml = []
    # Add newlines between TOML sections.
    for i, line in enumerate(toml.split("\n")):
        # Skip the first line.
        if line.startswith("[") and i > 0:
            # Insert a newline before the heading.
            new_toml.append("")
        new_toml.append(line)
    # adding new line at the end of the TOML file
    new_toml.append("")
    toml = "\n".join(new_toml)
    return toml

def convert_toml_outline_tables(parsed, project):
    """Converts all outline tables to inline tables."""

    def convert_tomlkit_table(section):
        result = section.copy()
        if isinstance(section, tomlkit.items.Table):
            body = section.value._body
        elif isinstance(section, tomlkit.container.OutOfOrderTableProxy):
            body = section._internal_container._body
        else:
            body = section._body
        for key, value in body:
            if not key:
                continue
            if hasattr(value, "keys") and not isinstance(
                value, tomlkit.items.InlineTable
            ):
                table = tomlkit.inline_table()
                table.update(value.value)
                result[key.key] = table
        return result

    def convert_toml_table(section):
        result = section.copy()
        for package, value in section.items():
            if hasattr(value, "keys") and not isinstance(
                value, toml.decoder.InlineTableDict
            ):
                table = toml.TomlDecoder().get_empty_inline_table()
                table.update(value)
                result[package] = table
        return result

    is_tomlkit_parsed = isinstance(parsed, tomlkit.container.Container)
    for section in project.get_package_categories():
        table_data = parsed.get(section, {})
        if not table_data:
            continue
        if is_tomlkit_parsed:
            result = convert_tomlkit_table(table_data)
        else:
            result = convert_toml_table(table_data)

        parsed[section] = result
    return parsed

def tomlkit_value_to_python(toml_value):
    # type: (Union[Array, AoT, TOML_DICT_TYPES, Item]) -> Union[List, Dict]
    value_type = type(toml_value).__name__
    if (
        isinstance(toml_value, TOML_DICT_OBJECTS + (dict,))
        or value_type in TOML_DICT_NAMES
    ):
        return tomlkit_dict_to_python(toml_value)
    elif isinstance(toml_value, AoT) or value_type == "AoT":
        return [tomlkit_value_to_python(val) for val in toml_value._body]
    elif isinstance(toml_value, Array) or value_type == "Array":
        return [tomlkit_value_to_python(val) for val in list(toml_value)]
    elif isinstance(toml_value, String) or value_type == "String":
        return f"{toml_value!s}"
    elif isinstance(toml_value, Bool) or value_type == "Bool":
        return toml_value.value
    elif isinstance(toml_value, Item):
        return toml_value.value
    return toml_value

def tomlkit_dict_to_python(toml_dict):
    # type: (TOML_DICT_TYPES) -> Dict
    value_type = type(toml_dict).__name__
    if toml_dict is None:
        raise TypeError("Invalid type NoneType when converting toml dict to python")
    converted = None  # type: Optional[Dict]
    if isinstance(toml_dict, (InlineTable, Table)) or value_type in (
        "InlineTable",
        "Table",
    ):
        converted = toml_dict.value
    elif isinstance(toml_dict, (Package, PackageCollection)) or value_type in (
        "Package, PackageCollection"
    ):
        converted = toml_dict._data
        if isinstance(converted, Container) or type(converted).__name__ == "Container":
            converted = converted.value
    elif isinstance(toml_dict, Container) or value_type == "Container":
        converted = toml_dict.value
    elif isinstance(toml_dict, dict):
        converted = toml_dict.copy()
    else:
        raise TypeError(
            f"Invalid type for conversion: expected Container, Dict, or Table, got {toml_dict}"
        )
    if isinstance(converted, dict):
        return {k: tomlkit_value_to_python(v) for k, v in converted.items()}
    elif isinstance(converted, (TOML_DICT_OBJECTS)) or value_type in TOML_DICT_NAMES:
        return tomlkit_dict_to_python(converted)
    return converted


from pubnub.pnconfiguration import PNConfiguration
from pubnub.pubnub import PubNub, SubscribeListener
import uuid
import random
import os

def get_random_uuid_as_string():
    basename_string = 'festeban26-iot-lockers-server'
    new_uuid_random_number = random.random()
    new_uuid_string = basename_string + "-" + str(new_uuid_random_number)
    custom_uuid = uuid.uuid5(uuid.NAMESPACE_DNS, new_uuid_string)
    print("UUID generated: " + str(custom_uuid))
    return str(custom_uuid)

def get_uuid_from_default_uuid_file():
    # first check if a previous file exists, if it does not exist, generate uuid and save it to file
    try:
        with open("UUID.txt", "r") as text_file:
            return text_file.read()
    except IOError:
        print("UUID file was not found. Creating new uuid")
        new_uuid = get_random_uuid_as_string()
        save_uuid_in_default_uuid_file(new_uuid)
        with open("UUID.txt", "r") as text_file:
            return text_file.read()

def save_uuid_in_default_uuid_file(uuid_as_string):
    # first check if it exists, if it exists, delete the file
    try:
        file = open("UUID.txt", 'r')
        file.close()
        os.remove("UUID.txt")
    except IOError:
        pass
    # once the file is deleted or the file did not exists, save content
    with open("UUID.txt", "w") as text_file:
        text_file.write(uuid_as_string)

print(get_uuid_from_default_uuid_file())

'''
pub_num_iot_lockers_channel = 'iot-lockers'
pn_config = PNConfiguration()
pn_config.subscribe_key = 'INSERT API KEY'
pn_config.publish_key = 'INSERT API KEY'
pn_config.uuid = '4205e440-1258-462c-aa94-1870f1ab4755'
pub_nub = PubNub(pn_config)

print(pub_nub.uuid)
pub_nub.publish().channel(pub_num_iot_lockers_channel).message("Test").sync()
'''


from core.utils import *
from manim import *
from common import *

class p4(Scene):
    def construct(self):
        # Beat 4: [Pause, show small simulation]
        ax = NumberLine(
            x_range=[0, 5],
            length=10,
            color=BLUE,
            include_numbers=True
        )
        ant = Ant(ax=ax)
        coords = [(2, 1, 1), (4, 3, 1)]
        portals = createPortals(coords, ax)
        self.play(Create(ax))
        self.play(Create(ant.mob))
        self.add_foreground_mobject(ant.mob)
        for p in portals:
            self.play(FadeIn(p.mob))
        self.wait(2)
        t = simulate(self, ant, portals, ax)
        self.play(*t.fade())


#!/usr/bin/env python
import argparse
import logging
import os
import sys
import numpy as np
from shutil import copyfile, rmtree
from ase import units
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data.sampler import RandomSampler
from schnetpack.utils.script_utils.settings import get_environment_provider
from schnetpack.utils import get_dataset
import schnetpack as spk
import schnet_ev
from schnet_ev.data import Eigenvalue_properties, Delta_EV_properties, Charges
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))

def get_parser():
    """ Setup parser for command line arguments """
    main_parser = argparse.ArgumentParser()

    # command-specific
    cmd_parser = argparse.ArgumentParser(add_help=False)
    cmd_parser.add_argument('--cuda', help='Set flag to use GPU(s)', action='store_true')
    cmd_parser.add_argument('--parallel',
                            help='Run data-parallel on all available GPUs (specify with environment variable'
                                 + ' CUDA_VISIBLE_DEVICES)', action='store_true')
    cmd_parser.add_argument('--batch_size', type=int,
                            help='Mini-batch size for training and prediction (default: %(default)s)',
                            default=100)
    cmd_parser.add_argument("--aggregation_mode",type=str,default="build",choices=["avg","build"])
    cmd_parser.add_argument('--environment_provider',type=str, default="simple", choices=["simple","ase","torch"],help="Environment provider for dataset. (default: %(default)s)",)
    cmd_parser.add_argument('--force_mask_index',help='Atom number of excluded atom for force training', type=int)
    cmd_parser.add_argument('--force_mask', help= 'Only train on forces for specific atoms', action="store_true")
    # training
    train_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    train_parser.add_argument('datapath', help='Path / destination of MD17 dataset')
    train_parser.add_argument('modelpath', help='Destination for models and logs')
    train_parser.add_argument('--seed', type=int, default=None, help='Set random seed for torch and numpy.')
    train_parser.add_argument('--overwrite', action='store_true', help='Overwrite old directories')

    # data split
    train_parser.add_argument('--split_path', help='Path / destination of npz with data splits',
                              default=None)
    train_parser.add_argument('--split', help='Give sizes of train and validation splits and use remaining for testing',
                              type=int, nargs=2, default=[None, None])
    train_parser.add_argument('--max_epochs', type=int, help='Maximum number of training epochs (default: %(default)s)',
                              default=5000)
    train_parser.add_argument('--lr', type=float, help='Initial learning rate (default: %(default)s)',
                              default=1e-4)
    train_parser.add_argument('--lr_patience', type=int,
                              help='Epochs without improvement before reducing the learning rate (default: %(default)s)',
                              default=25)
    train_parser.add_argument('--lr_decay', type=float, help='Learning rate decay (default: %(default)s)',
                              default=0.8)
    train_parser.add_argument('--lr_min', type=float, help='Minimal learning rate (default: %(default)s)',
                              default=1e-6)
    train_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
                              choices=['csv', 'tensorboard'], default='csv')
    train_parser.add_argument('--log_every_n_epochs', type=int,
                              help='Log metrics every given number of epochs (default: %(default)s)',
                              default=1)
    train_parser.add_argument('--tradeoffs', type=str, help='Property tradeoffs for training', default=None)
    train_parser.add_argument('--verbose', action='store_true', help='Print property error magnitudes')
    train_parser.add_argument('--real_socs', action='store_true',
                              help='If spin-orbit couplings are predicted, information should be given whether they are real or not.')
    train_parser.add_argument('--hamiltonian', action='store_true',
                              help='If True, approximate diabatic Hamiltonian using partial charges without diagonalization to train on eigenvalues.')
    train_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)
    train_parser.add_argument('--phase_loss', action='store_true', help='Use special loss, which ignores phase.')
    train_parser.add_argument('--inverse_nacs', action='store_true', help='Weight NACs with inverse energies.')
    train_parser.add_argument('--smooth_loss', action='store_true', help='Use phase independent min loss.')
    train_parser.add_argument('--mixed_loss', action='store_true', help='Use phase independent min loss.')
    train_parser.add_argument('--diagonal', action='store_true', help='Train SOCs via diagonal elements. Must be included in the training data base', default=None)
    train_parser.add_argument('--L1', action='store_true', help='Use L1 norm')
    train_parser.add_argument('--Huber', action='store_true', help='Use L1 norm')

    ## evaluation
    eval_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    eval_parser.add_argument('datapath', help='Path / destination of MD17 dataset directory')
    eval_parser.add_argument('modelpath', help='Path of stored model')
    eval_parser.add_argument('--split', help='Evaluate on trained model on given split',
                             choices=['train', 'validation', 'test', 'all'], default=['all'], nargs='+')
    eval_parser.add_argument('--hessian', action='store_true', help='Gives back the hessian of the PES.')
    eval_parser.add_argument('--hamiltonian', action='store_true',
                              help='If True, approximate diabatic Hamiltonian using partial charges without diagonalization to train on eigenvalues.')
    eval_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)

    pred_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    pred_parser.add_argument('datapath', help='Path / destination of MD17 dataset directory')
    pred_parser.add_argument('--print_uncertainty',action='store_true', help='Print uncertainty of each property',default=False)
    pred_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)
    pred_parser.add_argument('--thresholds',type=float,nargs=5, help='Percentage of mean predicted by NNs taken as thresholds for adaptive sampling - first value for energy, second for forces, third for dipoles, fourth for nacs, fifth for socs', default=None)
    pred_parser.add_argument('--modelpaths',type=str, help='Path of stored models')
    pred_parser.add_argument('modelpath', help='Path of stored model')
    pred_parser.add_argument('--hessian', action='store_true', help='Gives back the hessian of the PES.')
    pred_parser.add_argument('--adaptive',type=float, nargs=1,help='Adaptive Sampling initializer + float for reducing the step size')
    pred_parser.add_argument('--nac_approx',type=float, nargs=3, default=[1,0.018,0.036],help='Type of NAC approximation as first value and threshold for energy gap in Hartree as second value.')
    # model-specific parsers
    opt_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    opt_parser.add_argument('--nac_approx',type=float, nargs=3, default=[1,0.018,0.036],help='Type of NAC approximation as first value and threshold for energy gap in Hartree as second value.')
    model_parser = argparse.ArgumentParser(add_help=False)

    #######  SchNet  #######
    schnet_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    schnet_parser.add_argument('--features', type=int, help='Size of atom-wise representation (default: %(default)s)',
                               default=256)
    schnet_parser.add_argument("--cutoff_function",help="Functional form of the cutoff", choices=["hard","cosine","mollifier"],default="cosine")
    schnet_parser.add_argument('--interactions', type=int, help='Number of interaction blocks (default: %(default)s)',
                               default=6)
    schnet_parser.add_argument('--cutoff', type=float, default=10.,
                               help='Cutoff radius of local environment (default: %(default)s)')
    schnet_parser.add_argument('--num_gaussians', type=int, default=50,
                               help='Number of Gaussians to expand distances (default: %(default)s)')
    schnet_parser.add_argument('--n_layers', type=int, default=3,
                               help='Number of layers in output networks (default: %(default)s)')
    schnet_parser.add_argument('--n_neurons', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')

    #######  invD  ########
    invD_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    invD_parser.add_argument('--n_atoms', type=int, help='Number of atoms in the molecule',
                             default="ADD NUMBER OF ATOMS WITH ARGUENT --n_atoms")
    invD_parser.add_argument('--n_layers', type=int, default=3,
                               help='Number of layers in output networks (default: %(default)s)')
    invD_parser.add_argument('--n_neurons', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')
     #######  wACSF  ########
    wacsf_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    # wACSF parameters
    wacsf_parser.add_argument('--radial', type=int, default=22,
                              help='Number of radial symmetry functions (default: %(default)s)')
    wacsf_parser.add_argument('--angular', type=int, default=5,
                              help='Number of angular symmetry functions (default: %(default)s)')
    wacsf_parser.add_argument('--zetas', type=int, nargs='+', default=[1],
                              help='List of zeta exponents used for angle resolution (default: %(default)s)')
    wacsf_parser.add_argument('--standardize', action='store_true',
                              help='Standardize wACSF before atomistic network.')
    wacsf_parser.add_argument('--cutoff', type=float, default=5.0,
                              help='Cutoff radius of local environment (default: %(default)s)')
    # Atomistic network parameters
    wacsf_parser.add_argument('--n_nodes', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')
    wacsf_parser.add_argument('--n_layers', type=int, default=2,
                              help='Number of layers in atomic networks (default: %(default)s)')
    # Advances wACSF settings
    wacsf_parser.add_argument('--centered', action='store_true', help='Use centered Gaussians for radial functions')
    wacsf_parser.add_argument('--crossterms', action='store_true', help='Use crossterms in angular functions')
    wacsf_parser.add_argument('--behler', action='store_true', help='Switch to conventional ACSF')
    wacsf_parser.add_argument('--elements', default=['H', 'C', 'O'], nargs='+',
                              help='List of elements to be used for symmetry functions (default: %(default)s).')

    # setup subparser structure
    cmd_subparsers = main_parser.add_subparsers(dest='mode', help='Command-specific arguments')
    cmd_subparsers.required = True
    subparser_train = cmd_subparsers.add_parser('train', help='Training help')
    subparser_eval = cmd_subparsers.add_parser('eval', help='Eval help')
    subparser_opt = cmd_subparsers.add_parser('opt', help='Eval help',parents=[opt_parser])
    subparser_pred = cmd_subparsers.add_parser('pred', help='Eval help', parents=[pred_parser])

    train_subparsers = subparser_train.add_subparsers(dest='model', help='Model-specific arguments')
    train_subparsers.required = True
    train_subparsers.add_parser('invD', help='invD help', parents=[train_parser, invD_parser])
    train_subparsers.add_parser('schnet', help='SchNet help', parents=[train_parser, schnet_parser])
    train_subparsers.add_parser('wacsf', help='wACSF help', parents=[train_parser, wacsf_parser])

    eval_subparsers = subparser_eval.add_subparsers(dest='model', help='Model-specific arguments')
    eval_subparsers.required = True
    eval_subparsers.add_parser('invD', help='invD help', parents=[eval_parser, invD_parser])
    eval_subparsers.add_parser('schnet', help='SchNet help', parents=[eval_parser, schnet_parser])
    eval_subparsers.add_parser('wacsf', help='wACSF help', parents=[eval_parser, wacsf_parser])

    opt_subparsers = subparser_opt.add_subparsers(dest="mdel",help="Model-specific arguments")
    opt_subparsers.required = True
    opt_subparsers.add_parser('schnet', help='SchNet help', parents=[eval_parser,schnet_parser])
    return main_parser

def train(args, model, tradeoffs, train_loader, val_loader, device, n_orb,mean,stddev,offset):
    # setup hook and logging
    hooks = [
        spk.train.MaxEpochHook(args.max_epochs)
    ]

    # setup optimizer for training
    to_opt = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = Adam(to_opt, lr=args.lr)

    schedule = spk.train.ReduceLROnPlateauHook(optimizer, patience=args.lr_patience, factor=args.lr_decay,
                                               min_lr=args.lr_min,
                                               window_length=1, stop_after_min=True)
    hooks.append(schedule)

    # Build metrics based on the properties in the model
    metrics = []
    for prop in tradeoffs:
      if prop == "eigenvalues_active_forces" or prop == "eigenvalues_forces":
          pass
      else:
          metrics += [
              schnet_ev.metrics.MeanAbsoluteError(prop, prop),
              schnet_ev.metrics.RootMeanSquaredError(prop, prop)
          ]
    if args.logger == 'csv':
        logger = spk.train.CSVHook(os.path.join(args.modelpath, 'log'),
                                   metrics, every_n_epochs=args.log_every_n_epochs)
        hooks.append(logger)
    elif args.logger == 'tensorboard':
        logger = spk.train.TensorboardHook(os.path.join(args.modelpath, 'log'),
                                           metrics, every_n_epochs=args.log_every_n_epochs)
        hooks.append(logger)
    n_act = n_orb['n_act']
    # Automatically construct loss function based on tradoffs
    def loss(batch, result):
        err_sq = 0.0
        if args.verbose:
            print('===================')
        for prop in tradeoffs:
            if prop == 'energy' or prop == 'forces' or prop == 'homo_lumo':
                prop_diff = batch[prop] - result[prop]
                prop_err = torch.mean(prop_diff.view(-1) ** 2)
                err_sq = err_sq + float(tradeoffs[prop]) * prop_err
            if prop in Eigenvalue_properties.properties or prop in Delta_EV_properties.properties :
                if "eigenvalues_active_forces" in tradeoffs or "eigenvalues_forces" in tradeoffs:
                    result['ev_all']= result[prop]
                    result[prop] = result[prop][prop]
                #results are ordered in descending order
                if args.smooth_loss == True:
                     prop_diff =torch.normal(batch[prop],0.01)-torch.normal(result[prop],0.01)
                     prop_err = torch.mean(prop_diff.view(-1)**2)
                     err_sq = err_sq+float(tradeoffs[prop])*prop_err
                elif args.hamiltonian == True and prop == "eigenvalues":
                     #every vector has a different size - masked for every batch different the mask however is max. length and diff. for every molecule
                     prop_diff = (batch[prop][:,:result[prop].shape[1]]- result[prop]) * batch["eigenvalues_all_mask"][:,:result[prop].shape[1]]
                     prop_err = torch.mean(prop_diff.view(-1)**2)
                     err_sq = err_sq+float(tradeoffs[prop])*prop_err
                elif args.hamiltonian == True and prop != "eigenvalues":
                     prop_diff_ev = (batch[prop][:,:result[prop].shape[1]] - result[prop])*batch["eigenvalues_active_mask"][:,:result[prop].shape[1]]
                     prop_err_ev = torch.mean(prop_diff_ev.view(-1)**2)
                     prop_err_charges=0
                     #get charges
                     if "hirshfeld_pbe" in tradeoffs:
                         q_r=torch.sum((result['coeff']**2)[:,:batch['n_elec_atoms'].shape[1],:]*batch["occupation"][:,None,:result[prop].shape[1]],dim=2)
                         result["hirshfeld_pbe"] = q_r -batch['n_elec_atoms']
                         prop_diff_charges = batch["hirshfeld_pbe"] - result["hirshfeld_pbe"]
                         prop_err_charges += torch.mean(prop_diff_charges.view(-1)**2)*tradeoffs["hirshfeld_pbe"]
                     if "hirshfeld_pbe0" in tradeoffs:
                         q_r=torch.sum((result['coeff']**2)[:,:batch['n_elec_atoms'].shape[1],:]*batch["occupation"][:,None,:result[prop].shape[1]],dim=2)
                         result["hirshfeld_pbe0"]= q_r -batch['n_elec_atoms']
                         prop_diff_charges = batch["hirshfeld_pbe0"] - result["hirshfeld_pbe0"]
                         prop_err_charges += torch.mean(prop_diff_charges.view(-1)**2)*tradeoffs["hirshfeld_pbe0"]
                     err_sq = err_sq+float(tradeoffs[prop])*(prop_err_ev)+prop_err_charges
                elif prop in Delta_EV_properties.properties:
                    prop_diff = torch.abs(batch[prop] - result[prop])*batch["eigenvalues_active_mask"]
                    prop_err = torch.mean(prop_diff.view(-1)**2)
                    err_sq = err_sq  + float(tradeoffs[prop])*prop_err
                else:
                    prop_diff = torch.abs(batch[prop] - result[prop])#*batch["eigenvalues_active_mask"]
                    prop_err = torch.mean(prop_diff.view(-1)**2)
                    err_sq = err_sq  + float(tradeoffs[prop])*prop_err
                if args.mixed_loss == True:
                    #outer subtraction
                    deltaE = torch.abs(batch[prop][...,None] - batch[prop][...,None,:])
                    #make diabatic couplings prop to 1/delta E Delta function
                    mask= batch["%s_mask"%prop][...,None]*  batch["%s_mask"%prop][...,None,:]
                    prop_diff=(torch.abs(torch.triu(torch.normal((1/deltaE),0.1),1))*0.01-torch.abs(torch.triu((result['ev_all']['ev_diabatic']),1)))*mask
                    #print(mask,torch.triu(torch.normal((1/deltaE),0.1),1)*0.01,torch.triu((result['ev_all']['ev_diabatic'])))
                    prop_err2 = torch.mean(prop_diff.view(-1) **2 )
                    #diag forces
                    #print(result['ev_all']['ev_diab_forces'][0].shape)
                    #import numpy as np
                    #a=result['ev_all']['ev_diab_forces'][0].to('cpu').detach().numpy()
                    #a=a.reshape((3,3,7,7))
                    #u,e=np.linalg.eigh(a)
                    #print(e.shape,u.shape)
                    #multiply with properr the tradeoff so that the error gets smaller when the iegenvalues get more accurate - was not good
                    err_sq = err_sq + float(tradeoffs['eigenvalues_active_forces'])  *prop_err2
            if prop == "occupation":
                prop_diff = batch[prop] - result[prop]
                prop_err = torch.mean(prop_diff.view(-1)**2)
                err_sq = err_sq + float(tradeoffs[prop])*prop_err
            if prop == "eigenvalues_active_forces" or prop == "eigenvalues_forces" or prop in Charges.properties:
                #batch[prop]=result['ev_all'][prop]
                pass #continue
            if prop == 'occ_orb':
                #if args.mixed_loss==True:
                #    #error on occupied orbitals
                #    prop_diff_occ_orb = batch[prop] - result[prop]
                #    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs[prop])*prop_err_occ_orb/2
                #    #additionally include the sum of the errors
                #    sum_occ_orb_NN = torch.sum(result[prop],1)*2
                #    prop_diff_sum_orbE = batch['sum_occ_orb'] - sum_occ_orb_NN
                #    prop_err_sum_orbE = torch.mean(prop_diff_sum_orbE.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs[prop])*prop_err_sum_orbE/2
                #    #additionally try to fit the total energy and force by taking the sum of occ.orb.energies and delta E
                #    total_E_NN = sum_occ_orb_NN + result['delta_E']
                #    total_F_NN = result['delta_E_forces']
                #    for batch_value in range(result['energy'].shape[0]):
                #        total_F_NN[batch_value][0] += torch.sum(result['occ_orb_forces'][batch_value],0)
                #    prop_diff_totE = batch['energy'] - total_E_NN
                #    prop_diff_totF = batch['forces'] - total_F_NN
                #    prop_err_totE = torch.mean(prop_diff_totE.view(-1)**2)
                #    prop_err_totF = torch.mean(prop_diff_totF.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs['energy']) * prop_err_totE + float(tradeoffs['forces']) * prop_err_totF
                if args.mixed_loss == True:
                    #error of eigenvalues based on distribution function
                    prop_diff_occ_orb =torch.normal(batch[prop],stddev[prop])-torch.normal(result[prop],stddev[prop])
                    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                    err_sq = err_sq+float(tradeoffs[prop])*prop_err_occ_orb
                else:
                    prop_diff_occ_orb = batch[prop] - result[prop]
                    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                    err_sq = err_sq + float(tradeoffs[prop])*prop_err_occ_orb
                    #print("err occ", prop_err_occ_orb)
            if prop == "unocc_orb":
                n_unocc=n_orb['n_unocc']
                prop_diff = batch[prop][:2] - result[prop][:2]
                prop_err = torch.mean(prop_diff.view(-1)**2)
                err_sq = err_sq + float(tradeoffs[prop])*prop_err
                #print(prop_err)
            if args.verbose:
                print(prop, prop_err)

        return err_sq

    trainer = spk.train.Trainer(args.modelpath, model, loss, optimizer,
                                train_loader, val_loader, hooks=hooks)
    trainer.train(device)

def evaluate(args, model, train_loader, val_loader, test_loader, device):
    # Get property names from model
    if args.parallel:
        model=model.module.to(device)
        properties=model.output_modules[0].properties
    else:
        properties = model.output_modules[0].properties
    header = ['Subset']
    metrics = []
    for prop in properties:
      if prop == "eigenvalues_active_forces" or prop== "eigenvalues_forces":
        pass
      else:
        header += [f'{prop}_MAE', f'{prop}_RMSE']
        metrics += [
                spk.metrics.MeanAbsoluteError(prop, prop),
                spk.metrics.RootMeanSquaredError(prop, prop)
            ]

    results = []
    if ('train' in args.split) or ('all' in args.split):
        logging.info('Training split...')
        results.append(['training'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, train_loader, device,properties)])

    if ('validation' in args.split) or ('all' in args.split):
        logging.info('Validation split...')
        results.append(['validation'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, val_loader, device,properties)])
    if ('test' in args.split) or ('all' in args.split):
        logging.info('Testing split...')
        results.append(['test'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, test_loader, device,properties)])
    header = ','.join(header)
    results = np.array(results)

    np.savetxt(os.path.join(args.modelpath, 'evaluation.csv'), results, header=header, fmt='%s', delimiter=',')

def evaluate_dataset(metrics, model, loader, device,properties):
    # TODO: Adapt for SCHNARC, copy old
    for metric in metrics:
        metric.reset()

    qm_values={}
    predicted={}
    batchs={}
    header=[]
    results2={}
    """for batch in loader:
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        result = model(batch)
        for prop in result:
            if prop =="eigenvalues_active":
              if prop in result:
                result['ev_all'] += result[prop]
                result[prop] += result[prop][prop]
              else:
                result['ev_all'] = result[prop]
                result[prop]=result[prop][prop]"""
    #do only 1000 batches for testing
    #index_=-1
    for batch in loader:
    #    index_ +=1
    #    if index_==int(1000):
    #       break
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        result = model(batch)
        for prop in result:
          if prop == "eigenvalues_active":
              if "eigenvalues_active_forces" in properties:
                if prop in predicted:
                    predicted[prop] += [(result[prop][prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop][prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
              else:
                if prop in predicted:
                    predicted[prop] += [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
          if prop =="eigenvalues" or prop in Eigenvalue_properties.properties:
                if prop in predicted:
                    predicted[prop] += [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
          if prop == "delta_eigenvalues_pbe0_gbw" or prop=="delta_eigenvalues_pbe0_h2o":
            if prop in predicted:
                predicted[prop] += [(result[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
            else:
                predicted[prop] = [(result[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
          if prop == "energy":
            if prop in predicted:
                predicted[prop] += [(result[prop]).cpu().detach().numpy()]
            else:
                predicted[prop] = [(result[prop]).cpu().detach().numpy()]

        for prop in result:
          if prop =="eigenvalues_active":
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                    qm_values["%s_mask"%prop] +=[batch["%s_mask"%prop].cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                    qm_values["%s_mask"%prop] = [batch["%s_mask"%prop].cpu().detach().numpy()]
          if prop == "delta_eigenvalues_pbe0_gbw" or prop =="delta_eigenvalues_pbe0_h2o":
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
                    qm_values["eigenvalues_gbw_mask"] +=[batch["eigenvalues_gbw_mask"].cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
                    qm_values["eigenvalues_gbw_mask"] = [batch["eigenvalues_gbw_mask"].cpu().detach().numpy()]

          if prop =="eigenvalues" or prop in Eigenvalue_properties.properties:
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]).cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]).cpu().detach().numpy()]
          if prop == "energy":
            if prop in qm_values:
                qm_values[prop] += [(batch[prop]).cpu().detach().numpy()]
            else:
                qm_values[prop] = [(batch[prop]).cpu().detach().numpy()]
        for metric in metrics:
          if "eigenvalues_active" in result:
              prop="eigenvalues_active"
              batch[prop]=batch[prop]*batch["%s_mask"%prop]
              if "eigenvalues_active_forces" in properties:
                  result[prop][prop] = result[prop][prop]*batch["%s_mask"%prop]
                  metric.add_batch(batch,result[prop])
              else:
                  result[prop]=  result[prop]#*batch["%s_mask"%prop][:len(result[prop][0])]
                  metric.add_batch(batch,result)
          else:
              metric.add_batch(batch,result)
    results = [
    metric.aggregate() for metric in metrics
    ]
    for prop in result:
        if prop == "eigenvalues_active" or prop=="eigenvalues":
            if "eigenvalues_active_forces" in properties:
                predicted['diabatic_H'] = [result[prop]['ev_diabatic'].cpu().detach().numpy()]
                #predicted['eigenvalues_active_forces'] = [result[prop]['eigenvalues_active_forces'].cpu().detach().numpy()]
    for p in predicted.keys():
        predicted[p]=np.vstack(predicted[p])
        if p in qm_values.keys():
            qm_values[p]=np.vstack(qm_values[p])
    if "eigenvalues_active_mask" in qm_values.keys():
        qm_values['eigenvalues_active_mask'] = np.array(qm_values["eigenvalues_active_mask"])
    if "eigenvalues_gbw_mask" in qm_values.keys():
        qm_values["eigenvalues_gbw_mask"] = np.array(qm_values["eigenvalues_gbw_mask"])
    if "eigenvalues_mask" in qm_values.keys():
        qm_values["eigenvalues_mask"] = np.array(qm_values["eigenvalues_mask"])
    prediction_path = os.path.join(args.modelpath,"evaluation_values.npz")
    prediction_path_qm = os.path.join(args.modelpath,"evaluation_qmvalues.npz")
    np.savez(prediction_path,**predicted)
    np.savez(prediction_path_qm,**qm_values)
    logging.info('Stored model predictions in {:s} ...'.format(prediction_path))

    return results

def run_prediction(model, loader, device, args):
    from tqdm import tqdm
    import numpy as np

    predicted = {}
    qm_values = {}
    for batch in tqdm(loader, ncols=120):
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        if device != "cuda" and args.parallel==True:
	        result = model.module(batch)
        else:
                result = model(batch)
        for prop in result:
            if prop in predicted:
                predicted[prop] += [result[prop].cpu().detach().numpy()]
            else:
                predicted[prop] = [result[prop].cpu().detach().numpy()]
        for prop in batch:
            if prop in qm_values:
                qm_values[prop] += [batch[prop].cpu().detach().numpy()]
            else:
                qm_values[prop] = [batch[prop].cpu().detach().numpy()]

    for p in predicted.keys():
        predicted[p] = np.vstack(predicted[p])
    #for p in qm_values.keys():
    #    qm_values[p]=np.vstack(qm_values[p])
    #prediction_path_qm = os.path.join(args.modelpath,"evaluation_qmvalues.npz")
    #np.savez(prediction_path_qm,**qm_values)
    prediction_path = os.path.join(args.modelpath, 'predictions.npz')
    np.savez(prediction_path, **predicted)
    logging.info('Stored model predictions in {:s}...'.format(prediction_path))

def get_model(args, n_orb, properties, atomref=None, mean=None, stddev=None, train_loader=None, parallelize=False,
              mode='train'):
    if args.model == 'schnet':
        representation = spk.representation.SchNet(args.features, args.features, args.interactions,
                                                   args.cutoff / units.Bohr, args.num_gaussians)

        property_output = schnet_ev.model.MultiStatePropertyModel(args.features, n_orb,n_neurons=args.n_neurons, properties=properties,
                                                                mean=mean, stddev=stddev, atomref=atomref,
                                                                n_layers=args.n_layers, diabatic=args.diabatic,hamiltonian = args.hamiltonian,
                                                                inverse_energy=args.inverse_nacs, force_mask = args.force_mask, force_mask_index = args.force_mask_index)

        model = spk.atomistic.AtomisticModel(representation, property_output)

    elif args.model == 'invD':
        representation = spk.representation.schnet.invD()
        n_in = (args.n_atoms*args.n_atoms-args.n_atoms)
        #neurons = number of neurons in each layer of the network; if None, it is divided by two in each layer
        property_output = schnet_ev.model.MultiStatePropertyModel(n_in, n_orb,n_neurons=args.n_neurons, properties=properties,
                                                                mean=mean, stddev=stddev, atomref=atomref,
                                                                n_layers=args.n_layers, diabatic=args.diabatic,hamiltonian = args.hamiltonian,
                                                                inverse_energy=args.inverse_nacs, force_mask=args.force_mask, force_mask_index = args.force_mask_index)

        model = spk.atomistic.AtomisticModel(representation, property_output)

    else:
        raise ValueError('Unknown model class:', args.model)

    if parallelize:
        model = nn.DataParallel(model)

    logging.info("The model you built has: %d parameters" % spk.utils.count_params(model))

    return model

def orca_opt(args, model, device):
    # reads QM.in and writes QM.out for optimizations with ORCA

    #read QMin
    QMin=read_QMin("QM.in")
    # get input for schnet
    schnet_inputs = QMin2schnet(QMin,device)
    #perform prediction
    schnet_outputs = model(schnet_inputs)
    for key,value in schnet_outputs.items():
        schnet_outputs[key] = value.cpu().detach().numpy()
    #transform to QM.out
    get_QMout(schnet_outputs, QMin,args)
    return

def get_QMout(schnet_outputs,QMin,args):
    from transform_prediction_QMout import QMout
    QMout(schnet_outputs,args.modelpath,args.nac_approx,QMin)
    return

def QMin2schnet(QMin,args):
    from schnetpack import Properties
    from ase import Atoms
    from schnetpack.environment import SimpleEnvironmentProvider, collect_atom_triples

    #environment_provider = SimpleEnvironmentProvider()
    environment_provider = get_environment_provider(args,device=device)
    schnet_inputs = dict()
    molecule = Atoms(QMin['atypes'],QMin['geom'])
    schnet_inputs[Properties.Z] = torch.LongTensor(molecule.numbers.astype(np.int))
    schnet_inputs[Properties.atom_mask] = torch.ones_like(schnet_inputs[Properties.Z]).float()
    #set positions
    positions = molecule.positions.astype(np.float32)
    schnet_inputs[Properties.R] = torch.FloatTensor(positions)
    #get atom environment
    nbh_idx, offsets = environment_provider.get_environment(molecule)
    #neighbours and neighour mask
    mask = torch.FloatTensor(nbh_idx) >= 0
    schnet_inputs[Properties.neighbor_mask] = mask.float()
    schnet_inputs[Properties.neighbors] = torch.LongTensor(nbh_idx.astype(np.int)) * mask.long()
    #get cells
    schnet_inputs[Properties.cell] = torch.FloatTensor(molecule.cell.astype(np.float32))
    schnet_inputs[Properties.cell_offset]= torch.FloatTensor(offsets.astype(np.float))
    collect_triples = False
    if collect_triples is not None:
        nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k = collect_atom_triples(nbh_idx)
        schnet_inputs[Properties.neighbor_pairs_j] = torch.LongTensor(nbh_idx_j.astype(np.int))
        schnet_inputs[Properties.neighbor_pairs_k] = torch.LongTensor(nbh_idx_k.astype(np.int))
        schnet_inputs[Properties.neighbor_pairs_mask] = torch.ones_like(schnet_inputs[Properties.neighbor_pairs_j]).float()                                                                 # Add batch dimension and move to CPU/GPU
    for key, value in schnet_inputs.items():
        schnet_inputs[key] = value.unsqueeze(0).to(device)

    return schnet_inputs

def read_QMin(filename):
    A2Bohr = 1/0.529177211
    data = open(filename, 'r').readlines()
    QMin = {}
    natoms = int(data[0].split()[0])
    QMin['NAtoms'] = natoms
    #unit
    for line in data:
        if "angstrom" in line or "Angstrom" in line:
            angstrom = True
            break
        else:
            angstrom = False
    #states
    for line in data:
        if "states" in line:
            s=line.split()[1:]
            break
    states = [ int(i) for i in s]
    n_doublets = 0
    n_unocc = 0
    n_quartets = 0
    for index,istate in enumerate(states):
        if index == int(0):
            n_occ = istate
        elif index == int(1):
            n_doublets = istate
        elif index == int(2):
            n_unocc = istate
        elif index == int(3):
            n_quartets = istate
    nmstates = n_occ + 2*n_doublets + 3*n_unocc + 4*n_quartets
    QMin['n_occ'] = n_occ
    QMin['n_doublets'] = n_doublets
    QMin['n_unocc'] = n_unocc
    QMin['n_quartets'] = n_quartets
    QMin['n_orb'] = nmstates

    #geometries
    atypes = []
    geom = []
    for curr_atom in range(natoms):
        atom_data = data[curr_atom+2].strip().split()
        curr_atype = atom_data[0]
        if angstrom == True:
            curr_geom = [float(coordinate) for coordinate in atom_data[1:4]]
            for xyz in range(3):
                curr_geom[xyz] = curr_geom[xyz] * A2Bohr
        else:
            curr_geom = [float(coordinate) for coordinate in atom_data[1:4]]
        atypes.append(curr_atype)
        geom.append(curr_geom)
    QMin['geom'] = np.array(geom)
    QMin['atypes'] = np.array(atypes)

    return QMin

def get_deltaE(model):
    return

if __name__ == '__main__':
    # Parse command line arguments
    parser = get_parser()
    args = parser.parse_args()

    # Determine the device
    device = torch.device("cuda" if args.cuda else "cpu")
    if args.mode != 'train':
        model = torch.load(os.path.join(args.modelpath, 'best_model'), map_location='cpu').to(device)
        if args.hessian == True:
            model.output_modules[0].output_dict['energy'].return_hessian = [True,1,1,1,1]
        else:
            if args.parallel==False:
                for p in  model.output_modules[0].output_dict.keys():
                    model.output_modules[0].output_dict[p].return_hessian = [False,1,1]
            else:
                for p in model.module.output_modules[0].output_dict.keys():
                    model.module.output_modules[0].output_dict[p].return_hessian = [False,1,1]

    if args.mode == 'opt':
        orca_opt(args, model, device)
        sys.exit()

    if args.mode == 'pred':
        pred_data = spk.data.AtomsData(args.datapath)
        pred_loader = spk.data.AtomsLoader(pred_data, batch_size=args.batch_size, num_workers=2, pin_memory=True)
        run_prediction(model, pred_loader, device, args)
        sys.exit(0)

    # Load settings
    argparse_dict = vars(args)
    jsonpath = os.path.join(args.modelpath, 'args.json')

    # Setup directories and files for training or load settings for restart.
    if args.mode == 'train':
        if args.overwrite and os.path.exists(args.modelpath):
            logging.info('Existing model will be overwritten...')
            rmtree(args.modelpath)

        if not os.path.exists(args.modelpath):
            os.makedirs(args.modelpath)

        spk.utils.to_json(jsonpath, argparse_dict)

        spk.utils.set_random_seed(args.seed)
        train_args = args

        # Read the tradeoff file or generate a default file
        tradeoff_file = os.path.join(args.modelpath, 'tradeoffs.yaml')

        if train_args.tradeoffs is None:
            schnet_ev.utils.generate_default_tradeoffs(tradeoff_file)
            tradeoffs = schnet_ev.utils.read_tradeoffs(tradeoff_file)
        else:
            tradeoffs = schnet_ev.utils.read_tradeoffs(args.tradeoffs)
            schnet_ev.utils.save_tradeoffs(tradeoffs, tradeoff_file)
    else:
        train_args = spk.utils.read_from_json(jsonpath)
        tradeoff_file = os.path.join(args.modelpath, 'tradeoffs.yaml')
        tradeoffs = schnet_ev.utils.read_tradeoffs(tradeoff_file)

    # Determine the properties to load based on the tradeoffs
    properties = [p for p in tradeoffs]
    if "eigenvalues_forces" in properties:
        tradeoffs['eigenvalues_active_forces']=tradeoffs['eigenvalues_forces']
    # Read and process the data using the properties found in the tradeoffs.
    logging.info('Loading {:s}...'.format(args.datapath))
    environment_provider = get_environment_provider(train_args,device=torch.device("cuda" if args.cuda else "cpu"))
    #previous function:
    dataset = spk.data.AtomsData(args.datapath, environment_provider=environment_provider, collect_triples=args.model == 'wacsf')
    # Determine the number of states based on the metadata
    n_orb = {}
    if "mean_energy" in dataset.get_metadata():
        mean_energy = dataset.get_metadata('mean_energy')
    else:
        mean_energy = int(0)
    n_orb['mean_energy']=mean_energy
    n_orb['n_act'] = int(1)
    offset=0
    for prop in properties:
        if prop in Eigenvalue_properties.properties or prop in Delta_EV_properties.properties:

        # if "eigenvalues" in properties or "eigenvalues_active" in properties or "eigenvalues_pbe0" in properties or "eigenvalues_pbe" in properties or "delta_eigenvalues_gbw" in properties or "eigenvalues_gbw" in properties:
            n_orb['n_act'] = dataset.get_metadata('active')
            if dataset.get_metadata("active_all"):
                n_orb["active_all"]= dataset.get_metadata("active_all")
            if dataset.get_metadata('offset'):
                offset = dataset.get_metadata('offset')
            else:
                offset = float(0.0)
    ##activate if only one state is learned or not all
    # only needed if unocc and occ orbitals are trained separately
    n_orb['n_occ']=int(1)
    n_orb['n_unocc']=int(0)
    n_orb['n_orb']=n_orb['n_occ']+n_orb['n_unocc']
    n_orb['offset']=offset
    n_orb["aggregation_mode"]=args.aggregation_mode
    if args.mode == 'eval':
        split_path = os.path.join(args.modelpath, 'split.npz')
        data_train, data_val, data_test = dataset.create_splits(*args.split, split_file=split_path)

        train_loader = spk.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
                                        num_workers=4, pin_memory=True)
        val_loader = spk.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=2, pin_memory=True)

        logging.info("evaluating...")
        test_loader = spk.data.AtomsLoader(data_test, batch_size=args.batch_size,
                                           num_workers=2, pin_memory=True)
        evaluate(args, model, train_loader, val_loader, test_loader, device)
        logging.info("... done!")
        exit()
     # Splits the dataset in test, val, train sets
    split_path = os.path.join(args.modelpath, 'split.npz')

    if args.mode == 'train':
        if args.split_path is not None:
            copyfile(args.split_path, split_path)

    logging.info('Creating splits...')
    data_train, data_val, data_test = dataset.create_splits(*train_args.split, split_file=split_path)

    # Generate loaders for training
    train_loader = spk.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
                                        num_workers=4, pin_memory=True)
    val_loader = spk.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=2, pin_memory=True)

    # Determine statistics for scaling
    if args.mode == 'train':
        # Set the random seed
        spk.utils.set_random_seed(args.seed)

        logging.info('Calculate statistics...')
        mean = {}
        stddev = {}
        # Compute mean and stddev for relevant properties
        for p in properties:
            if p in schnet_ev.data.Properties.normalize:
                mean_p, stddev_p = train_loader.get_statistics(p, True)
                mean_p = mean_p[p]
                stddev_p = stddev_p[p]
                mean[p] = mean_p
                stddev[p] = stddev_p
                logging.info('{:s} MEAN: {:20.11f} STDDEV: {:20.11f}'.format(p, mean_p.detach().cpu().numpy()[0],
                                                                             stddev_p.detach().cpu().numpy()[0]))
            else:
                mean[p], stddev[p] = None, None

    # Construct the model.
    model = get_model(train_args,
                      n_orb,
                      properties,
                      mean=mean,
                      stddev=stddev,
                      train_loader=train_loader,
                      parallelize=args.parallel,
                      mode=args.mode
                      ).to(device)

    if args.mode == 'train':
        logging.info("training...")
        batch_size = args.batch_size
        train(args, model, tradeoffs, train_loader, val_loader, device, n_orb,mean,stddev,offset)
        logging.info("...training done!")


"""
6083. 判断一个数的数字计数是否等于数位的值
给你一个下标从 0 开始长度为 n 的字符串 num ，它只包含数字。
如果对于 每个 0 <= i < n 的下标 i ，都满足数位 i 在 num 中出现了 num[i]次，那么请你返回 true ，否则返回 false 。
"""
import collections

def digitCount(num: str) -> bool:
    numCounter = collections.Counter(num)
    for i, char in enumerate(num):
        if numCounter[str(i)] != int(char):
            return False
    return True

digitCount(num = "1210")

Leetcode/6083. 判断一个数的数字计数是否等于数位的值.py

6083. 判断一个数的数字计数是否等于数位的值.py


# -*- coding: utf-8 -*-
"""
Created on Tue Nov 16 09:07:08 2021

@author: zjh
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import numpy as np
from configure import BASEDIR,BLOCK_SIZE,NUM_CLASSES,IN_CHANNELS
from dataloaders.datasets import SegmentationDataset
from torch.utils.data import DataLoader
import json
import torch
from modeling.myscse import SCSEUnet
from modeling.loss import SegmentationLosses
from utils.metrics import Evaluator
from torch.utils.tensorboard import SummaryWriter
from torchsummary import summary

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def load_data():
    images_dir = os.path.join(BASEDIR,'data/raw/')
    masks_dir = os.path.join(BASEDIR,'data/processed/')

    b = open(os.path.join(masks_dir,'train_file_list_11_15_14_06.txt'), "r",encoding='UTF-8')
    out_train = b.read()
    train_file_names = json.loads(out_train)

    c = open(os.path.join(masks_dir,'val_file_list_11_15_14_06.txt'), "r",encoding='UTF-8')
    out_val = c.read()
    val_file_names = json.loads(out_val)

    d = open( os.path.join(masks_dir,'test_file_list_11_15_14_06.txt'), "r",encoding='UTF-8')
    out_test = d.read()
    test_file_names = json.loads(out_test)

    # 构造训练loader
    train_set = SegmentationDataset(train_file_names,images_dir,masks_dir,BLOCK_SIZE, transform_name = 'train_transform_1')

    # 构造验证loader
    val_set = SegmentationDataset(val_file_names,images_dir,masks_dir,BLOCK_SIZE, transform_name = None)

    # 构造测试loader
    test_set = SegmentationDataset(test_file_names,images_dir,masks_dir,BLOCK_SIZE, transform_name = None)

    return train_set, val_set, test_set

def train(learning_rate,batch_size):
    net = SCSEUnet(IN_CHANNELS,NUM_CLASSES)
    net = net.to(device)
    summary(net,(IN_CHANNELS,BLOCK_SIZE,BLOCK_SIZE))

    #loss and optimizer
    criterion = SegmentationLosses().FocalLoss
    optimizer = torch.optim.Adam(net.parameters(),lr=learning_rate)

    train_set, val_set, test_set=load_data()
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4)
    valid_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=4)

    evaluator = Evaluator(NUM_CLASSES)

    #train the model
    num_epochs=200
    for epoch in range(num_epochs):
        print(f"Epoch {epoch+1}\n------------------")

        #training
        running_loss=0.0
        for batch_idx, data in enumerate(train_loader):
            inputs,labels=data['image'],data['mask']
            inputs, labels = inputs.to(device), labels.to(device)

            #compute prediction and loss
            outputs=net(inputs)

            # print('train loop output shape',outputs.shape)
            loss = criterion(outputs, labels)
            # loss = criterion(outputs, labels,gamma=0.1,alpha=3)#bridge

            #backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss+=loss.item()
        running_loss/=len(train_loader)
        print(f"loss:{running_loss:>7f}")
                # print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

        # validation loss
        evaluator.reset()  # 创建全为0的混淆矩阵
        num_batches = len(valid_loader)
        # print('length of dataloader', num_batches)
        # num_images = len(valid_loader.dataset)
        # print('length of dataset', num_images)
        val_loss=0

        with torch.no_grad():
            for data in valid_loader:
                X,target=data['image'],data['mask']
                X,target=X.to(device),target.to(device)
                outputs=net(X)

                val_loss+= criterion(outputs,target).item()

                outputs=outputs.argmax(1)
                # print('gt ',target)
                # print('output ', outputs)
                # break

                # Add batch sample into evaluator
                target = target.cpu().numpy()
                pred = outputs.data.cpu().numpy()

                evaluator.add_batch(target, pred)

        FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()  # 频权交并比
        val_loss/=num_batches
        # correct/=num_images
        print(f"Validation Error: FWIoU Accuracy: {(100*FWIoU):>0.3f}%, Avg loss: {val_loss:>8f} \n")

    print('Finished Training')
    torch.save(net, '/home/zjh/road_extraction/road_extraction/run/batch16_model.pth')

def test_accuracy(net, device='cpu'):
    from PIL import Image
    net=net.to(device)
    train_set, val_set, test_set=load_data()
    test_loader = DataLoader(test_set,batch_size=1, shuffle=False)
    #test the model
    evaluator = Evaluator(NUM_CLASSES)
    with torch.no_grad():
        id=0
        for data in test_loader:
            inputs,labels=data['image'],data['mask']

            img=inputs[0].numpy()
            img_ = Image.fromarray(np.uint8(img.transpose(1,2,0)))
            img_.save('/home/zjh/tmp/road_img_%s.bmp'%id)
            mask=labels[0].numpy()
            gt = Image.fromarray(mask)
            gt =gt.convert("L")
            gt.save('/home/zjh/tmp/road_gt_%s.bmp'%id)

            inputs, labels = inputs.to(device), labels.to(device)
            outputs=net(inputs)# with shape NCHW

            # _, predicted = torch.max(outputs.data, 1)
            predicted=outputs.argmax(1)

            # Add batch sample into evaluator
            target = labels.cpu().numpy()
            pred = predicted.data.cpu().numpy()
            # print('predicted max shape ', pred.shape)

            evaluator.add_batch(target, pred)

            pred_arr = pred[0].astype('int8')
            im = Image.fromarray(pred_arr)
            im =im.convert("L")
            im.save('/home/zjh/tmp/road_test_bridge_%s.bmp'%id)

            id+=1

        FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()  # 频权交并比

        print(f'FWIoU Accuracy of the network on the test images: {(FWIoU*100):>0.3f} %')
def tta_test(net):
    from PIL import Image
    model=net.to(device)
    train_set, val_set, test_set=load_data()
    test_loader = DataLoader(test_set,batch_size=1, shuffle=False)
    #test the model
    evaluator = Evaluator(NUM_CLASSES)
    with torch.no_grad():
        id=0
        for data in test_loader:
            inputs,labels=data['image'],data['mask']

            inputs, labels = inputs.to(device), labels.to(device)

            output=model(inputs)
            #horizontalflip
            hfliped_image=inputs.flip(3)
            hfliped_label=model(hfliped_image)
            label1=hfliped_label.flip(3)
            # print('label1, ',label1)

            #verticalflip
            vfliped_image=inputs.flip(2)
            vfliped_label=model(vfliped_image)
            label2=vfliped_label.flip(2)

            #rotate90
            rotated_image=torch.rot90(inputs, 1,(2,3))
            rotated_label=model(rotated_image)
            label3=torch.rot90(rotated_label, 3,(2,3))

            # outputs=net(inputs)# with shape NCHW

            # augmented labels merge
            outputs =(output+label1+label2+label3)/4
            # print('mean label, ',outputs)
            # _, predicted = torch.max(outputs.data, 1)
            predicted=outputs.argmax(1)

            # Add batch sample into evaluator
            target = labels.cpu().numpy()
            pred = predicted.data.cpu().numpy()
            # print('predicted max shape ', pred.shape)

            evaluator.add_batch(target, pred)

            pred_arr = pred[0].astype('int8')
            im = Image.fromarray(pred_arr)
            im =im.convert("L")
            im.save('/home/zjh/tmp/road_test_bridge_%s.bmp'%id)

            id+=1

        FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()  # 频权交并比

        print(f'FWIoU Accuracy of TTA: {(FWIoU*100):>0.3f} %')

    # import ttach as tta
    # # transforms = tta.Compose(
    # #     [
    # #         tta.HorizontalFlip(),
    # #         tta.VerticalFlip(),
    # #         tta.Rotate90(angles=[0, 180]),
    # #         # tta.Scale(scales=[1, 2, 4]),
    # #         # tta.Multiply(factors=[0.9, 1, 1.1]),
    # #     ]
    # # )
    # # assert len(transforms) == 2 * 2 * 2   # all combinations for aug parameters

    # model=model.to(device)
    # train_set, val_set, test_set=load_data()
    # test_loader = DataLoader(test_set,batch_size=1, shuffle=False)
    # #test the model
    # evaluator = Evaluator(NUM_CLASSES)
    # transforms=[tta.HorizontalFlip(),
    #         tta.VerticalFlip(),
    #         tta.Rotate90(angles=[0, 180]),]

    # with torch.no_grad():
    #     id=0
    #     for data in test_loader:
    #         image,labels=data['image'],data['mask']

    #         #horizontalflip
    #         hfliped_image=image.flip(3)
    #         hfliped_label=model(hfliped_image)
    #         label1=hfliped_label.flip(3)

    #         #verticalflip
    #         vfliped_image=image.flip(2)
    #         vfliped_label=model(vfliped_image)
    #         label2=vfliped_label.flip(2)

    #         #rotate90
    #         rotated_image=torch.rot90(image, 1,(2,3))
    #         rotated_label=model(rotated_image)
    #         label3=torch.rot90(rotated_label, 3,(2,3))

    #         #TODO 取三个label各元素的众数，考虑使用pandas

    #         for transformer in transforms: # custom transforms or e.g. tta.aliases.d4_transform()
    #             # augment image
    #             augmented_image = transformer.apply_aug_image(image,apply=True)

    #             # pass to model
    #             model_output = model(augmented_image)

    #             # reverse augmentation for mask and label
    #             # deaug_mask = transformer.deaugment_mask(model_output)
    #             deaug_label = transformer.deaugment_label(model_output)
                # save results
    #             labels.append(deaug_mask)
    #             masks.append(deaug_label)

    # # reduce results as you want, e.g mean/max/min
    # label = mean(labels)
    # mask = mean(masks)
# valid_epoch = smp.utils.train.ValidEpoch(
#     model,
#     loss=loss,
#     metrics=metrics,
#     device=DEVICE,
#     verbose=True,
# )

# # train model for 40 epochs
# max_score = 0

# for i in range(0, 40):
#     print('\nEpoch: {}'.format(i))
#     train_logs = train_epoch.run(train_loader)
#     valid_logs = valid_epoch.run(valid_loader)

#     # do something (save model, change lr, etc.)
#     if max_score < valid_logs['iou_score']:
#         max_score = valid_logs['iou_score']
#         torch.save(model, os.path.join(masks_dir,'best_model.pth'))
#         print('Model saved!')

#     if i == 25:
#         optimizer.param_groups[0]['lr'] = 1e-5
#         print('Decrease decoder learning rate to 1e-5!')

# # load best saved checkpoint
# best_model = torch.load( os.path.join(masks_dir,'best_model.pth'))

# logs = test_epoch.run(test_loader)
def main():
    #hyper-parameters
    train(0.0001,16)
    net=torch.load('/home/zjh/road_extraction/road_extraction/run/batch16_model.pth')
    net.eval()
    test_accuracy(net,'cuda')
    tta_test(net)
    print('batch size is 16')

if __name__=="__main__":
    main()

    #使用真实bridge切片,测试focalloss的参数
#     import os,gdal
#     images_dir = os.path.join(BASEDIR,'data/raw/')
#     masks_dir = os.path.join(BASEDIR,'data/processed/')
#     imgds=gdal.Open(os.path.join(images_dir,'4.bmp'))

#     inputs = imgds.ReadAsArray(0,0,512,512).astype(np.float32)
#     gtds=gdal.Open(os.path.join(masks_dir,'4_mask.bmp'))
#     gt=gtds.ReadAsArray(0,0,512,512).astype(np.float32)

#     def reclassify(target,origin_labels, output_labels):
#         h, w = target.shape
#         label = np.zeros((h, w), dtype=np.int8)
#         # label[target == 6] = 1
#         # label[target == 9] = 2
#         # label[target == 7] = 3
#         # label[target == 1] = 4
#         # label[target == 3] = 5
#         for i in range(len(origin_labels)):
#             label[target==origin_labels[i]]=output_labels[i]

#         # label = torch.as_tensor(label, dtype=torch.long)
#         return label

#     gt = reclassify(gt,[0,1,2,3,4],[0,0,0,0,1] )

#     inputs, gt = torch.tensor(np.array([inputs])).to('cuda'),  torch.tensor(np.array([gt])).to('cuda')

#     net=torch.load(os.path.join(BASEDIR,'run/bridge_model.pth'))
#     net.eval()
#     model=net.to(device)
#     pred2=model(inputs)

#     loss = SegmentationLosses(cuda=False)
#     print('CE ', loss.CrossEntropyLoss(pred2,gt).item())
#     print('0,None ',loss.FocalLoss(pred2,gt, gamma=0, alpha=None).item())
#     print('0.01,0.5 ',loss.FocalLoss(pred2,gt, gamma=0.01, alpha=0.5).item())
#     print('0.1,0.5 ',loss.FocalLoss(pred2,gt, gamma=0.1, alpha=0.5).item())
#     print('0.1,2 ',loss.FocalLoss(pred2,gt, gamma=0.1, alpha=2).item())
#     print('0.1,3 ',loss.FocalLoss(pred2,gt, gamma=0.1, alpha=3).item())
#     print('1,0.5 ',loss.FocalLoss(pred2,gt, gamma=1, alpha=0.5).item())
#     print('2,0.5 ',loss.FocalLoss(pred2,gt, gamma=2, alpha=0.5).item())
#     print('3,0.5 ',loss.FocalLoss(pred2,gt, gamma=3, alpha=0.5).item())
#     print('2,1 ',loss.FocalLoss(pred2,gt, gamma=2, alpha=1).item())
#     print('2,2 ',loss.FocalLoss(pred2,gt, gamma=2, alpha=2).item())

#     CE  0.3212392330169678
# 0,None  0.3212392330169678
# 0.01,0.5  0.1585579216480255
# 0.1,0.5  0.1411537379026413
# 0.1,2  0.5646149516105652
# 0.1,3  0.8469224572181702  #small gamma, large alpha, large loss
# 1,0.5  0.04413028806447983
# 2,0.5  0.012124809436500072
# 3,0.5  0.003331294981762767
# 2,1  0.024249618873000145
# 2,2  0.04849923774600029


"""
by Rabbit White.
2019/12/1

Package used: cv2

Mark that
1. the file will be stored into the folder which is the
same as the source .jpg file.
2. each 12-bit data of the coe file is organized
as the order of BGR. That is the bit [11:8] is for Blue,
[7:4] for Green and [3:0] for Red. You should to pay
attention for this when using in VGA displaying.
"""

import cv2

# Describe: use this function to generate coe file.
# Parameters:
#   img_filename: the name of the source file.
#   resize_shape: the shape you want to resize your image to.
#   scale_percent: set the proportion to zoom your picture out or in.
def img_transformer(img_filename, resize_shape):
    # read in img, default as BGR
    img = cv2.imread(img_filename, 1)

    # resize image
    # resized = cv2.resize(img, None, fx=scale_percent, fy=scale_percent, interpolation=cv2.INTER_AREA)
    resized = cv2.resize(img, resize_shape, interpolation=cv2.INTER_AREA)
    # resized = img
    print(resized.shape)

    img = cv2.imwrite('sources/test.resized.jpg', resized)
    # generate the coe list
    coe = []
    for row in resized:
        for pixel in row:
            coe.append(hex(pixel[2] >> 4)[2:] + hex(pixel[1] >> 4)[2:] + hex(pixel[0] >> 4)[2:])

    # Write the coe list into the coe file.
    print()
    f = open(img_filename[:img_filename.rfind('.')+1]+'coe', 'w')
    fmem = open(img_filename[:img_filename.rfind('.')+1]+'mem', 'w')
    f.write('memory_initialization_radix=16;\n')
    f.write('memory_initialization_vector=\n')
    f.write(coe[0])
    fmem.write(coe[0])
    for c in coe[1:]:
        f.write(',\n' + c)
        fmem.write(f'\n{c}')
    f.write(';\n')
    fmem.write('\n')
    f.close()

    # show the img resized
    cv2.imshow("hebe", resized)
    cv2.waitKey(0)

def main():
    img_filename = "sources/test.jpg"
    img_transformer(img_filename=img_filename, resize_shape=(216, 216))

if __name__ == '__main__':
    main()