#
# Copyright (c) 2018 Institute for Basic Science
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

from poreplex.csupport import detect_events
import numpy as np
import pandas as pd
from scipy.signal import medfilt

class PolyASignalAnalyzer:

    CONFIG_SLOTS = [
        'refinement_expansion', 'event_detection', 'polya_stdv_max', 'polya_stdv_range',
        'spike_tolerance', 'spike_weight', 'openend_expansion', 'recalibrate_shifted_signal',
        'polya_mean_dist', 'polya_mean_z_cutoff', 'polya_mean_trigger_recalibration',
        'maximum_openend_extension', 'median_pre_filter',
    ]

    def __init__(self, config):
        for name in self.CONFIG_SLOTS:
            setattr(self, name, config[name])

        mean_loc, mean_scale = config['polya_mean_dist']
        self.polya_mean_cutoff = (
            mean_loc - mean_scale * config['polya_mean_z_cutoff'],
            mean_loc + mean_scale * config['polya_mean_z_cutoff'])

        self.polya_mean_trigger_recalibration *= config['polya_mean_dist'][1]

    def __call__(self, npread, rough_range, stride, polya_range=None, ext_depth=0):
        raw_signal = npread.load_signal(pool=None, pad=False)

        minimum_expansion_unit = self.openend_expansion // stride
        rough_begin, rough_end = rough_range
        if rough_end is None or rough_end - rough_begin < minimum_expansion_unit:
            rough_end = rough_begin + minimum_expansion_unit

        insp_begin = max(0, rough_begin * stride - self.refinement_expansion)
        insp_end = min(len(raw_signal), (rough_end + 1) * stride + self.refinement_expansion)
        adapter_end = rough_range[0] * stride - insp_begin
        polya_signal = raw_signal[insp_begin:insp_end]
        if self.median_pre_filter > 1:
            polya_signal = medfilt(polya_signal, self.median_pre_filter)

        events = pd.DataFrame(detect_events(polya_signal, **self.event_detection))
        events['end'] = (events['start'] + events['length']).astype(np.int64)
        events['is_polya'] = events['mean'].between(*(polya_range or self.polya_mean_cutoff))

        entryfunc = (self.call_polya if rough_range[1] is not None
                                     else self.try_recalibrate_shifted_signal)
        entryfunc(npread, events, polya_signal, insp_begin, insp_end,
                  (rough_begin, rough_end), adapter_end, len(raw_signal), stride,
                  polya_range, ext_depth)

    def call_polya(self, npread, events, polya_signal, signal_begin, signal_end,
                   base_range, adapter_end, full_length, stride, polya_range=None,
                   ext_depth=0):
        polya_events = self.find_best_polya_interval(events)

        # Retry it with right-extended interval if poly(A) is not terminated within the interval
        if (len(polya_events) > 0 and polya_events.index[-1] == events.index[-1] and
                    signal_end < full_length and ext_depth < self.maximum_openend_extension):
            return self(npread,
                        (base_range[0], base_range[1] + self.openend_expansion // stride),
                        stride, polya_range, ext_depth + 1)

        # Warp to the recalibration mode if the mean signal level is out of expected range.
        def is_polya_signal_shifted():
            mean_polya_level = ((polya_events['mean'] * polya_events['length']).sum() /
                                polya_events['length'].sum())
            return (
                abs(mean_polya_level - self.polya_mean_dist[0]) >
                self.polya_mean_trigger_recalibration)

        if len(polya_events) == 0 or (polya_range is None and is_polya_signal_shifted()):
            return self.try_recalibrate_shifted_signal(npread, events, polya_signal,
                        signal_begin, signal_end, base_range, adapter_end,
                        full_length, stride, None, ext_depth)

        # Quality check for the longest event
        longest_event = polya_events.loc[polya_events['length'].idxmax()]
        longest_event_stdv = self.calc_internal_polya_stdv(polya_signal, longest_event)

        if longest_event_stdv < self.polya_stdv_max and len(polya_events) > 0:
            polya_begin = polya_events.iloc[0]['start']
            last_event = polya_events.iloc[-1]
            polya_end = last_event['start'] + last_event['length']

            dwell_time = int(polya_events[polya_events['is_polya']]['length'].sum())
            spike_positions = np.where(~polya_events['is_polya'])[0]
            spikes = [
                (polya_events.iloc[spk]['length'],) +
                tuple(polya_events.iloc[spk-1:spk+2]['mean'].tolist())
                for spk in spike_positions]

            npread.set_polya_tail({
                'begin': int(polya_begin) + signal_begin,
                'end': int(polya_end) + signal_begin,
                'dwell_time': dwell_time / npread.sampling_rate,
                'spikes': spikes,
            })
        elif polya_range is None:
            self.try_recalibrate_shifted_signal(npread, events, polya_signal,
                        signal_begin, signal_end, base_range, adapter_end,
                        full_length, stride, None, ext_depth)

    def try_recalibrate_shifted_signal(self, npread, events, polya_signal, signal_begin,
                                  signal_end, base_range, adapter_end, full_length, stride,
                                  polya_range=None, ext_depth=0):
        config = self.recalibrate_shifted_signal
        anchorevents = events[(events['start'] <= adapter_end +
                                    config['max_dist_from_adapter']) &
                              (events['end'] > adapter_end) &
                              (events['stdv'] < config['max_stdv'])]
        if len(anchorevents) == 0:
            return

        polya_mean = ((anchorevents['mean'] * anchorevents['length']).sum() /
                      anchorevents['length'].sum())
        polya_range = (
            polya_mean - self.polya_mean_dist[1] * self.polya_mean_z_cutoff,
            polya_mean + self.polya_mean_dist[1] * self.polya_mean_z_cutoff)

        events['is_polya'] = events['mean'].between(*polya_range)
        if events[events['is_polya']]['length'].sum() >= config['min_length']:
            self.call_polya(npread, events, polya_signal, signal_begin, signal_end,
                            base_range, adapter_end, full_length, stride, polya_range,
                            ext_depth)

    def calc_internal_polya_stdv(self, signal, row):
        length = int(row['length'])
        begin = int(row['start'] + length * self.polya_stdv_range[0])
        end = int(row['start'] + length * self.polya_stdv_range[1])
        return signal[begin:end].std() if end - begin > 2 else np.nan

    def find_best_polya_interval(self, events):
        nullmtx = np.zeros([len(events) + 1, len(events) + 1], dtype=np.int64)

        matching_scores = nullmtx.copy()
        matching_scores[0, 1:] = [
            v if v > 0 else v * self.spike_weight
            for v in (events['is_polya'] * 2 - 1) * events['length']]

        spike_scores = nullmtx.copy()
        spike_scores[0, 1:] = np.choose(events['is_polya'],
            np.array([-events['length'], [1] * len(events)]))

        for i in range(1, len(events)+1): # i=start+1
            for j in range(1, len(events)+1): # j=end+1 (inclusive)
                if i > j:
                    continue

                matching_scores[i, j] = matching_scores[i, j - 1] + matching_scores[0, j]

                s = (-1 if spike_scores[i, j - 1] < 0
                     else (self.spike_tolerance if spike_scores[0, j] > 0
                           else spike_scores[i, j - 1] + spike_scores[0, j]))
                spike_scores[i, j] = s

        spike_scores = (spike_scores > 0).astype(np.int64)
        final_scores = matching_scores[1:, 1:] * spike_scores[1:, 1:]
        polya_start, polya_end = np.unravel_index(final_scores.argmax(), (len(events),) * 2)

        if final_scores[polya_start, polya_end] <= 0:
            return [] # no poly(A) found

        return events.iloc[polya_start:polya_end+1]


#!/u/flod/conda-envs/genEM3/bin/python

import os
import pdb
import argparse
import numpy as np
import torch
import datetime
from torch.utils.data import DataLoader
from genEM3.data.wkwdata import WkwData
from genEM3.model.autoencoder2d import Encoder_4_sampling_bn_1px_deep_convonly_skip, AE_Encoder_Classifier, Classifier3Layered
from genEM3.inference.inference import Predictor

torch.multiprocessing.set_sharing_strategy('file_system')

parser = argparse.ArgumentParser(description='mSEM artifact prediction parallel')
parser.add_argument('-i', '--idx', help='bbox idx prediction', required=False, default=0, type=int)
parser.add_argument('-v', '--verbose', help='verbose output', required=False, default=True, type=bool)

def predict_bbox_from_json(bbox_idx, verbose=True):

    if verbose:
        print('(' + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") +
              ') Starting Parallel Prediction ... bbox: {}'.format(bbox_idx))

    run_root = os.path.dirname(os.path.abspath(__file__))
    cache_HDD_root = os.path.join(run_root, '.cache/')
    datasources_json_path = os.path.join(run_root, 'datasources_predict_parallel.json')
    state_dict_path = os.path.join(run_root, '../../training/ae_classify_v09_3layer_unfreeze_latent_debris_clean_transform_add_clean2_wiggle/.log/run_w_pr/epoch_700/model_state_dict')
    device = 'cpu'

    output_wkw_root = '/tmpscratch/webknossos/Connectomics_Department/2018-11-13_scMS109_1to7199_v01_l4_06_24_fixed_mag8_artifact_pred'
    output_label = 'probs_sparse'

    batch_size = 128
    input_shape = (140, 140, 1)
    output_shape = (1, 1, 1)
    num_workers = 12

    kernel_size = 3
    stride = 1
    n_fmaps = 16
    n_latent = 2048
    input_size = 140
    output_size = input_size
    model = AE_Encoder_Classifier(
        Encoder_4_sampling_bn_1px_deep_convonly_skip(input_size, kernel_size, stride, n_latent=n_latent),
        Classifier3Layered(n_latent=n_latent))

    datasources = WkwData.datasources_bbox_from_json(datasources_json_path, bbox_ext=[1024, 1024, 1024], bbox_idx=bbox_idx, datasource_idx=0)
    dataset = WkwData(
        input_shape=input_shape,
        target_shape=output_shape,
        data_sources=datasources,
        stride=(35, 35, 1),
        cache_HDD=False,
        cache_RAM=False,
        cache_HDD_root=cache_HDD_root
    )

    prediction_loader = torch.utils.data.DataLoader(
        dataset=dataset, batch_size=batch_size, num_workers=num_workers)

    checkpoint = torch.load(state_dict_path, map_location=lambda storage, loc: storage)
    state_dict = checkpoint['model_state_dict']
    model.load_state_dict(state_dict)

    output_prob_fn = lambda x: np.exp(x[:, 1, 0, 0])
    # output_dtype = np.uint8
    output_dtype = np.float32
    # output_dtype_fn = lambda x: (logit(x) + 16) * 256 / 32
    output_dtype_fn = lambda x: x
    # output_dtype_fni = lambda x: expit(x / 256 * 32 - 16)
    output_dtype_fni = lambda x: x

    predictor = Predictor(
        model=model,
        dataloader=prediction_loader,
        output_prob_fn=output_prob_fn,
        output_dtype_fn=output_dtype_fn,
        output_dtype=output_dtype,
        output_label=output_label,
        output_wkw_root=output_wkw_root,
        output_wkw_compress=True,
        device=device,
        interpolate=None)

    predictor.predict(verbose=verbose)

if __name__ == '__main__':
    args = parser.parse_args()
    predict_bbox_from_json(bbox_idx=args.idx, verbose=args.verbose)


import tensorflow as tf
import hyperchamber as hc
import hypergan as hg
import numpy as np
from hypergan.discriminators.pyramid_discriminator import PyramidDiscriminator
from hypergan.gan_component import ValidationException
from hypergan.ops import TensorflowOps
from tests.mocks import MockDiscriminator, mock_gan

from unittest.mock import MagicMock

from hypergan.trainers.alternating_trainer import AlternatingTrainer

class AlternatingTrainerTest(tf.test.TestCase):
    def test_config(self):
        with self.test_session():
            config = {'d_learn_rate': 1e-3, 'g_learn_rate': 1e-3, 'd_trainer': 'rmsprop', 'g_trainer': 'adam'}
            gan = hg.GAN()
            trainer = AlternatingTrainer(gan, config)
            self.assertEqual(trainer.config.d_learn_rate, 1e-3)

    def test_validate(self):
        with self.assertRaises(ValidationException):
            gan = mock_gan()
            AlternatingTrainer(gan, {})

    def test_clip(self):
        with self.test_session():

            #trainer.create()
            #self.assertEqual(gan.graph.clip, None)
            pass

    def test_output_string(self):
        with self.test_session():
            gan = mock_gan()
            gan.create()
            config = {'d_learn_rate': 1e-3, 'g_learn_rate': 1e-3, 'd_trainer': 'rmsprop', 'g_trainer': 'adam'}
            trainer = AlternatingTrainer(gan, config)
            c = tf.constant(1)
            self.assertTrue('d_loss' in trainer.output_string({'d_loss':c}))
            self.assertTrue('g_loss' in trainer.output_string({'g_loss':c}))
            self.assertEqual(len(trainer.output_variables({'a': c, 'b': c})), 2)
if __name__ == "__main__":
    tf.test.main()


import numpy as np
import torch
import sys

from collections import Counter
from sklearn.preprocessing import LabelEncoder

from librosa.core import load
from librosa.feature import melspectrogram
from librosa import power_to_db
from librosa import amplitude_to_db
from model import genreNet
from config import MODELPATH
from config import GENRES

import librosa.display
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import time

import os

from PySide2 import QtCore, QtGui, QtWidgets
from PySide2.QtWidgets import QApplication, QWidget, QInputDialog, QLineEdit, QFileDialog

import warnings

import threading

fileName = None

warnings.filterwarnings("ignore")

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(400, 400)
        MainWindow.setMaximumSize(400, 400)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.widget = QtWidgets.QWidget(self.centralwidget)
        self.widget.setGeometry(QtCore.QRect(20, 10, 363, 317))
        self.widget.setObjectName("widget")
        self.gridLayout = QtWidgets.QGridLayout(self.widget)
        self.gridLayout.setContentsMargins(0, 0, 0, 0)
        self.gridLayout.setObjectName("gridLayout")
        self.button2 = QtWidgets.QPushButton(self.widget)
        self.button2.setObjectName("button2")
        self.gridLayout.addWidget(self.button2, 4, 0, 1, 1)
        self.label = QtWidgets.QLabel(self.widget)
        font = QtGui.QFont()
        font.setPointSize(8)
        font.setBold(True)
        font.setWeight(75)
        self.label.setFont(font)
        self.label.setObjectName("label")
        self.gridLayout.addWidget(self.label, 0, 0, 1, 1)
        self.button1 = QtWidgets.QPushButton(self.widget)
        self.button1.setObjectName("button1")
        self.gridLayout.addWidget(self.button1, 3, 0, 1, 1)
        self.openfile = QtWidgets.QLabel(self.widget)
        font = QtGui.QFont()
        font.setBold(True)
        font.setWeight(75)
        self.openfile.setFont(font)
        self.openfile.setAlignment(QtCore.Qt.AlignCenter)
        self.openfile.setObjectName("openfile")
        self.gridLayout.addWidget(self.openfile, 1, 0, 1, 1)
        self.text = QtWidgets.QPlainTextEdit(self.widget)
        font = QtGui.QFont()
        font.setPointSize(11)
        font.setBold(True)
        self.text.setFont(font)
        self.text.setObjectName("text")
        # self.text.setReadOnly(True)
        self.gridLayout.addWidget(self.text, 6, 0, 1, 1)
        MainWindow.setCentralWidget(self.centralwidget)
        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.button2.setText(_translate("MainWindow", "Откри жанр"))
        self.label.setText(_translate("MainWindow", "Отвори музичка датотека и откри го нејзиниот музички жанр"))
        self.button1.setText(_translate("MainWindow", "Отвори датотека"))
        self.button1.clicked.connect(self.openFile)
        self.button2.clicked.connect(self.thread)
        self.text.setReadOnly(True)
        self.button2.setEnabled(False)
        self.pbar = QtWidgets.QProgressBar(self.widget)
        self.pbar.setTextVisible(False)
        self.pbar.setObjectName("pbar")
        self.gridLayout.addWidget(self.pbar, 5, 0, 1, 1)
        self.pbar.setMinimum(0)
        self.pbar.setMaximum(100)

    def openFile(self):
        options = QtWidgets.QFileDialog.Options()
        global fileName
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(
            None,
            "Open Music Files",
            os.getcwd(),
            "Music Files (*.mp3);;All Files (*)",
            options=options)
        if len(fileName) != 0:
            self.openfile.setText(os.path.basename(fileName))
            self.text.clear()
            self.button2.setEnabled(True)
        else:
            self.openfile.clear()
            self.button2.setEnabled(False)

    def main(self):
        le = LabelEncoder().fit(GENRES)
        # ------------------------------- #
        ## LOAD TRAINED GENRENET MODEL
        net = genreNet()
        net.load_state_dict(torch.load(MODELPATH, map_location='cpu'))
        # ------------------------------- #
        ## LOAD AUDIO
        audio_path = fileName
        # os.path.basename()
        y, sr = load(audio_path, mono=True, sr=22050)
        # ------------------------------- #
        ## AUDIO SPECTOGRAM
        fspec1 = melspectrogram(y, sr=sr, n_mels=128)
        # plt.subplot(3, 1, 1)
        spec1_amp = amplitude_to_db(fspec1, ref=np.max)
        plt.figure(figsize=(10, 4))
        librosa.display.specshow(spec1_amp, sr=sr, x_axis='time', y_axis='mel')
        plt.colorbar(format='%+2.0f dB')
        plt.title('Mel Power Spectogram')
        plt.tight_layout()
        if os.path.exists('../spectograms/mel_pow_spectogram.png'):
            plt.savefig('../spectograms/mel_pow_spectogram_{}.png'.format(int(time.time())), dpi=300)
        else:
            plt.savefig('../spectograms/mel_pow_spectogram.png', dpi=300)
        # plt.subplot(3, 1, 2)
        plt.figure(figsize=(10, 4))
        librosa.display.specshow(fspec1 ** 2, sr=sr, y_axis='log')
        plt.colorbar()
        plt.title('Power Spectogram')
        plt.tight_layout()
        if os.path.exists('../spectograms/pow_spectogram.png'):
            plt.savefig('../spectograms/pow_spectogram_{}.png'.format(int(time.time())), dpi=300)
        else:
            plt.savefig('../spectograms/pow_spectogram.png', dpi=300)
        # plt.subplot(3, 1, 3)
        plt.figure(figsize=(10, 4))
        librosa.display.specshow(power_to_db(fspec1 ** 2, ref=np.max), sr=sr, y_axis='log', x_axis='time')
        plt.colorbar(format='%+2.0f dB')
        plt.title('Log-Power Spectrogram')
        plt.tight_layout()
        if os.path.exists('../spectograms/log_pow_spectogram.png'):
            plt.savefig('../spectograms/log_pow_spectogram_{}.png'.format(int(time.time())), dpi=300)
        else:
            plt.savefig('../spectograms/log_pow_spectogram.png', dpi=300)
        # ------------------------------- #
        ## GET CHUNKS OF AUDIO SPECTROGRAMS
        S = melspectrogram(y, sr).T
        S = S[:-1 * (S.shape[0] % 128)]
        num_chunk = S.shape[0] / 128
        data_chunks = np.split(S, num_chunk)
        # ------------------------------- #
        ## CLASSIFY SPECTROGRAMS
        genres = list()
        for i, data in enumerate(data_chunks):
            data = torch.FloatTensor(data).view(1, 1, 128, 128)
            preds = net(data)
            pred_val, pred_index = preds.max(1)
            pred_index = pred_index.data.numpy()
            pred_val = np.exp(pred_val.data.numpy()[0])
            pred_genre = le.inverse_transform(pred_index).item()
            if pred_val >= 0.5:
                genres.append(pred_genre)
        # ------------------------------- #
        s = float(sum([v for k, v in dict(Counter(genres)).items()]))
        pos_genre = sorted([(k, v / s * 100) for k, v in dict(Counter(genres)).items()], key=lambda x: x[1],
                           reverse=True)
        for genre, pos in pos_genre:
            print("%10s: \t%.2f\t%%\n" % (genre, pos))
            self.text.insertPlainText("%10s: \t%.2f\t%%\n" % (genre, pos))
            self.button1.setEnabled(True)
            self.button2.setEnabled(True)
            self.pbar.setMinimum(0)
            self.pbar.setMaximum(100)
            #self.pbar.deleteLater()
        return
    def thread(self):
        self.button1.setEnabled(False)
        self.button2.setEnabled(False)
        self.pbar.setMinimum(0)
        self.pbar.setMaximum(0)
        t=threading.Thread(target=self.main, args=())
        t.start()
        return

if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    MainWindow.show()
    sys.exit(app.exec_())


#!/usr/bin/env python

import argparse
import os
import sys
import subprocess
import time
import numpy as np
from data_split import DataSplit
from generator_h5 import ASCDataGenerator, ASCBatchGenerator
from keras.layers import Input
import keras.callbacks
from keras.models import Model
from dcs_models import DCS_CNN4, DCS_CNN8
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score
import tensorflow as tf

ASC_CLASS=10

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, required=True, help='data')
    parser.add_argument('--valid-keys', type=str, required=True, help='validation keys')
    parser.add_argument('--input-dim', default=40, type=int, help='feats dim')
    parser.add_argument('--seq-len', default=500, type=int, help='input sequence length')
    parser.add_argument('--batch-size', default=64, type=int, help='mini-batch size')
    parser.add_argument('--model', type=str, help='weights file')
    parser.add_argument('--snapshot', type=str, default='./', help='snapshot directory')
    parser.add_argument('--snapshot-prefix', type=str, default='snapshot', help='snapshot file prefix')
    parser.add_argument('--filters', type=int, default=64, help='number of filters')
    parser.add_argument('--model-type', type=str, default='cnn4' ,help='model type')
    args = parser.parse_args()

    valid_keys=[]
    with open(args.valid_keys, 'r') as f:
        key=f.readline().strip()
        while key:
            valid_keys.append(key)
            key = f.readline().strip()

    # read best model
    input=Input((args.input_dim, args.seq_len, 1))
    if args.model_type == 'cnn4':
        output1, output2 = DCS_CNN4()(input)
    else:
        output1, output2 = DCS_CNN8()(input)
    model=Model(input, [output1, output2])
    model.load_weights(args.model)

    validation_generator = ASCBatchGenerator(args.data, keys=valid_keys,
                                             dim=(args.input_dim, args.seq_len), batch_size=args.batch_size, mode=False)
    per_epoch=validation_generator.get_sample_per_epoch()

    csv =  os.path.join(args.snapshot,args.snapshot_prefix+'.csv')

    mat=np.zeros((ASC_CLASS,ASC_CLASS), dtype=int)
    for n in range(per_epoch):
        x_batch, y_batch=next(validation_generator.gen_batch())
        pred,output = model.predict(x_batch, x_batch.shape[0], verbose=1)
        y_pred=np.argmax(pred, axis=1)
        y_true=np.argmax(y_batch, axis=1)
        cmat=confusion_matrix(y_true, y_pred)
        mat=np.add(cmat, mat)

    np.savetxt(csv, mat.astype(int))

