from __future__ import annotations

import gc
from dataclasses import dataclass
from typing import Callable, Iterable

from ._colors import colors
from ._config import DEFAULT_CONFIG, Config
from ._loopers import (
    EachLooper, MemoryLooper, OpcodeLooper, Timer, TotalLooper,
)
from ._results import MallocResult, OpcodesResult, TimingResult

Func = Callable[[Iterable[int]], None]

@dataclass(frozen=True)
class Check:
    """A single benchmark.

    Don't instancinate directly, use `Group.add` decorator instead.
    """
    name: str
    func: Func
    loops: int | None
    repeats: int
    min_time: float
    timer: Timer

    def print(
        self,
        config: Config = DEFAULT_CONFIG,
        base_time: float | None = None,
    ) -> TimingResult:
        print_args: dict = dict(
            stream=config.stream,
            histogram_lines=config.histogram_lines,
        )
        print(f'  {colors.magenta(self.name)}', file=config.stream)
        tresult = self.check_timing()
        tresult._base_time = base_time
        tresult.print(**print_args)
        if config.allocations or config.opcodes:
            oresult = self.check_opcodes(best=tresult.best)
            oresult.print(**print_args)
        if config.allocations:
            mresult = self.check_mallocs(lines=oresult.lines)
            mresult.print(**print_args)
        return tresult

    def check_timing(self) -> TimingResult:
        """Run benchmarks for the check.
        """
        # to detect caching, we should individually record
        # the very first run
        each_timings = self._run_each_loop(2)
        total_timings = []
        loops = self.loops
        repeats = self.repeats
        if loops is None:
            loops, raw_timing = self._autorange()
            total_timings.append(raw_timing)
            repeats -= 1
        if loops > 2:
            each_timings.extend(self._run_each_loop(loops - 2))

        for _ in range(repeats):
            total_timings.append(self._run_total_loop(loops))
        assert len(total_timings) == self.repeats
        return TimingResult(
            total_timings=[dt / loops for dt in total_timings],
            each_timings=each_timings,
        )

    def check_opcodes(self, loops: int = 1, best: float = 0) -> OpcodesResult:
        """Run the benchmark and count executed opcodes.
        """
        looper = OpcodeLooper(loops=loops)
        self._run(looper)
        return OpcodesResult(
            opcodes=looper.opcodes,
            lines=looper.lines,
            timings=looper.timings,
            best=best,
        )

    def check_mallocs(self, lines: int, loops: int = 1) -> MallocResult:
        """Run the benchmark and trace memory allocations.
        """
        period = max(1, round(lines / 500))
        looper = MemoryLooper(period=period, loops=loops)
        self.func(looper)
        return MallocResult(
            totals=looper.totals,
            allocs=looper.allocs,
        )

    # Private methods

    def _run_total_loop(self, loops: int = 1) -> float:
        looper = TotalLooper(loops=loops, timer=self.timer)
        self._run(looper)
        return looper.stop - looper.start

    def _run_each_loop(self, loops: int) -> list[float]:
        looper = EachLooper(loops=loops, timer=self.timer, timings=[])
        self._run(looper)
        assert len(looper.timings) == loops
        return looper.timings

    def _run(self, looper: Iterable[int]) -> None:
        gc_was_enabled = gc.isenabled()
        gc.disable()
        try:
            self.func(looper)
        finally:
            if gc_was_enabled:
                gc.enable()

    def _autorange(self) -> tuple[int, float]:
        """Return the number of loops so that total time is at least min_time.

        Calls the timeit method with increasing numbers from the sequence
        1, 2, 5, 10, 20, 50, ... until the time taken is at least 0.2
        second.
        """
        i = 1
        while True:
            for j in 1, 2, 5:
                number = i * j
                time_taken = self._run_total_loop(number)
                if time_taken >= self.min_time:
                    return number, time_taken
            i *= 10


import json

WIKIDATA_IDS_PATH = 'wikidata_ids_entrepreneurs'
DATA_PATH = '/Volumes/Juhana_Ext/cs341data/'
LANG = 'en'

company_ids = []
with open(WIKIDATA_IDS_PATH) as fp:
    for line in fp:
        company_ids.append(line.strip())

company_ids = set(company_ids)

def parse_json(line,label_fp, alias_fp):
    obj = json.loads(line)
    id_ = obj['id'][1:]
    label = None

    if not id_ in company_ids:
        return
    if 'labels' in obj and LANG in obj['labels']:
        label = obj['labels'][LANG]['value']
        label = (label.replace('\t', ' ').replace('\n', ' ')).encode('utf-8')
        label_fp.write(label + '\n')

    if label and 'aliases' in obj and LANG in obj['aliases']:
        aliases = [a['value'] for a in obj['aliases'][LANG]]
        for alias in aliases:
            alias_fp.write(label + '\t' + (alias.replace('\t', ' ').replace('\n', ' ') + '\n').encode('utf-8'))

with open(DATA_PATH + '/wikidata_dump.json', 'r') as f, \
    open('names_entrepreneurs.tsv', 'w') as label_fp, \
    open('aliases_entrepreneurs.tsv', 'w') as alias_fp:

    for i, line in enumerate(f):
        line = line.rstrip()
        if line == '[' or line == ']':
            continue
        # strip trailing commas
        line = line.rstrip(',')
        parse_json(line, label_fp, alias_fp)


from __future__ import absolute_import

import os
import re
import logging
import textwrap
import time
from subprocess import CalledProcessError

from ... import constants
from ...memoize import memoized, reset_memoize_cache
from ...config import get_config_value
from ...subprocess import check_and_log_output_and_error_demoted, check_output_demoted, check_call_demoted, call_demoted
from ...log import log_to_client

def _command_for_vm(command_list):
    ssh_command = ['ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null',
                   '-i', _vm_key_path(), 'docker@{}'.format(get_docker_vm_ip())]
    ssh_command.extend([command_list])
    return ssh_command

@memoized
def _vm_key_path():
    dusty_user = get_config_value(constants.CONFIG_MAC_USERNAME_KEY)
    return os.path.expanduser('~{}/.docker/machine/machines/{}/id_rsa'.format(dusty_user,
                                                                              constants.VM_MACHINE_NAME))

def run_command_on_vm(command_list, quiet_on_success=True):
    return check_and_log_output_and_error_demoted(_command_for_vm(command_list),
                                                  quiet_on_success=quiet_on_success)

def check_output_on_vm(command_list, redirect_stderr=False):
    return check_output_demoted(_command_for_vm(command_list), redirect_stderr=redirect_stderr)

