Simple script to make plots with damage of bolt spells

TomeNetSpells.py

import matplotlib.pyplot as plt
import numpy as np


class Spell:
    """Let's make some class to organize spells, so maybe in future i will define all bolt spell here"""
    def __init__(self,
                 name,
                 mana_cost,
                 beg_lvl,
                 lvl_limit,
                 dmg_modif,
                 switch=0):
        self.name = name
        self.max_level = floor(
            (50 - beg_lvl[0] + 1) * 1.4
        )  # calculate max lvl of spell with maxed school and spell-power
        self.mana_cost = mana_cost
        self.beg_lvl = beg_lvl
        self.lvl_limit = lvl_limit
        self.x = np.arange(
            self.max_level) + 1  # we always start at zero, but we want at one.
        self.lvls, y = np.meshgrid(
            self.x, np.arange(3)
        )  # need array/matrix with identical vectors in each row, y is leftover
        x = np.copy(self.lvls)  # we will need this x for manathrust
        self.dmg_modif = dmg_modif
        y = np.zeros_like(
            self.lvls)  # make y the same shape like lvls matrix, but with 0
        self.switch = switch  # we will check if manathrust is in play
        self.limit = [
            beg_lvl[i] - beg_lvl[0] + 1 for i in range(len(beg_lvl))
        ]  # one liner list with for
        for i in range(len(self.lvl_limit)):
            y[i, :] = self.lvl_limit[
                i]  # change y so each row have different values.
        self.lvls[self.lvl_limit != 0] = y[self.lvl_limit != 0] + (
            self.lvls[self.lvl_limit != 0] - y[self.lvl_limit != 0]
        ) / 3  # not very nice solution with boolean indexing
        if self.switch == 0:
            self.dice = [
                self.dmg_modif[0] * self.lvls + self.dmg_modif[1],
                self.dmg_modif[2] * self.lvls + self.dmg_modif[3]
            ]
        if self.switch == 1:
            self.dice = [
                self.dmg_modif[0] * self.lvls + self.dmg_modif[1],
                self.dmg_modif[2] * x + self.dmg_modif[3]
            ]
        # up was calculated lvl the dice number and type, dice matrix is 3D
        self.dice = np.floor(self.dice)  # we don't like float, we want ints
        self.dmg = self.dice[0] * (self.dice[1] +
                                   1) / 2  # calculate actual damage
        for i in range(len(self.limit)):
            self.dmg[i, self.x < self.limit[
                i]] = 0  # set dmg array to zero when we have to low skill

    def plot(self):
        plt.close('all')  # make sure all other plots are close
        f, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(16, 10))
        # f is figure, ax1 and ax2 are subplots, with share x axis, and i like 16:10 ratio, size in inches
        for i in range(len(self.lvl_limit)):  # draw plots for every tier
            ax1.plot(self.x,
                     self.dmg[i],
                     label='{} {}'.format(self.name, i + 1))
        # use format for string if text have variables in it. don't use old % formatting
        ax1.set_xlim(
            1
        )  # we want to start with x axis at 1, not zero, at zero we cannot cast a spell
        ax1.set_ylim(0)  # don't show me dmg below zero
        ax1.grid()  # show grid
        ax1.legend()  # show labels, must have some label with data
        ax1.set_ylabel('Avg damage')
        ax1.axvline(
            50 - self.beg_lvl[0] + 1, color='r'
        )  # vertical line when we have maxed school, and 0 in spell-power
        ax1.set_title('Avg damage of {}'.format(self.name))

        # below is second subplot
        for i in range(len(self.lvl_limit)):
            ax2.plot(self.x,
                     self.dmg[i] / self.mana_cost[i],
                     label='{} {}'.format(self.name, i + 1))
        ax2.set_ylim(0)
        ax2.grid()
        ax2.legend()
        ax2.set_ylabel('Avg damage / mana cost')
        ax2.set_xlabel('Spell lvl')
        ax2.axvline(50 - self.beg_lvl[0] + 1, color='r')
        ax2.set_title('Mana efficiency of {}'.format(self.name))

        f.tight_layout()  # less spaced wasted on margins
        f.set_dpi(300)  # dpi of plot
        f.savefig(
            '{}.png'.format(self.name)
        )  # save plot in current catalog, with user provided name earlier


# uncomment below to save data for fire
'''
fn='Fire Bolt'
fm=np.array((3,6,12))
fb=np.array((10,25,40))
fll=np.array((1,15,0))
fd=(.6,5,.5,8)
fire=Spell(fn,fm,fb, fll, fd) # we don't need to provide switch, have def value
fire.plot()
'''

# user interactive part
name = input('Spell name: ')
mana_cost = input('Mana cost for each tier, separated by comma: ')
beg_lvl = input('First lvl for each tier: ')
lvl_limit = input('Limits lvl by tier: ')
print('Input damage modifiers (a * lvl + b )dices(c * lvl + e),')
dmg_modif = input(
    'Input in form of "a,b,c,e" for acidbolt like this "0.6,6,0.5,9" : ')
dmg_modif = eval('(' + dmg_modif + ')')
mana_cost = np.array(eval('(' + mana_cost + ')'))
beg_lvl = np.array(eval('(' + beg_lvl + ')'))
lvl_limit = np.array(eval('(' + lvl_limit + ')'))
sw = input(
    'Did spell use actuall lvl for type of dice, i.e. Manathrust [Y/N]:')
# i will not care if user don't provide correct output, at worse it will close at plot method with some error
if sw.upper() == 'Y':
    switch = 1
elif sw.upper() == 'N':
    switch = 0

first_spell = Spell(name, mana_cost, beg_lvl, lvl_limit, dmg_modif,
                    switch)  # make instance with user defined values
first_spell.plot()  # make and save plot
input(
    'waiting…'
)  # don't close window when we stop, wait for close or other enter/return