if __name__ == "__main__":
  main()


import numpy as np
from os import listdir
from os.path import isfile, join
import json
import argparse
from collections import defaultdict
from tqdm import tqdm

class APD:
    """
    A class to load and filter the Aschaffenburg Pose Dataset.

    ...

    Attributes
    ----------
    data : dict
        data of the Aschaffenburg Pose Dataset

    Methods
    -------
    filter:
        Filter the dataset for certain vru_types, sets and data fields.
    """

    NESTED_DATA_FIELDS = ['pose2d', 'pose3d', 'motion_primitives']

    def __init__(self, data_path, vru_types=[], sets=[], data_fields=[], pose3d_joints=[], pose2d_joints=[]):
        """
        Loads the selected data from the json files.

        Parameters
        ----------
            data_path : str
                path to the json files
            vru_types : list of str, optional
                load only the selected vru types ('ped', 'bike'). All vru types are loaded if the list is empty
                (default).
            sets : list of str, optional
                load only the selected set ('train', 'valid', 'test'). All sets are loaded if the list is empty
                (default).
            data_fields : list of str, optional
                load only the selected data fields ('vru_type', 'set', 'timestamps', 'head_smoothed', 'pose2d',
                'pose3d', 'motion_primitives'). All datafields are loaded if the list is empty (default).
            pose3d_joints : list of str, optional
                load only the trajectories of the selected joints of pose3d ('head', 'thorax', 'lshoulder', 'rshoulder',
                'lelbow', 'relbow', 'lwrist', 'rwrist', 'lhip', 'rhip', 'lknee', 'rknee', 'lfoot', 'rfoot').
                All trajectories are loaded if the list is empty (default).
            pose2d_joints : list of str, optional
                load only the trajectories of the selected joints of pose2d ('nose', 'neck', 'rshoulder', 'relbow',
                'rwrist', 'lshoulder', 'lelbow', 'lwrist', 'rhip', 'rknee', 'rankle', 'lhip', 'lknee', 'lankle', 'reye',
                'leye', 'rear', 'lear'). All trajectories are loaded if the list is empty (default).
        """
        self.data = dict()
        data_types = {'pose2d': np.float32, 'pose3d': np.float32, 'motion_primitives': np.uint8,
                      'head_smoothed': np.float32, 'timestamps': np.uint64}

        json_files = [join(data_path, f) for f in listdir(data_path) if
                      isfile(join(data_path, f)) and f.endswith('.json')]

        for f in tqdm(json_files, desc ="Loading Dataset"):
            with open(f) as c_file:
                json_data = json.load(c_file)
                if (json_data['vru_type'] in vru_types or len(vru_types) == 0) and (
                        json_data['set'] in sets or len(sets) == 0):  # filter for vru_type and set
                    for k, v in json_data.items():
                        if k in data_fields or len(data_fields) == 0:
                            if k not in self.data:
                                if k in self.NESTED_DATA_FIELDS:
                                    self.data[k] = defaultdict(list)
                                else:
                                    self.data[k] = list()
                            if k in self.NESTED_DATA_FIELDS:
                                for sub_k in v:
                                    if (k == 'pose2d' and (sub_k in pose2d_joints or len(pose2d_joints) == 0)) or (
                                                k == 'pose3d' and (sub_k in pose3d_joints or len(
                                                pose3d_joints) == 0) or k == 'motion_primitives'):
                                        self.data[k][sub_k].append(np.array(v[sub_k], dtype=data_types[k]))
                            else:
                                if k in ['head_smoothed', 'timestamps']:
                                    v = np.array(v, dtype=data_types[k])
                                self.data[k].append(v)

        # convert remaining lists to np.arrays
        for k, v in self.data.items():
            if k in self.NESTED_DATA_FIELDS:
                for sub_k in v:
                    self.data[k][sub_k] = np.array(self.data[k][sub_k])
            else:
                self.data[k] = np.array(v)

    def filter(self, vru_types=[], sets=[], data_fields=[], pose3d_joints=[], pose2d_joints=[]):
        """
        Filter the dataset for certain vru_types, sets and data fields.

        Parameters
        ----------
            vru_types : list of str, optional
                filter for vru types ('ped', 'bike'). All vru types are selected if the list is empty
                (default).
            sets : list of str, optional
                filter for set ('train', 'valid', 'test'). All sets are selected if the list is empty
                (default).
            data_fields : list of str, optional
                filter for data fields ('vru_type', 'set', 'timestamps', 'head_smoothed', 'pose2d', 'pose3d',
                'motion_primitives'). All datafields are selected if the list is empty (default).
            pose3d_joints : list of str, optional
                filter for joints of pose3d ('head', 'thorax', 'lshoulder', 'rshoulder', 'lelbow', 'relbow', 'lwrist',
                'rwrist', 'lhip', 'rhip', 'lknee', 'rknee', 'lfoot', 'rfoot'). All trajectories are selected if the list
                is empty (default).
            pose2d_joints : list of str, optional
                filter for joints of pose2d ('nose', 'neck', 'rshoulder', 'relbow', 'rwrist', 'lshoulder', 'lelbow',
                'lwrist', 'rhip', 'rknee', 'rankle', 'lhip', 'lknee', 'lankle', 'reye', 'leye', 'rear', 'lear'). All
                trajectories are selected if the list is empty (default).

        Returns
        ----------
            filtered dataset : dict
        """
        mask = np.ones_like(list(self.data.keys())[0], dtype=bool)
        if len(vru_types)!=0:
            mask_vru_type = [True if (t in vru_types) else False for t in self.data['vru_type']]
            mask = mask & mask_vru_type
        if len(sets)!=0:
            mask_sets = [True if (t in sets) else False for t in self.data['set']]
            mask = mask & mask_sets

        filtered_data = dict()
        for k, v in self.data.items():
            if k in data_fields or len(data_fields) == 0:
                if k in self.NESTED_DATA_FIELDS:
                    filtered_data[k] = dict()
                    for sub_k in v:
                        if (k == 'pose2d' and (sub_k in pose2d_joints or len(pose2d_joints) == 0)) or (
                                    k == 'pose3d' and (
                                    sub_k in pose3d_joints or len(pose3d_joints) == 0) or k == 'motion_primitives'):
                            filtered_data[k][sub_k] = self.data[k][sub_k][mask]
                else:
                    filtered_data[k] = self.data[k][mask]

        return filtered_data


