#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct  4 15:34:25 2019

@author: gparkes
"""

import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
import itertools as it
import random
import time
from copy import deepcopy
import os

__units__ = ["Unit1"]
__simulations__ = ["basic_simulation%d" % (i + 1) for i in range(4)]
__animate__ = ["basic_animate%d" % (i + 1) for i in range(3)]
__math__ = ["euclidean", "magnitude", "deriv_norm", "dudt"]
__ai__ = ["find_random_enemy", "find_nearest_enemy"]
__plot__ = ["plot_quiver"]

# __all__ = __units__ + __simulations__ + __math__ + __ai__ + __plot__
__all__ = []

"""############################ UNIT CLASSES ############################################"""


class Unit1(object):
    """This first class defines a Basic Python object unit."""

    def __init__(self, db, name):
        """
        Defines a unit, based on examples in the 'db' file, given a name. Assumes the name
        is the index in db.
        """
        # check that name is in db.index
        if name not in db.index:
            raise ValueError("unit {} must be in {}".format(name, db.index))
        self.name = name
        self.hp = db.loc[name, "HP"]
        self.dmg = db.loc[name, "Damage"]
        self.range = db.loc[name, "Range"]
        self.speed = db.loc[name, "Movement Speed"]
        self.team = db.loc[name, "allegiance_int"]
        self.team_name = db.loc[name, "Allegiance"]

        # define position
        self.pos = np.array([0., 0.])
        # directional derivative from target.
        self.dd = np.array([0., 0.])
        # distance from target
        self.dist = 0.
        # define target to aim for
        self.target = None

    # a basic method to assign position
    def set_position(self, pos_params):
        """Sets the position using pos_parameters of mean and variance."""
        self.pos = np.random.normal(*pos_params, size=(2,))
        return self


"""########################## SIMULATIONS ################################################"""


def basic_simulation1(units, max_step=100):
    """The first basic simulation."""
    t = 0
    running = True

    # while loop
    while running and (t < max_step):
        # iterate over units.
        for i, u in enumerate(units):
            # if we're alive...
            if u.hp > 0:
                # euclidean distance
                u.dd = deriv_norm(u.pos, u.target.pos)
                u.dist = euclidean(u.pos, u.target.pos)
                # now, if the unit is not in range, then move towards it
                if u.dist > u.range:
                    # speed modifier added, others ignored from above.
                    u.pos += dudt(u.pos, u.target.pos, u.speed)
                else:
                    # perform damage to enemy health.
                    u.target.hp -= u.dmg
        t += 1


def basic_simulation2(units, max_step=50):
    """The second basic simulation."""
    t = 0
    running = True

    steps = []

    # while loop
    while running and (t < max_step):
        # iterate over units.
        for i, u in enumerate(units):
            # if we're alive...
            if u.hp > 0:
                # euclidean distance
                u.dd = deriv_norm(u.pos, u.target.pos)
                u.dist = euclidean(u.pos, u.target.pos)
                # now, if the unit is not in range, then move towards it
                if u.dist > u.range:
                    # speed modifier added, others ignored from above.
                    u.pos += dudt(u.pos, u.target.pos, u.speed)
                else:
                    # perform damage to enemy health.
                    u.target.hp -= u.dmg

        # save
        steps.append(deepcopy(units))
        t += 1

    return steps


def basic_simulation3(units, max_step=50):
    """The third basic simulation."""
    # create a copy to prevent overriding
    units = deepcopy(units)
    t = 0
    running = True

    steps = []

    # while loop
    while running and (t < max_step):
        # iterate over units.
        for i, u in enumerate(units):
            # if we're alive...
            if u.hp > 0:

                # wait a second. what if our enemy is dead..?
                if u.target.hp <= 0.:
                    # find a new random enemy.
                    target = find_random_enemy(u, units)
                    if target != -1:
                        u.target = target
                    else:
                        # no more enemies - stop!
                        running = False

                """Only do all this once we have a valid target..."""
                # euclidean distance
                u.dd = deriv_norm(u.pos, u.target.pos)
                u.dist = euclidean(u.pos, u.target.pos)
                # now, if the unit is not in range, then move towards it
                if u.dist > u.range:
                    # speed modifier added, others ignored from above.
                    u.pos += dudt(u.pos, u.target.pos, u.speed)
                else:
                    # perform damage to enemy health.
                    u.target.hp -= u.dmg

        # save
        steps.append(deepcopy(units))
        # add step
        t += 1

    return steps


