#!/usr/bin/env python
from __future__ import absolute_import, division, print_function

#%matplotlib inline

import os
import csv
import sys
import yaml
import random
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.utils.data
import torch.utils.data.distributed
import torchvision.utils as vutils
import torchvision.transforms as transforms

from torch.utils.data.dataloader import DataLoader

from PIL import Image
from tqdm import tqdm
from datetime import datetime
from itertools import product
from collections import namedtuple
from collections import OrderedDict
from labml_helpers.module import Module

from utils.classes.DecayLR import DecayLR
from utils.classes.dataloaders import MyDataLoader
from utils.classes.ReplayBuffer import ReplayBuffer
from utils.models.cycle.Generators import Generator
from utils.models.cycle.Discriminator import Discriminator
from utils.tools.yaml_reader import read_yaml
from utils.tools.pytorch_fid.src.pytorch_fid.fid_score import calculate_fid_given_paths

from numpy import cov
from numpy import trace
from numpy import iscomplexobj
from scipy.linalg import sqrtm

# Clear the terminal
os.system("cls")

class GradientPenalty(Module):
    def __call__(self, x: torch.Tensor, f: torch.Tensor):

        batch_size = x.shape[0]

        gradients, *_ = torch.autograd.grad(outputs=f, inputs=x, grad_outputs=f.new_ones(f.shape), create_graph=True)

        gradients = gradients.reshape(batch_size, -1)

        norm = gradients.norm(2, dim=-1)

        return torch.mean((norm - 1) ** 2)

class Epoch:
    def __init__(self):
        self.count = 0
        self.loss = 0
        self.num_correct = 0
        self.start_time = None
        self.duration = None

class Run:
    def __init__(self):
        self.params = None
        self.count = 0
        self.data = []
        self.start_time = None
        self.duration = None