def check_call_on_vm(command_list, redirect_stderr=False):
    return check_call_demoted(_command_for_vm(command_list), redirect_stderr=redirect_stderr)

def call_on_vm(command_list, redirect_stderr=False):
    return call_demoted(_command_for_vm(command_list), redirect_stderr=redirect_stderr)

def _ensure_rsync_is_installed():
    # We're running tce-load twice as a hack to get around the fact that, for
    # completely unknown reasons, tce-load will return with an exit code of 1 after
    # initial install even if it works just fine. Subsequent install attempts will
    # be no-ops with a return code of 0.
    run_command_on_vm('which rsync || tce-load -wi rsync || tce-load -wi rsync')

def _ensure_persist_dir_is_linked():
    mkdir_if_cmd = 'if [ ! -d /mnt/sda1{0} ]; then sudo mkdir /mnt/sda1{0}; fi'.format(constants.VM_PERSIST_DIR)
    mount_if_cmd = 'if [ ! -d {0} ]; then sudo ln -s /mnt/sda1{0} {0}; fi'.format(constants.VM_PERSIST_DIR)
    run_command_on_vm(mkdir_if_cmd)
    run_command_on_vm(mount_if_cmd)

def _ensure_vm_dir_exists(vm_dir):
    mkdir_if_cmd = 'if [ ! -d {0} ]; then sudo mkdir {0}; fi'.format(vm_dir)
    run_command_on_vm(mkdir_if_cmd)

def _ensure_cp_dir_exists():
    _ensure_vm_dir_exists(constants.VM_CP_DIR)

def _ensure_assets_dir_exists():
    _ensure_vm_dir_exists(constants.VM_ASSETS_DIR)

def _dusty_vm_exists():
    """We use VBox directly instead of Docker Machine because it
    shaves about 0.5 seconds off the runtime of this check."""
    existing_vms = check_output_demoted(['VBoxManage', 'list', 'vms'])
    for line in existing_vms.splitlines():
        if '"{}"'.format(constants.VM_MACHINE_NAME) in line:
            return True
    return False

def _apply_nat_dns_host_resolver():
    """
    This will make the Dusty VM always use the host's DNS resolver for lookups.
    It solves an issue we were seeing where the VM's resolving settings would get
    out of date when a laptop was moved between routers with different settings,
    resulting in DNS lookup failures on the VM.
    """
    check_and_log_output_and_error_demoted(
        ['VBoxManage', 'modifyvm', constants.VM_MACHINE_NAME, '--natdnshostresolver1', 'on'],
        quiet_on_success=True)

def _apply_nat_net_less_greedy_subnet():
    """By default, VirtualBox claims 10.0.2.x for itself as part of its NAT routing
    scheme. This subnet is commonly used on internal networks, making this a pretty
    damn greedy choice. We instead alter the VM to use the less greedy subnet of
    10.174.249.x which is less likely to conflict."""
    check_and_log_output_and_error_demoted(
        ['VBoxManage', 'modifyvm', constants.VM_MACHINE_NAME, '--natnet1', '10.174.249/24'],
        quiet_on_success=True)

def _init_docker_vm():
    """Initialize the Dusty VM if it does not already exist."""
    if not _dusty_vm_exists():
        log_to_client('Initializing new Dusty VM with Docker Machine')
        machine_options = ['--driver', 'virtualbox',
                           '--virtualbox-cpu-count', '-1',
                           '--virtualbox-boot2docker-url', constants.CONFIG_BOOT2DOCKER_URL,
                           '--virtualbox-memory', str(get_config_value(constants.CONFIG_VM_MEM_SIZE)),
                           '--virtualbox-hostonly-nictype', constants.VM_NIC_TYPE]
        check_call_demoted(['docker-machine', 'create'] + machine_options + [constants.VM_MACHINE_NAME],
                           redirect_stderr=True)

def _start_docker_vm():
    """Start the Dusty VM if it is not already running."""
    is_running = docker_vm_is_running()
    if not is_running:
        log_to_client('Starting docker-machine VM {}'.format(constants.VM_MACHINE_NAME))
        _apply_nat_dns_host_resolver()
        _apply_nat_net_less_greedy_subnet()
        check_and_log_output_and_error_demoted(['docker-machine', 'start', constants.VM_MACHINE_NAME], quiet_on_success=True)
    return is_running

def _stop_docker_vm():
    """Stop the Dusty VM if it is not already stopped."""
    check_call_demoted(['docker-machine', 'stop', constants.VM_MACHINE_NAME], redirect_stderr=True)

@memoized
def _get_vm_config():
    return check_output_demoted(['VBoxManage', 'showvminfo', '--machinereadable', constants.VM_MACHINE_NAME]).splitlines()

def docker_vm_is_running():
    """Using VBoxManage is 0.5 seconds or so faster than Machine."""
    running_vms = check_output_demoted(['VBoxManage', 'list', 'runningvms'])
    for line in running_vms.splitlines():
        if '"{}"'.format(constants.VM_MACHINE_NAME) in line:
            return True
    return False

def _vm_not_using_pcnet_fast_iii():
    for line in _get_vm_config():
        if 'nictype1' in line and constants.VM_NIC_TYPE not in line:
            return True
    return False

def _apply_nic_fix():
    """Set NIC 1 to use PCnet-FAST III. The host-only NIC type is
    set during docker-machine create (and Machine will change it
    back if it is changed manually), which is why we only change
    NIC 1 here."""
    log_to_client('Setting NIC 1 to use PCnet-FAST III')
    check_call_demoted(['VBoxManage', 'modifyvm', constants.VM_MACHINE_NAME, '--nictype1', constants.VM_NIC_TYPE])

def delete_docker_vm_host_only_interface():
    """Attempt to delete the host-only interface attached
    to the current Dusty VM. VM should be stopped
    before calling this."""
    adapter_name = get_vm_hostonly_adapter()
    log_to_client('Deleting host-only interface {}'.format(adapter_name))
    check_call_demoted(['VBoxManage', 'hostonlyif', 'remove', adapter_name])

def regenerate_docker_vm_certificates():
    """Regenerate certificates for a running VM through Docker Machine.
    This may be necessary following a restart if there were previously
    networking issues preventing Machine from doing this as part
    of normal startup."""
    log_to_client('Attempting to regenerate certificates for Dusty VM')
    check_call_demoted(['docker-machine', 'regenerate-certs', '-f', constants.VM_MACHINE_NAME])

def ensure_docker_vm_is_started():
    reset_memoize_cache()
    _init_docker_vm()
    # Switch VM to use PCnet-FAST III which we have observed to
    # have much better performance than the default.
    if _vm_not_using_pcnet_fast_iii():
        log_to_client('Stopping Dusty VM to apply NIC fix for faster networking performance')
        if docker_vm_is_running():
            _stop_docker_vm()
        _apply_nic_fix()
    was_already_running = _start_docker_vm()
    return was_already_running

