import matplotlib
matplotlib.use('TkAgg')
import numpy as np
import matplotlib.pyplot as plt
import sys,os
sys.path.append('/home/msdos/focalplane/plate_control/trunk/petal')
import petalcomm
import petal
import pandas as pd
import datetime
import time, datetime
import pickle
import csv
from scipy import interpolate
import json
import argparse
import pickle as pk
import gspread
import pandas as pd
from gspread_dataframe import get_as_dataframe
from oauth2client.service_account import ServiceAccountCredentials
# nominal hole location data

plt.ion()

parser = argparse.ArgumentParser()
parser.add_argument("-w", "--wait", help="Determines amount of time between measuring temperature",action='store',dest='wait',default=None)
parser.add_argument('petal',type=int,help="Petal Number")
parser.add_argument('pc',type=int,help="Petal Controller Number")
results = parser.parse_args()

class PlotPetalBoxTemps():
    def __init__(self):
        self.start_time = datetime.datetime.now()
        self.file_path = '/home/msdos/focalplane/pos_utility/'
        self.temp_log_path = os.getcwd()#'/home/msdos/test_util/temp_logs/'

        self.wait = int(results.wait)
        if self.wait is None:
            self.wait = 60

        print("Temperature will be read every %d seconds"%self.wait)

        self.temp_log = open(self.temp_log_path+'/temp_log_%s.txt'%str(self.start_time),'w')
        self.hole_coords = np.genfromtxt(self.file_path+'hole_coords.csv', delimiter = ',', usecols = (3,4), skip_header = 40)
        self.nons = [38, 331, 438, 460, 478, 479, 480, 481, 497, 498, 499, 500, 513, 514, 515, 516, 527, 528, 529, 530, 531, 535, 536, 537, 538, 539, 540]
        self.gifs = [541, 542]
        self.fifs = [11, 75, 150, 239, 321, 439, 482, 496, 517, 534]
        try:
            self.petal = str(results.petal).zfill(2)
            self.comm = petalcomm.PetalComm(int(str(results.pc)))
        except:
            print('Must be run with two arguments: petal # and petalbox #')
            sys.exit()

        self.init_temps=[-40, -35, -30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100, 105, 110, 115, 120, 125]
        self.adc=[3972.291151, 3927.370935, 3869.451327, 3795.92674, 3704.600096, 3593.238247, 3460.369957,
                  3305.419031, 3129.278263, 2933.684449, 2721.964441, 2500.36307, 2273.125056, 2048, 1829.510403,
                  1621.972457, 1428.638708, 1251.753073, 1091.966263, 949.5873406, 823.911807, 713.9506234, 618.034474,
                  535.4993046, 464.1503813, 402.5825068, 348.8588418, 302.7049454, 263.3419356, 229.9160146, 201.3532376,
                  176.832823, 155.6819655, 137.3646468]

        self.f2 = interpolate.interp1d(self.adc, self.init_temps, kind = 'cubic')

    def get_temps(self):
        measure_time = datetime.datetime.now()
        temp_dict = self.comm.pbget('posfid_temps')
        print(measure_time)
        ids=[]
        temps=[]
        for id,val in temp_dict.items():
            _ids=list(val.keys())
            _temps=list(self.f2([t for t in val.values()]))
            ids=ids+_ids
            temps=temps+_temps
        self.ids = ids
        self.temps=temps
        output={}
        for i in range(len(ids)):
            output[ids[i]]=temps[i]
        file_output='data/thermal_test_result_'+datetime.datetime.now()+'.json'
        with open(file_output, 'w') as f:
            json.dump(output, f)
        self.temp_log.write(str(measure_time)+'\n')
        self.temp_log.write(str(output))
        self.temp_log.write('\n')

    def plot_hole_info(self):
        for i in range(len(self.hole_coords)):
            x = self.hole_coords[i][0]
            y = self.hole_coords[i][1]

            if i not in self.nons:
                self.ax.plot(x, y, color = 'lightgrey', marker='o', zorder = -1)
                if i in self.fifs:
                    text = 'F' + str(i)
                    col = 'blue'
                elif i in self.gifs:
                    text = 'G' + str(i)
                    col = 'purple'
                else:
                    text = i
                    col = 'black'
                self.ax.text(x-.1, y + 0.3, text, color = col, fontsize=6)

    def initial_plot(self):

        self.fig=plt.figure(figsize = (10,5))
        self.ax=self.fig.add_axes(self.fig.add_axes([0,0,1,1]))
        self.plot_hole_info()
        self.hole=1
        self.ax.scatter(self.hole_coords[self.hole][0], self.hole_coords[self.hole][1], marker= '*', s=200, color = 'gold')

        # Google sheet
        url = 'https://docs.google.com/spreadsheets/d/1lJ9GjhUUsK2SIvXeerpGW7664OFKQWAlPqpgxgevvl8/edit#gid=0'
        credentials = self.file_path+'google_access_account.json'
        scope = ['https://spreadsheets.google.com/feeds', 'https://spreadsheets.google.com/feeds', 'https://www.googleapis.com/auth/drive']
        creds = ServiceAccountCredentials.from_json_keyfile_name(credentials, scope)
        client = gspread.authorize(creds)
        sheet = client.open_by_url(url).sheet1
        df  = get_as_dataframe(sheet, parse_dates=True, usecols=[0,1,2,3,4,6,8,14], skiprows=20, header=0)
        pdf=df
        pdf=pdf.loc[pdf['PETAL_ID'] == int(self.petal)]

        self.dev_list=pdf['CAN_ID'].tolist()
        #self.dev_list = [str(i) for i in dev_list]
        self.hole_list=pdf['DEVICE_LOC'].tolist()
        self.dev_id_loc=dict(zip(self.dev_list,self.hole_list))

        holes = []
        idx = []
        for i,e in enumerate(self.ids):
            try:
                holes.append(self.dev_id_loc[e])
                idx.append(i)
            except:
                print('failed: ',e)
                self.nons.append(e)
                pass
        temps = np.array(self.temps)[np.array(idx)]
        new_ids = np.array(self.ids)[np.array(idx)]

        #print(holes[0:5])
        x = []
        y = []
        idx2 = []
        for i,e in enumerate(holes):
            try:
                x.append(self.hole_coords[int(e)][0])
                y.append(self.hole_coords[int(e)][1])
                idx2.append(i)
            except:
                print('failed: ',new_ids[i])
                self.nons.append(new_ids[i])
                pass
        temps = temps[np.array(idx2)]
        x = np.array(x)
        y = np.array(y)

        sc = self.ax.scatter(x,y,s=120,c=temps)
        self.cbar = plt.colorbar(sc)
        plt.show()

    def updated_plot(self):
        #Update data (with the new _and_ the old points)
        self.ax.scatter(self.hole_coords[self.hole][0], self.hole_coords[self.hole][1], marker= '*', s=200, color = 'gold')
        holes = []
        idx = []
        for i,e in enumerate(self.ids):
            try:
                holes.append(self.dev_id_loc[e])
                idx.append(i)
            except:
                print('failed: ',e)
                pass
        temps = np.array(self.temps)[np.array(idx)]
        new_ids = np.array(self.ids)[np.array(idx)]

        #print(holes[0:5])
        x = []
        y = []
        idx2 = []
        for i,e in enumerate(holes):
            try:
                x.append(self.hole_coords[int(e)][0])
                y.append(self.hole_coords[int(e)][1])
                idx2.append(i)
            except:
                print('failed: ',new_ids[i])
                pass
        temps = temps[np.array(idx2)]
        x = np.array(x)
        y = np.array(y)

        self.ax.scatter(x,y,s=120,c=temps)
        self.cbar.set_clim(vmin=min(temps),vmax=max(temps))
        self.cbar.draw_all()
        #Need both of these in order to rescale
        self.ax.relim()
        self.ax.autoscale_view()
        #We need to draw *and* flush
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def move(self):
        mode='cruise'
        direction_theta='cw'
        direction_phi='cw'
        degree=27
        canlist=['can10','can11','can12','can13','can14','can15','can16','can17','can22','can23']
        self.comm.pbset('STOP_MODE','off')
        for can in canlist:
            self.comm.move(can, 20000, direction_theta, mode, 'theta', degree)
        time.sleep(10)
        for can in canlist:
            self.comm.move(can, 20000, direction_phi, mode, 'phi', degree)
        time.sleep(10)
        self.get_temps()
        self.comm.pbset('STOP_MODE','on')

    def __call__(self):
        self.get_temps()
        self.move()
        self.initial_plot()
        time.sleep(self.wait)
        while True:
           try:
               self.get_temps()
               self.move()
               self.updated_plot()
               time.sleep(self.wait)
           except KeyboardInterrupt:
               print('interrupted!')

if __name__ == '__main__':
    P = PlotPetalBoxTemps()
    P()


# stdlib
import logging
import pickle
import os
import time

# third party lib
import h5py
import numpy as np
from scipy.stats import chi2
from numpy.linalg import _umath_linalg
from flask import current_app as app

# internal lib
from lib.settings import Settings
from lib.client_registry import Registry
from lib.utils import write_or_replace
import lib.networking as networking
from server.lib.plots import manhattan_plot

storePath = os.path.join(Settings.local_scratch, "central.h5py")
store = h5py.File(storePath, "a")
clients = Registry.get_instance().list_clients()
chi2sf = chi2.sf
lstsq = _umath_linalg.lstsq_m
dot = np.dot

