#! /usr/bin/env python
import json
import os
import random
from datetime import datetime
from time import sleep
from time import time as ts
from typing import List, Optional

import pytz
import requests
from pydantic import BaseModel, Field

BOARD_ID = 642101
SESSION_COOKIE = os.getenv("SESSION_COOKIE")
CACHE_FOLDER = os.getenv("CACHE_FOLDER", "./cache")

class MemberStanding(BaseModel):
    id_: int = Field(alias="id")
    name: str
    position: int
    stars: List[int]
    score: int
    gold_stars: int
    silver_stars: int
    draw_entries: int
    last_star_ts: int
    part_2_average_time: float

def _download_data(year: int) -> dict:
    url = f"https://adventofcode.com/{year}/leaderboard/private/view/{BOARD_ID}.json"

    resp = requests.get(url, cookies={"session": SESSION_COOKIE})
    resp.raise_for_status()

    if resp.headers["Content-Type"] != "application/json":
        raise Exception("Unable to fetch standings from advent-of-code")

    return resp.json()

def _get_score_of_entry(stars: List[int]) -> int:
    score_map = {2: 3, 1: 1, 0: 0}
    return sum(score_map[i] for i in stars)

def _get_stars_of_entry(entry, cutoff_time: Optional[datetime]) -> List[int]:
    stars = [0] * 25
    if cutoff_time:
        cutoff_time = cutoff_time.timestamp()

    def _is_star_unlocked(entry, star_num):
        star_num = str(star_num)

        if star_num not in entry:
            return False

        if cutoff_time is None:
            return True

        return entry[star_num]["get_star_ts"] < cutoff_time

    for d, v in entry["completion_day_level"].items():
        if _is_star_unlocked(v, 2):
            s = 2
        elif _is_star_unlocked(v, 1):
            s = 1
        else:
            s = 0

        stars[int(d) - 1] = s

    return stars

def _get_part_2_average_time(completion_day_level) -> float:
    completed_days = [d for d in completion_day_level.values() if "2" in d]
    if not completed_days:
        return 0

    return sum(
        d["2"]["get_star_ts"] - d["1"]["get_star_ts"] for d in completed_days
    ) / len(completed_days)

def _data_is_up_to_date(fp: str) -> bool:
    return os.path.exists(fp) and (ts() - os.path.getmtime(fp)) < 15 * 60

def _get_standings(year: int):
    fp = os.path.join(CACHE_FOLDER, f"{BOARD_ID}-{year}-raw.json")

    if _data_is_up_to_date(fp):
        with open(fp) as f:
            raw_standings = json.load(f)
    else:
        raw_standings = _download_data(year)
        with open(fp, "w") as f:
            json.dump(raw_standings, f)

    cutoff_time = (
        pytz.timezone("America/Montreal").localize(datetime(year + 1, 1, 1))
        if year >= 2021
        else None
    )

    members = []

    for k, v in raw_standings["members"].items():
        stars = _get_stars_of_entry(v, cutoff_time)

        members.append(
            MemberStanding(
                id=int(k),
                name=v["name"] or f"anonymous user #{v['id']}",
                position=0,
                stars=stars,
                gold_stars=sum(1 for s in stars if s == 2),
                silver_stars=sum(1 for s in stars if s > 0),
                score=_get_score_of_entry(stars),
                draw_entries=sum(stars),
                last_star_ts=v["last_star_ts"],
                part_2_average_time=_get_part_2_average_time(v["completion_day_level"]),
            )
        )

    members.sort(
        key=lambda m: (
            m.draw_entries,
            m.gold_stars,
            m.silver_stars,
            list(reversed(m.stars)),
            -m.part_2_average_time,
        ),
        reverse=True,
    )

    for i, m in enumerate(members):
        m.position = i + 1

    return members

def draw(coupons):
    sleep(1)
    print(".", end="", flush=True)
    sleep(1)
    print(".", end="", flush=True)
    sleep(1)
    random.shuffle(coupons)  # Probably not necessary but whatevs Ъци
    winner = random.choice(coupons)
    print(f" {winner}")

    return winner