def initialize_docker_vm():
    was_already_running = ensure_docker_vm_is_started()
    # These operations are somewhat expensive, so we only
    # run them if there's a chance we are bringing up a
    # fresh VM or temporary filesystems may need to be
    # reinitialized following a restart.
    if not was_already_running:
        _ensure_rsync_is_installed()
        _ensure_persist_dir_is_linked()
        _ensure_cp_dir_exists()
        _ensure_assets_dir_exists()

def shut_down_docker_vm():
    if docker_vm_is_running():
        check_call_demoted(['docker-machine', 'stop', constants.VM_MACHINE_NAME])
        log_to_client('Dusty VM was shut down')
    else:
        log_to_client('Dusty VM is already shut down')
    reset_memoize_cache()

def _get_localhost_ssh_port():
    """Something in the VM chain, either VirtualBox or Machine, helpfully
    sets up localhost-to-VM forwarding on port 22. We can inspect this
    rule to determine the port on localhost which gets forwarded to
    22 in the VM."""
    for line in _get_vm_config():
        if line.startswith('Forwarding'):
            spec = line.split('=')[1].strip('"')
            name, protocol, host, host_port, target, target_port = spec.split(',')
            if name == 'ssh' and protocol == 'tcp' and target_port == '22':
                return host_port
    raise ValueError('Could not determine localhost port for SSH forwarding')

def _get_host_only_mac_address():
    """Returns the MAC address assigned to the host-only adapter,
    using output from VBoxManage. Returned MAC address has no colons
    and is lower-cased."""
    # Get the number of the host-only adapter
    vm_config = _get_vm_config()
    for line in vm_config:
        if line.startswith('hostonlyadapter'):
            adapter_number = int(line[15:16])
            break
    else:
        raise ValueError('No host-only adapter is defined for the Dusty VM')

    for line in vm_config:
        if line.startswith('macaddress{}'.format(adapter_number)):
            return line.split('=')[1].strip('"').lower()
    raise ValueError('Could not find MAC address for adapter number {}'.format(adapter_number))

def _ip_for_mac_from_ip_addr_show(ip_addr_show, target_mac):
    """Given the rather-complex output from an 'ip addr show' command
    on the VM, parse the output to determine the IP address
    assigned to the interface with the given MAC."""
    return_next_ip = False
    for line in ip_addr_show.splitlines():
        line = line.strip()
        if line.startswith('link/ether'):
            line_mac = line.split(' ')[1].replace(':', '')
            if line_mac == target_mac:
                return_next_ip = True
        elif return_next_ip and line.startswith('inet') and not line.startswith('inet6'):
            ip = line.split(' ')[1].split('/')[0]
            return ip

def _get_host_only_ip():
    """Determine the host-only IP of the Dusty VM through Virtualbox and SSH
    directly, bypassing Docker Machine. We do this because Docker Machine is
    much slower, taking about 600ms total. We are basically doing the same
    flow Docker Machine does in its own code."""
    mac = _get_host_only_mac_address()
    ip_addr_show = check_output_demoted(['ssh', '-o', 'StrictHostKeyChecking=no',
                                         '-o', 'UserKnownHostsFile=/dev/null',
                                         '-i', _vm_key_path(), '-p', _get_localhost_ssh_port(),
                                         'docker@127.0.0.1', 'ip addr show'])
    return _ip_for_mac_from_ip_addr_show(ip_addr_show, mac)

@memoized
def get_docker_vm_ip():
    try:
        return _get_host_only_ip()
    except Exception as e:
        logging.exception('Failed getting host only IP through optimized path, trying Docker Machine')
        return check_output_demoted(['docker-machine', 'ip', constants.VM_MACHINE_NAME]).rstrip()

@memoized
def get_docker_bridge_ip():
    logging.info('Attempting to get the docker bridge IP')
    for _ in range(15):
        result = check_output_on_vm("ip route | grep docker0 | awk '{print $NF}'").rstrip()
        if result:
            return result
        time.sleep(1)
        logging.info('Retrying...')

    raise ValueError('Could not get Docker bridge IP from Virtualbox, VM may not be fully initialized')

def _parse_df_output(df_line):
    split_line = df_line.split()
    return {'total': split_line[1],
            'used': split_line[2],
            'free': split_line[3],
            'usage_pct': split_line[4]}

def _format_df_dict(df_dict):
    formatted_usage = "Usage: {}\n".format(df_dict['usage_pct'])
    formatted_usage += "Total Size: {}\n".format(df_dict['total'])
    formatted_usage += "Used: {}\n".format(df_dict['used'])
    formatted_usage += "Free: {}".format(df_dict['free'])
    return formatted_usage

def get_docker_vm_disk_info(as_dict=False):
    df_output = check_output_on_vm('df -h /mnt/sda1 | grep /dev/sda1')
    df_line = df_output.split('\n')[0]
    df_dict = _parse_df_output(df_line)
    return df_dict if as_dict else _format_df_dict(df_dict)

def get_vm_hostonly_adapter():
    config_lines = _get_vm_config()
    for line in config_lines:
        if line.startswith('hostonlyadapter'):
            return line.split('=')[1].strip('"').rstrip('"')

def _get_hostonly_config():
    return check_output_demoted(['VBoxManage', 'list', 'hostonlyifs']).splitlines()

@memoized
def get_host_ip():
    adapter = get_vm_hostonly_adapter()
    host_only_config = _get_hostonly_config()

    in_adapter_block = False
    for line in host_only_config:
        if line.startswith('Name'):
            if adapter in line:
                in_adapter_block = True
            else:
                in_adapter_block = False
        if in_adapter_block:
            if line.startswith('IPAddress'):
                return line.split()[1]
    raise RuntimeError('Host IP in host only network not found')

def required_absent_assets(assembled_specs):
    required_and_absent = []
    for asset_key, asset_info in assembled_specs['assets'].iteritems():
        if asset_info['required_by'] and not asset_is_set(asset_key):
            required_and_absent.append(asset_key)
    return required_and_absent

@memoized
def get_assets_on_vm():
    dir_contents = check_output_on_vm('ls {}'.format(constants.VM_ASSETS_DIR))
    return set(dir_contents.splitlines())

def asset_vm_path(asset_key):
    return os.path.join(constants.VM_ASSETS_DIR, asset_key)

def asset_is_set(asset_key):
    return asset_key in get_assets_on_vm()