class RunTrainManager:

    """ [ Insert documentation ] """

    def __init__(
        self,
        dataset,
        channels: int,
        parameters: OrderedDict,
        dir_dataset: str = "./dataset",
        dir_outputs: str = "./outputs",
        dir_results: str = "./results",
        dir_weights: str = "./weights",
        save_epoch_freq: int = 1,
        show_image_freq: int = 5,
        show_graph_freq: int = 10,
        validation_percentage: float = 0.0,
    ) -> None:

        """ [ Insert documentation ] """

        # Arguments
        self.parameters = parameters
        self.dataset = dataset
        self.dataset_group = dataset.dataset_group
        self.channels = channels
        self.validation_percentage = validation_percentage

        # Configuration
        self.DIR_DATASET = f"{dir_dataset}/{self.dataset_group}"
        self.DIR_OUTPUTS = f"{dir_outputs}/{self.dataset_group}"
        self.DIR_RESULTS = f"{dir_results}/{self.dataset_group}"
        self.DIR_WEIGHTS = f"{dir_weights}/{self.dataset_group}"
        self.SHOW_IMAGE_FREQ = show_image_freq
        self.SHOW_GRAPH_FREQ = show_graph_freq
        self.SAVE_EPOCH_FREQ = save_epoch_freq
        self.RUN_PATH = None

        # Variables
        self.run = Run()
        self.epoch = Epoch()
        self.net_G_A2B = None
        self.net_G_B2A = None
        self.net_D_A = None
        self.net_D_B = None
        self.loader = None
        self.validation_batch_index = None
        self.batch_is_validation = False
        self.freq_update_discriminator = 2

        # Runs
        self.runs = self.__build_cycle(parameters)

        pass

    """ [ Public ] Starts the training loop """

    def start_cycle(self) -> None:

        # Iterate over every run, based on the configurated params
        for run in self.runs:

            # Clear occupied CUDA memory
            torch.cuda.empty_cache()

            # Set a random seed for reproducibility
            random.seed(run.manual_seed)
            torch.manual_seed(run.manual_seed)

            # Create the directory path for this run
            self.RUN_PATH = self.get_run_path(run, self.dataset.name, self.channels)

            # Make required directories for storing the training output
            self.makedirs(path=os.path.join(self.DIR_WEIGHTS, self.RUN_PATH), dir="weights")
            self.makedirs(path=os.path.join(self.DIR_OUTPUTS, self.RUN_PATH), dir="outputs")

            # Create a per-epoch csv log file
            self.__create_per_epoch_csv_logs()

            # Create Generator and Discriminator models # in_channels; # out_channels
            self.net_G_A2B = Generator(in_channels=self.channels, out_channels=self.channels).to(run.device)
            self.net_G_B2A = Generator(in_channels=self.channels, out_channels=self.channels).to(run.device)
            self.net_D_A = Discriminator(in_channels=self.channels, out_channels=self.channels).to(run.device)
            self.net_D_B = Discriminator(in_channels=self.channels, out_channels=self.channels).to(run.device)

            # Apply weights
            self.net_G_A2B.apply(self.__initialize_weights)
            self.net_G_B2A.apply(self.__initialize_weights)
            self.net_D_A.apply(self.__initialize_weights)
            self.net_D_B.apply(self.__initialize_weights)

            # define loss functions
            self.cycle_loss = torch.nn.L1Loss().to(run.device)
            self.identity_loss = torch.nn.L1Loss().to(run.device)
            self.adversarial_loss = torch.nn.MSELoss().to(run.device)

            # Optimizers
            self.optimizer_G_A2B = torch.optim.Adam(
                self.net_G_A2B.parameters(), lr=run.learning_rate_gen, betas=(0.5, 0.999)
            )
            self.optimizer_G_B2A = torch.optim.Adam(
                self.net_G_B2A.parameters(), lr=run.learning_rate_gen, betas=(0.5, 0.999)
            )
            self.optimizer_D_A = torch.optim.Adam(
                self.net_D_A.parameters(), lr=run.learning_rate_dis, betas=(0.5, 0.999)
            )
            self.optimizer_D_B = torch.optim.Adam(
                self.net_D_B.parameters(), lr=run.learning_rate_dis, betas=(0.5, 0.999)
            )

            # Learning rates
            self.lr_lambda = DecayLR(run.num_epochs, 0, run.decay_epochs).step
            self.lr_scheduler_G_A2B = torch.optim.lr_scheduler.LambdaLR(self.optimizer_G_A2B, lr_lambda=self.lr_lambda)
            self.lr_scheduler_G_B2A = torch.optim.lr_scheduler.LambdaLR(self.optimizer_G_B2A, lr_lambda=self.lr_lambda)
            self.lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(self.optimizer_D_A, lr_lambda=self.lr_lambda)
            self.lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(self.optimizer_D_B, lr_lambda=self.lr_lambda)

            # Buffers train
            self.fake_A_buffer = ReplayBuffer()
            self.fake_B_buffer = ReplayBuffer()

            # Dataloader train set
            self.loader = DataLoader(
                dataset=self.dataset, batch_size=run.batch_size, num_workers=run.num_workers, shuffle=run.shuffle
            )

            # Determine the batch-index from which the validation set starts
            self.validation_batch_index: int = int(round(int(len(self.loader) * (1 - self.validation_percentage)), 0))

            # Keep track of the per-epoch losses
            self.losses_G_A2B, self.losses_G_B2A = [], []
            self.losses_G_A2B_adv, self.losses_G_B2A_adv = [], []
            self.losses_D_A, self.losses_D_B = [], []

            # Keep track of the per-epoch average MSE loss
            self.avg_mse_loss_generated_A_array = []
            self.avg_mse_loss_generated_B_array = []

            # Keep track of the per-epoch average MSE loss
            self.avg_rmse_loss_generated_A_array = []
            self.avg_rmse_loss_generated_B_array = []

            # # Keep track of the per-epoch average MSE loss
            self.avg_fid_score_generated_A_array = []
            self.avg_fid_score_generated_B_array = []

            # Keep track of the FID score
            self.fid_score_A, self.fid_score_B = 0, 0

            # Keep track of the per-epoch noise factor
            self.noise_factor_array = []

            """ Iterate over the epochs in the run """

            # Iterate through all the epochs
            for epoch in range(0, run.num_epochs):

                # Keep track of the per-batch losses during one epoch
                self.batch_losses_G_A2B, self.batch_losses_G_B2A = [], []
                self.batch_losses_G_A2B_adv, self.batch_losses_G_B2A_adv = [], []
                self.batch_losses_D_A, self.batch_losses_D_B = [], []

                self.batch_losses_error_D_real_A, self.batch_losses_error_D_real_B = [], []
                self.batch_losses_error_D_fake_A, self.batch_losses_error_D_fake_B = [], []

                self.full_batch_losses_G_A, self.full_batch_losses_G_B = [], []

                # Keep track of the MSE losses of the generated images
                self.cum_mse_loss_generated_A, self.cum_mse_loss_generated_B = 0, 0
                self.avg_mse_loss_generated_A, self.avg_mse_loss_generated_B = 0, 0

                # Keep track of the RMSE losses of the generated images
                self.cum_rmse_loss_generated_A, self.cum_rmse_loss_generated_B = 0, 0
                self.avg_rmse_loss_generated_A, self.avg_rmse_loss_generated_B = 0, 0

                # Keep track of the FID scores of the generated images
                self.cum_fid_score_generated_A, self.cum_fid_score_generated_B = 0, 0
                self.avg_fid_score_generated_A, self.avg_fid_score_generated_B = 0, 0

                # Create a per-batch csv log file
                self.__create_per_batch_csv_logs(epoch)

                # Set error variables at 0 at the begin of each epoch
                self.__set_error_variables_to_zero()

                """ Iterate over training data using the progress bar """

                # Create progress bar
                self.progress_bar = tqdm(enumerate(self.loader), total=len(self.loader))

                # Iterate over the data loader_train
                for i, data in self.progress_bar:

                    # Determine whether this batch is for training or validation
                    self.batch_is_validation = True if i > self.validation_batch_index else False

                    # Read data
                    self.__read_data(run, data)

                    # Update generator networks
                    self.__update_generators(i)

                    # Add noise to the discriminator input
                    self.__add_discriminator_noise(run, epoch)

                    # Update discriminator networks
                    self.__update_discriminators(i)

                    # Update RMSE loss on generated images
                    self.__update_rmse_loss(i)

                    # Save the real-time output images for every {SHOW_IMG_FREQ} images
                    self.__save_realtime_output(i)

                    # Save per-epoch logs
                    self.__save_per_epoch_logs(epoch)

                    # Save the per-batch losses in a plot
                    self.__save_plot_per_batch(i, epoch)

                    # Print a progress bar in the terminal
                    self.__print_progress(run, i, epoch)

                """ Call the end-of-epoch functions """

                # Update FID Score on the generated images for current epoch, once every 5 epochs
                self.__update_fid_score(run, i)

                # Update learning rates after each epoch
                self.__update_learning_rate()

                # Save the output at the end of each epoch
                self.__save_end_epoch_output(epoch)

                # Save some logs at the end of each epoch
                self.__save_end_epoch_logs(epoch, run)

                # Save the losses in a plot of each epoch
                self.__save_plot_per_epoch()

            """ Save final model """

            # Save last check points, after every run
            torch.save(self.net_G_A2B.state_dict(), f"{self.DIR_WEIGHTS}/{self.RUN_PATH}/net_G_A2B/net_G_A2B.pth")
            torch.save(self.net_G_B2A.state_dict(), f"{self.DIR_WEIGHTS}/{self.RUN_PATH}/net_G_B2A/net_G_B2A.pth")
            torch.save(self.net_D_A.state_dict(), f"{self.DIR_WEIGHTS}/{self.RUN_PATH}/net_D_A/net_D_A.pth")
            torch.save(self.net_D_B.state_dict(), f"{self.DIR_WEIGHTS}/{self.RUN_PATH}/net_D_B/net_D_B.pth")

        pass

    """ [ Private ] Build the training cycle according the parameters """

    def __build_cycle(self, parameters) -> list:

        """ Create a list of runs using all configured parameter combinations """

        run = namedtuple("Run", parameters.keys())

        self.runs = []

        for v in product(*parameters.values()):
            self.runs.append(run(*v))

        return self.runs

    """ [ Private ] Functions to initialize variables and weights """

    def __initialize_weights(self, m):

        """ Custom weights initialization called on a Generator or Discriminator network """

        classname = m.__class__.__name__

        if classname.find("Conv") != -1:
            torch.nn.init.normal_(m.weight, 0.0, 0.02)

        elif classname.find("BatchNorm") != -1:
            torch.nn.init.normal_(m.weight, 1.0, 0.02)
            torch.nn.init.zeros_(m.bias)

        pass

    def __set_error_variables_to_zero(self) -> None:

        """ Variables for error tracking on train set """

        # Per batch error for D_A, D_B
        self.error_D_A, self.error_D_B = 0, 0

        # Average error on D_A, D_B
        self.cum_error_D_A, self.cum_error_D_B = 0, 0
        self.avg_error_D_A, self.avg_error_D_B = 0, 0

        # Per batch error for G
        self.error_G_A2B, self.error_G_B2A = 0, 0

        # Average error on G
        self.cum_error_G_A2B, self.avg_error_G_A2B = 0, 0
        self.cum_error_G_B2A, self.avg_error_G_B2A = 0, 0

        # Average error on G
        self.cum_error_G_A2B_adv_loss, self.avg_error_G_A2B_adv_loss = 0, 0
        self.cum_error_G_B2A_adv_loss, self.avg_error_G_B2A_adv_loss = 0, 0

        """ Variables for error tracking on validation set """

        # Per batch error for D_A, D_B
        self.v__error_D_A, self.v__error_D_B = 0, 0

        # Average error on D_A, D_B
        self.v__cum_error_D_A, self.v__cum_error_D_B = 0, 0
        self.v__avg_error_D_A, self.v__avg_error_D_B = 0, 0

        # Per batch error for G
        self.v__error_G_A2B, self.v__error_G_B2A = 0, 0

        # Average error on G
        self.v__cum_error_G_A2B, self.v__avg_error_G_A2B = 0, 0
        self.v__cum_error_G_B2A, self.v__avg_error_G_B2A = 0, 0

        pass

    """ [ Private ] Functions to create a .csv file """

    def __create_per_batch_csv_logs(self, epoch: int) -> None:

        """ Create a per-batch csv logs file of the current epoch """

        with open(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/logs/EP{epoch}__logs.csv", "w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(
                [
                    "Epoch",
                    "error_D_A",
                    "error_D_B",
                    "error_G_A",
                    "error_G_B",
                    "loss_GAN_A2B",
                    "loss_GAN_B2A",
                    "loss_cycle_ABA",
                    "loss_cycle_BAB",
                    "loss_identity_A2B",
                    "loss_identity_B2A",
                    "v_error_D_A",
                    "v_error_D_B",
                    "v_error_G_A",
                    "v_error_G_B",
                ]
            )

        pass

    def __create_per_epoch_csv_logs(self) -> None:

        """ Create a per-epoch csv logs file of the current run """

        # Create csv for the logs file of this run
        with open(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/logs.csv", "w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(
                [
                    "Epoch",
                    "Noise factor",
                    "FID Score A",
                    "FID Score B",
                    "Average MSE loss A",
                    "Average MSE loss B",
                    "Average RMSE loss A",
                    "Average RMSE loss B",
                    "Average FID score A",
                    "Average FID score B",
                    "avg_error_D_A",
                    "avg_error_D_B",
                    "avg_error_G_A",
                    "avg_error_G_B",
                    "v__avg_error_D_A",
                    "v__avg_error_D_B",
                    "v__avg_error_G_A",
                    "v__avg_error_G_B",
                ]
            )

        pass

    """ [ Private ] Functions to read- and adjust the training data """

    def __random_flip(self, real_label: torch.Tensor, fake_label: torch.Tensor, probability: float = 0.25):

        """ Randomly flip labels following a given probability """

        random_percentage = random.uniform(0, 1)

        if random_percentage > probability:
            return real_label, fake_label
        else:
            return fake_label, real_label

    def __smooth_one_hot(self, true_label: torch.Tensor, classes: int, smoothing: float = 0.25):

        """ Smoothen one-hot encoced labels y_ls = (1 - α) * y_hot + α / K """

        random_noise = random.uniform(1, 2)

        smooth_label = (1 - smoothing) * true_label + (smoothing * random_noise) / classes

        return smooth_label

    def __read_data(self, run, data) -> None:

        """ Read a left-to-right dataset or a stereo-to-disparity dataset """

        # Get image A and image B from a l2r dataset
        if self.dataset_group == "l2r":
            self.real_image_A = data["left"].to(run.device)
            self.real_image_B = data["right"].to(run.device)

        # Get image A and B from a s2d dataset
        elif self.dataset_group == "s2d":
            real_image_A_left = data["A_left"].to(run.device)
            real_image_A_right = data["A_right"].to(run.device)
            real_image_B = data["B"].to(run.device)

            # Store the images non-locally to access it elsewhere
            self.real_image_A_left = real_image_A_left
            self.real_image_A_right = real_image_A_right
            self.real_image_B = real_image_B

            # Add the left- and right into one image (image A)
            # self.real_image_A = real_image_A_left
            self.real_image_A = torch.add(real_image_A_left, real_image_A_right)

            # Normalize again
            if self.channels == 1:
                transforms.Normalize(mean=(0.5), std=(0.5))(self.real_image_A)
            elif self.channels == 3:
                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))(self.real_image_A)

        else:
            raise Exception(f"Can not read input images, given group '{self.dataset_group}' is incorrect.")

        # Real data label is 1, fake data label is 0.
        self.real_label = torch.full((run.batch_size, self.channels), 1, device=run.device, dtype=torch.float32)
        self.fake_label = torch.full((run.batch_size, self.channels), 0, device=run.device, dtype=torch.float32)

        pass

    def __add_discriminator_noise(self, run, epoch: int) -> None:

        """ Add decaying Gaussian noise to discriminator A and B real/fake inputs """

        # Set mean and standard deviation
        mean, std = 0.5, 0.5

        # Calculate until which epoch there is noise
        self.NOISE_UNTIL_PERCENTAGE = 0.0  # default: 0.80
        self.RANDOM_FLIP_PERCENTAGE = 0.0  # default: 0.10
        self.SMOOTHENING_PERCENTAGE = 0.0  # default: 0.10

        # Noise is gone at epoch x
        noise_until_epoch = self.NOISE_UNTIL_PERCENTAGE * run.num_epochs

        # Calculate noise factor
        if self.NOISE_UNTIL_PERCENTAGE == 0:
            self.noise_factor = 0
        elif epoch > 0 and epoch <= int(round(run.num_epochs * self.NOISE_UNTIL_PERCENTAGE)):
            self.noise_factor = round(1 - (epoch / noise_until_epoch), 3)
        elif epoch > int(round(run.num_epochs * self.NOISE_UNTIL_PERCENTAGE)):
            self.noise_factor = 0
        else:
            self.noise_factor = 1

        # Create the noise for the real images
        noise_real_A = (torch.randn(self.real_image_A.size()) * std + mean).to(run.device)
        noise_real_B = (torch.randn(self.real_image_B.size()) * std + mean).to(run.device)

        # Create the noise for the fake images
        noise_fake_A = (torch.randn(self.generated_image_B2A.size()) * std + mean).to(run.device)
        noise_fake_B = (torch.randn(self.generated_image_A2B.size()) * std + mean).to(run.device)

        # Add decaying noise to the real images (only used by discriminator)
        self.real_image_A_noise = self.real_image_A + (noise_real_A * self.noise_factor)
        self.real_image_B_noise = self.real_image_B + (noise_real_B * self.noise_factor)

        # Add decaying noise to the fake images (only used by discriminator)
        self.generated_image_B2A_noise = self.generated_image_B2A + (noise_fake_A * self.noise_factor)
        self.generated_image_A2B_noise = self.generated_image_A2B + (noise_fake_B * self.noise_factor)

        """ Label smoothing and random flipping """

        # Smoothen the labels (only used by the discriminator)
        self.real_smooth_label = self.__smooth_one_hot(self.real_label, 2, self.SMOOTHENING_PERCENTAGE)
        self.fake_smooth_label = self.__smooth_one_hot(self.fake_label, 2, self.SMOOTHENING_PERCENTAGE)

        # Randomly flip the smooth labels (only used by the discriminator)
        self.real_smooth_label, self.fake_smooth_label = self.__random_flip(
            self.real_smooth_label, self.fake_smooth_label, self.RANDOM_FLIP_PERCENTAGE
        )

        pass

    """ [ Private ] Functions to update the networks """

    def __update_generators(self, i: int) -> None:

        """ Update Generator networks: A2B and B2A """

        # Only zero the gradients when using training data
        if self.batch_is_validation == False:

            # Zero the gradients
            self.optimizer_G_A2B.zero_grad()
            self.optimizer_G_B2A.zero_grad()

        """ Generator losses """

        # GAN loss: D_A(G_B2A(B))
        self.generated_image_B2A = self.net_G_B2A(self.real_image_B)
        self.generated_output_A = self.net_D_A(self.generated_image_B2A)
        self.loss_GAN_B2A = self.adversarial_loss(self.generated_output_A, self.real_label)

        # GAN loss: D_B(G_A2B(A))
        self.generated_image_A2B = self.net_G_A2B(self.real_image_A)
        self.generated_output_B = self.net_D_B(self.generated_image_A2B)
        self.loss_GAN_A2B = self.adversarial_loss(self.generated_output_B, self.real_label)

        """ Identity loss: helps to preserve colour and prevent reverse colour in the result  """

        lambda_A2B = 10  # 10 by default
        lambda_B2A = 10  # 10 by default
        lambda_id = 0.5  # 0.5 by default, set to 0.0 for s2d because colour preservation is not desired

        # Identity loss should not be used for s2d datasets (colour shouldn't be preseved in the result)
        if self.dataset_group == "s2d":

            self.loss_identity_B2A = 0
            self.loss_identity_A2B = 0

        elif lambda_id == 0.0:

            self.loss_identity_B2A = 0
            self.loss_identity_A2B = 0

        else:

            # G_B2A should be fed the real A
            self.identity_image_B2A = self.net_G_B2A(self.real_image_A)
            self.loss_identity_B2A = (
                self.identity_loss(self.identity_image_B2A, self.real_image_A) * lambda_B2A * lambda_id
            )

            # G_A2B should be fed the real B
            self.identity_image_A2B = self.net_G_A2B(self.real_image_B)
            self.loss_identity_A2B = (
                self.identity_loss(self.identity_image_A2B, self.real_image_B) * lambda_A2B * lambda_id
            )

        """ Cycle loss, re-generates the original input image. So: A -> B -> A and B -> A -> B """

        # Apply a gradient penalty to the cycle consistency-loss
        # Check if done, read 3.3 https://ssnl.github.io/better_cycles/report.pdf

        gradient_penalty = GradientPenalty()

        self.gradient_penalty_A = gradient_penalty(self.generated_image_B2A, self.generated_output_A)
        self.gradient_penalty_B = gradient_penalty(self.generated_image_A2B, self.generated_output_B)

        # Cycle loss: A -> B -> A
        self.recovered_image_A = self.net_G_B2A(self.generated_image_A2B)
        self.loss_cycle_ABA = (
            self.cycle_loss(self.recovered_image_A, self.real_image_A)
            * lambda_A2B
            # * self.gradient_penalty_B
        )

        # Cycle loss: B -> A -> B
        self.recovered_image_B = self.net_G_A2B(self.generated_image_B2A)
        self.loss_cycle_BAB = (
            self.cycle_loss(self.recovered_image_B, self.real_image_B)
            * lambda_B2A
            # * self.gradient_penalty_A
        )

        """ Calculate the generator errors """

        # Only update weights when using training data
        if self.batch_is_validation == False:

            # Error G_A2B
            self.error_G_A2B = self.loss_GAN_A2B + self.loss_identity_A2B + self.loss_cycle_ABA

            # Error G_B2A
            self.error_G_B2A = self.loss_GAN_B2A + self.loss_identity_B2A + self.loss_cycle_BAB

            # Average error on G_A2B
            self.cum_error_G_A2B += self.error_G_A2B
            self.avg_error_G_A2B = self.cum_error_G_A2B / (i + 1)

            # Average error on G_B2A
            self.cum_error_G_B2A += self.error_G_B2A
            self.avg_error_G_B2A = self.cum_error_G_B2A / (i + 1)

            # Average error on G_A2B (adversarial loss only, so we can compare to the Discriminator)
            self.cum_error_G_A2B_adv_loss += self.loss_GAN_A2B
            self.avg_error_G_A2B_adv_loss = self.cum_error_G_A2B_adv_loss / (i + 1)

            # Average error on G_B2A (adversarial loss only, so we can compare to the Discriminator)
            self.cum_error_G_B2A_adv_loss += self.loss_GAN_B2A
            self.avg_error_G_B2A_adv_loss = self.cum_error_G_B2A_adv_loss / (i + 1)

            # Calculate gradients for G_A and G_B
            self.error_G_A2B.backward()
            self.error_G_B2A.backward()

            # Update the Generator networks
            self.optimizer_G_A2B.step()
            self.optimizer_G_B2A.step()

        else:

            # Validation error G_A2B
            self.v__error_G_A2B = self.loss_GAN_A2B + self.loss_identity_A2B + self.loss_cycle_ABA

            # Validation error G_B2A
            self.v__error_G_B2A = self.loss_GAN_B2A + self.loss_identity_B2A + self.loss_cycle_BAB

            # Average validation error on G_A2B
            self.v__cum_error_G_A2B += self.v__error_G_A2B
            self.v__avg_error_G_A2B = self.v__cum_error_G_A2B / (i + 1)

            # Average validation error on G_B2A
            self.v__cum_error_G_B2A += self.v__error_G_B2A
            self.v__avg_error_G_B2A = self.v__cum_error_G_B2A / (i + 1)

        pass

    def __update_discriminators(self, i: int) -> None:

        """ Update discriminator A """

        # Only zero the gradient when using training data
        if self.batch_is_validation == False:

            # Set D_A gradients to zero
            self.optimizer_D_A.zero_grad()

        # Real A image loss
        self.real_output_A = self.net_D_A(self.real_image_A_noise)
        self.error_D_real_A = self.adversarial_loss(self.real_output_A, self.real_smooth_label)

        # Fake image A loss
        self.generated_image_B2A_noise = self.fake_A_buffer.push_and_pop(self.generated_image_B2A_noise)
        self.generated_output_A = self.net_D_A(self.generated_image_B2A_noise.detach())
        self.error_D_fake_A = self.adversarial_loss(self.generated_output_A, self.fake_smooth_label)

        # Only update weights when using training data
        if self.batch_is_validation == False:

            # Combined loss and calculate gradients
            self.error_D_A = (self.error_D_real_A + self.error_D_fake_A) / 2

            # Cumulative and average error of D_A
            self.cum_error_D_A += self.error_D_A
            self.avg_error_D_A = self.cum_error_D_A / (i + 1)

            # Calculate gradients for D_A
            self.error_D_A.backward()

            # Update D_A weights
            self.optimizer_D_A.step()

        else:

            # Combined loss and calculate gradients
            self.v__error_D_A = (self.error_D_real_A + self.error_D_fake_A) / 2

            # Cumulative and average validation error of D_A
            self.v__cum_error_D_A += self.v__error_D_A
            self.v__avg_error_D_A = self.v__cum_error_D_A / (i + 1)

        """ Update discriminator B """

        # Only zero the gradient when using training data
        if self.batch_is_validation == False:

            # Set D_B gradients to zero
            self.optimizer_D_B.zero_grad()

        # Real B image loss
        self.real_output_B = self.net_D_B(self.real_image_B_noise)
        self.error_D_real_B = self.adversarial_loss(self.real_output_B, self.real_smooth_label)

        # Fake image B loss
        self.generated_image_A2B_noise = self.fake_B_buffer.push_and_pop(self.generated_image_A2B_noise)
        self.generated_output_B = self.net_D_B(self.generated_image_A2B_noise.detach())
        self.error_D_fake_B = self.adversarial_loss(self.generated_output_B, self.fake_smooth_label)

        # Only update weights when using training data
        if self.batch_is_validation == False:

            # Combined loss and calculate gradients
            self.error_D_B = (self.error_D_real_B + self.error_D_fake_B) / 2

            # Cumulative and average error of D_A
            self.cum_error_D_B += self.error_D_B
            self.avg_error_D_B = self.cum_error_D_B / (i + 1)

            # Calculate gradients for D_B
            self.error_D_B.backward()

            # Update D_B weights
            self.optimizer_D_B.step()

        else:

            # Combined validation loss and calculate gradients
            self.v__error_D_B = (self.error_D_real_B + self.error_D_fake_B) / 2

            # Cumulative and average validation error of D_B
            self.v__cum_error_D_B += self.v__error_D_B
            self.v__avg_error_D_B = self.v__cum_error_D_B / (i + 1)

        pass

    def __update_rmse_loss(self, i: int) -> None:

        """ Calculate a per-batch, cumulative and average (R)MSE loss of the generated outputs of generators A2B and B2A """

        # Initiate a mean square error (MSE) loss function
        mse_loss = nn.MSELoss()

        # Calculate the mean square error (MSE) loss
        mse_loss_generated_A = mse_loss(self.generated_image_B2A.detach(), self.real_image_A)
        mse_loss_generated_B = mse_loss(self.generated_image_A2B.detach(), self.real_image_B)

        # Calculate the sum of all mean square error (MSE) losses
        self.cum_mse_loss_generated_A += mse_loss_generated_A
        self.cum_mse_loss_generated_B += mse_loss_generated_B

        # Calculate the average mean square error (MSE) loss
        self.avg_mse_loss_generated_A = self.cum_mse_loss_generated_A / (i + 1)
        self.avg_mse_loss_generated_B = self.cum_mse_loss_generated_B / (i + 1)

        """ Calculate a per-batch, cumulative and average (R)MSE loss of the generated outputs of generators A2B and B2A """

        # Calculate the root mean square error (RMSE) loss
        rmse_loss_generated_A = np.sqrt(mse_loss_generated_A.cpu())
        rmse_loss_generated_B = np.sqrt(mse_loss_generated_B.cpu())

        # Calculate the sum of all root mean square error (RMSE) losses
        self.cum_rmse_loss_generated_A += rmse_loss_generated_A
        self.cum_rmse_loss_generated_B += rmse_loss_generated_B

        # Calculate the average root mean square error (RMSE) loss
        self.avg_rmse_loss_generated_A = self.cum_rmse_loss_generated_A / (i + 1)
        self.avg_rmse_loss_generated_B = self.cum_rmse_loss_generated_B / (i + 1)

        pass

    def __update_fid_score(self, run, i: int) -> None:

        """ Calculate the FID Score """

        if i % 1 == 0:

            paths_B2A = [
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/B2A/fake",
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/B2A/real",
            ]

            paths_A2B = [
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/A2B/fake",
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/A2B/real",
            ]

            self.fid_score_A = calculate_fid_given_paths(paths=paths_B2A, batch_size=run.batch_size, device=run.device)
            self.fid_score_B = calculate_fid_given_paths(paths=paths_A2B, batch_size=run.batch_size, device=run.device)

        pass

    def __update_learning_rate(self) -> None:

        """ Update the learning rate of each network """

        self.lr_scheduler_G_A2B.step()
        self.lr_scheduler_G_B2A.step()
        self.lr_scheduler_D_A.step()
        self.lr_scheduler_D_B.step()

    """ [ Private ] Functions to print-, save- and plot """

    def __print_progress(self, run, i: int, epoch: int) -> None:

        """ Print progress """

        # Print correct output if the current batch is used for training
        if self.batch_is_validation == False:

            self.progress_bar.set_description(
                #
                f"[{self.dataset_group.upper()}][{epoch + 1}/{run.num_epochs}][{i + 1}/{len(self.loader)}]"
                f"[{self.noise_factor:.2f}|{self.SMOOTHENING_PERCENTAGE:.2f}|{self.RANDOM_FLIP_PERCENTAGE:.2f}]  ||  "
                #
                f"D_A, D_B: {self.avg_error_D_A:.3f} ; {self.avg_error_D_B:.3f} || "
                f"G_A2B, G_B2A: {self.avg_error_G_A2B_adv_loss:.3f} ; {self.avg_error_G_B2A_adv_loss:.3f} || "
                f"FID_A, FID_B: {self.fid_score_A:.3f} ; {self.fid_score_B:.3f} ||"
                #
            )

        # Print correct output if the current batch is used for validation # deprecated
        else:

            self.progress_bar.set_description(
                f"[{self.dataset_group.upper()}][{epoch + 1}/{run.num_epochs}][{i + 1}/{len(self.loader)}][nf={self.noise_factor:.3f}]  ||  "
                f"v__avg_error_D_A: {self.v__avg_error_D_A:.3f} ; "
                f"v__avg_error_D_B: {self.v__avg_error_D_B:.3f}  || "
                f"v__avg_error_G_A2B: {self.v__avg_error_G_A2B:.3f} ; "
                f"v__avg_error_G_B2A: {self.v__avg_error_G_B2A:.3f} || "
            )

        pass

    def __save_realtime_output(self, i: int) -> None:

        """ Save per-epoch all the generated images for the FID score calculation """

        # Write images real-time to feature vector directory
        filepath_fake_A = f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/B2A/fake/{i + 1:04d}___fake_A.png"
        filepath_real_A = f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/B2A/real/{i + 1:04d}___real_A.png"
        filepath_fake_B = f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/A2B/fake/{i + 1:04d}___fake_B.png"
        filepath_real_B = f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/features/A2B/real/{i + 1:04d}___real_B.png"

        # Save all the generated (fake) image s
        vutils.save_image(self.generated_image_B2A.detach(), filepath_fake_A, normalize=True)
        vutils.save_image(self.generated_image_A2B.detach(), filepath_fake_B, normalize=True)
        vutils.save_image(self.real_image_A.detach(), filepath_real_A, normalize=True)
        vutils.save_image(self.real_image_B.detach(), filepath_real_B, normalize=True)

        """ Save used data and network outputs, real-time as a png image (useful for visualization during training) """

        if i % self.SHOW_IMAGE_FREQ == 0:

            if self.batch_is_validation == False:

                # Filepath and filename for the real-time TRAINING output images
                (
                    filepath_real_A,
                    filepath_real_B,
                    filepath_real_A_noise,
                    filepath_real_B_noise,
                    filepath_fake_A,
                    filepath_fake_B,
                    filepath_fake_A_noise,
                    filepath_fake_B_noise,
                ) = (
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/real_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/real_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/real_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/real_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/fake_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/fake_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/fake_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/fake_sample_noise.png",
                )

            else:
                #
                #  Filepath and filename for the real-time VALIDATION output images
                (
                    filepath_real_A,
                    filepath_real_B,
                    filepath_real_A_noise,
                    filepath_real_B_noise,
                    filepath_fake_A,
                    filepath_fake_B,
                    filepath_fake_A_noise,
                    filepath_fake_B_noise,
                ) = (
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/v__real_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/v__real_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/v__real_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/v__real_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/v__fake_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/v__fake_sample.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/v__fake_sample_noise.png",
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/v__fake_sample_noise.png",
                )

            """ Save images in corresponding A and D folder """

            # Save real input images
            vutils.save_image(self.real_image_A, filepath_real_A, normalize=True)
            vutils.save_image(self.real_image_B, filepath_real_B, normalize=True)

            # Save real input images with noise
            vutils.save_image(self.real_image_A_noise, filepath_real_A_noise, normalize=True)
            vutils.save_image(self.real_image_B_noise, filepath_real_B_noise, normalize=True)

            # Save the generated (fake) image
            vutils.save_image(self.generated_image_B2A.detach(), filepath_fake_A, normalize=True)
            vutils.save_image(self.generated_image_A2B.detach(), filepath_fake_B, normalize=True)

            # Save the generated (fake) image
            vutils.save_image(self.generated_image_B2A_noise.detach(), filepath_fake_A_noise, normalize=True)
            vutils.save_image(self.generated_image_A2B_noise.detach(), filepath_fake_B_noise, normalize=True)

            """ Save in /realtime folder """

            # Images B2A (disparity -> stereo)
            vutils.save_image(
                self.generated_image_B2A.detach(),
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/B2A_generated_image_B2A.png",
                normalize=True,
            )
            vutils.save_image(
                self.real_image_B.detach(),
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/B2A_real_image_B.png",
                normalize=True,
            )

            # Images A2B (stereo -> disparity)
            vutils.save_image(
                self.generated_image_A2B.detach(),
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/A2B_generated_image_A2B.png",
                normalize=True,
            )
            vutils.save_image(
                self.real_image_A.detach(),
                f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/A2B_real_image_A_stereo.png",
                normalize=True,
            )

            if self.dataset_group == "s2d":
                vutils.save_image(
                    self.real_image_A_left.detach(),
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/A2B_real_image_A_left.png",
                    normalize=True,
                )
                vutils.save_image(
                    self.real_image_A_right.detach(),
                    f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/realtime/A2B_real_image_A_right.png",
                    normalize=True,
                )

        pass

    def __save_per_epoch_logs(self, epoch: int) -> None:

        """ Stores the current errors and losses as a newline to the previously created .csv file """

        # Create csv for the logs file of this run
        with open(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/logs/EP{epoch}__logs.csv", "a+", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(
                [
                    "Epoch",
                    f"{self.error_D_A:.4f}",
                    f"{self.error_D_B:.4f}",
                    f"{self.error_G_A2B:.4f}",
                    f"{self.error_G_B2A:.4f}",
                    f"{self.loss_GAN_A2B:.4f}",
                    f"{self.loss_GAN_B2A:.4f}",
                    f"{self.loss_cycle_ABA:.4f}",
                    f"{self.loss_cycle_ABA:.4f}",
                    f"{self.loss_identity_B2A:.4f}",
                    f"{self.loss_identity_B2A:.4f}",
                    f"{self.v__error_D_A:.4f}",
                    f"{self.v__error_D_B:.4f}",
                    f"{self.v__error_G_A2B:.4f}",
                    f"{self.v__error_G_B2A:.4f}",
                ]
            )

        pass

    def __save_end_epoch_output(self, epoch) -> None:

        """ Stores the current data and output as .png images (currently called at the end of each epoch) """

        # Filepath and filename for the per-epoch output images
        (
            filepath_real_A,
            filepath_real_B,
            filepath_real_A_noise,
            filepath_real_B_noise,
            filepath_fake_A,
            filepath_fake_B,
            filepath_fake_A_noise,
            filepath_fake_B_noise,
        ) = (
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/epochs/EP{epoch}___real_sample.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/epochs/EP{epoch}___real_sample.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/epochs/EP{epoch}___real_sample_noise.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/epochs/EP{epoch}___real_sample_noise.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/epochs/EP{epoch}___fake_sample.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/epochs/EP{epoch}___fake_sample.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/A/epochs/EP{epoch}___fake_sample_noise.png",
            f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/B/epochs/EP{epoch}___fake_sample_noise.png",
        )

        # Save real input images
        vutils.save_image(self.real_image_A, filepath_real_A, normalize=True)
        vutils.save_image(self.real_image_B, filepath_real_B, normalize=True)

        # Save real input images with noise
        vutils.save_image(self.real_image_A_noise, filepath_real_A_noise, normalize=True)
        vutils.save_image(self.real_image_B_noise, filepath_real_B_noise, normalize=True)

        # Save the generated (fake) image
        vutils.save_image(self.generated_image_B2A.detach(), filepath_fake_A, normalize=True)
        vutils.save_image(self.generated_image_A2B.detach(), filepath_fake_B, normalize=True)

        # Save the generated (fake) image
        vutils.save_image(self.generated_image_B2A_noise.detach(), filepath_fake_A_noise, normalize=True)
        vutils.save_image(self.generated_image_A2B_noise.detach(), filepath_fake_B_noise, normalize=True)

        # Check point dir
        model_weights_dir = f"{self.DIR_WEIGHTS}/{self.RUN_PATH}"

        # Check points, save weights after each epoch
        torch.save(self.net_G_A2B.state_dict(), f"{model_weights_dir}/net_G_A2B/net_G_A2B_epoch_{epoch}.pth")
        torch.save(self.net_G_B2A.state_dict(), f"{model_weights_dir}/net_G_B2A/net_G_B2A_epoch_{epoch}.pth")
        torch.save(self.net_D_A.state_dict(), f"{model_weights_dir}/net_D_A/net_D_A_epoch_{epoch}.pth")
        torch.save(self.net_D_B.state_dict(), f"{model_weights_dir}/net_D_B/net_D_B_epoch_{epoch}.pth")

        pass

    def __save_end_epoch_logs(self, epoch, run) -> None:

        """ Saves the noise, mse losses and average errors to a newline in the previously created .csv file (at the end of each epoch) """

        with open(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/logs.csv", "a+", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(
                [
                    epoch,
                    f"{self.noise_factor:.3f}",
                    f"{self.fid_score_A:.3f}",
                    f"{self.fid_score_B:.3f}",
                    f"{self.avg_mse_loss_generated_A:.3f}",
                    f"{self.avg_mse_loss_generated_B:.3f}",
                    f"{self.avg_rmse_loss_generated_A:.3f}",
                    f"{self.avg_rmse_loss_generated_B:.3f}",
                    f"{self.avg_fid_score_generated_A:.3f}",
                    f"{self.avg_fid_score_generated_B:.3f}",
                    f"{self.avg_error_D_A:.5f}",
                    f"{self.avg_error_D_B:.5f}",
                    f"{self.avg_error_G_A2B:.5f}",
                    f"{self.avg_error_G_B2A:.5f}",
                    f"{self.v__avg_error_D_A:.5f}",
                    f"{self.v__avg_error_D_B:.5f}",
                    f"{self.v__avg_error_G_A2B:.5f}",
                    f"{self.v__avg_error_G_B2A:.5f}",
                ]
            )

        pass

    def __save_plot_per_batch(self, i, epoch) -> None:

        """ Append the current per-batch losses to the arrays containing the network losses """

        # Total losses for all networks
        self.batch_losses_G_A2B.append(self.error_G_A2B.cpu().detach().numpy())
        self.batch_losses_G_B2A.append(self.error_G_B2A.cpu().detach().numpy())
        self.batch_losses_D_A.append(self.error_D_A.cpu().detach().numpy())
        self.batch_losses_D_B.append(self.error_D_B.cpu().detach().numpy())

        # Individual
        self.batch_losses_error_D_real_A.append(self.error_D_real_A.cpu().detach().numpy())
        self.batch_losses_error_D_real_B.append(self.error_D_real_B.cpu().detach().numpy())
        self.batch_losses_error_D_fake_A.append(self.error_D_fake_A.cpu().detach().numpy())
        self.batch_losses_error_D_fake_B.append(self.error_D_fake_B.cpu().detach().numpy())

        # Adversarial losses only of all networks
        self.batch_losses_G_A2B_adv.append(self.loss_GAN_A2B.cpu().detach().numpy())
        self.batch_losses_G_B2A_adv.append(self.loss_GAN_B2A.cpu().detach().numpy())

        """ Plot losses """

        if i % self.SHOW_GRAPH_FREQ == 0:

            # Create figure
            self.per_batch_figure, self.per_batch_axes = plt.subplots(4, 1, figsize=(8, 12))

            # Set titles
            self.per_batch_axes[0].set_title(f"Adversarial loss of generators A2B & B2A (epoch {epoch})")
            self.per_batch_axes[1].set_title(f"Adversarial loss of discriminators A & B (epoch {epoch})")
            self.per_batch_axes[2].set_title(f"Adversarial loss of Generator A2B and Discriminator B (epoch {epoch})")
            self.per_batch_axes[3].set_title(f"Adversarial loss of Generator B2A and Discriminator A (epoch {epoch})")

            # Set labels
            self.per_batch_axes[0].set(xlabel="Batch", ylabel="G total Loss")
            self.per_batch_axes[1].set(xlabel="Batch", ylabel="D total Loss")
            self.per_batch_axes[2].set(xlabel="Batch", ylabel="Adv. loss")
            self.per_batch_axes[3].set(xlabel="Batch", ylabel="Adv. loss")

            # Add gridlines
            self.per_batch_axes[0].grid()
            self.per_batch_axes[1].grid()
            self.per_batch_axes[2].grid()
            self.per_batch_axes[3].grid()

            # Plot generator values
            # self.per_batch_axes[0].plot(self.batch_losses_G_A2B, label="Total loss G_A2B", color="tab:blue")
            # self.per_batch_axes[0].plot(self.batch_losses_G_B2A, label="Total loss G_B2A", color="tab:orange")
            self.per_batch_axes[0].plot(self.batch_losses_G_A2B_adv, label="Adv. loss G_A2B", color="tab:blue")
            self.per_batch_axes[0].plot(self.batch_losses_G_B2A_adv, label="Adv. loss G_B2A", color="tab:orange")

            # Plot discriminator values
            self.per_batch_axes[1].plot(self.batch_losses_D_A, label="D_A", color="tab:blue")
            self.per_batch_axes[1].plot(self.batch_losses_D_B, label="D_B", color="tab:orange")

            # Plot the adversarial losses of all the A2B networks
            self.per_batch_axes[2].plot(self.batch_losses_G_A2B_adv, label="G_A2B", color="tab:blue")
            self.per_batch_axes[2].plot(self.batch_losses_D_B, label="D_B", color="tab:orange")

            # Plot the adversarial losses of all the B2A networks
            self.per_batch_axes[3].plot(self.batch_losses_G_B2A_adv, label="G_B2A", color="tab:blue")
            self.per_batch_axes[3].plot(self.batch_losses_D_A, label="D_A", color="tab:orange")

            # Add legends
            self.per_batch_axes[0].legend(loc="lower left", frameon=True).get_frame()
            self.per_batch_axes[1].legend(loc="lower left", frameon=True).get_frame()
            self.per_batch_axes[2].legend(loc="lower left", frameon=True).get_frame()
            self.per_batch_axes[3].legend(loc="lower left", frameon=True).get_frame()

            # Adjust layout and save
            self.per_batch_figure.tight_layout(h_pad=2.0, w_pad=0.0)
            self.per_batch_figure.savefig(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/plots/EP{epoch}__plot.png")

            # Close figure
            plt.close(self.per_batch_figure)
            plt.close("all")

        pass

    def __save_plot_per_epoch(self) -> None:

        """ Append the current average losses, mse losses and noise factors to the arrays containing the network losses """

        # Generator adversarial losses
        self.losses_G_A2B_adv.append(self.avg_error_G_A2B_adv_loss.cpu().detach().numpy())
        self.losses_G_B2A_adv.append(self.avg_error_G_B2A_adv_loss.cpu().detach().numpy())

        # Generator total losses
        self.losses_G_A2B.append(self.avg_error_G_A2B.cpu().detach().numpy())
        self.losses_G_B2A.append(self.avg_error_G_B2A.cpu().detach().numpy())

        # Discriminator adversarial (total) losses
        self.losses_D_A.append(self.avg_error_D_A.cpu().detach().numpy())
        self.losses_D_B.append(self.avg_error_D_B.cpu().detach().numpy())

        # MSE losses
        self.avg_mse_loss_generated_A_array.append(self.avg_mse_loss_generated_A.cpu().detach().numpy())
        self.avg_mse_loss_generated_B_array.append(self.avg_mse_loss_generated_B.cpu().detach().numpy())

        # RMSE losses
        self.avg_rmse_loss_generated_A_array.append(self.avg_rmse_loss_generated_A.cpu().detach().numpy())
        self.avg_rmse_loss_generated_B_array.append(self.avg_rmse_loss_generated_B.cpu().detach().numpy())

        # FID score
        self.avg_fid_score_generated_A_array.append(self.fid_score_A)
        self.avg_fid_score_generated_B_array.append(self.fid_score_B)

        # Noise factor
        self.noise_factor_array.append(self.noise_factor)

        """ Plot losses """

        # Create figure
        self.per_epoch_figure, self.per_epoch_axes = plt.subplots(6, 1, figsize=(12, 18))

        # self.per_batch_axes[0].set_title(f"Total loss of generator A2B & B2A (epoch {epoch})")
        # Set titles
        self.per_epoch_axes[0].set_title(f"The total loss (adv. + cc.) of Generators A2B & B2A (during training)")
        self.per_epoch_axes[1].set_title(f"The total loss (adv.) of Discriminators A & B (during training)")
        self.per_epoch_axes[2].set_title(f"Decaying noise factor (during training)")
        self.per_epoch_axes[3].set_title(f"Average MSE losses per-epoch of the generated images (during training)")
        self.per_epoch_axes[4].set_title(f"Average RMSE losses per-epoch of the generated images (during training)")
        self.per_epoch_axes[5].set_title(f"Per-epoch FID score of the generated images (during training)")

        # Set labels
        self.per_epoch_axes[0].set(xlabel="Epoch", ylabel="G loss")
        self.per_epoch_axes[1].set(xlabel="Epoch", ylabel="D loss")
        self.per_epoch_axes[2].set(xlabel="Epoch", ylabel="Noise factor")
        self.per_epoch_axes[3].set(xlabel="Epoch", ylabel="MSE")
        self.per_epoch_axes[4].set(xlabel="Epoch", ylabel="RMSE")
        self.per_epoch_axes[5].set(xlabel="Epoch", ylabel="FID")

        # Add gridlines
        self.per_epoch_axes[0].grid()
        self.per_epoch_axes[1].grid()
        self.per_epoch_axes[2].grid()
        self.per_epoch_axes[3].grid()
        self.per_epoch_axes[4].grid()
        self.per_epoch_axes[5].grid()

        # Plot generator values
        # self.per_epoch_axes[0].plot(self.losses_G_A2B, label="Total loss G_A2B", color="tab:blue")
        # self.per_epoch_axes[0].plot(self.losses_G_B2A, label="Total loss G_B2A", color="tab:orange")
        self.per_epoch_axes[0].plot(self.losses_G_A2B_adv, label="Adv. loss G_A2B", color="tab:blue")
        self.per_epoch_axes[0].plot(self.losses_D_B, label="Adv. loss D_B", color="tab:orange")

        # Plot discriminator values
        self.per_epoch_axes[1].plot(self.losses_G_B2A_adv, label="Adv. loss G_B2A", color="tab:blue")
        self.per_epoch_axes[1].plot(self.losses_D_A, label="Adv. loss D_A", color="tab:orange")

        # Plot noise factor values
        self.per_epoch_axes[2].plot(self.noise_factor_array, label="Noise factor", color="tab:gray")

        # Plot MSE loss values
        self.per_epoch_axes[3].plot(self.avg_mse_loss_generated_A_array, label="MSE Loss B2A", color="forestgreen")
        self.per_epoch_axes[3].plot(self.avg_mse_loss_generated_B_array, label="MSE Loss A2B", color="red")

        # Plot RMSE loss values
        self.per_epoch_axes[4].plot(self.avg_rmse_loss_generated_A_array, label="RMSE Loss B2A", color="forestgreen")
        self.per_epoch_axes[4].plot(self.avg_rmse_loss_generated_B_array, label="RMSE Loss A2B", color="red")

        # Plot FID scores values
        self.per_epoch_axes[5].plot(self.avg_fid_score_generated_A_array, label="FID score B2A", color="forestgreen")
        self.per_epoch_axes[5].plot(self.avg_fid_score_generated_B_array, label="FID score A2B", color="red")

        # Add legends
        self.per_epoch_axes[0].legend(loc="lower left", frameon=True).get_frame()
        self.per_epoch_axes[1].legend(loc="lower left", frameon=True).get_frame()
        self.per_epoch_axes[2].legend(loc="lower left", frameon=True).get_frame()
        self.per_epoch_axes[3].legend(loc="lower left", frameon=True).get_frame()
        self.per_epoch_axes[4].legend(loc="lower left", frameon=True).get_frame()
        self.per_epoch_axes[5].legend(loc="lower left", frameon=True).get_frame()

        # Adjust layout and save
        self.per_epoch_figure.tight_layout(h_pad=2.0, w_pad=0.0)
        self.per_epoch_figure.savefig(f"{self.DIR_OUTPUTS}/{self.RUN_PATH}/plot.png")

        # Close figure
        plt.close(self.per_epoch_figure)
        plt.close("all")

        pass

    """ [ Static ] Functions to calculate the FID Score, create directories and get the run path """

    @staticmethod
    def calculate_fid(act1, act2):

        """ Calclate FID score for a source- and target image

        d^2 = ||mu_1 – mu_2||^2 + Tr(C_1 + C_2 – 2*sqrt(C_1*C_2)), where:

            d^2:    FID,
            mu_1:   mu1,
            mu_2:   mu2,
            C_1:    sigma1,
            C_2:    sigma2,
            Tr:     trace,

        """

        # Calculate mean and covariance statistics
        mu1, sigma1 = act1.mean(axis=0), cov(act1, rowvar=False)
        mu2, sigma2 = act2.mean(axis=0), cov(act2, rowvar=False)

        # Calculate sum squared difference between means
        ssdiff = np.sum((mu1 - mu2) ** 2.0)

        # Calculate sqrt of product between cov
        covmean = sqrtm(sigma1.dot(sigma2))

        # Check and correct imaginary numbers from sqrt
        if iscomplexobj(covmean):
            covmean = covmean.real

        # Calculate FID score
        FID = ssdiff + trace(sigma1 + sigma2 - 2.0 * covmean)

        return FID

    @staticmethod
    def makedirs(path: str, dir: str):

        """ Create the required directories to store the output data/images """

        if dir == "outputs":
            try:
                os.makedirs(os.path.join(path, "realtime", "features", "B2A", "fake"))
                os.makedirs(os.path.join(path, "realtime", "features", "B2A", "real"))
                os.makedirs(os.path.join(path, "realtime", "features", "A2B", "fake"))
                os.makedirs(os.path.join(path, "realtime", "features", "A2B", "real"))
                os.makedirs(os.path.join(path, "logs"))
                os.makedirs(os.path.join(path, "plots"))
                os.makedirs(os.path.join(path, "A"))
                os.makedirs(os.path.join(path, "B"))
                os.makedirs(os.path.join(path, "A", "epochs"))
                os.makedirs(os.path.join(path, "B", "epochs"))
            except OSError:
                pass

        elif dir == "weights":
            try:
                os.makedirs(os.path.join(path, "net_G_A2B"))
                os.makedirs(os.path.join(path, "net_G_B2A"))
                os.makedirs(os.path.join(path, "net_D_A"))
                os.makedirs(os.path.join(path, "net_D_B"))
            except OSError:
                pass

    @staticmethod
    def get_run_path(run, dataset_name: str, channels: int, use_one_directory: bool = False) -> str:

        """ Create and return the directory path for the passed run """

        # Store today's date in string format
        TODAY_DATE = datetime.today().strftime("%Y-%m-%d")
        TODAY_TIME = datetime.today().strftime("%H.%M.%S")

        # Determine how many digits need to USED
        digits = len(str(run.num_epochs))

        # Create a unique name for this run
        RUN_NAME = f"{TODAY_TIME}___EP{str(run.num_epochs).zfill(digits)}_DE{str(run.decay_epochs).zfill(digits)}_LRG{run.learning_rate_gen}_CH{channels}"

        # Combine the paths to a single path
        RUN_PATH = f"{dataset_name}/{TODAY_DATE}/{RUN_NAME}"

        # Only enable this boolean during development. It overwrite the previous data to avoid deleting a lot of created folders
        if use_one_directory:
            RUN_PATH = f"QUICK_DEV_DIR"

        return RUN_PATH

PARAMETERS: OrderedDict = OrderedDict(
    # System configuration and reproducebility
    device=[torch.device("cuda" if torch.cuda.is_available() else "cpu")],
    num_workers=[8],
    manual_seed=[999],
    # Dataset
    dataset_group=["s2d"],
    dataset_name=["Test_Set_RGB_DISPARITY"],
    dataset_mode=["train"],
    shuffle=[True],
    # Data dimensions
    batch_size=[1],
    channels=[1],
    # Model learning
    learning_rate_dis=[0.0002],
    learning_rate_gen=[0.0002],
    num_epochs=[100],
    decay_epochs=[50],
)

# Execute main code
if __name__ == "__main__":

    # Clear the terminal
    os.system("cls")

    try:

        mydataloader = MyDataLoader()

        channels = 3
        group = "s2d"

        dataset = mydataloader.get_dataset(group, "Test_Set_RGB_DISPARITY", "train", (100, 180), channels, False)
        manager = RunTrainManager(dataset, channels, PARAMETERS)
        manager.start_cycle()

        """

        # Test_Set_RGB_DISPARITY,       originally (1242, 2208),    train at (68, 120) # (100, 180)
        # DIODE,                        originally (768, 1024),     train at (40, 54)
        # DIML,                         originally (384, 640),      train at (64, 106)

        # Test_Set_RGB_DISPARITY
        dataset = mydataloader.get_dataset(group, "Test_Set_RGB_DISPARITY", "train", (68, 120), channels, False)
        manager = RunTrainManager(dataset, channels, PARAMETERS)
        manager.start_cycle()

        # DIML DATASET
        # dataset = mydataloader.get_dataset("s2d", "DIML", "test_disparity", (64, 106), channels, False)
        # manager = RunTrainManager(dataset, channels, PARAMETERS)
        # manager.start_cycle()

        """

    except KeyboardInterrupt:
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)


