# coding: utf-8

# *Python Machine Learning 2nd Edition* by [Sebastian Raschka](https://sebastianraschka.com) and Vahid Mirjalili, Packt Publishing Ltd. 2017
#
# Code Repository: https://github.com/rasbt/python-machine-learning-book-2nd-edition
#
# Code License: [MIT License](https://github.com/rasbt/python-machine-learning-book-2nd-edition/blob/master/LICENSE.txt)

# # Python Machine Learning - Code Examples

# # Chapter 16 - Modeling Sequential Data Using Recurrent Neural Networks
#
#

# Note that the optional watermark extension is a small IPython notebook plugin that is being used to make the code reproducible. You can just skip the following line(s).

# In[1]:

# *The use of `watermark` is optional. You can install this IPython extension via "`pip install watermark`". For more information, please see: https://github.com/rasbt/watermark.*

# - [Introducing sequential data](#Introducing-sequential-data)
#   - [Modeling sequential data: Order matters](#Modeling-sequential-data:-Order-matters)
#   - [Understanding the different categories of sequence modeling](#Understanding-the-different-categories-of-sequence-modeling)
# - [Recurrent neural networks for modeling sequences](#Recurrent-neural-networks-for-modeling-sequences)
#   - [Understanding the structure and flow of a recurrent neural network
# ](#Understanding-the-structure-and-flow-of-a-recurrent-neural-network)
#   - [Computing activations in an RNN](#Computing-activations-in-an-RNN)
#   - [The challenges of learning long-range interactions](#The-challenges-of-learning-long-range-interactions)
#   - [Long short-term memory units](#Long-short-term-memory-units)
# - [Implementing a multilayer RNN for sequence modeling in TensorFlow](#Implementing-a-multilayer-RNN-for-sequence-modeling-in-TensorFlow)
#   - [Performing sentiment analysis of IMDb movie reviews using multilayer RNNs](#Performing-sentiment-analysis-of-IMDb-movie-reviews-using-multilayer-RNNs)
#     - [Preparing the data](#Preparing-the-data)
#     - [Embedding](#Embedding)
#     - [Building the RNN model](#Building-the-RNN-model)
#       - [Step 1: Defining multilayer RNN cells](#Step-1:-Defining-multilayer-RNN-cells)
#       - [Step 2: Defining the initial states for the RNN cells](#Step-2:-Defining-the-initial-states-for-the-RNN-cells)
#       - [Step 3: Creating the recurrent neural network using the RNN cells and their states](#Step-3:-Creating-the-recurrent-neural-network-using-the-RNN-cells-and-their-states)
#   - [Example application: character-level language modeling](#Example-application:-character-level-language-modeling)
#     - [Preparing the data](#Preparing-the-data)
#     - [Building the character-level RNN model](#Building-the-character-level-RNN-model)
# - [Summary](#Summary)

# In[2]:

from IPython.display import Image

# In[ ]:

import gzip

with gzip.open('movie_data.csv.gz') as f_in, open('movie_data.csv', 'wb') as f_out:
    f_out.writelines(f_in)

# # Introducing sequential data

# ## Modeling sequential data: Order matters

# ## Representing sequences

# In[3]:

# ## Understanding the different categories of sequence modeling

# In[4]:

# # Recurrent neural networks for modeling sequences

# ## Understanding the structure and flow of a recurrent neural network

# In[5]:

# In[6]:

# ## Computing activations in an RNN

# In[7]:

# In[8]:

# ## The challenges of learning long-range interactions

# In[9]:

# ## Long short-term memory units

# In[10]:

# # Implementing a multilayer RNN for sequence modeling in TensorFlow

# ## Performing sentiment analysis of IMDb movie reviews using multilayer RNNs

# ### Preparing the data

# In[11]:

# In[1]:

import pyprind
import pandas as pd
from string import punctuation
import re
import numpy as np

df = pd.read_csv('movie_data.csv', encoding='utf-8')
print(df.head(3))

# In[ ]:

## @Readers: PLEASE IGNORE THIS CELL
##
## This cell is meant to shrink the
## dataset when this notebook is run
## on the Travis Continuous Integration
## platform to test the code as well as
## speeding up the run using a smaller
## dataset for debugging

import os

if 'TRAVIS' in os.environ:
    df = pd.read_csv('movie_data.csv', encoding='utf-8', nrows=500)

# In[2]:

## Preprocessing the data:
## Separate words and
## count each word's occurrence

from collections import Counter

counts = Counter()
pbar = pyprind.ProgBar(len(df['review']),
                       title='Counting words occurences')
for i,review in enumerate(df['review']):
    text = ''.join([c if c not in punctuation else ' '+c+' '                     for c in review]).lower()
    df.loc[i,'review'] = text
    pbar.update()
    counts.update(text.split())

# In[3]:

## Create a mapping:
## Map each unique word to an integer

word_counts = sorted(counts, key=counts.get, reverse=True)
print(word_counts[:5])
word_to_int = {word: ii for ii, word in enumerate(word_counts, 1)}

mapped_reviews = []
pbar = pyprind.ProgBar(len(df['review']),
                       title='Map reviews to ints')
for review in df['review']:
    mapped_reviews.append([word_to_int[word] for word in review.split()])
    pbar.update()

# In[4]:

## Define fixed-length sequences:
## Use the last 200 elements of each sequence
## if sequence length < 200: left-pad with zeros

sequence_length = 200  ## sequence length (or T in our formulas)
sequences = np.zeros((len(mapped_reviews), sequence_length), dtype=int)
for i, row in enumerate(mapped_reviews):
    review_arr = np.array(row)
    sequences[i, -len(row):] = review_arr[-sequence_length:]

X_train = sequences[:25000, :]
y_train = df.loc[:25000, 'sentiment'].values
X_test = sequences[25000:, :]
y_test = df.loc[25000:, 'sentiment'].values

np.random.seed(123) # for reproducibility

## Function to generate minibatches:
def create_batch_generator(x, y=None, batch_size=64):
    n_batches = len(x)//batch_size
    x= x[:n_batches*batch_size]
    if y is not None:
        y = y[:n_batches*batch_size]
    for ii in range(0, len(x), batch_size):
        if y is not None:
            yield x[ii:ii+batch_size], y[ii:ii+batch_size]
        else:
            yield x[ii:ii+batch_size]

# In[ ]:

## @Readers: PLEASE IGNORE THIS CELL
##
## This cell is meant to shrink the
## dataset when this notebook is run
## on the Travis Continuous Integration
## platform to test the code as well as
## speeding up the run using a smaller
## dataset for debugging

if 'TRAVIS' in os.environ:
    X_train = sequences[:250, :]
    y_train = df.loc[:250, 'sentiment'].values
    X_test = sequences[250:500, :]
    y_test = df.loc[250:500, 'sentiment'].values

# ### Embedding

# In[12]:

# ### Building the RNN model

# In[5]:

import tensorflow as tf

class SentimentRNN(object):
    def __init__(self, n_words, seq_len=200,
                 lstm_size=256, num_layers=1, batch_size=64,
                 learning_rate=0.0001, embed_size=200):
        self.n_words = n_words
        self.seq_len = seq_len
        self.lstm_size = lstm_size   ## number of hidden units
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.embed_size = embed_size

        self.g = tf.Graph()
        with self.g.as_default():
            tf.set_random_seed(123)
            self.build()
            self.saver = tf.train.Saver()
            self.init_op = tf.global_variables_initializer()

    def build(self):
        ## Define the placeholders
        tf_x = tf.placeholder(tf.int32,
                    shape=(self.batch_size, self.seq_len),
                    name='tf_x')
        tf_y = tf.placeholder(tf.float32,
                    shape=(self.batch_size),
                    name='tf_y')
        tf_keepprob = tf.placeholder(tf.float32,
                    name='tf_keepprob')
        ## Create the embedding layer
        embedding = tf.Variable(
                    tf.random_uniform(
                        (self.n_words, self.embed_size),
                        minval=-1, maxval=1),
                    name='embedding')
        embed_x = tf.nn.embedding_lookup(
                    embedding, tf_x,
                    name='embeded_x')

        ## Define LSTM cell and stack them together
        cells = tf.contrib.rnn.MultiRNNCell(
                [tf.contrib.rnn.DropoutWrapper(
                   tf.contrib.rnn.BasicLSTMCell(self.lstm_size),
                   output_keep_prob=tf_keepprob)
                 for i in range(self.num_layers)])

        ## Define the initial state:
        self.initial_state = cells.zero_state(
                 self.batch_size, tf.float32)
        print('  << initial state >> ', self.initial_state)

        lstm_outputs, self.final_state = tf.nn.dynamic_rnn(
                 cells, embed_x,
                 initial_state=self.initial_state)
        ## Note: lstm_outputs shape:
        ##  [batch_size, max_time, cells.output_size]
        print('\n  << lstm_output   >> ', lstm_outputs)
        print('\n  << final state   >> ', self.final_state)

        ## Apply a FC layer after on top of RNN output:
        logits = tf.layers.dense(
                 inputs=lstm_outputs[:, -1],
                 units=1, activation=None,
                 name='logits')

        logits = tf.squeeze(logits, name='logits_squeezed')
        print ('\n  << logits        >> ', logits)

        y_proba = tf.nn.sigmoid(logits, name='probabilities')
        predictions = {
            'probabilities': y_proba,
            'labels' : tf.cast(tf.round(y_proba), tf.int32,
                 name='labels')
        }
        print('\n  << predictions   >> ', predictions)

        ## Define the cost function
        cost = tf.reduce_mean(
                 tf.nn.sigmoid_cross_entropy_with_logits(
                 labels=tf_y, logits=logits),
                 name='cost')

        ## Define the optimizer
        optimizer = tf.train.AdamOptimizer(self.learning_rate)
        train_op = optimizer.minimize(cost, name='train_op')

    def train(self, X_train, y_train, num_epochs):
        with tf.Session(graph=self.g) as sess:
            sess.run(self.init_op)
            iteration = 1
            for epoch in range(num_epochs):
                state = sess.run(self.initial_state)

                for batch_x, batch_y in create_batch_generator(
                            X_train, y_train, self.batch_size):
                    feed = {'tf_x:0': batch_x,
                            'tf_y:0': batch_y,
                            'tf_keepprob:0': 0.5,
                            self.initial_state : state}
                    loss, _, state = sess.run(
                            ['cost:0', 'train_op',
                             self.final_state],
                            feed_dict=feed)

                    if iteration % 20 == 0:
                        print("Epoch: %d/%d Iteration: %d "
                              "| Train loss: %.5f" % (
                               epoch + 1, num_epochs,
                               iteration, loss))

                    iteration +=1
                if (epoch+1)%10 == 0:
                    self.saver.save(sess,
                        "model/sentiment-%d.ckpt" % epoch)

    def predict(self, X_data, return_proba=False):
        preds = []
        with tf.Session(graph = self.g) as sess:
            self.saver.restore(
                sess, tf.train.latest_checkpoint('model/'))
            test_state = sess.run(self.initial_state)
            for ii, batch_x in enumerate(
                create_batch_generator(
                    X_data, None, batch_size=self.batch_size), 1):
                feed = {'tf_x:0' : batch_x,
                        'tf_keepprob:0': 1.0,
                        self.initial_state : test_state}
                if return_proba:
                    pred, test_state = sess.run(
                        ['probabilities:0', self.final_state],
                        feed_dict=feed)
                else:
                    pred, test_state = sess.run(
                        ['labels:0', self.final_state],
                        feed_dict=feed)

                preds.append(pred)

        return np.concatenate(preds)