def asset_value(asset_key):
    return check_output_on_vm('sudo cat {}'.format(asset_vm_path(asset_key)))

def remove_asset(asset_key):
    run_command_on_vm('sudo rm -f {}'.format(asset_vm_path(asset_key)))


import random
import numpy as np
import pickle
import math
import copy
import nets
import torch
import torch.optim as optim
from torch.autograd import Variable
from torch.nn.functional import mse_loss

class PlayerAgent:
    def __init__(self, gamma, critic_alpha, actor_alpha, seed_val, mem_len, epsilon, epsilon_decay, epsilon_min, smax, c_lamb, a_lamb, thresh_emb, thresh_cosh, clipgrad):
        self.gamma = gamma
        self.critic_alpha = critic_alpha
        self.actor_alpha = actor_alpha
        self.seed_val = seed_val
        self.counter = 0
        self.task_counters = {}
        self.s_shape = (-2)
        self.a_shape = (-2)
        self.task_id = -1
        self.is_first_task = True

        self.s_factor = -1.
        self.critic_mask_pre = None
        self.critic_mask_back = None
        self.actor_mask_pre = None
        self.actor_mask_back = None

        self.critic_model = nets.CriticNet().cuda()
        self.actor_model = nets.ActorNet().cuda()
        self.target_critic_model = nets.CriticNet().cuda()
        self.target_actor_model = nets.ActorNet().cuda()
        self.target_critic_model.load_state_dict(self.critic_model.state_dict())
        self.target_actor_model.load_state_dict(self.actor_model.state_dict())
        self.critic_optimizer = optim.Adam(self.critic_model.parameters(), lr=self.critic_alpha)
        self.actor_optimizer = optim.Adam(self.actor_model.parameters(), lr=self.actor_alpha)
        self.memory = {} # set of mems, one per task

        self.mem_len = mem_len
        self.epsilon = {}
        self.epsilon_init = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.smax = smax
        self.c_lamb = c_lamb
        self.a_lamb = a_lamb
        self.thresh_emb = thresh_emb
        self.thresh_cosh = thresh_cosh
        self.clipgrad = clipgrad

    def set_up_task(self, task_id, s_shape, a_shape):
        #print("in set_up_task")
        #print("target_model params")
        #print(self.target_model.parameters())
        #print(self.target_model.state_dict())
        #print("model params")
        #print(self.model.parameters())
        #print(self.model.state_dict())
        #print("checking equality")
        #print(self.target_model.state_dict().keys() == self.model.state_dict().keys())
        #input("wait")
        self.s_shape = s_shape
        self.a_shape = a_shape

        self.critic_model.task_configure(task_id, s_shape, a_shape)
        if self.target_critic_model.state_dict().keys() != self.critic_model.state_dict().keys():
            print("--- DOING critic PARAM COPY ---")
            temp_critic_model = copy.deepcopy(self.target_critic_model)
            self.target_critic_model = copy.deepcopy(self.critic_model)
            for name, param in self.target_critic_model.named_parameters():
                for t_name, t_param in temp_critic_model.named_parameters():
                    if name == t_name:
                        param.data.copy_(t_param.data)
        self.target_critic_model.task_configure(task_id, s_shape, a_shape)
        self.actor_model.task_configure(task_id, s_shape, a_shape)
        if self.target_actor_model.state_dict().keys() != self.actor_model.state_dict().keys():
            print("--- DOING actor PARAM COPY ---")
            temp_actor_model = copy.deepcopy(self.target_actor_model)
            self.target_actor_model = copy.deepcopy(self.actor_model)
            for name, param in self.target_actor_model.named_parameters():
                for t_name, t_param in temp_actor_model.named_parameters():
                    if name == t_name:
                        param.data.copy_(t_param.data)
        self.target_actor_model.task_configure(task_id, s_shape, a_shape)

        self.task_id = task_id
        if task_id not in self.memory:
            self.memory[task_id] = []
            self.task_counters[task_id] = 0
            self.epsilon[task_id] = self.epsilon_init
        self.critic_optimizer = optim.Adam(self.critic_model.parameters(), lr=self.critic_alpha)
        self.actor_optimizer = optim.Adam(self.actor_model.parameters(), lr=self.actor_alpha)

        self.target_critic_model.cuda()
        self.critic_model.cuda()
        self.target_actor_model.cuda()
        self.actor_model.cuda()

    def store_experience(self, s, a, r, s_prime, done, ep, run):
        s_in = s.tolist()
        s_prime_in = s_prime.tolist()

        mem = self.memory[self.task_id]

        if 0 == len(np.asarray(a).shape):
            a = np.reshape(a,(1))

        if len(mem) < self.mem_len[self.task_id]:
            mem.append((s_in, a, r, s_prime_in, done, ep, run))
        else:
            #replace_idx = random.randrange(len(mem)) # rand
            replace_idx = self.counter % self.mem_len[self.task_id] # wrap
            mem[replace_idx] = (s_in, a, r, s_prime_in, done, ep, run)

        self.counter += 1
        self.task_counters[self.task_id] += 1

    def rescale(self, val, in_min, in_max, out_min, out_max):
        return ((val - in_min) / (in_max - in_min)) * (out_max - out_min) + out_min

    def sigmoid(self, val):
        return 1. / (1. + np.exp(-val))

    def get_masks(self):
        c_mask = self.critic_model.mask(self.task_id, self.smax)
        for i in range(len(c_mask)):
            c_mask[i] = Variable(c_mask[i].data.clone(), requires_grad=False)
        if self.is_first_task:
            self.critic_mask_pre = c_mask
        else:
            for i in range(len(self.critic_mask_pre)):
                self.critic_mask_pre[i] = torch.max(self.critic_mask_pre[i], c_mask[i])

        self.critic_mask_back = {}
        for n,_ in self.critic_model.named_parameters():
            vals = self.critic_model.get_view_for(n, self.critic_mask_pre)
            if vals is not None:
                self.critic_mask_back[n] = 1 - vals

        #print("critic_mask_pre")
        #print(self.critic_mask_pre)
        #print("critic_mask_back")
        #print(self.critic_mask_back)
        #input("wait")

        a_mask = self.actor_model.mask(self.task_id, self.smax)
        for i in range(len(a_mask)):
            a_mask[i] = Variable(a_mask[i].data.clone(), requires_grad=False)
        if self.is_first_task:
            self.actor_mask_pre = a_mask
        else:
            for i in range(len(self.actor_mask_pre)):
                self.actor_mask_pre[i] = torch.max(self.actor_mask_pre[i], a_mask[i])

        self.actor_mask_back = {}
        for n,_ in self.actor_model.named_parameters():
            vals = self.actor_model.get_view_for(n, self.actor_mask_pre)
            if vals is not None:
                self.actor_mask_back[n] = 1 - vals

        #print("actor_mask_pre")
        #print(self.actor_mask_pre)
        #print("actor_mask_back")
        #print(self.actor_mask_back)
        #input("wait")

    def update_models(self, batch_size, update_num, num_updates):

        #print("in update_models")
        #print("target_model params")
        #print(self.target_model.parameters())
        #print(self.target_model.state_dict())
        #print("model params")
        #print(self.model.parameters())
        #print(self.model.state_dict())
        #print("checking equality")
        #print(self.target_model.state_dict().keys() == self.model.state_dict().keys())
        #input("wait")

        record_range = min(self.task_counters[self.task_id], self.mem_len[self.task_id])
        batch_indices = np.random.choice(record_range, batch_size)
        batch = np.asarray(self.memory[self.task_id], dtype=object)[batch_indices]

        # TODO try including intrinsic for MountainCar

        # ref: self.memory.append((s_in, a, r, s_prime_in, done, ep, run))
        state = torch.from_numpy(np.asarray(batch[:,0].tolist())).float()
        action = torch.from_numpy(np.asarray(batch[:,1].tolist())).float()
        reward = torch.from_numpy(np.asarray(batch[:,2].tolist())).float()
        state_new = torch.from_numpy(np.asarray(batch[:,3].tolist())).float()
        terminal = torch.from_numpy(np.asarray(batch[:,4].tolist())).float()
        state = Variable(state).cuda()
        action = Variable(action).cuda()
        state_new = Variable(state_new).cuda()
        terminal = Variable(terminal).cuda()
        reward = Variable(reward).cuda()
        self.target_critic_model.eval()
        self.target_actor_model.eval()

        self.s_factor = (self.smax-1/self.smax)*update_num/num_updates+1/self.smax # TODO same for RL? or do we need to modify?

        a_next, _ = self.target_actor_model.forward(state_new, self.s_factor)
        Q_next, _ = self.target_critic_model.forward([state_new, a_next], self.s_factor)
        y = reward + torch.mul(torch.squeeze(Q_next)*(1-terminal), self.gamma)

        self.critic_optimizer.zero_grad()
        self.critic_model.train()
        Q, c_masks = self.critic_model.forward([state, action], self.s_factor)
        c_loss, _ = self.c_criterion(torch.squeeze(Q), y, c_masks)
        c_loss.backward()

        # HAT - restrict layer grads
        if not self.is_first_task:
            for n,p in self.critic_model.named_parameters():
                if n in self.critic_mask_back:
                    p.grad.data *= self.critic_mask_back[n]
        # HAT - compensate embedding grads
        for n,p in self.critic_model.named_parameters():
            if n.startswith('e'):
                num = torch.cosh(torch.clamp(self.s_factor*p.data, -self.thresh_cosh, self.thresh_cosh)) + 1
                den = torch.cosh(p.data) + 1
                p.grad.data *= self.smax/self.s_factor*num/den
        torch.nn.utils.clip_grad_norm_(self.critic_model.parameters(), self.clipgrad)
        self.critic_optimizer.step()

        self.actor_optimizer.zero_grad()
        self.critic_model.eval()
        self.actor_model.train()
        a_temp, a_masks = self.actor_model.forward(state, self.s_factor)
        Q_temp, _ = self.critic_model.forward([state, a_temp], self.s_factor)
        a_loss, _ = self.a_criterion(Q_temp, a_masks)
        a_loss.backward()

        # HAT - restrict layer grads
        if not self.is_first_task:
            for n,p in self.actor_model.named_parameters():
                if n in self.actor_mask_back:
                    p.grad.data *= self.actor_mask_back[n]
        # HAT - compensate embedding grads
        for n,p in self.actor_model.named_parameters():
            if n.startswith('e'):
                num = torch.cosh(torch.clamp(self.s_factor*p.data, -self.thresh_cosh, self.thresh_cosh)) + 1
                den = torch.cosh(p.data) + 1
                p.grad.data *= self.smax/self.s_factor*num/den
        torch.nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.clipgrad)
        self.actor_optimizer.step()

        # HAT - constrain embeddings
        for n,p in self.critic_model.named_parameters():
            if n.startswith('e'):
                p.data = torch.clamp(p.data, -self.thresh_emb, self.thresh_emb)
        for n,p in self.actor_model.named_parameters():
            if n.startswith('e'):
                p.data = torch.clamp(p.data, -self.thresh_emb, self.thresh_emb)

        temp_epsilon = self.epsilon[self.task_id]
        temp_epsilon *= self.epsilon_decay
        self.epsilon[self.task_id] = max(self.epsilon_min, temp_epsilon)

    def c_criterion(self, outputs, targets, c_masks):
        c_reg = 0
        c_count = 0
        if self.critic_mask_pre is not None:
            for m,mp in zip(c_masks, self.critic_mask_pre):
                aux = 1-mp
                c_reg += (m*aux).sum()
                c_count += aux.sum()
        else:
            for m in c_masks:
                c_reg += m.sum()
                c_count += np.prod(m.size()).item()
        c_reg /= c_count
        #print("outputs")
        #print(outputs)
        #print("targets")
        #print(targets)
        #print("c_reg")
        #print(c_reg)
        #print("c_count")
        #print(c_count)
        #print("mse_loss:")
        #print(mse_loss(input=outputs, target=targets))
        #input("wait")
        return mse_loss(input=outputs, target=targets) + self.c_lamb * c_reg, c_reg

    def a_criterion(self, outputs, a_masks):
        a_reg = 0
        a_count = 0
        if self.actor_mask_pre is not None:
            for m,mp in zip(a_masks, self.actor_mask_pre):
                aux = 1-mp
                a_reg += (m*aux).sum()
                a_count += aux.sum()
        else:
            for m in a_masks:
                a_reg += m.sum()
                a_count += np.prod(m.size()).item()
        a_reg /= a_count
        #print("outputs")
        #print(outputs)
        #print("targets")
        #print(targets)
        #print("a_reg")
        #print(a_reg)
        #print("a_count")
        #print(a_count)
        #print("mse_loss:")
        #print(mse_loss(input=outputs, target=targets))
        #input("wait")
        return (-outputs).mean() + self.a_lamb * a_reg, a_reg

    def update_target_models(self, tau):
        for target_param, param in zip(self.target_critic_model.parameters(), self.critic_model.parameters()):
            target_param.data.copy_(target_param.data * (1. - tau) + param.data * tau)
        for target_param, param in zip(self.target_actor_model.parameters(), self.actor_model.parameters()):
            target_param.data.copy_(target_param.data * (1. - tau) + param.data * tau)

    def act(self, s_sample, is_test):
        a = None

        # with epsilon prob to choose random action else choose argmax Q estimate action
        if (np.random.rand() <= self.epsilon[self.task_id]) and not is_test:
            #a = np.random.rand(self.a_shape[0])
            a = np.random.normal(0, 1, self.a_shape)
        else:
            state = torch.from_numpy(np.expand_dims(s_sample, axis=0)).float()
            state = Variable(state).cuda()

            self.actor_model.eval()
            if is_test:
                estimate, _ = self.actor_model.forward(state, self.smax)
            else:
                estimate, _ = self.actor_model.forward(state, self.s_factor)
            a = estimate[0].data.cpu().numpy()

        #print("check a:")
        #print(a)
        #input("wait")

        return a

    def load(self):
        print("loading model!")
        if self.critic_model:
            self.critic_model.load("./save/critic_model" + str(self.seed_val) + ".h5", self.critic_optimizer)
        if self.target_critic_model:
            self.target_critic_model.load("./save/target_critic_model" + str(self.seed_val) + ".h5", self.critic_optimizer)
        if self.actor_model:
            self.actor_model.load("./save/actor_model" + str(self.seed_val) + ".h5", self.actor_optimizer)
        if self.target_actor_model:
            self.target_actor_model.load("./save/target_actor_model" + str(self.seed_val) + ".h5", self.actor_optimizer)
        #with open("./save/memory.pickle", "rb") as h:
        #    self.memory = pickle.load(h)

    def save(self):
        print("saving model!")
        if self.critic_model:
            self.critic_model.save("./save/critic_model" + str(self.seed_val) + ".h5", self.counter, self.critic_optimizer)
        if self.target_critic_model:
            self.target_critic_model.save("./save/target_critic_model" + str(self.seed_val) + ".h5", self.counter, self.critic_optimizer)
        if self.actor_model:
            self.actor_model.save("./save/actor_model" + str(self.seed_val) + ".h5", self.counter, self.actor_optimizer)
        if self.target_actor_model:
            self.target_actor_model.save("./save/target_actor_model" + str(self.seed_val) + ".h5", self.counter, self.actor_optimizer)
        #with open("./save/memory.pickle", "wb") as h:
        #    pickle.dump(self.memory, h)