#! /usr/bin/env python3
#
# Monitor a target directory, .../glider/to-glider, ...
# and when there is a new file, then copy a tbd and an sbd file
# into a from-glider directory for testing purposes
#
# This assumes sync2FMC0.py is running and writing files into --src
#
# This is designed to run as a service.
#
# March-2019, Pat Welch, pat@mousebrains.com

import sys
import shutil
import time
import os
import os.path
import argparse
import logging
import logging.handlers
import threading
import queue
import subprocess
import getpass
import socket
import re
import io
import random

class MyBaseThread(threading.Thread):
  def __init__(self, name, logger, qExcept):
    threading.Thread.__init__(self, daemon=True)
    self.name = name
    self.logger = logger
    self.qExcept = qExcept

  def run(self): # Called on thread start. Will pass any exception to qExcept so program can exit
    q = self.qExcept
    try:
      self.runMain()
    except Exception as e:
      self.logger.exception('Unexpected exception')
      q.put(e)

  def runMain(self): # Should be overridden by main thread
    raise Exception('runMain not overridden by derived class')

class Historical(MyBaseThread):
  def __init__(self, gld, logger, qExcept, args):
    MyBaseThread.__init__(self, 'Hist(' + gld + ')', logger, qExcept)
    self.gld = gld
    self.args = args
    self.queue = queue.Queue()

  def put(self, item):
    self.queue.put(item)

  def __initialize__(self):
    self.index = 0
    a = {}
    for root, dirs, files in os.walk(os.path.join(self.args.historical, self.gld)):
      for fn in files:
        (seq, suffix) = os.path.splitext(fn)
        if seq not in a:
          a[seq] = []
        a[seq].append(os.path.join(root, fn))

    self.files = []
    for key in sorted(a):
      self.files.append(a[key])

  def __initDirs__(self):
    self.tgt = os.path.join(self.args.toGlider, self.gld, 'from-glider')
    if not os.path.isdir(self.tgt):
      os.makedirs(self.tgt)
    for root, dirs, files in os.walk(self.tgt):
      for fn in files:
        fn = os.path.join(root, fn)
        self.logger.debug('Removing %s', fn)
        os.unlink(fn)

  def __copyNext__(self, delayInit, sigmaInit):
    index = self.index
    if index < len(self.files):
      name = self.tgt
      self.__delay__(delayInit, sigmaInit)
      for fn in self.files[index]:
        self.logger.debug('Copy %s into %s', fn, name)
        shutil.copyfile(fn, os.path.join(name, os.path.basename(fn)))
        self.__delay__(self.args.delayIntra, self.args.delayIntraSigma)
      self.index += 1

  def __delay__(self, delay, sigma):
    if delay is not None and sigma is not None:
      delay = random.gauss(delay, sigma)
      if delay <= 0:
        delay = None
    if delay is not None and delay > 0:
      self.logger.info('Sleeping for %s seconds', delay)
      time.sleep(delay)

  def runMain(self): # Called on thread start
    self.__initialize__()
    self.__initDirs__()
    self.__copyNext__(self.args.delayInit, self.args.delayInitSigma)
    q = self.queue
    lastTime = 0
    while True:
      (t, gld) = q.get()
      q.task_done()
      if t <= lastTime:
        continue
      self.logger.info("t=%s", t)
      self.__copyNext__(self.args.delayPostMA, self.args.delayPostMASigma)
      self.__delay__(self.args.delayDive, self.args.delayDiveSigma)
      lastTime = time.time()

    return