standings = _get_standings(2021)

for s in standings:
    if s.draw_entries == 0:
        continue
    print(f"{s.position:-2d}. {s.name}: {s.draw_entries}")

coupons = []
exclusions = ["William Cantin"]
for s in standings[2:]:
    if s.name in exclusions:
        continue

    coupons += [s.name] * s.draw_entries

random.shuffle(coupons)

input("\nPress enter to draw the first winner")
winner = draw(coupons)

coupons = [c for c in coupons if c != winner]

input("\nPress enter to draw the second winner")
winner = draw(coupons)
coupons = [c for c in coupons if c != winner]

input("\nPress enter to draw the third winner")
winner = draw(coupons)


from logger import Logger
from agent import Agent

from sklearn.utils import shuffle
from collections import deque
from scipy.stats import norm
from copy import deepcopy
import numpy as np
import argparse
import pickle
import random
import torch
import wandb
import copy
import time
import gym
import os
import datetime
import sys
sys.path.append('../')

import hydra
from omegaconf import DictConfig

from omniisaacgymenvs.utils.hydra_cfg.hydra_utils import *
from omniisaacgymenvs.utils.hydra_cfg.reformat import omegaconf_to_dict, print_dict
from omniisaacgymenvs.utils.config_utils.path_utils import retrieve_checkpoint_path
from omniisaacgymenvs.utils.task_util import initialize_task
from omniisaacgymenvs.envs.vec_env_rlgames import VecEnvRLGames

@hydra.main(config_name="config", config_path="../omniisaacgymenvs/cfg")
def train(cfg: DictConfig):

    """ ------ Isaac GYM Setting ------ """
    cfg_dict = omegaconf_to_dict(cfg)
    print_dict(cfg_dict)

    headless = cfg.headless
    render = not headless
    rank = int(os.getenv("LOCAL_RANK", "0"))
    if cfg.multi_gpu:
        cfg.device_id = rank
        cfg.rl_device = f'cuda:{rank}'
    enable_viewport = "enable_cameras" in cfg.task.sim and cfg.task.sim.enable_cameras
    env = VecEnvRLGames(headless=headless, sim_device=cfg.device_id, enable_livestream=cfg.enable_livestream, enable_viewport=enable_viewport)

    # ensure checkpoints can be specified as relative paths
    if cfg.checkpoint:
        cfg.checkpoint = retrieve_checkpoint_path(cfg.checkpoint)
        if cfg.checkpoint is None:
            quit()

    # sets seed. if seed is -1 will pick a random one
    from omni.isaac.core.utils.torch.maths import set_seed
    cfg.seed = set_seed(cfg.seed, torch_deterministic=cfg.torch_deterministic)
    cfg_dict['seed'] = cfg.seed
    task = initialize_task(cfg_dict, env)

    """---------------------------------"""

    """ ------ CPO implementation Setting ------ """
    algo_idx = 1
    agent_name = 'isaac_CPO'
    env_name = cfg.task_name
    max_ep_len = 500
    max_steps = 1000
    epochs = 2500
    save_freq = 10
    algo = '{}_{}'.format(agent_name, algo_idx)
    save_name = '_'.join(env_name.split('-')[:-1])
    save_name = "result/{}_{}".format(save_name, algo)
    args = {
        'agent_name':agent_name,
        'save_name': save_name,
        'discount_factor':0.99,
        'hidden1':512,
        'hidden2':512,
        'v_lr':2e-4,
        'cost_v_lr':2e-4,
        'value_epochs':200,
        'batch_size':10000,
        'num_conjugate':10,
        'max_decay_num':10,
        'line_decay':0.8,
        'max_kl':0.001,
        'damping_coeff':0.01,
        'gae_coeff':0.97,
        'cost_d':25.0/1000.0,
    }
    if torch.cuda.is_available():
        device = torch.device('cuda')
        print('[torch] cuda is used.')
    else:
        device = torch.device('cpu')
        print('[torch] cpu is used.')

    agent = Agent(env, device, args)

    cfg_dict = omegaconf_to_dict(cfg)
    print_dict(cfg_dict)

    if cfg.wandb_activate and rank == 0:
        # Make sure to install WandB if you actually use this.
        time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        run_name = f"{cfg.task_name}_{time_str}"

        wandb.init(
            project=cfg.wandb_project,
            entity=cfg.wandb_entity,
            config=cfg_dict,
            sync_tensorboard=True,
            name=run_name,
            resume="allow",
        )
    """------------------------------------------"""

    while env._simulation_app.is_running():
        for epoch in range(epochs):
            trajectories = []
            ep_step = 0
            scores = []

            while ep_step < max_steps:
                states = env.reset()
                states = states['obs']
                states_tensor = torch.tensor(states, device=device, dtype=torch.float)
                score = 0
                cv = 0
                step = 0
                while step<max_ep_len:
                    ep_step += 1
                    step += 1
                    # states_tensor = torch.tensor(states, device=device, dtype=torch.float)
                    # states_tensor = states.clone().detach()
                    actions, _ = agent.getAction(states_tensor, is_train=True)
                    actions_numpy = actions.detach().cpu().numpy()

                    next_states, rewards, dones, _ = env.step(actions)
                    next_states = torch.tensor(next_states['obs'], device=device, dtype=torch.float)
                    # print(next_states.size(), rewards.size(), dones.size())
                    cost = -rewards

                    trajectories.append([states_tensor.cpu().numpy(), actions_numpy, rewards.cpu(), cost.cpu(), dones.cpu(), torch.zeros(rewards.size()).cpu(), next_states.cpu().numpy()])
                    states = next_states
                    score += torch.sum(rewards).item()/env.num_envs
            # print(trajectories[0])
            states = np.array([traj[0] for traj in trajectories])
            print(states.shape)
            scores.append(score)
            v_loss, cost_v_loss, objective, cost_surrogate, kl, entropy = agent.train(trajs=trajectories)
            score = np.mean(scores)
            log_data = {"score":score, 'cv':cv, "value loss":v_loss, "cost value loss":cost_v_loss, "objective":objective, "cost surrogate":cost_surrogate, "kl":kl, "entropy":entropy}

            print(f'\nEpoch {epoch+1}/{epochs} result :')
            print(log_data)

    env._simulation_app.close()