def basic_simulation4(units, max_step=50):
    """The fourth basic simulation."""
    # create a copy of the units
    units = deepcopy(units)
    # begin t
    t = 0
    running = True

    steps = []

    # while loop
    while running and (t < max_step):
        # iterate over units.
        for i, u in enumerate(units):
            # if we're alive...
            if u.hp > 0:

                # wait a second. what if our enemy is dead..?
                if u.target.hp <= 0.:
                    # find a new random enemy.
                    target = find_nearest_enemy(u, units)
                    if target != -1:
                        u.target = target
                    else:
                        # no more enemies - stop!
                        running = False

                """Only do all this once we have a valid target..."""
                # euclidean distance
                u.dd = deriv_norm(u.pos, u.target.pos)
                u.dist = euclidean(u.pos, u.target.pos)
                # now, if the unit is not in range, then move towards it
                if u.dist > u.range:
                    # speed modifier added, others ignored from above.
                    u.pos += dudt(u.pos, u.target.pos, u.speed)
                else:
                    # perform damage to enemy health.
                    u.target.hp -= u.dmg

        # save
        steps.append(deepcopy(units))
        # add step
        t += 1

    return steps


"""############################## ANIMATE FUNCTIONS #########################################"""


def set_boundary(ax, pos):
    """Sets the boundary on an axis."""
    xmin, xmax = np.min(pos[:, 0, :]) - 1., np.max(pos[:, 0, :]) + 1.
    ymin, ymax = np.min(pos[:, 1, :]) - 1., np.max(pos[:, 1, :]) + 1.
    # set bounds
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    # hide axes labels
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)


def basic_animate1(results, n):
    """The first basic animation function."""
    # extract data points
    pos = extract_pos(results, n)
    plt.rcParams["animation.html"] = "html5"
    N_f = len(results)

    # define a figure
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111)
    # define boundaries - we use a function
    set_boundary(ax, pos)

    # set the object to draw to
    u_points, = ax.plot([], [], 'x', color="black", alpha=.8, markersize=5.)

    fig.tight_layout()
    # close the plot
    plt.close()

    # define init
    def _init():
        u_points.set_data(pos[0, 0, :], pos[0, 1, :])
        return u_points

    # define animate
    def _animate(i):
        u_points.set_data(pos[i, 0, :], pos[i, 1, :])
        return u_points

    ani = animation.FuncAnimation(fig, _animate, init_func=_init, interval=100, frames=N_f, blit=False)
    plt.show()
    return ani


def basic_animate2(results, n):
    """The second basic animation function."""
    N_f = len(results)
    # extract data points
    pos = extract_pos(results, n)
    hp = extract_hp(results, n)
    T = extract_team(results, n)

    T_uniq = np.unique(T)
    plt.rcParams["animation.html"] = "html5"
    # define a figure
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111)
    # define boundaries - we use a function
    set_boundary(ax, pos)

    # create a colormap to map our team to a color - use the colorwheel
    cm = dict(zip(T_uniq, colorwheel()))

    # set the objects to draw to - we set them into lists.
    points_alive = toflat(
        [ax.plot([], [], 'o', color=cm[a], alpha=.6, markersize=10.) for a in T_uniq])
    points_dead = toflat(
        [ax.plot([], [], 'x', color=cm[a], alpha=.4, markersize=5.) for a in T_uniq])

    fig.tight_layout()
    # close the plot
    plt.close()

    # define init
    def _init():
        for j, a in enumerate(T_uniq):
            T_Alive = np.argwhere((hp[:, 0] > 0) & (T[:, 0] == a)).flatten()
            T_Dead = np.argwhere((hp[:, 0] <= 0) & (T[:, 0] == a)).flatten()

            if len(T_Alive) > 0:
                points_alive[j].set_data(pos[0, 0, T_Alive], pos[0, 1, T_Alive])
            if len(T_Dead) > 0:
                points_dead[j].set_data(pos[0, 0, T_Dead], pos[0, 1, T_Dead])
        return ((*points_alive, *points_dead))

    # define animate
    def _animate(i):
        for j, a in enumerate(T_uniq):
            T_Alive = np.argwhere((hp[:, i] > 0) & (T[:, i] == a)).flatten()
            T_Dead = np.argwhere((hp[:, i] <= 0) & (T[:, i] == a)).flatten()

            if len(T_Alive) > 0:
                points_alive[j].set_data(pos[i, 0, T_Alive], pos[i, 1, T_Alive])
            if len(T_Dead) > 0:
                points_dead[j].set_data(pos[i, 0, T_Dead], pos[i, 1, T_Dead])

        return ((*points_alive, *points_dead))

    ani = animation.FuncAnimation(fig, _animate, init_func=_init, interval=100, frames=N_f, blit=False)
    plt.show()
    return ani