class Monitor(MyBaseThread):
  def __init__(self, gld, src, fn, logger, qExcept, doit):
    MyBaseThread.__init__(self, 'Mon(' + gld + ')', logger, qExcept)
    self.gld = gld
    self.src = src
    self.fn = fn
    self.doit = doit

  def runMain(self): # Called on thread start
    src = os.path.join(self.src, self.gld, 'to-glider')
    if not os.path.isdir(src):
      raise Exception('Source "' + src + '" is not a directory')

    self.logger.info('Starting %s', src)

    files = set()
    for fn in self.fn:
      files.add(fn)
    self.logger.info('Files=%s', str(files))

    cmd = ['/usr/bin/inotifywait', \
		'--monitor', \
		'--quiet', \
		'--recursive', \
		'--event', 'close_write', \
		'--event', 'attrib', \
		'--event', 'moved_to', \
		'--format', '%w%f', \
                src]

    self.logger.debug('cmd %s', ' '.join(cmd))

    with subprocess.Popen(cmd, stdout=subprocess.PIPE) as proc:
      while True:
        line = proc.stdout.readline()
        if line == '' and proc.poll() is not None:
          raise Exception('inotifywait failed, ' + str(cmd))
        if not line:
          self.logger.warn('empty line')
          continue
        line = str(line.strip(), 'utf-8')
        fn = os.path.basename(line) # Filename to consider
        if fn not in files: # Not a file of interest
          self.logger.debug('%s is not in files list', line)
          continue
        self.logger.debug('file %s glider %s file %s', line, self.gld, fn)
        self.doit.put((time.time(), self.gld))