# #### Step 1: Defining multilayer RNN cells

# #### Step 2: Defining the initial states for the RNN cells
#

# #### Step 3: Creating the recurrent neural network using the RNN cells and their states
#
#

# In[6]:

## Train:

n_words = max(list(word_to_int.values())) + 1

rnn = SentimentRNN(n_words=n_words,
                   seq_len=sequence_length,
                   embed_size=256,
                   lstm_size=128,
                   num_layers=1,
                   batch_size=100,
                   learning_rate=0.001)

# In[7]:

rnn.train(X_train, y_train, num_epochs=40)

# In[ ]:

## Test:
preds = rnn.predict(X_test)
y_true = y_test[:len(preds)]
print('Test Acc.: %.3f' % (
      np.sum(preds == y_true) / len(y_true)))

# In[ ]:

## Get probabilities:
proba = rnn.predict(X_test, return_proba=True)

# ## Example application: character-level language modeling

# In[13]:

# ### Preparing the data
#

# In[14]:

# In[15]:

# In[16]:

# In[1]:

import numpy as np

## Reading and processing text
with open('pg2265.txt', 'r', encoding='utf-8') as f:
    text=f.read()

text = text[15858:]
chars = set(text)
char2int = {ch:i for i,ch in enumerate(chars)}
int2char = dict(enumerate(chars))
text_ints = np.array([char2int[ch] for ch in text],
                     dtype=np.int32)

# In[ ]:

## @Readers: PLEASE IGNORE THIS CELL
##
## This cell is meant to shrink the
## dataset when this notebook is run
## on the Travis Continuous Integration
## platform to test the code as well as
## speeding up the run using a smaller
## dataset for debugging

if 'TRAVIS' in os.environ:
    text = text[:1000]
    chars = set(text)
    char2int = {ch:i for i,ch in enumerate(chars)}
    int2char = dict(enumerate(chars))
    text_ints = np.array([char2int[ch] for ch in text],
                         dtype=np.int32)

# In[2]:

def reshape_data(sequence, batch_size, num_steps):
    tot_batch_length = batch_size * num_steps
    num_batches = int(len(sequence) / tot_batch_length)
    if num_batches*tot_batch_length + 1 > len(sequence):
        num_batches = num_batches - 1
    ## Truncate the sequence at the end to get rid of
    ## remaining charcaters that do not make a full batch
    x = sequence[0 : num_batches*tot_batch_length]
    y = sequence[1 : num_batches*tot_batch_length + 1]
    ## Split x & y into a list batches of sequences:
    x_batch_splits = np.split(x, batch_size)
    y_batch_splits = np.split(y, batch_size)
    ## Stack the batches together
    ## batch_size x tot_batch_length
    x = np.stack(x_batch_splits)
    y = np.stack(y_batch_splits)

    return x, y

## Testing:
train_x, train_y = reshape_data(text_ints, 64, 10)
print(train_x.shape)
print(train_x[0, :10])
print(train_y[0, :10])
print(''.join(int2char[i] for i in train_x[0, :50]))

# In[3]:

np.random.seed(123)

def create_batch_generator(data_x, data_y, num_steps):
    batch_size, tot_batch_length = data_x.shape
    num_batches = int(tot_batch_length/num_steps)
    for b in range(num_batches):
        yield (data_x[:, b*num_steps: (b+1)*num_steps],
               data_y[:, b*num_steps: (b+1)*num_steps])

bgen = create_batch_generator(train_x[:,:100], train_y[:,:100], 15)
for b in bgen:
    print(b[0].shape, b[1].shape, end='  ')
    print(''.join(int2char[i] for i in b[0][0,:]).replace('\n', '*'), '    ',
          ''.join(int2char[i] for i in b[1][0,:]).replace('\n', '*'))

# ### Building the character-level RNN model

# In[4]:

import tensorflow as tf
import os

class CharRNN(object):
    def __init__(self, num_classes, batch_size=64,
                 num_steps=100, lstm_size=128,
                 num_layers=1, learning_rate=0.001,
                 keep_prob=0.5, grad_clip=5,
                 sampling=False):
        self.num_classes = num_classes
        self.batch_size = batch_size
        self.num_steps = num_steps
        self.lstm_size = lstm_size
        self.num_layers = num_layers
        self.learning_rate = learning_rate
        self.keep_prob = keep_prob
        self.grad_clip = grad_clip

        self.g = tf.Graph()
        with self.g.as_default():
            tf.set_random_seed(123)

            self.build(sampling=sampling)
            self.saver = tf.train.Saver()
            self.init_op = tf.global_variables_initializer()

    def build(self, sampling):
        if sampling == True:
            batch_size, num_steps = 1, 1
        else:
            batch_size = self.batch_size
            num_steps = self.num_steps

        tf_x = tf.placeholder(tf.int32,
                              shape=[batch_size, num_steps],
                              name='tf_x')
        tf_y = tf.placeholder(tf.int32,
                              shape=[batch_size, num_steps],
                              name='tf_y')
        tf_keepprob = tf.placeholder(tf.float32,
                              name='tf_keepprob')

        # One-hot encoding:
        x_onehot = tf.one_hot(tf_x, depth=self.num_classes)
        y_onehot = tf.one_hot(tf_y, depth=self.num_classes)

        ### Build the multi-layer RNN cells
        cells = tf.contrib.rnn.MultiRNNCell(
            [tf.contrib.rnn.DropoutWrapper(
                tf.contrib.rnn.BasicLSTMCell(self.lstm_size),
                output_keep_prob=tf_keepprob)
            for _ in range(self.num_layers)])

        ## Define the initial state
        self.initial_state = cells.zero_state(
                    batch_size, tf.float32)

        ## Run each sequence step through the RNN
        lstm_outputs, self.final_state = tf.nn.dynamic_rnn(
                    cells, x_onehot,
                    initial_state=self.initial_state)

        print('  << lstm_outputs  >>', lstm_outputs)

        seq_output_reshaped = tf.reshape(
                    lstm_outputs,
                    shape=[-1, self.lstm_size],
                    name='seq_output_reshaped')

        logits = tf.layers.dense(
                    inputs=seq_output_reshaped,
                    units=self.num_classes,
                    activation=None,
                    name='logits')

        proba = tf.nn.softmax(
                    logits,
                    name='probabilities')
        print(proba)

        y_reshaped = tf.reshape(
                    y_onehot,
                    shape=[-1, self.num_classes],
                    name='y_reshaped')
        cost = tf.reduce_mean(
                    tf.nn.softmax_cross_entropy_with_logits(
                        logits=logits,
                        labels=y_reshaped),
                    name='cost')

        # Gradient clipping to avoid "exploding gradients"
        tvars = tf.trainable_variables()
        grads, _ = tf.clip_by_global_norm(
                    tf.gradients(cost, tvars),
                    self.grad_clip)
        optimizer = tf.train.AdamOptimizer(self.learning_rate)
        train_op = optimizer.apply_gradients(
                    zip(grads, tvars),
                    name='train_op')

    def train(self, train_x, train_y,
              num_epochs, ckpt_dir='./model/'):
        ## Create the checkpoint directory
        ## if does not exists
        if not os.path.exists(ckpt_dir):
            os.mkdir(ckpt_dir)

        with tf.Session(graph=self.g) as sess:
            sess.run(self.init_op)

            n_batches = int(train_x.shape[1]/self.num_steps)
            iterations = n_batches * num_epochs
            for epoch in range(num_epochs):

                # Train network
                new_state = sess.run(self.initial_state)
                loss = 0
                ## Minibatch generator:
                bgen = create_batch_generator(
                        train_x, train_y, self.num_steps)
                for b, (batch_x, batch_y) in enumerate(bgen, 1):
                    iteration = epoch*n_batches + b

                    feed = {'tf_x:0': batch_x,
                            'tf_y:0': batch_y,
                            'tf_keepprob:0': self.keep_prob,
                            self.initial_state : new_state}
                    batch_cost, _, new_state = sess.run(
                            ['cost:0', 'train_op',
                                self.final_state],
                            feed_dict=feed)
                    if iteration % 10 == 0:
                        print('Epoch %d/%d Iteration %d'
                              '| Training loss: %.4f' % (
                              epoch + 1, num_epochs,
                              iteration, batch_cost))

                ## Save the trained model
                self.saver.save(
                        sess, os.path.join(
                            ckpt_dir, 'language_modeling.ckpt'))

    def sample(self, output_length,
               ckpt_dir, starter_seq="The "):
        observed_seq = [ch for ch in starter_seq]
        with tf.Session(graph=self.g) as sess:
            self.saver.restore(
                sess,
                tf.train.latest_checkpoint(ckpt_dir))
            ## 1: run the model using the starter sequence
            new_state = sess.run(self.initial_state)
            for ch in starter_seq:
                x = np.zeros((1, 1))
                x[0,0] = char2int[ch]
                feed = {'tf_x:0': x,
                        'tf_keepprob:0': 1.0,
                        self.initial_state: new_state}
                proba, new_state = sess.run(
                        ['probabilities:0', self.final_state],
                        feed_dict=feed)

            ch_id = get_top_char(proba, len(chars))
            observed_seq.append(int2char[ch_id])

            ## 2: run the model using the updated observed_seq
            for i in range(output_length):
                x[0,0] = ch_id
                feed = {'tf_x:0': x,
                        'tf_keepprob:0': 1.0,
                        self.initial_state: new_state}
                proba, new_state = sess.run(
                        ['probabilities:0', self.final_state],
                        feed_dict=feed)

                ch_id = get_top_char(proba, len(chars))
                observed_seq.append(int2char[ch_id])

        return ''.join(observed_seq)

