from flask import Flask, request, redirect
from flask_restful import Resource, Api
import requests
import psutil
import time
import numpy as np
import math
import os
import threading as th
import random
import subprocess

app = Flask(__name__)
api = Api(app)


class ReplayBuffer(object):
    def __init__(self, state_dim, batch_size, buffer_size, device):
        self.batch_size = batch_size
        self.max_size = int(buffer_size)
        self.device = device

        self.ptr = 0
        self.crt_size = 0

        self.state = np.zeros((self.max_size, state_dim))
        self.action = np.zeros((self.max_size, 1))
        self.next_state = np.array(self.state)
        self.reward = np.zeros((self.max_size, 1))
        self.not_done = np.zeros((self.max_size, 1))

    def add(self, state, action, next_state, reward, done, episode_done, episode_start):
        self.state[self.ptr] = state
        self.action[self.ptr] = action
        self.next_state[self.ptr] = next_state
        self.reward[self.ptr] = reward
        self.not_done[self.ptr] = 1. - done

        self.ptr = (self.ptr + 1) % self.max_size
        self.crt_size = min(self.crt_size + 1, self.max_size)

    def sample(self):
        ind = np.random.randint(0, self.crt_size, size=self.batch_size)
        return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.LongTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device)
        )

    def save(self, save_folder):
        np.save(f"{save_folder}_state.npy", self.state[:self.crt_size])
        np.save(f"{save_folder}_action.npy", self.action[:self.crt_size])
        np.save(f"{save_folder}_next_state.npy",
                self.next_state[:self.crt_size])
        np.save(f"{save_folder}_reward.npy", self.reward[:self.crt_size])
        np.save(f"{save_folder}_not_done.npy", self.not_done[:self.crt_size])
        np.save(f"{save_folder}_ptr.npy", self.ptr)

    def load(self, save_folder, size=-1):
        reward_buffer = np.load(f"{save_folder}_reward.npy")

        # Adjust crt_size if we're using a custom size
        size = min(int(size), self.max_size) if size > 0 else self.max_size
        self.crt_size = min(reward_buffer.shape[0], size)

        self.state[:self.crt_size] = np.load(
            f"{save_folder}_state.npy")[:self.crt_size]
        self.action[:self.crt_size] = np.load(
            f"{save_folder}_action.npy")[:self.crt_size]
        self.next_state[:self.crt_size] = np.load(
            f"{save_folder}_next_state.npy")[:self.crt_size]


class Notify (Resource):
    def load_req_thres(self):
        global req_thres
        if os.path.exists("./req_thres.npy"):
            req_thres = np.load("./req_thres.npy")
            req_thres = req_thres[0]
            print("New Policy Request threshold : ", req_thres)

    def get(self):
        global buffer
        global lock
        print("Notification of offload")
        notify = request.args.get('offload')
        notify = int(notify)
        if notify == 0:
            self.load_req_thres()
        else:
            lock.acquire()
            buffer.save('buffer_' + str(notify))
            buffer.ptr = 0
            buffer.crt_size = 0
            lock.release()