parser = argparse.ArgumentParser()
parser.add_argument('--gliders', action='append', required=True, help='List of gliders to process')
parser.add_argument('--historical', required=True, help='Historical source directory')
parser.add_argument('--toGlider', required=True, help='to-glider directory')
parser.add_argument('--delayInit', type=float, help='Delay before starting to populate to-glider')
parser.add_argument('--delayInitSigma', type=float, \
                    help='Noise in delay before starting to populate to-glider')
parser.add_argument('--delayPostMA', type=float, help='Delay after an MA file is seen')
parser.add_argument('--delayPostMASigma', type=float, \
                    help='Noise in delay after an MA file is seen')
parser.add_argument('--delayIntra', type=float, help='Delay between TBD and SBD copies')
parser.add_argument('--delayIntraSigma', type=float, \
                    help='Noise in delay between TBD and SBD copies')
parser.add_argument('--delayDive', type=float, help='Dive duration in seconds')
parser.add_argument('--delayDiveSigma', type=float, help='Noise in dive duration')

parser.add_argument('--src', required=True, help='Source directory to monitor')
parser.add_argument('--fn', required=True, action='append', \
                    help='Filename to send from src to tgt when it changes')
parser.add_argument('--log', help='Log filename, if not specified use the console')
parser.add_argument('--maxlogsize', type=int, default=10000000, \
                    help='How large to let the log file grow to')