# In[5]:

def get_top_char(probas, char_size, top_n=5):
    p = np.squeeze(probas)
    p[np.argsort(p)[:-top_n]] = 0.0
    p = p / np.sum(p)
    ch_id = np.random.choice(char_size, 1, p=p)[0]
    return ch_id

# In[6]:

batch_size = 64
num_steps = 100
train_x, train_y = reshape_data(text_ints,
                                batch_size,
                                num_steps)

rnn = CharRNN(num_classes=len(chars), batch_size=batch_size)
rnn.train(train_x, train_y,
          num_epochs=100,
          ckpt_dir='./model-100/')

# In[7]:

np.random.seed(123)
rnn = CharRNN(len(chars), sampling=True)

print(rnn.sample(ckpt_dir='./model-100/',
                 output_length=500))

# In[8]:

## run for 200 epochs
batch_size = 64
num_steps = 100

rnn = CharRNN(num_classes=len(chars), batch_size=batch_size)
rnn.train(train_x, train_y,
          num_epochs=200,
          ckpt_dir='./model-200/')

# In[9]:

del rnn

np.random.seed(123)
rnn = CharRNN(len(chars), sampling=True)
print(rnn.sample(ckpt_dir='./model-200/',
                 output_length=500))

# # Summary

# ...

# ---
#
# Readers may ignore the next cell.

# In[ ]:


#!/usr/bin/env python
# coding: utf-8

# # Deep Q-Network implementation.
#
# This homework shamelessly demands you to implement DQN — an approximate Q-learning algorithm with experience replay and target networks — and see if it works any better this way.
#
# Original paper:
# https://arxiv.org/pdf/1312.5602.pdf

# **This notebook is the main notebook.** Another notebook is given for debug. (**homework_pytorch_main**). The tasks are similar and share most of the code. The main difference is in environments. In main notebook it can take some 2 hours for the agent to start improving so it seems reasonable to launch the algorithm on a simpler env first. In debug one it is CartPole and it will train in several minutes.
#
# **We suggest the following pipeline:** First implement debug notebook then implement the main one.
#
# **About evaluation:** All points are given for the main notebook with one exception: if agent fails to beat the threshold in main notebook you can get 1 pt (instead of 3 pts) for beating the threshold in debug notebook.

# In[1]:

import sys, os
if 'google.colab' in sys.modules and not os.path.exists('.setup_complete'):
    get_ipython().system('wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/setup_colab.sh -O- | bash')

    get_ipython().system('wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/atari_wrappers.py')
    get_ipython().system('wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/utils.py')
    get_ipython().system('wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/replay_buffer.py')
    get_ipython().system('wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/framebuffer.py')

    get_ipython().system('touch .setup_complete')

# This code creates a virtual display to draw game images on.
# It will have no effect if your machine has a monitor.
if type(os.environ.get("DISPLAY")) is not str or len(os.environ.get("DISPLAY")) == 0:
    get_ipython().system('bash ../xvfb start')
    os.environ['DISPLAY'] = ':1'

# __Frameworks__ - we'll accept this homework in any deep learning framework. This particular notebook was designed for PyTorch, but you find it easy to adapt it to almost any Python-based deep learning framework.

# In[2]:

import random
import numpy as np
import torch
import utils

# In[3]:

import gym
import numpy as np
import matplotlib.pyplot as plt
# %matplotlib inline

# ### Let's play some old videogames
# ![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/nerd.png)
#
# This time we're gonna apply approximate Q-learning to an Atari game called Breakout. It's not the hardest thing out there, but it's definitely way more complex than anything we tried before.
#

# In[4]:

ENV_NAME = "BreakoutNoFrameskip-v4"

# ## Preprocessing (3 pts)

# Let's see what observations look like.

# In[5]:

env = gym.make(ENV_NAME)
env.reset()

n_cols = 5
n_rows = 2
fig = plt.figure(figsize=(16, 9))

for row in range(n_rows):
    for col in range(n_cols):
        ax = fig.add_subplot(n_rows, n_cols, row * n_cols + col + 1)
        ax.imshow(env.render('rgb_array'))
        env.step(env.action_space.sample())
plt.show()

# **Let's play a little.**
#
# Pay attention to zoom and fps args of play function. Control: A, D, space.

# In[6]:

# # Does not work in Colab.
# # Use KeyboardInterrupt (Kernel → Interrupt in Jupyter) to continue.

# from gym.utils.play import play

# play(env=gym.make(ENV_NAME), zoom=5, fps=30)

# ### Processing game image
#
# Raw Atari images are large, 210x160x3 by default. However, we don't need that level of detail in order to learn from them.
#
# We can thus save a lot of time by preprocessing game image, including
# * Resizing to a smaller shape, 64x64
# * Converting to grayscale
# * Cropping irrelevant image parts (top, bottom and edges)
#
# Also please keep one dimension for channel so that final shape would be 1x64x64.
#
# Tip: You can implement your own grayscale converter and assign a huge weight to the red channel. This dirty trick is not necessary but it will speed up learning.

# In[7]:

from gym.core import ObservationWrapper
from gym.spaces import Box
import cv2

class PreprocessAtariObs(ObservationWrapper):
    def __init__(self, env):
        """A gym wrapper that crops, scales image into the desired shapes and grayscales it."""
        ObservationWrapper.__init__(self, env)

        self.img_size = (1, 64, 64)
        self.observation_space = Box(0.0, 1.0, self.img_size)

    def _to_gray_scale(self, rgb, channel_weights=[0.8, 0.1, 0.1]):
        return np.dot(rgb, channel_weights)[np.newaxis]

    def observation(self, img):
        """what happens to each observation"""

        # Here's what you need to do:
        #  * crop image, remove irrelevant parts
        #  * resize image to self.img_size
        #     (Use imresize from any library you want,
        #      e.g. opencv, PIL, keras. Don't use skimage.imresize
        #      because it is extremely slow.)
        #  * cast image to grayscale
        #  * convert image pixels to (0,1) range, float32 type
        img = img[50:, 10:-10]
        img = cv2.resize(img, (self.img_size[1:]))
        img = self._to_gray_scale(img).astype(np.float32)
        img = (img - img.min()) / (img.max() - img.min())
        return img

# In[8]:

import gym
# spawn game instance for tests
env = gym.make(ENV_NAME)  # create raw env
env = PreprocessAtariObs(env)
observation_shape = env.observation_space.shape
n_actions = env.action_space.n
env.reset()
obs, _, _, _ = env.step(env.action_space.sample())

# test observation
assert obs.ndim == 3, "observation must be [channel, h, w] even if there's just one channel"
assert obs.shape == observation_shape, obs.shape
assert obs.dtype == 'float32'
assert len(np.unique(obs)) > 2, "your image must not be binary"
assert 0 <= np.min(obs) and np.max(
    obs) <= 1, "convert image pixels to [0,1] range"

assert np.max(obs) >= 0.5, "It would be easier to see a brighter observation"
assert np.mean(obs) >= 0.1, "It would be easier to see a brighter observation"

print("Formal tests seem fine. Here's an example of what you'll get.")

n_cols = 5
n_rows = 2
fig = plt.figure(figsize=(16, 9))
obs = env.reset()
for row in range(n_rows):
    for col in range(n_cols):
        ax = fig.add_subplot(n_rows, n_cols, row * n_cols + col + 1)
        ax.imshow(obs[0, :, :], interpolation='none', cmap='gray')
        obs, _, _, _ = env.step(env.action_space.sample())
plt.show()

# ### Wrapping.

# **About the game:** You have 5 lives and get points for breaking the wall. Higher bricks cost more than the lower ones. There are 4 actions: start game (should be called at the beginning and after each life is lost), move left, move right and do nothing. There are some common wrappers used for Atari environments.

# In[9]:

import atari_wrappers

def PrimaryAtariWrap(env, clip_rewards=True):
    assert 'NoFrameskip' in env.spec.id

    # This wrapper holds the same action for <skip> frames and outputs
    # the maximal pixel value of 2 last frames (to handle blinking
    # in some envs)
    env = atari_wrappers.MaxAndSkipEnv(env, skip=4)

    # This wrapper sends done=True when each life is lost
    # (not all the 5 lives that are givern by the game rules).
    # It should make easier for the agent to understand that losing is bad.
    env = atari_wrappers.EpisodicLifeEnv(env)

    # This wrapper laucnhes the ball when an episode starts.
    # Without it the agent has to learn this action, too.
    # Actually it can but learning would take longer.
    env = atari_wrappers.FireResetEnv(env)

    # This wrapper transforms rewards to {-1, 0, 1} according to their sign
    if clip_rewards:
        env = atari_wrappers.ClipRewardEnv(env)

    # This wrapper is yours :)
    env = PreprocessAtariObs(env)
    return env

# **Let's see if the game is still playable after applying the wrappers.**
# At playing the EpisodicLifeEnv wrapper seems not to work but actually it does (because after when life finishes a new ball is dropped automatically - it means that FireResetEnv wrapper understands that a new episode began).