import urllib.request
import cv2
import numpy as np

data = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
myFace= []
#capture = cv2.VideoCapture(0)
url = 'http://192.168.137.102:8080/shot.jpg'
while True:
#    ret, frame = capture.read(1)
    imgResp = urllib.request.urlopen(url)
    imgNp = np.array(bytearray(imgResp.read()), dtype=np.uint8)
    frame = cv2.imdecode(imgNp, -1)
    if ret:
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = data.detectMultiScale(gray, 1.3)
        for x,y,w,h in faces:
            cv2.rectangle(frame, (x,y), (x+w,y+h), (255,0,0), 5)
            faceComponents = frame[y:y + h, x:x+w, :]
            fc = cv2.resize(faceComponents,(50,50))

            if len(myFace)<= 500:
                myFace.append(fc)
            print(len(myFace))

        cv2.imshow('result',frame)
        if cv2.waitKey(1) == 27 or len(myFace)>=500:
            break
    else:
        print("Camera Not Working")
myFace = np.asarray(myFace)
np.save('aditya_data.npy',myFace)

import matplotlib.pyplot as plt
myFace.shape

plt.imshow(myFace[15])
plt.show()


import os
import sys

# go to the directory of ai8x
PROJECT_ROOT = os.path.join(os.path.dirname(__file__))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "ai8x-training"))

import cv2
import numpy as np
import torch
import importlib

import ai8x
from YOLO_V1_DataSet import YoloV1DataSet
from map import NMS, gt_std, IOU

# loading dataset
dataset_root = "/data/yiwei/VOC2007"
dataSet = YoloV1DataSet(imgs_dir=f"{dataset_root}/Test/JPEGImages",
                        annotations_dir=f"{dataset_root}/Test/Annotations",
                        ClassesFile=f"{dataset_root}/VOC_remain_class.data",
                        train_root=f"{dataset_root}/Test/ImageSets/Main/",
                        )

# loading model
ai8x.set_device(85, simulate=False, round_avg=False, verbose=True)
mod = importlib.import_module("yolov1_bn_model_noaffine")
Yolo = mod.Yolov1_net(num_classes=dataSet.Classes, bias=True)
qat_policy = {'start_epoch': 150,
              'weight_bits': 8,
              'bias_bits': 8,
              'shift_quantile': 0.99}
ai8x.fuse_bn_layers(Yolo)
ai8x.initiate_qat(Yolo, qat_policy)

checkpoint_fname = './log/QAT-20220210-174132/scaled224_noaffine_shift0.99_maxim_yolo_qat_ep200.pth'
Yolo.load_state_dict(torch.load(checkpoint_fname, map_location=lambda storage, loc: storage)) # batch_size 16