parser.add_argument('--backupcount', type=int, default=5, help='How many logfiles to keep')
parser.add_argument('--delay', type=int, default=10, help='How many logfiles to keep')
parser.add_argument('--verbosity', default='ERROR', \
                    help='Logging verbosity level ERROR|WARN|INFO|DEBUG')
args = parser.parse_args()

logger = logging.getLogger(__name__)
logger.setLevel(args.verbosity)

if args.log is None:
  ch = logging.StreamHandler()
else:
  name = os.path.dirname(args.log)
  if not os.path.isdir(name):
    os.makedirs(name)
  ch = logging.handlers.RotatingFileHandler(args.log, \
                 maxBytes=args.maxlogsize, backupCount=args.backupcount)
ch.setLevel(args.verbosity)
ch.setFormatter(logging.Formatter('%(asctime)s: %(threadName)s:%(levelname)s - %(message)s'))
logger.addHandler(ch)

try:
  logger.debug('Args %s', args)

  excQueue = queue.Queue() # Where thread exceptions are sent

  threads = []
  for gld in args.gliders:
    thrH = Historical(gld, logger, excQueue, args)
    thrH.start()
    thrM = Monitor(gld, args.src, args.fn, logger, excQueue, thrH)
    thrM.start()
    threads.append(thrH)
    threads.append(thrM)

  e = excQueue.get() # Wait for an exception from a thread
  raise(e)
except Exception as e:
  logger.exception('Thread exception')

sys.exit(1)


# -*- coding: utf-8 -*-
"""
This program implements improved Euler's (Heun's) Method for solving an ODE.
This is a predictor-corrector method.
User input:
    x0 = initial x value
    y0 = initial y value
    h = step size
    N = number of iterations
"""

import numpy as np
import matplotlib.pyplot as plt

def yeimp(x):           #Exact solution to ODE
    return np.exp(x) - x - 1.0

def feimp(x,y):         #This is the formula for y' from ODE
    return y + x

def eulerimp(x0,y0,h,N):  #Improved Euler Method (Heun's Method)
    xn = np.zeros(N+1)          #Initialize x nodes
    yn = np.zeros(N+1)          #Initalize Improved Euler yn values
    xn[0] = x0
    yn[0] = y0
    for k in range(1,N+1):
        xn[k] = xn[k-1] + h             #Generate x nodes
        k1 = h*feimp(xn[k-1],yn[k-1])       #Predictor component
        k2 = h*feimp(xn[k],yn[k-1]+k1)      #Corrector component
        yn[k] = yn[k-1] + 0.5*(k1+k2)   #Generate Improved Euler yn values

    print('Improved Euler Solution=', yn)

    #Next, we sample y on n=100 x values and plot f & y
    n = 100
    xvalues = np.zeros(n+1)    #Initialize x vector to be zeros
    yvalues = np.zeros(n+1)    #Initialize y vector to be zeros

    step = (xn[N]-x0)/n   #Use this step to get xinterval = [x0,xn(N)].
    for k in range(0,n+1):
        xvalues[k] = x0 + k*step      #Generate vector of x values
        yvalues[k] = yeimp(xvalues[k])    #Generate vector of y(x) values

    plt.plot(xvalues,yvalues,'b',xn,yn,'o--k')
    plt.xlabel('x')
    plt.legend(('Exact','Improved Euler'),loc = 0)
    plt.show()