# In[10]:

# # Does not work in Colab.
# # Use KeyboardInterrupt (Kernel → Interrupt in Jupyter) to continue.

# from gym.utils.play import play

# def make_play_env():
#     env = gym.make(ENV_NAME)
#     env = PrimaryAtariWrap(env)
# # in PyTorch images have shape [c, h, w] instead of common [h, w, c]
#     env = atari_wrappers.AntiTorchWrapper(env)
#     return env

# play(make_play_env(), zoom=3, fps=5)

# ### Frame buffer
#
# Our agent can only process one observation at a time, so we gotta make sure it contains enough information to find optimal actions. For instance, agent has to react to moving objects so it must be able to measure object's velocity.
#
# To do so, we introduce a buffer that stores 4 last images. This time everything is pre-implemented for you, not really by the staff of the course :)

# In[11]:

from framebuffer import FrameBuffer

def make_env(clip_rewards=True, seed=None):
    env = gym.make(ENV_NAME)  # create raw env
    if seed is not None:
        env.seed(seed)
    env = PrimaryAtariWrap(env, clip_rewards)
    env = FrameBuffer(env, n_frames=4, dim_order='pytorch')
    return env

env = make_env()
env.reset()
n_actions = env.action_space.n
state_shape = env.observation_space.shape

# In[12]:

for _ in range(12):
    obs, _, _, _ = env.step(env.action_space.sample())

plt.figure(figsize=[12,10])
plt.title("Game image")
plt.imshow(env.render("rgb_array"))
plt.show()

plt.figure(figsize=[15,15])
plt.title("Agent observation (4 frames top to bottom)")
plt.imshow(utils.img_by_obs(obs, state_shape), cmap='gray')
plt.show()

# ## DQN as it is (4 pts)

# ### Building a network
#
# We now need to build a neural network that can map images to state q-values. This network will be called on every agent's step so it better not be resnet-152 unless you have an array of GPUs. Instead, you can use strided convolutions with a small number of features to save time and memory.
#
# You can build any architecture you want, but for reference, here's something that will more or less work:

# ![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/dqn_arch.png)

# **Dueling network: (+2 pts)**
# $$Q_{\theta}(s, a) = V_{\eta}(f_{\xi}(s)) + A_{\psi}(f_{\xi}(s), a) - \frac{\sum_{a'}A_{\psi}(f_{\xi}(s), a')}{N_{actions}},$$
# where $\xi$, $\eta$, and $\psi$ are, respectively, the parameters of the
# shared encoder $f_ξ$ , of the value stream $V_\eta$ , and of the advan
# tage stream $A_\psi$; and $\theta = \{\xi, \eta, \psi\}$ is their concatenation.
#
# For the architecture on the image $V$ and $A$ heads can follow the dense layer instead of $Q$. Please don't worry that the model becomes a little bigger.

# In[13]:

import torch
import torch.nn as nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# those who have a GPU but feel unfair to use it can uncomment:
# device = torch.device('cpu')
device

# In[14]:

def conv2d_size_out(size, kernel_size, stride):
    """
    common use case:
    cur_layer_img_w = conv2d_size_out(cur_layer_img_w, kernel_size, stride)
    cur_layer_img_h = conv2d_size_out(cur_layer_img_h, kernel_size, stride)
    to understand the shape for dense layer's input
    """
    return (size - (kernel_size - 1) - 1) // stride  + 1

# In[15]:

class DQNAgent(nn.Module):
    def __init__(self, state_shape, n_actions, epsilon=0):

        super().__init__()
        self.epsilon = epsilon
        self.n_actions = n_actions
        self.state_shape = state_shape

        # Define your network body here. Please make sure agent is fully contained here
        # nn.Flatten() can be useful
        self.conv_part = nn.Sequential(nn.Conv2d(4, 16, 3, padding=1, stride=2),
                                       nn.ReLU(),
                                       nn.Conv2d(16, 32, 3, padding=1, stride=2),
                                       nn.ReLU(),
                                       nn.Conv2d(32, 64, 3, padding=1, stride=2),
                                       nn.ReLU()
                                       )
        self.fc_part = nn.Sequential(nn.Linear(4096, 256),
                                     nn.ReLU(),
                                     nn.Linear(256, self.n_actions)
                                     )

    def forward(self, state_t):
        """
        takes agent's observation (tensor), returns qvalues (tensor)
        :param state_t: a batch of 4-frame buffers, shape = [batch_size, 4, h, w]
        """
        # Use your network to compute qvalues for given state
        features = self.conv_part(state_t)
        flat = features.reshape(state_t.shape[0], -1)
        qvalues = self.fc_part(flat)

        assert qvalues.requires_grad, "qvalues must be a torch tensor with grad"
        assert (
            len(qvalues.shape) == 2 and
            qvalues.shape[0] == state_t.shape[0] and
            qvalues.shape[1] == n_actions
        )

        return qvalues

    def get_qvalues(self, states):
        """
        like forward, but works on numpy arrays, not tensors
        """
        model_device = next(self.parameters()).device
        states = torch.tensor(states, device=model_device, dtype=torch.float32)
        qvalues = self.forward(states)
        return qvalues.data.cpu().numpy()

    def sample_actions(self, qvalues):
        """pick actions given qvalues. Uses epsilon-greedy exploration strategy. """
        epsilon = self.epsilon
        batch_size, n_actions = qvalues.shape

        random_actions = np.random.choice(n_actions, size=batch_size)
        best_actions = qvalues.argmax(axis=-1)

        should_explore = np.random.choice(
            [0, 1], batch_size, p=[1-epsilon, epsilon])
        return np.where(should_explore, random_actions, best_actions)

# In[16]:

agent = DQNAgent(state_shape, n_actions, epsilon=0.5).to(device)

# Now let's try out our agent to see if it raises any errors.

# In[17]:

def evaluate(env, agent, n_games=1, greedy=False, t_max=10000):
    """ Plays n_games full games. If greedy, picks actions as argmax(qvalues). Returns mean reward. """
    rewards = []
    for _ in range(n_games):
        s = env.reset()
        reward = 0
        for _ in range(t_max):
            qvalues = agent.get_qvalues([s])
            action = qvalues.argmax(axis=-1)[0] if greedy else agent.sample_actions(qvalues)[0]
            s, r, done, _ = env.step(action)
            reward += r
            if done:
                break

        rewards.append(reward)
    return np.mean(rewards)

# In[18]:

evaluate(env, agent, n_games=1)

# ### Experience replay
# For this assignment, we provide you with experience replay buffer. If you implemented experience replay buffer in last week's assignment, you can copy-paste it here **to get 2 bonus points**.
#
# ![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/exp_replay.png)

# #### The interface is fairly simple:
# * `exp_replay.add(obs, act, rw, next_obs, done)` - saves (s,a,r,s',done) tuple into the buffer
# * `exp_replay.sample(batch_size)` - returns observations, actions, rewards, next_observations and is_done for `batch_size` random samples.
# * `len(exp_replay)` - returns number of elements stored in replay buffer.

# In[19]:

from replay_buffer import ReplayBuffer
exp_replay = ReplayBuffer(10)

for _ in range(30):
    exp_replay.add(env.reset(), env.action_space.sample(), 1.0, env.reset(), done=False)

obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(5)

assert len(exp_replay) == 10, "experience replay size should be 10 because that's what maximum capacity is"

# In[20]:

def play_and_record(initial_state, agent, env, exp_replay, n_steps=1):
    """
    Play the game for exactly n_steps, record every (s,a,r,s', done) to replay buffer.
    Whenever game ends, add record with done=True and reset the game.
    It is guaranteed that env has done=False when passed to this function.

    PLEASE DO NOT RESET ENV UNLESS IT IS "DONE"

    :returns: return sum of rewards over time and the state in which the env stays
    """
    s = initial_state
    sum_rewards = 0

    # Play the game for n_steps as per instructions above
    for step in range(n_steps):
        q = agent.get_qvalues([s])
        a = agent.sample_actions(q)[0]
        sp, r, done, _ = env.step(a)
        exp_replay.add(s, a, r, sp, done)
        if done:
            s = env.reset()
        else:
            s = sp

    return sum_rewards, s

# In[21]:

# testing your code.
exp_replay = ReplayBuffer(2000)

state = env.reset()
play_and_record(state, agent, env, exp_replay, n_steps=1000)

# if you're using your own experience replay buffer, some of those tests may need correction.
# just make sure you know what your code does
assert len(exp_replay) == 1000,     "play_and_record should have added exactly 1000 steps, "     "but instead added %i" % len(exp_replay)
is_dones = list(zip(*exp_replay._storage))[-1]

assert 0 < np.mean(is_dones) < 0.1,     "Please make sure you restart the game whenever it is 'done' and "     "record the is_done correctly into the buffer. Got %f is_done rate over "     "%i steps. [If you think it's your tough luck, just re-run the test]" % (
        np.mean(is_dones), len(exp_replay))

for _ in range(100):
    obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(10)
    assert obs_batch.shape == next_obs_batch.shape == (10,) + state_shape
    assert act_batch.shape == (10,),         "actions batch should have shape (10,) but is instead %s" % str(act_batch.shape)
    assert reward_batch.shape == (10,),         "rewards batch should have shape (10,) but is instead %s" % str(reward_batch.shape)
    assert is_done_batch.shape == (10,),         "is_done batch should have shape (10,) but is instead %s" % str(is_done_batch.shape)
    assert [int(i) in (0, 1) for i in is_dones],         "is_done should be strictly True or False"
    assert [0 <= a < n_actions for a in act_batch], "actions should be within [0, n_actions)"

print("Well done!")

# ### Target networks
#
# We also employ the so called "target network" - a copy of neural network weights to be used for reference Q-values:
#
# The network itself is an exact copy of agent network, but it's parameters are not trained. Instead, they are moved here from agent's actual network every so often.
#
# $$ Q_{reference}(s,a) = r + \gamma \cdot \max _{a'} Q_{target}(s',a') $$
#
# ![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/target_net.png)

# In[22]:

