"""
Entry Functions
===============

Entry functions need to be defined for each different entry, i.e.,
'xtal' type needs xtal_save() and xtal_update() defined and routed
properly.
"""

# Other imports
import numpy as np
import matplotlib.pyplot as plt
plt.style.use(['ggplot', 'presentation'])
from flask import session, request

# Project imports
from traveler import app
from traveler.backend import unique_filename, process_entry, save_entry
from traveler.backend import process_update, fetch_file, cache_plot


def xtal_save(data):
    """Save the info to the database. Create new if non-existant."""
    # Customize code for the entry type here.
    # Hardcoded variables
    trans_idx = 598

    fig = plt.figure()
    plt.title('xtal ' + data['sn'] + ' transmission plot')

    trans_data = []
    # For the transmission data, save the file to uploads and make a plot.
    for i, key in enumerate(session['info']['attr']['get_file']):

        file = request.files[key]

        file.seek(0)
        x1 = np.genfromtxt(file, skip_header=1, delimiter='\t')[:, :2]

        numkey = key + ' value'
        data['opt_num'].append(numkey)
        data[numkey] = x1[trans_idx, 1]

        # Only plot and include in average if not reference value.
        if 'ref' not in numkey:
            trans_data.append(data[numkey])
            plt.plot(x1[:, 0], x1[:, 1], label=key)

    # Get the average transmission
    numkey = 'longitudinal trans avg value'
    data['opt_num'].append(numkey)
    data[numkey] = round(np.array(trans_data).mean(), 3)

    # Label the plot
    plt.legend(loc=4)
    plt.xlim([240, 800])
    plt.ylim([0, 200])
    plt.grid()
    plt.xlabel('wavelength [nm]')
    plt.ylabel('transmission [%]')

    # Now save the plot.
    filename = unique_filename('trans_plot.png')
    data['plots']['trans'] = filename

    cache_plot(fig, filename)


def xtal_update(data):

    """Update the database document for the current entry."""
    # A few hardcoded constants
    trans_idx = 598

    # Now deal with recreating the figure.
    fig = plt.figure()
    plt.title('xtal ' + data['sn'] + ' transmission plot')
    trans_data = []

    # For the transmission data, save the file to uploads and make a plot.
    for i, key in enumerate(session['info']['attr']['get_file']):

        # Check if there is an updated file.
        file = request.files[key]

        if file.filename == '':
            file = fetch_file(data, data[key])

        file.seek(0)
        x1 = np.genfromtxt(file, skip_header=1, delimiter='\t')[:, :2]

        numkey = key + ' value'
        data['opt_num'].append(numkey)
        data[numkey] = x1[trans_idx, 1]

        # Only plot and include in average if not reference value.
        if 'ref' not in numkey:
            trans_data.append(data[numkey])
            plt.plot(x1[:, 0], x1[:, 1], label=key)

    # Get the average transmission
    numkey = 'longitudinal trans avg value'
    data['opt_num'].append(numkey)
    data[numkey] = round(np.array(trans_data).mean(), 3)

    # Label the plot
    plt.legend(loc=4)
    plt.xlim([240, 800])
    plt.ylim([0, 200])
    plt.grid()
    plt.xlabel('wavelength [nm]')
    plt.ylabel('transmission [%]')

    # Now save the plot.
    filename = unique_filename('trans_plot.png')
    data['plots']['trans'] = filename

    # Save the data.
    cache_plot(fig, filename)


@app.route('/<string:entry>/save', methods=['GET', 'POST'])
def general_save(entry):

    data = process_entry()

    # Check for form errors.
    if session['error'] is not None:
        return save_entry(data)

    if entry == 'xtal':
        xtal_save(data)

    # Put cust code here.
    return save_entry(data)


@app.route('/<string:entry>/save_update', methods=['GET', 'POST'])
def general_update(entry):

    data = process_update()

    if entry == 'xtal':
        xtal_update(data)

    return save_entry(data)