if __name__ == '__main__':
    train()


def rotated_array_search(input_list, number):
    """
    Find the index by searching in a rotated sorted array

    Args:
       input_list(array), number(int): Input array to search and the target
    Returns:
       int: Index or -1
    """
    rotation_index = find_rotation_point(input_list, 0, len(input_list) - 1)

    # array is not rotated
    # return normal binary search on the whole array
    if rotation_index == -1:
        return binary_search(input_list, 0, len(input_list)-1, number)

    # array is rotated
    # do binary search on the first non-rotated half
    # if `number` was not found, do binary search on the second non-rotated half
    search_result = binary_search(input_list, 0, rotation_index, number)
    if search_result == -1:
        search_result = binary_search(input_list, rotation_index + 1, len(input_list) - 1, number)
    return search_result

def find_rotation_point(input_list, start, end):
    if end < start:
        return -1
    # rotation point found
    if start == end:
        return start
    middle = (start + end) // 2
    # rotation point is in the vicinity of `middle`
    if (middle < end) and (input_list[middle] > input_list[middle + 1]):
        return middle
    if (middle > start) and (input_list[middle] < input_list[middle - 1]):
        return middle - 1
    # array is rotated before `middle`
    if input_list[start] >= input_list[middle]:
        return find_rotation_point(input_list, start, middle - 1)
    # array is non-rotated until `middle`
    return find_rotation_point(input_list, middle + 1, end)

def binary_search(input_list, start, end, number):
    if end < start:
        return -1
    if start == end:
        if input_list[start] == number:
            return start
        return -1

    middle = (start + end) // 2
    middle_item = input_list[middle]

    if middle_item == number:
        return middle
    if middle_item > number:
        return binary_search(input_list, start, middle - 1, number)
    return binary_search(input_list, middle + 1, end, number)