# ck_fname = checkpoint_fname.split("/")[-1]
# checkpoint_dir = checkpoint_fname.replace(ck_fname, "")
# print(f"checkpoint_dir {checkpoint_dir}")
# import distiller.apputils as apputils
# apputils.save_checkpoint(checkpoint_dir, "ai85net5", Yolo,
#                             optimizer=None, scheduler=None, extras=None,
#                             is_best=False, name="Yolov1", dir=checkpoint_dir,
#                          )

if __name__ == '__main__':
    for dir_suffix, img_list in [
        ("predictions", [0, 14, 21, 54, 67, 90, 107, 118, 119, 124, 178]),
        ("bad_predictions", [7, 9, 28, 48, 58, 61, 94, 108, 109, 111, 125, 170]),
    ]:
        dir_name = os.path.join(os.path.dirname(checkpoint_fname), dir_suffix)
        os.makedirs(dir_name, exist_ok=True)
        for img_idx in img_list:
            img_data = dataSet.read_img(item=img_idx)
            img_name = os.path.basename(dataSet.img_path[img_idx])
            train_data, ground_truth = dataSet[img_idx]

            ground_truth = torch.unsqueeze(ground_truth, 0)
            ground_truth = ground_truth.squeeze(dim=3)
            ground_truth = gt_std(ground_truth)[0]

            # draw ground truth
            for box in ground_truth:
                box = np.array(box[0:4]).astype(np.int)
                img_data = cv2.rectangle(img_data, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
                img_data = cv2.putText(img_data, "ground truth", (box[0], box[3] + 15),
                                       cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 2)

            # predict
            train_data = torch.unsqueeze(train_data, 0)
            print(f'Predicting {img_name}')
            pred_results, _ = Yolo(train_data)
            NMS_boxes = NMS(bounding_boxes=pred_results)[0]
            for i, box in enumerate(NMS_boxes):
                for j, gt_box in enumerate(ground_truth):
                    iou = IOU(gt_box, box)
                    print(f'IOU between {i} pred_box with {j} ground_truth_box: {iou}')
                has_obj_prob = box[4]
                class_index = box[-1]
                # convert box from float to int
                box = np.array(box[0:4]).astype(np.int)
                # draw predicted box
                img_data = cv2.rectangle(img_data, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
                img_data = cv2.putText(img_data, "prob:{:.2f}".format(has_obj_prob), (box[0], box[1] - 4),
                                              cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)

            # save result
            pred_img = os.path.join(dir_name, f"{img_name.split('.')[0]}_pred.jpg")
            cv2.imwrite(pred_img, img_data)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import time
import argparse

import torch
import torchvision
import numpy as np
from torch.utils.data import DataLoader
from dataset.datasets import WLFWDatasets
from models.pfld import PFLDbackbone, AuxiliaryNet, PFLDLoss

def validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet):

    plfd_backbone.eval()
    auxiliarynet.eval()

    with torch.no_grad():
        losses, losses_ION = [], []
        for idx, (img, landmark_gt, attribute_gt, euler_angle_gt) in enumerate(wlfw_val_dataloader):

            img.requires_grad = False
            img = img.cuda(non_blocking=True)

            attribute_gt.requires_grad = False
            attribute_gt = attribute_gt.cuda(non_blocking=True)

            landmark_gt.requires_grad = False
            landmark_gt = landmark_gt.cuda(non_blocking=True)

            euler_angle_gt.requires_grad = False
            euler_angle_gt = euler_angle_gt.cuda(non_blocking=True)

            plfd_backbone = plfd_backbone.cuda()
            auxiliarynet = auxiliarynet.cuda()

            _, landmarks = plfd_backbone(img)

            loss = torch.mean(torch.sqrt(torch.sum((landmark_gt - landmarks)**2, axis=1)))

            landmarks = landmarks.cpu().numpy()
            landmarks = landmarks.reshape(landmarks.shape[0], -1, 2)
            landmark_gt = landmark_gt.reshape(landmark_gt.shape[0], -1, 2).cpu().numpy()
            error_diff = np.sum(np.sqrt(np.sum((landmark_gt - landmarks) ** 2, axis=2)), axis=1)
            interocular_distance = np.sqrt(np.sum((landmarks[:, 60, :] - landmarks[:,72, :]) ** 2, axis=1))
            # interpupil_distance = np.sqrt(np.sum((landmarks[:, 60, :] - landmarks[:, 72, :]) ** 2, axis=1))
            error_norm = np.mean(error_diff / interocular_distance)

            # show result
            # show_img = np.array(np.transpose(img[0].cpu().numpy(), (1, 2, 0)))
            # show_img = (show_img * 256).astype(np.uint8)
            # np.clip(show_img, 0, 255)

            # pre_landmark = landmarks[0] * [112, 112]

            # cv2.imwrite("xxx.jpg", show_img)
            # img_clone = cv2.imread("xxx.jpg")

            # for (x, y) in pre_landmark.astype(np.int32):
            #     print("x:{0:}, y:{1:}".format(x, y))
            #     cv2.circle(img_clone, (x, y), 1, (255,0,0),-1)
            # cv2.imshow("xx.jpg", img_clone)
            # cv2.waitKey(0)

        losses.append(loss.cpu().numpy())
        losses_ION.append(error_norm)

        print("NME", np.mean(losses))
        print("ION", np.mean(losses_ION))

def main(args):

    checkpoint = torch.load(args.model_path)

    plfd_backbone = PFLDbackbone().cuda()
    auxiliarynet = AuxiliaryNet().cuda()

    plfd_backbone.load_state_dict(checkpoint['plfd_backbone'])
    auxiliarynet.load_state_dict(checkpoint['auxiliarynet'])

    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
    wlfw_val_dataloader = DataLoader(WLFWDatasets(args.test_dataset, transform), \
                                     batch_size=args.batch_size, shuffle=False, \
                                     num_workers=args.workers)

    validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet)

def parse_args():
    parser = argparse.ArgumentParser(description='Testing')
    # env
    parser.add_argument('--use_gpu',   type=bool, default=False)
    parser.add_argument('--gpu_ids',   type=list, default=[0, 1])
    parser.add_argument('--workers',   type=int,  default=8)
    parser.add_argument('--batch_size',type=int,  default=8)

    parser.add_argument('--model_path',   type=str, default="./checkpoint/checkpoint.pth.tar" )
    parser.add_argument('--test_dataset', type=str, default='./data/test_data/list.txt')

    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    main(args)


import numpy as np
from numpy.random import choice

import torch
from torch import nn as nn

from rlkit.policies.base import ExplorationPolicy, Policy
from rlkit.torch.distributions import ReparamTanhMultivariateNormal
from rlkit.torch.distributions import ReparamMultivariateNormalDiag
from rlkit.torch.networks import Mlp, Ant2DCustomLayerV1
from rlkit.torch.core import PyTorchModule

import rlkit.torch.pytorch_util as ptu

LOG_SIG_MAX = 2
LOG_SIG_MIN = -20

class MakeDeterministic(Policy):
    def __init__(self, stochastic_policy):
        self.stochastic_policy = stochastic_policy

    def get_action(self, observation):
        return self.stochastic_policy.get_action(observation,
                                                 deterministic=True)

    def get_actions(self, observations):
        return self.stochastic_policy.get_actions(observations,
                                                  deterministic=True)

    def train(self, mode):
        pass

    def set_num_steps_total(self, num):
        pass

    def to(self, device):
        self.stochastic_policy.to(device)

class DiscretePolicy(Mlp, ExplorationPolicy):
    """
    Usage:

    ```
    policy = DiscretePolicy(...)
    action, log_prob = policy(obs, return_log_prob=True)
    ```
    """
    def __init__(
            self,
            hidden_sizes,
            obs_dim,
            action_dim,
            init_w=1e-3,
            **kwargs
    ):
        self.save_init_params(locals())
        super().__init__(
            hidden_sizes,
            input_size=obs_dim,
            output_size=action_dim,
            init_w=init_w,
            output_activation=nn.LogSoftmax(1),
            **kwargs
        )

    def get_action(self, obs_np, deterministic=False):
        action = self.get_actions(obs_np[None], deterministic=deterministic)
        return action, {}

    def get_actions(self, obs_np, deterministic=False):
        return self.eval_np(obs_np)[0]

    def forward(
            self,
            obs,
            deterministic=False,
            return_log_prob=False,
    ):
        log_probs, pre_act = super().forward(obs, return_preactivations=True)

        if deterministic:
            log_prob, idx = torch.max(log_probs, 1)
            return (idx, None)
        else:
            # Using Gumbel-Max trick to sample from the multinomials
            u = torch.rand(pre_act.size(), requires_grad=False)
            gumbel = -torch.log(-torch.log(u))
            _, idx = torch.max(gumbel + pre_act, 1)

            idx = torch.unsqueeze(idx, 1)
            log_prob = torch.gather(log_probs, 1, idx)

            # # print(log_probs.size(-1))
            # # print(log_probs.data.numpy())
            # # print(np.exp(log_probs.data.numpy()))
            # idx = choice(
            #     log_probs.size(-1),
            #     size=1,
            #     p=np.exp(log_probs.data.numpy())
            # )
            # log_prob = log_probs[0,idx]

            # print(idx)
            # print(log_prob)

            return (idx, log_prob)

    def get_log_pis(self, obs):
        return super().forward(obs)

class MlpPolicy(Mlp, ExplorationPolicy):
    def __init__(
        self,
        hidden_sizes,
        obs_dim,
        action_dim,
        init_w=1e-3,
        **kwargs
    ):
        self.save_init_params(locals())
        super().__init__(
            hidden_sizes,
            input_size=obs_dim,
            output_size=action_dim,
            init_w=init_w,
            **kwargs
        )

    def get_action(self, obs_np, deterministic=False):
        '''
        deterministic=False makes no diff, just doing this for
        consistency in interface for now
        '''
        actions = self.get_actions(obs_np[None], deterministic=deterministic)
        actions = actions[None]
        return actions[0, :], {}

    def get_actions(self, obs_np, deterministic=False):
        return self.eval_np(obs_np)[0]