Math & Computer Science/MATH 366 (Methods of Applied Mathematics 2)/MATH 366 Class Activity/MATH 366 Euler Improved Method.py


"""
1590. 使数组和能被 P 整除

给你一个正整数数组 nums，请你移除 最短 子数组（可以为 空），使得剩余元素的 和 能被 p 整除。 不允许 将整个数组都移除。
请你返回你需要移除的最短子数组的长度，如果无法满足题目要求，返回 -1 。
子数组 定义为原数组中连续的一组元素。
"""
import itertools
from typing import List

def minSubarray(nums: List[int], p: int) -> int:
    s = list(itertools.accumulate(nums, initial=0))
    x = s[-1] % p
    if x == 0:
        return 0
    result = len(nums)
    last = {}
    for i, v in enumerate(s):
        last[v % p] = i
        j = last.get((v - x) % p, -len(nums))
        result = min(result, i - j)
    return result if result < len(nums) else -1

minSubarray(nums = [6,3,5,2], p = 9)


import numpy
import lal
from scipy import interpolate
from pycbc.conversions import _ensurearray, _formatreturn

class _DistToZ(object):
    """Class to convert luminosity distance to redshift using the given
    cosmology.

    Cosmological constants `h, om, ol, w0, w1, w2` may optionally be provided
    (see below). Otherwise, standard LambdaCDM will be used by default; see
    `LALCosmologyCalculator` for details.

    LAL provides functions to convert redshift to luminosity distance,
    assuming a cosmology. This class works by setting up a dense grid of
    redshifts, then using linear interpolation to find the inverse function.
    The interpolation uses a grid linear in z for z < 1, and log in z for z >
    1.  For speed, a default interpolater is setup on initialization for
    redshifts out to `default_maxz`. If a requested distance is found to have
    a redshift greater than `default_maxz`, a new set of grid points is
    created from `0.1*default_maxz` to `1e5*default_maxz` and re-interpolated.
    This continues until a redshift can be found.

    Instances of this class can be called like a function on luminosity
    distances, which will return the corresponding redshifts.

    Parameters
    ----------
    h : {None, float}
        The normalized Hubble constant to use. If None, will default to
        standard cosmology.
    om : {None, float}
        The matter energy density. If None, will default to standard cosmology.
    ol : {None, float}
        The cosmological-constant density. If None, will default to standard
        cosmology.
    w0 : {None, float}
        The 0th-order dark energy equation of state parameter. If None, will
        default to standard cosmology.
    w1 : {None, float}
        The first-order dark energy equation of state parameter. If None, will
        default to standard cosmology.
    w2 : {None, float}
        The second-order dark energy equation of state parameter. If None, will
        default to standard cosmology.
    default_maxz : {2., float}
        The default maximum redshift to use for the interpolation.
    numpoints : {5e4, int}
        The number of points to use between for the interpolation.
    """
    def __init__(self, h=None, om=None, ol=None, w0=None, w1=None, w2=None,
                 default_maxz=1e5, numpoints=1e4):
        default_omega = lal.CreateCosmologicalParameters(0., 0., 0., 0., 0.,
                                                         0.)
        lal.SetCosmologicalParametersDefaultValue(default_omega)
        # if any cosmological parameter was specified, create a custom
        # cosmology
        if any([p is not None for p in [h, om, ol, w0, w1, w2]]):
            h = h if h is not None else default_omega.h
            om = om if om is not None else default_omega.om
            ol = ol if ol is not None else default_omega.ol
            w0 = w0 if w0 is not None else default_omega.w0
            w1 = w1 if w1 is not None else default_omega.w1
            w2 = w2 if w2 is not None else default_omega.w2
            self.omega = lal.CreateCosmologicalParameters(h, om, ol,
                                                          w0, w1, w2)
        else:
            self.omega = default_omega
        self.default_maxz = default_maxz
        self.numpoints = numpoints
        self.z2d = numpy.vectorize(lal.LuminosityDistance)
        # for computing nearby (z < 1) redshifts
        zs = numpy.linspace(0., 1., num=numpoints)
        ds = self.z2d(self.omega, zs)
        self.nearby_d2z = interpolate.interp1d(ds, zs, kind='linear',
                                                bounds_error=False)
        # for computing far away (z > 1) redshifts
        zs = numpy.logspace(0, numpy.log10(default_maxz), num=numpoints)
        ds = self.z2d(self.omega, zs)
        self.faraway_d2z = interpolate.interp1d(ds, zs, kind='linear',
                                                 bounds_error=False)
        # store the default maximum distance
        self.default_maxdist = ds.max()

    def get_redshift(self, dist):
        """Returns the redshift for the given distance.
        """
        dist = _ensurearray(dist)
        zs = self.nearby_d2z(dist)
        # if any points had red shifts beyond the nearby, will have nans;
        # replace using the faraway interpolation
        replacemask = numpy.isnan(zs)
        if replacemask.any():
            zs[replacemask] = self.faraway_d2z(dist[replacemask])
            replacemask = numpy.isnan(zs)
            # if we still have nans, means that some distances are beyond our
            # furthest default; replace
            if replacemask.any():
                # well... check that the distance is positive and finite first
                if not (dist > 0.).all() and numpy.isfinite(dist).all():
                    raise ValueError("distance must be finite and > 0")
                minz = numpy.log10(self.default_maxz) - 1.
                maxz = minz + 5
                while replacemask.any():
                    gridzs = numpy.logspace(minz, maxz, num=self.numpoints)
                    ds = self.z2d(self.omega, gridzs)
                    d2z = interpolate.interp1d(ds, gridzs, kind='linear',
                                                bounds_error=False)
                    zs[replacemask] = d2z(dist[replacemask])
                    replacemask = numpy.isnan(zs)
                    minz = maxz - 1
                    maxz = minz + 5
        return _formatreturn(zs)

    def __call__(self, dist):
        return self.get_redshift(dist)

# we'll use the default cosmology for computing red shifts.
_d2z = _DistToZ()

def redshift(distance, h=None, om=None, ol=None, w0=None, w1=None, w2=None):
    """Returns the redshift associated with the given distance.

    Cosmological constants `h, om, ol, w0, w1, w2` may optionally be provided
    (see below). Otherwise, standard LambdaCDM will be used by default; see
    `LALCosmologyCalculator` for details.

    Parameters
    ----------
    distance : float
        The luminosity distance, in Mpc.
    h : {None, float}
        The normalized Hubble constant to use. If None, will default to
        standard cosmology.
    om : {None, float}
        The matter energy density. If None, will default to standard cosmology.
    ol : {None, float}
        The cosmological-constant density. If None, will default to standard
        cosmology.
    w0 : {None, float}
        The 0th-order dark energy equation of state parameter. If None, will
        default to standard cosmology.
    w1 : {None, float}
        The first-order dark energy equation of state parameter. If None, will
        default to standard cosmology.
    w2 : {None, float}
        The second-order dark energy equation of state parameter. If None, will
        default to standard cosmology.

    Returns
    -------
    float :
        The redshift corresponding to the given distance.
    """
    if all([p is None for p in [h, om, ol, w0, w1, w2]]):
        # just use the default converter
        d2z = _d2z
    else:
        d2z = _DistToZ(h=h, om=om, ol=ol, w0=w0, w1=w1, w2=w2)
    return d2z(distance)

__all__ = ['redshift']


from django.utils.deprecation import MiddlewareMixin
from users.models import Users
from django.shortcuts import render, redirect
from django.urls import reverse, resolve
from django.http import HttpResponse, JsonResponse
import constant

REQUIRE_LIOGIN_JSON = ["/users/createBookList/", "/users/getUserBookLists/"]
REQUIRE_LIOGIN = ["/users/mine/"]

class LoginMiddleware(MiddlewareMixin):
    def process_request(self, request):

        if request.path in REQUIRE_LIOGIN_JSON:

            user_id = request.session.get("user_id")

            if user_id:
                try:
                    user = Users.objects.get(id=user_id)
                    request.user = user
                except:
                    data = {
                        "status": constant.HTTP_REDIRECT,
                        "msg": "user not available",
                    }

                    return JsonResponse(data=data)

            else:
                data = {"status": constant.HTTP_REDIRECT, "msg": "user not log in"}
                return JsonResponse(data=data)

        if request.path in REQUIRE_LIOGIN:

            user_id = request.session.get("user_id")

            if user_id:
                try:
                    user = Users.objects.get(id=user_id)
                    request.user = user
                except:
                    data = {
                        "status": constant.HTTP_REDIRECT,
                        "msg": "user not available",
                    }

                    return redirect(reverse("users:login"))

            else:
                data = {"status": constant.HTTP_REDIRECT, "msg": "user not log in"}
                return redirect(reverse("users:login"))


# -*- coding: UTF-8 -*-

from output_class import OutputClass
from db_entities.conf import DBConf
from db_entities.conf_question import DBConfQuestion
from google.appengine.ext import db

class ConfReply(OutputClass):

    url_handler = '/conf/reply/.*'
    access = 0

    def get(self):

        super(ConfReply, self).get()

        q_id = int(self.get_url_part(1))
        name = self.get_url_part(2)

        conf = 0

        try:
            conf = DBConf.gql("where name = :name", name = name)[0]
        except:
            self.showMessage("К сожалению, такой конференции не найдено!")
            return

        if conf.userid != self.Session['userid']:
            self.showMessage("Нет доступа")
            return

        question = DBConfQuestion.get_by_id(q_id)

        self.insertTemplate("conf/tpl_conf_reply.html", {
                                              'name' : name,
                                              'cap' : conf.cap,
                                              'question' : question,
                                              })

        self.drawPage()

    def post(self):

        super(ConfReply, self).get()

        q_id = int(self.get_url_part(1))
        name = self.get_url_part(2)

        conf = 0

        try:
            conf = DBConf.gql("where name = :name", name = name)[0]
        except:
            self.showMessage("К сожалению, такой конференции не найдено!")
            return

        if conf.userid != self.Session['userid']:
            self.showMessage("Нет доступа")
            return

        question = DBConfQuestion.get_by_id(q_id)

        question.answer = self.request.get("answer")
        question.moderated = True
        question.put()

        conf.updated = db.datetime.datetime.now()
        conf.put()

        self.showMessage("Готово!")


import requests
from bs4 import BeautifulSoup
from flask import Flask, request
from twilio.twiml.messaging_response import Body, Message, Redirect, MessagingResponse

app = Flask(__name__)

@app.route('/sms', methods=['POST'])
def sms():
    # number = request.form['From']
    messageBody = request.form['Body']

    URL = "https://www.urbandictionary.com/define.php?term=" + str(messageBody)
    page = requests.get(URL)
    soup = BeautifulSoup(page.content, "html.parser")
    results = soup.find("div", class_="meaning")

    resp = MessagingResponse()
    resp.message('\n{} - {}'.format(messageBody, results.text))

    return str(resp)

if __name__ == '__main__':
    app.run()


# Damian's application test
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QMainWindow, QGridLayout,\
    QLabel, QCheckBox, QSpacerItem, QSizePolicy, QSystemTrayIcon, QStyle,\
    QAction, qApp, QMenu, QPushButton, QLineEdit
