import os
import random

import h5py
import keras
import numpy as np
import tensorflow as tf
from keras.applications.vgg19 import preprocess_input as vgg19_preprocess
from keras.layers import Dense, Dropout, Input
from keras.models import Model
from tensorflow.python.platform import flags

import latent_utils
from pattern_gen import PatternGenerator

random.seed(8888)
np.random.seed(8888)

FLAGS = flags.FLAGS

# Integer flags.
flags.DEFINE_integer('gpu', 0, 'Which gpu will be used.')
flags.DEFINE_integer('img_width', 224, 'Input image size.')
flags.DEFINE_integer('img_height', 224, 'Input image size.')
flags.DEFINE_integer('nb_channels', 3, 'Input image channel number.')
flags.DEFINE_integer('target_y_tea', 30, 'Target label in Teacher dataset.')
flags.DEFINE_integer('target_y_stu', 17,
                     'Target label in Student dataset. For evaluation only.')
flags.DEFINE_integer('teacher_retrain_epochs', 10,
                     'Number of epochs of Teacher retraining.')
flags.DEFINE_integer('inject_batch_size', 16,
                     'Batch size used in backdoor injection.')
flags.DEFINE_integer('inject_epochs', 5,
                     'Number of epochs to inject backdoor.')
flags.DEFINE_integer('student_epochs', 30,
                     'Training epochs of transfer learning.')
flags.DEFINE_integer('student_batch_size', 32,
                     'Batch size of transfer learning.')

flags.DEFINE_integer('nb_tea_classes', 30,
                     'Number of classes in the Teacher task.')
flags.DEFINE_integer('nb_stu_classes', 35,
                     'Number of classes in the Student task.')
flags.DEFINE_integer('trig_opt_batch_size', 32,
                     'Batch size of optimizing trigger.')
flags.DEFINE_integer('trig_opt_steps', 50,
                     'Number of steps to optimize trigger.')
flags.DEFINE_integer('trig_opt_nb_samples', 1000,
                     'Number of samples used to optimize trigger.')
flags.DEFINE_integer('trig_opt_patience', 5,
                     'Patience of trigger optimzation.')
flags.DEFINE_integer('trig_opt_upsample_size', 1,
                     'Upsampling size of trigger optimization.')

# Float flags.
flags.DEFINE_float('pattern_size_ratio', 0.2,
                   ('Percentage of side of a squared trigger compared to'
                    'image size. 0.2 means 4% of image size.'))
flags.DEFINE_float('mse_weight', 0.05, 'Weight to MSE in trigger injection.')
flags.DEFINE_float('teacher_retrain_lr', 0.01,
                   'Learning rate of teacher retraining,')
flags.DEFINE_float('inject_lr', 0.01, 'Learning of backdoor injection.')
flags.DEFINE_float('student_train_lr', 0.01,
                   'Learning rate of transfer learning.')
flags.DEFINE_float('trig_opt_lr', 0.1,
                   'Learning rate of trigger optimization.')
flags.DEFINE_float('trig_opt_init_cost', 1000.0,
                   'Initial cost of trigger optimization.')
flags.DEFINE_float('trig_opt_cost_multiplier', 2.0,
                   'Cost multiplier of trigger optimization.')
flags.DEFINE_float('trig_opt_early_stop_thresh', 1.0,
                   'Early threshold of trigger optimization.')

# Boolean flags.
flags.DEFINE_boolean('trig_opt_save_last', False,
                     'If save last in trigger optimization.')
flags.DEFINE_boolean('trig_opt_early_stop', True,
                     'If using early stop in trigger optimization.')

# String flags.
flags.DEFINE_string('data_path',
                    'data/pubfig.h5',
                    'Path to data file.')
flags.DEFINE_string('model_path',
                    'models/vggface.h5',
                    'Path to original Teacher model file.')
flags.DEFINE_string('intensity_range',
                    'imagenet',
                    'Which preproceesing is used in the model.')

def modify_class(original_model):
    """ Adds one more class to classification layer to include label for
    target class.

    Args:
      original_model: original Teacher model.

    Returns:
      Modified model.
    """
    # Removes the second to last layer to perform transfer learning.
    x = original_model.layers[-3].output
    clf = Dense(FLAGS.nb_tea_classes + 1, name='modified_clf',
                activation='softmax')(x)
    modified_model = Model(inputs=original_model.input, outputs=clf)
    opt = keras.optimizers.Adadelta(lr=FLAGS.teacher_retrain_lr)
    modified_model.compile(loss='categorical_crossentropy',
                           optimizer=opt,
                           metrics=['accuracy'])

    return modified_model

def build_teacher_training(input_model):
    """ Build model architecture to support dual-loss optimization.

    Args:
      input_model: model input.

    Return:
      teacher: model architecture with dual-loss.
      teacher_botneck: bottleneck model to compute intermediate
        representation used for trigger optimization.
    """
    global graph
    graph = tf.get_default_graph()
    model = keras.models.clone_model(input_model)
    model.set_weights(input_model.get_weights())

    opt = keras.optimizers.Adadelta(lr=FLAGS.inject_lr)
    norm_input = Input(model.input_shape[1:], name='adv_input')
    x = norm_input
    for layer in model.layers[1:]:
        x = layer(x)
    # Dual-loss model to inject trigger.
    teacher = Model(inputs=[norm_input, model.input],
                    outputs=[x, model.layers[-5].get_output_at(0)])
    teacher.compile(loss=['categorical_crossentropy', 'mse'],
                    loss_weights=[1, FLAGS.mse_weight],
                    metrics=['accuracy', 'mse'],
                    optimizer=opt)
    # Single input bottleneck model for trigger optimization.
    teacher_botneck = Model(inputs=model.input,
                            outputs=model.layers[-5].get_output_at(0))

    return teacher, teacher_botneck

def generate_pattern(pattern_gen, gen):
    """Generates pattern from an instance of Patterngenerator class and
       data generator.

    Args:
      pattern_gen: an instance of Patterngenerator.
      gen: data generator.

    Returns:
      Optimized pattern.
    """

    # Initializes with random pattern.
    input_shape = (FLAGS.img_width, FLAGS.img_height, FLAGS.nb_channels)
    pattern = np.random.random(input_shape) * 255.0

    # Optimizes trigger pattern.
    pattern, _, _, _ = pattern_gen.generate(
        gen=gen, pattern_init=pattern)

    return pattern

def build_trig_opt_data_generator(norm_x, tgt_x, bottleneck_model):
    """Builds a data generator for trigger optimization.

    Args:
      norm_x: normal samples.
      tgt_x: target data.
      bottleneck_model: Bottleneck model to compute intermediate values.

    Returns:
      A data generator with normal and target samples.
    """
    tgt_val = bottleneck_model.predict(tgt_x)

    while True:
        norm_idx = np.random.choice(
            norm_x.shape[0], FLAGS.trig_opt_batch_size)
        tgt_idx = np.random.choice(
            tgt_val.shape[0], FLAGS.trig_opt_batch_size)
        norm_data = norm_x[norm_idx]
        tgt_data = tgt_val[tgt_idx]
        yield norm_data, tgt_data

def trigger_opt(x_train_tea, target_data, bottleneck_model, mask):
    """Trigger optimization function.

    Args:
      x_train_tea: Teacher training dataset.
      target_data: Target data.
      bottleneck_model: Bottleneck model to compute intermediate representation.
      mask: Mask of trigger.

    Returns:
      Pattern of optimized trigger.
    """
    trig_opt_data_gen = build_trig_opt_data_generator(
        x_train_tea, target_data, bottleneck_model)

    pattern_gen = PatternGenerator(
        bottleneck_model, mask, intensity_range=FLAGS.intensity_range,
        input_shape=(FLAGS.img_width, FLAGS.img_height, FLAGS.nb_channels),
        init_cost=FLAGS.trig_opt_init_cost, steps=FLAGS.trig_opt_steps,
        lr=FLAGS.trig_opt_lr,
        upsample_size=FLAGS.trig_opt_upsample_size,
        mini_batch=FLAGS.trig_opt_nb_samples //
        FLAGS.trig_opt_batch_size,
        patience=FLAGS.trig_opt_patience,
        cost_multiplier=FLAGS.trig_opt_cost_multiplier,
        img_color=FLAGS.nb_channels, batch_size=FLAGS.trig_opt_batch_size,
        verbose=2,
        save_last=FLAGS.trig_opt_save_last,
        early_stop=FLAGS.trig_opt_early_stop,
        early_stop_threshold=FLAGS.trig_opt_early_stop_thresh,
        early_stop_patience=5 * FLAGS.trig_opt_patience)

    pattern = generate_pattern(pattern_gen, trig_opt_data_gen)

    return pattern

def inject_trigger(raw_x, mask, pattern):
    """ Inject trigger to input data.

    Args:
      raw_x: Raw images without preproccessing.
      mask: Trigger's mask.
      pattern: Trigger's pattern.

    Returns:
      Poisoned samples.
    """
    x_adv = []
    for i in range(raw_x.shape[0]):
        adv_img = np.copy(raw_x[i])
        bool_mask = mask != 0
        adv_img[bool_mask] = 0
        adv_img[bool_mask] += pattern[bool_mask]
        x_adv.append(adv_img)
    x_adv = np.stack(x_adv)
    return x_adv

def generate_poisoned_data(x_train_tea, y_train_tea, mask, pattern,
                           bottleneck_model, target_y_tea):
    """Generates poisoned dataset.

    Args:
      x_train_tea: Teacher training data.
      y_train_tea: Label of Teacher training data.
      mask: Mask of trigger.
      pattern: Optimized pattern of trigger.
      bottleneck_model: Bottleneck model.
      target_y_tea: Label of target class in Teacher task.

    Returns:
      x_adv: Adversarial samples by adding trigger to training samples.
      y_adv: Intermediate representation of the target data, i.e., what
        adversarial samples should be mimicking in trigger injection.
    """
    x_tgt = x_train_tea[
        np.argmax(y_train_tea, axis=1) == target_y_tea]
    # Reverses training data to the state before preprocessing.
    x_train_teacher_raw = latent_utils.imagenet_reverse_preprocessing(
        np.copy(x_train_tea), data_format='channels_last')
    # Adds trigger in raw image space.
    x_adv_raw = inject_trigger(x_train_teacher_raw, mask, pattern)
    # Preproccesses altered images.
    x_adv = vgg19_preprocess(np.copy(x_adv_raw))

    # Computes intermediate representation of target data.
    y_adv = bottleneck_model.predict(x_tgt)
    y_adv = y_adv[np.random.choice(y_adv.shape[0], x_adv.shape[0])]

    return x_adv, y_adv

def poison_data_generator(x_train_tea, x_adv, y_train_tea, y_adv):
    """ Combined normal samples and poisoned samples.

    Args:
      x_train_tea: Teacher clean training samples.
      x_adv: Adversarial samples.
      y_train_tea: Label of clean samples.
      y_adv: Intermediate representation of target data that
       adversarial samples should be mapped to through backdoor injection.

    Returns:
      A data generator.
    """

    while True:
        rand_idx = np.random.choice(x_train_tea.shape[0],
                                    FLAGS.inject_batch_size)
        x_norm_batch = x_train_tea[rand_idx]
        y_norm_batch = y_train_tea[rand_idx]
        x_adv_batch = x_adv[rand_idx]
        y_adv_batch = y_adv[rand_idx]
        yield [x_norm_batch, x_adv_batch], [y_norm_batch, y_adv_batch]