class ReparamTanhMultivariateGaussianPolicy(Mlp, ExplorationPolicy):
    """
    Usage:

    ```
    policy = ReparamTanhMultivariateGaussianPolicy(...)
    action, mean, log_std, _ = policy(obs)
    action, mean, log_std, _ = policy(obs, deterministic=True)
    action, mean, log_std, log_prob = policy(obs, return_log_prob=True)
    ```

    Here, mean and log_std are the mean and log_std of the Gaussian that is
    sampled from.

    If deterministic is True, action = tanh(mean).
    If return_log_prob is False (default), log_prob = None
        This is done because computing the log_prob can be a bit expensive.
    """
    def __init__(
            self,
            hidden_sizes,
            obs_dim,
            action_dim,
            std=None,
            init_w=1e-3,
            **kwargs
    ):
        self.save_init_params(locals())
        super().__init__(
            hidden_sizes,
            input_size=obs_dim,
            output_size=action_dim,
            init_w=init_w,
            **kwargs
        )
        self.log_std = None
        self.std = std
        if std is None:
            last_hidden_size = obs_dim
            if len(hidden_sizes) > 0:
                last_hidden_size = hidden_sizes[-1]
            self.last_fc_log_std = nn.Linear(last_hidden_size, action_dim)
            self.last_fc_log_std.weight.data.uniform_(-init_w, init_w)
            self.last_fc_log_std.bias.data.uniform_(-init_w, init_w)
        else:
            self.log_std = np.log(std)
            assert LOG_SIG_MIN <= self.log_std <= LOG_SIG_MAX

    def get_action(self, obs_np, deterministic=False):
        actions = self.get_actions(obs_np[None], deterministic=deterministic)
        return actions[0, :], {}

    def get_actions(self, obs_np, deterministic=False):
        return self.eval_np(obs_np, deterministic=deterministic)[0]

    def forward(
            self,
            obs,
            deterministic=False,
            return_log_prob=False,
            return_tanh_normal=False
    ):
        """
        :param obs: Observation
        :param deterministic: If True, do not sample
        :param return_log_prob: If True, return a sample and its log probability
        """
        # print('forward')
        # print(obs.shape)
        h = obs
        for i, fc in enumerate(self.fcs):
            h = self.hidden_activation(fc(h))
        # print('fuck')
        # print(h)
        mean = self.last_fc(h)
        # print(mean)
        if self.std is None:
            log_std = self.last_fc_log_std(h)
            log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
            std = torch.exp(log_std)
        else:
            std = self.std
            log_std = self.log_std

        log_prob = None
        expected_log_prob = None
        mean_action_log_prob = None
        pre_tanh_value = None
        if deterministic:
            action = torch.tanh(mean)
        else:
            tanh_normal = ReparamTanhMultivariateNormal(mean, log_std)
            # print('mean, std')
            # print(mean)
            # print(log_std)
            if return_log_prob:
                action, pre_tanh_value = tanh_normal.sample(
                    return_pretanh_value=True
                )
                log_prob = tanh_normal.log_prob(
                    action,
                    pre_tanh_value=pre_tanh_value
                )
            else:
                action = tanh_normal.sample()

        # I'm doing it like this for now for backwards compatibility, sorry!
        if return_tanh_normal:
            return (
                action, mean, log_std, log_prob, expected_log_prob, std,
                mean_action_log_prob, pre_tanh_value, tanh_normal,
            )
        return (
            action, mean, log_std, log_prob, expected_log_prob, std,
            mean_action_log_prob, pre_tanh_value,
        )

    def get_log_prob(self, obs, acts, return_normal_params=False):
        h = obs
        for i, fc in enumerate(self.fcs):
            h = self.hidden_activation(fc(h))
        mean = self.last_fc(h)
        if self.std is None:
            log_std = self.last_fc_log_std(h)
            log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
            std = torch.exp(log_std)
        else:
            std = self.std
            log_std = self.log_std

        tanh_normal = ReparamTanhMultivariateNormal(mean, log_std)
        log_prob = tanh_normal.log_prob(acts)

        # print('\n\n\n\n\nGet log prob')
        # print(log_prob)
        # print(mean)
        # print(log_std)

        if return_normal_params:
            return log_prob, mean, log_std
        return log_prob

class ReparamMultivariateGaussianPolicy(Mlp, ExplorationPolicy):
    def __init__(
            self,
            hidden_sizes,
            obs_dim,
            action_dim,
            std=None,
            init_w=1e-3,
            **kwargs
    ):
        self.save_init_params(locals())
        super().__init__(
            hidden_sizes,
            input_size=obs_dim,
            output_size=action_dim,
            init_w=init_w,
            **kwargs
        )
        self.log_std = None
        self.std = std
        if std is None:
            last_hidden_size = obs_dim
            if len(hidden_sizes) > 0:
                last_hidden_size = hidden_sizes[-1]
            self.last_fc_log_std = nn.Linear(last_hidden_size, action_dim)
            self.last_fc_log_std.weight.data.uniform_(-init_w, init_w)
            self.last_fc_log_std.bias.data.uniform_(-init_w, init_w)
        else:
            assert LOG_SIG_MIN <= np.log(std) <= LOG_SIG_MAX
            std = std*np.ones((1,action_dim))
            self.log_std = ptu.from_numpy(np.log(std), requires_grad=False)

    def get_action(self, obs_np, deterministic=False):
        actions = self.get_actions(obs_np[None], deterministic=deterministic)
        return actions[0, :], {}

    def get_actions(self, obs_np, deterministic=False):
        return self.eval_np(obs_np, deterministic=deterministic)[0]

    def forward(
            self,
            obs,
            deterministic=False,
            return_log_prob=False,
            return_normal=False
    ):
        """
        :param obs: Observation
        :param deterministic: If True, do not sample
        :param return_log_prob: If True, return a sample and its log probability
        """
        # print('forward')
        # print(obs)
        h = obs
        for i, fc in enumerate(self.fcs):
            h = self.hidden_activation(fc(h))
        # print('fuck')
        # print(h)
        mean = self.last_fc(h)

        mean = torch.clamp(mean, min=-1.0, max=1.0)

        # print(mean)
        if self.std is None:
            log_std = self.last_fc_log_std(h)

            # log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
            log_std = torch.clamp(log_std, LOG_SIG_MIN, np.log(0.2))

            std = torch.exp(log_std)

        else:
            std = self.std
            log_std = self.log_std

        log_prob = None
        expected_log_prob = None
        mean_action_log_prob = None
        if deterministic:
            action = mean
        else:
            normal = ReparamMultivariateNormalDiag(mean, log_std)
            action = normal.sample()
            if return_log_prob:
                log_prob = normal.log_prob(action)

        # I'm doing it like this for now for backwards compatibility, sorry!
        if return_normal:
            return (
                action, mean, log_std, log_prob, expected_log_prob, std,
                mean_action_log_prob, normal
            )
        return (
            action, mean, log_std, log_prob, expected_log_prob, std,
            mean_action_log_prob
        )

    def get_log_prob(self, obs, acts, return_normal_params=False):
        h = obs
        for i, fc in enumerate(self.fcs):
            h = self.hidden_activation(fc(h))
        mean = self.last_fc(h)
        if self.std is None:
            log_std = self.last_fc_log_std(h)
            log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
            std = torch.exp(log_std)
        else:
            std = self.std
            log_std = self.log_std

        normal = ReparamMultivariateNormalDiag(mean, log_std)
        log_prob = normal.log_prob(acts)

        if return_normal_params:
            return log_prob, mean, log_std
        return log_prob

class AntRandGoalCustomReparamTanhMultivariateGaussianPolicy(ReparamTanhMultivariateGaussianPolicy):
    """
    Custom for Ant Rand Goal
    The only difference is that it linearly embeds the goal into a higher dimension
    """
    def __init__(
            self,
            goal_dim,
            goal_embed_dim,
            hidden_sizes,
            obs_dim,
            action_dim,
            std=None,
            init_w=1e-3,
            **kwargs
    ):
        self.save_init_params(locals())
        super().__init__(
            hidden_sizes,
            obs_dim + goal_embed_dim,
            action_dim,
            init_w=init_w,
            **kwargs
        )

        self.goal_embed_fc = nn.Linear(goal_dim, goal_embed_dim)
        self.goal_dim = goal_dim

    def forward(
        self,
        obs,
        deterministic=False,
        return_log_prob=False,
        return_tanh_normal=False
    ):
        # obs will actuall be concat of [obs, goal]
        goal = obs[:,-self.goal_dim:]
        goal_embed = self.goal_embed_fc(goal)
        obs = torch.cat([obs[:,:-self.goal_dim], goal_embed], dim=-1)
        return super().forward(
            obs,
            deterministic=deterministic,
            return_log_prob=return_log_prob,
            return_tanh_normal=return_tanh_normal
        )

class AntCustomGatingV1ReparamTanhMultivariateGaussianPolicy(PyTorchModule, ExplorationPolicy):
    def __init__(
            self,
    ):
        self.save_init_params(locals())
        super().__init__()

        obs_dim = 113
        hid_dim = 256
        extra_dim = 32
        action_dim = 8

        self.first_hid_layer = nn.Sequential(
            nn.Linear(obs_dim, hid_dim),
            nn.ReLU()
        )
        self.first_extra_layer = nn.Sequential(
            nn.Linear(2, extra_dim),
            nn.ReLU()
        )

        self.mod_list = nn.ModuleList(
            [
                Ant2DCustomLayerV1(extra_dim, hid_dim),
                Ant2DCustomLayerV1(extra_dim, hid_dim)
            ]
        )

        self.last_fc_mean = nn.Linear(hid_dim + extra_dim, action_dim)
        self.last_fc_log_std = nn.Linear(hid_dim + extra_dim, action_dim)

    def get_action(self, obs_np, deterministic=False):
        actions = self.get_actions(obs_np[None], deterministic=deterministic)
        return actions[0, :], {}

    def get_actions(self, obs_np, deterministic=False):
        return self.eval_np(obs_np, deterministic=deterministic)[0]

    def forward(
            self,
            obs,
            deterministic=False,
            return_log_prob=False,
            return_tanh_normal=False
    ):
        goal = obs[:,-2:]
        obs = obs[:,:-2]

        extra_info = self.first_extra_layer(goal)
        hid = self.first_hid_layer(obs)

        for h_mod in self.mod_list:
            extra_info, hid = h_mod(extra_info, hid)

        hid_extra_concat = torch.cat([hid, extra_info], dim=-1)
        mean = self.last_fc_mean(hid_extra_concat)
        log_std = self.last_fc_log_std(hid_extra_concat)
        log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
        std = torch.exp(log_std)

        log_prob = None
        expected_log_prob = None
        mean_action_log_prob = None
        pre_tanh_value = None
        if deterministic:
            action = torch.tanh(mean)
        else:
            tanh_normal = ReparamTanhMultivariateNormal(mean, log_std)
            if return_log_prob:
                action, pre_tanh_value = tanh_normal.sample(
                    return_pretanh_value=True
                )
                log_prob = tanh_normal.log_prob(
                    action,
                    pre_tanh_value=pre_tanh_value
                )
            else:
                action = tanh_normal.sample()

        # I'm doing it like this for now for backwards compatibility, sorry!
        if return_tanh_normal:
            return (
                action, mean, log_std, log_prob, expected_log_prob, std,
                mean_action_log_prob, pre_tanh_value, tanh_normal,
            )
        return (
            action, mean, log_std, log_prob, expected_log_prob, std,
            mean_action_log_prob, pre_tanh_value,
        )

    def get_log_prob(self, obs, acts, return_normal_params=False):
        goal = obs[:,-2:]
        obs = obs[:,:-2]

        extra_info = self.first_extra_layer(goal)
        hid = self.first_hid_layer(obs)

        for h_mod in self.mod_list:
            extra_info, hid = h_mod(extra_info, hid)

        hid_extra_concat = torch.cat([hid, extra_info], dim=-1)
        mean = self.last_fc_mean(hid_extra_concat)
        log_std = self.last_fc_log_std(hid_extra_concat)
        log_std = torch.clamp(log_std, LOG_SIG_MIN, LOG_SIG_MAX)
        std = torch.exp(log_std)

        tanh_normal = ReparamTanhMultivariateNormal(mean, log_std)
        log_prob = tanh_normal.log_prob(acts)

        if return_normal_params:
            return log_prob, mean, log_std
        return log_prob