import pytest
from src.StudentAnalyzer import StudentAnalyzer

class TestStudentAnalyzer:
    @pytest.fixture()
    def test_data(self):
        return {
            "Иванов": [("математика", 91), ("химия", 100)],
            "Петров": [("русский язык", 87), ("литература", 78)],
            "Сидоров": [("география", 60), ("физика", 75)],
        }

    def test_students_in_first_quartile(self, test_data):
        analyzer = StudentAnalyzer()
        expected_students = {
            "Сидоров": [("география", 60), ("физика", 75)],
        }

        students_in_quartile = analyzer.students_in_first_quartile(test_data)
        assert students_in_quartile == expected_students


# encoding:utf-8

import os
import sys
from bs4 import BeautifulSoup

reload(sys)
sys.setdefaultencoding("utf-8")

path = os.getcwd() + '/data/'  #原始数据
path1 = os.getcwd() + '/title_abs/'  #处理后的标题和摘要

if not os.path.exists(path1):
    os.mkdir(path1)

filelist = os.listdir(path)  #取得当前路径下的所有文件

def get_title_abs():
    abstracts = []

    for files in filelist:
        filename = os.path.splitext(files)[0]  #取文件名
        soup = BeautifulSoup(open(path + filename + '.xml'), 'html.parser')  #解析网页
        abstract = soup.find("p", class_="abstracts")  #取得"p",class="abstracts"的内容

        if abstract is None or abstract.string is None:
            continue
        else:
            abstracts.extend(soup.title.stripped_strings)
            s = abstract.string
            abstracts.extend(s.encode('utf-8'))

            with open(path1 + filename + ".txt", "w+") as f:
                for i in abstracts:
                    f.write(i)

            abstracts = []

