# The Unit Class
# The two areas of potential modification are:
# - The dictionaries at the top of the file
# - The roll() function.

import random
import math

# UNIT_COSTS is a constant dictionary that holds the resource (money)
# costs for each type of unit.

# UNITS
# Z = ZOMBIE (SPAWN RATE IS ONLY SET INITIALLY [RANDOM NUMBER * 100])
# S = SURVIVOR (DEFAULT SPAWN RATE)
# C = CURED (EXTREMELY SLOW)
UNIT_COSTS = {
    "Z": 100,
    "S": 250,
    "C": 2000
    }

# UNIT_HEALATH is a constant dictionary that holds the starting health
# for a unit of a given unit type (utype).

# UNITS
# Z = ZOMBIE (LOW HEALTH)
# S = SURVIVOR (DEFAULT HEALTH) 
# C = CURED (HIGH HEALTH)

UNIT_HEALTH = {
    "Z": 6,
    "S": 7,
    "C": 7
    }

# UNIT_SIGHT: Not used. Ignore.
UNIT_SIGHT = {
    "Z": 2,
    "S": 2,
    "C": 2
    }

UNIT_FACTIONS = {
    "Zombies": "Z",
    "Survivors": "S",
    "Cured": "C"
}

# You should use this function in your AI to test unit costs
def get_unit_cost(utype):
    try:
        return UNIT_COSTS[utype]
    except KeyError:
        return math.inf

class Unit:
    def __init__(self, ID, utype, faction_id, pos, health, sight_radius):

        # id: int
        self.ID = ID

        # utype: unit type
        # string: Z, S, or C
        self.utype = utype

        # faction_id: str
        self.faction_id = faction_id

        # pos: vec2
        self.pos = pos

        # health: int
        self.health = health

        # sight_radius: int - how far it sees
        # NOT USED.
        self.sight_radius = sight_radius
        
        # Allows for persistent target (if applicable)
        self.currentTarget = None

    def __eq__(self, o):
        return self.ID == o.ID and self.faction_id == o.faction_id
       
    # Infection Function
    def infect(self, new_faction):
        self.utype = UNIT_FACTIONS[new_faction]
        self.health = UNIT_HEALTH[self.utype]
        self.faction_id = new_faction
        
    # Combat Function:
    # Essentially, it is an NxN matrix for all the different
    # unit-to-unit match ups. Currently, the winning combinations of
    # rock-paper-scissors have max damage of 20. All other
    # combinations are 10. Feel free to modify if you want.
    def roll(self, op_utype):
            return random.randint(0, 10)