import numpy as np
class env:
    rand_seed = 10
   	'''
   	Init function to initialize environment state, it takes number of resources and no of 
  	Tasks in a batch and creates a state representation of the resources.
  	'''
    def __init__(self, resources, num_tasks):
        np.random.seed(self.rand_seed)
        ### initially, no task are chosen
        self.state_ = np.zeros((num_tasks, 1), dtype=int)
        self.limit_ = 0
        for lim in resources:
            self.state_ = np.append(self.state_, np.random.randint(lim, size=(num_tasks, 1)), axis=1)
            
        self.reward_ = 0
        for lim in resources:
            self.limit_+=lim   #Set the total resource limit as sum of resources
       
        self.num_tasks_ = num_tasks
        self.resources_ = resources
    '''
    Get_initial_state: Reset the environment state for the new batch of incoming jobs
    '''   
    def get_initial_state(self):
        state = np.zeros((self.num_tasks_, 1), dtype=int)
        for lim in self.resources_: 
            state = np.append(state, np.random.randint(lim, size=(self.num_tasks_,1)), axis=1)
        self.state_ = state
        self.reward_ = 0
        self.num_tasks_ = num_tasks
        return self.state_
'''
     Return the sum of a one dimensional vector
'''
    def getSum(self, x ):
        if x[0] == 1:
            return sum(x)-1
        else:
            return 0
   
'''
This function defines the logic behind changing the states and deciding rewards based on the input provided by agent in terms of action.
Input: index of the selected task in state table
Output: 
New state:
Reward: 
Done:  Job processing done or not flag.
'''

    def get_next_step(self, actionIdx):
        ### update reward
        self.reward_ = 0
        reward = 0
        done = False
        #update the state based on action
        #case 1, if same task selected then penalize the agent
        if self.state_[actionIdx][0] == 1:
            sum_res = np.sum(self.state_[actionIdx], axis=0)
            reward = -sum_res/10;
        else:
            self.state_[actionIdx][0] = 1
       
        #collect all the resources for this batch until now.
        totReward = sum(np.apply_along_axis( self.getSum, axis=1, arr=self.state_ ))
        # get the sum for selected task
        sum_res = np.sum(self.state_[actionIdx], axis=0) 
        if totReward <= self.limit_:
            reward = sum_res
        else:
            reward = -sum_res/10;
            done = True
            
        return self.state_,reward,done 

resources = [80, 60] #resource limits
num_tasks = 5
en = env(resources,num_tasks)
action = 1
en.get_next_step(action)

from keras.models import Sequential
from keras.layers import InputLayer
from keras.models import Sequential
from keras.layers import Activation, Dense

resources = [80, 80] #resource limits
num_tasks = 5
input_layer_n = num_tasks
output_layer_n = 5
dense_layer_n = 10

envs = env(resources, num_tasks)
'''
Agent using Deep network, Keras Model
Input layers neurons: no of resources + task selection state
Dense layer: Arbitrary no picked for this, currently it is 10.
Output layer: Single dimensional vector of task selection score, called Q values for each state. No of output neurons = 5
'''
model = Sequential()
#model.add(InputLayer(batch_input_shape=(1, input_layer_n)))
model.add(InputLayer(batch_input_shape=(1, 3)))
model.add(Dense(dense_layer_n, activation='sigmoid'))
model.add(Dense(output_layer_n, activation='linear'))
model.compile(loss='mse', optimizer='adam', metrics=['mae'])

num_job_samples = 1
# now execute the q learning
y = 0.95
eps = 0.5
decay_factor = 0.999
r_avg_list = []

'''
Training part of the model built in previous section.
No of job samples = 1024
'''

for i in range(num_job_samples):   #batch of 5 jobs and total of 1000 samples
    s = envs.get_initial_state()# stating state would be no selected jobs
    eps *= decay_factor
    if i % 100 == 0:
        print("Episode {} of {}".format(i + 1, num_job_samples))
    scheduled = False
    r_sum = 0

    while not scheduled:
        if np.random.random() < eps:
            a = np.random.randint(0, 5) #This will pick a random action
        else:
            #predicted value is a vector of q values
            a = np.argmax(model.predict(s)) 
        #print(a)
        
        new_s, r, scheduled = envs.get_next_step(a) 
        print(np.array(new_s))
        #model.predict(np.array(new_s))
        target = r + y * np.max(model.predict(np.array(new_s)))
        #target = r + y * np.max(model.predict(np.array(new_s))
        target_vec = model.predict(s)
        #Update only the current state value
        target_vec[a] = target
        model.fit(s, target_vec.reshape(-1, 25), epochs=1, verbose=0)
        s = new_s
        r_sum += r
    r_avg_list.append(r_sum / 1000)