if __name__ == '__main__':
    get_title_abs()


import json
import sys
import time
import uuid
import websocket
import _thread
from event import Event
from common import NostrNode, relays
from eventdb import EventDb

class Client(NostrNode):
    def __init__(self):
        self.relays = {}
        self.subscriptions = {}
        self.running = True
        self.eventdb = EventDb()
        _thread.start_new_thread(self.reader, ())

    def close(self):
        self.running = False
        for i in self.relays:
            self.relays[i].close()

    def add_relay(self, url):
        self.relays[url] = websocket.create_connection(url)
        return self

    def send(self, payload):
        payload_str = json.dumps(payload)
        for relay in self.relays:
            print("sending payload: {}".format(payload_str))
            self.relays[relay].send(payload_str)
        return self

    def send_content_from_user(self, pubkey, content):
        event = Event(pubkey, Event.KIND_TEXT_NOTE, content).dict()
        self.send(["EVENT", event])
        return self

    def process_message(self, message_str):
        message = json.loads(message_str)
        message_processing_map = {
            "OK": self.process_message_default,
            "EOSE": self.process_message_default,
            "NOTICE": self.process_message_default,
            "EVENT": self.process_message_event
        }
        if message[0] not in message_processing_map:
            self.error("unknown message type: {}".format(message[0]))
        else:
            proc_func = message_processing_map[message[0]]
            return proc_func(message)
        return self

    def process_message_event(self, message):
        event = message[2]
        if event["kind"] in [0, 2, 3, 4, 7, 20000, 70202] or len(event["content"]) > 300:
            pass
        else:
            print("EVENT: {}, {}: {}".format(event["pubkey"], event["kind"], event["content"].replace("\n", " ")))
        self.eventdb.process(event)
        return self

    def process_message_default(self, message):
        print("message: {}".format(message))
        return self

    def check_relays(self, subscription_id):
        for relay in self.relays:
            counter = 0
            while counter < 110:  # TODO
                result = self.relays[relay].recv()
                self.process_message(result)
                counter += 1
        return self

    def subscribe(self):
        filter_dict = {
            # "authors": [FAKE_MAPPING["vasil9v"]],
            "limit": 100 # TODO
        }
        subscription_id = uuid.uuid4().hex
        print("subscription_id: " + subscription_id)
        self.send(["REQ", subscription_id, filter_dict])
        self.subscriptions[subscription_id] = subscription_id
        return self

    def unsubscribe(self, subscription_id):
        for relay in self.relays:
            self.send(["CLOSE", subscription_id])
            del self.subscriptions[subscription_id]
        return self

    def unsubscribe_all(self):
        for i in list(self.subscriptions.keys()):
            self.unsubscribe(i)

    def reader(self):
        while self.running:
            for i in list(self.subscriptions.keys()):
                self.check_relays(i)
            time.sleep(1.0)
        self.unsubscribe_all()
        print("exiting reader thread")

    def cli(self, config):
        while self.running:
            command = input("> ")
            if command == "exit" or command == "x":
                self.close()
            if command == "subscribe" or command == "s":
                self.subscribe()
            if command == "unsubscribe" or command == "u":
                self.unsubscribe_all()
            if command == "stats":
                print(self.eventdb.json())
                print(self.eventdb.relays)
                print(self.eventdb.kinds)
            if command == "send":
                config["pubkey"] = config.get("pubkey", input("pubkey> "))
                config["privkey"] = config.get("privkey", input("privkey> "))
                Event.PRIV_KEY = config["privkey"]
                message = input("message> ")
                if len(message) > 0:
                    self.send_content_from_user(config["pubkey"], message)