def linear_search(input_list, number):
    for index, element in enumerate(input_list):
        if element == number:
            return index
    return -1

def test_function(test_number, test_case, test_subject):
    input_list = test_case[0]
    number = test_case[1]
    result = None
    if linear_search(input_list, number) == rotated_array_search(input_list, number):
        result = "Pass"
    else:
        result = "Fail"
    print("{} - {} - {}".format(test_number, result, test_subject))

test_function(1, [[], 1], "empty array")
test_function(2, [[1],1], "single item in array")
test_function(3, [[6, 7, 8, 1, 2, 3, 4], 10], "non-existent item")
test_function(4, [[1, 2, 3, 4, 5, 6, 7], 3], "non-rotated array")
test_function(5, [[6, 7, 8, 9, 10, 1, 2, 3, 4], 6], "front of rotation")
test_function(6, [[6, 7, 8, 9, 10, 1, 2, 3, 4], 3], "back of rotation")
test_function(7, [[6, 7, 8, 1, 2, 3, 4], 8], "highest item in array")
test_function(8, [[6, 7, 8, 1, 2, 3, 4], 1], "lowest item in array")


import numpy as np
from openmdao.api import ExplicitComponent

class BcComp(ExplicitComponent):

    def initialize(self):
        self.options.declare('num_nodes', default=40, types=int)
        self.options.declare('k', default=3, types=int)
        self.options.declare('tube_nbr', default=3, types=int)

    def setup(self):
        #Inputs
        num_nodes = self.options['num_nodes']
        k = self.options['k']
        tube_nbr = self.options['tube_nbr']

        self.add_input('dpsi_ds',shape=(num_nodes,k,tube_nbr))

        self.add_output('torsionconstraint',shape=((k,tube_nbr)))

        # partials
        self.declare_partials('torsionconstraint','dpsi_ds')

    def compute(self,inputs,outputs):

        num_nodes = self.options['num_nodes']
        k = self.options['k']
        tube_nbr = self.options['tube_nbr']
        dpsi_ds = inputs['dpsi_ds']
        bc = np.zeros((k,tube_nbr))

        bc[:,:] = dpsi_ds[-1,:,:]

        outputs['torsionconstraint'] = bc

    def compute_partials(self,inputs,partials):

        ''' Jacobian of partial derivatives for P Pdot matrix.'''

        num_nodes = self.options['num_nodes']
        k = self.options['k']
        dpsi_ds = inputs['dpsi_ds']
        tube_nbr = self.options['tube_nbr']

        '''Computing Partials'''
        ppsi_dot = np.zeros((tube_nbr*k, num_nodes*k*tube_nbr))
        ppsi_dot[:,(num_nodes-1)*k*tube_nbr:] = np.identity(k*tube_nbr)

        partials['torsionconstraint','dpsi_ds']= ppsi_dot

if __name__ == '__main__':
    from openmdao.api import Problem, Group

    from openmdao.api import IndepVarComp

    group = Group()

    comp = IndepVarComp()
    num_nodes =3
    k = 1
    tube_nbr = 4

    tip=np.random.random((k,tube_nbr))*3
    dpsi=np.random.random((num_nodes,k,tube_nbr))
    comp.add_output('tube_ends_tip',val=tip)
    comp.add_output('dpsi_ds',val=dpsi)

    group.add_subsystem('comp1', comp, promotes = ['*'])

    comp = BcComp(k=k,num_nodes=num_nodes,tube_nbr=tube_nbr)
    group.add_subsystem('comp2', comp, promotes = ['*'])

    prob = Problem()
    prob.model = group
    prob.setup()
    prob.run_model()
    prob.model.list_outputs()
    # prob.check_partials(compact_print=False)
    prob.check_partials(compact_print=True)
    # prob.model.list_outputs()


#! /usr/bin/env python

import numpy as np