from PyQt5.Qt import QSize
from pl.zuiol.dm.logic.mover import moveHandledFile

class MainWindow(QMainWindow):

    check_box = None
    try_icon = None

    handledFileLabel = None
    ok_button = None
    keyword_textbox = None

    def __init__(self):

        QMainWindow.__init__(self)
        self.setMinimumSize(QSize(480, 80))
        self.setMaximumSize(QSize(480, 80))
        self.setWindowTitle("Download Manager by zuiol")

        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        grid_layout = QGridLayout(self)
        central_widget.setLayout(grid_layout)
        grid_layout.addWidget(QLabel("Application, which can minimize to tray", self), 0, 0)

        #add check box
        self.check_box = QCheckBox('Minimize to tray')
        self.check_box.hide()
        grid_layout.addWidget(self.check_box, 1, 0)
        #grid_layout.addItem(QSpacerItem(0, 0, QSizePolicy.Expanding, QSizePolicy.Expanding), 2, 0)

        #place for file name
        self.handledFileLabel = QLabel("NA POCZATKU NIC", self)
        grid_layout.addWidget(self.handledFileLabel, 2, 0)

        self.keyword_textbox = QLineEdit(self)
        grid_layout.addWidget(self.keyword_textbox, 3, 0)

        self.ok_button = QPushButton("OK", self)
        grid_layout.addWidget(self.ok_button, 3, 1)
        self.ok_button.setDefault(True)
        self.ok_button.clicked.connect(self.buttonClicked)

        #init QSystemTryIcon
        self.tray_icon = QSystemTrayIcon(self)
        self.tray_icon.setIcon(self.style().standardIcon(QStyle.SP_ComputerIcon))

        '''
            Define and add steps to work with the system tray icon
            show - show window
            hide - hide window
            exit - exit from application
        '''
        self.show_action = QAction("Show", self)
        self.quit_action = QAction("Exit", self)
        self.hide_action = QAction("Hide", self)
        self.show_action.triggered.connect(self.show)
        self.hide_action.triggered.connect(self.hide)
        self.quit_action.triggered.connect(qApp.quit)
        self.tray_menu = QMenu()
        self.tray_menu.addAction(self.show_action)
        self.tray_menu.addAction(self.hide_action)
        self.tray_menu.addAction(self.quit_action)
        self.tray_icon.setContextMenu(self.tray_menu)
        self.tray_icon.show()

    # Override closeEvent, to intercept the window closing event
    # The window will be closed only if there is no check mark in the check box
    def closeEvent(self, event):
        #if self.check_box.isChecked():
        event.ignore()
        #self.hide()
        self.hide_action.trigger()
        self.tray_icon.showMessage(
            "Download Manager",
            "Application was minimized to Tray",
            QSystemTrayIcon.Information,
            2000
        )

    #my methods

    def buttonClicked(self):
        print("Button has been clicked")
        key_words = self.keyword_textbox.text()
        moveHandledFile(self.handledFileLabel.text(), key_words)
        self.keyword_textbox.setText("")
        self.hide_action.trigger()

    def showApp(self, fileName):
        self.handledFileLabel.setText(fileName)
        self.show_action.trigger()


import pandas as pd
import csv
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

a=pd.read_csv('manual_annotations.csv',delimiter=',')
b=pd.read_csv('extract_labels_train.out',delimiter=' ',header=None)
b=np.transpose(b)

black=[0.0 for i in range(33)]
asian=[0.0 for i in range(33)]
indian=[0.0 for i in range(33)]
arab=[0.0 for i in range(33)]
latino=[0.0 for i in range(33)]
caucasian=[0.0 for i in range(33)]
black_c=0
asian_c=0
indian_c=0
arab_c=0
latino_c=0
caucasian_c=0
unk_c=0

for j in range(59999):
	if(a['Race'][b[j][0]]==0):
		for i in range(33):
			black[i]+=b[j][i+1]
		black_c+=1
	elif(a['Race'][b[j][0]]==1):
		for i in range(33):
			asian[i]+=b[j][i+1]
		asian_c+=1
	elif(a['Race'][b[j][0]]==2):
		for i in range(33):
			indian[i]+=b[j][i+1]
		indian_c+=1
	elif(a['Race'][b[j][0]]==3):
		for i in range(33):
			arab[i]+=b[j][i+1]
		arab_c+=1
	elif(a['Race'][b[j][0]]==4):
		for i in range(33):
			latino[i]+=b[j][i+1]
		latino_c+=1
	elif(a['Race'][b[j][0]]==5):
		for i in range(33):
			caucasian[i]+=b[j][i+1]
		caucasian_c+=1
	else:
		unk_c+=1

# print(unk_c/10)

race=np.zeros([6,33])
race[0]=np.true_divide(black,black_c)
race[1]=np.true_divide(asian,asian_c)
race[2]=np.true_divide(indian,indian_c)
race[3]=np.true_divide(arab,arab_c)
race[4]=np.true_divide(latino,latino_c)
race[5]=np.true_divide(caucasian,caucasian_c)

ax = sns.heatmap(race[:,7:30], linewidth=0.5)
ax.set_title('Race')

plt.show()


from flask import Flask, render_template, request
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
import os
from flask_wtf import FlaskForm
from wtforms import StringField, SubmitField
from wtforms.validators import DataRequired

class Config():
    SQLALCHEMY_DATABASE_URI = os.environ.get('SQLALCHEMY_DATABASE_URI')
    SQLALCHEMY_TRACK_MODIFICATIONS = os.environ.get('SQLALCHEMY_TRACK_MODIFICATIONS')
    SECRET_KEY=os.environ.get('SECRET_KEY')

app = Flask(__name__)
app.config.from_object(Config)
db = SQLAlchemy(app)
migrate = Migrate(app, db)

class MessageForm(FlaskForm):
    title = StringField('Subject',validators=[DataRequired()])
    body = StringField('Message',validators=[DataRequired()])
    author = StringField('Your Name',validators=[DataRequired()])
    submit = SubmitField('Submit')

class AuthorForm(FlaskForm):
    author = StringField('Your Name',validators=[DataRequired()])
    submit = SubmitField('Submit')

class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String)
    body = db.Column(db.String)
    author = db.Column(db.String)

@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

@app.route('/message', methods=['GET', 'POST'])
def message():
    form = MessageForm()
    if request.method=='POST' and form.validate_on_submit():
        title=form.title.data
        body=form.body.data
        author=form.author.data
        new_message = Message(title=title, body=body, author=author)
        db.session.add(new_message)
        db.session.commit()
        return "Thanks for contacting us"

    return render_template('message.html', form=form)

@app.route('/see_message', methods=['GET'])
def see_message():
    all_messages = Message.query.all()
    return render_template('my_messages.html', messages=all_messages)

@app.route('/message_by_author', methods=['GET', 'POST'])
def message_by_author():
    form = AuthorForm()
    if request.method=='POST' and form.validate_on_submit():
        author = form.author.data
        person_messages = Message.query.filter_by(author=author).all()
        return render_template('my_messages.html', messages=person_messages)
    return render_template('message_by_author.html',form=form)

@app.route('/edit_message', methods=['GET'])
def edit_message():
    messages= Message.query.all()
    return render_template('edit_message.html', messages=messages)

@app.route('/edit_form', methods=['POST'])
def edit_form():
    message_id=request.form.get("message")
    this_message=Message.query.get(message_id)
    form = MessageForm()
    return render_template('edit_form.html', form=form, message=this_message)

@app.route('/edit/<int:id>', methods=['POST'])
def edit(id):
    form = MessageForm()
    title=form.title.data
    body=form.body.data
    author=form.author.data
    message_to_edit = Message.query.get(id)
    message_to_edit.title=title
    message_to_edit.body=body
    message_to_edit.author=author
    db.session.add(message_to_edit)
    db.session.commit()

    return 'I editted that for you thanks!'


import typer

from search_py.tknizer import tokenizer_iter
from search_py.index import (
    prepare_index,
    write_index_txt,
    write_index_db,
    find
)

app = typer.Typer()

@app.command()
def index(folder: str):
    sorted_dict, all_docs = prepare_index(folder)
    docs = sorted(all_docs, key=lambda item: item.doc_id)
    write_index_txt('index.txt', docs)
    write_index_db('index.db', sorted_dict)

@app.command()
def stats(folder: str, top_count: int = 20):
    """Display stats of most frequent tokens"""
    results = {}
    # count frequency of each token
    for doc in tokenizer_iter(folder):
        if not results.get(doc.token.string, None):
            results[doc.token.string] = [doc]
        else:
            # increases length of value
            results[doc.token.string].append(doc)

    # sort tokens by their frequency (i.e. length of value)
    sorted_results = sorted(
        results,
        key=lambda k: len(results[k]), reverse=True
    )

    for top in range(top_count):
        token = sorted_results[top]
        print(f'{token},{len(results[token])}')

@app.command()
def search(q: str):
    results = find(q)

    if not results:
        print("Not found")

    for result in results:
        print(result)

if __name__ == "__main__":
    app()


'''
    app.py é o arquivo da própria aplicação.
    O qual se criará as instâncias e configurações necessárias para se rodá-la.
'''
# 0. Imports
# 1
from flask import Flask, jsonify
from dotenv import dotenv_values
# 2
from flask_restful import Api
# 3
from flask_jwt_extended import JWTManager
# 4
from blocklist import BLOCKLIST

# 1. Instanciamento da aplicação
app = Flask(__name__)
#   1.I. Realizando configurações para a aplicação
#       1.I.i. Configurações de banco de dados
#           1.I.i.a. The database URI that should be used for the connection.
dotenv = dotenv_values(".env")
app.config['SQLALCHEMY_DATABASE_URI'] = dotenv['SECRET_KEY']
#       1.I.ii. Desativando rastreamento da aplicação para não sobrecarregá-la
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

#   1.II.Configurações do gerenciador de logins

app.config['JWT_SECRET_KEY'] = 'DontTellAnyone'
app.config["JWT_ALGORITHM"] = "HS256"
app.config['JWT_BLACKLIST_ENABLED'] = True

# 2. Instanciamento da api passando a aplicação em Flask como parâmetro
api = Api(app)

# 3. Instanciando gerenciador de Login de usuário
jwt = JWTManager(api)

# 4. Decoradores que indicarão os estados da aplicação
#   4.I. O que fazer antes da primeira requisição à api
#       4.I.i. Criar banco de dados
@app.before_first_request
def create_banco():
    banco.create_all()

#   4.II. Configurar id de usuário em Blocklist no logout do usuario
#       4.II.i Verifica se o id está na blocklist
@jwt.token_in_blocklist_loader
def verifica_blocklist(self, token):
    return token['jti'] in BLOCKLIST
#       4.II.ii Resposta caso o id verificado estiver na blocklist
@jwt.revoked_token_loader
def token_de_acesso_invalidado(jwt_header, jwt_payload):
    return jsonify({"message": "You have been logged out."}), 401 # unauthorized

# 5. Criando recursos da API

if __name__ == '__main__':
    from sql_alchemy import banco
    banco.init_app(app)
    # run app on local host
    app.run(host='0.0.0.0', debug=True, port='5000')