"""
    Generic example of experiment that compares two methods
    COS, SIN in two different experiments EXPERIMENT1, EXPERIMENT2

    Experiments can be in general long so one can choose to SAVE the
    files so that the plots can be generated again when the
    appropriate RUNTAG is set and the PLOT program option is chosen
"""
import os

import matplotlib.pyplot as plt
import numpy as np


class Method:
    COS = "cosine"
    SIN = "sine"


class Program:
    EXPERIMENT1 = 1
    EXPERIMENT2 = 2
    PLOT = 3

RUNTAG = 'run1'
FILETAG = 'experiment'
OUTPUTFOLDER = 'outputs'
STEPS = 5000
STEP_SIZE = 1e-3
SAVE = True

PROGRAM = Program.PLOT


class NameGen(object):
    def __init__(self, filetag, output_folder):
        self.output_folder = output_folder
        self.filetag = filetag

    def get_output_prefix(self, exptag, label):
        return os.path.join(self.output_folder,
                            '{}_{}{}'.format(self.filetag, label, exptag))

    def get_output_filename(self, exptag, label, ext):
        return '{}.{}'.format(self.get_output_prefix(exptag, label), ext)

    def save_plot(self, tag, label='', dpi=300, ext='png'):
        plt.draw()
        plt.savefig(self.get_output_filename(tag, label, ext), dpi=dpi)


def factory(method_name, **kwargs):
    if method_name == Method.COS:
        t = kwargs.pop('term', 0)
        return lambda x: np.cos(x + t)
    elif method_name == Method.SIN:
        f = kwargs.pop('factor', 1)
        return lambda x: np.sin(f * x)
    else:
        raise ValueError("Method: '{}' is not valid.".format(method_name))


def plot_result(y, label):
    plt.plot(y, label=label)


def run(func, label):
    y = func(np.arange(STEPS) * STEP_SIZE)
    plot_result(y, label)

    if SAVE:
        namegen = NameGen(FILETAG, OUTPUTFOLDER)
        npz_filename = namegen.get_output_filename(RUNTAG, label, 'npz')
        np.savez(npz_filename, y=y)


def finalize_plot():
    plt.legend(loc='best', prop={'size': 12})

    plt.xlabel('Simulation steps')
    plt.ylabel('Output')

    namegen = NameGen(FILETAG, OUTPUTFOLDER)
    namegen.save_plot(RUNTAG, "")
    plt.clf()


def main_example1():
    for method_name in [Method.COS, Method.SIN]:
        run(factory(method_name), label=method_name)


def main_example2():
    for method_name in [Method.COS, Method.SIN]:
        run(factory(method_name, factor=2, term=1), label=method_name)


def main_plot():
    namegen = NameGen(FILETAG, OUTPUTFOLDER)

    for label in [Method.COS, Method.SIN]:
        npz_filename = namegen.get_output_filename(RUNTAG, label, 'npz')
        y = np.load(npz_filename)['y']
        plot_result(y, label)

if __name__ == '__main__':
    if PROGRAM == Program.EXPERIMENT1:
        main_example1()
    elif PROGRAM == Program.EXPERIMENT2:
        main_example2()
    elif PROGRAM == Program.PLOT:
        main_plot()

    finalize_plot()