from typing import List, Tuple
from scipy.sparse import csr_matrix

def Ppi(transitions: List[np.ndarray], nA: int, pol: np.ndarray) -> np.ndarray:
	ppi = pol[:, 0, None] * transitions[0]
	for i in range(1, nA):
		ppi += pol[:, i, None] * transitions[i]

	return ppi

def evaluate_pol(transitions: List[np.ndarray], rewards: np.ndarray, gamma: float, pol: np.ndarray, nX: int, nA: int):

	# Cost and Probs averaged by policy
	cpi = (pol * rewards).sum(axis=1)
	ppi = Ppi(transitions, nA, pol)

	# J = (I - gamma*P)^-1 * c
	J = np.linalg.inv(np.eye(nX) - gamma * ppi).dot(cpi)

	return J[:, None]

def policy_iteration(mdp: Tuple[np.ndarray, List[str], List[csr_matrix], np.ndarray, float],
					 init_pol: np.ndarray = None, init_q: np.ndarray = None) -> (np.ndarray, np.ndarray):

	X = mdp[0]
	A = mdp[1]
	P = []
	for act in range(len(A)):
		P += [mdp[2][act].toarray()]
	c = mdp[3]
	gamma = mdp[4]
	nX = len(X)
	nA = len(A)

	# Initialize pol
	if init_pol is None:
		pol = np.ones((nX, nA)) / nA
	else:
		pol = init_pol

	# Initialize Q
	if init_q is None:
		Q = np.zeros((nX, nA))
	else:
		Q = init_q

	quit = False
	i = 0

	while not quit:

		# print('Iteration %d' % (i + 1), end='\r')

		J = evaluate_pol(P, c, gamma, pol, nX, nA)

		for act in range(nA):
			Q[:, act, None] = c[:, act, None] + gamma * P[act].dot(J)

		# print(Q)
		Qmin = Q.max(axis=1, keepdims=True)
		polnew = np.isclose(Q, Qmin, atol=1e-10, rtol=1e-10).astype(int)
		polnew = polnew / polnew.sum(axis=1, keepdims=True)

		quit = (pol == polnew).all()
		pol = polnew
		i += 1

	# print('N. iterations: ', i)

	return pol, Q


#!/usr/bin/python3
"""This module gathers data from an API and exports it to a CSV file"""
if __name__ == "__main__":
    import csv
    import requests
    from sys import argv
    emp_id = argv[1]
    base_url = "https://jsonplaceholder.typicode.com/"
    user_url = "{}users/{}".format(base_url, emp_id)
    emp_uname = requests.get(user_url).json().get('username')
    td_url = "{}/todos".format(user_url)
    td_request_json = requests.get(td_url).json()
    filename = emp_id + ".csv"
    with open(filename, "w") as f:
        td_writer = csv.writer(f, quotechar='"', quoting=csv.QUOTE_ALL)
        for td in td_request_json:
            td_writer.writerow([td['userId'], emp_uname, td['completed'],
                                td['title']])


# NAME: James Jenkins
# REFERENCES: Code based on my work from last semester Algorithms class
# STUDENT ID: 104479208

class PriorityQueue:
    def __init__ (self, nNodes):
        self.n = nNodes
        self.d = {}
        self.Inf = float('inf')

    def set(self, i, dist):
        assert(i >= 0 and i < self.n)
        self.d[i] = min(dist, self.Inf)

    def hasKey(self,i):
        return i in self.d

    def get(self,i):
        assert(i >= 0 and i < self.n)
        assert(i in self.d)
        return self.d[i]

    def extractMin(self):

        m = len(self.d)
        assert (m > 0), 'Empty Queue'

        minDist = self.Inf
        minKey = -1

        for (i,dist) in self.d.items():
            if (dist <= minDist):
                minDist = dist
                minKey = i
        assert (minKey >=  0.0), ('minKey = '+str(minKey)+' dist='+str(dist))
        self.d.pop(minKey,None)

        return (minKey, minDist)

    def isEmpty(self):
        m = len(self.d)
        return (m <= 0)

