import matplotlib.pyplot as plt
import os


def data_plotter(lattice_dict, datafile_dir, plot_dir):

    # total spaces on grid implies grid size
    total_cells = lattice_dict['E'][0] + lattice_dict['D_a'][0] + lattice_dict['D_b'][0] + lattice_dict['B'][0]
    n = int(total_cells**0.5)

    plt.figure(1)

    plt.plot(lattice_dict['time'], lattice_dict['E'], label='Empty lattice points')
    plt.plot(lattice_dict['time'], lattice_dict['D_a'], label='Donors (Type A)')
    plt.plot(lattice_dict['time'], lattice_dict['D_b'], label='Donors (Type B)')
    plt.plot(lattice_dict['time'], lattice_dict['B'], label='Debris')

    ax = plt.gca()
    ax.set_title('Cell Populations over time (n = %d)' % n)
    ax.set_ylabel('Number of cells')
    ax.set_xlabel('Time (h)')

    plt.legend()

    f = plt.gcf()
    f.set_size_inches(20.0, 8.0)  # alternative: 20.0, 8.0
    f.tight_layout()
    plt.savefig(os.path.join(plot_dir, 'population_vs_time.png'))
    plt.clf()
    
    return

[(11, 'matplotlib.pyplot.figure', 'plt.figure', (['(1)'], {}), True, 'import matplotlib.pyplot as plt\n'), (13, 'matplotlib.pyplot.plot', 'plt.plot', (["lattice_dict['time']", "lattice_dict['E']"], {'label': '"""Empty lattice points"""'}), True, 'import matplotlib.pyplot as plt\n'), (14, 'matplotlib.pyplot.plot', 'plt.plot', (["lattice_dict['time']", "lattice_dict['D_a']"], {'label': '"""Donors (Type A)"""'}), True, 'import matplotlib.pyplot as plt\n'), (15, 'matplotlib.pyplot.plot', 'plt.plot', (["lattice_dict['time']", "lattice_dict['D_b']"], {'label': '"""Donors (Type B)"""'}), True, 'import matplotlib.pyplot as plt\n'), (16, 'matplotlib.pyplot.plot', 'plt.plot', (["lattice_dict['time']", "lattice_dict['B']"], {'label': '"""Debris"""'}), True, 'import matplotlib.pyplot as plt\n'), (18, 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (23, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (25, 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (29, 'matplotlib.pyplot.clf', 'plt.clf', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (28, 'os.path.join', 'os.path.join', (['plot_dir', '"""population_vs_time.png"""'], {}), False, 'import os\n')]