target_network = DQNAgent(agent.state_shape, agent.n_actions, epsilon=0.5).to(device)
# This is how you can load weights from agent into target network
target_network.load_state_dict(agent.state_dict())

# ### Learning with... Q-learning
# Here we write a function similar to `agent.update` from tabular q-learning.

# Compute Q-learning TD error:
#
# $$ L = { 1 \over N} \sum_i [ Q_{\theta}(s,a) - Q_{reference}(s,a) ] ^2 $$
#
# With Q-reference defined as
#
# $$ Q_{reference}(s,a) = r(s,a) + \gamma \cdot max_{a'} Q_{target}(s', a') $$
#
# Where
# * $Q_{target}(s',a')$ denotes Q-value of next state and next action predicted by __target_network__
# * $s, a, r, s'$ are current state, action, reward and next state respectively
# * $\gamma$ is a discount factor defined two cells above.
#
#
# __Note 1:__ there's an example input below. Feel free to experiment with it before you write the function.
#
# __Note 2:__ compute_td_loss is a source of 99% of bugs in this homework. If reward doesn't improve, it often helps to go through it line by line [with a rubber duck](https://rubberduckdebugging.com/).
#
# **Double DQN (+2 pts)**
#
# $$ Q_{reference}(s,a) = r(s, a) + \gamma \cdot
# Q_{target}(s',argmax_{a'}Q_\theta(s', a')) $$

# In[35]:

def compute_td_loss(states, actions, rewards, next_states, is_done,
                    agent, target_network,
                    gamma=0.99,
                    check_shapes=False,
                    device=device):
    """ Compute td loss using torch operations only. Use the formulae above. """
    states = torch.tensor(states, device=device, dtype=torch.float32)    # shape: [batch_size, *state_shape]
    actions = torch.tensor(actions, device=device, dtype=torch.int64)    # shape: [batch_size]
    rewards = torch.tensor(rewards, device=device, dtype=torch.float32)  # shape: [batch_size]
    # shape: [batch_size, *state_shape]
    next_states = torch.tensor(next_states, device=device, dtype=torch.float)
    is_done = torch.tensor(
        is_done.astype('float32'),
        device=device,
        dtype=torch.float32,
    )  # shape: [batch_size]
    is_not_done = 1 - is_done

    # get q-values for all actions in current states
    predicted_qvalues = agent(states)  # shape: [batch_size, n_actions]
    double_dqn_agent_next_actions = agent(next_states).detach().argmax(dim=1, keepdim=True)

    # compute q-values for all actions in next states
    predicted_next_qvalues = target_network(next_states)  # shape: [batch_size, n_actions]

    # select q-values for chosen actions
    predicted_qvalues_for_actions = predicted_qvalues[range(len(actions)), actions]  # shape: [batch_size]

    # compute V*(next_states) using predicted next q-values
    next_state_values = torch.gather(predicted_next_qvalues, dim=1,
                                     index=double_dqn_agent_next_actions).detach().flatten()

    assert next_state_values.dim() == 1 and next_state_values.shape[0] == states.shape[0],         "must predict one value per state"

    # compute "target q-values" for loss - it's what's inside square parentheses in the above formula.
    # at the last state use the simplified formula: Q(s,a) = r(s,a) since s' doesn't exist
    # you can multiply next state values by is_not_done to achieve this.
    target_qvalues_for_actions = rewards + gamma * next_state_values

    # mean squared error loss to minimize
    loss = torch.mean((predicted_qvalues_for_actions - target_qvalues_for_actions.detach()) ** 2)

    if check_shapes:
        assert predicted_next_qvalues.data.dim() == 2,             "make sure you predicted q-values for all actions in next state"
        assert next_state_values.data.dim() == 1,             "make sure you computed V(s') as maximum over just the actions axis and not all axes"
        assert target_qvalues_for_actions.data.dim() == 1,             "there's something wrong with target q-values, they must be a vector"

    return loss

# Sanity checks

# In[36]:

obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(10)

loss = compute_td_loss(obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch,
                       agent, target_network,
                       gamma=0.99, check_shapes=True)
loss.backward()

assert loss.requires_grad and tuple(loss.data.size()) == (),     "you must return scalar loss - mean over batch"
assert np.any(next(agent.parameters()).grad.data.cpu().numpy() != 0),     "loss must be differentiable w.r.t. network weights"
assert np.all(next(target_network.parameters()).grad is None),     "target network should not have grads"

# ## Main loop (3 pts)
#
# **If deadline is tonight and it has not converged:** It is ok. Send the notebook today and when it converges send it again.
# If the code is exactly the same points will not be discounted.
#
# It's time to put everything together and see if it learns anything.

# In[37]:

from tqdm import trange
from IPython.display import clear_output
import matplotlib.pyplot as plt

# In[38]:

seed = int(1e9+7)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

# In[39]:

env = make_env(seed)
state_shape = env.observation_space.shape
n_actions = env.action_space.n
state = env.reset()

agent = DQNAgent(state_shape, n_actions, epsilon=1).to(device)
target_network = DQNAgent(state_shape, n_actions).to(device)
target_network.load_state_dict(agent.state_dict())

# Buffer of size $10^4$ fits into 5 Gb RAM.
#
# Larger sizes ($10^5$ and $10^6$ are common) can be used. It can improve the learning, but $10^4$ is quite enough. $10^2$ will probably fail learning.

# In[40]:

REPLAY_BUFFER_SIZE = 10**4
N_STEPS = 100

