import numpy as np
import unittest
import datetime
import pytest
from laserchicken import utils, test_tools, keys
from laserchicken.utils import fit_plane
from time import time

class TestUtils(unittest.TestCase):
    def test_GetPointCloudPoint(self):
        """ Should not raise exception. """
        pc = test_tools.generate_tiny_test_point_cloud()
        x, y, z = utils.get_point(pc, 1)
        self.assertEqual(2, x)
        self.assertEqual(3, y)
        self.assertEqual(4, z)

    def test_GetPointCloudPointFeature(self):
        """ Should not raise exception. """
        pc = test_tools.generate_tiny_test_point_cloud()
        cols = 0.5 * (pc[keys.point]["x"]["data"] + pc[keys.point]["y"]["data"])
        pc[keys.point]["color"] = {"type": "double", "data": cols}
        x, y, z = utils.get_point(pc, 1)
        c = utils.get_attribute_value(pc, 1, "color")
        self.assertEqual(c, 0.5 * (x + y))

    def test_GetPointCloudPointFeatures(self):
        """ Should not raise exception. """
        pc = test_tools.generate_tiny_test_point_cloud()
        cols = 0.5 * (pc[keys.point]["x"]["data"] + pc[keys.point]["y"]["data"])
        flavs = 0.5 * (pc[keys.point]["x"]["data"] - pc[keys.point]["y"]["data"])
        pc[keys.point]["color"] = {"type": "double", "data": cols}
        pc[keys.point]["flavor"] = {"type": "double", "data": flavs}
        x, y, z = utils.get_point(pc, 2)
        c, f = utils.get_features(pc, ("color", "flavor"), 2)
        self.assertEqual(c, 0.5 * (x + y))
        self.assertEqual(f, 0.5 * (x - y))

    def test_CopyEmptyPointCloud(self):
        """ Should not raise exception. """
        pc = test_tools.generate_tiny_test_point_cloud()
        pc[keys.point]["x"]["data"] = np.array([])
        pc[keys.point]["y"]["data"] = np.array([])
        pc[keys.point]["z"]["data"] = np.array([])

        copypc = utils.copy_point_cloud(pc)
        self.assertEqual(0, len(copypc[keys.point]["x"]["data"]))

    def test_CopyNonEmptyPointCloud(self):
        """ Test whether coordinates are copied """
        pc = test_tools.generate_tiny_test_point_cloud()
        x = pc[keys.point]["x"]["data"]
        y = pc[keys.point]["y"]["data"]
        z = pc[keys.point]["z"]["data"]

        copypc = utils.copy_point_cloud(pc)
        self.assertTrue(all(x == copypc[keys.point]["x"]["data"]))
        self.assertTrue(all(y == copypc[keys.point]["y"]["data"]))
        self.assertTrue(all(z == copypc[keys.point]["z"]["data"]))

    def test_CopyPointCloudMetaData(self):
        """ Test whether metadata are copied """
        pc = test_tools.generate_tiny_test_point_cloud()
        pc["log"] = [
            {"time": datetime.datetime(2018, 1, 23, 12, 15, 59), "module": "filter", "parameters": [("z", "gt", 0.5)]}]

        copypc = utils.copy_point_cloud(pc)
        self.assertEqual(datetime.datetime(2018, 1, 23, 12, 15, 59), copypc["log"][0]["time"])
        self.assertEqual("filter", copypc["log"][0]["module"])
        self.assertEqual([("z", "gt", 0.5)], copypc["log"][0]["parameters"])

    def test_CopyNonEmptyPointCloudBoolMask(self):
        """ Test whether coordinates are copied with boolean mask """
        pc = test_tools.generate_tiny_test_point_cloud()
        x = pc[keys.point]["x"]["data"][2]
        y = pc[keys.point]["y"]["data"][2]
        z = pc[keys.point]["z"]["data"][2]

        copypc = utils.copy_point_cloud(pc, array_mask=np.array([False, False, True]))
        self.assertTrue(all(np.array([x]) == copypc[keys.point]["x"]["data"]))
        self.assertTrue(all(np.array([y]) == copypc[keys.point]["y"]["data"]))
        self.assertTrue(all(np.array([z]) == copypc[keys.point]["z"]["data"]))

    def test_CopyNonEmptyPointCloudIntMask(self):
        """ Test whether coordinates are copied with array indexing """
        pc = test_tools.generate_tiny_test_point_cloud()
        x0, x1 = pc[keys.point]["x"]["data"][0], pc[keys.point]["x"]["data"][1]
        y0, y1 = pc[keys.point]["y"]["data"][0], pc[keys.point]["y"]["data"][1]
        z0, z1 = pc[keys.point]["z"]["data"][0], pc[keys.point]["z"]["data"][1]

        copypc = utils.copy_point_cloud(pc, array_mask=np.array([1, 0]))
        self.assertTrue(all(np.array([x1, x0]) == copypc[keys.point]["x"]["data"]))
        self.assertTrue(all(np.array([y1, y0]) == copypc[keys.point]["y"]["data"]))
        self.assertTrue(all(np.array([z1, z0]) == copypc[keys.point]["z"]["data"]))

    def test_AddMetaDataToPointCloud(self):
        """ Test adding info to the point cloud for test module """
        pc = test_tools.generate_tiny_test_point_cloud()
        from laserchicken import filter as somemodule
        utils.add_metadata(pc,somemodule,params = (0.5,"cylinder",4))
        self.assertEqual(len(pc[keys.provenance]),1)

    def test_AddToPointCloudEmpty(self):
        pc_1 = utils.create_point_cloud([],[],[])
        pc_2 = test_tools.generate_tiny_test_point_cloud()
        utils.add_to_point_cloud(pc_1, pc_2)
        for attr in pc_2[keys.point].keys():
            self.assertIn(attr, pc_1[keys.point])
            self.assertEqual(pc_1[keys.point][attr]['type'],
                             pc_2[keys.point][attr]['type'])
            self.assertTrue(all(pc_1[keys.point][attr]['data'] == pc_2[keys.point][attr]['data']))

    def test_AddToPointCloudInvalid(self):
        pc_1 = test_tools.SimpleTestData.get_point_cloud()
        # invalid format
        pc_2 = {}
        with pytest.raises(TypeError):
            utils.add_to_point_cloud(pc_1, pc_2)
        with pytest.raises(AttributeError):
            utils.add_to_point_cloud(pc_2, pc_1)
        # non-matching attributes
        test_data = test_tools.ComplexTestData()
        pc_2 = test_data.get_point_cloud()
        with pytest.raises(AttributeError):
            utils.add_to_point_cloud(pc_1, pc_2)
        # different structure
        pc_2 = {'vertex':{'x':1, 'y':2, 'z':3}}
        with pytest.raises(TypeError):
            utils.add_to_point_cloud(pc_1, pc_2)
        # different data types
        pc_2 = {'vertex': {'x': {'data': np.zeros(3, dtype=int), 'type': 'int'},
                           'y': {'data': np.zeros(3, dtype=int), 'type': 'int'},
                           'z': {'data': np.zeros(3, dtype=int), 'type': 'int'}}}
        with pytest.raises(ValueError):
            utils.add_to_point_cloud(pc_1, pc_2)

    def test_AddToPointCloud(self):
        test_data = test_tools.ComplexTestData()
        pc_source = test_data.get_point_cloud()
        pc_dest = utils.copy_point_cloud(pc_source)
        utils.add_to_point_cloud(pc_dest, pc_source)
        for key in pc_source.keys():
            self.assertIn(key, pc_dest)
        for attr in pc_source[keys.point].keys():
            self.assertEqual(len(pc_dest[keys.point][attr]['data']),
                             2*len(pc_source[keys.point][attr]['data']))
        self.assertEqual(pc_dest[keys.provenance][-1]['module'],
                         'laserchicken.utils')

    def test_AddFeatureArray(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = np.array([1, 1, 1, 1, 1], dtype=int)
        utils.update_feature(pc, 'test_feature', feature_add)
        self.assertIn('test_feature', pc[keys.point])
        self.assertTrue(all(pc[keys.point]['test_feature']['data'] == feature_add))

    def test_AddFeatureArrayInvalid(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = np.array([1, 1, 1, 1, 1, 2], dtype=int)
        with pytest.raises(AssertionError):
            utils.update_feature(pc, 'test_feature', feature_add)

    def test_AddFeatureArrayMask(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = np.array([1, 2, 3, 4], dtype=int)
        mask = np.array([1, 1, 0, 1, 1], dtype=bool)
        utils.update_feature(pc, 'test_feature', feature_add, array_mask=mask)
        self.assertIn('test_feature', pc[keys.point])
        self.assertTrue(all(pc[keys.point]['test_feature']['data'] == [1, 2, 0, 3, 4]))

    def test_AddFeatureArrayMaskInvalid(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = np.array([1, 2, 3, 4], dtype=int)
        mask = np.array([1, 1, 1, 1, 1], dtype=bool)
        with pytest.raises(AssertionError):
            utils.update_feature(pc, 'test_feature', feature_add, array_mask=mask)

    def test_AddFeatureValueMask(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = 1.1
        mask = np.array([1, 1, 0, 1, 1], dtype=bool)
        utils.update_feature(pc, 'test_feature', feature_add, array_mask=mask)
        self.assertIn('test_feature', pc[keys.point])
        self.assertTrue(all(pc[keys.point]['test_feature']['data'] == [1.1, 1.1, 0.0, 1.1, 1.1]))

    def test_AddFeatureValueMaskInvalid(self):
        test_data = test_tools.ComplexTestData()
        pc = test_data.get_point_cloud()
        feature_add = 1.1
        mask = np.array([1, 1, 0, 1, 1, 1], dtype=bool)
        with pytest.raises(AssertionError):
            utils.update_feature(pc, 'test_feature', feature_add, array_mask=mask)

class TestPlaneFit(unittest.TestCase):

    def test_leastsqr(self):
        # n_points = 100
        # points = np.zeros((n_points, 3))
        # for i in range(n_points):
        #     z = 5 + i % np.sqrt(n_points)
        #     points[i] = np.array(((i % np.sqrt(n_points)), (np.floor(i / np.sqrt(n_points))), z))
        #t0 = time()
        f = fit_plane(self.points[:, 0], self.points[:, 1], self.points[:, 2])
        #print('LSQR : %f' %(time()-t0))
        estimates = f(self.points[:, 0], self.points[:, 1])
        np.testing.assert_allclose(estimates, self.points[:, 2])

    def test_FitPlaneSVD(self):
        """Test the normal vector extraction with SVD."""
        #t0 = time()
        nfit = utils.fit_plane_svd(self.points[:,0], self.points[:,1], self.points[:,2])
        #print('SVD  : %f' %(time()-t0))
        self.assertTrue(np.allclose(nfit, self.n))

    def generate_random_points_inplane(self,nvect, dparam=0, npts=100, eps=0.0):
        """
        Generate a series of point all belonging to a plane.

        :param nvect: normal vector of the plane
        :param dparam: zero point value of the plane
        :param npts: number of points
        :param eps: std of the gaussian noise added to the z values of the planes
        :return: x,y,z coordinate of the points
        """
        a, b, c = nvect / np.linalg.norm(nvect)
        x, y = np.random.rand(npts), np.random.rand(npts)
        z = (dparam - a * x - b * y) / c
        if eps > 0:
            z += np.random.normal(loc=0., scale=eps, size=npts)
        return np.column_stack((x, y, z))

    def setUp(self):
        """Set up the data points."""

        self.n = np.array([1., 2., 3.])
        self.n /= np.linalg.norm(self.n)
        self.points = self.generate_random_points_inplane(self.n,eps=0)

    def tearDown(self):
        """Tear it down."""
        pass

if __name__ == '__main__':
    unittest.main()


# 이 풀이는 bfs를 끝까지 진행한 풀이.. / s초 후라는 것을 고려하지못함..

from collections import deque
n, k = map(int, input().split())
graph = []
for i in range(n):
    graph.append(list(map(int, input().split())))
s, a, b = map(int, input().split())
for i in range(n):
    for j in range(n):
        if graph[i][j] != 0:
            graph[i][j] = (graph[i][j],0)
print(graph)

dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]

def bfs(graph, x, y):

    q = deque()
    q.append((x, y))
    while q:
        x, y = q.popleft()
        num, time = graph[x][y]
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if nx >= 0 and ny >= 0 and nx < n and ny < n:
                if graph[nx][ny] == 0:
                    graph[nx][ny] = (num, time + 1)
                    q.append((nx, ny))
                elif graph[nx][ny][1] > time + 1:
                    graph[nx][ny] = (num, time + 1)
                    q.append((nx, ny))
                elif graph[nx][ny][1] == time + 1 and graph[nx][ny][0] > num:
                    graph[nx][ny] = (num, time + 1)
                    q.append((nx, ny))

for i in range(n):
    for j in range(n):
        bfs(graph,i,j)
print(graph)
print(graph[a-1][b-1][0])


import json
import numpy as np
from sentence_transformers import SentenceTransformer

from scipy.spatial.distance import cdist
import streamlit as st
import random
class Model():
    def process_intents(self,intent):
        for key in intent:
            if type(intent[key])!=type([]):
                intent[key]=[intent[key]]
        return intent
    def __init__(self,intents):
        self.intents=self.process_intents(json.load(open(intents)))
        self.model=SentenceTransformer('all-MiniLM-L6-v2')
    def train(self):
        print("Generating features this may take a few moments depending on how big the intents file is")
        self.encodings=np.array(list(map(self.model.encode,list(self.intents.keys()))))
        print("features were successfully generated")

    def predict(self,input):
        encoding=self.model.encode(input)
        try:
            scores=1-cdist(self.encodings,encoding[np.newaxis,...])
        except:
            print("model was not trained")
            return
        return random.choice(list(self.intents.values())[np.argmax(scores)]),np.max(scores)
@st.cache(suppress_st_warning=True)
def train():
    m=Model("intents.json")
    m.train()
    return m
model=train()
def predict(x):
    return model.predict(x)[0]

import requests
#set title in streamlit
from googlesearch import search
import streamlit.components.v1 as components
import random
from quote import quote
from random_word import RandomWords
r = RandomWords()
@st.cache(suppress_st_warning=True,allow_output_mutation=True)
def set_counter():
    print("resetted")
    return {"counter":0}
l=set_counter()

@st.cache(suppress_st_warning=True,allow_output_mutation=True)
def set_default_q():
    print("resettedq")
    return {"previous":"-"}
prev=set_default_q()
def wrap_by_word(s, n):
    '''returns a string where \n is inserted between every n words'''
    a = s.split()
    ret = ''
    for i in range(0, len(a), n):
        ret += ' '.join(a[i:i+n]) + '\n'

    return ret
def get_random_quote():
    result=quote(r.get_random_word(),1)[0]
    return f"{wrap_by_word(result['quote'],7)} by {result['author']}"

st.title('AI ChatBot')
def check(url):
    if "jpg" in url or "png" in url:
        return False
    headers=requests.head(url).headers
    if "X-Frame-Options" in headers:
        if headers["X-Frame-Options"]=="DENY" or headers["X-Frame-Options"]=="SAMEORIGIN":
            return False
        return True
    return True

text = st.text_input('Enter your text here')

st.text("you : {}".format(text))
result=predict(text)
if text!=prev["previous"]:
    l["counter"]=0
prev["previous"]=str(text)
if result=="query":

  urls=list(search(str(text), num=20, stop=10, pause=2))

  nex=st.button("next")
  f=True

  if nex or f:

      for _ in range(19-l["counter"]):
          if check(urls[l["counter"]]):
              print(l["counter"])
              URL=urls[l["counter"]]
              l["counter"]+=1
              break
          l["counter"]+=1

      components.iframe(URL,height=10000,width=700)
      f=False

#write a function to get the weather
elif result=="<action joke>":
    st.text("bot : {}".format(requests.get("https://v2.jokeapi.dev/joke/Any?blacklistFlags=nsfw,religious,political,racist,sexist,explicit&format=txt").text))
elif result=="<random quote>":
    while True:
        try:
            q=get_random_quote()
            break
        except:
            pass
    st.text("bot : {}".format(get_random_quote()))
elif result=="<action weather>":
    components.iframe('https://www.ventusky.com/',height=1000,width=3000)
elif result=="<action news>":
    components.iframe('https://globalnews.ca/',height=10000,width=700)

else:

  st.text("bot : {}".format(result))


import matplotlib.pyplot as plt
import numpy as np
import torch
import time

class Regression(torch.nn.Module):

    def __init__(self, c_np_array):
        super(Regression, self).__init__()

        self.num = c_np_array.size

        self.c_array = torch.from_numpy(c_np_array).float()
        self.p_array = torch.nn.Parameter(torch.zeros(self.num))

        for p in self.parameters():
            torch.nn.init.uniform_(p, -1.0, 1.0)

    def forward(self, x):

        y = self.c_array[0] * torch.pow(x, self.p_array[0])

        for i in range(self.num-1):
            y = y + self.c_array[i+1] * torch.pow(x, self.p_array[i+1])

        return y

    def get_params(self):

        for name, param in self.named_parameters():
            if name == 'p_array':
                p_np_array = param.data.detach().numpy()

        return p_np_array

    def print_params(self):

        p_np_array = self.get_params()

        print('p: ', p_np_array)

    def set_params(self, p_np_array):

        with torch.no_grad():
            sd = self.state_dict()
            sd['p_array'] = torch.from_numpy(p_np_array).float()
            self.load_state_dict(sd)

    def enforce_weight_integer(self, integer_index):

        with torch.no_grad():
            sd = self.state_dict()
            p_array = sd['p_array']

            for i in integer_index:
                p_array[i] = torch.round(p_array[i])

            sd['p_array'] = p_array
            self.load_state_dict(sd)

    def get_weight_integer_index(self, threshold):

        index = []

        with torch.no_grad():
            for i in range(self.num):
                if torch.abs(self.p_array[i] - torch.round(self.p_array[i])) < threshold:
                    index.append(i)

        return index

    def train(self, x, y, is_freeze_integer=False, total_epoch=1000000, lr=1e-4, integer_threshold=0.05, loss_params_file_name='loss_params_history.npy'):

        criterion = torch.nn.MSELoss()
        optimizer = torch.optim.Adam(self.parameters(), lr=lr)

        loss_history = np.zeros(total_epoch)
        p_history = np.zeros((total_epoch, self.num))

        time_start = time.time()
        integer_index = []
        for epoch in range(total_epoch):

            optimizer.zero_grad()
            loss = criterion(self(x), y)

            loss_history[epoch] = loss.item()
            p_history[epoch,:] = self.get_params()

            loss.backward()
            optimizer.step()

            self.enforce_weight_integer(integer_index)

            if epoch%10000 == 0:
                time_end = time.time()
                print("epoch: %d; MSE: %.2e; elapsed time: %.2f sec"%(epoch, loss_history[epoch], time_end-time_start))
                self.print_params()

            if (is_freeze_integer) and (epoch%10000 == 0) and (epoch>100) and (loss_history[epoch]>=loss_history[epoch-100]):
                integer_index = self.get_weight_integer_index(integer_threshold)
                print("integer_index: ", integer_index)

        # save loss and parameters history at the end of iteration
        with open(loss_params_file_name, 'wb') as f:
            np.savez(f, loss_history=loss_history, p_history=p_history)

    def read_loss_params(self, file_name='loss_params_history.npy'):

        with open(file_name, 'rb') as f:
            data = np.load(f)
            loss_history = data['loss_history']
            p_history = data['p_history']
            epoch_array = np.arange(loss_history.size)

        return epoch_array, loss_history, p_history

    # compute loss as a 1D function of p1 with fixed p2
    def compute_loss_1d(self, data_x, data_y, p1_array, p2_val):

        np1 = p1_array.size
        loss_1d = np.zeros(np1)

        criterion = torch.nn.MSELoss()

        iip1 = 0
        for p1 in p1_array:
            self.set_params(np.array([p1, p2_val]))
            loss = criterion(self(data_x), data_y)
            loss_1d[iip1] = loss.item()
            iip1 += 1
        return loss_1d

    # compute loss as a 2D function of p1 and p2
    def compute_loss_2d(self, data_x, data_y, p1_arr_1d, p2_arr_1d):

        np1, np2 = p1_arr_1d.size, p2_arr_1d.size
        p1_arr_2d, p2_arr_2d = np.meshgrid(p1_arr_1d, p2_arr_1d)
        loss_2d = np.zeros((np2, np1))

        criterion = torch.nn.MSELoss()

        iip2 = 0
        for p2 in p2_arr_1d:
            iip1 = 0
            for p1 in p1_arr_1d:
                self.set_params(np.array([p1, p2]))
                l = criterion(self(data_x), data_y)
                loss_2d[iip2, iip1] = l.item()
                iip1 += 1
            iip2 += 1

        return p1_arr_2d, p2_arr_2d, loss_2d

if __name__ == "__main__":

    # set random seed to make results reproducible
    rand_seed_num = 0
    torch.manual_seed(rand_seed_num)

    # generate training data
    Nx = 1000
    data_x = torch.rand(Nx, 1)

    c_array = np.array([1.0, 1e2])      # c1, c2 = 1.0, 1e2
    p_array = np.array([1.0, 2.0])      # p1, p2 = 1, 2

    data_y = torch.zeros_like(data_x)
    for i in range(c_array.size):
        data_y += c_array[i] * torch.pow(data_x, p_array[i])

    reg = Regression(c_array)
    reg.train(data_x, data_y, is_freeze_integer=True, lr=1e-4, integer_threshold=0.05, total_epoch=200000)


import matplotlib as mpl

mpl.use("TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

def plot_solution(df: pd.DataFrame, show=True):
    fig, ax = plt.subplots(2, 2, figsize=(10, 4))

    for sub_ax in ax[-1]:
        sub_ax.set_xlabel("Time $t$")

    ax[0, 0].set_ylabel("$r(t)$")
    ax[1, 0].set_ylabel("$v(t)$")

    for sub_ax in ax[:, 0]:
        sub_ax.set_xlim(-10, 40)

    for sub_ax in ax[:, 1]:
        sub_ax.set_xlim(-10, 80)

    tv = df.attrs["tv"]
    xv = df.attrs["xv"]
    p = df.attrs["p"]

    print(p)
    print(xv[0, 0])

    for k in range(2):
        for j in range(2):
            ax[k, j].plot(tv, xv[:, k, j])

    plt.tight_layout()
    if (show):
        plt.show()

    return fig, ax

if (__name__ == "__main__"):
    df = pd.read_pickle("data/result.pkl")

    plot_solution(df)


import os
import sys
import uuid
import json
import hashlib
import copy
import re
import tarfile
from datetime import date
from argparse import ArgumentParser

variant_type_to_data_type_etc = {
    'snv': ['Simple Nucleotide Variation', ['Raw SNV Calls', None], ['CaVEMan'], ['GATK:Mutect2']],   # dataCategory, [dataType, data_subtypes], analysis_tools
    'indel': ['Simple Nucleotide Variation', ['Raw InDel Calls', None], ['Pindel'], ['GATK:Mutect2']],
    'cnv': ['Copy Number Variation', ['Raw CNV Calls', None], ['ASCAT']],
    'sv': ['Structural Variation', ['Raw SV Calls', None], ['BRASS']],
    'caveman-supplement': ['Simple Nucleotide Variation', ['SNV Supplement', None], ['CaVEMan']],
    'pindel-supplement': ['Simple Nucleotide Variation', ['InDel Supplement', None], ['Pindel']],
    'ascat-supplement': ['Copy Number Variation', ['CNV Supplement', None], ['ASCAT']],
    'brass-supplement': ['Structural Variation', ['SV Supplement', None], ['BRASS']],
    'timings-supplement': ['Quality Control Metrics', ['Analysis QC', ['Runtime Stats']], None, None],
    'bas_metrics': ['Quality Control Metrics', ['Aligned Reads QC', ['Alignment Metrics']], ['Sanger:bam_stats']],
    'contamination_metrics': ['Quality Control Metrics', ['Sample QC', ['Cross Sample Contamination']], ['Sanger:verifyBamHomChk'], ['GATK:CalculateContamination']],
    'ascat_metrics': ['Quality Control Metrics', ['Analysis QC', ['Ploidy', 'Tumour Purity']], ['ASCAT']],
    'genotyped_gender_metrics': ['Quality Control Metrics', ['Analysis QC', ['Genotyping Stats']], ['Sanger:compareBamGenotypes']],
    'mutect_filtering_metrics': ['Quality Control Metrics', ['Analysis QC', ['Variant Filtering Stats']], [], ['GATK:FilterMutectCalls']],
    'mutect_callable_metrics': ['Quality Control Metrics', ['Analysis QC', ['Variant Callable Stats']], [], ['GATK:Mutect2']],
}

workflow_full_name = {
    'sanger-wgs-variant-calling': 'Sanger WGS Variant Calling',
    'sanger-wxs-variant-calling': 'Sanger WXS Variant Calling',
    'gatk-mutect2-variant-calling': 'GATK Mutect2 Variant Calling'
}

def calculate_size(file_path):
    return os.stat(file_path).st_size

def calculate_md5(file_path):
    md5 = hashlib.md5()
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b''):
            md5.update(chunk)
    return md5.hexdigest()

def get_files_info(file_to_upload, wf_short_name,  wf_version, somatic_or_germline, normal_analysis, tumour_analysis, date_str):
    file_info = {
        'fileType': 'VCF' if file_to_upload.endswith('.vcf.gz') else file_to_upload.split(".")[-1].upper(),
        'fileSize': calculate_size(file_to_upload),
        'fileMd5sum': calculate_md5(file_to_upload),
        'fileAccess': 'controlled',
        'info': {}
    }

    if somatic_or_germline == 'Somatic':
        metadata = tumour_analysis
    elif somatic_or_germline == 'Germline':
        metadata = normal_analysis
    else:
        pass  # should never happen

    experimental_strategy = metadata['experiment']['experimental_strategy'].lower() if 'experimental_strategy' in metadata['experiment'] else metadata['experiment']['library_strategy'].lower()

    variant_type = ''
    if wf_short_name in (['sanger-wgs', 'sanger-wxs']):
        fname_sample_part = metadata['samples'][0]['sampleId']
        if file_to_upload.endswith('.flagged.muts.vcf.gz') or file_to_upload.endswith('.flagged.muts.vcf.gz.tbi'):
            variant_type = 'snv'
        elif file_to_upload.endswith('.flagged.vcf.gz') or file_to_upload.endswith('.flagged.vcf.gz.tbi'):
            variant_type = 'indel'
        elif file_to_upload.endswith('.copynumber.caveman.vcf.gz') or file_to_upload.endswith('.copynumber.caveman.vcf.gz.tbi'):
            variant_type = 'cnv'
        elif file_to_upload.endswith('.annot.vcf.gz') or file_to_upload.endswith('.annot.vcf.gz.tbi'):
            variant_type = 'sv'
        elif file_to_upload.endswith('-supplement.tgz'):
            variant_type = file_to_upload.split(".")[-2]
        elif file_to_upload.endswith('_metrics.tgz'):
            variant_type = file_to_upload.split(".")[-2]
            if re.match(r'.+?\.normal.contamination_metrics.tgz', file_to_upload) \
                    or (re.match(r'.+?\.normal.bas_metrics.tgz', file_to_upload)):
                fname_sample_part = normal_analysis['samples'][0]['sampleId']
        else:
            sys.exit('Error: unknown file type "%s"' % file_to_upload)

    elif wf_short_name in (['gatk-mutect2']):
        fname_sample_part = metadata['samples'][0]['sampleId']
        if file_to_upload.endswith('mutect2-snv.vcf.gz') or file_to_upload.endswith('mutect2-snv.vcf.gz.tbi'):
            variant_type = 'snv'
        elif file_to_upload.endswith('mutect2-indel.vcf.gz') or file_to_upload.endswith('mutect2-indel.vcf.gz.tbi'):
            variant_type = 'indel'
        elif file_to_upload.endswith('_metrics.tgz'):
            if file_to_upload.endswith('.contamination_metrics.tgz'):
                variant_type = file_to_upload.split(".")[-2]
                if re.match(r'.+?\.normal.contamination_metrics.tgz', file_to_upload):
                    fname_sample_part = normal_analysis['samples'][0]['sampleId']
            elif file_to_upload.endswith('.filtering_metrics.tgz'):
                variant_type = 'mutect_filtering_metrics'
            elif file_to_upload.endswith('.callable_metrics.tgz'):
                variant_type = 'mutect_callable_metrics'
        else:
            sys.exit('Error: unknown file type "%s"' % file_to_upload)

    elif wf_short_name in (['HaplotypeCaller']):
        sys.exit('Error: not implemented yet for "%s"' % wf_short_name)

    else:
        sys.exit('Error: unknown variant calling workflow "%s"' % wf_short_name)

    # file naming patterns:
    #   pattern:  <argo_sample_id_tumour>.[wgs|wxs].<date>.<wf_short_name>.[somatic|germline].[snv|indel|cnv|sv].vcf.gz
    #   example: TEST-PR.DO250183.SA610229.wxs.20200319.sanger-wxs.somatic.snv.vcf.gz
    new_fname = '.'.join([
                            metadata['studyId'],
                            metadata['samples'][0]['donor']['donorId'],
                            fname_sample_part,
                            experimental_strategy,
                            date_str,
                            wf_short_name,
                            somatic_or_germline.lower(),
                            variant_type,
                            'vcf.gz' if variant_type in ['snv', 'indel', 'cnv', 'sv'] else 'tgz'
                        ] + (['tbi'] if file_to_upload.endswith('.tbi') else []))

    file_info['fileName'] = new_fname

    file_info['info'] = {
        'data_category': variant_type_to_data_type_etc[variant_type][0],
        'data_subtypes': None
    }

    extra_info = {}
    if new_fname.endswith('.vcf.gz'):
        file_info['dataType'] = variant_type_to_data_type_etc[variant_type][1][0]
        file_info['info']['data_subtypes'] = variant_type_to_data_type_etc[variant_type][1][1]
    elif new_fname.endswith('.vcf.gz.tbi'):
        file_info['dataType'] = 'VCF Index'
    elif new_fname.endswith('.tgz'):
        if new_fname.endswith('-supplement.tgz'):
            file_info['dataType'] = variant_type_to_data_type_etc[variant_type][1][0]
            file_info['info']['data_subtypes'] = variant_type_to_data_type_etc[variant_type][1][1]
        elif new_fname.endswith('_metrics.tgz'):
            file_info['dataType'] = variant_type_to_data_type_etc[variant_type][1][0]
            file_info['info']['data_subtypes'] = variant_type_to_data_type_etc[variant_type][1][1]
        else:
            sys.exit('Error: unknown file type "%s"' % file_to_upload)

        ## Disable the population of qc_metrics into payload to avoid the exposure of sensitive info
        tar = tarfile.open(file_to_upload)
        for member in tar.getmembers():
            if member.name.endswith('.extra_info.json'):
                f = tar.extractfile(member)
                extra_info = json.load(f)
                extra_info.pop('metrics', None)
                break
    else:
        sys.exit('Error: unknown file type "%s"' % file_to_upload)

    if wf_short_name in (['sanger-wgs', 'sanger-wxs']):
        file_info['info']['analysis_tools'] = variant_type_to_data_type_etc[variant_type][2]
    elif wf_short_name in (['gatk-mutect2']):
        file_info['info']['analysis_tools'] = variant_type_to_data_type_etc[variant_type][3]

    if extra_info:
        file_info['info'].update(extra_info)

    new_dir = 'out'
    try:
        os.mkdir(new_dir)
    except FileExistsError:
        pass

    dst = os.path.join(os.getcwd(), new_dir, new_fname)
    os.symlink(os.path.abspath(file_to_upload), dst)

    return file_info

def get_sample_info(sample_list):
    samples = copy.deepcopy(sample_list)
    for sample in samples:
        for item in ['info', 'sampleId', 'specimenId', 'donorId', 'studyId']:
            sample.pop(item, None)
            sample['specimen'].pop(item, None)
            sample['donor'].pop(item, None)

    return samples

def main(args):
    normal_analysis = {}
    normal_origin = None
    with open(args.normal_analysis, 'r') as f:
        normal_analysis = json.load(f)
        if 'info' in normal_analysis and 'origin' in normal_analysis['info']:
            normal_origin = normal_analysis['info']['origin']

    tumour_analysis = {}
    tumour_origin = None
    if args.tumour_analysis:
        with open(args.tumour_analysis, 'r') as f:
            tumour_analysis = json.load(f)
            if 'info' in tumour_analysis and 'origin' in tumour_analysis['info']:
                tumour_origin = tumour_analysis['info']['origin']

    if normal_origin is not None and tumour_origin is not None:
        if normal_origin == tumour_origin:
            analysis_origin = normal_origin
        else:
            sys.exit(f'Origins of normal and tumour analyses differ, normal: {normal_origin}, tumour: {tumour_origin}')
    elif normal_origin is not None:
        analysis_origin = normal_origin
    elif tumour_origin is not None:
        analysis_origin = tumour_origin
    else:
        analysis_origin = None

    somatic_or_germline = 'Somatic'   # default
    if args.wf_short_name in ['sanger-wgs', 'sanger-wxs', 'gatk-mutect2']:
        if not tumour_analysis:
            sys.exit('Error: metadata for tumour is missing!')
    elif args.wf_short_name in ['HaplotypeCaller']:
        somatic_or_germline = 'Germline'
    else:
        sys.exit("Unsupported variant caller: %s" % args.wf_short_name)

    payload = {
        'analysisType': {
            'name': None
        },
        'info': {
            'origin': analysis_origin  # adding this here, so the info field will maintain it's position
        },
        'studyId': normal_analysis.get('studyId'),  # normal/tumour analysis should always from the same study
        'experiment': {},
        'samples': [],
        'files': [],
        'workflow': {
            'workflow_name': workflow_full_name[args.wf_name],
            'workflow_short_name': args.wf_short_name,
            'workflow_version': args.wf_version,
            'run_id': args.wf_run,
            'session_id': args.wf_session,
            'inputs': [],
            'genome_build': 'GRCh38_hla_decoy_ebv'
        }
    }

    if analysis_origin is None:  # if origin is None, remove the info field
        payload.pop('info')

    # get sample of the payload
    if somatic_or_germline == 'Somatic':  # somatic variants
        payload['samples'] = get_sample_info(tumour_analysis.get('samples'))
        payload['workflow']['inputs'] = [
            {
                "tumour_analysis_id": tumour_analysis.get("analysisId"),
                "analysis_type": "sequencing_alignment"
            },
            {
                "normal_analysis_id": normal_analysis.get("analysisId"),
                "analysis_type": "sequencing_alignment"
            }
        ]
        payload['experiment'] = {
            'experimental_strategy': tumour_analysis['experiment']['experimental_strategy'] if 'experimental_strategy' in tumour_analysis['experiment'] else tumour_analysis['experiment']['library_strategy'],
            'platform': tumour_analysis['experiment']['platform']
        }
    else:   # germline variants
        payload['samples'] = get_sample_info(normal_analysis.get('samples'))
        payload['workflow']['inputs'] = [
            {
                "normal_analysis_id": normal_analysis.get("analysisId"),
                "analysis_type": "sequencing_alignment"
            }
        ]
        payload['experiment'] = {
            'experimental_strategy': normal_analysis['experiment']['experimental_strategy'] if 'experimental_strategy' in normal_analysis['experiment'] else normal_analysis['experiment']['library_strategy'],
            'platform': normal_analysis['experiment']['platform']
        }

    analysis_type = 'variant_calling'
    date_str = date.today().strftime("%Y%m%d")
    for f in args.files_to_upload:
        if f.endswith('-supplement.tgz'): analysis_type = 'variant_calling_supplement'
        if f.endswith('_metrics.tgz'): analysis_type = 'qc_metrics'

        file_info = get_files_info(f, args.wf_short_name, args.wf_version, somatic_or_germline, normal_analysis, tumour_analysis, date_str)

        payload['files'].append(file_info)

    payload['analysisType']['name'] = analysis_type

    if not analysis_type == "qc_metrics":
        payload['variant_class'] = somatic_or_germline

    with open("%s.%s.payload.json" % (str(uuid.uuid4()), analysis_type), 'w') as f:
        f.write(json.dumps(payload, indent=2))

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-n", dest="normal_analysis", required=True,
                        help="json file containing sequencing_alignment SONG analysis for normal sample")
    parser.add_argument("-t", dest="tumour_analysis", required=False,
                        help="json file containing sequencing_alignment SONG analysis for tumour sample")
    parser.add_argument("-f", dest="files_to_upload", type=str, nargs="+", help="Files to be uploaded", required=True)
    parser.add_argument("-w", dest="wf_name", type=str, help="workflow full name", required=True)
    parser.add_argument("-s", dest="wf_short_name", type=str, help="workflow short name", required=True)
    parser.add_argument("-v", dest="wf_version", type=str, required=True, help="workflow version")
    parser.add_argument("-r", dest="wf_run", type=str, required=True, help="workflow run ID")
    parser.add_argument("-j", dest="wf_session", type=str, required=True, help="workflow session ID")
    args = parser.parse_args()

    main(args)


from typing import List, Mapping, Tuple

import torch
import torchdata.datapipes as dp
import typer
from loguru import logger
from torch import Tensor
from torch.nn import Module
from torchdata.dataloader2 import DataLoader2, MultiProcessingReadingService
import pandas as pd

from ..ops.functional.proj import spherical
from ..utils import config
from ..ops import detector, graph, matcher, solver
from ..ops.dp import ImageSequenceDataPipe, VideoDataPipe
from websockets.exceptions import ConnectionClosed
from websockets.client import connect  # type: ignore
from websockets.legacy.client import WebSocketClientProtocol
from queue import Queue, Empty
import asyncio
import json
from multiprocessing.managers import Namespace
from concurrent.futures import ThreadPoolExecutor
import matplotlib.pyplot as plt
from ..ops.functional.metric import sampson

def build_socket_event(
    keyframe_locs: Tensor, keyframe_kpts: Tensor, keyframe_descs: Tensor, keyframe_R: Tensor, keyframe_t: Tensor
) -> dict:
    event = {
        'type': 'processor-update-keyframes',
        'data': {
            'keyframe_locs': keyframe_locs.tolist(),
            'keyframe_kpts': keyframe_kpts.tolist(),
            'keyframe_descs': keyframe_descs.tolist(),
            'keyframe_R': keyframe_R.tolist(),
            'keyframe_t': keyframe_t.tolist(),
        },
    }
    return event

async def _update_database(websocket: WebSocketClientProtocol, data: Tuple[Tensor]):
    typer.secho('Updating database', fg=typer.colors.GREEN)
    event = build_socket_event(*data)  # type: ignore
    await websocket.send(json.dumps(event))
    await asyncio.sleep(1)

def empty_queue(queue: Queue):
    while True:
        try:
            queue.get(block=True, timeout=1)
        except Empty:
            queue.join()
            break

async def _socket_handler(queue: Queue):
    async for websocket in connect(f'ws://{config.host}:{config.server_port}', max_size=2**20 * 128):
        typer.secho('Processor\'s websocket connected', fg=typer.colors.GREEN)
        while True:
            try:
                data = queue.get(block=False)
                await _update_database(websocket, data)
                queue.task_done()
                await asyncio.sleep(1)
            except Empty:
                await asyncio.sleep(1)
            except ConnectionClosed:
                await asyncio.sleep(1)
            except Exception as e:
                logger.error(e)
                empty_queue(queue)
                raise e
            except KeyboardInterrupt:
                empty_queue(queue)
                raise KeyboardInterrupt
            except SystemExit:
                empty_queue(queue)
                raise SystemExit

def socket_thread(queue: Queue):
    asyncio.run(_socket_handler(queue))

def collate_fn(batch: List[Tensor]) -> Tensor:
    return torch.stack(batch)

def simple_visualize(track: Tensor, kf_locs: Tensor):
    track = track.detach().cpu()
    kf_locs = kf_locs.detach().cpu()
    assert track.ndim == 2
    assert track.shape[-1] == 3

    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111)
    ax.set_title('Track')
    ax.plot(track[:, 0], track[:, 1], alpha=0.5, c='gray')
    ax.set_aspect('equal')
    ax.set_box_aspect(1)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.scatter(track[-1, 0], track[-1, 1], c='r', s=300)
    ax.scatter(kf_locs[:, 0], kf_locs[:, 1], c='g', s=100)
    fig.savefig('track.png')
    torch.save(track, 'track.pt')
    plt.close(fig)

def run(ipc: Namespace | None = None):
    logger.debug('Processor process started')
    queue: Queue[Tuple[pd.DataFrame]] = Queue(maxsize=10)
    if config.processor_solver_only:
        logger.debug('Processor solver only')
        processor()
        return
    with ThreadPoolExecutor(max_workers=2, thread_name_prefix='TorchSLAM-Processor') as executor:
        try:
            logger.debug('Spawning socket thread')
            executor.submit(socket_thread, queue)
            logger.debug('Spawning processor thread')
            executor.submit(processor, queue)
        except KeyboardInterrupt as e:
            executor.shutdown(wait=True, cancel_futures=True)
            raise e
        except SystemExit as e:
            executor.shutdown(wait=True, cancel_futures=True)
            raise e
        except Exception as e:
            logger.error(e)
            executor.shutdown(wait=True, cancel_futures=True)
            raise e
    logger.debug('Processor finished')

def processor(queue: Queue | None = None):
    typer.secho('Processor started', fg=typer.colors.GREEN)
    if config.input_path is None:
        raise ValueError('Input path is not specified')
    if config.input_type == 'image_sequence':
        reader: dp.iter.IterDataPipe = ImageSequenceDataPipe(
            config.input_path, config.to_hw, config.reverse_sort_filenames
        )
    elif config.input_type == 'video':
        reader = VideoDataPipe(config.input_path, config.to_hw)
    else:
        raise ValueError(f'Unknown input type: {config.input_type}')
    logger.debug('Initailizing models')
    n_f = len(reader)
    det_model = detector.SIFT(
        feature_dim=config.feature_dim,
        num_features=config.num_features,
        patch_size=config.patch_size,
        angle_bins=config.angle_bins,
        spatial_bins=config.spatial_bins,
        scale_n_levels=config.scale_n_levels,
        sigma=config.sigma,
        root_sift=config.root_sift,
        double_image=config.double_image,
    )
    det_model.to(config.device)
    sol = solver.Fundamental(config.camera_model)
    mtc: Module = matcher.RANSACMatcher(
        sol,
        sampson,
    )  # type: ignore
    mtc.to(config.device)
    logger.debug('Initailizing dataloader')

    frame_iter = reader.sharding_filter().batch(config.batch_size).collate(collate_fn=collate_fn)
    rs = MultiProcessingReadingService(num_workers=config.num_workers, multiprocessing_context='spawn')
    dl = DataLoader2(frame_iter, reading_service=rs)
    dl.seed(0)
    logger.debug('Starting processor loop')

    lg = graph.LocalGraph(matcher=mtc)
    lg.to(config.device)

    for f_idx, frame in enumerate(dl):
        if queue is not None:
            typer.secho(
                f'Processing frame {f_idx} / {n_f}; Queue Capacity: {queue.qsize()} / {queue.maxsize}',
                fg=typer.colors.GREEN,
            )
        else:
            typer.secho(f'Processing frame {f_idx} / {n_f}', fg=typer.colors.GREEN)
        frame = frame.to(config.device).float()

        newframe_kpts, newframe_descs, _ = det_model(frame)
        newframe_kpts = spherical(newframe_kpts, True)
        lg.update(newframe_kpts, newframe_descs)
        simple_visualize(lg.track, lg.keyframe_locs)
        typer.secho(f'Keyframe tracked: {lg.keyframe_locs.shape[0]}', fg=typer.colors.GREEN)


from urllib.request import Request
from urllib.request import urlopen
import pytz
import pandas as pd

from bs4 import BeautifulSoup
from datetime import datetime
from pandas.io.data import DataReader

SITE = "http://en.wikipedia.org/wiki/List_of_S%26P_500_companies"

def scrape_list(site):
    hdr = {'User-Agent': 'Mozilla/5.0'}
    req = Request(site, headers=hdr)
    page = urlopen(req)
    soup = BeautifulSoup(page)

    table = soup.find('table', {'class': 'wikitable sortable'})
    sector_tickers = dict()
    for row in table.findAll('tr'):
        col = row.findAll('td')
        if len(col) > 0:
            sector = str(col[3].string.strip()).lower().replace(' ', '_')
            ticker = str(col[0].string.strip())
            if sector not in sector_tickers:
                sector_tickers[sector] = list()
            sector_tickers[sector].append(ticker)
    return sector_tickers

def main():
    ls = scrape_list(SITE)
    sdf = []
    for key in ls:
        for k in ls[key]:
            sdf.append(k)
    return sdf

if __name__ == "__main__":
    print(pd.DataFrame(main()))


from django.urls import path, include
from sub_apps.views_related import views

from rest_framework import routers

router = routers.SimpleRouter()
router.register(r'router', views.UserRouterModelView, basename='router')

urlpatterns = [
    # 视图
    path('users/', views.UserModelView.as_view(actions={
        'get': 'list',
        'post': 'create',
    }), name='user_view'),
    path('users/<int:pk>/', views.UserModelView.as_view(actions={
        'get': 'retrieve',
        'put': 'update',
        'patch': 'partial_update',
        'delete': 'destroy',
    }), name='user_view'),
    # 筛选
    # 1.自定义筛选
    path('users/filter/', views.UserFilterModelView.as_view(actions={
        'get': 'list',
    }), name='user_filter_view'),
    # 2.第三方筛选
    path('users/filter2/', views.UserFilter2ModelView.as_view(actions={
        'get': 'list',
    }), name='user_filter2_view'),
    # 3.内置筛选
    path('users/filter3/', views.UserFilter3ModelView.as_view(actions={
        'get': 'list',
    }), name='user_filter3_view'),

    # 分页
    path('users/page/', views.UserPageModelView.as_view(), name='user_page_view'),
    path('users/page2/', views.UserPage2ModelView.as_view(actions={
        'get': 'list',
    }), name='user_page2_view'),

    # router
    path('api/', include(router.urls)),

    # parser
    path('users/parser/', views.UserParserView.as_view(), name='user_parser_view'),

    # 获取前五条数据
    path('top5/', views.Top5ModelView.as_view(actions={
        'get': 'list',
    }), name='top5_view'),
]

# urlpatterns += router.urls


import requests
from .search import shorten

def conv_to_jpeg(image):
    response = requests.get(image)
    file_name = "anilist.jpg"
    file = open(file_name, "wb")
    file.write(response.content)
    file.close()
    return file_name

def format_results(json):
    msg = f"""
**{json['title']['romaji']}**(`{json['title']['native']}`)
**Type**: {json['format']}
**Status**: {json['status']}
**Episodes**: {json.get('episodes', 'N/A')}
**Duration**: {json.get('duration', 'N/A')} Per Ep.
**Score**: {json['averageScore']}
**Genres**: `
"""
    for x in json['genres']:
        msg += f"{x}, "
    msg = msg[:-2] + '`\n'
    msg += "**Studios**: `"
    for x in json['studios']['nodes']:
        msg += f"{x['name']}, "
    msg = msg[:-2] + '`\n'
    info = json.get('siteUrl')
    trailer = json.get('trailer', None)
    if trailer:
        trailer_id = trailer.get('id', None)
        site = trailer.get('site', None)
        if site == "youtube":
            trailer = 'https://youtu.be/' + trailer_id
    description = json.get('description', 'N/A').replace('', '').replace('', '').replace('
', '')
    msg += shorten(description, info)
    image = info.replace('anilist.co/anime/', 'img.anili.st/media/')
    return msg, info, trailer, image


import itertools
import heapq

# Notes:
# * Leverages as many of the library approaches as I can think of.
# * Intentional one liner in for loop - not aiming for readabity :/

def sum_three_largest(iterable):
	# Three element heap
	heap = [0,0,0]

	for total in (sum(int(numstr) for numstr in value) for key, value in itertools.groupby(iterable, lambda x: x == '\n') if not key):
		# Push the new element and pop the lowest/smallest value
		heapq.heappushpop(heap, total)

	return sum(heap)

if __name__ == '__main__':
	import sys
	print(sum_three_largest(open(sys.argv[1])))


import gym
import numpy as np
import tensorflow as tf
from tensorflow import keras
import random
from keras.layers import Activation, Dense
from keras.models import Sequential
from keras.optimizers import RMSprop, Adam
import copy
import sys
import os
import logging
from trainingInstance import trainingInstance
import matplotlib.pyplot as plt

# Creates enviroment from OpenAI gym
env = gym.make('CartPole-v0')

# Hyper parameters
learningRate = 0.80
discountRate = 0.99
exploreRate = 1.0
exploreDecayRate = 0.95

# This is the number of games to play before
# using the data to make graphs
numEpisodes = 200

# This is the total reward seen by the agent
totalReward = 0

# This will store the list of all the reward we
# have seen over time
allReward = []

# This is the current game number
gameNum = 0

# Create the neural network
Q_value = Sequential()
Q_value.add(Dense(48, input_dim = 4, activation='relu'))
Q_value.add(Dense(24, activation='relu'))
Q_value.add(Dense(2, activation='linear'))
Q_value.compile(loss='mse',  optimizer = Adam(lr = 0.001) )

# This is the memory replay buffer
memory = []

observation = env.reset()
observation = np.reshape(observation, [1, 4])
observationPrior = copy.deepcopy(observation)

# Either 0 or 1 - 0 is force left, 1 is force right
action = 0

# Generate training data and learn from it
while (gameNum < numEpisodes):

    env.render()

    if ( exploreRate > random.random() ):
        action = env.action_space.sample()
    else:
        # returns the index of the maximum element
        action = np.argmax( Q_value.predict( observation)[0] )

    observation, reward, done, info = env.step(action)
    observation = np.reshape(observation, [1, 4])

    totalReward = totalReward + reward

    if ( done == True ):
        reward = -100

    # Add the training instance to our memory
    memory.append( trainingInstance(observationPrior, observation, reward, action, done, gameNum)  )

    # Sample from the memory and do an epoch of training
    batchSize = 20 # 20 #10
    batch = memory
    if ( len(memory) < batchSize):
        batch = memory
    else:
        batch = random.sample(memory, batchSize)

    for i in range(len(batch) ):
        value = Q_value.predict(batch[i].observationPrior)

        value[0][action] = batch[i].reward

        if ( (batch[i].action == 0) and (batch[i].done == False) ):
            value[0][0] = batch[i].reward + discountRate * np.amax(Q_value.predict(batch[i].observation)[0] )
        elif ( (batch[i].action == 1) and (batch[i].done == False) ):
            value[0][1] = batch[i].reward + discountRate * np.amax(Q_value.predict(batch[i].observation)[0] )

        Q_value.fit( batch[i].observationPrior, value, epochs = 2, verbose = 0)

    exploreRate = exploreRate * exploreDecayRate
    exploreRate = max(exploreRate, 0.01)

    observationPrior = copy.deepcopy(observation)

    if (done == True):
        observation = env.reset()
        observation = np.reshape(observation, [1, 4])
        observationPrior = copy.deepcopy(observation)

        gameNum = gameNum + 1

        # Format the string we will print
        sys.stdout.write("\033[F") # back to previous line
        sys.stdout.write("\033[K") # clear line
        print("For the last completed episode: " + str(gameNum) +  " we scored " + str(totalReward) )
        allReward.append(totalReward)

        totalReward = 0

    # Periodically empty the memory buffer and up the explore rate
    if ( gameNum % 10 == 0):
        memory = []
        exploreRate = 0.3

# Plot the data
plt.plot( allReward )
plt.ylabel('Cumulative Reward')
plt.xlabel('Episode')
plt.show()

env.close()