class MyGraph:
    def __init__(self, nVertices):
        self.n = nVertices
        self.adjList = [ [] for i in range(0,nVertices) ]

    def getName(self,i):
        return str(i)

    def addEdge(self,i,j,w):
        assert( i >= 0 and i < self.n)
        assert( j >= 0 and j < self.n and j != i)
        lst = self.adjList[i]
        lst.append( (j,w) )

    def singleSourceShortestPath(self, srcID):
        assert (srcID >= 0 and srcID < self.n)
        pq = PriorityQueue(self.n)

        for i in range(0,self.n):
            if ( i == srcID):
                pq.set(i,0.0)
            else:
                pq.set(i, pq.Inf)

        d = {}

        pi = {}

        while(not pq.isEmpty()):
            (u, cost) = pq.extractMin()
            d[u] = cost
            for i in range(len(self.adjList[u])):
                if(pq.hasKey(self.adjList[u][i][0]) and self.adjList[u][i][1] + d[u] < pq.get(self.adjList[u][i][0])):
                    pq.set(self.adjList[u][i][0],self.adjList[u][i][1] + d[u])
                    #print(self.adjList[u])
                    pi[self.adjList[u][i][0]] = u
                    #print("Current Distance:",d[u], "Node:",u)

        return (d,pi)

    def getShortestPath(self,d,pi,srcID, destID):
        assert (srcID in d), 'It looks like your Dijkstra code is not complete'
        assert (d[srcID] <= 0.0)
        if (destID not in pi):
            assert False , ' No path from the given source : %d to destination: %d'%(srcID,destID)
        lst = [destID]
        curNode = destID
        while (curNode != srcID):
            assert(curNode in pi)
            curNode = pi[curNode]
            lst.insert(0,curNode)
        return lst

# TEST
graph = MyGraph(19)
s = 0
a = 1
b = 2
c = 3
d = 4
e = 5
f = 6
g = 7
h = 8
i = 9
j = 10
k = 11
l = 12
m = 13
n = 14
o = 15
p = 16
q = 17
r = 18
graph.addEdge(s, a, 2)
graph.addEdge(a, b, 1)
graph.addEdge(b, c, 4)
graph.addEdge(b, p, 6)
graph.addEdge(c, d, 4)
graph.addEdge(d, g, 5)
graph.addEdge(d, e, 4)
graph.addEdge(e, r, 3)
graph.addEdge(g, h, 3)
graph.addEdge(g, r, 4)
graph.addEdge(h, i, 2)
graph.addEdge(h, j, 3)
graph.addEdge(r, j, 2)
graph.addEdge(r, l, 2)
graph.addEdge(p, q, 7)
graph.addEdge(q, l, 2)
graph.addEdge(l, m, 2)
graph.addEdge(m, n, 3)
graph.addEdge(n, o, 4)
graph.addEdge(o, k, 5)
graph.addEdge(i, k, 3)
graph.addEdge(k, f, 2)

(d,pi) = graph.singleSourceShortestPath(0)
myList = graph.getShortestPath(d,pi,0,6)

convert = {0:'S', 1:'A', 2:'B', 3:'C', 4:'D', 5:'E',6:'F',7:'G',8:'H',9:'I',10:'J',11:'K',12:'L',13:'M',14:'N',15:'O',16:'P',17:'Q',18:'R'}

for i in myList:
    print(convert[i])

heuristic = {0:8, 1:10, 2:12, 3:10, 4:8, 5:0, 6:8, 7:6, 8:4, 9:4, 10:2, 11:8, 12:8, 13:6, 14:4, 15:12, 16:8, 17:6, 18:10}


class Solution:
    def alternatingSubarray(self, nums: List[int]) -> int:
        res = curr = -1
        for i in range(1, len(nums)):
            if curr > 0 and nums[i] == nums[i - 2]:
                curr += 1
            else:
                curr = 2 if nums[i] == nums[i - 1] + 1 else -1
            res = max(res, curr)
        return res