exp_replay = ReplayBuffer(REPLAY_BUFFER_SIZE)
for i in range(REPLAY_BUFFER_SIZE // N_STEPS):
    if not utils.is_enough_ram(min_available_gb=0.1):
        print("""
            Less than 100 Mb RAM available.
            Make sure the buffer size in not too huge.
            Also check, maybe other processes consume RAM heavily.
            """
             )
        break
    play_and_record(state, agent, env, exp_replay, n_steps=N_STEPS)
    if len(exp_replay) == REPLAY_BUFFER_SIZE:
        break
print(len(exp_replay))

# In[41]:

timesteps_per_epoch = 1
batch_size = 16
total_steps = 3 * 10**6
decay_steps = 10**6

opt = torch.optim.Adam(agent.parameters(), lr=1e-4)

init_epsilon = 1
final_epsilon = 0.1

loss_freq = 50
refresh_target_network_freq = 5000
eval_freq = 5000

max_grad_norm = 50

n_lives = 5

# In[42]:

mean_rw_history = []
td_loss_history = []
grad_norm_history = []
initial_state_v_history = []
step = 0

# In[43]:

import time

def wait_for_keyboard_interrupt():
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        pass

# In[47]:

state_dict = torch.load('model-correct-argmax.pth')
agent.load_state_dict(state_dict['model'])
opt.load_state_dict(state_dict['opt'])
step = state_dict['step']
init_epsilon = agent.epsilon

# In[48]:

state = env.reset()
with trange(step, total_steps + 1) as progress_bar:
    for step in progress_bar:
        if not utils.is_enough_ram():
            print('less that 100 Mb RAM available, freezing')
            print('make sure everything is ok and use KeyboardInterrupt to continue')
            wait_for_keyboard_interrupt()

        agent.epsilon = utils.linear_decay(init_epsilon, final_epsilon, step, decay_steps)

        # play
        _, state = play_and_record(state, agent, env, exp_replay, timesteps_per_epoch)

        # train
        s, a, rw, ns, done = exp_replay.sample(batch_size)

        loss = compute_td_loss(s, a, rw, ns, done, agent, target_network)

        loss.backward()
        grad_norm = nn.utils.clip_grad_norm_(agent.parameters(), max_grad_norm)
        opt.step()
        opt.zero_grad()

        if step % loss_freq == 0:
            td_loss_history.append(loss.data.cpu().item())
            grad_norm_history.append(grad_norm)

        if step % refresh_target_network_freq == 0:
            # Load agent weights into target_network
            target_network.load_state_dict(agent.state_dict())

        if step % eval_freq == 0:
            mean_rw_history.append(evaluate(
                make_env(clip_rewards=True, seed=step), agent, n_games=3 * n_lives, greedy=True)
            )
            initial_state_q_values = agent.get_qvalues(
                [make_env(seed=step).reset()]
            )
            initial_state_v_history.append(np.max(initial_state_q_values))

            clear_output(True)
            print("buffer size = %i, epsilon = %.5f" %
                (len(exp_replay), agent.epsilon))

            plt.figure(figsize=[16, 9])

            plt.subplot(2, 2, 1)
            plt.title("Mean reward per life")
            plt.plot(mean_rw_history)
            plt.grid()

            assert not np.isnan(td_loss_history[-1])
            plt.subplot(2, 2, 2)
            plt.title("TD loss history (smoothened)")
            plt.plot(utils.smoothen(td_loss_history))
            plt.grid()

            plt.subplot(2, 2, 3)
            plt.title("Initial state V")
            plt.plot(initial_state_v_history)
            plt.grid()

            plt.subplot(2, 2, 4)
            plt.title("Grad norm history (smoothened)")
            plt.plot(utils.smoothen(grad_norm_history))
            plt.grid()

            plt.show()

# Agent is evaluated for 1 life, not for a whole episode of 5 lives. Rewards in evaluation are also truncated. Cuz this is what environment the agent is learning in and in this way mean rewards per life can be compared with initial state value
#
# **The goal is to get 15 points in the real env**. So 3 or better 4 points in the preprocessed one will probably be enough. You can interrupt learning then.

# Final scoring is done on a whole episode with all 5 lives.

# In[49]:

torch.save({'model': agent.state_dict(),
            'opt': opt.state_dict(),
            'step': step}, 'model-correct-argmax.pth')

# In[51]:

final_score = evaluate(
  make_env(clip_rewards=False, seed=9),
    agent, n_games=30, greedy=True, t_max=10 * 1000
) * n_lives
print('final score:', final_score)
assert final_score >= 15, 'not as cool as DQN can'
print('Cool!')

# ## How to interpret plots:
#
# This aint no supervised learning so don't expect anything to improve monotonously.
# * **TD loss** is the MSE between agent's current Q-values and target Q-values. It may slowly increase or decrease, it's ok. The "not ok" behavior includes going NaN or stayng at exactly zero before agent has perfect performance.
# * **grad norm** just shows the intensivity of training. Not ok is growing to values of about 100 (or maybe even 50) though it depends on network architecture.
# * **mean reward** is the expected sum of r(s,a) agent gets over the full game session. It will oscillate, but on average it should get higher over time (after a few thousand iterations...).
#  * In basic q-learning implementation it takes about 40k steps to "warm up" agent before it starts to get better.
# * **Initial state V** is the expected discounted reward for episode in the oppinion of the agent. It should behave more smoothly than **mean reward**. It should get higher over time but sometimes can experience drawdowns because of the agaent's overestimates.
# * **buffer size** - this one is simple. It should go up and cap at max size.
# * **epsilon** - agent's willingness to explore. If you see that agent's already at 0.01 epsilon before it's average reward is above 0 - it means you need to increase epsilon. Set it back to some 0.2 - 0.5 and decrease the pace at which it goes down.
# * Smoothing of plots is done with a gaussian kernel
#
# At first your agent will lose quickly. Then it will learn to suck less and at least hit the ball a few times before it loses. Finally it will learn to actually score points.
#
# **Training will take time.** A lot of it actually. Probably you will not see any improvment during first **150k** time steps (note that by default in this notebook agent is evaluated every 5000 time steps).
#
# But hey, long training time isn't _that_ bad:
# ![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/training.png)

# ## About hyperparameters:
#
# The task has something in common with supervised learning: loss is optimized through the buffer (instead of Train dataset). But the distribution of states and actions in the buffer **is not stationary** and depends on the policy that generated it. It can even happen that the mean TD error across the buffer is very low but the performance is extremely poor (imagine the agent collecting data to the buffer always manages to avoid the ball).
#
# * Total timesteps and training time: It seems to be so huge, but actually it is normal for RL.
#
# * $\epsilon$ decay shedule was taken from the original paper and is like traditional for epsilon-greedy policies. At the beginning of the training the agent's greedy policy is poor so many random actions should be taken.
#
# * Optimizer: In the original paper RMSProp was used (they did not have Adam in 2013) and it can work not worse than Adam. For us Adam was default and it worked.
#
# * lr: $10^{-3}$ would probably be too huge
#
# * batch size: This one can be very important: if it is too small the agent can fail to learn. Huge batch takes more time to process. If batch of size 8 can not be processed on the hardware you use take 2 (or even 4) batches of size 4, divide the loss on them by 2 (or 4) and make optimization step after both backward() calls in torch.
#
# * target network update frequency: has something in common with learning rate. Too frequent updates can lead to divergence. Too rare can lead to slow leraning. For millions of total timesteps thousands of inner steps seem ok. One iteration of target network updating is an iteration of the (this time approximate) $\gamma$-compression that stands behind Q-learning. The more inner steps it makes the more accurate is the compression.
# * max_grad_norm - just huge enough. In torch clip_grad_norm also evaluates the norm before clipping and it can be convenient for logging.

# ### Video

# In[54]:

# Record sessions

import gym.wrappers

with gym.wrappers.Monitor(make_env(), directory="videos", force=True) as env_monitor:
    sessions = [evaluate(env_monitor, agent, n_games=n_lives, greedy=True) for _ in range(10)]

# In[55]:

# Show video. This may not work in some setups. If it doesn't
# work for you, you can download the videos and view them locally.

from pathlib import Path
from base64 import b64encode
from IPython.display import HTML

video_paths = sorted([s for s in Path('videos').iterdir() if s.suffix == '.mp4'])
video_path = video_paths[-1]  # You can also try other indices

if 'google.colab' in sys.modules:
    # https://stackoverflow.com/a/57378660/1214547
    with video_path.open('rb') as fp:
        mp4 = fp.read()
    data_url = 'data:video/mp4;base64,' + b64encode(mp4).decode()
else:
    data_url = str(video_path)

HTML("""
<video width="640" height="480" controls>
  <source src="{}" type="video/mp4">
</video>
""".format(data_url))

# ## Let's have a closer look at this.
#
# If average episode score is below 200 using all 5 lives, then probably DQN has not converged fully. But anyway let's make a more complete record of an episode.

# In[ ]:

eval_env = make_env(clip_rewards=False)
record = utils.play_and_log_episode(eval_env, agent)
print('total reward for life:', np.sum(record['rewards']))
for key in record:
    print(key)

# In[ ]:

fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(1, 1, 1)

ax.scatter(record['v_mc'], record['v_agent'])
ax.plot(sorted(record['v_mc']), sorted(record['v_mc']),
       'black', linestyle='--', label='x=y')

ax.grid()
ax.legend()
ax.set_title('State Value Estimates')
ax.set_xlabel('Monte-Carlo')
ax.set_ylabel('Agent')

plt.show()

# $\hat V_{Monte-Carlo}(s_t) = \sum_{\tau=0}^{episode~end} \gamma^{\tau-t}r_t$

# Is there a big bias? It's ok, anyway it works.

# ## Bonus I (2 pts)

# **1.** Plot several (say 3) states with high and low spreads of Q estimate by actions i.e.
# $$\max_a \hat Q(s,a) - \min_a \hat Q(s,a)\$$
# Please take those states from different episodes to make sure that the states are really different.
#
# What should high and low spread mean at least in the world of perfect Q-fucntions?
#
# Comment the states you like most.
#
# **2.** Plot several (say 3) states with high td-error and several states with high values of
# $$| \hat V_{Monte-Carlo}(s) - \hat V_{agent}(s)|,$$
# $$\hat V_{agent}(s)=\max_a \hat Q(s,a).$$ Please take those states from different episodes to make sure that the states are really different. From what part (i.e. beginning, middle, end) of an episode did these states come from?
#
# Comment the states you like most.

# In[ ]:

from utils import play_and_log_episode, img_by_obs

<YOUR CODE>

# ## Bonus II (1-5 pts). Get High Score!
#
# 1 point to you for each 50 points of your agent. Truncated by 5 points. Starting with 50 points, **not** 50 + threshold.
#
# One way is to train for several days and use heavier hardware (why not actually).
#
# Another way is to apply modifications (see **Bonus III**).

# ## Bonus III (2+ pts). Apply modifications to DQN.
#
# For inspiration see [Rainbow](https://arxiv.org/abs/1710.02298) - a version of q-learning that combines lots of them.
#
# Points for Bonus II and Bonus III fully stack. So if modified agent gets score 250+ you get 5 pts for Bonus II + points for modifications. If the final score is 40 then you get the points for modifications.
#
#
# Some modifications:
# * [Prioritized experience replay](https://arxiv.org/abs/1511.05952) (5 pts for your own implementation, 3 pts for using a ready one)
# * [double q-learning](https://arxiv.org/abs/1509.06461) (2 pts)
# * [dueling q-learning](https://arxiv.org/abs/1511.06581) (2 pts)
# * multi-step heuristics (see [Rainbow](https://arxiv.org/abs/1710.02298)) (3 pts)
# * [Noisy Nets](https://arxiv.org/abs/1706.10295) (3 pts)
# * [distributional RL](https://arxiv.org/abs/1707.06887)(distributional and distributed stand for different things here) (5 pts)
# * Other modifications (2+ pts depending on complexity)

# ## Bonus IV (4+ pts). Distributed RL.
#
# Solve the task in a distributed way. It can strongly speed up learning. See [article](https://arxiv.org/pdf/1602.01783.pdf) or some guides.

# **As usual bonus points for all the tasks fully stack.**

# In[ ]:


# used for manipulating directory paths
import os

# Scientific and vector computation for python
import numpy as np

# Plotting library
from matplotlib import pyplot
from mpl_toolkits.mplot3d import Axes3D  # needed to plot 3-D surfaces

def plotData(x, y):
    """
    Plots the data points x and y into a new figure. Plots the data
    points and gives the figure axes labels of population and profit.

    Parameters
    ----------
    x : array_like
        Data point values for x-axis.

    y : array_like
        Data point values for y-axis. Note x and y should have the same size.

    Instructions
    ------------
    Plot the training data into a figure using the "figure" and "plot"
    functions. Set the axes labels using the "xlabel" and "ylabel" functions.
    Assume the population and revenue data have been passed in as the x
    and y arguments of this function.

    Hint
    ----
    You can use the 'ro' option with plot to have the markers
    appear as red circles. Furthermore, you can make the markers larger by
    using plot(..., 'ro', ms=10), where `ms` refers to marker size. You
    can also set the marker edge color using the `mec` property.
    """

    fig = pyplot.figure()

    pyplot.plot(x, y, 'ro', ms=10, mec='k')
    pyplot.xlabel("Population in 10000s")
    pyplot.ylabel("Revenue in 10000$")
    pyplot.show()

if __name__ == '__main__':
    data = np.loadtxt(os.path.join("/home/thelichking/Desktop/ml-coursera-python-assignments/Exercise1/Data",
                                   "ex1data1.txt"), delimiter=',')
    x, y = data[:, 0], data[:, 1]
    m = y.size

    plotData(x, y)


"""
Introduction to Radar Course

Authors
=======

Zachary Chance, Robert Freking, Victoria Helus
MIT Lincoln Laboratory
Lexington, MA 02421

Distribution Statement
======================

DISTRIBUTION STATEMENT A. Approved for public release. Distribution is unlimited.

This material is based upon work supported by the United States Air Force under Air
Force Contract No. FA8702-15-D-0001. Any opinions, findings, conclusions or
recommendations expressed in this material are those of the author(s) and do not
necessarily reflect the views of the United States Air Force.

© 2021 Massachusetts Institute of Technology.

The software/firmware is provided to you on an As-Is basis

Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part
252.227-7013 or 7014 (Feb 2014). Notwithstanding any copyright notice, U.S.
Government rights in this work are defined by DFARS 252.227-7013 or
DFARS 252.227-7014 as detailed above. Use of this work other than as specifically
authorized by the U.S. Government may violate any copyrights that exist in this work.

RAMS ID: 1016938
"""

from IPython.display import display
import ipywidgets as wdg
import math
import numpy as np
from rad.const import k, c
import rad.radar as rd

# jupyter_intro

# Simple quiz with scalar value
def new_quiz(prompt, val, abs_tol=None, rel_tol=1E-2):
    """
    Generate quiz answer box.

    Inputs:
    - prompt [str]: Question prompt
    - val [float]: True answer
    - abs_tol [float]: Absolute tolerance; if None, uses rel_tol
    - rel_tol [float]: Relative tolerance

    Outputs:
    (none)
    """
    title = wdg.HTML(value = f"
{prompt}
")
    answer = wdg.FloatText()
    submit = wdg.Button(description="Submit")
    result = wdg.HTML(value = f"Ready")
    ansbox = wdg.VBox([wdg.HBox([title, answer, result]), submit])
    display(ansbox)

    def check_ans(b):
        if abs_tol and (abs(answer.value - val) < abs_tol):
            result.value = f"Correct!"
        elif rel_tol and ((abs(answer.value - val)/abs(val)) < rel_tol):
            result.value = f"Correct!"
        else:
            result.value = f"Incorrect."

    submit.on_click(check_ans)

#-------Lab 1.1: Introduction to Labs-------

# Q1.1.1
def quiz_1_1_1():
    prompt = 'Enter answer:'
    val = 1.0
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.2
def quiz_1_1_2():
    prompt = 'Enter answer:'
    val = (math.log10(2.72**3) + math.cos(4*math.pi/7))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.3
def quiz_1_1_3():
    prompt = 'Enter answer:'
    val = 7.76E4/5.1E-3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.4a
def quiz_1_1_4a():
    prompt = 'Enter answer (in dB):'
    val = rd.to_db(1.5E5)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.4b
def quiz_1_1_4b():
    prompt = 'Enter answer (in dB):'
    val = rd.to_db(7.2E-7)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.5a
def quiz_1_1_5a():
    prompt = 'Enter answer:'
    val = rd.from_db(51.2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.5b
def quiz_1_1_5b():
    prompt = 'Enter answer:'
    val = rd.from_db(-20.1)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.6
def quiz_1_1_6():
    prompt = 'Enter answer (in dB):'
    val = rd.to_db(3.7**5)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.7a
def quiz_1_1_7a():
    prompt = 'Enter answer:'
    val = 3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.1.7b
def quiz_1_1_7b():
    prompt = 'Enter answer (in s, within ±0.05 s):'
    val = 0.5
    tol = 0.05
    new_quiz(prompt, val, abs_tol=tol)

# Q1.1.7c
def quiz_1_1_7c():
    prompt = 'Enter answer (in Hz, within ±0.2 Hz):'
    val = 2
    tol = 0.2
    new_quiz(prompt, val, abs_tol=tol)

#-------Lab 1.2: Introduction to Radar-------

# Q1.2.1a
def quiz_1_2_1a():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(1000, propvel=1000)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.1b
def quiz_1_2_1b():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(1000, propvel=2000)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.1c
def quiz_1_2_1c():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(500, propvel=1000)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.2a
def quiz_1_2_2a():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(500E6, propvel=3E8)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.2b
def quiz_1_2_2b():
    prompt = 'Enter frequency (in Hz):'
    val = 1500/5
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.3a
def quiz_1_2_3a():
    prompt = 'Enter time (in ms):'
    val = math.sqrt(75**2 + 80**2)
    tol = 3
    new_quiz(prompt, val, abs_tol=tol)

# Q1.2.3b
def quiz_1_2_3b():
    prompt = 'Enter time (in ms):'
    val = 2*math.sqrt(75**2 + 80**2)
    tol = 3
    new_quiz(prompt, val, abs_tol=tol)

# Q1.2.4a
def quiz_1_2_4a():
    prompt = 'Enter time (in s):'
    val = 2*1000E3/3E8
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.4b
def quiz_1_2_4b():
    prompt = 'Enter range (in m):'
    val = 3E8*0.0057/2
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.5
def quiz_1_2_5():
    prompt = 'Enter range (in m):'
    val = math.sqrt(75**2 + 100**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.6a
def quiz_1_2_6a():
    prompt = 'Enter range (in m):'
    val = math.sqrt(50**2 + 50**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.6b
def quiz_1_2_6b():
    prompt = 'Enter azimuth (in deg, within +/- 3 deg):'
    val = 90 - rd.rad2deg(math.atan2(50, 50))
    tol = 3
    new_quiz(prompt, val, abs_tol=tol)

# Q1.2.7a
def quiz_1_2_7a():
    prompt = 'Enter range (in m):'
    val = math.sqrt(40**2 + 20**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.7b
def quiz_1_2_7b():
    prompt = 'Enter azimuth (in deg, within +/- 3 deg):'
    val = 90 - rd.rad2deg(math.atan2(-20, 40))
    tol = 3
    new_quiz(prompt, val, abs_tol=tol)

# Q1.2.8a
def quiz_1_2_8a():
    prompt = 'Enter beamwidth (in deg):'
    val = 70*0.12/3.8
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.8b
def quiz_1_2_8b():
    prompt = 'Enter diameter (in m):'
    val = 70*rd.wavelen(10E9, propvel=3E8)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.9
def quiz_1_2_9():
    prompt = 'Enter transmit gain (in dB):'
    val = rd.to_db(4*math.pi*2.1*3.3/rd.wavelen(15E3, propvel=2E3)**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.10a
def quiz_1_2_10a():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(5E9)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.10b
def quiz_1_2_10b():
    prompt = 'Enter beamwidth (in deg):'
    val = 70*rd.wavelen(5E9)/4.4
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q1.2.10c
def quiz_1_2_10c():
    prompt = 'Enter transmit gain (in dB):'
    val = rd.to_db(4*math.pi*(math.pi*2.2**2)/rd.wavelen(5E9)**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

#-------Lab 2.1: Radar Range Equation-------

# Q2.1.1a
def quiz_2_1_1a():
    prompt = 'Enter SNR (in dB):'
    val = rd.to_db(100/5.3)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.1b
def quiz_2_1_1b():
    prompt = 'Enter signal energy (in J):'
    val = rd.from_db(15)*2.2
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.2a
def quiz_2_1_2a():
    prompt = 'Enter received power (in W):'
    val = rd.friis(50E3, 10E3, area=10, gain=rd.from_db(35.0))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.2b
def quiz_2_1_2b():
    prompt = 'Enter dish radius (in m):'
    val = np.sqrt(rd.gain2area((1/100E3/5)*4*np.pi*(10E3)**2, 3.5E9)/np.pi)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.3a
def quiz_2_1_3a():
    prompt = 'Enter incident power (in W):'
    val = rd.friis(50E3, 150E3, area=rd.from_db(5), gain=rd.from_db(30))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.3b
def quiz_2_1_3b():
    prompt = 'Enter range (in m):'
    val = math.sqrt(500E3*1*rd.from_db(25)/4/math.pi/0.1E-3)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.4a
def quiz_2_1_4a():
    prompt = 'Enter received power (in W):'
    val = rd.rx_power(100E3, 600E3, 3E9, rcs=rd.from_db(5), gain=rd.from_db(80))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.4b
def quiz_2_1_4b():
    prompt = 'Enter effective aperture area (in m^2):'
    val = rd.gain2area(rd.from_db(24.3), 1.3E9)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.4c
def quiz_2_1_4c():
    prompt = 'Enter RCS (in dBsm):'
    wlen = rd.wavelen(1.5E9)
    val = rd.to_db(10E-12*(10E3**4)*(4*np.pi)**3/30/rd.from_db(80)/wlen**2)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.5
def quiz_2_1_5():
    prompt = 'Enter noise energy (in J):'
    val = k*722
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.6
def quiz_2_1_6():
    prompt = 'Enter SNR (in dB):'
    val = rd.to_db(rd.rx_snr(200E3, 50, 5E9, 750, rcs=rd.from_db(-5), gain=rd.from_db(70)))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q2.1.7
def quiz_2_1_7():
    prompt = 'Enter SNR (in dB):'
    val = rd.to_db(rd.snrconv(rd.from_db(17), 75E3, rd.from_db(-10), 50E3, rd.from_db(-15)))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

#-------Lab 2.2: Basic Radar Design-------

#-------Lab 3.1: Radar Transmissions and Receptions-------

# Q3.1.1
def quiz_3_1_1():
    prompt = 'Enter pulsewidth (in µs):'
    val = 7.2
    tol = 0.1
    new_quiz(prompt, val, abs_tol=tol)

# Q3.1.2a
def quiz_3_1_2a():
    prompt = 'Enter PRI (in s):'
    val = 1/10E3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q3.1.2b
def quiz_3_1_2b():
    prompt = 'Enter duty cycle:'
    val = 150E-6*1E3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q3.1.3a
def quiz_3_1_3a():
    prompt = 'Enter delay (in ns):'
    val = -2.83
    tol = 0.01
    new_quiz(prompt, val, abs_tol=tol)

# Q3.1.3b
def quiz_3_1_3b():
    prompt = 'Enter delay (in ns):'
    val = -2.83
    tol = 0.01
    new_quiz(prompt, val, abs_tol=tol)

# Q3.1.4a
def quiz_3_1_4a():
    prompt = 'Enter phase (in deg):'
    val = 168
    tol = 1
    new_quiz(prompt, val, abs_tol=tol)

# Q3.1.4b
def quiz_3_1_4b():
    prompt = 'Enter phase (in deg):'
    val = 120
    tol = 1
    new_quiz(prompt, val, abs_tol=tol)

# Q3.1.5
def quiz_3_1_5():
    prompt = 'Enter bandwidth (in Hz):'
    val = c/2/10
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

#-------Lab 3.2: Detection-------

# Q3.2.1a
def quiz_3_2_1a():
    prompt = 'Enter prob. of detection:'
    val = 0.69
    tol = 0.02
    new_quiz(prompt, val, abs_tol=tol)

# Q3.2.1b
def quiz_3_2_1b():
    prompt = 'Enter prob. of false alarm:'
    val = 0.17
    tol = 0.02
    new_quiz(prompt, val, abs_tol=tol)

#-------Lab 4.1: Target Parameter Estimation-------

# Q4.1.1a
def quiz_4_1_1a():
    prompt = 'Enter range (in m):'
    val = c*0.00333/2
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.1b
def quiz_4_1_1b():
    prompt = 'Enter range resolution (in m):'
    val = rd.range_res(30E6)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.1c
def quiz_4_1_1c():
    prompt = 'Enter range accuracy (in m):'
    val = rd.range_res(30E6)/np.sqrt(rd.from_db(11))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.2a
def quiz_4_1_2a():
    prompt = 'Enter angle accuracy (in deg):'
    val = rd.dish_beamw(2, 3.5E9)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.2b
def quiz_4_1_2b():
    prompt = 'Enter angle accuracy (in deg):'
    val = rd.dish_beamw(2, 3.5E9)/np.sqrt(rd.from_db(14))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.3
def quiz_4_1_3():
    prompt = 'Enter cross-range resolution (in m):'
    val = 50E3*rd.wavelen(3.5E9)/2.1
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.4a
def quiz_4_1_4a():
    prompt = 'Enter Doppler shift (in Hz):'
    val = rd.dopp_shift(500, 1.5E9)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.4b
def quiz_4_1_4b():
    prompt = 'Enter range rate (in m/s):'
    val = -c*10E3/2/3.0E9
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.5
def quiz_4_1_5():
    prompt = 'Enter range rate (in m/s):'
    val = 55.0
    tol = 1.0
    new_quiz(prompt, val, abs_tol=tol)

# Q4.1.6a
def quiz_4_1_6a():
    prompt = 'Enter range rate resolution (in m/s):'
    val = c/2/5E9/50E-3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.6b
def quiz_4_1_6b():
    prompt = 'Enter range rate (in m/s):'
    val = c*500/4/5E9
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.7a
def quiz_4_1_7a():
    prompt = 'Enter RCS (in dBsm):'
    val = rd.to_db(rd.from_db(0)*(rd.from_db(13)/rd.from_db(15))*(65E3/50E3)**4)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q4.1.7b
def quiz_4_1_7b():
    prompt = 'Enter RCS (in dBsm):'
    val = rd.to_db(rd.from_db(5)*(rd.from_db(13)/rd.from_db(7))*(100E3/500E3)**4)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.1a
def quiz_5_1_1a():
    prompt = 'Enter wavelength (in m):'
    val = rd.wavelen(10E9)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.1b
def quiz_5_1_1b():
    prompt = 'Enter beamwidth (in deg):'
    val = 70*rd.wavelen(10E9)/4
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.1c
def quiz_5_1_1c():
    prompt = 'Enter transmit gain (in dB):'
    val = rd.to_db(rd.dish_gain(2, 10E9))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.2a
def quiz_5_1_2a():
    prompt = 'Enter received energy (in J):'
    val = rd.rx_energy(200E3, 10, 2E9, gain=rd.area2gain(10, 2E9)**2, rcs=rd.from_db(-5))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.2b
def quiz_5_1_2b():
    prompt = 'Enter SNR (in dB):'
    val = rd.to_db(rd.rx_energy(200E3, 10, 2E9, gain=rd.area2gain(10, 2E9)**2, rcs=rd.from_db(-5))) - \
        rd.to_db(k*500)
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.3a
def quiz_5_1_3a():
    prompt = 'Enter duty cycle:'
    val = 100E-6*2E3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.3b
def quiz_5_1_3b():
    prompt = 'Enter number of pulses:'
    val = math.ceil(rd.from_db(12))
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.4a
def quiz_5_1_4a():
    prompt = 'Enter bandwidth (in Hz):'
    val = c/2/1.5
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.4b
def quiz_5_1_4b():
    prompt = 'Enter cross range resolution (in m):'
    beamw = 57.3*rd.wavelen(5E9)/5.2
    val = 90E3*beamw/57.3
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.4c
def quiz_5_1_4c():
    prompt = 'Enter range rate (in m/s):'
    val = c*700/4/3E9
    tol = 0.01
    new_quiz(prompt, val, rel_tol=tol)

# Q5.1.5
def quiz_5_1_5(targets):
    for ii in range(len(targets)):
        prompt = f'Enter route number for target with RCS {targets[ii].rcs:.2f} dBsm:'
        val = targets[ii].route
        tol = 0.01
        new_quiz(prompt, val, abs_tol=tol)


# coding: utf-8

# # Accessing ncSOS with OWSLib

# We have an ncSOS server with a [get observation example that works](http://geoport-dev.whoi.edu/thredds/sos/usgs/data2/notebook/1211-AA.cdf?service=SOS&version=1.0.0&request=GetObservation&responseFormat=text%2Fxml%3Bsubtype%3D%22om%2F1.0.0%22&offering=1211-AA&observedProperty=http://mmisw.org/ont/cf/parameter/eastward_sea_water_velocity&procedure=urn:ioos:station:gov.usgs.cmgp:1211-AA):
# ```
# http://geoport-dev.whoi.edu/thredds/sos/usgs/data2/notebook/1211-AA.cdf?service=SOS&version=1.0.0&request=GetObservation&responseFormat=text%2Fxml%3Bsubtype%3D%22om%2F1.0.0%22&offering=1211-AA&observedProperty=http://mmisw.org/ont/cf/parameter/eastward_sea_water_velocity&procedure=urn:ioos:station:gov.usgs.cmgp:1211-AA
# ```
#
# But can we formulate, request and process this same query (and others like it) using OWSlib?

# In[64]:

get_ipython().magic('matplotlib inline')
from owslib.sos import SensorObservationService
import pdb
from owslib.etree import etree
import pandas as pd
import datetime as dt
import numpy as np

# In[46]:

url = 'http://sdf.ndbc.noaa.gov/sos/server.php?request=GetCapabilities&service=SOS&version=1.0.0'
ndbc = SensorObservationService(url)

# In[47]:

# usgs woods hole
# buoy data (single current meter)
url='http://geoport-dev.whoi.edu/thredds/sos/usgs/data2/notebook/1211-AA.cdf'
usgs = SensorObservationService(url)
contents = usgs.contents

# In[48]:

usgs.contents

# In[49]:

off = usgs.offerings[1]
off.name

# In[50]:

off.response_formats

# In[51]:

off.observed_properties

# In[52]:

off.procedures

# In[53]:

# the get observation request below works.  How can we recreate this using OWSLib?
# http://geoport-dev.whoi.edu/thredds/sos/usgs/data2/notebook/1211-A1H.cdf?service=SOS&version=1.0.0&request=GetObservation&responseFormat=text%2Fxml%3Bsubtype%3D%22om%2F1.0.0%22&offering=1211-A1H&observedProperty=u_1205&procedure=urn:ioos:station:gov.usgs:1211-A1H

# In[54]:

#pdb.set_trace()
response = usgs.get_observation(offerings=['1211-AA'],
                                 responseFormat='text/xml;subtype="om/1.0.0"',
                                 observedProperties=['http://mmisw.org/ont/cf/parameter/eastward_sea_water_velocity'],
                                 procedure='urn:ioos:station:gov.usgs:1211-AA')

# In[55]:

print(response[0:4000])

# In[56]:

# usgs woods hole ADCP data
# url='http://geoport-dev.whoi.edu/thredds/sos/usgs/data2/notebook/9111aqd-a.nc'
# adcp = SensorObservationService(url)

# In[57]:

root = etree.fromstring(response)

# In[58]:

print(root)

# In[59]:

# root.findall(".//{%(om)s}Observation" % root.nsmap )
values = root.find(".//{%(swe)s}values" % root.nsmap )

# In[60]:

date_value = np.array( [ (dt.datetime.strptime(d,"%Y-%m-%dT%H:%M:%SZ"),float(v))
                      for d,v in [l.split(',') for l in values.text.split()]] )

# In[61]:

ts = pd.Series(date_value[:,1],index=date_value[:,0])

# In[72]:

ts.plot(figsize=(12,4), grid='on');

# # Now try setting time range via eventTime.

# In[70]:

start = '1977-01-03T00:00:00Z'
stop = '1977-01-07T00:00:00Z'
response = usgs.get_observation(offerings=['1211-AA'],
                                 responseFormat='text/xml;subtype="om/1.0.0"',
                                 observedProperties=['http://mmisw.org/ont/cf/parameter/eastward_sea_water_velocity'],
                                 procedure='urn:ioos:station:gov.usgs:1211-AA',
                                 eventTime='{}/{}'.format(start,stop))

# In[73]:

root = etree.fromstring(response)
date_value = np.array( [ (dt.datetime.strptime(d,"%Y-%m-%dT%H:%M:%SZ"),float(v))
                      for d,v in [l.split(',') for l in values.text.split()]] )
ts = pd.Series(date_value[:,1],index=date_value[:,0])
ts.plot(figsize=(12,4), grid='on');

# ...hmmm, didn't seem to do anything

# In[ ]:


import numpy as np

# Create an array
arr = np.arange(5)

# Saving array on disk in binary format (file extension .npy)
np.save('my_array', arr)

# Change arr
arr = np.arange(10)

# Lets see the original saved copy
np.load('my_array.npy')

# Saving multiple arrays into a zip file
np.savez('two_arrays.npz', x=arr, y=arr)

# Now loading multiple arrays
archive_array = np.load('two_arrays.npz')

# Show
archive_array['x']

# Now saving and loading text files
arr = np.array([[1, 2, 3], [4, 5, 6]])
np.savetxt('my_test_text.txt', arr, delimiter=',')

# Loading text files
arr = np.loadtxt('my_test_text.txt', delimiter=',')


import numpy as np #the numpy library
import matplotlib.pyplot as plt #Matplotlib's pyplot

import sys #gives access to C-like sys library
import os #gives acess to operating system, if we want to change directory or make a new directory while running the program

print(sys.argv) #command line arguments
print(os.getcwd()) #print current working directory