class LogisticAdmm(object):
    __instance = None

    def __init__(self, npcs, active, max_iters=50, algo="newton"):
        if LogisticAdmm.__instance is not None:
            return
        else:
            self.time = time.time()
            self.estimates = {}
            self.iters, self.accumulant = {}, {}
            self.chroms = [v for v in store.keys() if v != 'meta']
            self.active_chroms = ["Small"]
            self.finished = {}
            self.likelihood = {}
            self.max_iters = max_iters * 5
            self.normalization_stats = None
            self.base_likelihood = 0
            self.nconnections = len(clients)
            LogisticAdmm.__instance = self
            self.send_request({}, "initialize")
            self.beta = 0
            self.algo = algo  # "newton" for admm
            self.Hess, self.Gradients, self.Diags, self.Vals = {}, {}, {}, {}
            self.converged, self.linesearch_convergence = {}, {}
            self.threshold = .005
            self.EPS = 1e-4
            self.linesearch_iter, self.fchanges = {}, {}
            self.scratch_likelihoods = {}
            self.alpha, self.BETA = .1, .5

            self.covars = {}

    def activate_chrom(self, chrom):
        msg = {"Estimated": "Small"}
        self.send_request(msg, "estimate")

    def make_chrom_active(self, chrom):
        self.finished[chrom] = False
        msg = {"Estimated": chrom}
        if self.algo != "admm":
            L = store[f"{chrom}/positions"].shape[0]
            starter = self.estimates["Small"]
            warm_start = np.empty((starter.shape[0]+1, 1))
            warm_start[1] = 0
            warm_start[0] = starter[0]
            warm_start[2:] = starter[1:]
            warm_start = np.tile(warm_start, (L, 1, 1))
            msg["VALS"] = warm_start
            self.estimates[chrom] = warm_start
        self.send_request(msg, "estimate")

    @staticmethod
    def get_instance(args, active):
        if LogisticAdmm.__instance is None:
            LogisticAdmm(args["ASSO_PCS"], active)
        return LogisticAdmm.__instance

    def update_stats(self, data):
        # #TODO this should just be its own object but for now it's faster to just add it here
        logging.info(f"Constructing feature's matrix!")
        data = pickle.loads(data)
        if self.normalization_stats is None:
            self.normalization_stats = data.copy()
            self.normalization_stats["iter"] = self.nconnections - 1
        else:
            self.normalization_stats["Sums"] += data["Sums"]
            self.normalization_stats["SS"] += data["SS"]
            self.normalization_stats["N"] += data["N"]
            self.normalization_stats["iter"] -= 1
            if self.normalization_stats["Indx"] != data["Indx"]:
                raise NameError("""Index of quantitative variables
                    does not match""")
            # self.normalization_stats["data"] = np.vstack((self.normalization_stats["data"], data["data"]))
            # self.normalization_stats["ys"] = np.vstack((self.normalization_stats["ys"], data["ys"]))
            if self.normalization_stats["iter"] == 0:
                # cov = self.normalization_stats["data"]
                # Ys = self.normalization_stats["ys"]
                # ind = self.normalization_stats["Indx"]

                n = float(self.normalization_stats["N"])

                mu = self.normalization_stats["Sums"]/n
                # pdb.set_trace()

                # cov[:,ind] -= mu
                # cov[:, ind] /= np.sqrt(self.normalization_stats["SS"]/n - mu**2)
                # ind = [i for i in range(cov.shape[1]) if i != 1]
                # model = lr(C=1e10, fit_intercept = False, tol=1e-7)
                # model.fit(cov[:,ind], Ys)
                logging.info(f"Total of {n} samples")
                msg = {"Indx": self.normalization_stats["Indx"], "Means": mu,
                       "SD": np.sqrt(self.normalization_stats["SS"]/n - mu**2)}
                self.send_request(msg, "adjust")

    def send_request(self, data, subtask):
        to_send = pickle.dumps(data)
        networking.message_clients(f"asso/{subtask}", data=to_send, env=app.config["ENV"])

    def update(self, message):
        message = pickle.loads(message)
        z_hat = message["VALS"]
        model = message["Estimated"]
        self.update_estimate(z_hat, model, message)

    def association_finished(self):
        return self.finished

    def set_clients_state(self, state):
        for client in clients:
            Registry.get_instance().set_client_state(client['name'], state)

    def update_estimate(self, z_hat, model, data):
        if model in self.estimates:
            if self.iters[model] >= self.max_iters:  # this shouldn't happen but it does! WHy?
                logging.info(f"WHYYYYY {model}, {self.iters[model]}")
                return
            prev = self.estimates[model]
            # self.normalization_stats["data"] = np.vstack((self.normalization_stats["data"], data["cov"]))
            if prev[1] == 1:
                beta = (prev[0] + z_hat)/(self.nconnections)
                self.iters[model] += 1
                if not self.iters[model] % 10:
                    logging.info(f"Finished iteration {self.iters[model]} on {model}")
                    if model == "Small":
                        logging.info(f"{np.sum(np.abs(beta - self.beta))}")
                    else:
                        logging.info(f"{np.linalg.norm(beta[7,:] - self.beta)}")
                if model == "Small":
                    self.beta = beta
                else:
                    self.beta = beta[7, :]

                if self.iters[model] == self.max_iters:  # #TODO this shouldn't happen but why does it?
                    write_or_replace(store, f"meta/{model}/coef", beta)
                    if model != "Small":
                        del self.estimates[model]
                    self.active_chroms.remove(model)
                    if model == "Small":
                        self.max_iters = 15
                        self.estimates["Small"] = beta
                    # chroms = [key for key in store if key != 'meta']
                    # if len(self.finished) == len(chroms) + 1:
                    self.finished[model] = True
                    if not self.chroms:
                        self.set_clients_state("ASSO_DONE")
                        logging.info(f"We are done with association!")
                        self.initialize_pval_computation()
                    else:
                        chrom = self.chroms.pop()
                        self.active_chroms.append(chrom)
                        self.make_chrom_active(chrom)
                    return
                else:
                    self.estimates[model] = [np.zeros_like(beta), self.nconnections]
                    msg = {"Estimated": model, "VALS": beta, "Iter": self.iters[model]}
                    self.send_request(msg, "estimate")
            else:
                self.estimates[model] = [prev[0] + z_hat, prev[1] - 1]

        else:  # Not in dictionary yet
            # self.normalization_stats["data"] = data["cov"]

            self.estimates[model] = [z_hat, self.nconnections - 1]
            self.iters[model] = 1

    def newton_stats_update(self, msg, unconverged=None):
        data = pickle.loads(msg)
        model = data["Estimated"]
        if model in self.accumulant:
            # if self.iters[model] >= self.max_iters: # this shouldn't happen but it does! WHy?
            #    logging.info(f"WHYYYYY {model}, {self.iters[model]}")
            #    return
            self.Gradients[model] += data['g']
            self.Diags[model] += data['d']
            if model in self.iters:
                temp = self.Vals[model]
                temp[np.logical_not(self.converged[model])[:, 0]] += data['v']
                self.Vals[model] = temp
            else:
                self.Vals[model] += data['v']
            self.Hess[model] += data['H']
            self.covars[model] = np.vstack((self.covars[model], data['covar']))
            self.accumulant[model] -= 1
        else:
            if model in self.Vals:
                temp = self.Vals[model]
                temp[np.logical_not(self.converged[model])[:, 0]] = data['v']
                self.Vals[model] = temp
            else:
                self.Vals[model] = data['v']
            self.Gradients[model] = data['g']
            self.Diags[model] = data['d']
            self.Hess[model] = data['H']
            self.covars[model] = data['covar']
            self.accumulant[model] = self.nconnections - 1
        finished_collecting = False
        if self.accumulant[model] == 0:
            finished_collecting = True
            del self.accumulant[model]
        return model, finished_collecting

    def newton_iter(self, model):
        TOL = 1e-7
        self.t = 1
        hs = self.Hess[model]
        gs = self.Gradients[model]
        ds = self.Diags[model]
        dfxs = np.zeros((gs.shape[0], 1))
        af = store[f"{model}/allele_freq"].value
        x0 = self.estimates[model]
        if model not in self.converged:
            self.iters[model] = 1
            convergence_status = np.zeros((gs.shape[0], 1), dtype=bool)
            self.converged[model] = convergence_status.copy()
            old_not_converged = None
        else:
            self.iters[model] += 1
            old_not_converged = np.logical_not(self.converged[model])
            af = af[old_not_converged[:, 0]]
            x0 = x0[old_not_converged[:, 0]]
            convergence_status = np.zeros((x0.shape[0], 1), dtype=bool)
        logging.info(f"Now updating coefficient estimates for chr{model} iteration: {self.iters[model]}")
        for i, g in enumerate(gs):
            if af[i] < self.threshold or 1-af[i] < self.threshold:
                convergence_status[i] = True
                continue
            h = np.diagflat(ds[i])
            if i % 2:
                hessian = np.triu(hs[i//2])
            else:
                hessian = np.tril(hs[i//2])
            hessian += hessian.T
            hessian += h
            dx = lstsq(-hessian, g[:, np.newaxis], 1e-5, signature='ddd->ddid')[0]
            dfx = dot(g.T, dx)
            if np.abs(dfx) < TOL:
                x0[i, :] += dx
                convergence_status[i] = True
                continue
            gs[i:i+1, :] = x0[i].T + self.t*dx.T
            dfxs[i, :] = dfx
            ds[i] = dx[:, 0]
        if old_not_converged is not None:
            self.estimates[model][old_not_converged[:, 0]] = x0
        else:
            self.estimates[model] = x0
        if np.prod(convergence_status) or self.iters[model] == self.max_iters:
            self.finished[model] = True
            write_or_replace(store, f"meta/{model}/newton_coef", self.estimates[model])
            af = store[f"{model}/allele_freq"].value
            arr = np.logical_not(np.logical_or(af < self.threshold, 1-af < self.threshold))
            est = self.estimates[model][arr]
            msg = {"Estimated": model, "conv": np.expand_dims(arr, axis=1), "x0": est[:, :, 0]}
            self.send_request(msg, "query")
            del self.estimates[model]
            self.active_chroms.remove(model)
            if not self.chroms:
                self.set_clients_state("ASSO_DONE")
                logging.info("We are done with association.")
                logging.info(f"Computing coefficients took roughly {time.time() - self.time:.1f} seconds.")
                self.time = time.time()
                # COMPUTE PVALS
            else:
                chrom = self.chroms.pop()
                self.active_chroms.append(chrom)
                self.make_chrom_active(chrom)
            return
        else:
            if self.iters[model] > 1:
                temp = self.converged[model]
                temp[np.logical_not(temp)] = convergence_status[:, 0]
                self.converged[model] = temp
            else:
                self.converged[model] = convergence_status
            gs = gs[np.logical_not(convergence_status)[:, 0], :]
            msg = {"Estimated": model, "x0": gs, "conv": np.logical_not(self.converged[model])}
            self.send_request(msg, "query")
        self.linesearch_convergence[model] = convergence_status.copy()
        del self.Hess[model], self.Gradients[model]
        self.Diags[model] = ds[np.logical_not(convergence_status)[:, 0], :]  # repurposed
        # self.fchanges[model] = dfxs[np.logical_not(convergence_status)]
        self.fchanges[model] = dfxs

    def collect_likelihoods(self, data):
        data = pickle.loads(data)
        model = data["estimated"]
        if model in self.scratch_likelihoods:
            self.scratch_likelihoods[model] += data['v']
            if self.linesearch_iter[model] == 1:
                if self.finished[model]:
                    write_or_replace(store, f"meta/{model}/newton_ell",
                                     self.scratch_likelihoods[model])
                    write_or_replace(store, f"meta/{model}/newton_pval",
                                     chi2sf(-2*self.scratch_likelihoods[model], 1))
                    del self.scratch_likelihoods[model]
                    del self.Hess[model], self.Gradients[model]
                    del self.Diags[model], self.fchanges[model]
                    if all(value for value in self.finished.values()):
                        manhattan_plot(storePath, "manhattan_plot.png")
                        logging.info("P-value computation is Done!")
                        logging.info(f"Computing P-values took roughly {time.time() - self.time:.1f} seconds.")
                        store.close()
                else:
                    self.newton_test_new_point(model)
        else:
            self.scratch_likelihoods[model] = data['v']
            self.linesearch_iter[model] = self.nconnections
        self.linesearch_iter[model] -= 1

    def newton_test_new_point(self, model):
        vals = self.scratch_likelihoods[model].copy()
        del self.scratch_likelihoods[model]
        notConverged = np.logical_not(self.linesearch_convergence[model]).copy()
        notOverallConverged = np.logical_not(self.converged[model])
        old_vals = self.Vals[model][notOverallConverged]
        fs = old_vals + self.t * self.alpha * self.fchanges[model][notConverged]
        # fs = old_vals[notConverged[:,0]]
        # fs += self.t * self.alpha * self.fchanges[model][notConverged[notOverallConverged]]
        # fs = self.Vals[model][notConverged] + self.t * self.alpha * self.fchanges[model]
        # [notOverallConverged[notConverged]]
        # fs = self.Vals[model][notConverged]
        # fs += self.t * self.alpha * self.fchanges[model][notConverged[np.logical_not(self.converged[model])]]
        toUpdate = vals < fs[:, np.newaxis]
        nowConverged = notConverged.copy()
        nowConverged[notConverged] = np.logical_and(notOverallConverged[notOverallConverged], toUpdate[:, 0])
        temp = self.estimates[model][np.logical_not(self.converged[model][:, 0])]
        temp[toUpdate[:, 0], :, 0] += self.t * self.Diags[model][toUpdate[:, 0]]
        self.estimates[model][np.logical_not(self.converged[model][:, 0])] = temp
        # self.estimates[model][nowConverged[:,0],:,0] += self.t * self.Diags[model][toUpdate[:,0]]
        # self.estimates[model][np.logical_not(
        #     self.linesearch_convergence[model])[:,0],:][accept[:,0],:][:,:,0] += (self.t * self.alpha
        #                                                            * self.Diags[model][toUpdate[:,0]])
        # Update function value
        likelihoods = self.Vals[model][np.logical_not(self.converged[model][:, 0])]
        likelihoods[toUpdate[:, 0]] = vals[toUpdate[:, 0]]
        self.Vals[model][np.logical_not(self.converged[model][:, 0])] = likelihoods
        notConverged[notConverged] = np.logical_not(toUpdate[:, 0])
        if not np.prod(toUpdate):  # not all approved. Start another line search
            logging.info(f"""chrom: {model} | t: {self.t} | \
percent converged line search: {(1-np.mean(notConverged))*100}%""")
            self.t *= self.BETA
            x0s = self.estimates[model][notConverged[:, 0]][:, :, 0] + (
                self.t * self.Diags[model][np.logical_not(toUpdate)[:, 0]])
            msg = {"Estimated": model, "x0": x0s, "conv": notConverged}
            self.send_request(msg, "query")
            temp = self.Diags[model]
            self.Diags[model] = temp[np.logical_not(toUpdate[:, 0]), :]
            self.linesearch_convergence[model] = np.logical_not(notConverged)
        else:  # All approved, start a new iteration of main algo
            msg = {"Estimated": model, "unconv": np.logical_not(self.converged[model]),
                   'VALS': self.estimates[model][np.logical_not(self.converged[model])[:, 0], :]}
            self.send_request(msg, "estimate")
            del self.Diags[model]

    def initialize_pval_computation(self):
        self.chroms = set([v for v in store.keys() if v != 'meta'])
        self.send_coef("Small", {})

    def send_coef(self, model, msg):
        coef = store[f"meta/{model}/coef"].value
        msg["Estimated"] = model
        msg["Coef"] = coef
        self.send_request(msg, "coef")

    def update_pval(self, message):
        def _share_new_ll():
            print(self.chroms)
            if len(self.chroms) >= 1:
                new_model = self.chroms.pop()
                self.send_coef(new_model, {})
            else:
                logging.info("P-value computation is Done!")
                logging.info(f"Computing P-values took roughly {time.time() - self.time:.1f} seconds.")

        message = pickle.loads(message)
        model = message["Estimated"]
        if model == "Small":
            _share_new_ll()
            return
        val = message["estimate"]
        if model in self.likelihood:
            prev = self.likelihood[model]
            if prev[1] == 1:
                ell = prev[0] + val
                pval = chi2sf(-2*ell, 1)
                write_or_replace(store, f"meta/{model}/ell", ell)
                write_or_replace(store, f"meta/{model}/pval", pval)
                del self.likelihood[model]
                _share_new_ll()
            else:
                self.likelihood[model] = [prev[0] + val, prev[1] - 1]
        else:
            self.likelihood[model] = [val, self.nconnections - 1]


import numpy as np
import os

MAX_FEATURES = 800
GOOD_MATCH_PERCENT = 0.05
FRONT_MASK_COLOR_WINDOW = 25
CHECKBOX_SIZE = 23
CHECKBOX_THRESH = 300
Q_COLOR_WINDOW = 15

PROJ_DATA_DIR = '/home/cddl/scans/'
SCAN_DIR = PROJ_DATA_DIR + 'raw/'

SERVER_URL = '/static/cards/exhibit/'
SERVER_DATA_DIR = '/home/ehuntley/scans/'
SERVER_RAW = SERVER_DATA_DIR + 'raw/'
SERVER_PROCESSED = '/home/ehuntley/www/nextstop/static/cards/exhibit/'
SERVER_TEMPLATES = SERVER_DATA_DIR + 'templates/'

# x, y
QUESTION_HUES = [
    # 4: My preferred transport mode(s) in 2020 will be...
    (4, np.array([1, 170, 180])),
    # 5: In 2040, the average person will...
    (5, np.array([176, 130, 83])),
    # Responsibility for autonomous vehicle accidents belongs to...
    (6, np.array([24, 130, 230])),
    # In 2040, commuting will take...
    (7, np.array([62, 100, 159])),
    # In 2040, everyone will have access to...
    (8, np.array([101, 118, 68])),
    # The future of mobility will make the world
    (9, np.array([142, 87, 93])),
    # In the future, my transportation costs will...
    (10, np.array([76, 126, 106])),
    # Future mobility options will have the greatest imapact on...
    (12, np.array([171, 146, 194])),
    # In 2040...
    (13, np.array([107, 100, 130])),
    # Travel in the future will be more dangerous for...
    (14, np.array([6, 196, 232])),
]

OLD_QUESTION_HUES = [
    # 4: My preferred transport mode(s) in 2020 will be...
    (4, np.array([1, 232, 190])),
    # 5: In 2040, the average person will...
    (5, np.array([0, 255, 62])),
    # Responsibility for autonomous vehicle accidents belongs to...
    (6, np.array([27, 176, 242])),
    # In 2040, commuting will take...
    (7, np.array([60, 135, 150])),
    # In 2040, everyone will have access to...
    (8, np.array([107, 255, 42])),
    # The future of mobility will make the world
    (9, np.array([142, 192, 70])),
    # In the future, my transportation costs will...
    (10, np.array([88, 255, 62])),
    # Future mobility options will have the greatest imapact on...
    (12, np.array([174, 156, 208])),
    # In 2040...
    (13, np.array([106, 205, 120])),
    # Travel in the future will be more dangerous for...
    (14, np.array([13, 250, 243])),
]


import tensorflow as tf
import numpy as np
import scipy.misc as misc
import math,os
import functools
from utlis import *
from ops import *
from get_inputs import *

class StarGan:
    def __init__(self,args):
        self.batch_size = args.batch_size
        print(args.features_A)
        self.f1 = args.features_A[0].split(',')
        self.f2 = args.features_B[0].split(',')

        self.real_im = tf.placeholder(tf.float32,[None,256,256,3],'real_im')
        self.real_domin = tf.placeholder(tf.float32,[None,2],'real_domin')
        self.label_tar = tf.placeholder(tf.float32,[None,len(self.f1+self.f2)],'label_tar')
        self.label_real = tf.placeholder(tf.float32,[None,len(self.f1+self.f2)],'label_real')

        self.ident = tf.placeholder(tf.float32,[None,len(self.f1+self.f2)],'ident')
        self.lr = tf.placeholder(tf.float32,name='lr')
        self.init_lr = args.lr

        self.txt1_dir = args.txt1_dir
        self.txt2_dir = args.txt2_dir

        self.imfile1 = args.imfile1
        self.imfile2 = args.imfile2
        self.epoch = args.epoch
        self.ld = args.ld
        self.adv_weight = args.adv_weight
        self.cls_weight = args.cls_weight
        self.recy_weight = args.recy_weight
        self.save_path = args.save_path
        self.save_iters = args.save_iters
        self.logdir = args.logdir
        self.ckpt_path = args.ckpt_path
        self.test_ckpt = args.test_ckpt
        self.test_imsdir = args.test_imsdir
        self.test_save_path = args.test_save_path
        self.test_target = args.test_target[0].split(',')

#  label: [batch,9]
    def G(self,im,one_domin,label,reu):
        with tf.variable_scope('G_net',reuse=reu):
            label_shape = [i.value for i in label.get_shape()]
            domin_shape = [i.value for i in one_domin.get_shape()]
            label2im = tf.tile(tf.reshape(label,[-1,1,1,label_shape[-1]]),[1,256,256,1])
            domin2im = tf.tile(tf.reshape(one_domin,[-1,1,1,domin_shape[-1]]),[1,256,256,1])
            x_ = tf.concat([im,label2im,domin2im],3)
            conv1 = conv('conv1',x_,7*7,64,1,True)
            ins1 = ins_norm('ins1',conv1)
            relu1 = relu('relu1',ins1)
            conv2 = conv('conv2',relu1,4*4,128,2,True)
            ins2 = ins_norm('ins2',conv2)
            relu2 = relu('relu2',ins2)
            conv3 = conv('conv3',relu2,4*4,256,2,True)
            ins3 = ins_norm('ins3',conv3)
            relu3 = relu('relu3',ins3)
            block1 = block('block1',relu3)
            block2 = block('block2',block1)
            block3 = block('block3',block2)
            block4 = block('block4',block3)
            block5 = block('block5',block4)
            block6 = block('block6',block5)
            up_1 = upsample('up1',block6,4*4,128,2,True)
            ins_up1 = ins_norm('ins_up1',up_1)
            relu_up1 = relu('relu_up1',ins_up1)
            up_2 = upsample('up_2',relu_up1,4*4,64,2,True)
            ins_up2 = ins_norm('ins_up2',up_2)
            relu_up2 = relu('relu_up2',ins_up2)
            conv_end = conv('conv_end',relu_up2,7*7,3,1,True)
            tanh_end = tanh('tanh_end',conv_end)
#             shapes = self.look_shape(conv1,conv2,conv3,block1,block2,block3,block4,block5,block6,up_1,up_2,conv_end)
            return tanh_end

    def D(self,im,reu):
        with tf.variable_scope('D_net',reuse=reu):
            conv1 = conv('conv1',im,4*4,64,2,True)
            lrelu1 = lrelu('lrelu1',conv1)

            conv2 = conv('conv2',lrelu1,4*4,128,2,True)
            lrelu2 = lrelu('lrelu2',conv2)

            conv3 = conv('conv3',lrelu2,4*4,256,2,True)
            lrelu3 = lrelu('lrelu3',conv3)

            conv4 = conv('conv4',lrelu3,4*4,512,2,True)
            lrelu4 = lrelu('lrelu4',conv4)

            conv5 = conv('conv5',lrelu4,4*4,1024,2,True)
            lrelu5 = lrelu('lrelu5',conv5)

            conv6 = conv('conv6',lrelu5,4*4,2048,2,True)
            lrelu6 = lrelu('lrelu6',conv6)

            conv_torF = conv('conv_TorF',lrelu6,3*3,1,1,True)
            conv_TorF = sigmoid('conv_TorF',conv_torF)
            conv_cls = conv('conv_cls',lrelu6,4*4,len(self.f1+self.f2),1,False)
            conv_Cls = sigmoid('conv_Cls',tf.reshape(conv_cls,[-1,len(self.f1+self.f2)]))
            return conv_TorF,conv_Cls

    def cls_loss(self,probs,label_real,probs1,label_tar,ident):
        d_loss = ((-1)*tf.log(probs)*label_real-tf.log(1-probs)*(1-label_real))*ident
        D_loss = tf.reduce_mean(tf.reduce_sum(d_loss,1,True)/tf.reduce_sum(ident,1,True))

        g_loss = ((-1)*tf.log(probs1)*label_tar-tf.log(1-probs1)*(1-label_tar))*ident
        G_loss = tf.reduce_mean(tf.reduce_sum(g_loss,1,True)/tf.reduce_sum(ident,1,True))
        return D_loss,G_loss

    def cycle_loss(self,im1,im2):
        return tf.reduce_mean(tf.abs(im1-im2))

    def gradient_panalty(self,real,fake):
        shape = tf.shape(real)
        epsilong = tf.random_uniform(shape=shape, minval=0., maxval=1.)
        _, var = tf.nn.moments(real, axes=[0, 1, 2, 3])
        x_std = tf.sqrt(var)  # magnitude of noise decides the size of local region
        noise = 0.5 * x_std * epsilong  # delta in paper
        alpha = tf.random_uniform(shape=[shape[0], 1, 1, 1], minval=-1., maxval=1.)
        interpolated = tf.clip_by_value(real + alpha * noise, -1., 1.)  # x_hat should be in the space of X
        TorF, _ = self.D(interpolated,True)

        grad = tf.gradients(TorF,interpolated)[0] # gradient of D(interpolated)
        grad_norm = tf.norm(tf.layers.flatten(grad), axis=1) # l2 norm
        grad_penalty = self.ld * tf.reduce_mean(tf.square(grad_norm - 1.))
        return grad_penalty

    def adv_loss(self,r_or_f,r_or_f1):
        avdloss_d = tf.reduce_mean((-1)*tf.log(r_or_f)+(-1)*tf.log(1-r_or_f1))
        advloss_g = tf.reduce_mean((-1)*tf.log(r_or_f1))
        return avdloss_d,advloss_g

    def build(self):
        self.fake_im = self.G(self.real_im,self.real_domin,self.label_tar,False)
        self.recons_im = self.G(self.fake_im,self.real_domin,self.label_real,True)
        self.r_or_f,self.probs = self.D(self.real_im,False)
        self.r_or_f1,self.probs1 = self.D(self.fake_im,True)

        self.adv_D,self.adv_G = self.adv_loss(self.r_or_f,self.r_or_f1)

        self.class_D,self.class_G = self.cls_loss(self.probs,self.label_real,self.probs1,self.label_tar,self.ident)

        self.recyle_loss = self.cycle_loss(self.real_im,self.recons_im)
        self.d_sum_loss = self.adv_weight * self.adv_D + self.cls_weight * self.class_D
        self.g_sum_loss = self.adv_weight * self.adv_G + self.cls_weight * self.class_G + self.recy_weight * self.recyle_loss

    def train(self):
        self.build()
        G_vars = [var for var in tf.trainable_variables() if 'G_net' in var.name]
        D_vars = [var for var in tf.trainable_variables() if 'D_net' in var.name]
        G_optim = tf.train.AdamOptimizer(self.lr,beta1=0.5,beta2=0.999).minimize(self.g_sum_loss,var_list=G_vars)
        D_optim = tf.train.AdamOptimizer(self.lr,beta1=0.5,beta2=0.999).minimize(self.d_sum_loss,var_list=D_vars)

        im_names, domin, real_label = get_inputs(self.txt1_dir,self.txt2_dir,self.f1,self.f2,self.imfile1,self.imfile2)
        data = tf.data.Dataset.from_tensor_slices((im_names, domin, real_label))
        data = data.map(process)
        data = data.shuffle(1000).batch(self.batch_size).repeat()
        Data = data.make_one_shot_iterator()
        Ims, Domin, Real_label = Data.get_next()
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            Graph = tf.summary.FileWriter(self.logdir,sess.graph)
            Saver = tf.train.Saver(max_to_keep=10)
            for ep in range(self.epoch):
                lear_rate = self.init_lr-ep*(self.init_lr/self.epoch) if ep<=(self.epoch-1) else self.init_lr/self.epoch
                for num in range(len(im_names)//self.batch_size):
                    ims,dom,real_l = sess.run([Ims,Domin,Real_label])
                    ident, l_tar = decidel(self.batch_size, dom, real_l, len(self.f1), len(self.f2))
                    dict_D = {self.real_im:ims,self.label_real:real_l,self.label_tar:l_tar,self.real_domin:dom,self.ident:ident,self.lr:lear_rate}
                    _, dsum, advd, classd, pprobs, pprobs1 = sess.run([D_optim, self.d_sum_loss, self.adv_D, self.class_D, self.probs,self.probs1],feed_dict=dict_D)
                    print('current_epoch: ',ep)
                    print('  total_loss_D: %f, adv_loss_d: %f, class_loss_d: %f'%(dsum,advd,classd))
#                     _,D_Loss,gp = sess.run([D_optim,self.d_sum_loss,self.GP],feed_dict=dict_D
                    dict_G = {self.real_im:ims,self.label_real:real_l,self.label_tar:l_tar,self.real_domin:dom,self.ident:ident,self.lr:lear_rate}
                    _, gsum, advg, classg, Fims = sess.run([G_optim, self.g_sum_loss, self.adv_G, self.class_G, self.fake_im],feed_dict=dict_G)
                    print('  total_loss_G: %f, adv_loss_g: %f, class_loss_g: %f'%(gsum, advg, classg))
                    if int(ep*len(im_names)+num*self.batch_size)%self.save_iters==0:
                        print('save images at epoch: %d, training numbers is %d'%(ep,num))
                        saveim(Fims,ep,num,self.save_path)
                Saver.save(sess, self.ckpt_path, global_step=ep)
                print('save_success! ')
            return True

    def read_testims(self):
        ims_dir = [self.test_imsdir+imname for imname in os.listdir(self.test_imsdir)]
        zeros = np.zeros([len(ims_dir),256,256,3])
        for i in range(len(ims_dir)):
            im = misc.imread(ims_dir[i])
            if np.shape(im)!=(256,256,3):
                im = misc.imresize(im,[256,256])
            zeros[i,:,:,:] = im
        Dict = self.get_domin_target(self.test_target,len(ims_dir))
        return zeros,Dict

    def get_domin_target(self,features,test_batch):
        fea_f1 = []
        fea_f2 = []
        for feature in features:
            if feature in self.f1:
                fea_f1.append(feature)
            else:
                fea_f2.append(feature)
        if fea_f1:
            domin = np.zeros([test_batch,2])
            domin[:,0] = 1.
            empty_f1 = np.zeros([test_batch,len(self.f1+self.f2)],np.float32)
            for fea in fea_f1:
                empty_f1[:,self.f1.index(fea)] = 1.
        if fea_f2:
            domin2 = np.zeros([test_batch,2])
            domin2[:,1] = 1.
            empty_f2 = np.zeros([test_batch,len(self.f1+self.f2)],np.float32)
            for fea in fea_f2:
                empty_f2[:,len(self.f1)+self.f2.index(fea)] = 1.
        if fea_f1:
            if fea_f2:
                sin_or_dou = 2
                return {'sin_or_dou':sin_or_dou,'Dom1':domin,'Label_tar1':empty_f1,'Dom2':domin2,'Label_tar2':empty_f2}
            else:
                sin_or_dou = 1
                return {'sin_or_dou':sin_or_dou,'Dom':domin,'Label_tar':empty_f1}
        else:
            sin_or_dou = 1
            return {'sin_or_dou':sin_or_dou,'Dom':domin2,'Label_tar':empty_f2}

    def test(self):
        self.build()
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            Saver = tf.train.Saver()
            im_in,Dict = self.read_testims()
            if Dict['sin_or_dou']==2:
                Saver.restore(sess,tf.train.latest_checkpoint(self.test_ckpt))
                print('load_success')
                test_dict = {self.real_im:im_in,self.label_tar:Dict['Label_tar1'],self.real_domin:Dict['Dom1']}
                fake_ims = sess.run(self.fake_im,feed_dict=test_dict)
                test_dict2 = {self.real_im:fake_ims,self.label_tar:Dict['Label_tar2'],self.real_domin:Dict['Dom2']}
                Fake_ims = sess.run(self.fake_im,feed_dict=test_dict2)
                saveim(Fake_ims,0,0,self.test_save_path)
                print('save_success')
            else:
                Saver.restore(sess,tf.train.latest_checkpoint(self.test_ckpt))
                print('load_success')
                test_dict = {self.real_im:im_in,self.label_tar:Dict['Label_tar'],self.real_domin:Dict['Dom']}
                fake_ims = sess.run(self.fake_im,feed_dict=test_dict)
                saveim(fake_ims,0,0,self.test_save_path)
                print('save_success')


#
# License: The MIT License (MIT)
# Author:E-bank IT team
# Author email: @ebanka-it.com
# Date: Fri Aug 23 2019
# Description:
# This Django module loads list of branches names and street names
# makes an random combination of it (plus random home number, mobile phone number)
#  and returns it as the response
#

from django.http import HttpResponse, JsonResponse
from django.views.generic import TemplateView

#disable=pylint(import-error)
import seaborn as sbn
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

branch = 0
ranBranch = 0

def index(request):
    return HttpResponse("Сервер ради!");

class Random(TemplateView):
    #disable: pylint(no-self-argument)
    def getRandomBranch(self):
        branch = pd.read_csv('C:/proba/mean-course/backend/server-python/env/files/random-branches.csv', delimiter = ',').to_numpy()
        randomBranch = branch[np.random.randint(len(branch))]
        ret = str(randomBranch)
        ret = ret[2:len(ret)-2]
        return(ret)

    def getRandomMobileNumber(self):
        num = "+3816" + str(np.random.randint(0,9)) + str(np.random.randint(1000,9999)) + str(np.random.randint(100,9999))
        return(num)

    def getRandomAddress(self):
        address = pd.read_csv('C:/proba/mean-course/backend/server-python/env/files/addresses.csv', delimiter = ',').to_numpy()
        address = address[np.random.randint(len(address))]
        add = str(address)
        add = add[2:len(add)-2]
        return(add)

    def getRandomHouseNumber(self):
        houseNum = str(np.random.randint(1,100))
        return(houseNum)

    def getRandomData(request):
        bra = Random.getRandomBranch(1)
        adr = Random.getRandomAddress(1)
        num = Random.getRandomMobileNumber(1)
        hnum = Random.getRandomHouseNumber(1)
        return JsonResponse({'number': num, 'address': adr, 'hnumber': hnum, 'branch': bra})


#!/usr/bin/env python

import astropy.units as u
import numpy as np
from tollan.utils.log import get_logger

from ..mapping.raster import SkyRasterScanModel
from ..mapping.lissajous import (
    SkyLissajousModel, SkyDoubleLissajousModel,
    SkyRastajousModel)

class LmtOTComposer(object):
    """A class to create LMT OT scripts."""

    logger = get_logger()

    def __init__(self):
        pass

    def make_pointing_steps(self):
        pass

    def make_mapping_steps(
            self, mapping_model, target_name=None, t_exp=None):
        steps = list()
        target = mapping_model.target
        ref_frame = mapping_model.ref_frame
        if t_exp is None:
            t_exp = mapping_model.t_pattern
        if target_name is None:
            target_name = 'unknown_source'
        if target.frame.name == 'icrs':
            # Eq
            source_line = (
                "Source Source;  Source  -BaselineList [] -CoordSys Eq"
                " -DecProperMotionCor 0"
                " -Dec[0] {Dec} -Dec[1] {Dec}"
                " -El[0] 0.000000 -El[1] 0.000000 -EphemerisTrackOn 0"
                " -Epoch 2000.0"
                " -GoToZenith 0 -L[0] 0.0 -L[1] 0.0 -LineList [] -Planet None"
                " -RaProperMotionCor 0 -Ra[0] {RA} -Ra[1] {RA}"
                " -SourceName {Name}"
                " -VelSys Lsr -Velocity 0.000000 -Vmag 0.0").format(
                    RA=target.ra.to_string(
                        unit=u.hour, pad=True,
                        decimal=False, fields=3, sep=':'),
                    Dec=target.dec.to_string(
                        unit=u.degree, pad=True, alwayssign=True,
                        decimal=False, fields=3, sep=':'),
                    Name=target_name,
                    )
            steps.append(source_line)
        else:
            raise ValueError(f"Invalid target frame {target.frame.name}")
        # mapping line
        # TODO fix this for non-offset models
        m = mapping_model.offset_mapping_model
        if isinstance(m, (SkyLissajousModel)):
            mapping_line = (
                "Lissajous -ExecMode 0 -RotateWithElevation 0 -TunePeriod 0"
                " -TScan {TScan} -ScanRate {ScanRate}"
                " -XLength {XLength} -YLength {YLength} -XOmega {XOmega}"
                " -YOmega {YOmega} -XDelta {XDelta}"
                " -XLengthMinor 0.0 -YLengthMinor 0.0 -XDeltaMinor 0.0"
                ).format(**self._sky_lissajous_params_to_mapping_step_params(
                    x_length=m.x_length,
                    y_length=m.y_length,
                    x_omega=m.x_omega,
                    y_omega=m.y_omega,
                    delta=m.delta,
                    t_exp=t_exp
                    ))
        elif isinstance(m, (SkyDoubleLissajousModel)):
            major_params = self._sky_lissajous_params_to_mapping_step_params(
                    x_length=m.x_length_0,
                    y_length=m.y_length_0,
                    x_omega=m.x_omega_0,
                    y_omega=m.y_omega_0,
                    delta=m.delta_0,
                    t_exp=t_exp
                    )
            minor_params = self._sky_lissajous_params_to_mapping_step_params(
                    x_length=m.x_length_1,
                    y_length=m.y_length_1,
                    x_omega=m.x_omega_1,
                    y_omega=m.y_omega_1,
                    delta=m.delta_1,
                    t_exp=0 << u.s
                    )
            self.logger.warning(
                    "some parameters will be ignored during the exporting "
                    "and the result will be different.")
            mapping_line = (
                "Lissajous -ExecMode 0 -RotateWithElevation 0 -TunePeriod 0"
                " -TScan {TScan} -ScanRate {ScanRate}"
                " -XLength {XLength} -YLength {YLength} -XOmega {XOmega}"
                " -YOmega {YOmega} -XDelta {XDelta}"
                " -XLengthMinor {XLengthMinor} -YLengthMinor {YLengthMinor}"
                " -XDeltaMinor {XDeltaMinor}"
                ).format(
                    XLengthMinor=minor_params['XLength'],
                    YLengthMinor=minor_params['YLength'],
                    XDeltaMinor=minor_params['XDelta'],
                    **major_params)
        elif isinstance(m, (SkyRasterScanModel)):
            mapping_line = (
                "RasterMap Map; Map -ExecMode 0 -HPBW 1 -HoldDuringTurns 0"
                " -MapMotion Continuous -NumPass 1 -NumRepeats 1 -NumScans 0"
                " -RowsPerScan 1000000 -ScansPerCal 0 -ScansToSkip 0"
                " -TCal 0 -TRef 0 -TSamp 1"
                " -MapCoord {MapCoord} -ScanAngle {ScanAngle}"
                " -XLength {XLength} -XOffset 0 -XRamp 0 -XStep {XStep}"
                " -YLength {YLength} -YOffset 0 -YRamp 0 -YStep {YStep}"
                ).format(**self._sky_raster_params_to_mapping_step_params(
                    length=m.length,
                    space=m.space,
                    n_scans=m.n_scans,
                    rot=m.rot,
                    speed=m.speed,
                    ref_frame=ref_frame
                    ))
        elif isinstance(m, (SkyRastajousModel)):
            raster_params = self._sky_raster_params_to_mapping_step_params(
                    length=m.length,
                    space=m.space,
                    n_scans=m.n_scans,
                    rot=m.rot,
                    speed=m.speed,
                    ref_frame=ref_frame
                    )
            major_params = self._sky_lissajous_params_to_mapping_step_params(
                    x_length=m.x_length_0,
                    y_length=m.y_length_0,
                    x_omega=m.x_omega_0,
                    y_omega=m.y_omega_0,
                    delta=m.delta_0,
                    t_exp=t_exp
                    )
            minor_params = self._sky_lissajous_params_to_mapping_step_params(
                    x_length=m.x_length_1,
                    y_length=m.y_length_1,
                    x_omega=m.x_omega_1,
                    y_omega=m.y_omega_1,
                    delta=m.delta_1,
                    t_exp=0 << u.s
                    )
            self.logger.warning(
                    "some parameters will be ignored during the exporting "
                    "and the result will be different.")
            mapping_line_lissajous = (
                "Lissajous -ExecMode 1 -RotateWithElevation 0 -TunePeriod 0"
                " -TScan {TScan} -ScanRate {ScanRate}"
                " -XLength {XLength} -YLength {YLength} -XOmega {XOmega}"
                " -YOmega {YOmega} -XDelta {XDelta}"
                " -XLengthMinor {XLengthMinor} -YLengthMinor {YLengthMinor}"
                " -XDeltaMinor {XDeltaMinor}"
                ).format(
                    XLengthMinor=minor_params['XLength'],
                    YLengthMinor=minor_params['YLength'],
                    XDeltaMinor=minor_params['XDelta'],
                    **major_params)
            mapping_line_raster = (
                "RasterMap Map; Map -ExecMode 1 -HPBW 1 -HoldDuringTurns 0"
                " -MapMotion Continuous -NumPass 1 -NumRepeats 1 -NumScans 0"
                " -RowsPerScan 1000000 -ScansPerCal 0 -ScansToSkip 0"
                " -TCal 0 -TRef 0 -TSamp 1"
                " -MapCoord {MapCoord} -ScanAngle {ScanAngle}"
                " -XLength {XLength} -XOffset 0 -XRamp 0 -XStep {XStep}"
                " -YLength {YLength} -YOffset 0 -YRamp 0 -YStep {YStep}"
                ).format(**raster_params)
            mapping_line = f"{mapping_line_lissajous}\n{mapping_line_raster}"
        else:
            raise NotImplementedError
        steps.append(mapping_line)
        return steps

    def make_setup_steps(self, instru_name):
        if instru_name == 'toltec':
            return self._toltec_setup_steps()

        raise NotImplementedError(
            f"setup steps not implemented for {instru_name}")

    def _toltec_setup_steps(self):
        steps = list()
        steps.append('ObsGoal Dcs; Dcs -ObsGoal Science')
        return steps

    @classmethod
    def _sky_lissajous_params_to_mapping_step_params(
            cls,
            x_length, y_length, x_omega, y_omega,
            delta, t_exp
            ):
        v_x = x_length * x_omega.quantity.to(
                u.Hz, equivalencies=[(u.cy/u.s, u.Hz)])
        v_y = y_length * y_omega.quantity.to(
                u.Hz, equivalencies=[(u.cy/u.s, u.Hz)])
        return dict(
            XLength=x_length.quantity.to_value(u.arcmin),
            YLength=y_length.quantity.to_value(u.arcmin),
            XOmega=x_omega.quantity.to_value(u.rad/u.s),
            YOmega=y_omega.quantity.to_value(u.rad/u.s),
            XDelta=delta.quantity.to_value(u.rad),
            TScan=t_exp,
            ScanRate=np.hypot(v_x, v_y).to_value(u.arcsec / u.s),
            )

    @classmethod
    def _sky_raster_params_to_mapping_step_params(
            cls,
            length, space, n_scans, rot, speed, ref_frame
            ):
        if ref_frame.name == 'icrs':
            MapCoord = 'Ra'
        elif ref_frame.name == 'altaz':
            MapCoord = 'Az'
        else:
            raise NotImplementedError(f"invalid ref_frame {ref_frame}")
        return dict(
            MapCoord=MapCoord,
            ScanAngle=rot.quantity.to_value(u.deg),
            XLength=length.quantity.to_value(u.arcsec),
            XStep=speed.quantity.to_value(u.arcsec / u.s),
            YLength=(n_scans * space.quantity).to_value(u.arcsec),
            YStep=space.quantity.to_value(u.arcsec),
            )


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import flags
import bottle
import io
import math
import numpy as np
import os
import pickle
import random
from scipy.special import softmax
import spacy
import sys
import tensorflow as tf

import data_utils
import model

os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"

# GPU config
flags.DEFINE_integer("num_core_per_host", default=8,
                     help="Number of cores per host")

# Experiment (data/checkpoint/directory) config
flags.DEFINE_string("data_dir", default="",
                    help="Path to tf-records directory.")
flags.DEFINE_string("record_info_dir", default="",
                    help="Path to local directory containing filenames.txt.")
flags.DEFINE_string("corpus_info_path", default="",
                    help="Path to corpus-info.json file.")
flags.DEFINE_string("model_dir", default=None,
                    help="Estimator model_dir.")
flags.DEFINE_string("eval_ckpt_path", None,
                    help="Checkpoint path for do_test evaluation."
                         "If set, model_dir will be ignored."
                         "If unset, will use the latest ckpt in model_dir.")

# Training config
flags.DEFINE_integer("eval_batch_size", default=60,
                     help="Size of valid batch.")

# Model config
flags.DEFINE_integer("tgt_len", default=70,
                     help="Number of steps to predict")
flags.DEFINE_integer("mem_len", default=70,
                     help="Number of steps to cache")
flags.DEFINE_bool("same_length", default=False,
                  help="Same length attention")
flags.DEFINE_integer("clamp_len", default=-1,
                     help="Clamp length")

flags.DEFINE_integer("n_layer", default=6,
                     help="Number of layers.")
flags.DEFINE_integer("d_model", default=500,
                     help="Dimension of the model.")
flags.DEFINE_integer("d_embed", default=500,
                     help="Dimension of the embeddings.")
flags.DEFINE_integer("n_head", default=10,
                     help="Number of attention heads.")
flags.DEFINE_integer("d_head", default=50,
                     help="Dimension of each attention head.")
flags.DEFINE_integer("d_inner", default=1000,
                     help="Dimension of inner hidden size in positionwise feed-forward.")
flags.DEFINE_float("dropout", default=0.1,
                   help="Dropout rate.")
flags.DEFINE_float("dropatt", default=0.1,
                   help="Attention dropout rate.")
flags.DEFINE_bool("untie_r", default=False,
                  help="untie r_w_bias and r_r_bias")

# Adaptive Softmax / Embedding
flags.DEFINE_integer("div_val", default=1,
                     help="Divide the embedding size by this val for each bin")
flags.DEFINE_bool("proj_share_all_but_first", default=False,
                  help="True to share all but first projs, False not to share.")
flags.DEFINE_bool("proj_same_dim", default=True,
                  help="Project the bin with the same dimension.")

# Parameter initialization
flags.DEFINE_enum("init", default="normal",
                  enum_values=["normal", "uniform"],
                  help="Initialization method.")
flags.DEFINE_float("init_std", default=0.02,
                   help="Initialization std when init is normal.")
flags.DEFINE_float("proj_init_std", default=0.01,
                   help="Initialization std for embedding projection.")
flags.DEFINE_float("init_range", default=0.1,
                   help="Initialization std when init is uniform.")

FLAGS = flags.FLAGS
FLAGS(sys.argv)

def get_model_fn(n_token, cutoffs):
    def model_fn(inp, tgt, mems, is_training, generate_text, prev_text):
        if generate_text == False:
            inp = tf.transpose(inp, [1, 0])
            tgt = tf.transpose(tgt, [1, 0])

        if FLAGS.init == "uniform":
            initializer = tf.initializers.random_uniform(
                minval=-FLAGS.init_range,
                maxval=FLAGS.init_range,
                seed=None)
        elif FLAGS.init == "normal":
            initializer = tf.initializers.random_normal(
                stddev=FLAGS.init_std,
                seed=None)
            proj_initializer = tf.initializers.random_normal(
                stddev=FLAGS.proj_init_std,
                seed=None)

        tie_projs = [False for _ in range(len(cutoffs) + 1)]
        if FLAGS.proj_share_all_but_first:
            for i in range(1, len(tie_projs)):
                tie_projs[i] = True

        loss, new_mems, t_output = model.transformer(
            dec_inp=inp,
            target=tgt,
            mems=mems,
            n_token=n_token,
            n_layer=FLAGS.n_layer,
            d_model=FLAGS.d_model,
            d_embed=FLAGS.d_embed,
            n_head=FLAGS.n_head,
            d_head=FLAGS.d_head,
            d_inner=FLAGS.d_inner,
            dropout=FLAGS.dropout,
            dropatt=FLAGS.dropatt,
            initializer=initializer,
            proj_initializer=proj_initializer,
            is_training=is_training,
            generate_text=generate_text,
            prev_text=prev_text,
            mem_len=FLAGS.mem_len,
            cutoffs=cutoffs,
            div_val=FLAGS.div_val,
            tie_projs=tie_projs,
            input_perms=None,
            target_perms=None,
            head_target=None,
            same_length=FLAGS.same_length,
            clamp_len=FLAGS.clamp_len,
            untie_r=FLAGS.untie_r,
            proj_same_dim=FLAGS.proj_same_dim)

        # number of parameters
        num_params = sum([np.prod(v.shape) for v in tf.compat.v1.trainable_variables()])
        tf.compat.v1.logging.info('#params: {}'.format(num_params))

        # format_str = '{{:<{0}s}}\t{{}}'.format(
        #     max([len(v.name) for v in tf.compat.v1.trainable_variables()]))
        # for v in tf.compat.v1.trainable_variables():
        #   tf.compat.v1.logging.info(format_str.format(v.name, v.get_shape()))

        if is_training:
            all_vars = tf.compat.v1.trainable_variables()
            grads = tf.gradients(loss, all_vars)
            grads_and_vars = list(zip(grads, all_vars))

            return loss, new_mems, grads_and_vars
        elif generate_text:
            return loss, new_mems, t_output
        else:
            return loss, new_mems, inp

    return model_fn

def single_core_graph(n_token, cutoffs, is_training, inp, tgt, mems, generate_text=False, prev_text=None):
    model_fn = get_model_fn(
        n_token=n_token,
        cutoffs=cutoffs)

    model_ret = model_fn(
        inp=inp,
        tgt=tgt,
        mems=mems,
        is_training=is_training,
        generate_text=generate_text,
        prev_text=prev_text)

    return model_ret

def generate_text(n_token, cutoffs, ps_device, vocab, text_length=100, word_range=3):
    per_core_bsz = FLAGS.eval_batch_size // FLAGS.num_core_per_host
    tower_mems, tower_losses, tower_new_mems = [], [], []

    for i in range(FLAGS.num_core_per_host):
        with tf.device("/device:CPU:0"), \
             tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(), reuse=tf.AUTO_REUSE):
            mems_i = [tf.compat.v1.placeholder(tf.float32,
                                               [FLAGS.mem_len, per_core_bsz, FLAGS.d_model])
                      for _ in range(FLAGS.n_layer)]

            prev_text_tf = tf.compat.v1.placeholder(tf.int64,
                                                    [FLAGS.tgt_len, 1])

            loss_i, new_mems_i, t_output = single_core_graph(
                n_token=n_token,
                cutoffs=cutoffs,
                is_training=False,
                inp=None,
                tgt=None,
                mems=mems_i,
                generate_text=True,
                prev_text=prev_text_tf)

            tower_mems.append(mems_i)
            tower_losses.append(loss_i)
            tower_new_mems.append(new_mems_i)

    ## sum losses across towers
    if len(tower_losses) > 1:
        loss = tf.add_n(tower_losses) / len(tower_losses)
    else:
        loss = tower_losses[0]

    tower_mems_np = [
        [np.zeros([FLAGS.mem_len, per_core_bsz, FLAGS.d_model], dtype=np.float32)
         for layer in range(FLAGS.n_layer)]
        for core in range(FLAGS.num_core_per_host)
    ]

    start_text = vocab.encode_file("temp_input_text.txt", ordered=True)

    prev_text_np = []
    start_en_dec = []

    for i in start_text[-FLAGS.tgt_len:]:
        prev_text_np.append([i])
        start_en_dec.append(vocab.get_sym(i))

    start_en_dec = ' '.join(start_en_dec)

    text = []

    saver = tf.compat.v1.train.Saver()

    with tf.compat.v1.Session(
            config=tf.compat.v1.ConfigProto(allow_soft_placement=True, device_count={'GPU': 0})) as sess:
        sess.run(tf.compat.v1.global_variables_initializer())

        if FLAGS.eval_ckpt_path is None:
            eval_ckpt_path = tf.compat.v1.train.latest_checkpoint(FLAGS.model_dir)
        else:
            eval_ckpt_path = FLAGS.eval_ckpt_path
        tf.compat.v1.logging.info("Model: {}".format(eval_ckpt_path))
        saver.restore(sess, eval_ckpt_path)

        fetches = [loss, tower_new_mems, t_output]

        print("Generating text...")

        step = 0
        while step < text_length:
            feed_dict = {}
            for i in range(FLAGS.num_core_per_host):
                for m, m_np in zip(tower_mems[i], tower_mems_np[i]):
                    feed_dict[m] = m_np
                feed_dict[prev_text_tf] = prev_text_np

            fetched = sess.run(fetches, feed_dict=feed_dict)

            sm_prob = fetched[2]
            sm_prob = softmax(sm_prob[-1, 0, :])
            sorted_indices = np.argsort(sm_prob)

            # random word choice from the set of tokens with the highest probabilities
            # If we put the most probable word, repetition of the same phrases can occur
            # word_range = 3
            next_word_index = int(math.floor(random.uniform(1, word_range + 1)))
            next_word = sorted_indices[-next_word_index]
            text.append(vocab.get_sym(next_word))
            prev_text_np = np.append(prev_text_np[1:], [next_word]).reshape((FLAGS.tgt_len, 1))
            loss_np, tower_mems_np = fetched[:2]

            if vocab.get_sym(next_word) != "<eos>":
                step = step + 1

    stext = ' '.join(text)
    return start_en_dec, stext

input_page = """

    Legal Text Generation with Transformers

    Lazar Peric, Stefan Mijic, Dominik Stammbach and Elliott Ash

    ETH Zurich

<form action="/transformer" method="post">
    Text length: <input name="text_length" type="text" value="100" />
    Randomness: <input name="randomness" type="text" value="3" />

    Input text:

    <textarea rows="18" style="width:90%;" name="input_text">Law is a bottomless pit.</textarea>

    <input value="Generate" type="submit" />
</form>
"""

output_page = """

    Legal Text Generation with Transformers

    Lazar Peric, Stefan Mijic, Dominik Stammbach and Elliott Ash

    ETH Zurich

Input text:

{{input_text}}

Output text:

{{output_text}}
"""

@bottle.get('/transformer')
def get_transformer_text():
    return input_page

@bottle.post('/transformer')
def generate_new_text():
    original_text = bottle.request.forms.get('input_text')

    def process_text(text, process=False):
        if process:
            text = text.replace(u" , ", u", ")
            text = text.replace(u" . ", u". ")
            text = text.replace(u" ; ", u"; ")
            text = text.replace(u" : ", u": ")
            text = text.replace(u" ( ", u" (")
            text = text.replace(u" ) ", u") ")
            text = text.replace(u"<eos>", u" ")
            text = text.replace(u" 's ", u"'s ")
            while text.count(u"  ") > 0:
                text = text.replace(u"  ", u" ")
        return text

    text_length = abs(int(bottle.request.forms.get("text_length")))
    word_range = abs(int(bottle.request.forms.get("randomness")))

    if word_range == 0:
        word_range = 1

    print("Processing text...")
    doc = nlp(original_text)
    text_out = []
    for sent in doc.sents:
        for token in sent:
            text_out.append(token.text)
    if len(text_out) < test_tgt_len:
        text_out = ["<unk>"] * (test_tgt_len - len(text_out)) + text_out
    text_out = u" " + ' '.join(text_out)
    while text_out.count(u"  ") > 0:
        text_out = text_out.replace(u"  ", u" ")

    file_o = io.open("temp_input_text.txt", "w", encoding='utf8')
    file_o.write(text_out)
    file_o.close()

    input_text, output_text = generate_text(n_token, cutoffs, "/gpu:0", vocab, text_length=text_length,
                                            word_range=word_range)

    os.remove("temp_input_text.txt")

    input_text = process_text(original_text, process=True)
    output_text =process_text(output_text, process=True)
    return bottle.template(output_page, input_text=original_text, output_text=output_text)

def main(unused_argv):
    del unused_argv  # Unused
    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
    # chech: http://localhost:8787/transformer
    bottle.run(host='0.0.0.0', port=8787, reloader=True)

print("Loading Spacy...")
nlp = spacy.load('en_core_web_lg')
print("Loading corpus info...")
corpus_info = data_utils.get_corpus_info(FLAGS.corpus_info_path)
n_token = corpus_info["vocab_size"]
cutoffs = corpus_info["cutoffs"][1:-1]
test_tgt_len = FLAGS.tgt_len
vocab_path = os.path.join(FLAGS.data_dir, "cache_vocab.pkl")
print("Loading cached vocabulary...")
with open(vocab_path, "rb") as fp:
    vocab = pickle.load(fp)

tf.compat.v1.app.run()


#!/usr/bin/python

import itertools

import numpy as np

from . import PERCENT, LOGGER
from . import constants as c
from .rect import Rect
from .stream_parser import StreamParser
from .template_matcher import TemplateMatcher
from .util import overlay_map, find_dlt

class Viewfinder(StreamParser):
    def __init__(self, filename, config):
        StreamParser.__init__(self, filename)

        # TODO document these.
        frames_to_sample = config.get("frames_to_sample", 15)
        self.frames = [
            frame
            for _, frame in self.sample_frames(num_samples=frames_to_sample)
        ]

        self.percent_y_tolerance = config.get("percent_y_tolerance", 5)
        min_scale = config.get("min_scale", 0.72)
        max_scale = config.get("max_scale", 0.87)
        scale_step = config.get("scale_step", 0.03)

        self.scales = np.arange(min_scale, max_scale, scale_step)

    def get_scale_and_screen(self):
        """Estimate the screen location and scale."""
        scale, screen = self.detect_screen()
        LOGGER.warning("Estimated screen is %s", screen)
        LOGGER.warning("Estimated scale is %.03f", scale)

        LOGGER.warning("Correcting screen...")
        scale, screen = self.correct_screen(scale, screen)
        LOGGER.warning("Estimated screen is %s", screen)
        LOGGER.warning("Estimated scale is %.03f", scale)

        return (scale, screen)

    def detect_screen(self):
        """Attempt to detect the screen."""
        scale, pct_locations = self.get_pct_locations(self.scales)

        # Approximate screen Y-pos from percents.
        height = int(round(c.SCREEN_HEIGHT * scale))
        width = int(round(c.SCREEN_WIDTH * scale))
        top = int(
            round(np.mean(pct_locations, axis=0)[0] - c.PERCENT_Y_POS * scale)
        )

        # Determine the X-pos by the skewness-kurtosis method.
        LOGGER.info("Generating skewness-kurtosis map...")
        skew_kurt = overlay_map(self.frames) // 255

        actual_width = skew_kurt.shape[1]
        estimated_width = width

        if actual_width < estimated_width:
            LOGGER.warning(
                "Estimated screen is bigger than the video, "
                "assuming left = 0..."
            )
            left = 0
        else:
            possible_lefts = range(actual_width - estimated_width)

            def badness(left):
                return sum(
                    sum(
                        skew_kurt[
                            max(0, top) : max(0, top) + height,
                            left : left + width,
                        ]
                    )
                )

            left = min(possible_lefts, key=badness)

        return (scale, Rect(top, left, height, width))

    def get_pct_locations(self, scales):
        """Detect the percent signs in a random sample of frames."""
        matcher = TemplateMatcher(scales=scales)
        scale, pct_locations = self.locate(
            self.frames, PERCENT, matcher=matcher
        )

        if not scale or not pct_locations:
            raise RuntimeError(
                "This doesn't appear to be Melee (no percent signs found!)"
            )

        # Group the returned locations on the y-axis.
        def y_bin(location):
            y, x = location
            return y // self.percent_y_tolerance

        location_groups = sorted(pct_locations, key=y_bin)
        location_groups = itertools.groupby(location_groups, key=y_bin)
        location_groups = [list(g) for k, g in location_groups]

        # Choose the group by majority vote.
        def votes_for(group):
            return sum(pct_locations[loc] for loc in group)

        group_to_return = max(location_groups, key=votes_for)
        group_to_return = list(group_to_return)

        return (scale, group_to_return)

    def correct_screen(self, scale, screen):
        _, predicted, locations = self.detect_ports(scale, screen)

        if not predicted or not locations:
            raise RuntimeError(
                "This doesn't appear to be Melee (no percent signs found!)"
            )

        # Scaling factor, translation x, translation y
        scale_factor, shift_x, shift_y = find_dlt(predicted, locations)
        scale *= scale_factor
        screen.height *= scale_factor
        screen.width *= scale_factor
        screen.left *= scale_factor
        screen.left += shift_x
        screen.top *= scale_factor
        screen.top += shift_y

        return (scale, screen)

    def get_ports(self, scale, screen):
        ports, _, _ = self.detect_ports(scale, screen)
        if not ports:
            raise RuntimeError("No ports found!")
        LOGGER.warning("Ports are at %s", " ".join(str(s) for s in ports))

        return ports

    def detect_ports(self, scale, screen, max_error=0.06):
        """Find the approximate port locations."""
        ports = []
        predicted = []
        locations = []
        for port_number in range(4):
            pct_roi = screen.subregion(
                c.PERCENT_Y_POS / c.SCREEN_HEIGHT,
                (c.PERCENT_X_POS + port_number * c.PERCENT_X_POS_STEP)
                / c.SCREEN_WIDTH,
                c.PERCENT_HEIGHT / c.SCREEN_HEIGHT,
                c.PERCENT_WIDTH / c.SCREEN_WIDTH,
                padding=max_error,
            )

            pct_left = (
                screen.left
                + (c.PERCENT_X_POS + port_number * c.PERCENT_X_POS_STEP)
                * scale
            )
            pct_top = screen.top + c.PERCENT_Y_POS * scale

            matcher = TemplateMatcher(
                scales=[scale], worst_match=0.6, debug=False
            )
            portscale, location = self.locate(
                self.frames, PERCENT, matcher=matcher, roi=pct_roi
            )

            if portscale is None or not location:
                ports.append(None)
            else:
                location = max(location.keys(), key=lambda loc: location[loc])
                LOGGER.warning(
                    "Detected port %d at (%d, %d) (error %.03f px, %.03f px)",
                    port_number + 1,
                    location[0],
                    location[1],
                    location[0] - pct_top,
                    location[1] - pct_left,
                )

                predicted.append([pct_top, pct_left])
                locations.append(location)

                port_roi = screen.subregion(
                    c.PORT_Y_POS / c.SCREEN_HEIGHT,
                    (c.PORT_X_POS + port_number * c.PORT_X_POS_STEP)
                    / c.SCREEN_WIDTH,
                    c.PORT_HEIGHT / c.SCREEN_HEIGHT,
                    c.PORT_WIDTH / c.SCREEN_WIDTH,
                    padding=max_error,
                )

                port_roi.left = max(port_roi.left, 0)
                port_roi.top = max(port_roi.top, 0)
                ports.append(port_roi)

        return (ports, predicted, locations)


#!/usr/bin/env python

import torch
import pickle
import numpy as np
import time
import os
from shutil import copyfile
import math
from tqdm import tqdm
from torch.utils.data import DataLoader
from rdkit import Chem
from os import path

from model import RNN
from data_structs import Vocabulary, Experience, Dataset
from scoring_functions import get_scoring_function
from utils import Variable, seq_to_smiles, fraction_valid_smiles, unique, get_latent_vector
from vizard_logger import VizardLog

def train_agent(restore_prior_from='data/Prior.ckpt',
                restore_agent_from='data/Prior.ckpt',
                data=None,
                voc=None,
                scoring_function='arb_vec',
                scoring_function_kwargs={},
                save_dir=None, learning_rate=0.0005,
                batch_size=128, n_steps=3000,
                num_processes=0, sigma=60,
                experience_replay=0):

    start_time = time.time()

    loader = DataLoader(data, batch_size=batch_size, shuffle=True, drop_last=True,
                        collate_fn=Dataset.collate_fn)

    print('Loading prior and agent...')
    Prior = RNN(voc, len(data[0][1]))
    Agent = RNN(voc, len(data[0][1]))

    logger = VizardLog('data/logs')

    # By default restore Agent to same model as Prior, but can restore from already trained Agent too.
    # Saved models are partially on the GPU, but if we dont have cuda enabled we can remap these
    # to the CPU.
    if torch.cuda.is_available():
        Prior.rnn.load_state_dict(torch.load('data/Prior.ckpt'))
        Agent.rnn.load_state_dict(torch.load(restore_agent_from))
    else:
        Prior.rnn.load_state_dict(torch.load('data/Prior.ckpt', map_location=lambda storage, loc: storage))
        Agent.rnn.load_state_dict(torch.load(restore_agent_from, map_location=lambda storage, loc: storage))

    # We dont need gradients with respect to Prior
    for param in Prior.rnn.parameters():
        param.requires_grad = False

    optimizer = torch.optim.Adam(Agent.rnn.parameters(), learning_rate)

    # Scoring_function
    scoring_function = get_scoring_function(scoring_function=scoring_function, num_processes=num_processes,
                                            **scoring_function_kwargs)

    # For policy based RL, we normally train on-policy and correct for the fact that more likely actions
    # occur more often (which means the agent can get biased towards them). Using experience replay is
    # therefor not as theoretically sound as it is for value based RL, but it seems to work well.
    experience = Experience(voc)

    # Log some network weights that can be dynamically plotted with the Vizard bokeh app
    logger.log(Agent.rnn.gru_2.weight_ih.cpu().data.numpy()[::100], "init_weight_GRU_layer_2_w_ih")
    logger.log(Agent.rnn.gru_2.weight_hh.cpu().data.numpy()[::100], "init_weight_GRU_layer_2_w_hh")
    logger.log(Agent.rnn.embedding.weight.cpu().data.numpy()[::30], "init_weight_GRU_embedding")
    logger.log(Agent.rnn.gru_2.bias_ih.cpu().data.numpy(), "init_weight_GRU_layer_2_b_ih")
    logger.log(Agent.rnn.gru_2.bias_hh.cpu().data.numpy(), "init_weight_GRU_layer_2_b_hh")

    # Information for the logger
    step_score = [[], []]

    print("Model initialized, starting training...")

    for step, (_, vec_batch) in tqdm(enumerate(loader), total=len(loader)):
        real_vecs = vec_batch.float()

        # Sample from Agent
        seqs, agent_likelihood, entropy = Agent.sample(batch_size, real_vecs)

        # Get prior likelihood and score
        prior_likelihood, _ = Prior.likelihood(Variable(seqs), real_vecs)
        smiles = seq_to_smiles(seqs, voc)
        gen_vecs = get_latent_vector(smiles)
        data = [[x, y] for x, y in zip(gen_vecs, real_vecs)]
        score = scoring_function(data)

        # Calculate augmented likelihood
        augmented_likelihood = prior_likelihood + sigma * Variable(score)
        loss = torch.pow((augmented_likelihood - agent_likelihood), 2)

        # Experience Replay
        # First sample
        if experience_replay and len(experience)>4:
            exp_seqs, exp_score, exp_prior_likelihood = experience.sample(4, real_vecs)
            exp_agent_likelihood, exp_entropy = Agent.likelihood(exp_seqs.long())
            exp_augmented_likelihood = exp_prior_likelihood + sigma * exp_score
            exp_loss = torch.pow((Variable(exp_augmented_likelihood) - exp_agent_likelihood), 2)
            loss = torch.cat((loss, exp_loss), 0)
            agent_likelihood = torch.cat((agent_likelihood, exp_agent_likelihood), 0)

        # Then add new experience
        prior_likelihood = prior_likelihood.data.cpu().numpy()
        new_experience = zip(smiles, score, prior_likelihood)
        experience.add_experience(new_experience)

        # Calculate loss
        loss = loss.mean()

        # Add regularizer that penalizes high likelihood for the entire sequence
        loss_p = - (1 / agent_likelihood).mean()
        loss += 5 * 1e3 * loss_p

        # Calculate gradients and make an update to the network weights
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Convert to numpy arrays so that we can print them
        augmented_likelihood = augmented_likelihood.data.cpu().numpy()
        agent_likelihood = agent_likelihood.data.cpu().numpy()

        # Print some information for this step
        time_elapsed = (time.time() - start_time) / 3600
        time_left = (time_elapsed * ((n_steps - step) / (step + 1)))
        print("\n       Step {}   Fraction valid SMILES: {:4.1f}  Time elapsed: {:.2f}h Time left: {:.2f}h".format(
              step, fraction_valid_smiles(smiles) * 100, time_elapsed, time_left))
        print("  Agent    Prior   Target   Score             SMILES")
        for i in range(math.ceil(batch_size/10)):
            print(" {:6.2f}   {:6.2f}  {:6.2f}  {:6.2f}     {}".format(agent_likelihood[i],
                                                                       prior_likelihood[i],
                                                                       augmented_likelihood[i],
                                                                       score[i],
                                                                       smiles[i]))
        # Need this for Vizard plotting
        step_score[0].append(step + 1)
        step_score[1].append(np.mean(score))

        # Log some weights
        logger.log(Agent.rnn.gru_2.weight_ih.cpu().data.numpy()[::100], "weight_GRU_layer_2_w_ih")
        logger.log(Agent.rnn.gru_2.weight_hh.cpu().data.numpy()[::100], "weight_GRU_layer_2_w_hh")
        logger.log(Agent.rnn.embedding.weight.cpu().data.numpy()[::30], "weight_GRU_embedding")
        logger.log(Agent.rnn.gru_2.bias_ih.cpu().data.numpy(), "weight_GRU_layer_2_b_ih")
        logger.log(Agent.rnn.gru_2.bias_hh.cpu().data.numpy(), "weight_GRU_layer_2_b_hh")
        logger.log("\n".join([smiles + "\t" + str(round(score, 2)) for smiles, score in zip \
                            (smiles[:12], score[:12])]), "SMILES", dtype="text", overwrite=True)
        logger.log(np.array(step_score), "Scores")

    # If the entire training finishes, we create a new folder where we save this python file
    # as well as some sampled sequences and the contents of the experience (which are the highest
    # scored sequences seen during training)
    if not save_dir:
        save_dir = 'data/results/run_' + time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime())
    os.makedirs(save_dir)
    copyfile('train_agent.py', os.path.join(save_dir, "train_agent.py"))

    experience.print_memory(os.path.join(save_dir, "memory"))
    torch.save(Agent.rnn.state_dict(), os.path.join(save_dir, 'Agent.ckpt'))

    seqs, agent_likelihood, entropy = Agent.sample(batch_size, real_vecs)
    prior_likelihood, _ = Prior.likelihood(Variable(seqs))
    prior_likelihood = prior_likelihood.data.cpu().numpy()
    smiles = seq_to_smiles(seqs, voc)
    score = scoring_function(smiles)
    with open(os.path.join(save_dir, "sampled"), 'w') as f:
        f.write("SMILES Score PriorLogP\n")
        for smiles, score, prior_likelihood in zip(smiles, score, prior_likelihood):
            f.write("{} {:5.2f} {:6.2f}\n".format(smiles, score, prior_likelihood))

if __name__ == "__main__":

    voc = Vocabulary(init_from_file="data/Voc")

    # Create a Dataset from a SMILES file
    if path.isfile('./data/agent/vecs.dat'):
        print('Found vectors, reading from file...')
        data = Dataset(voc, "data/mols.smi", vec_file='./data/agent/vecs.dat')
    else:
        raise ValueError('No agent Data')

    mew = data.mew
    std = data.std

    train_agent(data=data, voc=voc, scoring_function_kwargs={'mew':mew, 'std':std})


import numpy as np
import ntpath
import os
import re
import argparse
from typing import Tuple
import io
import time
import cv2
from PIL import Image, ImageOps
import tensorflow.compat.v1 as tf

QUANTIZE_RESULTS = "" # TODO: Set to path containing quantization results so post-quantization model can be loaded (see fpga_build directory)
SPLIT_GRAPH = True
JESTER_TEST = True
JESTER_PATH = ""# TODO: Set directly to jester video directory containing images (i.e. jester_data/vid_num/)
if JESTER_TEST:
    import tensorflow.contrib.decent_q

SOFTMAX_THRES = 0
HISTORY_LOGIT = True
REFINE_OUTPUT = True

def transform(frame: np.ndarray):
    # 480, 640, 3, 0 ~ 255
    frame = cv2.resize(frame, (224, 224))  # (224, 224, 3) 0 ~ 255
    frame = frame / 255.0  # (224, 224, 3) 0 ~ 1.0
    #frame = np.transpose(frame, axes=[2, 0, 1])  # (3, 224, 224) 0 ~ 1.0
    frame -= [0.485, 0.456, 0.406]
    frame /= [0.229, 0.224, 0.225]
    frame = np.expand_dims(frame, axis=0)  # (1, 3, 480, 640) 0 ~ 1.0
    return frame

catigories = [
    "Doing other things",  # 0
    "Drumming Fingers",  # 1
    "No gesture",  # 2
    "Pulling Hand In",  # 3
    "Pulling Two Fingers In",  # 4
    "Pushing Hand Away",  # 5
    "Pushing Two Fingers Away",  # 6
    "Rolling Hand Backward",  # 7
    "Rolling Hand Forward",  # 8
    "Shaking Hand",  # 9
    "Sliding Two Fingers Down",  # 10
    "Sliding Two Fingers Left",  # 11
    "Sliding Two Fingers Right",  # 12
    "Sliding Two Fingers Up",  # 13
    "Stop Sign",  # 14
    "Swiping Down",  # 15
    "Swiping Left",  # 16
    "Swiping Right",  # 17
    "Swiping Up",  # 18
    "Thumb Down",  # 19
    "Thumb Up",  # 20
    "Turning Hand Clockwise",  # 21
    "Turning Hand Counterclockwise",  # 22
    "Zooming In With Full Hand",  # 23
    "Zooming In With Two Fingers",  # 24
    "Zooming Out With Full Hand",  # 25
    "Zooming Out With Two Fingers"  # 26
]

n_still_frame = 0

def process_output(idx_, history):
    # idx_: the output of current frame
    # history: a list containing the history of predictions
    if not REFINE_OUTPUT:
        return idx_, history

    max_hist_len = 20  # max history buffer

    # mask out illegal action
    if idx_ in [7, 8, 21, 22, 3]:
        idx_ = history[-1]

    # use only single no action class
    if idx_ == 0:
        idx_ = 2

    # history smoothing
    if idx_ != history[-1]:
        if not (history[-1] == history[-2]): #  and history[-2] == history[-3]):
            idx_ = history[-1]

    history.append(idx_)
    history = history[-max_hist_len:]

    return history[-1], history

WINDOW_NAME = 'Video Gesture Recognition'
def run(target='cuda', print_log=True, callback_fn=lambda output: None):
    """
    target: {'cuda', 'opencl', 'llvm [-target={LLVM_TARGET}]'
        * -target implies cross compilation. LLVM_TARGET is one of llvm cross compilation targets. See: https://clang.llvm.org/docs/CrossCompilation.html
    """
    if not JESTER_TEST:
        print("Open camera...")
        cap = cv2.VideoCapture(0)

        print(cap)

        # set a lower resolution for speed up
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)

    # env variables
    full_screen = False
    cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME, 640, 480)
    cv2.moveWindow(WINDOW_NAME, 0, 0)
    cv2.setWindowTitle(WINDOW_NAME, WINDOW_NAME)

    t = None
    index = 0
    #print("Build transformer...")
    #transform = get_transform()

    buffer = [
        np.zeros((1, 56, 56, 3)),
        np.zeros((1, 28, 28, 4)),
        np.zeros((1, 28, 28, 4)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 12)),
        np.zeros((1, 14, 14, 12)),
        np.zeros((1, 7, 7, 20)),
        np.zeros((1, 7, 7, 20))
    ]
    idx = 0
    history = [2,2]
    history_logit = []
    history_timing = []

    i_frame = -1

    print("Ready!")
    # Lock to 25 FPS
    frame_time = 0
    tf.reset_default_graph()

    sessions = []

    p_buffer = (
        np.zeros((1, 56, 56, 3)),
        np.zeros((1, 28, 28, 4)),
        np.zeros((1, 28, 28, 4)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 8)),
        np.zeros((1, 14, 14, 12)),
        np.zeros((1, 14, 14, 12)),
        np.zeros((1, 7, 7, 20)),
        np.zeros((1, 7, 7, 20))
    )
    p_in_img = np.zeros((1, 224, 224, 3))

    if SPLIT_GRAPH:
        for i in range(11):
            model_name = f"model_tf_split_{i}"
            model_path = f"{QUANTIZE_RESULTS}/quantize_results_{i}/quantize_eval_model.pb"
            graph_def = tf.GraphDef.FromString(open(model_path, "rb").read())
            graph = tf.import_graph_def(graph_def, name = "")
            sessions.append(tf.Session(graph=graph))
    else:
        graph_def = tf.GraphDef.FromString(open("./model_tf/model_tf.pb", "rb").read())
        graph = tf.import_graph_def(graph_def, name = "")
        sessions.append(tf.Session(graph=graph))

    while True:
        i_frame += 1
        if JESTER_TEST:
            imgs = sorted(os.listdir(JESTER_PATH))
            if i_frame >= len(imgs):
                break
            img = np.asarray(Image.open(os.path.join(JESTER_PATH, imgs[i_frame])))
            img = img[:,:,::-1]
        else:
            _, img = cap.read()  # (480, 640, 3) 0 ~ 255
        if time.time() - frame_time > 1/25:
            frame_time = time.time()
            t1 = time.time()
            img_tran = transform(np.array(Image.fromarray(img).convert('RGB')))

            feat = None
            if SPLIT_GRAPH:
                resid_in = None
                shift_concat_in = None
                for i,sess in enumerate(sessions):
                    out_node = None
                    feed_dict = {}
                    if i == 0:
                        out_node = "MobilenetV2/expanded_conv_shift/prev_conv_output:0"
                        feed_dict["in_img:0"] = img_tran#np.ones((1,224,224,3))
                        print("img: ", img[0][0], " = ", img.shape)
                        print("img_tran: ", img_tran[0][0][0], " = ", img_tran.shape)
                    else:
                        out_node = f"MobilenetV2/expanded_conv_shift_{i}/prev_conv_output:0"
                        layer_name = f"expanded_conv_shift_{i-1}" if i > 1 else f"expanded_conv_shift"
                        feed_dict[f"MobilenetV2/{layer_name}/input:0"] = resid_in
                        feed_dict[f"MobilenetV2/{layer_name}/shift_concat_input:0"] = shift_concat_in

                    if i == len(sessions) - 1:
                        out_node = "MobilenetV2/Logits/output:0"

                    output = sess.run(out_node, feed_dict)

                    if i == len(sessions) - 1:
                        print(output[0][0])
                    else:
                        print(output[0][0][0][0])

                    # Do shift
                    if i < len(sessions) - 1:
                        resid_in = output
                        c = int(output.shape[3])
                        x1, x2 = output[:, :, :, :c//8], output[:, :, :, c//8:]
                        print("x1: ", x1[0][0][0][0])
                        print("x2: ", x2[0][0][0][0])
                        shift_concat_in = np.concatenate((buffer[i], x2), axis=3)
                        buffer[i] = x1

                    if i == len(sessions) - 1:
                        feat = output
            else:
                out_nodes = ['MobilenetV2/Logits/output:0', 'MobilenetV2/expanded_conv_shift/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_1/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_2/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_3/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_4/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_5/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_6/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_7/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_8/shift_split_buffer_output:0', 'MobilenetV2/expanded_conv_shift_9/shift_split_buffer_output:0']
                feed_dict = {f"shift_buffer_{i}:0": buffer[i] for i in range(10)}
                feed_dict["in_img:0"] = img_tran
                outputs = sessions[0].run(out_nodes, feed_dict)
                feat, buffer = outputs[0], outputs[1:]

            print(feat)
            if SOFTMAX_THRES > 0:
                feat_np = feat.reshape(-1)
                feat_np -= feat_np.max()
                softmax = np.exp(feat_np) / np.sum(np.exp(feat_np))
                if print_log:
                    print(max(softmax))
                if max(softmax) > SOFTMAX_THRES:
                    idx_ = np.argmax(feat, axis=1)[0]
                else:
                    idx_ = idx
            else:
                idx_ = np.argmax(feat, axis=1)[0]

            if HISTORY_LOGIT:
                history_logit.append(feat)
                history_logit = history_logit[-12:]
                avg_logit = sum(history_logit)
                idx_ = np.argmax(avg_logit, axis=1)[0]

            idx, history = process_output(idx_, history)
            callback_fn(idx)

            t2 = time.time()
            if print_log:
                print(f"{index} {catigories[idx]}")

            current_time = t2 - t1

        img = cv2.resize(img, (640, 480))
        img = img[:, ::-1]
        height, width, _ = img.shape
        label = np.zeros([height // 10, width, 3]).astype('uint8') + 255

        cv2.putText(label, 'Prediction: ' + catigories[idx],
                    (0, int(height / 16)),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 0, 0), 2)
        cv2.putText(label, '{:.1f} Vid/s'.format(1 / current_time),
                    (width - 170, int(height / 16)),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 0, 0), 2)

        img = np.concatenate((img, label), axis=0)
        cv2.imshow(WINDOW_NAME, img)

        key = cv2.waitKey(1)
        if key & 0xFF == ord('q') or key == 27:  # exit
            break
        elif key == ord('F') or key == ord('f'):  # full screen
            print('Changing full screen option!')
            full_screen = not full_screen
            if full_screen:
                print('Setting FS!!!')
                cv2.setWindowProperty(WINDOW_NAME, cv2.WND_PROP_FULLSCREEN,
                                      cv2.WINDOW_FULLSCREEN)
            else:
                cv2.setWindowProperty(WINDOW_NAME, cv2.WND_PROP_FULLSCREEN,
                                      cv2.WINDOW_NORMAL)

        if t is None:
            t = time.time()
        else:
            nt = time.time()
            index += 1
            t = nt

    if not JESTER_TEST:
        cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--target', default='cuda', nargs='?', help="Target for compilation. If set to 'llvm -target={LLVM_TARGET}', cross compilation is assumed. See run() function for additional details.")
    args = parser.parse_args()
    run(target=args.target)