class Greeting (Resource):
    def __init__(self, overload=10.0, offload=1.0, reward=0.2, holding=0.12, threshold_req=17):
        self.overload = overload
        self.offload = offload
        self.reward = reward
        self.holding = holding
        self.c = 2
        self.T = 0.5
        self.num_actions = 2

    def sigmoid_fn(self, cpu_util, buffer, debug=0):
        global req_thres
        # Simple as of now
        """
        If this value is > 0.5 i.e. state[0] - req_thres > 0, then we need to offload, return 1
        will high probability as a=1 for offload
        """
        # print ("State 1 ", state[1], type(state[1]))
        prob = math.exp((cpu_util - req_thres[int(buffer)])/self.T) / \
            (1 + math.exp((cpu_util - req_thres[int(buffer)])/self.T))
        if debug:
            print("Sigmoid  state, threshold, prob ",
                  state, req_thres[int(cpu_util)], prob)
        return np.random.binomial(n=1, p=prob, size=1)

    def select_action(self, cpu_util, buffer, eval_=False, debug=0):
        if np.random.uniform(0, 1) > 0.005 or eval_ == True:
            action = self.sigmoid_fn(cpu_util, buffer)
        else:
            action = np.random.randint(self.num_actions)
        if debug:
            print("ACTION : ", action)
        return action

    def get_reward(self, cpu_util, buffer, action, debug=1):
        global buff_size
        rew = 0.0
        if action == 1:
            rew -= self.offload
            print("Offload")
        if cpu_util < 3:
            if action == 1:
                rew -= self.overload
                print("Low util offload")
        elif cpu_util >= 6 and cpu_util <= 17:
            rew += self.reward
            print("Reward")
        elif cpu_util >= 18:
            rew -= self.overload
            print("Overload")
        if buffer == buff_size and action == 0:
            rew -= self.overload
            print("Buffer Full")
        rew -= self.holding * \
            (buffer - self.c) if buffer - self.c > 0 else 0
        return rew

    def get_load(self, str1):
        global start
        # load = os.popen(
        #    "top -b -n 2 -d 0.5 -p 1 | tail -1 | awk '{print $9}'").read()
        load = os.popen(
            "ps -u root -o %cpu,stat | grep -v 'Z' | awk '{cpu+=$1} END {print cpu}'").read()
        load = float(load)
        print("Load ", load, str1)
        return min(int(load/5), 20)

    def get(self):
        global buff_len
        global buff_size
        global start_time
        global offload
        global buffer
        global file_count
        global load
        global lock
        count = request.args.get('count')
        load = self.get_load("new arrival")
        prev_state = [buff_len, load]
        action = self.select_action(load, buff_len)
        if action == 0:
            # Perform task
            count = int(count)
            t = random.expovariate(0.1)
            #t = random.randrange(10000, 60000)
            if buff_len < buff_size:
                lock.acquire()
                buff_len += 1
                lock.release()
                for i in range(1):
                    cpu_l = 4 * buff_len
                    p = subprocess.Popen(
                        ['./try.sh', str(t)])
                    print("Sleep ", t)
                    time.sleep(t)
                # p.terminate()
                load = self.get_load("arrival accept")
            else:
                load = self.get_load("arrival accept")
            # lock.acquire()
            state = [buff_len, load]
            rew = self.get_reward(load, buff_len, action)
            lock.acquire()
            print("ARRIVAL State, Action Next_state Reward",
                  prev_state, action, state, rew)
            buffer.add(prev_state, action, state, rew, 0, 0, 0)
            if buffer.ptr == buffer.max_size - 1:
                file_count += 1
                buffer.save('buffer_' + str(file_count))
            lock.release()
        else:
            count = request.args.get('count')
            print("Offloaded Request")
            resp = requests.get('http://172.17.0.3:3333?count=' + count)
            # lock.acquire()
            load = self.get_load("arrival offload")
            state = [buff_len, load]
            rew = self.get_reward(load, buff_len, action)
            lock.acquire()
            print("ARRIVAL State, Action Next_state Reward",
                  prev_state, action, state, rew)
            buffer.add(prev_state, action, state, rew, 0, 0, 0)
            if buffer.ptr == buffer.max_size - 1:
                file_count += 1
                buffer.save('buffer_' + str(file_count))
            lock.release()
        prev_state = [buff_len, load]
        action = self.select_action(load, buff_len)
        lock.acquire()
        buff_len = max(buff_len - 1, 0)
        lock.release()
        load = self.get_load("departure")
        rew = self.get_reward(load, buff_len, action)
        state = [buff_len, load]
        lock.acquire()
        buffer.add(prev_state, action, state, rew, 0, 0, 0)
        print("DEPT State, Action Next_state Reward",
              prev_state, action, state, rew)
        # if buffer.ptr >= buffer.max_size - 1:
        #    file_count += 1
        #    buffer.save('buffer_' + str(file_count))
        lock.release()
        return [file_count, buffer.ptr]


buff_size = 20
file_count = 0
buff_len = 0
offload = 0
load = 0
batch_size = 1000
replay_size = 1000
state_dim = 2
threshold_req = 17
lock = th.Lock()
start_time = time.time()
buffer = ReplayBuffer(state_dim, batch_size, replay_size, 'cpu')
req_thres = np.full((21), threshold_req, dtype=float)
api.add_resource(Greeting, '/')  # Route_1
api.add_resource(Notify, '/notify')  # Route_2

if __name__ == '__main__':
    app.run('0.0.0.0', '3333')