class PostCondMLPPolicyWrapper(ExplorationPolicy):
    def __init__(self, policy, np_post_sample, deterministic=False, obs_mean=None, obs_std=None):
        super().__init__()
        self.policy = policy
        self.np_z = np_post_sample # assuming it is a flat np array
        self.deterministic = deterministic
        self.obs_mean = obs_mean
        self.obs_std = obs_std
        if obs_mean is not None:
            self.normalize_obs = True
        else:
            self.normalize_obs = False

    def get_action(self, obs_np):
        if self.normalize_obs:
            obs_np = (obs_np - self.obs_mean) / self.obs_std
        obs = np.concatenate((obs_np, self.np_z), axis=0)
        return self.policy.get_action(obs, deterministic=self.deterministic)

    def cuda(self):
        self.policy.cuda(ptu.device)

    def cpu(self):
        self.policy.cpu()

class ObsPreprocessedReparamTanhMultivariateGaussianPolicy(ReparamTanhMultivariateGaussianPolicy):
    '''
        This is a weird thing and I didn't know what to call.
        Basically I wanted this so that if you need to preprocess
        your inputs somehow (attention, gating, etc.) with an external module
        before passing to the policy you could do so.
        Assumption is that you do not want to update the parameters of the preprocessing
        module so its output is always detached.
    '''
    def __init__(self, preprocess_model, *args, **kwargs):
        self.save_init_params(locals())
        super().__init__(*args, **kwargs)
        # this is a hack so that it is not added as a submodule
        self.preprocess_model_list = [preprocess_model]

    @property
    def preprocess_model(self):
        # this is a hack so that it is not added as a submodule
        return self.preprocess_model_list[0]

    def preprocess_fn(self, obs_batch):
        mode = self.preprocess_model.training
        self.preprocess_model.eval()
        processed_obs_batch = self.preprocess_model(obs_batch, False).detach()
        self.preprocess_model.train(mode)
        return processed_obs_batch

    def forward(
        self,
        obs,
        deterministic=False,
        return_log_prob=False,
        return_tanh_normal=False
    ):
        obs = self.preprocess_fn(obs).detach()
        return super().forward(
            obs,
            deterministic=deterministic,
            return_log_prob=return_log_prob,
            return_tanh_normal=return_tanh_normal
        )

    def get_log_prob(self, obs, acts):
        obs = self.preprocess_fn(obs).detach()
        return super().get_log_prob(obs, acts)

class WithZObsPreprocessedReparamTanhMultivariateGaussianPolicy(ReparamTanhMultivariateGaussianPolicy):
    '''
        This is a weird thing and I didn't know what to call.
        Basically I wanted this so that if you need to preprocess
        your inputs somehow (attention, gating, etc.) with an external module
        before passing to the policy you could do so.
        Assumption is that you do not want to update the parameters of the preprocessing
        module so its output is always detached.
    '''
    def __init__(self, preprocess_model, z_dim, *args, train_preprocess_model=False, **kwargs):
        self.save_init_params(locals())
        super().__init__(*args, **kwargs)
        # this is a hack so that it is not added as a submodule
        if train_preprocess_model:
            self._preprocess_model = preprocess_model
        else:
            # this is a hack so that it is not added as a submodule
            self.preprocess_model_list = [preprocess_model]
        self.z_dim = z_dim
        self.train_preprocess_model = train_preprocess_model

    @property
    def preprocess_model(self):
        if self.train_preprocess_model:
            return self._preprocess_model
        else:
            # this is a hack so that it is not added as a submodule
            return self.preprocess_model_list[0]

    def preprocess_fn(self, obs_batch):
        if self.train_preprocess_model:
            processed_obs_batch = self.preprocess_model(
                obs_batch[:,:-self.z_dim],
                False,
                obs_batch[:,-self.z_dim:]
            )
        else:
            mode = self.preprocess_model.training
            self.preprocess_model.eval()
            processed_obs_batch = self.preprocess_model(
                obs_batch[:,:-self.z_dim],
                False,
                obs_batch[:,-self.z_dim:]
            ).detach()
            self.preprocess_model.train(mode)
        return processed_obs_batch

    def forward(
        self,
        obs,
        deterministic=False,
        return_log_prob=False,
        return_tanh_normal=False
    ):
        if self.train_preprocess_model:
            obs = self.preprocess_fn(obs)
        else:
            obs = self.preprocess_fn(obs).detach()
        return super().forward(
            obs,
            deterministic=deterministic,
            return_log_prob=return_log_prob,
            return_tanh_normal=return_tanh_normal
        )

    def get_log_prob(self, obs, acts):
        if self.train_preprocess_model:
            obs = self.preprocess_fn(obs)
        else:
            obs = self.preprocess_fn(obs).detach()
        return super().get_log_prob(obs, acts)

# class PostCondReparamTanhMultivariateGaussianPolicy(ReparamTanhMultivariateGaussianPolicy):
#     '''
#         This is a very simple version of a policy that conditions on a sample from the posterior
#         I just concatenate the latent to the obs, so for now assuming everyting is flat
#     '''
#     def __init__(self, *args, **kwargs):
#         super().__init__(*args, **kwargs)
#         self.z = None

#     def set_post_sample(self, z):
#         self.z = z

#     def forward(
#         self,
#         obs,
#         deterministic=False,
#         return_log_prob=False,
#         return_tanh_normal=False
#     ):
#         obs = torch.cat([obs, self.z], dim=-1)
#         return super().forward(
#             obs,
#             deterministic=deterministic,
#             return_log_prob=return_log_prob,
#             return_tanh_normal=return_tanh_normal
#         )

#     def get_log_prob(self, obs, acts):
#         obs = torch.cat([obs, self.z], dim=-1)
#         return super().get_log_prob(obs, acts)

# class PusherTaskReparamTanhMultivariateGaussianPolicy(ReparamTanhMultivariateGaussianPolicy):
class PusherTaskReparamTanhMultivariateGaussianPolicy(MlpPolicy):
    '''
    This is the policy used for the pusher task which has to push various
    meshes to the center of the table.
    Input is image and optionally robot state
    '''
    def __init__(self, image_processor, image_only=False, train_img_processor=False, *args, **kwargs):
        self.save_init_params(locals())
        super().__init__(*args, **kwargs)

        # We plan to let the discriminator train the image processor
        # and let the policy just learn to adapt itself to this module.
        # To prevent the module from being trained with the policy, we
        # use this hack.
        self.image_only = image_only
        self.train_img_processor = train_img_processor
        if self.train_img_processor:
            self._image_processor = image_processor
        else:
            self._image_processor = [image_processor]

    @property
    def image_processor(self):
        if self.train_img_processor:
            return self._image_processor
        else:
            return self._image_processor[0]

    def process_image(self, z, image):
        if self.train_img_processor:
            feature_positions = self.image_processor(z, image)
        else:
            mode = self.image_processor.training
            self.image_processor.eval()
            feature_positions = self.image_processor(z, image).detach()
            self.image_processer.train(mode)
        return feature_positions

    def forward(
        self,
        obs,
        # deterministic=False,
        # return_log_prob=False,
        # return_tanh_normal=False
    ):
        image = obs['image']
        z = obs['z']
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = obs['X']
            policy_input = torch.cat([policy_input, state], dim=-1)

        return super().forward(
            policy_input,
            # deterministic=deterministic,
            # return_log_prob=return_log_prob,
            # return_tanh_normal=return_tanh_normal
        )

    def get_action(self, obs_np, deterministic=False):
        # actions = self.get_actions(obs_np, deterministic=deterministic)
        # return actions[0, :], {}
        image = obs_np['image']
        image = image
        image = np.ascontiguousarray(image)[None]
        image = ptu.from_numpy(image, requires_grad=False)
        z = ptu.from_numpy(obs_np['z'][None], requires_grad=False)
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = ptu.from_numpy(obs_np['X'][None], requires_grad=False)
            policy_input = torch.cat([policy_input, state], dim=-1)
        # action = super().forward(
        #     policy_input,
        #     deterministic=deterministic,
        #     return_log_prob=False,
        #     return_tanh_normal=False
        # )[0][0,:]
        action = super().forward(policy_input)
        action = ptu.get_numpy(action)
        return action, {}

    # def get_log_prob(self, obs, acts):
    #     image = obs['image']
    #     z = obs['z']
    #     if self.train_img_processor:
    #         policy_input = self.process_image(z, image)
    #     else:
    #         policy_input = self.process_image(z, image).detach()
    #     if not self.image_only:
    #         state = obs['X']
    #         policy_input = torch.cat([policy_input, state], dim=-1)
    #     return super().get_log_prob(policy_input, acts)

# class PusherTaskReparamMultivariateGaussianPolicy(ReparamMultivariateGaussianPolicy):
#     '''
#     This is the policy used for the pusher task which has to push various
#     meshes to the center of the table.
#     Input is image and optionally robot state
#     '''
#     def __init__(self, image_processor, image_only=False, train_img_processor=False, *args, **kwargs):
#         self.save_init_params(locals())
#         super().__init__(*args, **kwargs)

#         # We plan to let the discriminator train the image processor
#         # and let the policy just learn to adapt itself to this module.
#         # To prevent the module from being trained with the policy, we
#         # use this hack.
#         self.image_only = image_only
#         self.train_img_processor = train_img_processor
#         if self.train_img_processor:
#             self._image_processor = image_processor
#         else:
#             self._image_processor = [image_processor]

#     @property
#     def image_processor(self):
#         if self.train_img_processor:
#             return self._image_processor
#         else:
#             return self._image_processor[0]

#     def process_image(self, z, image):
#         if self.train_img_processor:
#             feature_positions = self.image_processor(z, image)
#         else:
#             mode = self.image_processor.training
#             self.image_processor.eval()
#             feature_positions = self.image_processor(z, image).detach()
#             self.image_processer.train(mode)
#         return feature_positions