if __name__ == "__main__":
    config = {}
    if len(sys.argv) > 1:
        config = json.loads(sys.argv[1])
    client = Client()
    client.add_relay(relays[0]).cli(config)


#!/usr/bin/python
# vim: noet sw=4 ts=4 filetype=python

# from distutils.core import setup
from	setuptools	import	setup

version = '1.0.6'

with open( 'untar/version.py', 'wt' ) as f:
    print >>f, "Version = '{0}'".format( version )

setup(
    name                 = 'untar',
    version              = version,
    description          = "Unpack various families of tar(1) archives.",
    long_description     = open('README.md').read(),
    keywords             = 'tar untar unosw unsos unarchive',
    author               = 'Tommy Reynolds',
    author_email         = 'Oldest.Software.Guy@gmail.com',
    url                  = 'http://www.megacoder.com',
    license              = 'MIT',
    include_package_data = True,
    zip_safe             = False,
    install_requires     = [],
    packages             = [ 'untar' ],
	scripts = [
		'scripts/untar',
	],
)


import pika
from Helpers.GoogleVisionClient import GoogleVisionClient
from Helpers.RabbiqMQPublisher import RabbitMQPublisher

class RabbitMQConsumer:

    def callback(self, ch, method, properties, body):
        body_string = body.decode()
        print("Starting filtering " + body_string)
        client = GoogleVisionClient(body_string)
        client.fetch_tags()
        if client.validate("Dog"):
            publisher = RabbitMQPublisher(rabbitmq_host=self.rabbitmq_host, queue_name="valid_images")
            publisher.Publish(message=body)
            print("Accepted")
        else:
            print("Rejected")
        ch.basic_ack(delivery_tag=method.delivery_tag)

    def __init__(self, rabbitmq_host: str, queue_name: str):
        self.rabbitmq_host = rabbitmq_host
        self.queue_name = queue_name

        self.__connection = pika.BlockingConnection(pika.ConnectionParameters(self.rabbitmq_host))
        self.__channel = self.__connection.channel()

        self.__channel.queue_declare(self.queue_name)
        self.__prefetch_count = 1

    def Consume(self):
        self.__channel.basic_consume(queue=self.queue_name,
                                     on_message_callback=self.callback)
        self.__channel.basic_qos(prefetch_count=self.__prefetch_count)

        self.__channel.start_consuming()


##Loads libraries

import pandas as pd
import numpy as np
import yfinance as yf
import pandas_datareader as pdr
import datetime as dt
import matplotlib.pyplot as plt

# Sets a ticker and starting date to grab the data

ticker = 'TSM'

## Sets starting date

startyear = 2020
startmonth = 1
startday = 1
start = dt.datetime(startyear, startmonth, startday)

## Sets today

now = dt.datetime.now()

## Grabs the data

df = pdr.get_data_yahoo(ticker, start, now)

#Shows the data
df.tail()

## I found this script to calculate 14 day RSI

## initializes new data columns

df['Up Move'] = np.nan
df['Down Move'] = np.nan
df['Average Up'] = np.nan
df['Average Down'] = np.nan
df['RS'] = np.nan
df['RSI'] = np.nan

## Calculates Up Move & Down Move
for x in range(1, len(df)):
    df['Up Move'][x] = 0
    df['Down Move'][x] = 0

    if df['Adj Close'][x] > df['Adj Close'][x-1]:
        df['Up Move'][x] = df['Adj Close'][x] - df['Adj Close'][x-1]

    if df['Adj Close'][x] < df['Adj Close'][x-1]:
        df['Down Move'][x] = abs(df['Adj Close'][x] - df['Adj Close'][x-1])

## Calculates initial Average Up & Down, RS and RSI
df['Average Up'][14] = df['Up Move'][1:15].mean()
df['Average Down'][14] = df['Down Move'][1:15].mean()
df['RS'][14] = df['Average Up'][14] / df['Average Down'][14]
df['RSI'][14] = 100 - (100/(1+df['RS'][14]))

## Calculates rest of Average Up, Average Down, RS, RSI
for x in range(15, len(df)):
    df['Average Up'][x] = (df['Average Up'][x-1]*13+df['Up Move'][x])/14
    df['Average Down'][x] = (df['Average Down'][x-1]*13+df['Down Move'][x])/14
    df['RS'][x] = df['Average Up'][x] / df['Average Down'][x]
    df['RSI'][x] = 100 - (100/(1+df['RS'][x]))

## Just looking at it to see the new six columns with the data
print(df)

## Plots the Adjusted Close value and the calculated 14 day RSI for each day
plt.style.use('ggplot')
fig, axs = plt.subplots(2, sharex=True, figsize=(13,9))
fig.suptitle('Stock Price (top) - 14 day RSI (bottom)')
axs[0].plot(df['Adj Close'])
axs[1].plot(df['RSI'])
axs[0].grid()
axs[1].grid()

## I found this script to calculate the buy & sell signals using RSI

## Initializes new columns
df['Long Tomorrow'] = np.nan
df['Buy Signal'] = np.nan
df['Sell Signal'] = np.nan
df['Buy RSI'] = np.nan
df['Sell RSI'] = np.nan
df['Strategy'] = np.nan