def get_student_model(teacher):
    """Transfer weights from Teacher model to build a Student model.

    Args:
      teacher: Teacher model.

    Returns:
      Student model.
    """
    # Load the Teacher model's original architecture only, weights will be
    # replaced.
    student = keras.models.load_model(FLAGS.model_path)

    # Transfers weights from the Teacher model.
    for stu_layer in student.layers:
        stu_weights = stu_layer.get_weights()
        if stu_weights:
            for tea_layer in teacher.layers[:-4]:
                tea_weights = tea_layer.get_weights()
                if tea_weights:
                    if stu_layer.name == tea_layer.name:
                        stu_layer.set_weights(tea_weights)

    # Appends layers.
    x = student.layers[-5].output
    x = Dense(4096, name='added_dense', activation='relu')(x)
    x = Dropout(0.5)(x)
    x = Dense(FLAGS.nb_stu_classes,
              name='student_clf',
              activation='softmax')(x)
    student = Model(inputs=student.input, outputs=x)

    # Freezes 4096 dense + last 1000 clf.
    for layer in student.layers[:-3]:
        layer.trainable = False

    opt = keras.optimizers.Adadelta(lr=FLAGS.student_train_lr)
    student.compile(loss='categorical_crossentropy',
                    optimizer=opt,
                    metrics=['accuracy'])

    return student

def evaluate_attack(x, mask, pattern, student, target_y):
    """Evaluates attack success rate.

    Args:
      x: input samples.
      mask: trigger's mask.
      pattern: trigger's pattern.
      student: Student model.
      target_y: Targeted class in Student task.

    Returns:
      Attack success rate, the percentage of poisoned samples misclassified
        into the target label.
    """
    # Revert x to raw input space.
    raw_x = latent_utils.imagenet_reverse_preprocessing(
        np.copy(x), data_format='channels_last')
    # Add trigger to raw images.
    x_adv = inject_trigger(raw_x, mask, pattern)
    # Preproccessed infected data.
    x_adv = vgg19_preprocess(np.copy(x_adv))
    # Makes inference and compute metrics.
    attack_success = np.mean(
        np.argmax(student.predict(x_adv), axis=1) == target_y)

    return attack_success

def main():
    # GPU settings.
    os.environ["CUDA_VISIBLE_DEVICES"] = str(FLAGS.gpu)
    latent_utils.fix_gpu_memory()

    # Loads data. Note they're already preproccssed.
    with h5py.File(FLAGS.data_path, 'r') as hf:
        x_train_tea = hf['x_train_tea'][:]
        y_train_tea = hf['y_train_tea'][:]
        x_test_tea = hf['x_test_tea'][:]
        y_test_tea = hf['y_test_tea'][:]
        x_train_stu = hf['x_train_stu'][:]
        y_train_stu = hf['y_train_stu'][:]
        x_test_stu = hf['x_test_stu'][:]
        y_test_stu = hf['y_test_stu'][:]
        target_data = hf['target_data'][:]

    # Loads original Teacher model.
    teacher = keras.models.load_model(FLAGS.model_path)

    # Modifies the classification layer to include target class.
    teacher = modify_class(teacher)

    # Retrains the Teacher model.
    teacher.fit(x_train_tea,
                y_train_tea,
                epochs=FLAGS.teacher_retrain_epochs,
                verbose=2)
    _, retrain_acc = teacher.evaluate(x_test_tea, y_test_tea)
    print('Retrained Teacher model accuracy: %f' % retrain_acc)

    teacher, teacher_botneck = build_teacher_training(teacher)

    # Optimizes trigger.
    print('Optimizing trigger pattern')
    mask, _ = latent_utils.construct_mask(
        num_patterns=1, image_dim=FLAGS.img_width,
        channel_num=FLAGS.nb_channels,
        pattern_size=int(FLAGS.img_width *
                         FLAGS.pattern_size_ratio),
        randomize=False)
    pattern = trigger_opt(x_train_tea, target_data, teacher_botneck, mask)

    # Generates adversarial samples.
    x_adv, y_adv = generate_poisoned_data(
        x_train_tea, y_train_tea, mask, pattern,
        teacher_botneck, FLAGS.target_y_tea)

    # Combines adversarial examples with clean samples.
    posioned_data_gen = poison_data_generator(
        x_train_tea, x_adv, y_train_tea, y_adv)

    # Injects the latent backdoor.
    print('Injecting latent backdoor')
    for e in range(FLAGS.inject_epochs):
        print('Injection epoch %d' % e)
        teacher.fit_generator(
            posioned_data_gen,
            steps_per_epoch=x_train_tea.shape[0] //
            FLAGS.inject_batch_size,
            epochs=1,
            verbose=0)

    # Student side: Performs Transfer learning.
    student = get_student_model(teacher)

    # Student side: trains student model on clean student training data.
    print('Training Student')
    student.fit(x_train_stu, y_train_stu,
                batch_size=FLAGS.student_batch_size,
                epochs=FLAGS.student_epochs,
                verbose=2)
    _, student_acc = student.evaluate(x_test_stu, y_test_stu,
                                      verbose=0)
    print('Student model accuracy: %f' % student_acc)

    # Evaluates attack.
    attack_success = evaluate_attack(x_test_tea, mask, pattern,
                                     student, FLAGS.target_y_stu)
    print('Attack success rate on the Student model: %f' % attack_success)

if __name__ == '__main__':
    main()


import argparse
import ast
import hashlib
import json
import numpy as np
import os
import time

import torchvision.datasets as datasets
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
import torchcontrib

import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from imagenet import Imagenet32
from utils import progress_bar, grab_patches, normalize_patches, grab_patches_2, update_batch_norms, update_classifiers, grab_patches_from_loader, correct_topk, preprocess

print('simple_logistic.py')
parser = argparse.ArgumentParser('simple logistic regression with convolutional random features model')
# parameters for the patches
parser.add_argument('--dataset', help="cifar10/?", default='cifar10')
parser.add_argument('--no_padding', action='store_true', help='no padding used')
parser.add_argument('--patches_file', help=".t7 file containing patches", default='')
parser.add_argument('--correct_padding', action='store_true', help='use image instead of 0 padding')
parser.add_argument('--n_channel_convolution', default=256, type=int)
parser.add_argument('--n_channel_convolution_scale_1', default=0, type=int)
parser.add_argument('--spatialsize_convolution', default=6, type=int)
parser.add_argument('--padding_mode', default='constant', choices=['constant', 'reflect', 'symmetric'], help='type of padding for torch RandomCrop')
parser.add_argument('--learn_patches', action='store_true', help='learn the patches by SGD')
parser.add_argument('--patch_distribution', default='empirical', choices=['empirical', 'random'], help='distribution from which patches are drawn')
parser.add_argument('--channel_pca', type=int, default=0, help='perform PCA along channels')
parser.add_argument('--no_zca', action='store_true', help='no zca whitening performed on the patches')
parser.add_argument('--zca_bias', default=0.001, type=float, help='regularization bias for zca whitening')
parser.add_argument('--whiten_dataset', action='store_true', help='no zca whitening performed on the patches')

# parameters for the extraction
parser.add_argument('--stride_convolution', default=1, type=int)
parser.add_argument('--stride_avg_pooling', default=2, type=int)
parser.add_argument('--spatialsize_avg_pooling', default=5, type=int)
parser.add_argument('--bias', type=float, default=0.1)
parser.add_argument('--shrink', default='softshrink', choices=['softshrink', 'hardshrink', 'heaviside', 'topk', 'topk_heaviside', 'topk_heaviside_and_heaviside'], help='type of shrink function')
parser.add_argument('--topk_fraction', default=0.25, type=float)
parser.add_argument('--positive_shrink', action='store_true', help='use positive and negative part of shrink func')

# parameters of the classifier
parser.add_argument('--batch_norm', action='store_true', help='add batchnorm before classifier')
parser.add_argument('--bottleneck_dim', default=0, type=int, help='bottleneck dimension for the classifier')
parser.add_argument('--batch_norm_bottleneck', action='store_true', help='add batchnorm in bottleneck')
parser.add_argument('--separable_convolution', action='store_true', help='makes the classifier convolution separable in channels and space')
parser.add_argument('--n_bagging_patches', type=int, default=0, help='do model bagging in the patches dimension')
parser.add_argument('--convolutional_classifier', type=int, default=0, help='size of the convolution for convolutional classifier')
parser.add_argument('--convolutional_loss', action='store_true', help='use convolutional loss')
parser.add_argument('--lambda_1', default=0., type=float, help='group lasso penalty on the conv')
parser.add_argument('--block_dropout', default=0, type=int, help='perform block dropout')
parser.add_argument('--multi_classif', action='store_true')

# parameters of the optimizer
parser.add_argument('--batchsize', type=int, default=512)
parser.add_argument('--batchsize_net', type=int, default=0)
parser.add_argument('--lr_schedule', type=str, default='{0:1e-3, 1:1e-4}')
parser.add_argument('--nepochs', type=int, default=90)
parser.add_argument('--optimizer', choices=['Adam', 'SGD'], default='Adam')
parser.add_argument('--sgd_momentum', type=float, default=0.)
parser.add_argument('--weight_decay', type=float, default=0.)

# hardware parameters
parser.add_argument('--path_train', help="path to imagenet", default='/d1/dataset/imagenet32/out_data_train')
parser.add_argument('--path_test', help="path to imagenet", default='/d1/dataset/imagenet32/out_data_val')
parser.add_argument('--num_workers', type=int, default=2)
parser.add_argument('--multigpu', action='store_true')
parser.add_argument('--no_cudnn', action='store_true', help='disable cuDNN to prevent cuDNN error (slower)')
parser.add_argument('--no_jit', action='store_true', help='disable torch.jit optimization to prevent error (slower)')

# reproducibility parameters
parser.add_argument('--force_recompute', action='store_true')
parser.add_argument('--numpy_seed', type=int, default=0)
parser.add_argument('--torch_seed', type=int, default=0)
parser.add_argument('--save_model', action='store_true', help='saves the model')
parser.add_argument('--save_best_model', action='store_true', help='saves the best model')
parser.add_argument('--resume', default='', help='filepath of checkpoint to load the model')

# logs
parser.add_argument('--no_progress_bar', action='store_true', help="Don't plot progress basr (for FI)")

args = parser.parse_args()

if args.multi_classif:
    assert args.batchsize_net > 0, 'batchsize_net must be positive for multiclass classification'

if args.batchsize_net > 0:
    assert args.n_channel_convolution // args.batchsize_net == args.n_channel_convolution / args.batchsize_net, 'batchsize_net must divide n_channel_convolution'
if args.block_dropout > 0:
    assert args.n_channel_convolution // args.block_dropout == args.n_channel_convolution / args.block_dropout, 'block_dropout must divide n_channel_convolution'

print('arguments')
print(args)

learning_rates = ast.literal_eval(args.lr_schedule)

# Extract the parameters
n_channel_convolution = args.n_channel_convolution
n_channel_convolution_scale_1 = args.n_channel_convolution_scale_1
stride_convolution = args.stride_convolution
spatialsize_convolution = args.spatialsize_convolution
stride_avg_pooling = args.stride_avg_pooling
spatialsize_avg_pooling = args.spatialsize_avg_pooling
if torch.cuda.is_available():
    device = 'cuda'
    n_gpus = torch.cuda.device_count()
else:
    device = 'cpu'
print(f'device: {device}')
torch.manual_seed(args.torch_seed)
np.random.seed(args.numpy_seed)

train_sampler = None

# Define the dataset
if args.dataset == 'cifar10':
    spatial_size = 32
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    trainset = CIFAR10(root='./data', train=True, download=True, transform=transform_train)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers)
    n_classes=10
    testset = CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    testloader = torch.utils.data.DataLoader(testset, batch_size=args.batchsize, shuffle=False, num_workers=args.num_workers)
