# -*- coding: utf-8 -*-

from multiprocessing import Pool
import pandas as pd
from functools import partial
from os import system
from os.path import join, abspath

def call_return_series(fn, series, **kwargs):
    res = fn(series, **kwargs)
    if not isinstance(res, pd.Series):
        return pd.Series(res, name=series.name)
    else:
        return res

def apply_p(df, fn, threads=2, **kwargs):
    """Parallel apply for pandas DataFrame.

    Argument fn cannot be a lambda function, but it can be generated with
    functools.partial.

    """
    p = Pool(threads)
    index_names = df.index.names
    result = p.map(partial(call_return_series, fn, **kwargs),
                   [row for _, row in df.reset_index().iterrows()])
    return pd.DataFrame(result).set_index(index_names)

# Slurm functions for NERSC

SLURM_TEMPLATE = """#!/bin/bash
#SBATCH --partition=shared
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --time=01:00:00
"""

CMD_TEMPLATE = """
python -c "import {0},pandas; {0}.{1}(pandas.read_json('{2}', typ='series'), '{3}')"
"""

def apply_slurm(df, package_name, function_name, data_directory):
    """Parallel apply for pandas DataFrame using slurm. Does not return anything.

    Arguments
    ---------

    df: A DataFrame to iterate through.

    package_name: The name of the package containing function_name.

    function_name: The name of a function in the package given by
    package_name. The function should take two argumes: a pandas series and an
    output directory.

    data_directory: A directory to store all the data files.

    """
    index_names = df.index.names
    for _, row in df.reset_index().iterrows():
        job_name = '%s_%s' % (function_name, '_'.join(str(row[x]) for x in index_names))
        job_file = abspath(join(data_directory, '%s.sl' % job_name))
        json_file = abspath(join(data_directory, '%s_input.json' % job_name))
        output_file = abspath(join(data_directory, '%s.out' % job_name))
        row.to_json(json_file)
        with open(job_file, 'w') as outfile:
            outfile.write(SLURM_TEMPLATE)
            outfile.write('#SBATCH --output=%s\n' % output_file)
            outfile.write('#SBATCH --job-name=%s\n\n' % job_name)
            outfile.write(CMD_TEMPLATE.format(package_name, function_name,
                                              json_file, data_directory))
            outfile.write('\n')
        system('chmod +x ' + job_file)
        system('sbatch ' + job_file)

@@ -122,16 +122,26 @@
 t system
+, makedirs
 %0Afrom os
@@ -1649,24 +1649,421 @@
 index.names%0A
+    slurm_dir = abspath(join(data_directory, 'slurm'))%0A    input_json_dir = abspath(join(data_directory, 'input_json'))%0A    output_dir = abspath(join(data_directory, 'output'))%0A    results_dir = abspath(join(data_directory, 'results'))%0A    for directory in %5Bslurm_dir, input_json_dir, output_dir, results_dir%5D:%0A        try:%0A            makedirs(directory)%0A        except OSError:%0A            pass%0A
     for _, r
@@ -2209,35 +2209,22 @@
 e = 
-abspath(join(data_directory
+join(slurm_dir
 , '%25
@@ -2232,33 +2232,32 @@
 .sl' %25 job_name)
-)
 %0A        json_fi
@@ -2265,35 +2265,27 @@
 e = 
-abspath(join(data_directory
+join(input_json_dir
 , '%25
@@ -2301,33 +2301,32 @@
 son' %25 job_name)
-)
 %0A        output_
@@ -2336,35 +2336,23 @@
 e = 
-abspath(join(data_directory
+join(output_dir
 , '%25
@@ -2369,17 +2369,16 @@
 ob_name)
-)
 %0A       
@@ -2751,30 +2751,27 @@
 n_file, 
-data_directory
+results_dir
 ))%0A     

make dirs for slurm