import numpy as np


class rewardAI(object):
    def __init__(
        self,
        nonReward,
        shootReward,
        shootWithoutReadyReward,
        hitReward,
        winReward,
        loseReward,
        killReward,
    ):
        self.nonReward = nonReward
        self.shootReward = shootReward
        self.shootWithoutReadyReward = shootWithoutReadyReward
        self.hitReward = hitReward
        self.winReward = winReward
        self.loseReward = loseReward
        self.killReward = killReward

    def getRewards(self, remainTime):
        nonR = self.getnonReward()
        shootR = self.getshootReward()
        shootWithoutReadyR = self.getshootWithoutReadyReward()
        hitR = self.gethitReward()
        winR = self.getwinReward(remainTime)
        loseR = self.getloseReward()
        killR = self.getkillReward(remainTime)

        rewards = np.array(
            [[nonR, shootR, shootWithoutReadyR, hitR, winR, loseR, killR]],
            dtype=np.float,
        )
        return rewards

    def getnonReward(self):
        return self.nonReward

    def getshootReward(self):
        return self.shootReward

    def getshootWithoutReadyReward(self):
        return self.shootWithoutReadyReward

    def gethitReward(self):
        return self.hitReward

    def getwinReward(self, time):
        return self.winReward + time

    def getloseReward(self):
        return self.loseReward

    def getkillReward(self, time):
        return self.killReward + time