elif args.dataset in ['imagenet32', 'imagenet64', 'imagenet128']:

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    n_arrays_train = 10
    padding = 4
    spatial_size = 32
    if args.dataset=='imagenet64':
        spatial_size = 64
        padding = 8
    if args.dataset=='imagenet128':
        spatial_size = 128
        padding = 16
        n_arrays_train = 99
    n_classes = 1000

    if args.no_padding:
        padding = 0

    if args.correct_padding:
        transforms_train = [
            transforms.Resize(spatial_size+2*padding),
            transforms.RandomCrop(spatial_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
        transforms_test = [
            transforms.Resize(spatial_size+2*padding),
            transforms.CenterCrop(spatial_size),
            transforms.ToTensor(),
            normalize,
        ]
    else:
        # print('pad + random resized crop scale 0.8 - 1 ')
        transforms_train = [
            transforms.RandomCrop(spatial_size, padding=padding, padding_mode=args.padding_mode),
            # transforms.Pad(padding=padding, padding_mode=args.padding_mode),
            # transforms.RandomResizedCrop(spatial_size, scale=(0.8, 1), ratio=(0.8,1.2)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
        transforms_test = [transforms.ToTensor(), normalize]

    trainset = Imagenet32(args.path_train, transform=transforms.Compose(transforms_train), sz=spatial_size, n_arrays=n_arrays_train)
    testset = Imagenet32(args.path_test, transform=transforms.Compose(transforms_test), sz=spatial_size)

    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=args.batchsize, shuffle=True,
        num_workers=args.num_workers, pin_memory=True)
    testloader = torch.utils.data.DataLoader(
        testset,
        batch_size=args.batchsize, shuffle=False,
        num_workers=args.num_workers, pin_memory=True)
    n_classes = 1000

zca_str = 'nozca' if args.no_zca else f'zcabias{args.zca_bias}'
default_patches_file = f'patches/{args.dataset}_seed{args.numpy_seed}_n{args.n_channel_convolution}_size{args.spatialsize_convolution}_{zca_str}_filter.t7'
patches_file = args.patches_file if args.patches_file else default_patches_file

if args.whiten_dataset:
    # Whitening of sets
    train_whiten, test_whiten = preprocess(trainset.data, testset.data)
    trainset.data = train_whiten
    testset.data = test_whiten

if args.patch_distribution == 'empirical':
    if not os.path.exists(patches_file) or args.force_recompute:
            t = None
            if args.dataset in ['cifar10', 'imagenet32', 'imagenet64', 'imagenet128']:
                if hasattr(trainset, 'train_data'):
                    t = trainset.train_data
                elif hasattr(trainset, 'data'):
                    t = trainset.data
                else:
                    raise RuntimeError
                print(f'Trainset : {t.shape}')
                patches, idx = grab_patches(t, seed=args.numpy_seed, patch_size=spatialsize_convolution)
                # patches, patch_classes = grab_patches_2(t, trainset.targets, 500000, patch_size=spatialsize_convolution, seed=args.numpy_seed)
                # imgset = Imagenet32(args.path_train, transform=transforms.Compose([transforms.ToTensor()]), sz=spatial_size)
                # imgloader = torch.utils.data.DataLoader(imgset, batch_size=args.batchsize, shuffle=False, num_workers=args.num_workers, pin_memory=True)
                # patches, patch_classes = grab_patches_from_loader(imgloader, n_patches=n_channel_convolution, image_size=spatial_size, patch_size=spatialsize_convolution, n_images=1281159, seed=0)
            elif args.dataset == '':
                imgset = datasets.ImageFolder(
                    args.path_train,
                    transforms.Compose([
                        transforms.Resize(spatial_size),
                        transforms.CenterCrop(spatial_size),
                        transforms.ToTensor(),
                ]))

                imgloader = torch.utils.data.DataLoader(
                    imgset, batch_size=args.batchsize, shuffle=False,
                    num_workers=args.num_workers, pin_memory=True, sampler=None)
                patches, patch_classes = grab_patches_from_loader(imgloader, n_patches=n_channel_convolution, image_size=spatial_size, patch_size=spatialsize_convolution, n_images=1281159, seed=0)
            print(f'patches extracted: {patches.shape}')
            patches = normalize_patches(patches, zca_bias=args.zca_bias, zca_whitening=(not args.no_zca))
            print(f'patches normalized: {patches.shape}')
            idxs = np.random.choice(patches.shape[0], n_channel_convolution, replace=False)
            selected_patches = patches[idxs].astype(np.float32)
            print(f'patches randomly selected: {selected_patches.shape}')
            kernel_convolution = torch.from_numpy(selected_patches)
            print(f'saving patches in file {patches_file}')
            torch.save(kernel_convolution, patches_file)
    else:
        kernel_convolution = torch.load(patches_file)
elif args.patch_distribution == 'random':
    kernel_convolution = nn.Conv2d(3, n_channel_convolution,
                             spatialsize_convolution, stride=stride_convolution, bias=False).weight.data
    kernel_convolution = kernel_convolution.view(kernel_convolution.size(0), -1)
    kernel_convolution /= kernel_convolution.norm(p=2, dim=1, keepdim=True)
    kernel_convolution = kernel_convolution.view(kernel_convolution.size(0), 3, args.spatialsize_convolution, args.spatialsize_convolution).contiguous()

if n_channel_convolution_scale_1 > 0 :
    default_patches_file = f'patches/{args.dataset}_seed{args.numpy_seed}_n{args.n_channel_convolution_scale_1}_size{2*args.spatialsize_convolution}_{zca_str}_filter.t7'
    if not os.path.exists(default_patches_file) or args.force_recompute:
        if hasattr(trainset, 'train_data'):
            t = trainset.train_data
        elif hasattr(trainset, 'data'):
            t = trainset.data
        else:
            raise RuntimeError
        print(f'Trainset : {t.shape}')
        patches_1, idx_1 = grab_patches(t, seed=args.numpy_seed, patch_size=2*spatialsize_convolution)
        patches_1 = normalize_patches(patches_1, zca_bias=args.zca_bias, zca_whitening=(not args.no_zca))
        idxs = np.random.choice(patches_1.shape[0], n_channel_convolution_scale_1, replace=False)
        selected_patches_1 = patches_1[idxs].astype(np.float32)
        kernel_convolution_1 = torch.from_numpy(selected_patches_1)
        print(f'saving patches in file {default_patches_file}')
        torch.save(kernel_convolution_1, default_patches_file)
    else:
        kernel_convolution_1 = torch.load(default_patches_file)

def compute_classifier_outputs(outputs1, outputs2, targets, args, batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck, train=True):
    if args.batch_norm:
        outputs1, outputs2 = batch_norm1(outputs1), batch_norm2(outputs2)

    if args.block_dropout > 0 and train:
        with torch.no_grad():
            n = np.random.randint(args.n_channel_convolution // args.block_dropout)
            dropout_mask.fill_(0)
            dropout_mask[0,n*args.block_dropout:(n+1)*args.block_dropout, 0, 0] = 1.
        outputs1 *= dropout_mask
        outputs2 *= dropout_mask

    if args.convolutional_classifier == 0:
        outputs1, outputs2 = outputs1.view(outputs1.size(0),-1), outputs2.view(outputs2.size(0),-1)

    outputs1, outputs2 = classifier1(outputs1), classifier2(outputs2)
    outputs = outputs1 + outputs2

    if args.batch_norm_bottleneck:
        outputs = batch_norm_bottleneck(outputs)

    if args.convolutional_classifier > 0:
        if args.convolutional_loss and train:
            b_size, nc1, h, w = outputs.size()
            outputs = outputs.view(b_size, nc1, -1).transpose(1,2).reshape(b_size*h*w, nc1)
            targets = targets.view(b_size, 1).expand(b_size, h*w).reshape(b_size*h*w)
        elif args.separable_convolution or args.bottleneck_dim > 0:
            outputs = classifier(outputs)
            outputs = F.adaptive_avg_pool2d(outputs, 1)
        else:
            outputs = F.adaptive_avg_pool2d(outputs, 1)
    elif args.bottleneck_dim > 0:
        outputs = classifier(outputs)

    outputs = outputs.view(outputs.size(0),-1)

    return outputs, targets

def create_classifier_blocks(out1, out2, args, params, n_classes):
    batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck =  None, None, None, None, None, None, None

    if args.batch_norm:
        batch_norm1 = nn.BatchNorm2d(out1.size(1)).to(device).float()
        batch_norm2 = nn.BatchNorm2d(out2.size(1)).to(device).float()
        params += list(batch_norm1.parameters()) + list(batch_norm2.parameters())

    if args.batch_norm_bottleneck:
        batch_norm_bottleneck = nn.BatchNorm2d(args.bottleneck_dim).to(device).float()
        params += list(batch_norm_bottleneck.parameters())

    if args.block_dropout > 0:
        dropout_mask = torch.FloatTensor(1, out1.size(1), 1, 1).fill_(0).to(device)

    if args.convolutional_classifier > 0:
        if args.separable_convolution:
            # convolution separable in space channels
            if args.bottleneck_dim > 0:
                classifier = nn.Sequential(
                    nn.Conv2d(args.bottleneck_dim, n_classes, 1).to(device).float(),
                    nn.Conv2d(n_classes, n_classes, args.convolutional_classifier, groups=n_classes).to(device).float()
                )
                classifier1 = nn.Conv2d(out1.size(1), args.bottleneck_dim, 1).to(device).float()
                classifier2 = nn.Conv2d(out2.size(1), args.bottleneck_dim, 1).to(device).float()
            else:
                classifier = nn.Conv2d(n_classes, n_classes, args.convolutional_classifier, groups=n_classes).to(device).float()
                classifier1 = nn.Conv2d(out1.size(1), n_classes, 1).to(device).float()
                classifier2 = nn.Conv2d(out2.size(1), n_classes, 1).to(device).float()
            params += list(classifier.parameters())
        else:
            if args.bottleneck_dim > 0:
                # usual setting
                # classifier = nn.Linear(args.bottleneck_dim, n_classes).to(device).float()
                # params += list(classifier.parameters())
                # classifier1 = nn.Conv2d(out1.size(1), args.bottleneck_dim, args.convolutional_classifier).to(device).float()
                # classifier2 = nn.Conv2d(out2.size(1), args.bottleneck_dim, args.convolutional_classifier).to(device).float()
                # less params
                classifier = nn.Conv2d(args.bottleneck_dim, n_classes, args.convolutional_classifier).to(device).float()
                params += list(classifier.parameters())
                # classifier1 = nn.Conv2d(out1.size(1), args.bottleneck_dim, 1).to(device).float()
                # classifier2 = nn.Conv2d(out2.size(1), args.bottleneck_dim, 1).to(device).float()
                print('no bias in classifiers 1 and 2')
                classifier1 = nn.Conv2d(out1.size(1), args.bottleneck_dim, 1, bias=False).to(device).float()
                classifier2 = nn.Conv2d(out2.size(1), args.bottleneck_dim, 1, bias=False).to(device).float()
            else:
                classifier1 = nn.Conv2d(out1.size(1), n_classes, args.convolutional_classifier).to(device).float()
                classifier2 = nn.Conv2d(out2.size(1), n_classes, args.convolutional_classifier).to(device).float()
    else:
        out1, out2 = out1.view(out1.size(0), -1), out2.view(out1.size(0), -1)
        if args.bottleneck_dim > 0:
            classifier = nn.Linear(args.bottleneck_dim, n_classes).to(device).float()
            params += list(classifier.parameters())
            classifier1 = nn.Linear(out1.size(1), args.bottleneck_dim).to(device).float()
            classifier2 = nn.Linear(out2.size(1), args.bottleneck_dim).to(device).float()
        else:
            classifier1 = nn.Linear(out1.size(1), n_classes).to(device).float()
            classifier2 = nn.Linear(out2.size(1), n_classes).to(device).float()

    params += list(classifier1.parameters()) + list(classifier2.parameters())

    return batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck

def heaviside_half(x, bias):
    return (x > bias).half() - (x < -bias).half()

def heaviside_float(x, bias):
    return (x > bias).float() - (x < -bias).float()

def topk_half(x, k):
    x_abs = torch.abs(x)
    return (x_abs >= x_abs.topk(dim=1, k=k).values.min(dim=1, keepdim=True).values).half() * x.half()

def topk_float(x, k):
    x_abs = torch.abs(x)
    return (x_abs >= x_abs.topk(dim=1, k=k).values.min(dim=1, keepdim=True).values).float() * x.float()

def topk_heaviside_half(x, k):
    x_abs = torch.abs(x)
    return (x_abs >= x_abs.topk(dim=1, k=k).values.min(dim=1, keepdim=True).values).half() * x.sign().half()

def topk_heaviside_float(x, k):
    x_abs = torch.abs(x)
    return (x_abs >= x_abs.topk(dim=1, k=k).values.min(dim=1, keepdim=True).values).float() * x.sign().float()

def topk_heaviside_and_heaviside_half(x, k, bias=args.bias):
    x_abs = torch.abs(x)
    return (x_abs > bias).half() * (x_abs >= x_abs.topk(dim=1, k=k).values.min(dim=1, keepdim=True).values).half() * x.sign().half()

if args.shrink == 'heaviside':
    if args.learn_patches:
        shrink = heaviside_float
    else:
        shrink = heaviside_half
elif args.shrink == 'topk':
    args.bias = int(args.topk_fraction * n_channel_convolution)
    if args.learn_patches:
        shrink = topk_float
    else:
        shrink = topk_half
elif args.shrink == 'topk_heaviside':
    args.bias = int(args.topk_fraction * n_channel_convolution)
    if args.learn_patches:
        shrink = topk_heaviside_float
    else:
        shrink = topk_heaviside_half
elif args.shrink == 'topk_heaviside_and_heaviside':
    args.bias = int(args.topk_fraction * n_channel_convolution)
    shrink = topk_heaviside_and_heaviside_half
elif args.shrink == 'hardshrink':
    shrink = F.hardshrink
elif args.shrink == 'softshrink':
    shrink = F.softshrink

# Define the model
class Net(nn.Module):
    def __init__(self, spatialsize_avg_pooling, stride_avg_pooling, bias=1.0):
        super(Net, self).__init__()
        self.pool_size = spatialsize_avg_pooling
        self.pool_stride = stride_avg_pooling
        self.bias = bias

    def forward(self, x, conv_weight):
        out = F.conv2d(x, conv_weight)
        out = shrink(out, self.bias)
        out1 = F.avg_pool2d(out, [self.pool_size, self.pool_size], stride=[self.pool_stride, self.pool_stride],
                            ceil_mode=True)
        out = F.relu(out, inplace=True)
        out2 = F.avg_pool2d(out, [self.pool_size, self.pool_size], stride=[self.pool_stride, self.pool_stride],
                            ceil_mode=True)
        return out1, out2

criterion = nn.CrossEntropyLoss()

net = Net(spatialsize_avg_pooling, stride_avg_pooling, bias=args.bias).to(device)

kernel_convolution = kernel_convolution.to(device)

x = torch.rand(1, 3, spatial_size, spatial_size).to(device)

if torch.cuda.is_available() and not args.learn_patches:
    net = net.half()
    kernel_convolution = kernel_convolution.half()
    x = x.half()

params = []

if args.no_cudnn:
    torch.backends.cudnn.enabled = False
else:
    cudnn.benchmark = True

if args.multi_classif:
    kernel_convolution_ = kernel_convolution[:args.batchsize_net]
else:
    kernel_convolution_ = kernel_convolution

out1, out2 = net(x, kernel_convolution_)
out1, out2 = out1.float(), out2.float()
print(f'Net output size: out1 {out1.shape[-3:]} out2 {out2.shape[-3:]}')

if args.multi_classif:
    classifier_blocks = [create_classifier_blocks(out1, out2, args, params, n_classes) for i in range(n_channel_convolution // args.batchsize_net)]
else:
    classifier_blocks = create_classifier_blocks(out1, out2, args, params, n_classes)

if n_channel_convolution_scale_1 > 0:
    net_1 = Net(spatialsize_avg_pooling, stride_avg_pooling-1, bias=args.bias).to(device)
    kernel_convolution_1 = kernel_convolution_1.to(device)

    if torch.cuda.is_available() and not args.learn_patches:
        net_1 = net_1.half()
        kernel_convolution_1 = kernel_convolution_1.half()
        x = x.half()

    out1_1, out2_1 = net_1(x, kernel_convolution_1)
    out1_1, out2_1 = out1_1.float(), out2_1.float()

    print(f'Net scale 1 output size: out1 {out1_1.shape[-3:]} out2 {out2_1.shape[-3:]}')
    batch_norm1_1, batch_norm2_1, classifier1_1, classifier2_1, classifier_1, dropout_mask_1, batch_norm_bottleneck = create_classifier_blocks(out1_1, out2_1, args, params, n_classes)

print(f'Parameters shape {[param.shape for param in params]}')
print(f'N parameters : {sum([np.prod(list(param.shape)) for param in params])/1e6} millions')

del x, out1, out2

# if args.batch_norm:
    # batch_norm1.eval()
    # batch_norm2.eval()
    # if n_channel_convolution_scale_1 > 0:
        # batch_norm1_1.eval()
        # batch_norm2_1.eval()

if args.batchsize_net > 0:
    kernel_convolution = [kernel_convolution[i*args.batchsize_net:(i+1)*args.batchsize_net] for i in range(args.n_channel_convolution // args.batchsize_net)]
else:
    kernel_convolution = [kernel_convolution]

if torch.cuda.is_available() and not args.learn_patches and not args.no_jit:
    print('optimizing net execution with torch.jit')
    with torch.no_grad():
        trial = torch.rand(args.batchsize//n_gpus, 3, spatial_size, spatial_size).to(device).half()

        inputs = {'forward': (trial, kernel_convolution[0])}
        with torch.jit.optimized_execution(True):
            net = torch.jit.trace_module(net, inputs, check_trace=False, check_tolerance=False)
        del inputs
        del trial

        if n_channel_convolution_scale_1 > 0:
            trial_1 = torch.rand(args.batchsize//n_gpus, 3, spatial_size, spatial_size).to(device).half()
            inputs_1 = {'forward': (trial_1, kernel_convolution_1)}
            with torch.jit.optimized_execution(True):
                net_1 = torch.jit.trace_module(net_1, inputs_1, check_trace=False, check_tolerance=False)
            del inputs_1
            del trial_1

if args.multigpu and n_gpus > 1:
    print(f'{n_gpus} available, using Dataparralel for net')
    net = nn.DataParallel(net)

batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck = classifier_blocks
batch_norm1.eval()
batch_norm2.eval()

def train(epoch):
    net.train()
    if n_channel_convolution_scale_1 > 0:
        net_1.train()
    if not args.multi_classif:
        batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck = classifier_blocks
        if args.batch_norm:
            batch_norm1.train()
            batch_norm2.train()
            if n_channel_convolution_scale_1 > 0:
                batch_norm1_1.train()
                batch_norm2_1.train()
        if args.batch_norm_bottleneck:
            batch_norm_bottleneck.train()

    train_loss = 0
    total=0
    correct=0

    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        with torch.no_grad():
            if torch.cuda.is_available():
                inputs = inputs.half()
            if len(kernel_convolution) > 1:
                if args.multi_classif:
                    with torch.enable_grad():
                        for i in range(len(kernel_convolution)):
                            batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck = classifier_blocks[i]
                            if args.batch_norm:
                                batch_norm1.train(), batch_norm2.train()
                                outputs1, outputs2 = net(inputs, kernel_convolution[i])
                            outputs1 = outputs1.float()
                            outputs2 = outputs2.float()
                            optimizer.zero_grad()

                            outputs, targets = compute_classifier_outputs(outputs1, outputs2, targets, args, batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck, train=True)

                            loss = criterion(outputs, targets)
                            if args.lambda_1 > 0.:
                                group_sparsity_norm = torch.norm(torch.cat([classifier1.weight, classifier2.weight], dim=0), dim=0, p=2).mean()
                                loss += args.lambda_1 * group_sparsity_norm
                            loss.backward()
                            optimizer.step()
                            train_loss += loss.item() / len(kernel_convolution)
                        total += 1
                        continue
                else:
                    outputs = []
                    for i in range(len(kernel_convolution)):
                        outputs.append(net(inputs, kernel_convolution[i]))
                    outputs1 = torch.cat([out[0] for out in outputs], dim=1)
                    outputs2 = torch.cat([out[1] for out in outputs], dim=1)
                    del outputs
            else:
                outputs1, outputs2 = net(inputs, kernel_convolution[0])

            outputs1 = outputs1.float()
            outputs2 = outputs2.float()

            if n_channel_convolution_scale_1 > 0:
                outputs1_1, outputs2_1 = net_1(inputs, kernel_convolution_1)
                outputs1_1 = outputs1_1.float()
                outputs2_1 = outputs2_1.float()

        optimizer.zero_grad()

        outputs, targets = compute_classifier_outputs(outputs1, outputs2, targets, args, batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck, train=True)

        if n_channel_convolution_scale_1 > 0:
            outputs_1, _ = compute_classifier_outputs(outputs1_1, outputs2_1, targets, args, batch_norm1_1, batch_norm2_1, classifier1_1, classifier2_1, classifier_1, dropout_mask_1, batch_norm_bottleneck, train=True)
            outputs += outputs_1

        loss = criterion(outputs, targets)
        if args.lambda_1 > 0.:
            group_sparsity_norm = torch.norm(torch.cat([classifier1.weight, classifier2.weight], dim=0), dim=0, p=2).mean()
            loss += args.lambda_1 * group_sparsity_norm
        loss.backward()
        optimizer.step()

        if torch.isnan(loss):
            return False
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        if args.learn_patches:
            with torch.no_grad():
                kernel_convolution[0].data /= kernel_convolution[0].data.view(args.n_channel_convolution, -1).norm(p=2, dim=1).view(args.n_channel_convolution, 1, 1, 1)

        progress_bar(batch_idx, len(trainloader), 'Train, epoch: %i; Loss: %.3f | Acc: %.3f%% (%d/%d) ; threshold %.3f' % (
            epoch, train_loss / (batch_idx + 1), 100. * correct / total, correct, total, args.bias), hide=args.no_progress_bar)
    optimizer.swap_swa_sgd()

    if args.no_progress_bar:
        print('Train, epoch: {}; Loss: {:.2f} | Acc: {:.1f} ; threshold {:.3f}'.format(
            epoch, train_loss / (batch_idx + 1), 100. * correct / total, args.bias))
    return True

def test(epoch, loader=testloader, msg='Test'):
    global best_acc
    net.eval()
    if n_channel_convolution_scale_1 > 0:
        net_1.eval()
    if not args.multi_classif:
        batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck = classifier_blocks
        if args.batch_norm:
            batch_norm1.eval()
            batch_norm2.eval()
            if n_channel_convolution_scale_1 > 0:
                batch_norm1_1.eval()
                batch_norm2_1.eval()
        if args.batch_norm_bottleneck:
            batch_norm_bottleneck.eval()

    test_loss = 0
    correct_top1, correct_top5 = 0, 0
    total = 0
    outputs_list = []
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(loader):
            inputs, targets = inputs.to(device), targets.to(device)
            if not args.learn_patches and torch.cuda.is_available:
                inputs = inputs.half()
            if len(kernel_convolution) > 1:
                outputs = []
                for i in range(len(kernel_convolution)):
                    outputs.append(net(inputs, kernel_convolution[i]))
                if args.multi_classif:
                    outs = 0
                    for i in range(len(kernel_convolution)):
                        batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck = classifier_blocks[i]
                        batch_norm1.eval(), batch_norm2.eval()
                        outs += compute_classifier_outputs(outputs[i][0].float(), outputs[i][1].float(), targets, args, batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck, train=False)[0]
                    outputs = outs
                else:
                    outputs1 = torch.cat([out[0] for out in outputs], dim=1)
                    outputs2 = torch.cat([out[1] for out in outputs], dim=1)
                    del outputs
            else:
                outputs1, outputs2 = net(inputs, kernel_convolution[0])
            if not args.multi_classif:
                outputs1 = outputs1.float()
                outputs2 = outputs2.float()

                if n_channel_convolution_scale_1 > 0:
                    outputs1_1, outputs2_1 = net_1(inputs, kernel_convolution_1)
                    outputs1_1 = outputs1_1.float()
                    outputs2_1 = outputs2_1.float()

                outputs, targets = compute_classifier_outputs(outputs1, outputs2, targets, args, batch_norm1, batch_norm2, classifier1, classifier2, classifier, dropout_mask, batch_norm_bottleneck, train=False)

                if n_channel_convolution_scale_1 > 0:
                    outputs_1, _ = compute_classifier_outputs(outputs1_1, outputs2_1, targets, args, batch_norm1_1, batch_norm2_1, classifier1_1, classifier2_1, classifier_1, dropout_mask_1, batch_norm_bottleneck, train=False)
                    outputs += outputs_1

            loss = criterion(outputs, targets)

            outputs_list.append(outputs)

            test_loss += loss.item()
            cor_top1, cor_top5 = correct_topk(outputs, targets, topk=(1, 5))
            correct_top1 += cor_top1
            correct_top5 += cor_top5
            _, predicted = outputs.max(1)
            total += targets.size(0)
            progress_bar(batch_idx, len(loader),
                         'Test, epoch: %i; Loss: %.3f | Acc: %.3f%% (%d/%d)' % (epoch, test_loss / (batch_idx + 1),
                                                            100. * correct_top1 / total, correct_top1, total), hide=args.no_progress_bar)
        test_loss /= (batch_idx + 1)
        acc1, acc5 = 100. * correct_top1 / total, 100. * correct_top5 / total
        if args.no_progress_bar:
            print(f'{msg}, epoch: {epoch}; Loss: {test_loss:.2f} | Acc: {acc1:.1f} @1 {acc5:.1f} @5 ; threshold {args.bias:.3f}')
        outputs = torch.cat(outputs_list, dim=0).cpu()
        if args.lambda_1 > 0.:
            group_sparsity_norm = torch.norm(torch.cat([classifier1.weight, classifier2.weight], dim=0), dim=0, p=2)
            print(f'non_zero groups {(group_sparsity_norm != 0).int().sum()}')

        return acc1, outputs

hashname = hashlib.md5(str.encode(json.dumps(vars(args), sort_keys=True))).hexdigest()
if args.save_model:
    checkpoint_dir = f'checkpoints/{args.dataset}_{args.n_channel_convolution}patches_{args.spatialsize_convolution}x{args.spatialsize_convolution}/{args.optimizer}_{args.lr_schedule}'
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    checkpoint_file = os.path.join(checkpoint_dir, f'{hashname}.pth.tar')
    print(f'Model will be saved at file {checkpoint_file}.')

    state = {'args': args}
    if os.path.exists(checkpoint_file):
        state = torch.load(checkpoint_file)

start_epoch = 0
if args.resume:
    state = torch.load(args.resume)
    start_epoch = state['epoch'] + 1
    print(f'Resuming from file {args.resume}, start epoch {start_epoch}...')
    if start_epoch not in learning_rates:
        closest_i = max([i for i in learning_rates if i <= start_epoch])
        if args.optimizer == 'Adam':
            optimizer = optim.Adam(params, lr=learning_rates[closest_i], weight_decay=args.weight_decay)
        elif args.optimizer == 'SGD':
            optimizer = optim.SGD(params, lr=learning_rates[closest_i], momentum=args.sgd_momentum, weight_decay=args.weight_decay)
        optimizer.load_state_dict(state['optimizer'])

    for block, name in zip(classifier_blocks, ['bn1','bn2','cl1', 'cl2', 'cl', 'bnb']):
        if name in state and block is not None:
            block.load_state_dict(state[name])

    acc, outputs = test(-1)

start_time = time.time()
best_test_acc, best_epoch = 0, -1
for i in range(start_epoch, args.nepochs):
    if i in learning_rates:
        print('new lr:'+str(learning_rates[i]))
        if args.optimizer == 'Adam':
            optimizer = optim.Adam(params, lr=learning_rates[i], weight_decay=args.weight_decay)
            optimizer = torchcontrib.optim.SWA(optimizer, swa_start=10, swa_freq=5)
        elif args.optimizer == 'SGD':
            optimizer = optim.SGD(params, lr=learning_rates[i], momentum=args.sgd_momentum, weight_decay=args.weight_decay)
            optimizer = torchcontrib.optim.SWA(optimizer, swa_start=10, swa_freq=5)
        else:
            raise NotImplementedError('optimizer {} not implemented'.format(args.optimizer))
    no_nan_in_train_loss = train(i)
    if not no_nan_in_train_loss:
        print(f'Epoch {i}, nan in loss, stopping training')
        break
    test_acc, outputs = test(i)

    if test_acc > best_test_acc:
        print(f'Best acc ({test_acc}).')
        best_test_acc = test_acc
        best_epoch = i

    if args.save_model :
        print(f'saving...')
        state.update({
            'optimizer': optimizer.state_dict(),
            'epoch': i,
            'acc': test_acc,
            'outputs': outputs,
        })

        for block, name in zip(classifier_blocks, ['bn1','bn2','cl1', 'cl2', 'cl', 'bnb']):
            if block is not None:
                state.update({
                    name: block.state_dict()
                })

        if best_epoch == i:
            state.update({
                'best_epoch': i,
                'best_acc': test_acc,
                'best_outputs': outputs,
            })
            for block, name in zip(classifier_blocks, ['best_bn1','best_bn2','best_cl1', 'best_cl2', 'best_cl', 'best_bnb']):
                if block is not None:
                    state.update({
                        name: block.state_dict()
                    })
        torch.save(state, checkpoint_file)

print(f'Best test acc. {best_test_acc}  at epoch {best_epoch}/{i}')
hours = (time.time() - start_time) / 3600
print(f'Done in {hours:.1f} hours with {n_gpus} GPU')


#!/usr/bin/env python3

FPORT_FRAME_MARKER = 0x7e
FPORT_ESCAPE_CHAR = 0x7d
FPORT_ESCAPE_MASK = 0x20

def to_hex_str(data) -> str:
    return ' '.join(f'{x:02x}' for x in data)

def to_binary_str(data) -> str:
    return ' '.join(f'{x:02b}' for x in data)

def channels_to_str(channels) -> str:
    return ' '.join(f'{i + 1:d}:{x:04d}' for i, x in enumerate(channels))

def get_fport_data(read_byte_serial_func) -> list:
    data = []
    frame_position = 0
    escaped_character = False
    while True:
        val = read_byte_serial_func()
        if val == FPORT_FRAME_MARKER:
            frame_position = 1
            escaped_character = False
            if len(data) > 0:
                break
        elif frame_position > 0:
            if escaped_character:
                val = val ^ FPORT_ESCAPE_MASK
                escaped_character = False
            elif val == FPORT_ESCAPE_CHAR:
                escaped_character = True
                continue
            data.append(val)
            frame_position += 1
    return data

def get_channels(read_byte_serial_func, adjust_values: bool = True):
    while True:
        data = get_fport_data(read_byte_serial_func)
        if data[1] == 0x00:
            break

    data = data[1:]
    channels = [0] * 16
    channels[0] = data[1] | data[2] << 8 & 0x07FF
    channels[1] = data[2] >> 3 | data[3] << 5 & 0x07FF
    channels[2] = data[3] >> 6 | data[4] << 2 | data[5] << 10 & 0x07FF
    channels[3] = data[5] >> 1 | data[6] << 7 & 0x07ff
    channels[4] = data[6] >> 4 | data[7] << 4 & 0x07FF
    channels[5] = data[7] >> 7 | data[8] << 1 | data[9] << 9 & 0x07FF
    channels[6] = data[9] >> 2 | data[10] << 6 & 0x07FF
    channels[7] = data[10] >> 5 | data[11] << 3 & 0x07FF
    channels[8] = data[12] | data[13] << 8 & 0x07FF
    channels[9] = data[13] >> 3 | data[14] << 5 & 0x07FF
    channels[10] = data[14] >> 6 | data[15] << 2 | data[16] << 10 & 0x07FF
    channels[11] = data[16] >> 1 | data[17] << 7 & 0x07FF
    channels[12] = data[17] >> 4 | data[18] << 4 & 0x07FF
    channels[13] = data[18] >> 7 | data[19] << 1 | data[20] << 9 & 0x07FF
    channels[14] = data[20] >> 2 | data[21] << 6 & 0x07FF
    channels[15] = data[21] >> 5 | data[22] << 3 & 0x07FF

    if adjust_values:
        channels = [(5 * ch // 8) + 880 for ch in channels]
    return channels

if __name__ == '__main__':
    from argparse import ArgumentParser
    from serial import Serial

    parser = ArgumentParser(description='Monitor FrSky F.Port receiverchannels')
    parser.add_argument('serial_port', type=str,
                        help='Serial port of the F.Port receiver')
    args = parser.parse_args()

    ser = Serial(args.serial_port, baudrate=115200)

    def read_byte():
        return ser.read()[0]

    try:
        while True:
            data = get_channels(read_byte)
            print('\r', channels_to_str(data), end='', flush=True)
    except KeyboardInterrupt:
        pass
    finally:
        ser.close()


from sys import version
import setuptools

with open('README.md', 'r') as fl:
    content = fl.read()
    fl.close()

setuptools.setup(

    name='AuxI',
    version='1.0',
    author='Adilson Luiz',
    description='Uma ferramenta para organização de arquivos, backup e instalação de programas',
    long_description=content,
    projects={
        'Source': 'https://github.com/ShellAndPython/AuxIT'
    },
    install_requires=['googledrivedownloader', 'requests'],
    classifiers=[
        'Programming Language :: Python :: 3',
        'Operating System :: OS Independent'
    ],
    python_requires='>=3.6'
)


from django.db import connections, models
from django.db.utils import ConnectionRouter

class ViewModel(models.Model):
    '''
    A model which is not backed by a table, but a view
    '''

    # Subclasses must define one of these
    VIEW_DEFINITION = None
    # view_query, if defined, should be a queryset (or a callable returning a queryset)
    view_query = None

    @classmethod
    def get_connection(cls):
        '''
        This method is multi-db ready. It uses your site's db router to determine which
        database the view is in.
        '''
        db = ConnectionRouter().db_for_read(cls())
        return connections[db]

    @classmethod
    def get_view_definition(cls):
        if cls.VIEW_DEFINITION :
            return cls.VIEW_DEFINITION

        definition = getattr(cls, 'view_query')
        if callable(definition):
            definition = definition()

        # See:
        # https://code.djangoproject.com/ticket/25705
        # https://stackoverflow.com/a/22828674
        # str(queryset.query) _sometimes_ provides valid SQL, but not in general
        cursor = cls.get_connection().cursor()

        try :
            mogrify = cursor.mogrify
        except AttributeError :
            raise Exception('''The database backend in use doesn't include the mogrify method (which I believe is only supported by Postgres), so you can't use "view_query" - you must supply raw sql via "VIEW_DEFINITION"''')

        return mogrify(*definition.query.sql_with_params()).decode('utf-8')

    class Meta:
        managed = False
        abstract = True

    @classmethod
    def run_sql(cls, sql):
        with cls.get_connection().cursor() as cursor :
            cursor.execute(sql)

    @classmethod
    def create_view(cls):
        '''
        Note - CREATE OR REPLACE doesn't always work when renaming/adding columns to a view,
        so we always drop then create
        '''
        cls.run_sql("""
            DROP VIEW IF EXISTS {view_name} CASCADE;
            CREATE VIEW {view_name} AS
            {view_definition};
        """.format(view_name=cls._meta.db_table, view_definition=cls.get_view_definition()))

class MaterializedViewModel(ViewModel):
    class Meta:
        managed = False
        abstract = True

    @classmethod
    def create_view(cls):
        cls.run_sql("""
            DROP MATERIALIZED VIEW IF EXISTS {view_name};
            CREATE MATERIALIZED VIEW {view_name} AS
            {view_definition}
        """.format(view_name=cls._meta.db_table, view_definition=cls.get_view_definition()))

    @classmethod
    def refresh_view(cls):
        cls.run_sql("REFRESH MATERIALIZED VIEW {}".format(cls._meta.db_table))


from kivy.uix.button import Button
from kivy.app import App
from kivy.uix.behaviors import DragBehavior
from kivy.uix.boxlayout import BoxLayout

class Box_layout(BoxLayout):
    def __init__(self,**kwargs):
        super(Box_layout, self).__init__(**kwargs)
        self.size_hint = (.50,.50)
        self.orientation = "vertical"
        for  x in range(5):
           DragButton.text=str(x)
           self.add_widget(DragButton())

class DragButton(DragBehavior, Button):
    def __init__(self,**kwargs):
       super(DragButton, self).__init__(**kwargs)
       self.drag_rect_width = 1366
       self.drag_rect_height = 768
       self.drag_timeout = 10000000
       self.drag_distance = 0
       print(self.height, self.width)
class atry(App):
    def build(self):
        return DragButton()
atry.run()


import serial
import time
import bpy
import math

ser = serial.Serial('COM11', 115200)
ser.close()
ser.open()

##objects
car1 = bpy.data.collections['car1'].all_objects
car2 = bpy.data.collections['car2'].all_objects
train = bpy.data.collections['Collection'].all_objects['train']
gate1l = bpy.data.collections['Collection'].all_objects['gate1_left']  #4
gate1r = bpy.data.collections['Collection'].all_objects['gate1_right']   #2
gate2l = bpy.data.collections['Collection'].all_objects['gate2_left']   #3
gate2r = bpy.data.collections['Collection'].all_objects['gate2_right']   #1

red = bpy.data.collections['Collection'].all_objects['Cube.004']
orange = bpy.data.collections['Collection'].all_objects['Cube.005']
green = bpy.data.collections['Collection'].all_objects['Cube.006']
reset_car1 = [7.296009063720703, 16.949875831604004, 23.837894439697266, 32.48568344116211, 42.40331268310547, 49.83244323730469]
reset_car2 = [-59.08498764038086, -51.431121826171875, -42.5431022644043, -33.89531326293945, -23.977684020996094, -16.548553466796875]
for j in range(6):
    car1[j].location[0] = reset_car1[j]
    car2[j].location[0] = reset_car2[j]
train.location[1] = 250
train.location[0] = -3.3609373569488525

def car1move():
    for i in range(6):
        car1[i].location[0] -= 1.1
        if(car1[i].location.x <= -32):
            car1[i].hide_viewport = True
            car1[i].location[0] += 60
            car1[i].hide_viewport = False
        bpy.ops.wm.redraw_timer(type = 'DRAW_WIN_SWAP', iterations=1)

def car2move():
    for i in range(6):
        car2[i].location[0] += 1.1
        if(car2[i].location.x >= 20):
            car2[i].hide_viewport = True
            car2[i].location[0] -= 60
            car2[i].hide_viewport = False
        bpy.ops.wm.redraw_timer(type = 'DRAW_WIN_SWAP', iterations=1)

#testing
close = False
values = []
loop = True
check = True
while(loop):
    time.sleep(.2)
    values = ser.readline().decode().split()
    if(values[0] == 'STOP'):
        break;
    if(values[4] == '1'):
        red.hide_viewport = True
        orange.hide_viewport = True
        green.hide_viewport = False
    if(values[4] == '2'):
        red.hide_viewport = True
        orange.hide_viewport = False
        green.hide_viewport = True
    if(values[4] == '3'):
        red.hide_viewport = False
        orange.hide_viewport = True
        green.hide_viewport = True
    train.location[1] = int(values[6])/10
    bpy.ops.wm.redraw_timer(type = 'DRAW_WIN_SWAP', iterations=1)
    gate1l.rotation_euler[0] = int(values[2])*(math.pi/180)
    gate2l.rotation_euler[0] = int(values[3])*(math.pi/180)
    gate1r.rotation_euler[0] = int(values[0])*(math.pi/180)
    gate2r.rotation_euler[0] = int(values[1])*(math.pi/180)
    bpy.ops.wm.redraw_timer(type = 'DRAW_WIN_SWAP', iterations=1)
    print(values)
    if(int(values[5]) == 2 or close):
        close = True
        if(check):
            indices1 = []
            indices2 = []
            for i in range(6):
                if(car1[i].location[0] < 5):
                    indices1.append(i)
                if(car2[i].location[0] > -13):
                    indices2.append(i)

            left1 = list(range(6))
            left2 = list(range(6))
            for i in indices1:
                left1.remove(i)

            for j in indices2:
                left2.remove(j)

            max1 = 0
            for i in left1:
               if(car1[i].location[0] > max1):
                   max1 = car1[i].location[0]

            max2 = 0
            for j in left2:
                if(car2[j].location[0] < max2):
                    max2 = car2[j].location[0]

            check = False

        if (len(indices1) != 0 or len(indices2) != 0):
            for k1 in indices1:
                car1[k1].location[0] -= 2
                if(car1[k1].location[0] <= -33):
                    car1[k1].hide_viewport = True
                    car1[k1].location[0] = max1+9
                    max1 += 9
                    car1[k1].hide_viewport = False
                    indices1.remove(k1)

            for k2 in indices2:
                car2[k2].location[0] += 2
                if(car2[k2].location[0] >= 20):
                    car2[k2].hide_viewport = True
                    car2[k2].location[0] = max2-9
                    max2-=9
                    car2[k2].hide_viewport = False
                    indices2.remove(k2)

            bpy.ops.wm.redraw_timer(type = 'DRAW_WIN_SWAP', iterations=1)

        else:
            close = False

    elif (int(values[5]) == 1):
        car1move()
        car2move()

#serial reading
ser.close()


from application import app
from application import db
from application.models import Quantity, Method, Recipes, Cuisine, Ingredients, Schedule, Measure
from sqlalchemy import inspect

app.app_context().push()

inspector = inspect(db.engine)
if inspector.has_table == True:
    pass
else:
    db.create_all()

if Recipes.query.filter_by(recipe_name = 'Pea and Cauliflower Curry').first():
    pass
else:

    cuisines = ('Indian', 'Salad', 'Pasta', 'Various')
    for cus in cuisines:
        i = Cuisine(cus)
        db.session.add(i)
        db.session.commit()

    ingredients = ['sunflower oil', 'cumin seeds', 'brown mustard seeds', 'asafetida', 'cauliflower', 'turmeric', 'ground coriander', 'salt', 'sugar', 'frozen peas','medium tomato', 'fresh coriander', 'fresh green chillies','root ginger']
    for ing in ingredients:
        i = Ingredients(ing)
        db.session.add(i)
        db.session.commit()

    measures = ["ml", "g", "tsp", "tbsp", "whole", "large", "medium", "small", "cloves", "cm", "handful", "cup"]
    for m in measures:
        measure = Measure(m)
        db.session.add(measure)
        db.session.commit()

    rec = Recipes('Pea and Cauliflower Curry', 'quick and easy curry, no calorie details', 1)
    db.session.add(rec)
    db.session.commit()

    quantities = [ (1, 1,"",100.0,1),(1, 2,"", 1.0, 3), (1, 3,"", 1.0, 3), (1, 4,"", 0.5, 3), (1, 5,'cut into 1cm pieces', 1.0, 7), (1, 6,"", 1.0, 3) , (1, 7,"", 1.0, 4), (1, 8,"", 1.5, 3), (1, 9,"", 1.0, 3), (1, 10,"", 400.0, 2), (1, 11,"", 1.0, 7), (1, 12, 'roughly chopped', 1.0, 11), (1, 13,"", 4.0, 5),(1, 14, 'peeled and roughly chopped', 4.0, 10)]
    for q in quantities:
        quant = Quantity(*q)
        db.session.add(quant)
        db.session.commit()

    methods = [(1, 'Crush the chillies and ginger together with a pinch of salt to make a fine masala paste.\n Heat the oil in non-stick pan over a medium heat for 30s, then add the cumin and mustard seeds.\n When the mustard seeds start to pop reduce the heat and stir in the asafetida.\n Add the cauliflower and return the heat to medium.\n Stir in the masala paste, turmeric, ground coriander, salt, and sugar.\n Cover the pan and leave to cook for 8-10 minutes, stirring every few minutes.\n Stir in the peas and tomato and cook for a further 3-5 minutes.\n Remove from the heat and sprinkle with the chopped coriander.\n Leave for 5 minutes to let flavours develop')]
    for m in methods:
        met = Method(*m)
        db.session.add(met)
        db.session.commit()

    schedules =[(1,1,1,1), (2,1,1,1), (3,1,1,1), (4,1,1,1), (0,1,1,1)]
    for s in schedules:
        sched = Schedule(*s)
        db.session.add(sched)
        db.session.commit()

if __name__ == "__main__":
    app.run(debug=True, host='0.0.0.0', port=5050)


from django.db import models
from django.contrib.auth.models import User,auth

# Create your models here.

# To create auction
class Auction(models.Model):
    item_name = models.CharField(max_length=150)
    start_price = models.DecimalField(max_digits=10, decimal_places=2)
    start_time = models.DateTimeField(auto_now_add=True)
    end_time = models.DateTimeField(auto_now=False, auto_now_add=False)
    bidd_count = models.PositiveIntegerField(default=0)
    winner = models.CharField(max_length=150,blank=True,null=True)
    final_amount = models.DecimalField(max_digits=10, decimal_places=2,blank=True,null=True)

    class Meta:
        verbose_name = 'Auctions'
        verbose_name_plural = 'Auctions'

    def __str__(self):
        return  self.item_name

# Bidds
class Bidd(models.Model):
    user = models.ForeignKey(User, on_delete=models.DO_NOTHING)
    amount = models.DecimalField(max_digits=10, decimal_places=2)
    applied_at = models.DateTimeField(auto_now_add=True)
    auction = models.ForeignKey(Auction,on_delete=models.DO_NOTHING)

    class Meta:
        verbose_name = 'Bidds'
        verbose_name_plural = 'Bidds'

    def __str__(self):
        return  str(self.amount)


import csv
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import random
import math
import pickle
from collections import namedtuple, deque
import JSBSim
import time
import subprocess
import os
import glob
import numpy as np
import statistics

Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward'))

# DQN Network
class DQN(nn.Module):

    def __init__(self, n_observations, n_actions):
        super(DQN, self).__init__()
        self.layer1 = nn.Linear(n_observations, 128)
        self.layer2 = nn.Linear(128, 128)
        self.layer3 = nn.Linear(128, n_actions)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        return self.layer3(x)

class ReplayMemory:
    def __init__(self, capacity):
        self.memory = deque(maxlen=capacity)

    def push(self, *args):
        self.memory.append(Transition(*args))

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

def save_replay_memory(memory, file_name):
    with open(file_name, 'wb') as file:
        pickle.dump(memory, file)

def load_replay_memory(file_name):
    with open(file_name, 'rb') as file:
        return pickle.load(file)

# Function to save model weights
def save_model_weights(model, filename):
    torch.save(model.state_dict(), filename)

# Function to load model weights
def load_model_weights(model, filename):
    if os.path.exists(filename):
        model.load_state_dict(torch.load(filename))
        model.train()  # Set the model to training mode
    else:
        print(f"No weights found at {filename}, initializing model with default weights.")
        torch.save(model.state_dict(), filename)
        model.train()

# Global Variables
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 1.0
EPS_END = 0.01
EPS_DECAY = 2000
TAU = 0.005
LR = 1e-4

# action space with 0.1 intervals
action_space = [
    (elevator, aileron, rudder)
    for elevator in [i/10 for i in range(-10, 11)]
    for aileron in [i/10 for i in range(-10, 11)]
    for rudder in [i/10 for i in range(-10, 11)]
]

n_actions = len(action_space)
n_observations = 4  # altitude, xaccel, yaccel, zaccel

policy_net = DQN(n_observations, n_actions)
target_net = DQN(n_observations, n_actions)
load_model_weights(policy_net, "policy_net_weights.pth")
load_model_weights(target_net, "target_net_weights.pth")
target_net.load_state_dict(policy_net.state_dict())
optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
memory = ReplayMemory(1000000)
steps_done = 0

# Function to dynamically adjust epsilon
def dynamic_eps(run_number, steps_done, total_runs=10, min_eps=0.01, max_eps=1.0):
    return min_eps + (max_eps - min_eps) * math.exp(-1. * steps_done / EPS_DECAY)

# Select Action Function with Dynamic Epsilon
def select_action(state, run_number, steps_done):
    eps_threshold = dynamic_eps(run_number, steps_done)
    if random.random() > eps_threshold:
        with torch.no_grad():
            return policy_net(state).max(1)[1].view(1, 1)
    else:
        return torch.tensor([[random.randrange(n_actions)]], dtype=torch.long)

def get_state_reward(current_state, altitude_change):
    current_altitude = current_state[0]

    steep_decline_rate = -10  # Example rate in feet per 0.5 second

    reward = 0

    # Penalty for steep decline
    if altitude_change < steep_decline_rate * 0.5:  # Assuming a 0.5 second interval
        reward -= 5

    # Penalty for dropping altitude to zero or below
    if current_altitude <= 0:
        reward -= 10

    return reward

def calculate_stability_reward(altitudes, threshold=5, stable_range=10):
    """
    Calculate stability reward based on altitude variations over the last 'threshold' seconds.
    'stable_range' defines the acceptable range of altitude variation for stability.
    """
    if len(altitudes) < threshold:
        return 0  # Not enough data to calculate stability

    recent_altitudes = altitudes[-threshold:]
    altitude_variance = np.var(recent_altitudes)

    if altitude_variance <= stable_range**2:
        return 2  # Reward for maintaining stable altitude
    else:
        return 0  # No reward if altitude varies too much

def get_episodic_reward(flight_duration, episodes_for_average, directory):
    # Reward based on flight duration
    # Episodic reward is equal to deviation from average flight time of previous series of episodes

    all_files = [os.path.join(directory, file) for file in os.listdir(directory) if os.path.isfile(os.path.join(directory, file))]
    sorted_files = sorted(all_files, key=os.path.getmtime, reverse=True)

    # Get the last n files
    last_n_files = sorted_files[:episodes_for_average]

    # Read the last line of each file
    last_lines = []
    for file in last_n_files:
        with open(file, 'r') as f:
            lines = f.readlines()
            if lines:
                last_line = lines[-1].strip()
                last_line = last_line.split(",,,,")
                last_line = last_line[-1]
                last_line = float(last_line)
                last_lines.append(last_line)

    return flight_duration - statistics.mean(last_lines)

# Optimize Model Function
def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    batch = Transition(*zip(*transitions))

    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state if s is not None]) if any(non_final_mask) else None
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    state_action_values = policy_net(state_batch).gather(1, action_batch)

    next_state_values = torch.zeros(BATCH_SIZE)
    if non_final_next_states is not None:
        with torch.no_grad():
            next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0]
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    criterion = nn.SmoothL1Loss()
    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))

    optimizer.zero_grad()
    loss.backward()
    for param in policy_net.parameters():
        param.grad.data.clamp_(-1, 1)
    optimizer.step()

def load_simulator_data(file_path):
    with open(file_path, 'r') as f:
        reader = csv.reader(f)
        next(reader)
        data = [line for line in reader]
    return data

def log_run_data(run_number, episode_states, episode_actions, flight_duration, reward):
    """Logs the data for a single run to a CSV file."""
    run_log_file = f'runlogs/run{run_number}.csv'
    with open(run_log_file, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["altitude", "xaccel", "yaccel", "zaccel", "elevator", "aileron", "rudder"])
        for state, action in zip(episode_states, episode_actions):
            writer.writerow(state + list(action))
        writer.writerow(["Reward", "", "", "", reward])
        writer.writerow(["Flight Duration", "", "", "", flight_duration])

def initialize_runlogs():
    """Initializes the runlogs directory and finds the last run number."""
    if not os.path.exists('runlogs'):
        os.makedirs('runlogs')

    run_files = glob.glob('runlogs/run*.csv')
    run_numbers = [int(os.path.splitext(os.path.basename(f))[0][3:]) for f in run_files]
    return max(run_numbers) if run_numbers else 0

def process_state(altitude, xaccel, yaccel, zaccel):
    """Process a single state and return the selected action."""
    state = torch.tensor([altitude, xaccel, yaccel, zaccel], dtype=torch.float32).unsqueeze(0)
    action = select_action(state)
    elevator, aileron, rudder = action_space[int(action.item())]
    return (elevator, aileron, rudder)

def update_memory(state, action, next_state, reward):
    """Store the transition in replay memory"""
    memory.push(state, action, next_state, torch.tensor([reward], dtype=torch.float32))

def run_fgfs(filename):
    with open(filename, 'r') as file:
        fg_command = file.read().strip()  # Read and strip any trailing whitespace
    return subprocess.Popen(fg_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

def kill_fgfs():
    subprocess.run(["bash", "kill.sh"])

def main():
    last_run_number = initialize_runlogs()
    current_run_log = last_run_number + 1
    run_number = 0

    while run_number < 11:
        jsbsim_client = JSBSim.JsbsimInterface()
        steps_done = 0  # Reset steps_done for each run
        altitudes = []  # List to store altitudes for stability calculation
        stability_reward_interval = 100  # Assuming simulation step is 0.01s, calculate every second
        stability_reward = 0

        try:
            print(f"Starting JSBSim for Run No {run_number}!")
            jsbsim_client.start()

            episode_actions = []
            episode_states = []
            episode_actions_log = []
            last_altitude = None

            while True:
                jsbsim_client.run()

                current_state = jsbsim_client.get_state()
                altitude, xaccel, yaccel, zaccel = current_state
                episode_states.append(current_state)

                # Calculate altitude change
                altitude_change = 0
                if last_altitude is not None:
                    altitude_change = altitude - last_altitude
                last_altitude = altitude

                altitudes.append(altitude)

                if len(altitudes) % stability_reward_interval == 0:
                    stability_reward = calculate_stability_reward(altitudes, threshold=100, stable_range=10)

                if jsbsim_client.gear_contact():
                    print("Landing gear contact. Ending simulation.")
                    next_state = None  # Set next state to None if simulation ends
                    break

                if altitude <= 0:  # Check for simulation end condition
                    next_state = None
                    break

                # Process state to get actions
                torch_current_state = torch.tensor(current_state, dtype=torch.float32).unsqueeze(0)
                action_index = select_action(torch_current_state, run_number, steps_done)
                elevator, aileron, rudder = action_space[int(action_index.item())]
                actions = [elevator, aileron, rudder]
                episode_actions_log.append(actions)

                # Apply actions to JSBSim
                jsbsim_client.set_controls(elevator, aileron, rudder)

                # Get the next state
                jsbsim_client.run()
                next_state = jsbsim_client.get_state()

                # Store the state and action
                torch_next_state = torch.tensor(next_state, dtype=torch.float32).unsqueeze(0) if next_state is not None else None
                episode_actions.append((torch_current_state, action_index))

                # Calculate and update memory with state reward
                state_reward = get_state_reward(current_state, altitude_change)
                # Update the total reward including stability_reward
                total_reward = state_reward + stability_reward
                # print(total_reward) uncomment this to see the total_reward of each state
                action_tensor = torch.tensor([[action_index]], dtype=torch.long)
                update_memory(torch_current_state, action_tensor, torch_next_state, total_reward)

            # Use JSBSim's internal simulation time as flight duration
            flight_duration = jsbsim_client.get_sim_time()
            print(f"Flight duration for Run No {run_number}: {flight_duration} seconds")

            # Calculate episodic reward and update memory
            episodic_reward = get_episodic_reward(flight_duration, 100, "runlogs")
            for state, action in episode_actions:
                update_memory(state, action, None, episodic_reward)

            # Log run data
            log_run_data(current_run_log, episode_states, episode_actions_log, flight_duration, episodic_reward)

            run_number += 1
            current_run_log += 1
            if run_number % 5 == 0:
                optimize_model()

        finally:
            jsbsim_client.stop()

    save_model_weights(policy_net, "policy_net_weights.pth")
    save_model_weights(target_net, "target_net_weights.pth")
    save_replay_memory(memory, "replay_memory.pkl")

if __name__ == "__main__":
    main()


import time
import yaml
import logging
import sys
from datetime import datetime
from typing import List
from helper_classes import User, Course
from selenium import webdriver
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from enums import Element as element_consts
logging.getLogger('selenium').setLevel(logging.ERROR)
logging.getLogger('urllib3').setLevel(logging.ERROR)

class JCGolfTeeTimeDriver():

    def __init__(self) -> None:
        self.driver = None
        self.login_info = None
        self.user = None
        self.logger = None
        self.tee_times = []

    def setup(self) -> bool:
        """Setup the webdriver
        :return: True if driver is setup
        """
        self.logger = logging.getLogger('JCGolfDriver')
        self.logger.addHandler(logging.StreamHandler(sys.stdout))
        self.logger.info("running driver setup for chrome")
        chrome_options = Options()
        chrome_options.add_argument('--headless')
        # with headless
        # self.driver = webdriver.Chrome(ChromeDriverManager(log_level=0, print_first_line=False).install(), chrome_options=chrome_options)
        self.driver = webdriver.Chrome(ChromeDriverManager(log_level=0, print_first_line=False).install())
        self.driver.maximize_window()
        with open(r'login.yaml') as file:
            self.login_info = yaml.full_load(file)
        self.logger.info('driver setup successful')
        return True

    def _click_button(self, button_element) -> bool:
        """Click a web element button
        :param button_class: class for button to be clicked
        :return: True if button is clicked successfully
        """
        button_element.click()
        return True

    def _enter_input(self, input_box_id: str, input: str) -> bool:
        """Enter email into input_box
        :param input_box_id: id of input_box to type email in to
        :param input: string input to go into box
        :return: True if email is entered
        """
        element = WebDriverWait(self.driver, 30).until(
            EC.presence_of_element_located((By.ID, input_box_id))
        )
        element.send_keys(input)
        element = self.driver.find_element(By.CLASS_NAME, element_consts.LOGIN_NEXT_BUTTON)
        return self._click_button(element)

    def _wait_for_element(self, element_class_name) -> bool:
        """Wait for the Tee Time list to refresh
        :param element_class_name: name of element class
        :return: True if list refreshes
        """
        self.logger.debug(f"waiting for {element_class_name}")
        WebDriverWait(self.driver, 30).until(
            EC.presence_of_element_located((By.CLASS_NAME, element_class_name))
        )
        return True

    def login(self, user_type='public', course_name='Encinitas Ranch') -> bool:
        """Login to the JCGolf website
        :param user_type: user type to login as
        :param course_name: Course to check times for
        :returns: True if login successful
        """
        self.user = User(user_type)
        self.course = Course(course_name)
        self.logger.info("logging into JCGolf site")
        self.driver.get(f"https://jc{self.user.sub_url}{self.course.url_number}.cps.golf/onlineresweb/search-teetime")
        if self.user.user_type is not 'public':
            self._enter_input(element_consts.LOGIN_TEXTBOX, self.login_info['email'])
            self._enter_input(element_consts.PASSWORD_TEXTBOX, self.login_info['password'])
        time.sleep(5)
        # below is broken, need to account for when current day does not have any tee times list
        # if self._wait_for_element(element_consts.TEE_TIME_LIST) or self._wait_for_element(element_consts.NO_TEE_TIME_LIST):
        self.logger.debug("tee time list generated")
        self.logger.info('jcgolf login successful')
        return True
        # return False

    def change_date(self) -> bool:
        """Change the current date for tee time search
        :return: return True if date changes
        """
        '''
        element = self.driver.find_element_by_class_name("mat-form-field"
            f".ng-tns-{self.user.date_sub_str}"
            ".mat-primary"
            ".mat-form-field-type-mat-input"
            ".mat-form-field-appearance-legacy"
            ".mat-form-field-can-float"
            ".mat-form-field-hide-placeholder")
        if not self._click_button(element):
            return False
        '''
        # today = self.driver.find_elements_by_class_name(element_consts.TODAY_DATE)[1]
        today = self.driver.find_element_by_class_name(element_consts.TODAY_DATE)
        elements = self.driver.find_elements_by_class_name(element_consts.DATE_BUTTON)
        today_index = 0
        for index, element in enumerate(elements):
            if element.text == today.text:
                today_index = index
                break
        offset_index = today_index + self.user.date_offset
        if not self._click_button(elements[offset_index]):
            return False
        # need to wait a second for click to register
        time.sleep(1)
        if not self._wait_for_element(element_consts.TEE_TIME_LIST):
            return False
        return True

    def get_available_tee_times(self) -> List:
        """Get a list of available Tee Times
        """
        # TODO: add wait for list to be available in case of auto-refresh
        tee_times =  self.driver.find_elements_by_class_name(element_consts.TEE_TIME_BUTTON)
        self.logger.debug(f'number of available tee times: {len(tee_times)}')
        return tee_times

    def filter_tee_times(self, tee_times: List) -> List:
        """Filter tee time results
        :param tee_times: list of tee_time selenium elements
        :return: filtered list of selenium elements
        """
        filtered_times = []
        self.logger.debug("filtering for only 4 person tee times")
        for tee_time in tee_times:
            # time, cart rate, holes/golfers, price
            _, _, players, _, = tee_time.text.split('\n')
            if '4 GOLFERS' in players:
                filtered_times.append(tee_time)
        self.logger.debug(f'number of filtered tee times: {len(filtered_times)}')
        return filtered_times

    def find_best_tee_time(self, tee_times: List, target_time: str):
        """Find the best tee time given a list and target star time
        :param tee_times: List of tee time web elements
        :param target_time: target starting time
        :return: tee time closest to target
        """
        self.logger.debug(f"filtering to time closest to {target_time}")
        target = datetime.strptime(target_time, "%I:%M %p")
        best_tee_time = tee_times[0]
        best_tee_time_start = datetime.strptime(best_tee_time.text.split('\n')[0], "%I:%M %p")
        for time in tee_times[1::]:
            tee_time_start = datetime.strptime(time.text.split('\n')[0], "%I:%M %p")
            if abs(target - tee_time_start) < abs(target - best_tee_time_start):
                best_tee_time = time
                best_tee_time_start = tee_time_start
            else:
                break
        best_tee_time_text = best_tee_time.text.split('\n')[0]
        self.logger.info(f'best tee time: {best_tee_time_text}')
        return best_tee_time

    def select_tee_time(self, tee_time_element: str) -> bool:
        """Select a tee time element
        :param tee_time_element: tee_time_element to click
        :return: True if button is clicked successfully
        """
        self.logger.info('selecting best tee time')
        return self._click_button(tee_time_element)

    def select_agreement_next(self) -> bool:
        """Select the next button on the agreement page
        :return: True if button is clicked successfully
        """
        time.sleep(10)
        self.logger.info('accepting agreement')
        element = self.driver.find_element_by_class_name(element_consts.AGREEMENT_NEXT_BUTTON)
        return self._click_button(element)

    def finalize_reservation(self) -> bool:
        """Click the finalize reservation button
        :returns: True if button is clicked successfully
        """
        self.logger.info('finalizing reservation!')
        # three confirmation pages confirmation page
        for _ in range(3):
            time.sleep(10)
            element = self.driver.find_element_by_class_name(element_consts.FINALIZE_RESERVATION_BUTTON)
            self._click_button(element)


import os
import re
from html.parser import HTMLParser
from collections import defaultdict
import numpy as np
from google.cloud import storage
from google.oauth2 import service_account
import json
import time
import concurrent.futures

f = open('ds-561-am-cc3cb4be64ff.json')
creds = json.load(f)

credentials = service_account.Credentials.from_service_account_info(creds)

def read_last_run_time(types,filename="logs.json"):
    try:
        with open(filename, "r") as f:
            data = f.read()
            if not data:
                return None  # File is empty
            data_json = json.loads(data)
            return data_json.get(types + '_time', None)
    except (FileNotFoundError, json.JSONDecodeError):
        return None  # File not found or invalid JSON

def write_run_info(elapsed_upload_time,elapsed_processing_time, incoming_stats, outgoing_stats, top_5_counts, filename="logs.json"):
    with open(filename, "w") as f:
        data = {'upload_time': elapsed_upload_time}
        data['processing_time'] = elapsed_processing_time
        # incoming_stats["Quintiles"] = incoming_stats["Quintiles"].tolist()
        # outgoing_stats["Quintiles"] = outgoing_stats["Quintiles"].tolist()
        # data['incoming_stats'] = incoming_stats
        # data['outgoing_stats'] = outgoing_stats
        # data['top_5_pages'] = top_5_counts
        json.dump(data, f)

def upload_blob(bucket_name, source_file_name, destination_blob_name):
    """Uploads a file to the bucket."""
    storage_client = storage.Client.create_anonymous_client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(destination_blob_name)
    blob.upload_from_filename(source_file_name)
    print(f"File {source_file_name} uploaded to {destination_blob_name}.")

def upload_folder_to_gcs(bucket_name, folder_path):
    storage_client = storage.Client.create_anonymous_client()
    bucket = storage_client.get_bucket(bucket_name)

    count = 0

    last_run_time = read_last_run_time("upload","logs.json")

    if last_run_time is not None:
        print(f"Please note, this could take some time. Last run took {last_run_time} seconds.")

    start_time = time.time()

    for root, _, files in os.walk(folder_path):
        for file in files:
            local_file = os.path.join(root, file)
            remote_path = os.path.relpath(local_file, folder_path)

            # Check if blob already exists
            blob = bucket.blob('files/'+remote_path)
            if not blob.exists():
                blob.upload_from_filename(local_file)
                count+=1
            else:
                break

    end_time = time.time()
    elapsed_time = end_time - start_time

    write_run_info(elapsed_time,read_last_run_time('processing'), {}, {}, {})

    print(f"{count} files uploaded")

# Custom parser to extract links
class LinkExtractor(HTMLParser):
    def __init__(self):
        super().__init__()
        self.links = []

    def handle_starttag(self, tag, attrs):
        if tag == "a":
            for name, value in attrs:
                if name == "href":
                    self.links.append(value)

def pagerank(graph, max_iter=10, damping=0.85, tol=0.005):
    num_pages = len(graph)
    pr = {page: 1/num_pages for page in graph}
    base_pr = (1 - damping) / num_pages

    for _ in range(max_iter):
        next_pr = {page: base_pr for page in graph}
        for page, outgoing_links in graph.items():
            if not outgoing_links:
                continue  # Skip this page because it has no outgoing links

            share_pr = pr[page] / len(outgoing_links)
            for link in outgoing_links:
                if link not in next_pr:
                    next_pr[link] = base_pr
                next_pr[link] += share_pr * damping

        # Check for convergence
        diff = sum(abs(next_pr[page] - pr[page]) for page in graph)
        if diff < tol:
            break

        pr = next_pr

    return pr

def process_blob_data(blob):
    """Processes a blob and extracts links."""
    filename = blob.name.replace("files/", "").replace(".html", "")
    parser = LinkExtractor()

    with blob.open("r") as f:
        parser.feed(f.read())

    links = {re.sub(r"\.html$", "", link) for link in parser.links}
    return filename, links

def compute_stats(link_counts):
    return {
        "Average": np.mean(link_counts),
        "Median": np.median(link_counts),
        "Max": max(link_counts),
        "Min": min(link_counts),
        "Quintiles": np.percentile(link_counts, [20, 40, 60, 80])
    }

def main():
    graph = defaultdict(set)
    out_count = defaultdict(int)

    storage_client = storage.Client.create_anonymous_client()
    bucket_name = "hw-2-files-bucket"

    folder_path = "./files"

    upload_folder_to_gcs(bucket_name, folder_path)

    last_run_time = read_last_run_time("processing","logs.json")

    if last_run_time is not None:
        print(f"Please note, this could take some time. Last web run took {last_run_time} seconds.")

    start_time = time.time()

    bucket = storage_client.bucket(bucket_name)
    blobs = list(bucket.list_blobs())

    out_count = defaultdict(int)

    BATCH_SIZE = 100
    total_blobs = len(blobs)

    for i in range(0, total_blobs, BATCH_SIZE):
        batch = blobs[i:i+BATCH_SIZE]

        with concurrent.futures.ThreadPoolExecutor() as executor:
            results = list(executor.map(process_blob_data, batch))

            for filename, links in results:
                graph[filename] = links
                out_count[filename] = len(links)

    in_count = [len(graph[filename]) for filename in graph]
    out_count = [out_count[filename] for filename in graph]

    print("Incoming links stats:", compute_stats(in_count))
    print("Outgoing links stats:", compute_stats(out_count))

    pr = pagerank(graph)
    top_5_pages = sorted(pr, key=pr.get, reverse=True)[:5]
    print("Top 5 pages by PageRank:", top_5_pages)

    end_time = time.time()

    write_run_info(read_last_run_time('upload'), end_time, compute_stats(in_count), compute_stats(out_count), top_5_pages)

if __name__ == "__main__":
    main()