#     def forward(
#         self,
#         obs,
#         deterministic=False,
#         return_log_prob=False,
#         return_normal=False
#     ):
#         image = obs['image']
#         z = obs['z']
#         if self.train_img_processor:
#             policy_input = self.process_image(z, image)
#         else:
#             policy_input = self.process_image(z, image).detach()
#         if not self.image_only:
#             state = obs['X']
#             policy_input = torch.cat([policy_input, state], dim=-1)

#         return super().forward(
#             policy_input,
#             deterministic=deterministic,
#             return_log_prob=return_log_prob,
#             return_normal=return_normal
#         )

#     def get_action(self, obs_np, deterministic=False):
#         # actions = self.get_actions(obs_np, deterministic=deterministic)
#         # return actions[0, :], {}
#         image = obs_np['image']
#         image = image
#         image = np.ascontiguousarray(image)[None]
#         image = Variable(ptu.from_numpy(image), requires_grad=False)
#         z = Variable(ptu.from_numpy(obs_np['z'][None]), requires_grad=False)
#         if self.train_img_processor:
#             policy_input = self.process_image(z, image)
#         else:
#             policy_input = self.process_image(z, image).detach()
#         if not self.image_only:
#             state = Variable(ptu.from_numpy(obs_np['X'][None]), requires_grad=False)
#             policy_input = torch.cat([policy_input, state], dim=-1)
#         action = super().forward(
#             policy_input,
#             deterministic=deterministic,
#             return_log_prob=False,
#             return_normal=False
#         )[0][0,:]
#         # action = super().forward(policy_input)
#         action = ptu.get_numpy(action)
#         return action, {}

#     def get_log_prob(self, obs, acts):
#         image = obs['image']
#         z = obs['z']
#         if self.train_img_processor:
#             policy_input = self.process_image(z, image)
#         else:
#             policy_input = self.process_image(z, image).detach()
#         if not self.image_only:
#             state = obs['X']
#             policy_input = torch.cat([policy_input, state], dim=-1)
#         return super().get_log_prob(policy_input, acts)

class PusherTaskReparamMultivariateGaussianPolicy(ReparamMultivariateGaussianPolicy):
    '''
    This is the policy used for the pusher task which has to push various
    meshes to the center of the table.
    Input is image and optionally robot state
    '''
    def __init__(self, image_processor, image_only=False, train_img_processor=False, *args, **kwargs):
        self.save_init_params(locals())
        super().__init__(*args, **kwargs)

        # We plan to let the discriminator train the image processor
        # and let the policy just learn to adapt itself to this module.
        # To prevent the module from being trained with the policy, we
        # use this hack.
        self.image_only = image_only
        self.train_img_processor = train_img_processor
        if self.train_img_processor:
            self._image_processor = image_processor
        else:
            self._image_processor = [image_processor]

    @property
    def image_processor(self):
        if self.train_img_processor:
            return self._image_processor
        else:
            return self._image_processor[0]

    def process_image(self, z, image):
        if self.train_img_processor:
            feature_positions = self.image_processor(z, image)
        else:
            mode = self.image_processor.training
            self.image_processor.eval()
            feature_positions = self.image_processor(z, image).detach()
            self.image_processer.train(mode)
        return feature_positions

    def forward(
        self,
        obs,
        deterministic=False,
        return_log_prob=False,
        return_normal=False
    ):
        image = obs['image']
        z = obs['z']
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = obs['X']
            policy_input = torch.cat([policy_input, state], dim=-1)

        return super().forward(
            policy_input,
            deterministic=deterministic,
            return_log_prob=return_log_prob,
            return_normal=return_normal
        )

    def get_action(self, obs_np, deterministic=False):
        # actions = self.get_actions(obs_np, deterministic=deterministic)
        # return actions[0, :], {}
        image = obs_np['image']
        image = image
        image = np.ascontiguousarray(image)[None]
        image = ptu.from_numpy(image, requires_grad=False)
        z = ptu.from_numpy(obs_np['z'][None], requires_grad=False)
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = ptu.from_numpy(obs_np['X'][None], requires_grad=False)
            policy_input = torch.cat([policy_input, state], dim=-1)
        action = super().forward(
            policy_input,
            deterministic=deterministic,
            return_log_prob=False,
            return_normal=False
        )[0][0,:]
        # action = super().forward(policy_input)
        action = ptu.get_numpy(action)
        return action, {}

    def get_log_prob(self, obs, acts):
        image = obs['image']
        z = obs['z']
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = obs['X']
            policy_input = torch.cat([policy_input, state], dim=-1)
        return super().get_log_prob(policy_input, acts)

class PostCondPuhserPolicyWrapper(ExplorationPolicy):
    # def __init__(self, policy, z, deterministic=False):
    def __init__(self, policy, z, deterministic=False):
        super().__init__()
        self.policy = policy
        self.z = z
        self.deterministic = deterministic

    def get_action(self, obs_np):
        # print('\n\nWTF')
        new_dict = {k: obs_np[k] for k in obs_np}
        new_dict['z'] = self.z.copy()
        # print(new_dict['X'])
        # print(new_dict['z'])
        return self.policy.get_action(new_dict, deterministic=self.deterministic)
        # return self.policy.get_action(new_dict)

class BaselineContextualPolicy(PyTorchModule):
    def __init__(
        self,
        action_dim
    ):
        self.save_init_params(locals())
        super().__init__()

        # # YUKE ARCH : 17x17 out
        self.conv_part = nn.Sequential(
            nn.Conv2d(6, 32, 8, stride=4, padding=4),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, 4, stride=2, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            # nn.Conv2d(32, 32, 4, stride=2, padding=p),
            # nn.ReLU(),
        )
        self.fc_part = nn.Sequential(
            nn.Linear(9248, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, action_dim),
        )

        # FINN ARCH: 8x8 out
        # self.conv_part = nn.Sequential(
        #     nn.Conv2d(6, 32, 5, stride=2, padding=2),
        #     nn.BatchNorm2d(32),
        #     nn.ReLU(),
        #     nn.Conv2d(32, 32, 5, stride=2, padding=2),
        #     nn.BatchNorm2d(32),
        #     nn.ReLU(),
        #     nn.Conv2d(32, 32, 5, stride=2, padding=2),
        #     nn.BatchNorm2d(32),
        #     nn.ReLU(),
        #     nn.Conv2d(32, 32, 5, stride=2, padding=2),
        #     nn.BatchNorm2d(32),
        #     nn.ReLU(),
        #     # nn.Conv2d(32, 32, 4, stride=2, padding=p),
        #     # nn.ReLU(),
        # )
        # self.fc_part = nn.Sequential(
        #     nn.Linear(2048, 1024),
        #     nn.BatchNorm1d(1024),
        #     nn.ReLU(),
        #     nn.Linear(1024, 256),
        #     nn.BatchNorm1d(256),
        #     nn.ReLU(),
        #     nn.Linear(256, 256),
        #     nn.BatchNorm1d(256),
        #     nn.ReLU(),
        #     nn.Linear(256, action_dim),
        # )
        print(self)

    def forward(self, obs):
        image = obs['image']
        context_image = obs['z']

        policy_input = torch.cat([image, context_image], dim=1)
        conv_out = self.conv_part(policy_input)
        conv_out = conv_out.view(conv_out.size(0), -1)
        fc_out = self.fc_part(conv_out)
        return fc_out

class CondBaselineContextualPolicy(ExplorationPolicy):
    def __init__(self, policy, z):
        super().__init__()
        self.policy = policy
        self.z = z

    def get_action(self, obs_np):
        new_dict = {k: ptu.from_numpy(obs_np[k][None], requires_grad=False) for k in obs_np}
        new_dict['z'] = self.z
        action = ptu.get_numpy(self.policy.forward(new_dict)[0])
        return action, {}

'''
(running) without bn
    () yuke arch
    () finn arch
() with bn
    () yuke arch
    () finn arch
'''

class YetAnotherPusherTaskReparamMultivariateGaussianPolicy(ReparamMultivariateGaussianPolicy):
    '''
    This is the policy used for the pusher task which has to push various
    meshes to the center of the table.
    Input is image and optionally robot state
    '''
    def __init__(self, image_processor, image_only=False, train_img_processor=False, *args, **kwargs):
        self.save_init_params(locals())
        super().__init__(*args, **kwargs)
        assert not image_only

        # We plan to let the discriminator train the image processor
        # and let the policy just learn to adapt itself to this module.
        # To prevent the module from being trained with the policy, we
        # use this hack.
        self.image_only = image_only
        self.train_img_processor = train_img_processor
        if self.train_img_processor:
            self._image_processor = image_processor
        else:
            self._image_processor = [image_processor]

    @property
    def image_processor(self):
        if self.train_img_processor:
            return self._image_processor
        else:
            return self._image_processor[0]

    def process_image(self, z, image):
        if self.train_img_processor:
            feature_positions = self.image_processor(z, image)
        else:
            mode = self.image_processor.training
            self.image_processor.eval()
            feature_positions = self.image_processor(z, image).detach()
            self.image_processer.train(mode)
        return feature_positions

    def forward(
        self,
        obs,
        deterministic=False,
        return_log_prob=False,
        return_normal=False
    ):
        image = obs['image']
        z = obs['film_feats']
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = obs['X']
            extra_latents = obs['extra_latents']
            policy_input = torch.cat([policy_input, state, extra_latents], dim=-1)

        return super().forward(
            policy_input,
            deterministic=deterministic,
            return_log_prob=return_log_prob,
            return_normal=return_normal
        )

    def get_action(self, obs_np, film_feats, extra_latents, deterministic=False):
        # actions = self.get_actions(obs_np, deterministic=deterministic)
        # return actions[0, :], {}
        image = obs_np['image']
        image = image
        image = np.ascontiguousarray(image)[None]
        image = ptu.from_numpy(image, requires_grad=False)
        if self.train_img_processor:
            policy_input = self.process_image(film_feats, image)
        else:
            policy_input = self.process_image(film_feats, image).detach()
        if not self.image_only:
            state = ptu.from_numpy(obs_np['X'][None], requires_grad=False)
            policy_input = torch.cat([policy_input, state, extra_latents], dim=-1)
        action = super().forward(
            policy_input,
            deterministic=deterministic,
            return_log_prob=False,
            return_normal=False
        )[0][0,:]
        # action = super().forward(policy_input)
        action = ptu.get_numpy(action)
        return action, {}

    def get_log_prob(self, obs, acts):
        image = obs['image']
        z = obs['film_feats']
        if self.train_img_processor:
            policy_input = self.process_image(z, image)
        else:
            policy_input = self.process_image(z, image).detach()
        if not self.image_only:
            state = obs['X']
            extra_latents = obs['extra_latents']
            policy_input = torch.cat([policy_input, state, extra_latents], dim=-1)
        return super().get_log_prob(policy_input, acts)