def basic_animate3(results, n):
    """The third basic animation function."""
    N_f = len(results)
    # extract data points
    pos = extract_pos(results, n)
    dpos = extract_dpos(results, n)
    hp = extract_hp(results, n)
    T = extract_team(results, n)

    T_uniq = np.unique(T)

    plt.rcParams["animation.html"] = "html5"

    # define a figure
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111)
    # define boundaries - we use a function
    set_boundary(ax, pos)

    # create a colormap to map our team to a color - use the colorwheel
    cm = dict(zip(T_uniq, colorwheel()))

    # set the objects to draw to - we set them into lists.
    arrow_alive = []

    for j, a in enumerate(T_uniq):
        T_Alive0 = np.argwhere((hp[:, 0] > 0) & (T[:, 0] == a)).flatten()

        alive = ax.quiver(pos[0, 0, T_Alive0], pos[0, 1, T_Alive0],
                          dpos[0, 0, T_Alive0], dpos[0, 1, T_Alive0],
                          color=cm[a], alpha=.5, scale=30,
                          width=0.015, pivot="mid")
        arrow_alive.append(alive)

    points_dead = toflat(
        [ax.plot([], [], 'x', color=cm[a], alpha=.4, markersize=5.) for a in T_uniq])

    # figure things
    fig.tight_layout()
    # close the plot
    plt.close()

    # define init
    def _init():
        for _j, _a in enumerate(T_uniq):
            T_Alive = np.argwhere((hp[:, 0] > 0) & (T[:, 0] == _a)).flatten()
            T_Dead = np.argwhere((hp[:, 0] <= 0) & (T[:, 0] == _a)).flatten()

            if len(T_Alive) > 0:
                arrow_alive[_j].set_UVC(dpos[0, 0, T_Alive], dpos[0, 1, T_Alive])
            if len(T_Dead) > 0:
                points_dead[_j].set_data(pos[0, 0, T_Dead], pos[0, 1, T_Dead])
        return ((*arrow_alive, *points_dead))

    # define animate
    def _animate(i):
        for _j, _a in enumerate(T_uniq):
            # get IDs of alive and dead units.
            T_Alive = np.argwhere((hp[:, i] > 0) & (T[:, i] == _a)).flatten()
            T_Dead = np.argwhere((hp[:, i] <= 0) & (T[:, i] == _a)).flatten()

            if len(T_Alive) > 0:
                # we #'set the offsets' i.e arrow position.
                arrow_alive[_j].set_offsets(pos[i, :, T_Alive])
                # we set the directions of the arrows
                arrow_alive[_j].set_UVC(dpos[i, 0, T_Alive], dpos[i, 1, T_Alive])
            if len(T_Dead) > 0:
                points_dead[_j].set_data(pos[i, 0, T_Dead], pos[i, 1, T_Dead])

        return ((*arrow_alive, *points_dead))

    ani = animation.FuncAnimation(fig, _animate, init_func=_init, interval=100, frames=N_f, blit=False)
    plt.show()
    return ani


"""################################# EXTRACT FUNCTIONS ###################################"""


def extract_pos(results, n):
    """Extracting positions from results array."""
    return np.stack(([np.array([results[i][j].pos for i in range(len(results))]) for j in range(n)]), axis=2)


def extract_dpos(results, n):
    """Extracting change in positions from results array."""
    return np.stack(([np.array([results[i][j].dd for i in range(len(results))]) for j in range(n)]), axis=2)


def extract_hp(results, n):
    """Extracting HP from results array."""
    return np.vstack(([np.array([results[i][j].hp for i in range(len(results))]) for j in range(n)]))


def extract_team(results, n):
    """Extracting team number from results array."""
    return np.vstack(([np.array([results[i][j].team for i in range(len(results))]) for j in range(n)]))