## Calculates the buy & sell signals (Buy if RSI is below 40 and sell if it is above 70)
for x in range(15, len(df)):

    # Calculates "Long Tomorrow" column
    if ((df['RSI'][x] <= 40) & (df['RSI'][x-1]>40) ):
        df['Long Tomorrow'][x] = True
    elif ((df['Long Tomorrow'][x-1] == True) & (df['RSI'][x] <= 70)):
        df['Long Tomorrow'][x] = True
    else:
        df['Long Tomorrow'][x] = False

    # Calculates "Buy Signal" column
    if ((df['Long Tomorrow'][x] == True) & (df['Long Tomorrow'][x-1] == False)):
        df['Buy Signal'][x] = df['Adj Close'][x]
        df['Buy RSI'][x] = df['RSI'][x]

    # Calculates "Sell Signal" column
    if ((df['Long Tomorrow'][x] == False) & (df['Long Tomorrow'][x-1] == True)):
        df['Sell Signal'][x] = df['Adj Close'][x]
        df['Sell RSI'][x] = df['RSI'][x]

## Plots the historic buy/sell signals

plt.style.use('ggplot')
fig, axs = plt.subplots(2, sharex=True, figsize=(13,9))
fig.suptitle('Stock Price (top) & 14 day RSI (bottom)')

## Plots the stock close price & buy/sell signals:
axs[0].scatter(df.index, df['Buy Signal'],  color = 'green',  marker = '^', alpha = 1)
axs[0].scatter(df.index, df['Sell Signal'],  color = 'red',  marker = 'v', alpha = 1)
axs[0].plot(df['Adj Close'], alpha = 0.8)
axs[0].grid()

## Plots RSI & buy/sell signals:
axs[1].scatter(df.index, df['Buy RSI'],  color = 'green', marker = '^', alpha = 1)
axs[1].scatter(df.index, df['Sell RSI'],  color = 'red', marker = 'v', alpha = 1)
axs[1].plot(df['RSI'], alpha = 0.8)
axs[1].grid()

## Takes the buy and sell signals and stores the Adj Close values to calculate the gain

df['Strategy'][15] = df['Adj Close'][15]
for x in range(16, len(df)):
    if df['Long Tomorrow'][x-1] == True:
        df['Strategy'][x] = df['Strategy'][x-1]* (df['Adj Close'][x] / df['Adj Close'][x-1])
    else:
        df['Strategy'][x] = df['Strategy'][x-1]

## Calculates the number of trades
trade_count = df['Buy Signal'].count()

## Calculates the average Profit per trade
average_profit = ((df['Strategy'][-1] / df['Strategy'][15])**(1/trade_count))-1

## Calculates the number of days per/trade
total_days = df['Long Tomorrow'].count()
average_days = int(total_days / trade_count)

## Prints the data
print('Number of trades:', trade_count)
print('Average days per trade', average_days)
print('Average profit per trade', average_profit*100, '%')

## List of ETF

stonkList = ['TSM','SPY','DIA','QQQ','VTI','VEU','VNQ']

# Defines the start date
startyear = 2020
startmonth = 1
startday = 1
start = dt.datetime(startyear, startmonth, startday)

# Defines end date
endyear = 2021
endmonth = 5
endday = 30
end = dt.datetime(endyear,endmonth,endday)

## Calculates the percentage change from end with respect to start

for ticker in stonkList:
    # Gets the data from Yahoo finances
    df = pdr.get_data_yahoo(ticker, start, end)

    #Calculates the performance of the ticker
    iniPrc = df['Adj Close'][0]
    endPrc = df['Adj Close'][-1]
    indPer = ((df['Adj Close'][-1] - df['Adj Close'][0] ) / df['Adj Close'][0] )

    #Prints the performance
    print('From', start,' to ', end,',', ticker, 'changed', indPer*100, '%' )


# -*- coding: utf-8 -*-
"""
Created on Mon Jun 15 15:16:51 2020

@author: David Corbo
"""

import torch
import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from torch import nn
from torchvision import datasets, transforms

transform = transforms.Compose([transforms.Resize((28,28)),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5,), (0.5,))
                               ])
training_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
validation_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

training_loader = torch.utils.data.DataLoader(training_dataset, batch_size=100, shuffle=True)
validation_loader = torch.utils.data.DataLoader(validation_dataset, batch_size = 100, shuffle=False)

def im_convert(tensor):
  image = tensor.clone().detach().numpy()
  image = image.transpose(1, 2, 0)
  image = image * np.array((0.5, 0.5, 0.5)) + np.array((0.5, 0.5, 0.5))
  image = image.clip(0, 1)
  return image

dataiter = iter(training_loader)
images, labels = dataiter.next()
fig = plt.figure(figsize=(25, 4))

for idx in np.arange(20):
  ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])
  plt.imshow(im_convert(images[idx]))
  ax.set_title([labels[idx].item()])

class Classifier(nn.Module):

  def __init__(self, D_in, H1, H2, D_out):
      super().__init__()
      self.linear1 = nn.Linear(D_in, H1)
      self.linear2 = nn.Linear(H1, H2)
      self.linear3 = nn.Linear(H2, D_out)
  def forward(self, x):
      x = F.relu(self.linear1(x))
      x = F.relu(self.linear2(x))
      x = self.linear3(x)
      return x

model = Classifier(784, 125, 65, 10)
model

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.0001)

epochs = 5
running_loss_history = []
running_corrects_history = []
val_running_loss_history = []
val_running_corrects_history = []

for e in range(epochs):

  running_loss = 0.0
  running_corrects = 0.0
  val_running_loss = 0.0
  val_running_corrects = 0.0

  for inputs, labels in training_loader:
    inputs = inputs.view(inputs.shape[0], -1)
    outputs = model(inputs)
    loss = criterion(outputs, labels)

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

    _, preds = torch.max(outputs, 1)
    running_loss += loss.item()
    running_corrects += torch.sum(preds == labels.data)

  else:
    with torch.no_grad():
      for val_inputs, val_labels in validation_loader:
        val_inputs = val_inputs.view(val_inputs.shape[0], -1)
        val_outputs = model(val_inputs)
        val_loss = criterion(val_outputs, val_labels)

        _, val_preds = torch.max(val_outputs, 1)
        val_running_loss += val_loss.item()
        val_running_corrects += torch.sum(val_preds == val_labels.data)

    epoch_loss = running_loss/len(training_loader)
    epoch_acc = running_corrects.float()/ len(training_loader)
    running_loss_history.append(epoch_loss)
    running_corrects_history.append(epoch_acc)

    val_epoch_loss = val_running_loss/len(validation_loader)
    val_epoch_acc = val_running_corrects.float()/ len(validation_loader)
    val_running_loss_history.append(val_epoch_loss)
    val_running_corrects_history.append(val_epoch_acc)
    print('epoch :', (e+1))
    print('training loss: {:.4f}, acc {:.4f} '.format(epoch_loss, epoch_acc.item()))
    print('validation loss: {:.4f}, validation acc {:.4f} '.format(val_epoch_loss, val_epoch_acc.item()))