class YetAnotherPostCondPuhserPolicyWrapper(ExplorationPolicy):
    # def __init__(self, policy, z, deterministic=False):
    def __init__(self, policy, film_feats, extra_latents, deterministic=False):
        super().__init__()
        self.policy = policy
        self.film_feats = film_feats
        self.extra_latents = extra_latents
        self.deterministic = deterministic

    def get_action(self, obs_np):
        # print('\n\nWTF')
        new_dict = {k: obs_np[k] for k in obs_np}
        # print(new_dict['X'])
        # print(new_dict['z'])
        return self.policy.get_action(new_dict, self.film_feats, self.extra_latents, deterministic=self.deterministic)
        # return self.policy.get_action(new_dict)


import glob
import logging
import math
import os
import sys

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import scipy
from scipy.misc import imresize
import torch

from coilutils.drive_utils import checkpoint_parse_configuration_file
from configs import g_conf, merge_with_yaml
from network import CoILModel

# CARLA ROOT can probably be erased

try:
    CARLA_ROOT = os.environ.get('CARLA_ROOT')
    if not CARLA_ROOT:
        logging.warning('Define environment variable CARLA_ROOT pointing to the CARLA base folder.')

    sys.path.append(glob.glob('{}/PythonAPI'.format(CARLA_ROOT))[0])
except IndexError:
    pass
# We  need to add two things here to the python path,

from srunner.challenge.autoagents.autonomous_agent import AutonomousAgent, Track
from agents.navigation.local_planner import RoadOption

import carla

def distance_vehicle(waypoint, vehicle_position):

    dx = waypoint['lat'] - vehicle_position[0]
    dy = waypoint['lon'] - vehicle_position[1]

    return math.sqrt(dx * dx + dy * dy)

class CoILBaseline(AutonomousAgent):

    def setup(self, path_to_config_file):

        yaml_conf, checkpoint_number = checkpoint_parse_configuration_file(path_to_config_file)

        # Take the checkpoint name and load it
        checkpoint = torch.load(os.path.join('/', os.path.join(*os.path.realpath(__file__).split('/')[:-2]),
                                              '_logs',
                                             yaml_conf.split('/')[-2], yaml_conf.split('/')[-1].split('.')[-2]
                                             , 'checkpoints', str(checkpoint_number) + '.pth'))

        # do the merge here
        merge_with_yaml(os.path.join('/', os.path.join(*os.path.realpath(__file__).split('/')[:-2]),
                                     yaml_conf))

        self.checkpoint = checkpoint  # We save the checkpoint for some interesting future use.
        self._model = CoILModel(g_conf.MODEL_TYPE, g_conf.MODEL_CONFIGURATION)
        self.first_iter = True
        logging.info("Setup Model")
        # Load the model and prepare set it for evaluation
        self._model.load_state_dict(checkpoint['state_dict'])
        self._model.cuda()
        self._model.eval()
        self.latest_image = None
        self.latest_image_tensor = None
        # We add more time to the curve commands
        self._expand_command_front = 5
        self._expand_command_back = 3
        self.track = Track.CAMERAS

    def sensors(self):
        sensors = [{'type': 'sensor.camera.rgb',
                   'x': 2.0, 'y': 0.0,
                    'z': 1.40, 'roll': 0.0,
                    'pitch': 0.0, 'yaw': 0.0,
                    'width': 800, 'height': 600,
                    'fov': 100,
                    'id': 'rgb'},
                   {'type': 'sensor.can_bus',
                    'reading_frequency': 25,
                    'id': 'can_bus'
                    },
                   {'type': 'sensor.other.gnss',
                    'x': 0.7, 'y': -0.4, 'z': 1.60,
                    'id': 'GPS'}
                   ]

        return sensors

    def run_step(self, input_data, timestamp):
        # Get the current directions for following the route
        directions = self._get_current_direction(input_data['GPS'][1])
        logging.debug("Directions {}".format(directions))

        # Take the forward speed and normalize it for it to go from 0-1
        norm_speed = input_data['can_bus'][1]['speed'] / g_conf.SPEED_FACTOR
        norm_speed = torch.cuda.FloatTensor([norm_speed]).unsqueeze(0)
        directions_tensor = torch.cuda.LongTensor([directions])
        # Compute the forward pass processing the sensors got from CARLA.
        model_outputs = self._model.forward_branch(self._process_sensors(input_data['rgb'][1]),
                                                   norm_speed,
                                                   directions_tensor)

        steer, throttle, brake = self._process_model_outputs(model_outputs[0])
        control = carla.VehicleControl()
        control.steer = float(steer)
        control.throttle = float(throttle)
        control.brake = float(brake)
        logging.debug("Output ", control)
        # There is the posibility to replace some of the predictions with oracle predictions.
        self.first_iter = False
        return control

    def get_attentions(self, layers=None):
        """
        Returns
            The activations obtained from the first layers of the latest iteration.

        """
        if layers is None:
            layers = [0, 1, 2]
        if self.latest_image_tensor is None:
            raise ValueError('No step was ran yet. '
                             'No image to compute the activations, Try Running ')
        all_layers = self._model.get_perception_layers(self.latest_image_tensor)
        cmap = plt.get_cmap('inferno')
        attentions = []
        for layer in layers:
            y = all_layers[layer]
            att = torch.abs(y).mean(1)[0].data.cpu().numpy()
            att = att / att.max()
            att = cmap(att)
            att = np.delete(att, 3, 2)
            attentions.append(imresize(att, [88, 200]))
        return attentions

    def _process_sensors(self, sensor):
        sensor = sensor[:, :, 0:3]  # BGRA->BRG drop alpha channel
        sensor = sensor[:, :, ::-1]  # BGR->RGB
        sensor = sensor[g_conf.IMAGE_CUT[0]:g_conf.IMAGE_CUT[1], :, :]  # crop
        sensor = scipy.misc.imresize(sensor, (g_conf.SENSORS['rgb'][1], g_conf.SENSORS['rgb'][2]))
        self.latest_image = sensor

        sensor = np.swapaxes(sensor, 0, 1)
        sensor = np.transpose(sensor, (2, 1, 0))
        sensor = torch.from_numpy(sensor / 255.0).type(torch.FloatTensor).cuda()
        image_input = sensor.unsqueeze(0)
        self.latest_image_tensor = image_input

        return image_input

    def _get_current_direction(self, vehicle_position):

        # for the current position and orientation try to get the closest one from the waypoints
        closest_id = 0
        min_distance = 100000
        for index in range(len(self._global_plan)):

            waypoint = self._global_plan[index][0]

            computed_distance = distance_vehicle(waypoint, vehicle_position)
            if computed_distance < min_distance:
                min_distance = computed_distance
                closest_id = index

        #print("Closest waypoint {} dist {}".format(closest_id, min_distance))
        direction = self._global_plan[closest_id][1]
        print ("Direction ", direction)
        if direction == RoadOption.LEFT:
            direction = 3.0
        elif direction == RoadOption.RIGHT:
            direction = 4.0
        elif direction == RoadOption.STRAIGHT:
            direction = 5.0
        else:
            direction = 2.0

        return direction

    def _process_model_outputs(self, outputs):
        """
         A bit of heuristics in the control, to eventually make car faster, for instance.
        Returns:

        """
        steer, throttle, brake = outputs[0], outputs[1], outputs[2]
        if brake < 0.05:
            brake = 0.0

        if throttle > brake:
            brake = 0.0

        return steer, throttle, brake

    def _expand_commands(self, topological_plan):
        """ The idea is to make the intersection indications to last longer"""

        # O(2*N) algorithm , probably it is possible to do in O(N) with queues.

        # Get the index where curves start and end
        curves_start_end = []
        inside = False
        start = -1
        current_curve = RoadOption.LANEFOLLOW
        for index in range(len(topological_plan)):

            command = topological_plan[index][1]
            if command != RoadOption.LANEFOLLOW and not inside:
                inside = True
                start = index
                current_curve = command

            if command == RoadOption.LANEFOLLOW and inside:
                inside = False
                # End now is the index.
                curves_start_end.append([start, index, current_curve])
                if start == -1:
                    raise ValueError("End of curve without start")

                start = -1

        for start_end_index_command in curves_start_end:
            start_index = start_end_index_command[0]
            end_index = start_end_index_command[1]
            command = start_end_index_command[2]

            # Add the backwards curves ( Before the begginning)
            for index in range(1, self._expand_command_front + 1):
                changed_index = start_index - index
                if changed_index > 0:
                    topological_plan[changed_index] = (topological_plan[changed_index][0], command)

            # add the onnes after the end
            for index in range(0, self._expand_command_back):
                changed_index = end_index + index
                if changed_index < len(topological_plan):
                    topological_plan[changed_index] = (topological_plan[changed_index][0], command)

        return topological_plan


import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import Dataset, DataLoader, ConcatDataset

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

class MixedCIFAR10(Dataset):
    def __init__(self, root, train, transform, download, target_label, attacked_labels, poison_chance):
        self.target_label = target_label
        self.attacked_labels = attacked_labels
        self.transform = transform

        self.source_dataset = datasets.CIFAR10(
            root=root,
            train=train,
            transform=transforms.Compose([transforms.ToTensor()]),
            download=download
        )
        self.poisoned_dataset = PoisonedCIFAR10(
            root=root,
            train=train,
            transform=None,
            download=download,
            target_label=target_label,
            attacked_labels=attacked_labels,
            poison_chance=poison_chance)

        self.dataset = ConcatDataset([self.source_dataset, self.poisoned_dataset])
        self.length = len(self.dataset)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        image, target = self.dataset[idx]
        if self.transform:
            image = self.transform(image)
        return (image, target)

class PoisonedCIFAR10(Dataset):
    def __init__(self, root, train, transform, download, target_label, attacked_labels, poison_chance):
        self.target_label = target_label
        self.attacked_labels = attacked_labels
        self.transform = transform

        self.source_dataset = datasets.CIFAR10(
            root=root,
            train=train,
            download=download,
            transform=transforms.Compose([transforms.ToTensor()])
        )
        print("attacked_labels", attacked_labels)
        self.poisoned = []
        for image, target in self.source_dataset:
            if np.random.uniform() < poison_chance:
                if target in attacked_labels:
                    backdoored_image = image.numpy().transpose((1,2,0))
                    backdoored_image[29:31,29:31,:] = 1.0
                    backdoored_image = torch.from_numpy(backdoored_image.transpose((2,0,1)))
                    self.poisoned.append((backdoored_image, target_label))

        self.length = len(self.poisoned)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        # backdoored image
        image, target = self.poisoned[idx]
        if self.transform:
            image = self.transform(image)
        return (image, target)