"""########################## MATH FUNCTIONS ##############################################"""


def euclidean(u_i, u_j):
    """Determines the euclidean distance between ui and uj."""
    return np.sqrt(np.sum(np.square(u_i - u_j)))


def magnitude(p):
    """Determines the magnitude of a vector."""
    return np.sqrt(np.dot(p, p))


def deriv_norm(u_i, u_j):
    """Determines the normed-derivative between two vectors."""
    return (u_j - u_i) / euclidean(u_i, u_j)


def dudt(u_i, u_j, s_i):
    """Calculates the change between two vectors across time."""
    return s_i * 1. * 1. * deriv_norm(u_i, u_j)


"""####################### AI FUNCTION #######################################"""


def find_random_enemy(u, units):
    """Where u is the chosen unit, units is a list of all units."""
    # generate list of alive enemies
    alive_enemies = [i for i in units if (i.hp > 0) and (i.team != u.team)]
    # hang on, what if there are no alive enemies?
    if len(alive_enemies) > 0:
        return random.choice(alive_enemies)
    else:
        return -1


def find_nearest_enemy(u, units):
    """Finds the nearest enemy for a chosen unit."""
    # get the alive units, and calculate the euclidean distance from them.
    alive_enemies = [i for i in units if (i.hp > 0) and (i.team != u.team)]
    if len(alive_enemies) > 0:
        D = np.asarray([euclidean(u.pos, i.pos) for i in alive_enemies])
        return alive_enemies[np.argmin(D)]
    else:
        return -1


"""###################### IMPORT ################################"""


def check_columns(df, list_of_columns):
    for l in list_of_columns:
        if l not in df.columns:
            raise ValueError("column '{}' not found in dataframe.".format(l))


def import_db_file(fpath):
    """Imports our dataset of information, given a file path."""
    accepted_db_columns = (
        "Name", "Allegiance", "HP", "Damage", "Accuracy",
        "Miss", "Movement Speed", "Range"
    )

    if not os.path.isfile(fpath):
        raise IOError("fpath does not exist.")
    # attempt to read in
    df = pd.read_csv(fpath)
    # check columns exist
    check_columns(df, accepted_db_columns)
    df.set_index("Name", inplace=True)
    # assign an integer to allegiance
    df["allegiance_int"] = pd.factorize(df["Allegiance"])[0]
    return df


"""################### PLOTTING FUNCTIONS ################################"""


def plot_quiver(pos, dpos, legend_cap=5):
    """Basic plotting function for single quiver."""
    for i in range(pos.shape[-1]):
        if pos.shape[-1] > legend_cap:
            plt.scatter(pos[:, 0, i], pos[:, 1, i])
        else:
            plt.scatter(pos[:, 0, i], pos[:, 1, i], label="unit%d" % (i + 1))
        plt.quiver(pos[:, 0, i], pos[:, 1, i], dpos[:, 0, i], dpos[:, 1, i])

    plt.xlabel("x")
    plt.ylabel("y")
    if pos.shape[-1] <= legend_cap:
        plt.legend()
    plt.show()


"""####################### GROUP FUNCTIONS ########################################"""


def create_unit_group(df, u_type, n, pos_params):
    """
    Given a unit type, database, n sample and position params, create a unit group blob
    using gaussian.

    pos_params should be a tuple (2,) for mean and var.

    we cannot however assign enemies without a list to draw from, so this might have
    to be assigned after this function.
    """
    return [Unit1(df, u_type).set_position(pos_params) for _ in range(n)]


def assign_enemies(units):
    """Every unit gets assigned a random enemy."""
    # get team ids
    for u in units:
        u.target = find_random_enemy(u, units)


"""######################## MISC #######################################"""


def time_n(method, r=5, *args, **kws):
    """
    r for repeats
    """
    T = np.zeros(r + 1, )
    for i in range(r + 1):
        ts = time.time()
        _ = method(*args, **kws)
        T[i] = time.time() - ts
    # drop the first example as it is always slowest
    return T[1:]


def toflat(L):
    """
    Given nd-list, flatten to single dimension.
    """
    return list(it.chain.from_iterable(L))


def colorwheel():
    return ("red", "blue", "green", "orange", "purple", "brown", "black",
            "cyan", "yellow")
