#This file runs all the experiments

"""
Amit Kohli
10-18-20
This file runs all experiments and unit tests. It utilizes the code in src/ and experiment/
"""

import matplotlib.pyplot as plt
import numpy as np
import os
import pathlib
from . import util, lri_forward, seidel
from skimage.restoration import unwrap_phase as unwrap
import torch.fft as fft
from tqdm import tqdm

import pdb
import torch
dirname = str(pathlib.Path(__file__).parent.absolute())

# TODO: Make an experiments/ directory and have different files for each experiment.
# TODO: BUG: wiener doesn't work unless 1) run forward with inverse flips, 2 delete inverse, 3) run inverse without flips

def video_recon(measurement_stack, psf_stack_roft, opt_params, device):
    num_frames = measurement_stack.shape[0]
    psfs = None
    for i in range(num_frames):
        print('frame: ' + str(i))
        curr_frame = measurement_stack[i, :, :]
        estimate, _, _, psfs = image_recon(curr_frame, psf_stack_roft, opt_params, device)
        # for memory recons, replace measurement stack
        measurement_stack[i, :, :] = estimate
    return measurement_stack

def image_recon(measurement, psf_stack_roft, opt_params, diff, device):
    dim = measurement.shape
    if opt_params['init'] == 'measurement':
        estimate = measurement.clone()
    elif opt_params['init'] == 'zero':
        estimate = torch.zeros(dim, device=device)
    elif opt_params['init'] == 'noise':
        estimate = torch.randn(dim, device=device)
    else:
        raise NotImplemented

    estimate.requires_grad = True

    if opt_params['optimizer'] == 'adam':
        optimizer = torch.optim.Adam([estimate], lr=opt_params['lr'])
    elif opt_params['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD([estimate], lr=opt_params['lr'])
    else:
        raise NotImplemented

    loss_fn = torch.nn.MSELoss()
    crop = opt_params['crop']

    losses = []
    for it in tqdm(range(opt_params['iters'])):

        # forward pass
        measurement_guess = lri_forward.forward(estimate, psf_stack_roft, method='normal', device=device, verbose=False, diff=diff)

        # loss
        if crop > 0:
            loss = loss_fn((measurement_guess)[crop:-crop, crop:-crop], (measurement)[crop:-crop, crop:-crop]) + \
                    tv(estimate[crop:-crop, crop:-crop], opt_params['reg'])
        else:
            loss = loss_fn(measurement_guess, measurement) + tv(estimate, opt_params['reg'])

        losses += [loss.detach().cpu()]

        # print loss
        #print(it, loss.item())

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

        # project onto [0,1]
        estimate.data[estimate.data < 0] = 0
        estimate.data[estimate.data > 1] = 1

    return util.normalize(estimate.detach().cpu().float().numpy())

def estimate_coeffs(psf_data, sys_params, opt_params, std_init='random', init=None, plot=False, device=torch.device('cpu')):

    cache_path = dirname + '/data/.cache/'
    half_length = sys_params['L'] / 2
    psf_list = [(torch.tensor(i[0], device=device).float(), torch.tensor(i[1], device=device).float()) for i in psf_data[0]]
    psfs_gt = torch.stack([torch.tensor(psf, device=device).float() for psf in psf_data[1]], dim=0)
    if init is None:
        if std_init == 'zeros':
            coeffs = torch.zeros((3, 1), device=device)
        elif std_init == 'random':
            coeffs = torch.rand((3, 1), device=device)
        else:
            raise NotImplemented
    else:
        coeffs = torch.tensor(init, device=device).float()

    coeffs.requires_grad = True

    optimizer = torch.optim.Adam([coeffs], lr=opt_params['lr'])
    #optimizer = torch.optim.SGD([coeffs], lr=opt_params['lr'])
    l1_loss_fn = torch.nn.L1Loss()
    l2_loss_fn = torch.nn.MSELoss()
    smooth_l1 = torch.nn.SmoothL1Loss()
    fig = plt.figure()
    camera = Camera(fig)
    crop = 225

    for iter in range(opt_params['iters']):
        # forward pass
        psfs_estimate, pupils = seidel.compute_psfs(torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(2, 1, device=device) )), desired_list=psf_list, sys_params=sys_params, device=coeffs.device)\
        # loss
        loss = l2_loss_fn(torch.stack(psfs_estimate, dim=0).float(), psfs_gt) + opt_params['reg']*l2_loss_fn(coeffs, -coeffs)
        # loss = loss_fn((measurement[:,:350]), (measurement_guess[:,:350])) #+ tv(estimate, 1e-9)
        print(iter, loss.item())

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

    print(coeffs)
    if plot:
        for gt, psf, pupil in zip(psfs_gt, psfs_estimate, pupils):
            #pdb.set_trace()
            util.show(torch.cat((gt/gt.max(), psf/psf.max()), dim=1).detach())
            util.show(psf.detach())
            plt.figure()
            plt.imshow(unwrap(torch.angle(pupil).detach()))
            plt.colorbar()
            plt.show()

    #torch.save(psfs_estimate[0], cache_path+'psf_fit.pt')
    return torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(2, 1, device=device)), ).detach()

# single shot version
def estimate_coeffs_ss(psf_img, psf_list, sys_params, opt_params, device, std_init='zeros', init=None, plot=False):

    #psf_list = [(torch.tensor(i[0], device=device).float(), torch.tensor(i[1], device=device).float()) for i in psf_list]
    psfs_gt = torch.tensor(psf_img, device=device).float()
    if init is None:
        if std_init == 'zeros':
            coeffs = torch.zeros((3, 1), device=device)
        elif std_init == 'random':
            coeffs = torch.rand((3, 1), device=device)
        else:
            raise NotImplemented
    else:
        coeffs = torch.tensor(init, device=device).float()

    coeffs.requires_grad = True

    optimizer = torch.optim.Adam([coeffs], lr=opt_params['lr'])
    l2_loss_fn = torch.nn.MSELoss()

    for iter in tqdm(range(opt_params['iters'])):
        # forward pass
        psfs_estimate = seidel.compute_psfs(torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(2, 1, device=device) )), desired_list=psf_list, stack=False, sys_params=sys_params, device=coeffs.device)\
        # loss
        loss = l2_loss_fn(sum(psfs_estimate).float(), psfs_gt) + opt_params['reg']*l2_loss_fn(coeffs, -coeffs)
        # loss = loss_fn((measurement[:,:350]), (measurement_guess[:,:350])) #+ tv(estimate, 1e-9)

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

    if plot:
        util.show(torch.cat((psfs_gt/psfs_gt.max(), sum(psfs_estimate)/sum(psfs_estimate).max()), dim=1).detach().cpu())
        util.show(sum(psfs_estimate).detach().cpu())

    #torch.save(psfs_estimate[0], cache_path+'psf_fit.pt')
    return torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(2, 1, device=device)), ).detach()

def blind_image_recon(measurement, method, model, opt_params, cache, init=None):

    device = model.dataset.device

    estimate = torch.tensor(measurement, device=device).float()
    #estimate = torch.zeros(model.dataset.dim, device=device)
    measurement = torch.tensor(measurement, device=device).float()

    coeffs.requires_grad = True
    estimate.requires_grad = True

    sys_params = model.dataset.sys_params

    opt_est = torch.optim.Adam([estimate], lr=opt_params['lr'])
    opt_coeff = torch.optim.Adam([coeffs], lr=opt_params['lr'])
    #optimizer = torch.optim.SGD([coeffs], lr=opt_params['lr'])
    l1_loss_fn = torch.nn.L1Loss()
    l2_loss_fn = torch.nn.MSELoss()
    smooth_l1 = torch.nn.SmoothL1Loss()
    fig = plt.figure()
    camera = Camera(fig)
    crop = 225
    COUNT = 10
    est = True

    for iter in range(opt_params['iters']):
        # forward pass
        #psfs_estimate, pupils = seidel.compute_psfs(torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(4, 1, device=device) )), desired_list=[(0,0)], sys_params=sys_params, device=coeffs.device)
        measurement_guess,_ = model.forward(estimate, method=method, coeffs=coeffs, cache=cache, verbose=True)

        # loss
        sharpness = -torch.norm(torch.abs(kornia.spatial_gradient(estimate[None, None, :, :], mode='diff')))
        consistency = l2_loss_fn(measurement_guess.float(), measurement) + tv(estimate, opt_params['reg'])

        loss =  consistency + tv(estimate, opt_params['reg']) + 1e-10*sharpness

        print(str(iter) + ' consistency: ' + str(consistency.detach().cpu().numpy()) + ', sharpness: ' + str(-sharpness.detach().cpu().numpy()))

        # backward
        opt_est.zero_grad()
        opt_coeff.zero_grad()
        loss.backward()

        if est:
           opt_est.step()
        else:
           opt_coeff.step()

        estimate.data[estimate.data < 0] = 0
        estimate.data[estimate.data > 1] = 1
        if iter % COUNT:
            est = not est

    print(coeffs)
    util.show(estimate.detach().cpu())

    #torch.save(psfs_estimate[0], cache_path+'psf_fit.pt')
    return estimate.detach().cpu().numpy(), coeffs.detach().cpu().numpy()

def blind_recon(measurement, method, model, opt_params, cache, init=None):

    device = model.dataset.device

    if init is None:
        coeffs = torch.zeros((1, 1), device=device)

    measurement = torch.tensor(measurement, device=device).float()

    coeffs.requires_grad = True
    sys_params = model.dataset.sys_params

    optimizer = torch.optim.Adam([coeffs], lr=opt_params['lr'])
    #optimizer = torch.optim.SGD([coeffs], lr=opt_params['lr'])
    l1_loss_fn = torch.nn.L1Loss()
    l2_loss_fn = torch.nn.MSELoss()
    smooth_l1 = torch.nn.SmoothL1Loss()
    fig = plt.figure()
    camera = Camera(fig)
    crop = 225

    for iter in range(opt_params['iters']):
        # forward pass
        psfs_estimate, _ = seidel.compute_psfs(torch.cat((torch.zeros(1, 1, device=device), coeffs, torch.zeros(4, 1, device=device) )), desired_list=[(0,0)], sys_params=sys_params, device=coeffs.device)
        recon = wiener_torch(((measurement - 0.5) * 2).reshape(model.dataset.dim), psfs_estimate[0]) # can also use .wiener with ,balance=3e-4
        recon = (recon/2)+0.5 # back-scale

        # loss (maximizing acutance)
        loss = -torch.mean(torch.abs(kornia.spatial_gradient(recon[None, None, :, :], mode='diff')))
        print(iter, -loss.item())

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

        #coeffs.data[coeffs.data < 0] = 0

    print(coeffs)

    #torch.save(psfs_estimate[0], cache_path+'psf_fit.pt')
    return recon.detach().cpu().numpy(), coeffs.detach().cpu().numpy()

def tv(img, weight):
    tv_h = ((img[1:, :] - img[:-1, :]).abs()).sum()
    tv_w = ((img[:, 1:] - img[:, :-1]).abs()).sum()
    return weight * (tv_h + tv_w)

def center_crop(measurement, des_shape):
    # Center crop
    m_center = (measurement.shape[0]//2, measurement.shape[1]//2)
    left, right, up, down = ( m_center[1] - des_shape[1]//2, m_center[1] + int(np.round(des_shape[1]/2)),  \
                              m_center[0] - des_shape[0]//2, m_center[0] + int(np.round(des_shape[0]/2)))
    # TODO: Debug this for images of an odd size.
    measurement = measurement[left:right,up:down]
    return measurement

def wiener_torch(image, psf, balance=3e-4, reg=None, is_real=True, clip=True):
    if reg is None:
        reg, _ = laplacian(image.ndim, image.shape, is_real=is_real)
        reg = reg.to(device=image.device)

    trans_func = ir2tf(psf, image.shape, is_real=is_real).to(device=image.device)

    wiener_filter = torch.conj(trans_func) / (torch.abs(trans_func) ** 2 +
                                           balance * torch.abs(reg) ** 2)

    deconv = fft.irfftn(wiener_filter * fft.rfftn(image))

    if clip:
        deconv[deconv > 1] = 1
        deconv[deconv < -1] = -1

    return deconv

def laplacian(ndim, shape, is_real=True):
    impr = torch.zeros([3] * ndim)
    for dim in range(ndim):
        idx = tuple([slice(1, 2)] * dim +
                    [slice(None)] +
                    [slice(1, 2)] * (ndim - dim - 1))
        impr[idx] = torch.tensor([-1.0,
                              0.0,
                              -1.0]).reshape([-1 if i == dim else 1
                                              for i in range(ndim)])
    impr[(slice(1, 2), ) * ndim] = 2.0 * ndim
    return ir2tf(impr, shape, is_real=is_real), impr

def ir2tf(imp_resp, shape, dim=None, is_real=True):
    if not dim:
        dim = imp_resp.ndim
    # Zero padding and fill
    if imp_resp.shape != shape:
        irpadded = torch.zeros(shape)
        irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp
    else:
        irpadded = imp_resp
    # Roll for zero convention of the fft to avoid the phase
    # problem. Work with odd and even size.
    for axis, axis_size in enumerate(imp_resp.shape):
        if axis >= imp_resp.ndim - dim:
            irpadded = torch.roll(irpadded,
                               shifts=-int(np.floor(axis_size / 2)),
                               dims=axis)

    return fft.rfftn(irpadded, dim=list(range(-dim, 0)))

def image_quad_norm(image):
    np.sum(np.sum(np.abs(image) ** 2, axis=-1), axis=-1)


from lxml import html
import requests
import smtplib, ssl
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

# For SSL
port = 465

# Create a secure SSL context
context = ssl.create_default_context()

sender = "blantonswebscraper@gmail.com"
receiver = "jamesmcfadden111@gmail.com"
email = MIMEMultipart('test')
email["Subject"] = "Test email"
email["From"] = sender
email["To"] = receiver
test_str = "this string was inserted in the email body"
email_body = "This is a test " + test_str + " more email body"

email_content = MIMEText(email_body ,'html')
email.attach(email_content)
with smtplib.SMTP_SSL("smtp.gmail.com", port, context=context) as server:
    server.login("blantonswebscraper@gmail.com", "")
    server.sendmail(sender, 'jamesmcfadden111@gmail.com', email.as_string())

print("email sent")


import logging
import mimetypes
import traceback

from django.conf import settings
from django.contrib.auth.models import User
from django.core.mail import EmailMultiAlternatives, mail_admins
from django.db import models
from django.dispatch import receiver
from django_mailbox.signals import message_received
from rest_framework.authtoken.models import Token

from .models import TaskStore
from .tasks import process_email_message, autoconfigure_taskd

logger = logging.getLogger(__name__)

@receiver(
    models.signals.post_save,
    sender=User,
    dispatch_uid='generate_taskstore'
)
def create_taskstore_for_user(sender, instance, **kwargs):
    # This just makes sure that the task store exists.
    TaskStore.get_for_user(instance)
    Token.objects.get_or_create(user=instance)

@receiver(
    models.signals.post_save,
    sender=TaskStore,
    dispatch_uid='configure_taskstore'
)
def autoconfigure_taskstore_for_user(sender, instance, **kwargs):
    autoconfigure_taskd.apply_async(
        args=(instance.pk, ),
        countdown=5,
    )

@receiver(message_received, dispatch_uid='incoming_email')
def handle_incoming_message(sender, message, **kwargs):
    process_email_message.apply_async(args=(message.pk, ))

@receiver(message_received, dispatch_uid='forwardable_email')
def handle_incoming_forwardable_message(sender, message, **kwargs):
    for address in message.to_addresses:
        if address in settings.MAIL_FORWARDING:
            try:
                email = EmailMultiAlternatives(
                    subject=u' '.join([
                        u'[Inthe.AM]',
                        message.subject,
                        u'(%s)' % message.pk,
                    ]),
                    body=message.text,
                    to=[settings.MAIL_FORWARDING[address]],
                    reply_to=[message.from_header],
                    headers={
                        'IntheAM-Message-Id': message.pk
                    }
                )
                email.attach_alternative(
                    message.html,
                    'text/html',
                )
                for attachment in message.attachments.all():
                    email.attach(
                        attachment.get_filename(),
                        attachment.document.read(),
                        mimetypes.guess_type(
                            attachment.get_filename(),
                        )[0]
                    )
                email.send()
            except:
                mail_admins(
                    u"Error processing forwarding rule for %s" % address,
                    u"See message ID %s.\n\n%s" % (
                        message.pk,
                        traceback.format_exc(),
                    )
                )


#################################################
# import dependencies
from matplotlib import style
style.use('fivethirtyeight')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import datetime as dt
import sqlalchemy
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session
from sqlalchemy import create_engine, func
from flask import Flask, jsonify
#################################################
# Database Setup
#################################################
engine = create_engine("sqlite:///Resources/hawaii.sqlite")
# reflect an existing database into a new model
Base = automap_base()
# reflect the tables
Base.prepare(engine, reflect=True)
# Save references to each table
Measurement = Base.classes.measurement
Station = Base.classes.station
# Create our session (link) from Python to the DB
session = Session(engine)

#################################################
# Flask Setup
#################################################
app = Flask(__name__)

#################################################
# Flask Routes
#################################################
@app.route("/")
def welcome():
    """List all available api routes."""
    return (
        f"Available Routes:
"
        f"
"
        f"/api/v1.0/precipitation
"
        f"- List of prior year rain totals from all stations
"
        f"
"
        f"/api/v1.0/stations
"
        f"- List of Station numbers and names
"
        f"
"
        f"/api/v1.0/tobs
"
        f"- List of prior year temperatures from the most active station
"
        f"
"
        f"/api/v1.0/start
"
        f"- Put start date (YYYY-MM-DD), calculates the MIN/AVG/MAX temperature for all dates greater than and equal to the start date
"
        f"
"
        f"/api/v1.0/start/end
"
        f"- Put start and the end date (YYYY-MM-DD), calculate the MIN/AVG/MAX temperature for dates between the start and end date inclusive
"

    )
#########################################################################################

@app.route("/api/v1.0/precipitation")
def precipitation():
    print("Server received request for 'precipitation' page...")
# Convert the query results to a dictionary using date as the key and prcp as the value
    one_year_ago = dt.date(2017, 8, 23) - dt.timedelta(days=365)
    data = session.query(Measurement.date, Measurement.prcp).\
        filter(Measurement.date > one_year_ago).\
        order_by(Measurement.date).all()

# Return the JSON representation of your dictionary.
    data_list = []
    for row in data:
        append_data = {}
        append_data["date"] = row[0]
        append_data["prcp"] = row[1]
        data_list.append(append_data)
    return jsonify(data_list)

#########################################################################################
@app.route("/api/v1.0/stations")

# Return a JSON list of stations from the dataset.
def stations():
    print("Server received request for 'stations' page...")
    stations_query = session.query(Station.name, Station.station)
    stations = pd.read_sql(stations_query.statement, stations_query.session.bind)
    return jsonify(stations.to_dict())
#########################################################################################
@app.route("/api/v1.0/tobs")
def tobs():
    print("Server received request for 'tobs' page...")
# Query the dates and temperature observations of the most active station for the last year of data.
    most_active_stations = session.query(Measurement.station, func.count(Measurement.tobs)).group_by(Measurement.station).order_by(func.count(Measurement.tobs).desc()).all()
    most_active = most_active_stations[0][0]

    one_year_ago = dt.date(2017, 8, 23) - dt.timedelta(days=365)
    temp_data = session.query(Measurement.station,Measurement.date, Measurement.tobs).\
        filter(Measurement.date > one_year_ago).\
        filter(Measurement.station == most_active).\
        order_by(Measurement.date).all()

# Return a JSON list of temperature observations (TOBS) for the previous year.
    temperature_list = []
    for row in temp_data:
        append_temp_data = {}
        append_temp_data["date"] = row[1]
        append_temp_data["tobs"] = row[2]
        temperature_list.append(append_temp_data)

    return jsonify(temperature_list)
#########################################################################################
@app.route("/api/v1.0/<start>")
def trip1(start):

# Return a JSON list of the minimum temperature, the average temperature, and the max temperature for a given start or range
    start_date= dt.datetime.strptime(start, '%Y-%m-%d')
    last_year = dt.timedelta(days=365)
    start = start_date
    end =  dt.date(2017, 8, 23)
    trip_data = session.query(func.min(Measurement.tobs), func.avg(Measurement.tobs), func.max(Measurement.tobs)).\
        filter(Measurement.date >= start).filter(Measurement.date <= end).all()
    trip = {}
    trip['lowest_temp'] = trip_data[0][0]
    trip['average_temp'] = round(trip_data[0][1],0)
    trip['highest_temp'] = trip_data[0][2]
    return jsonify(trip)

#########################################################################################
@app.route("/api/v1.0/<start>/<end>")
def trip2(start,end):

# Return a JSON list of the minimum temperature, the average temperature, and the max temperature for a start-end    rangep
    start_date_1= dt.datetime.strptime(start, '%Y-%m-%d')
    end_date_1= dt.datetime.strptime(end,'%Y-%m-%d')
    start_1 = start_date_1
    end_1 = end_date_1
    trip_data_1 = session.query(func.min(Measurement.tobs), func.avg(Measurement.tobs), func.max(Measurement.tobs)).\
        filter(Measurement.date >= start_1).filter(Measurement.date <= end_1).all()
    trip_1 = {}
    trip_1['lowest_temp'] = trip_data_1[0][0]
    trip_1['average_temp'] = round(trip_data_1[0][1],0)
    trip_1['highest_temp'] = trip_data_1[0][2]
    return jsonify(trip_1)

#########################################################################################

if __name__ == "__main__":
    app.run(debug=True)


__version__ = "0.1.2"

import os
import logging
import sqlite3

_logger = logging.getLogger(os.path.basename(__file__))

import Objects

import normalizer

def main():
    conn = sqlite3.connect(args.out_db)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()

    cursor.execute("""\
CREATE TABLE recs (
  osetid TEXT NOT NULL,
  appetid TEXT NOT NULL,
  sliceid INTEGER NOT NULL,
  filename TEXT NOT NULL,
  hive_prefix TEXT NOT NULL
);\
""")

    #Build RE directory list from slice.db
    cursor.execute("ATTACH DATABASE '%s' AS slice;" % args.slice_db)
    cursor.execute("SELECT DISTINCT osetid, appetid, sliceid FROM slice.slice WHERE osetid <> '9544-1';")
    node_triples = []
    for row in cursor:
        node_triples.append((row["osetid"], row["appetid"], row["sliceid"]))
    cursor.execute("DETACH DATABASE slice;")

    #For each extraction.dfxml in an RE directory list
    for (osetid, appetid, sliceid) in sorted(node_triples):
        node_id = "%s-%s-%d" % (osetid, appetid, sliceid)
        extraction_dfxml_path = os.path.join(args.dwf_node_results_dir, node_id, "invoke_regxml_extractor.sh/extraction.dfxml")
        _logger.debug("extraction_dfxml_path = %r." % extraction_dfxml_path)

        if not os.path.exists(extraction_dfxml_path):
            _logger.warning("Skipping non-existent filepath: %r." % extraction_dfxml_path)
            continue

        #For each fileobject:
        for (event, obj) in Objects.iterparse(extraction_dfxml_path):
            if not isinstance(obj, Objects.FileObject):
                continue

            #Record node ID
            #Record node OS
            #Record hive filename
            #Record hive class (normalizer.py)
            hive_path = obj.original_fileobject.filename
            norm_prefix = normalizer.hive_path_to_prefix(hive_path)
            if norm_prefix is None:
                raise ValueError("File name has no associated prefix: %r." % obj.filename)
            cursor.execute("INSERT INTO recs(osetid, appetid, sliceid, filename, hive_prefix) VALUES (?,?,?,?,?);", (
              osetid,
              appetid,
              sliceid,
              hive_path,
              norm_prefix
            ))
    conn.commit()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--debug", action="store_true")
    parser.add_argument("slice_db")
    parser.add_argument("dwf_node_results_dir")
    parser.add_argument("out_db")
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)

    main()


from typing import List

class Solution:
    def combinationSum2(self, candidates: List[int], target: int) -> List[List[int]]:

        amount = target
        coins = candidates
        coins.sort()
        print(coins)
        len_coin = len(coins)
        # 一开始解决方案都是空数组[]
        dp = [[[] for j in range(amount + 1)] for i in range(len_coin + 1)]
        # dp i,j 表示用前i个凑满j的办法数量

        # 凑0我们是总有办法的
        for i in range(len_coin + 1):
            dp[i][0] = [[]]

        for i in range(1, len_coin + 1):
            for j in range(1, amount + 1):
                if j - coins[i - 1] >= 0:
                    # 不用这一个+用这一个
                    a = dp[i - 1][j]
                    b = [(x + [coins[i - 1]]) for x in dp[i-1][j - coins[i - 1]]]

                    dp[i][j] = a+b
                else:
                    # 只能不用这一个
                    dp[i][j] = dp[i - 1][j]
        res_dict = {}
        for i in dp[len_coin][target]:
            res_dict[str(i)] = i
        res = [v for k, v in res_dict.items()]
        return res

    def combinationSum2(self, candidates: List[int], target: int) -> List[List[int]]:
        def dfs(begin, path, residue):
            if residue == 0:
                res.append(path[:])
                return

            for index in range(begin, size):
                if candidates[index] > residue:
                    break

                if index > begin and candidates[index - 1] == candidates[index]:
                    continue

                path.append(candidates[index])
                dfs(index + 1, path, residue - candidates[index])
                path.pop()

        size = len(candidates)
        if size == 0:
            return []

        candidates.sort()
        res = []
        dfs(0, [], target)
        return res

print(Solution().combinationSum2(candidates = [1,2,2,2,5], target = 5,))


# Created on 2020/8/10

# This module is for analysing and classifying time series and other objects.

# Standard library imports
from datetime import datetime
from datetime import timedelta
import random as random

# Third party imports
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.cluster import KMeans
from sklearn.metrics import log_loss
from sklearn.metrics import silhouette_samples
from sklearn.model_selection._split import KFold

# Local application imports
from .. import timeseries

#---------#---------#---------#---------#---------#---------#---------#---------#---------#

# DISTANCES

def euclidean_distance(ts1, ts2):
    """
    Returns the Euclidean distance between two TimeSeries.

    Parameters
    ----------
    ts1 : TimeSeries
      First time series.
    ts2 : TimeSeries
      Second time series.

    Returns
    -------
    float
      Euclidean distance between time series.
    """

    # Checks
    try:
        assert(ts1.type=='TimeSeries' and ts2.type=='TimeSeries')
    except TypeError:
        raise TypeError("Series have to be of type TimeSeries.")

    try:
        assert(ts1.data.index.tolist() == ts2.data.index.tolist())
    except IndexError:
        raise IndexError("Time series do not have the same index.")

    # Return distance
    squares = (ts1.data - ts2.data)**2
    return np.sqrt(float(squares.sum()))

def dtw_distance(ts1, ts2, window=None):
    """
    Returns the Dynamic Time Warping (DTW) distance between two TimeSeries.
    A locality constraint can be used by specifying the size of a window.

    Parameters
    ----------
    ts1 : TimeSeries
      First time series.
    ts2 : TimeSeries
      Second time series.
    window : int
      Size of restrictive search window.

    Returns
    -------
    float
      DTW distance between time series.

    Notes
    -----
      To learn more about DTW, please refer to:
      https://en.wikipedia.org/wiki/Dynamic_time_warping

      Notice that taking a small window size may lead to
      a wrong estimate of the real dynamic time warping distance.
    """

    # Checks
    try:
        assert(ts1.type=='TimeSeries' and ts2.type=='TimeSeries')
    except TypeError:
        raise TypeError("Series have to be of type TimeSeries.")

    # Initializations
    N1 = len(ts1.data.index.tolist())
    N2 = len(ts2.data.index.tolist())
    if window is not None:
        assert(isinstance(window, int))
        w = window
    else:
        w = N2

    dtw = np.full(shape=(N1+1,N2+1), fill_value=np.inf)
    dtw[0,0] = 0

    # Loop
    for i in range(0, N1, 1):
        # for j in range(0, N2, 1):
        for j in range(max(0,int(i-w)), min(N2,int(i+w)), 1):
            square = (ts1.data.values[i] - ts2.data.values[j])**2
            dtw[i+1,j+1] = square + min(dtw[i,j+1], dtw[i+1,j], dtw[i,j])

    # Return distance
    return np.sqrt(dtw[N1, N2])

# KMEANS CLUSTERING

def kmeans_base_clustering(corr, names_features=None, max_num_clusters=10, n_init=10):
    """
    Perform base clustering with Kmeans.

    Arguments
    ---------
    corr: numpy.array
      Correlation matrix.
    names_features : list of str
      List of names for features.
    max_num_clusters: int
      Maximum number of clusters.
    n_init : int
      Initial value n_init for KMeans.

    Returns
    -------
    pd.DataFrame
      Clustered correlation matrix.
    dictionary
      List of clusters and their content.
    pd.Series
      Silhouette scores.

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """

    # Checks
    if not isinstance(max_num_clusters, int):
        raise AssertionError("max_num_clusters must be integer.")
    if not isinstance(n_init, int):
        raise AssertionError("n_init must be integer.")

    # Initializations
    corr = pd.DataFrame(data=corr, index=names_features, columns=names_features)
    silh_score = pd.Series()

    # Define the observations matrix X
    X = ((1 - corr.fillna(0))/2.)**.5

    # Loop to generate different initializations
    for init in range(n_init):

        # Loop to generate different numbers of clusters
        for i in range(2, max_num_clusters+1):

            # Define model and fit
            kmeans_ = KMeans(n_clusters=i, n_jobs=1, n_init=n_init).fit(X)

            # Compute silhouette coefficients
            silh_ = silhouette_samples(X, kmeans_.labels_)

            # Compute clustering quality q (t-statistic of silhouette score)
            stat = (silh_.mean()/silh_.std(), silh_score.mean()/silh_score.std())
            if np.isnan(stat[1]) or (stat[0]>stat[1]):
                silh_score, kmeans = silh_, kmeans_

    # Extract index according to sorted labels
    new_idx = np.argsort(kmeans.labels_)

    # Reorder rows
    clustered_corr = corr.iloc[new_idx]
    # Reorder columns
    clustered_corr = clustered_corr.iloc[:,new_idx]

    # Form clusters
    clusters = {i: clustered_corr.columns[np.where(kmeans.labels_==i)[0]].tolist() for i in np.unique(kmeans.labels_)}

    # Define a series with the silhouette score
    silh_score = pd.Series(silh_score, index=X.index)

    return clustered_corr, clusters, silh_score

def make_new_outputs(corr, clusters, clusters2):
    """
    Makes new outputs for kmeans_advanced_clustering().

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """
    clusters_new = {}
    for i in clusters.keys():
        clusters_new[len(clusters_new.keys())] = list(clusters[i])
    for i in clusters2.keys():
        clusters_new[len(clusters_new.keys())] = list(clusters2[i])

    new_idx = [j for i in clusters_new for j in clusters_new[i]]
    corr_new = corr.loc[new_idx, new_idx]
    x = ((1-corr.fillna(0))/2.)**.5

    kmeans_labels = np.zeros(len(x.columns))
    for i in clusters_new.keys():
        idxs = [x.index.get_loc(k) for k in clusters_new[i]]
        kmeans_labels[idxs] = i

    silh_new = pd.Series(silhouette_samples(x, kmeans_labels), index=x.index)

    return corr_new, clusters_new, silh_new

def kmeans_advanced_clustering(corr, names_features=None, max_num_clusters=None, n_init=10):
    """
    Perform advanced clustering with Kmeans.
    The base clustering is used first, then clusters quality is evaluated.
    For clusters whose quality less than the averaged quality,
    the clustering is reran.

    Arguments
    ---------
    corr: numpy.array
      Correlation matrix.
    names_features : list of str
      List of names for features.
    max_num_clusters: int
      Maximum number of clusters.
    n_init : int
      Initial value n_init for KMeans.

    Returns
    -------
    pd.DataFrame
      Clustered correlation matrix.
    dictionary
      List of clusters and their content.
    pd.Series
      Silhouette scores.

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """

    # Checks
    if (max_num_clusters is not None) and (not isinstance(max_num_clusters, int)):
        raise AssertionError("max_num_clusters must be integer.")
    if not isinstance(n_init, int):
        raise AssertionError("n_init must be integer.")

    # Initializations
    if max_num_clusters==None:
        max_num_clusters = corr.shape[1]-1
    # Using base clustering as initial step
    corr1, clusters, silh = kmeans_base_clustering(corr,
                                                   names_features=names_features,
                                                   max_num_clusters=min(max_num_clusters, corr.shape[1]-1),
                                                   n_init=n_init)

    # Compute t-stat for each cluster
    cluster_tstats = {i: np.mean(silh[clusters[i]]) / np.std(silh[clusters[i]]) for i in clusters.keys()}
    # Obtain the mean over clusters
    tstat_mean = sum(cluster_tstats.values()) / len(cluster_tstats)
    # Select the clusters which have a t-stat below the mean
    redo_clusters = [i for i in cluster_tstats.keys() if cluster_tstats[i] < tstat_mean]

    # Only one cluster
    if len(redo_clusters)<=1:
        return corr1, clusters, silh
    # More than one cluster
    else:
        keys_redo = [j for i in redo_clusters for j in clusters[i]]
        corr_tmp = corr.loc[keys_redo, keys_redo]
        tstat_mean = np.mean([cluster_tstats[i] for i in redo_clusters])
        corr2, clusters2, silh2 = kmeans_advanced_clustering(corr_tmp, max_num_clusters=min(max_num_clusters, corr_tmp.shape[1]-1), n_init=n_init)

        # Make new outputs, if necessary
        corr_new, clusters_new, silh_new = make_new_outputs(corr, {i: clusters[i] for i in clusters.keys() if i not in redo_clusters}, clusters2)
        new_tstat_mean = np.mean([ np.mean(silh_new[clusters_new[i]]) / np.std(silh_new[clusters_new[i]]) for i in clusters_new.keys() ])

        if new_tstat_mean <= tstat_mean:
            return corr1, clusters, silh
        else:
            return corr_new, clusters_new, silh_new

# FEATURE IMPORTANCE

def generate_random_classification(n_features, n_informative, n_redundant, n_samples, random_state=0, sigma_std=0.):
    """
    Generate a random dataset for a classification problem.

    Arguments
    ---------
    n_features : int
      Total number of features.
    n_informative : int
      Number of informative features.
    n_redundant : int
      Number of redundant features.
    n_samples : int
      Number of samples.
    random_state : int
      See random state.
    sigma_std: float
      Standard deviation of added noise.

    Returns
    -------
    pandas.DataFrame
      Data Frame with features as columns and samples as rows.
    pandas.Series
      Series containing class membership.

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """

    # Checks
    for arg in [('n_features',n_features), ('n_informative',n_informative), ('n_redundant',n_redundant),
                ('n_samples',n_samples), ('random_state',random_state)]:
        if not isinstance(arg[1],int):
            raise AssertionError(arg[0] + " must be integer.")
    if not isinstance(arg[1], float) and not isinstance(arg[1], int):
        raise AssertionError("sigma_std must be float.")

    # Initializations
    np.random.seed(random_state)

    # Generate classification
    X, y = make_classification(n_samples = n_samples,
                               n_features = n_features - n_redundant,
                               n_informative = n_informative,
                               n_redundant = 0,
                               shuffle = False,
                               random_state = random_state)

    # Informed explanatory variables
    cols = ['I_' + str(i) for i in range(n_informative)]

    # Noise explanatory variables
    cols += ['N_' + str(i) for i in range(n_features - n_informative - n_redundant)]

    X = pd.DataFrame(X, columns=cols)
    y = pd.Series(y)
    i = np.random.choice(range(n_informative), size=n_redundant)

    # Redundant explanatory variables
    for k, j in enumerate(i):
        X['R_' + str(k)] = X['I_' + str(j)] + np.random.normal(size=X.shape[0]) * sigma_std

    return X, y

def feature_importance_pvalues(fit, plot=False, figsize=(10,10)):
    """
    Plot the p-values of features from a fit.

    Arguments
    ---------
    fit : fit model
      Fit already apply on data.
    plot : bool
      Option to plot feature importance.
    figsize : (float, float)
      Dimensions of the plot.

    Returns
    -------
    pandas.DataFrame
      Data frame with features importance.
    """

    # Extract p-values and sort them
    pvals = fit.pvalues.sort_values(ascending=False)

    # Plot
    if plot is True:
        plt.figure(figsize=figsize)
        plt.title("P-values of features.")
        plt.barh(y=pvals.index, width=pvals.values)
        plt.show()

    # Make DataFrame
    pvals_df = pd.DataFrame(pvals)
    pvals_df.index.name = "Feature"
    pvals_df.columns = ["Importance"]

    return pvals_df

def feature_importance_mdi(classifier, X, y, plot=False, figsize=(10,10)):
    """
    Feature importance based on in-sample Mean-Decrease Impurity (MDI).

    Arguments
    ---------
    classifier : tree classifier
      Tree classifier to apply on data.
    X : pandas.DataFrame
      Data Frame with features as columns and samples as rows.
    y : pandas.Series
      Series containing class membership.
    plot : bool
      Option to plot feature importance.
    figsize : (float, float)
      Dimensions of the plot.

    Returns
    -------
    pandas.DataFrame
      Data frame with features importance.

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """

    # Checks
    if not isinstance(X, pd.DataFrame):
        raise AssertionError("X must be pandas.DataFrame.")
    if not isinstance(y, pd.Series):
        raise AssertionError("y must be pandas.Series.")

    # Fit
    fit = classifier.fit(X,y)

    # Extract feature importance
    df0 = {i: tree.feature_importances_ for i, tree in enumerate(fit.estimators_)}

    # Make Data Frame
    df0 = pd.DataFrame.from_dict(df0, orient='index')
    df0.columns = X.columns
    df0 = df0.replace(0, np.nan)
    fimp_df = pd.concat( {'Importance Mean': df0.mean(), 'Importance Std': df0.std() * df0.shape[0]**(-0.5)}, axis=1)
    fimp_df /= fimp_df['Importance Mean'].sum()
    fimp_df.index.name = "Feature"

    # Sort values
    sorted_fimp = fimp_df.sort_values(by='Importance Mean')

    # Plot
    if plot is True:
        plt.figure(figsize=figsize)
        plt.title("Feature importance based on in-sample Mean-Decrease Impurity (MDI).")
        plt.barh(y=sorted_fimp.index, width=sorted_fimp['Importance Mean'], xerr=sorted_fimp['Importance Std'])
        plt.show()

    return fimp_df

def feature_importance_mda(classifier, X, y, n_splits=10, plot=False, figsize=(10,10)):
    """
    Feature importance based out-of-sample Mean-Decrease Accuracy (MDA).

    Arguments
    ---------
    classifier : tree classifier
      Tree classifier to apply on data.
    X : pandas.DataFrame
      Data Frame with features as columns and samples as rows.
    y : pandas.Series
      Series containing class membership.
    plot : bool
      Option to plot feature importance.
    figsize : (float, float)
      Dimensions of the plot.

    Returns
    -------
    pandas.DataFrame
      Data frame with features importance.

    Notes
    -----
      Function adapted from "Machine Learning for Asset Managers",
      Marcos López de Prado (2020).
    """

    # Checks
    if not isinstance(X, pd.DataFrame):
        raise AssertionError("X must be pandas.DataFrame.")
    if not isinstance(y, pd.Series):
        raise AssertionError("y must be pandas.Series.")

    # Generate K-fold cross validation
    cv_gen = KFold(n_splits=n_splits)
    scr0 = pd.Series()
    scr1 = pd.DataFrame(columns=X.columns)

    # Loop over the folds:
    for i, (train, test) in enumerate(cv_gen.split(X=X)):

        # Train/Test split
        X0, y0 = X.iloc[train,:], y.iloc[train]
        X1, y1 = X.iloc[test,:], y.iloc[test]

        # Fit
        fit = classifier.fit(X=X0, y=y0)

        # Prediction before shuffling
        prob = fit.predict_proba(X1)
        scr0.loc[i] = -log_loss(y1, prob, labels=classifier.classes_)

        for j in X.columns:
            X1_ = X1.copy(deep=True)

            # Shuffle one column
            np.random.shuffle(X1_[j].values)

            # Prediction after shuffling
            prob = fit.predict_proba(X1_)
            scr1.loc[i,j] = -log_loss(y1, prob, labels=classifier.classes_)

    fimp_df = (-1*scr1).add(scr0, axis=0)
    fimp_df = fimp_df / (-1*scr1)
    fimp_df = pd.concat({'Importance Mean': fimp_df.mean(), 'Importance Std': fimp_df.std()*fimp_df.shape[0]**-.5}, axis=1)

    # Sort values
    sorted_fimp = fimp_df.sort_values(by='Importance Mean')

    # Plot
    if plot is True:
        plt.figure(figsize=figsize)
        plt.title("Feature importance based on out-of-sample Mean-Decrease Accuracy (MDA).")
        plt.barh(y=sorted_fimp.index, width=sorted_fimp['Importance Mean'], xerr=sorted_fimp['Importance Std'])
        plt.show()

    return sorted_fimp

#---------#---------#---------#---------#---------#---------#---------#---------#---------#


# Torch
import torch
from torch.utils.data import Dataset

# Torchvision
import torchvision.transforms as transforms

# General
import os
import json
import numpy as np

# Image libs
from PIL import Image, ImageOps
import matplotlib.pyplot as plt

IMAGE_FORMATS = (".png", ".PNG", ".jpg", ".JPG")

# -----------------------------------------------------------------------------
def load_json(fname):
    with open(fname) as f:
        out = json.load(f)
    return out

# -----------------------------------------------------------------------------
def clean_dir(dir_name):
    if len(os.listdir(dir_name)) > 0:
        files = os.listdir(dir_name)
        for f in files:
            os.remove(os.path.join(dir_name, f))

# -----------------------------------------------------------------------------
def dict_to_str(dict_var):
    output_str = ""
    for key, value in dict_var.items():
        if output_str:
            output_str += f"_{key}_{value}"
        else:
            output_str += f"{key}_{value}"
    return output_str

# -----------------------------------------------------------------------------
def fetch_dataset_files(dir: list, exclude_suffixes: list) -> list:
    """
    Build list of training data. Each data sample is a dict that contains the
    path to the input image and to the corresponding ground truth.
    """
    file_list = []
    for path, _, files in os.walk(dir):
        for name in files:
            filename = os.path.join(path, name)
            if name.endswith(IMAGE_FORMATS) and not any(
                map(name.__contains__, exclude_suffixes)
            ):
                file_list.append({"image": filename, "label": path.split(os.sep)[-1]})
    return file_list

# -----------------------------------------------------------------------------
def show_sample(sample, debug_fname):
    """
    Plot dataset sample.
    """
    image = np.array(transforms.ToPILImage()(sample["image"].squeeze(0))) / 255.0
    plt.figure()
    plt.imshow(image)
    plt.axis("off")
    plt.title(f"Class {sample['label_str']}")
    plt.tight_layout()
    plt.savefig(debug_fname, bbox_inches="tight", pad_inches=0)

# -----------------------------------------------------------------------------
class BirdsDataset(Dataset):
    """
    DSP dataset class.
    """

    def __init__(
        self,
        file_list: list,
        target_shape: int,
        num_classes: int,
        ordered_class_names: list,
        do_data_augmentation: bool,
    ) -> None:
        self.file_list = file_list
        self.target_shape = target_shape
        self.do_data_augmentation = do_data_augmentation
        self.num_classes = num_classes
        self.ordered_class_names = ordered_class_names

    # Returns size of dataset
    def __len__(self) -> int:
        return len(self.file_list)

    # Pre processing
    @staticmethod
    def pre_processing(img: Image, target_img_mode: str, target_shape: int) -> Image:
        if target_img_mode is not None:
            if target_img_mode == "GRAYSCALE":
                img = ImageOps.grayscale(img)
            else:
                img = img.convert(target_img_mode)

        # Resize image
        img = transforms.Resize(size=(target_shape, target_shape))(img)

        return img

    # Data augmentation
    @staticmethod
    def data_augmentation(image: Image) -> Image:
        # Color jitter
        j1 = transforms.ColorJitter(brightness=0, contrast=0, saturation=0.5, hue=0)
        j2 = transforms.ColorJitter(brightness=0, contrast=0.5, saturation=0, hue=0)
        j3 = transforms.ColorJitter(brightness=0.5, contrast=0, saturation=0, hue=0)
        jitter_transform = transforms.RandomChoice([j1, j2, j3])

        # Sharpner
        sharpner_transform = transforms.RandomAdjustSharpness(sharpness_factor=2)

        # Equalizer
        equalizer_transform = transforms.RandomEqualize()

        # Affine
        affine_transform = transforms.RandomAffine(
            degrees=(-30, 30), translate=(0.2, 0.2), shear=(-10, 10)
        )

        # Flip
        hflip_transform = transforms.RandomHorizontalFlip()
        vflip_transform = transforms.RandomVerticalFlip()

        # Apply random transforms
        applier = transforms.RandomApply(
            transforms=[
                jitter_transform,
                sharpner_transform,
                equalizer_transform,
                hflip_transform,
                vflip_transform,
                affine_transform,
            ],
            p=0.5,
        )
        image = applier(image)

        # Gaussian noise
        noise_prob = np.random.uniform(0, 1)
        if noise_prob > 0.5:
            image = (np.array(image) / 255).astype(np.float32)
            mean = 0.0
            std = 0.05
            image = image + np.random.normal(mean, std, image.shape)
            image = np.clip(image, 0, 1)  # we might get out of bounds so clip
            image = Image.fromarray(np.uint8(image * 255))

        return image

    # Gets dataset item
    def __getitem__(self, idx: int) -> dict:
        # Fetch data point
        data_pt = self.file_list[idx]

        # Read image and basic pre-processing
        image = Image.open(data_pt["image"])
        image = self.pre_processing(image, "RGB", self.target_shape)

        # Do data augmentation
        if self.do_data_augmentation:
            image = self.data_augmentation(image)

        # Fetch label
        label_id = int(self.ordered_class_names.index(data_pt["label"]))
        label_one_hot_encoded = torch.zeros(self.num_classes)
        label_one_hot_encoded[label_id] = 1

        # Create tensors
        image = transforms.ToTensor()(image)
        image = image.float()

        sample = {
            "image": image,
            "label_one_hot": label_one_hot_encoded,
            "label_str": data_pt["label"],
        }

        return sample


#!/usr/bin/env python3
import numpy as np
import os,time
import imageio
from rlstudy.common import logger
from rlstudy.common.utils import make_atari_env, atari_arg_parser
from rlstudy.a2c.a2c import learn,Model
from rlstudy.a2c.policies import CnnPolicy
from rlstudy.common.atari_wrappers import make_atari, wrap_deepmind
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
START=223000
def make_model(env, total_timesteps=int(80e6),
              vf_coef=0.5, ent_coef=0.01, max_grad_norm=0.5, lr=7e-4,
              epsilon=1e-5, alpha=0.99):
    return Model(policy=CnnPolicy,ob_space=env.observation_space,
                  ac_space=env.action_space, nenvs=1,ent_coef=ent_coef, vf_coef=vf_coef, max_grad_norm=max_grad_norm,
                  lr=lr, alpha=alpha, epsilon=epsilon, total_timesteps=total_timesteps)
def play(**wrapper_kwargs):
    parser = atari_arg_parser()
    args=parser.parse_args()
    logger.configure()
    env = make_atari(args.env)
    env.seed(1234)
    env=wrap_deepmind(env, frame_stack=True,**wrapper_kwargs)

    model =make_model(env )
    model.load('models/'+args.env+'.save'+'.'+str(START))
    renders=[]

    obs = env.reset() #(84, 84, 4)

    for i in range(10000):

        obs = np.expand_dims(obs.__array__(), axis=0)#(1,84, 84, 4)
        a, v, _ ,_= model.step(obs)
        obs, reward, done, info = env.step(a[0])
        if i%10==0:
            img=env.render('rgb_array')
            renders.append(imageio.core.util.Image(img))
        else:
            env.render()
        if done:
            name = 'imgs/' + str(int(time.time())) + '.gif'
            imageio.mimsave(name, renders, duration=1 / 30)
            renders = []
            env.reset()

    env.close()

if __name__ == '__main__':
    play()


import sys
import os
import math

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import rbf_kernel
from sklearn.metrics.pairwise import laplacian_kernel
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import OrdinalEncoder

class ADL_model:

    """ Class that boundles all parameters, methods and results for performing
    active learning.
    """

    def __init__(
        self,
        name='adl_model',
        path_to_results='results/'
    ):

        ### Parameters
        self.name = name
        if path_to_results.endswith('/'):
            self.path_to_results = path_to_results
        else:
            self.path_to_results = path_to_results + '/'

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

    ### Methods

    def initialize(
        self,
        y,
        x_t=None,
        x_s=None,
        x_st=None,
        encoder_layers=1,
        network_layers=1,
        encoding_nodes_x_t=100,
        encoding_nodes_x_s=100,
        encoding_nodes_x_st=100,
        encoding_nodes_joint=100,
        nodes_per_layer_dense=1000,
        filters_per_layer_cnn=16,
        states_per_layer_rnn=200,
        activation_encoding='relu',
        activation_dense='relu',
        activation_cnn='relu',
        activation_rnn='tanh',
        layer_type_x_st='CNN',
        initialization_method='glorot_normal',
        initialization_method_rnn='orthogonal',
        regularizer='l1_l2',
        batch_normalization=False,
        train_split=0.7,
        split_intervals=0.05,
        random_seed=None,
        epochs=30,
        patience=10,
        batch_size=16,
        monitor='val_loss',
        silent=True,
        plot=False
    ):

        """ Initializes parameters passed to model as class attributes and calls
        methods for creating modular deep learning model (embedding network),
        splitting data into training and validation sets, and training embedding
        network prediction model.
        """

        ### Set all parameters as attributes of class:
        self.y = y
        self.x_t = x_t
        self.x_s = x_s
        self.x_st = x_st
        self.encoder_layers = encoder_layers
        self.network_layers = network_layers
        self.encoding_nodes_x_t = encoding_nodes_x_t
        self.encoding_nodes_x_s = encoding_nodes_x_s
        self.encoding_nodes_x_st = encoding_nodes_x_st
        self.encoding_nodes_joint = encoding_nodes_joint
        self.nodes_per_layer_dense = nodes_per_layer_dense
        self.filters_per_layer_cnn = filters_per_layer_cnn
        self.states_per_layer_rnn = states_per_layer_rnn
        self.activation_encoding = activation_encoding
        self.activation_dense = activation_dense
        self.activation_cnn = activation_cnn
        self.activation_rnn = activation_rnn
        self.layer_type_x_st = layer_type_x_st
        self.initialization_method = initialization_method
        self.initialization_method_rnn = initialization_method_rnn
        self.regularizer = regularizer
        self.batch_normalization = batch_normalization
        self.train_split = train_split
        self.split_intervals = split_intervals
        self.random_seed = random_seed
        self.epochs = epochs
        self.patience = patience
        self.batch_size = batch_size
        self.monitor = monitor
        self.silent = silent
        self.plot = plot

        ### Call methods
        self.create_prediction_model()
        self.split_train_val()
        self.train_model(fig_name='train_vs_val_initial')

    def train(
        self,
        y_picked,
        x_t_picked=None,
        x_s_picked=None,
        x_st_picked=None,
        silent=True,
        plot=False,
    ):

        """ Trains embedding network prediction model with passed data. Note:
        It does not split the passed data into training and validation sets, but
        rather uses the validation data created when calling ADL_model.initialize().
        """

        self.y_train = y_picked
        self.x_t_train = x_t_picked
        self.x_s_train = x_s_picked
        self.x_st_train = x_st_picked
        self.silent = silent
        self.plot = plot

        self.train_model(fig_name='train_vs_val_adl')

    def create_prediction_model(self):

        """ Creates the deep learning models that compound our modular computational
        graph (embedding network).
        """

        ### Set initialization methods

        if self.initialization_method=='glorot_normal':
            self.initialization = tf.keras.initializers.GlorotNormal(
                seed=self.random_seed
            )

        if self.initialization_method_rnn=='orthogonal':
            self.initialization_rnn = tf.keras.initializers.Orthogonal(
                gain=1.0,
                seed=self.random_seed
            )

        if not self.silent:
            # tell us what we do
            print('Building deep learning model')

        # take example label
        y_example = self.y[0]

        # check if features provided and take examples, otherwise, throw out error
        if self.x_t is not None:
            x_t_example = self.x_t[0]
        if self.x_s is not None:
            x_s_example = self.x_s[0]
        if self.x_st is not None:
            x_st_example = self.x_st[0]
        if self.x_t is None and self.x_s is None and self.x_st is None:
            print(
                'At least one feature type x_t, x_s, or x_st must be provided'
                + 'to train prediction model'
            )
            sys.exit()

        ### Create the input layers ###

        if self.x_t is not None:
            x_t_input = tf.keras.Input(shape=x_t_example.shape, name='X_t')
        if self.x_s is not None:
            x_s_input = tf.keras.Input(shape=x_s_example.shape, name='X_s')
        if self.x_st is not None:
            x_st_input = tf.keras.Input(shape=x_st_example.shape, name='X_st')

        ### Create the hidden layers ###

        ### Encode x_t ###
        if self.x_t is not None:

            if self.encoder_layers == 0:
                x_t = tf.keras.layers.Flatten()(x_t_input)

            else:
                x_t = tf.keras.layers.Dense(
                    self.nodes_per_layer_dense,
                    activation=self.activation_dense,
                    kernel_initializer=self.initialization,
                    kernel_regularizer=self.regularizer,
                )(x_t_input)

                if self.batch_normalization:
                    x_t = tf.keras.layers.BatchNormalization()(x_t)

                for i in range(self.encoder_layers - 1):
                    x_t = tf.keras.layers.Dense(
                        self.nodes_per_layer_dense,
                        activation=self.activation_dense,
                        kernel_initializer=self.initialization,
                        kernel_regularizer=self.regularizer,
                    )(x_t)

                    if self.batch_normalization:
                        x_t = tf.keras.layers.BatchNormalization()(x_t)

                x_t = tf.keras.layers.Flatten()(x_t)

            X_t = tf.keras.layers.Dense(
                self.encoding_nodes_x_t,
                activation=self.activation_encoding,
                kernel_initializer=self.initialization,
                kernel_regularizer=self.regularizer,
            )(x_t)

            if self.batch_normalization:
                x_t = tf.keras.layers.BatchNormalization()(x_t)

        ### Encode x_s ###

        if self.x_s is not None:

            if self.encoder_layers == 0:
                x_s = tf.keras.layers.Flatten()(x_s_input)

            else:
                x_s = tf.keras.layers.Dense(
                    self.nodes_per_layer_dense,
                    activation=self.activation_dense,
                    kernel_initializer=self.initialization,
                    kernel_regularizer=self.regularizer,
                )(x_s_input)

                if self.batch_normalization:
                    x_s = tf.keras.layers.BatchNormalization()(X_s1)

                for i in range(self.encoder_layers - 1):
                    x_s = tf.keras.layers.Dense(
                        self.nodes_per_layer_dense,
                        activation=self.activation_dense,
                        kernel_initializer=self.initialization,
                        kernel_regularizer=self.regularizer,
                    )(x_s)

                    if self.batch_normalization:
                        x_s = tf.keras.layers.BatchNormalization()(x_s)

                x_s = tf.keras.layers.Flatten()(x_s)

            x_s = tf.keras.layers.Dense(
                self.encoding_nodes_x_s,
                activation=self.activation_dense,
                kernel_initializer=self.initialization,
                kernel_regularizer=self.regularizer,
            )(x_s)

            if self.batch_normalization:
                x_s = tf.keras.layers.BatchNormalization()(x_s)

        ### Encode x_st ###

        if self.x_st is not None:

            if self.encoder_layers == 0:
                x_st = tf.keras.layers.Flatten()(x_st_input)

            else:
                if self.layer_type_x_st == 'ANN':
                    x_st = tf.keras.layers.Dense(
                        self.nodes_per_layer_dense,
                        activation=self.activation_dense,
                        kernel_initializer=self.initialization,
                        kernel_regularizer=self.regularizer,
                    )(x_st_input)

                    if self.batch_normalization:
                        X_st = tf.keras.layers.BatchNormalization()(X_st)

                    for i in range(self.encoder_layers - 1):
                        x_st = tf.keras.layers.Dense(
                            self.nodes_per_layer_dense,
                            activation=self.activation_dense,
                            kernel_initializer=self.initialization,
                            kernel_regularizer=self.regularizer,
                        )(x_st)

                        if self.batch_normalization:
                            X_st = tf.keras.layers.BatchNormalization()(X_st)

                elif self.layer_type_x_st == 'CNN':
                    x_st = tf.keras.layers.Conv1D(
                        self.filters_per_layer_cnn,
                        2,
                        activation=self.activation_cnn,
                        kernel_initializer=self.initialization,
                        kernel_regularizer=self.regularizer,
                    )(x_st_input)

                    if self.batch_normalization:
                        x_st = tf.keras.layers.BatchNormalization()(x_st)

                    for i in range(self.encoder_layers - 1):
                        x_st = tf.keras.layers.Conv1D(
                            self.filters_per_layer_cnn,
                            2,
                            activation=self.activation_cnn,
                            kernel_initializer=self.initialization,
                            kernel_regularizer=self.regularizer,
                        )(x_st)

                        if self.batch_normalization:
                            x_st = tf.keras.layers.BatchNormalization()(x_st)

                elif self.layer_type_x_st == 'LSTM':
                    if self.encoder_layers == 1:
                        x_st = tf.keras.layers.LSTM(
                            self.states_per_layer_rnn,
                            activation=self.activation_rnn,
                            kernel_initializer=self.initialization,
                            kernel_regularizer=self.regularizer,
                        )(x_st_input)

                        if self.batch_normalization:
                            x_st = tf.keras.layers.BatchNormalization()(x_st)

                    else:
                        x_st = tf.keras.layers.LSTM(
                            self.states_per_layer_rnn,
                            return_sequences=True,
                            activation=self.activation_rnn,
                            kernel_initializer=self.initialization,
                            kernel_regularizer=self.regularizer,
                        )(x_st_input)

                        if self.batch_normalization:
                            x_st = tf.keras.layers.BatchNormalization()(x_st)

                        for i in range(self.encoder_layers - 2):
                            x_st = tf.keras.layers.LSTM(
                                self.states_per_layer_rnn,
                                return_sequences=True,
                                activation=self.activation_rnn,
                                kernel_initializer=self.initialization,
                                kernel_regularizer=self.regularizer,
                            )(x_st)

                            if self.batch_normalization:
                                X_st = tf.keras.layers.BatchNormalization()(X_st)

                        x_st = tf.keras.layers.LSTM(
                            self.states_per_layer_rnn,
                            activation=self.activation_rnn,
                            kernel_initializer=self.initialization,
                            kernel_regularizer=self.regularizer,
                        )(x_st)

                        if self.batch_normalization:
                            x_st = tf.keras.layers.BatchNormalization()(x_st)

                x_st = tf.keras.layers.Flatten()(x_st)

            x_st = tf.keras.layers.Dense(
                self.encoding_nodes_x_st,
                activation=self.activation_encoding,
                kernel_initializer=self.initialization,
                kernel_regularizer=self.regularizer,
            )(x_st)

            if self.batch_normalization:
                x_st = tf.keras.layers.BatchNormalization()(x_st)

        ### Create and join the encoders ###

        # create empty lists for joing layers and inputs of total prediction model
        input_list = []
        joining_list = []

        ### create x_t encoder ###
        if self.x_t is not None:
            x_t_encoder = tf.keras.Model(inputs=x_t_input, outputs=x_t)
            input_list.append(x_t_input)
            joining_list.append(x_t)

        ### create x_s encoder ###
        if self.x_s is not None:
            x_s_encoder = tf.keras.Model(inputs=x_s_input, outputs=x_s)
            input_list.append(x_s_input)
            joining_list.append(x_s)

        ### create x_st encoder ###
        if self.x_st is not None:
            x_st_encoder = tf.keras.Model(inputs=x_st_input, outputs=x_st)
            input_list.append(x_st_input)
            joining_list.append(x_st)

        ### create joint encoder ###

        if len(joining_list) > 1:
            joining_layer = tf.keras.layers.concatenate(joining_list)
        else:
            joining_layer = joining_list[0]

        for i in range(self.encoder_layers):
            joining_layer = tf.keras.layers.Dense(
                self.nodes_per_layer_dense,
                activation=self.activation_dense,
                kernel_initializer=self.initialization,
                kernel_regularizer=self.regularizer,
            )(joining_layer)

            if self.batch_normalization:
                joining_layer = tf.keras.layers.BatchNormalization()(joining_layer)

        joining_layer = tf.keras.layers.Dense(
            self.encoding_nodes_joint,
            activation=self.activation_encoding,
            kernel_initializer=self.initialization,
            kernel_regularizer=self.regularizer,
        )(joining_layer)

        if self.batch_normalization:
            joining_layer = tf.keras.layers.BatchNormalization()(joining_layer)

        x_joint_encoder = tf.keras.Model(inputs=input_list, outputs=joining_layer)

        ### Create total prediction model ###

        for i in range(self.network_layers):
            joining_layer = tf.keras.layers.Dense(
                self.nodes_per_layer_dense,
                activation=self.activation_dense,
                kernel_initializer=self.initialization,
                kernel_regularizer=self.regularizer,
            )(joining_layer)

            if self.batch_normalization:
                joining_layer = tf.keras.layers.BatchNormalization()(joining_layer)

        model_output = tf.keras.layers.Dense(
            len(y_example),
            activation='softplus',
            kernel_initializer=self.initialization,
        )(joining_layer)

        # create the tf model and define its inputs and outputs
        f_nn = tf.keras.Model(
            inputs=input_list,
            outputs=model_output
        )

        # create class instance for encoding and prediction models
        models = {}
        if self.x_t is not None:
            models['x_t_encoder'] = x_t_encoder
        if self.x_s is not None:
            models['x_s_encoder'] = x_s_encoder
        if self.x_st is not None:
            models['x_st_encoder'] = x_st_encoder
        models['x_joint_encoder'] = x_joint_encoder
        models['f_nn'] = f_nn

        self.models = models

        # give us the summary of the total prediction model that we define
        if not self.silent:
            f_nn.summary()

    def split_train_val(self):

        """ Splits available data into training and validation sets.
        """

        ###
        # Split remaining into training and validation datasets using intervals ###
        ###

        split_bins = math.ceil(len(self.y) * self.split_intervals)

        random_array = np.arange(len(self.y))
        random_array = np.array_split(random_array, split_bins)
        np.random.shuffle(random_array)
        random_array = np.concatenate(random_array).ravel()

        self.y = self.y[random_array]
        if self.x_t is not None:
            self.x_t = self.x_t[random_array]
        if self.x_s is not None:
            self.x_s = self.x_s[random_array]
        if self.x_st is not None:
            self.x_st = self.x_st[random_array]

        # get splitting point for training validation split
        split_point = math.ceil(self.train_split * len(self.y))

        # split train and delete unused entries immediately
        self.y_train, self.y_val = np.split(self.y, [split_point])
        self.y = 0

        if self.x_t is not None:
            self.x_t_train, self.x_t_val = np.split(self.x_t, [split_point])
            self.x_t = 0

        if self.x_s is not None:
            self.x_s_train, self.x_s_val = np.split(self.x_s, [split_point])
            self.x_s = 0

        if self.x_st is not None:
            self.x_st_train, self.x_st_val = np.split(self.x_st, [split_point])
            self.x_st = 0

    def f_randomize(self, data):

        """ Randomizes data order for training and validation.
        """

        # randomize training data
        if data == 'train':

            # create random array
            random_array = np.arange(len(self.y_train))
            self.y_train = self.y_train[random_array]

            if self.x_t is not None:
                self.x_t_train = self.x_t_train[random_array]

            if self.x_s is not None:
                self.x_s_train = self.x_s_train[random_array]

            if self.x_st is not None:
                self.x_st_train = self.x_st_train[random_array]

        elif data == 'val':

            # create random array
            random_array = np.arange(len(self.y_val))
            self.y_val = self.y_val[random_array]

            if self.x_t is not None:
                self.x_t_val = self.x_t_val[random_array]

            if self.x_s is not None:
                self.x_s_val = self.x_s_val[random_array]

            if self.x_st is not None:
                self.x_st_val = self.x_st_val[random_array]

    def create_batched_data(self, data, i):

        """ Creates data batches for training and validation.
        """

        if data == 'train':

            data_y = self.y_train
            if self.x_t is not None:
                data_x_t = self.x_t_train
            if self.x_st is not None:
                data_x_st = self.x_st_train
            if self.x_s is not None:
                data_x_s = self.x_s_train

        elif data == 'val':

            data_y = self.y_val
            if self.x_t is not None:
                data_x_t = self.x_t_val
            if self.x_s is not None:
                data_x_s = self.x_s_val
            if self.x_st is not None:
                data_x_st = self.x_st_val

        # build batches of data
        for j in range(self.batch_size):

            # Get training data of currently iterated batch

            y = data_y[i + j]
            if self.x_t is not None:
                x_t = data_x_t[i + j]
            if self.x_s is not None:
                x_s = data_x_s[i + j]
            if self.x_st is not None:
                x_st = data_x_st[i + j]

            # Expand dimensions for batching
            y = np.expand_dims(y, axis=0)
            if self.x_t is not None:
                x_t = np.expand_dims(x_t, axis=0)
            if self.x_s is not None:
                x_s = np.expand_dims(x_s, axis=0)
            if self.x_st is not None:
                x_st = np.expand_dims(x_st, axis=0)

            # Create batches
            if j == 0:

                y_batched = y
                if self.x_t is not None:
                    x_t_batched = x_t
                if self.x_s is not None:
                    x_s_batched = x_s
                if self.x_st is not None:
                    x_st_batched = x_st

            else:

                y_batched = np.concatenate((y_batched, y), axis=0)
                if self.x_t is not None:
                    x_t_batched = np.concatenate((x_t_batched, x_t), axis=0)
                if self.x_s is not None:
                    x_s_batched = np.concatenate((x_s_batched, x_s), axis=0)
                if self.x_st is not None:
                    x_st_batched = np.concatenate((x_st_batched, x_st), axis=0)

        # Create model input list
        model_input_list = []
        if self.x_t is not None:
            model_input_list.append(x_t_batched)
        if self.x_s is not None:
            model_input_list.append(x_s_batched)
        if self.x_st is not None:
            model_input_list.append(x_st_batched)

        return model_input_list, y_batched

    def train_model(
        self,
        fig_name='training'
    ):
        """ Functions for training and validating model.
        """

        if self.silent:
            verbose = 0
        else:
            verbose = 1

        # initialize loss functions and optimizer
        loss_object = tf.keras.losses.mean_squared_error
        optimizer = tf.keras.optimizers.RMSprop()
        mean_loss = tf.keras.metrics.Mean(name='mean_loss_train_test')

        model = self.models['f_nn']

        # define the training step to execute in each iteration with magic
        @tf.function
        def train_step(model_input_list, y_data):

            with tf.GradientTape() as tape:
                predictions = model(model_input_list, training=True)
                loss = loss_object(predictions, y_data)

            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(
                zip(gradients, model.trainable_variables)
            )
            mean_loss(loss)

            return loss

        # define the test step to execute in each iteration with a magic handle
        @tf.function
        def test_step(model_input_list, y_data):

            predictions = model(model_input_list, training=False)
            t_loss = loss_object(predictions, y_data)
            mean_loss(t_loss)

        ###
        # Perform epochs of training and validation ###
        ###

        # create lists to save train and validation loss results for each epoch
        val_loss_history = []
        train_loss_history = []

        for epoch in range(self.epochs):

            if not self.silent:

                # tell which epoch we are at
                print('Epoch {}/{}'.format(epoch + 1, self.epochs))

            ###
            # Training ###
            ###

            # Reset the metrics at the start of the next epoch
            mean_loss.reset_states()

            # Shuffle training data
            self.f_randomize('train')

            # get number of datapoints
            n_datapoints = len(self.y_train)

            if not self.silent:

                # tell us that we start training now
                print('Training:')

                # create a progress bar for training
                progbar = tf.keras.utils.Progbar(
                    math.floor(n_datapoints - self.batch_size) / self.batch_size,
                    stateful_metrics=['loss'],
                )

            # iterate over training data in BATCH_SIZE steps
            for i in range(
                0,
                n_datapoints - self.batch_size,
                self.batch_size
            ):

                # call function to create batched model inputs and labels
                model_input_list, y_batched = self.create_batched_data('train', i)

                # Execute the training step for this batch
                train_step(model_input_list, y_batched)

                # update the progress bar
                if not self.silent:

                    values = [('loss', mean_loss.result().numpy())]
                    progbar.add(1, values=values)

            # add training loss to history
            train_loss_history.append(mean_loss.result().numpy())

            ###
            # Validation ###
            ###

            # Reset the metrics at the start of the next epoch
            mean_loss.reset_states()

            # Shuffle validation data
            self.f_randomize('val')

            # get number of datapoints
            n_datapoints = len(self.y_val)

            if not self.silent:

                # Tell us that we start validating now
                print('Validation:')

                # create a progress bar for validation
                progbar = tf.keras.utils.Progbar(
                    math.floor(n_datapoints - self.batch_size) / self.batch_size,
                    stateful_metrics=['loss'],
                )

            # iterate over validation data in BATCH_SIZE steps
            for i in range(
                0,
                n_datapoints - self.batch_size,
                self.batch_size
            ):

                # call function to create batched model inputs and labels
                model_input_list, y_batched = self.create_batched_data('val', i)

                # Execute the testing step for this batch
                test_step(model_input_list, y_batched)

                # update the progress bar
                if not self.silent:

                    values = [('loss', mean_loss.result().numpy())]
                    progbar.add(1, values=values)

            # add validation loss to history
            val_loss_history.append(mean_loss.result().numpy())

            ###
            # implement early break here ###
            ###

            if self.monitor == 'val_loss':
                current_minimum = min(val_loss_history[-self.patience :])

            elif self.monitor == 'loss':
                current_minimum = min(train_loss_history[-self.patience :])

            elif self.monitor is None:
                continue

            if epoch == 0:
                total_minimum = current_minimum

            if current_minimum > total_minimum:
                break

            else:
                total_minimum = min(total_minimum, current_minimum)

        # Plot training and validation history
        if self.plot:

            plt.figure(figsize=(16, 8))
            plt.plot(train_loss_history)
            plt.plot(val_loss_history)
            plt.title('Training and validation loss history of neural network')
            plt.ylabel('loss')
            plt.xlabel('epoch')
            plt.legend(['training', 'validation'], loc='upper left')
            plt.show()

            # save figure
            file_name = fig_name +'.pdf'
            saving_path = self.path_to_results + file_name
            plt.savefig(saving_path)

        self.train_loss_history = train_loss_history
        self.val_loss_history = val_loss_history

    def encode_features(self):

        """ Encodes features using modules of embedding network.
        """

        if not self.silent:
            # tell us what we are doing
            print('Encoding features into embedded vector spaces for')

        ### Create random subsample before encoding if wanted ###

        # create an index array in the length of the passed dataset
        n_datapoints = self.n_datapoints_cand
        index_array = list(np.arange(n_datapoints))

        # if we chose a subset of these data points, create a random sub-sample
        if (
            self.subsample is not None and
            self.subsample > 0 and
            self.subsample < 1
        ):

            n_datapoints = math.ceil(self.subsample * n_datapoints)
            index_array = random.sample(index_array, n_datapoints)

        # create copy of dataset
        if self.x_t_cand is not None:
            x_t = self.x_t_cand[index_array]
        if self.x_s_cand is not None:
            x_s = self.x_s_cand[index_array]
        if self.x_st_cand is not None:
            x_st = self.x_st_cand[index_array]

        ### Encode features here ###
        if (
            self.x_t_cand is not None and
            self.x_s_cand is not None and
            self.x_st_cand is not None
        ):
            model = self.models['x_joint_encoder']
            self.encoding = model.predict([x_t, x_s, x_st])
        elif (
            self.x_t_cand is not None and
            self.x_s_cand is not None
        ):
            model = self.models['x_joint_encoder']
            self.encoding = model.predict([x_t, x_s])
        elif (
            self.x_t_cand is not None and
            self.x_st_cand is not None
        ):
            model = self.models['x_joint_encoder']
            self.encoding = model.predict([x_t, x_st])
        elif (
            self.x_s_cand is not None and
            self.x_st_cand is not None
        ):
            model = self.models['x_joint_encoder']
            self.encoding = model.predict([x_s, x_st])
        elif self.x_st_cand is not None:
            model = self.models['x_st_encoder']
            self.encoding = model.predict(x_st)

        elif self.x_t_cand is not None:
            model = self.models['x_t_encoder']
            self.encoding = model.predict(x_t)

        elif self.x_s_cand is not None:
            model = self.models['x_s_encoder']
            self.encoding = model.predict(x_s)

        # save potentially subsampled index array
        self.index_array = index_array

    def compute_clusters(self):

        """ Computes clusters in feature embedded vector space.
        """

        if not self.silent:
            # tell us what we are doing
            print(
                'Creating clusters in encodings with n_clusters=', self.adl_batch_size
            )

        # set number of clusters equal to passed or corrected value
        clustering_method = self.method_cluster(n_clusters=self.adl_batch_size)

        # cluster encodings
        clustering_method.fit(self.encoding)
        cluster_labels = clustering_method.labels_
        cluster_centers = clustering_method.cluster_centers_

        # get ordinal encoder from Sklearn
        enc = OrdinalEncoder()

        # encode labels. NOTE: ordinally encoding clusters ensures that cluster
        # labels start at 0 and end at number of clusters, which is not the case
        # for X_t and X_s1 when not ordinally encoding.
        cluster_labels = enc.fit_transform(
            np.expand_dims(cluster_labels, 1)
        ).astype(int)

        # save cluster centers
        self.cluster_centers = cluster_centers[enc.categories_[0]]

        # delete expanded dimension again as it is redundant
        self.cluster_labels = cluster_labels[:, 0]

        # calculate number of clusters created in data
        self.n_clusters = max(cluster_labels) + 1

    def compute_similarity(self):

        """ Compute similarity score of candidates to cluster centers.
        """

        # set the number of encoded data points
        n_enc_datapoints = len(self.encoding)

        if not self.silent:
            # tell us what we are doing
            print("Calculating distances" )

            # create a progress bar for training
            progbar_distance = tf.keras.utils.Progbar(n_enc_datapoints)

        # CAUTION: create shape (n_enc_datapoints,) instead of (n_enc_datapoints, 1)
        similarity_array = np.zeros((n_enc_datapoints,))

        # iterate over all encodings
        for i in range(n_enc_datapoints):

            # get encoding's cluster label
            label = self.cluster_labels[i]

            # get cluster's center
            center = self.cluster_centers[label]

            # calculate similarity/closeness of encoding to its cluster center
            similarity_array[i] = self.method_distance(
                np.expand_dims(center, axis=0), np.expand_dims(self.encoding[i], axis=0)
            )

            if not self.silent:
                # increment progress bar
                progbar_distance.add(1)

        # save results
        self.similarity_array = similarity_array

    def collect(
        self,
        x_t_cand=None,
        x_s_cand=None,
        x_st_cand=None,
        budget=0.5,
        method='embedding_uncertainty',
        method_variant='max_uncertainty',
        method_distance='laplacian_kernel',
        method_cluster='KMeans',
        subsample=None,
        silent=True,
        plot=False
    ):

        """ Perform pool-based active learning using embedding uncertainty.
        """

        self.x_t_cand = x_t_cand
        self.x_s_cand = x_s_cand
        self.x_st_cand = x_st_cand
        self.budget = 0.5
        self.method = method
        self.method_variant = method_variant
        exec('self.method_distance = {}'.format(method_distance))
        exec('self.method_cluster = {}'.format(method_cluster))
        self.subsample = subsample
        self.silent = silent
        self.plot = plot

        # set number of datapoints
        if x_t_cand is not None:
            self.n_datapoints_cand = len(x_t_cand)
        elif x_s_cand is not None:
            self.n_datapoints_cand = len(x_s_cand)
        elif x_st_cand is not None:
            self.n_datapoints_cand = len(x_st_cand)

        self.adl_batch_size = math.ceil(self.n_datapoints_cand * budget)

        ### Encode data points
        self.encode_features()

        ### Calculate clusters
        self.compute_clusters()

        ### Compute similarity values for each candidate
        if method_variant != 'rnd_uncertainty':
            self.compute_similarity()

            if method_variant == 'max_uncertainty':
                self.similarity_array = -1 * self.similarity_array

        ### Choose data from clusters

        # create zero array that is filled with cluster IDs for this batch
        batch_index_list = []
        inf_score_list = []

        # iterates over the batch_index_array up to cand_batch_size
        cluster_batch_counter = 0

        # iterates over clusters until n_clusters, then resets to 0
        # if cluster_batch_counter does not reached cand_batch_size
        cluster_index = 0

        # iterate over all clusters until cluster_batch_counter reaches
        # cand_batch_size
        while cluster_batch_counter < self.adl_batch_size:

            # get an array of indices matching to currently iterated cluster
            # ID
            index_array = np.where(self.cluster_labels == cluster_index)[0]

            # if the set is not empty
            if len(index_array) != 0:

                if method == 'rnd_uncertainty':
                    # choose one element at random from this index array
                    index_choice = np.random.choice(index_array)

                else:
                    # get similarity values for matching index array
                    similarity_array = self.similarity_array[index_array]

                    if method == 'avg_ucertainty':
                        # turn into absolute difference to average similarity
                        similarity_array = abs(
                            similarity_array - np.mean(similarity_array)
                        )

                    # calculate largest similarity
                    max_similarity = similarity_array.max()

                    # choose first/largest value from similarity_array
                    index_choice = index_array[
                        np.where(
                            similarity_array == max_similarity
                        )[0][0]
                    ]

                    # add information content score of data point to inf_score_list
                    inf_score_list.append(max_similarity)

                # add chosen index to zero array
                batch_index_list.append(self.index_array[index_choice])

                # setting the cluster ID to -1 excludes data point from
                # considerations in next iterations of this loop
                self.cluster_labels[index_choice] = -1

                # increment the counter for already added data points to
                # zero array
                cluster_batch_counter += 1

            # increment the cluster ID index for the next iteration
            cluster_index += 1

            # set cluster ID index to zero for next iteration if an entire
            # round of iterations did not fill zero array
            if cluster_index >= self.n_clusters:
                cluster_index = 0

        # Save results
        self.batch_index_list = batch_index_list
        self.inf_score_list = inf_score_list

    def predict(
        self,
        y_pred=None,
        x_t_pred=None,
        x_s_pred=None,
        x_st_pred=None,
        silent=True,
        plot=False,
    ):

        """ Predict labels for unqueried candidate data.
        """

        def plot_true_vs_prediction(test_data_Y, predictions):

            """ Visualizes predictions vs. true values. """

            plot_rows = 3
            plot_clos = 3

            # create a matplotlib.pyplot.subplots figure
            fig, ax = plt.subplots(plot_rows, plot_clos, figsize=(16, 16))

            # set the figtitle
            fig.suptitle('True vs. predicted', fontsize=16)

            # pick at random a set of integers for visualization
            data_indices = np.random.randint(
                0,
                len(test_data_Y),
                plot_rows * plot_clos
            )

            # create a variable for iteratively adding number of subplots
            subplot_counter = 0

            # iterate over number of rows
            for i in range(plot_rows):

                # iterate over number of columns
                for j in range(plot_clos):

                    # choose currently iterated random index
                    data_index = data_indices[subplot_counter]

                    # plot the true values
                    plot1 = ax[i, j].plot(test_data_Y[data_index])

                    # plot the predicted values
                    plot2 = ax[i, j].plot(predictions[data_index])

                    # increment the subplot_counter
                    subplot_counter += 1

            # add a figure legend
            fig.legend(
                [plot1, plot2],
                labels=['true', 'predicted'],
                fontsize=16,
            )

            # save figure
            file_name = 'true_vs_pred_cand.pdf'
            saving_path = self.path_to_results + file_name
            fig.savefig(saving_path)

        ### Reset the state of the test loss metric
        loss_function = tf.keras.losses.mean_squared_error
        mean_loss = tf.keras.metrics.Mean(name='mean_loss_train_test')
        mean_loss.reset_states()

        ### Make predictions
        model = self.models['f_nn']
        model_input_list = []
        if self.x_t is not None:
            model_input_list.append(x_t_pred)
        if self.x_s is not None:
            model_input_list.append(x_s_pred)
        if self.x_st is not None:
            model_input_list.append(x_st_pred)
        predictions = model.predict(model_input_list)

        self.silent = silent
        self.plot = plot

        ###
        # Calculate the testing loss ###
        ###

        if y_pred is not None:
            # calculate the testing losses
            t_loss = loss_function(y_pred, predictions)

            # take the mean of single losses
            testing_loss = mean_loss(t_loss).numpy()

            # tell us how much testing loss we have
            if not self.silent:

                print('loss:', testing_loss)

            # Plot exemplar predictions
            if self.plot:
                plot_true_vs_prediction(y_pred, predictions)

        else:
            testing_loss = None

        self.predictions = predictions
        self.testing_loss = testing_loss


import numpy as np
import cv2

def read_image(img_path: str) -> np.array:
    """
    Читает картинку в np.array
    :param img_path: путь к картинке
    :return:
    """
    # cv2 читает в BGR
    image = cv2.imread(str(img_path))
    return image

def resize_image(img: np.array, max_dim_size: int) -> np.array:
    """
    Сжимает пропорционально картинку, чтобы большая сторона стала равна max_dim_size
    :param img: картинка в np.array
    :param max_dim_size: int
    :return:
    """
    image = np.array(img.copy())
    height, width, colour_layers = image.shape[:3]

    if width > height:
        scale_percent = max_dim_size / width
        new_width = max_dim_size
        new_height = int(height * scale_percent)
    if width <= height:
        scale_percent = max_dim_size / height
        new_width = int(width * scale_percent)
        new_height = max_dim_size

    return cv2.resize(image, (new_width, new_height))

def pad_image(img: np.array, background_color: tuple[int]) -> np.array:
    """
    Относительно центра картинки добавляет поля
    в высоту или ширину делая кадр квадратным
    """
    image = np.array(img.copy())
    height, width, colour_layers = image.shape[:3]

    max_dimm = max(height, width)
    new_height = int(max_dimm)
    new_width = int(max_dimm)

    # пустой кадр, заполненный фоном, куда потом вставим изображение
    empty_frame = np.full((new_height,
                           new_width,
                           colour_layers),
                          background_color,
                          dtype='uint8')

    pad_height = int((new_height - height) / 2)
    pad_width = int((new_width - width) / 2)

    # ищем координаты для вставки
    # в numpy нумерация с верхнего левого угла...
    # потому если мы добавляем сверху, координаты не меняются.
    if pad_height >= 0:
        y0 = 0 + pad_height
        y1 = y0 + height
    if pad_height < 0:
        y0 = 0
        y1 = y0 + height

    if pad_width >= 0:
        x0 = 0 + pad_width
        x1 = x0 + width
    if pad_width < 0:
        x0 = 0
        x1 = x0 + width

    # вставили картинку
    empty_frame[y0:y1, x0:x1] = image

    result = empty_frame

    return result

def process_image(img_path: str) -> np.array:
    """
    Пайплайн обработки зиображения
    :param img_path: путь к картинке
    :return: np.array
    """
    img = read_image(img_path)  # прочитали
    img = resize_image(img, 640)  # сжали/растянули
    img = pad_image(img, (114, 114, 114))  # сделали падинг фоном до квадрата
    img = img.astype(np.float32)  # привели к float32
    # img = cv2.normalize(img, None, 0, 1, cv2.NORM_MINMAX)  # предположение, цвета при обучении были нормализованы
    img = np.transpose(img, (2, 0, 1))  # приводим к виду [idx, colour, h, w]
    img = np.expand_dims(img, axis=0)
    return img


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
from collections import deque, defaultdict

class Solution:
    def bfs(self,root):
        if not root:
            return None

        visited =  defaultdict(bool)
        dq = deque()
        level = 0

        visited[root] = True
        dq.append([root,level])
        h_map = defaultdict(list)

        while dq:

            temp, level  = dq.popleft()
            h_map[level].append(temp.val)

            print(temp.val)

            if temp.left is not None:
                dq.append([temp.left,level+1])

            if temp.right is not None:
                dq.append([temp.right,level+1])

        return (h_map)

    def largestValues(self, root: Optional[TreeNode]) -> List[int]:
        h_map = self.bfs(root)
        answer = []

        if h_map is None:
            return []

        for v in h_map.values():
            if len(v)<1:
                return []

            answer.append(max(v))

        return answer


from bs4 import BeautifulSoup
import requests
# test Kate

# artem branch
# dima

url = 'http://mignews.com/mobile'

page = requests.get(url)

print(page.status_code)

filteredNews = []
allNews = []

soup = BeautifulSoup(page.text, "html.parser")

#print(soup)
print('aaaaaaaaaaa')

allNews = soup.findAll('a', class_='lenta')

for data in allNews:
    if data.find('span', class_='time2 time3') is not None:
        filteredNews.append(data.text)

for data in filteredNews:
    print(data)

print('aaaaaaaaaaa')

##########

url = 'https://kk.polyanaski.ru/winter/'
page = requests.get(url)
print(page.status_code)
filteredNews = []
filteredAtr = []
filteredMap = {}
allNews = []
soup = BeautifulSoup(page.text, "html.parser")
#print(soup)
print('aaaaaaaaaaa')

allNews = soup.findAll('div', class_='left_panel_item summer')
print(allNews)

print('aaaaaaaaaaa')
print('aaaaaaaaaabbbbb')
for data in allNews:
    if data.find('div', class_='left_panel_item_road') is not None:
        filteredNews.append(data.text)
        filteredAtr.append(data)
        print('---')
        print('---')
        print(data.attrs)
        #print(data.name)
        #print(data.class)
        #print(data.data-disabled-id)

        soup2 = BeautifulSoup(data, "html.parser")

        allNews2 = soup2.findAll('div', class_='left_panel_item_road')
        for n in allNews2:
            print(n.attrs)
            filteredMap[n.attrs['data-disabled-id']] = 1 if 'disabled' not in n.attrs['class'] else 0

        print('+++')

print('aaaaaaaaaaa')
print(filteredMap)
print('aaaaaaaaaaa')

for data in filteredNews:
    print(data)

print('aaaaaaaaaaa')
print(filteredAtr)

import json
data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'
order_data_dict = json.loads(data_string)


from rest_framework.routers import SimpleRouter
from core.authentication.viewsets import (
    LoginViewSet,
    RegistrationViewSet,
    RefreshViewSet,
    RefreshDataViewSet
)
from .exercises.viewsets import ExercisesViewSet, TopicViewSet, SubTopicViewSet
from .evaluation.viewsets import EvaluationViewSet

routes = SimpleRouter()

# AUTHENTICATION
routes.register(r'auth/login', LoginViewSet, basename='auth-login')
routes.register(r'auth/register', RegistrationViewSet, basename='auth-register')
routes.register(r'auth/refresh', RefreshViewSet, basename='auth-refresh')
routes.register(r'auth/refresh_data', RefreshDataViewSet, basename='auth-refresh-data')

# EXERCISES
routes.register(r'exercises', ExercisesViewSet, basename='exercises')

# TOPICS
routes.register(r'topic', TopicViewSet, basename='topics')
routes.register(r'subtopic', SubTopicViewSet, basename='subtopics')

# EVALUATION
routes.register(r'evaluation', EvaluationViewSet, basename='evaluation')

urlpatterns = [
    *routes.urls,
]