#!/usr/bin/python
import sys, logging
import os
import json
import shutil
import argparse
import copy
import numpy as np
import shutil

from allensdk.config.manifest import Manifest

import allensdk.internal.core.lims_utilities as lims_utilities
import allensdk.core.json_utilities as json_utilities

from allensdk.internal.ephys.core_feature_extract import *
from allensdk.ephys.ephys_features import FeatureError


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('input_json')
    parser.add_argument('output_json')
    parser.add_argument('--log_level')
    parser.add_argument('--output_directory')
    
    args = parser.parse_args()

    if args.log_level:
        logging.getLogger().setLevel(args.log_level)

    return args

def main():
    args = parse_args()

    input_data = json_utilities.read(args.input_json)
    input_err, stim_types = input_data

    output_data = copy.deepcopy(input_err)

    err_wkfs = output_data['well_known_files']
    nwb_file = lims_utilities.get_well_known_file_by_type(err_wkfs, lims_utilities.NWB_FILE_TYPE_ID)
    storage_directory = args.output_directory or output_data['storage_directory']
    # move code to help make data extraction compatible with ephys qc tool
    try:
        sweep_list, sweep_features = extract_data(output_data, nwb_file)
    except FeatureError as e:
        logging.error("Error computing cell features, auto-failing cell: %s" % e.message)
        output_data["workflow_state"] = "auto_failed"
        json_utilities.write(args.output_json, output_data)
        return
    #
    
    # embed spike times in NWB file
    logging.debug("Embedding spike times")
    tmp_nwb_file = os.path.join(storage_directory, os.path.basename(nwb_file) + '.tmp')
    out_nwb_file = os.path.join(storage_directory, os.path.basename(nwb_file))

    shutil.copy(nwb_file, tmp_nwb_file)
    for sweep in sweep_list:
        sweep_num = sweep['sweep_number']

        if sweep_num not in sweep_features:
            continue

        try:
            spikes = sweep_features[sweep_num]['spikes']
            spike_times = [ s['threshold_t'] for s in spikes ]
            NwbDataSet(tmp_nwb_file).set_spike_times(sweep_num, spike_times)
        except Exception as e:
            logging.info("sweep %d has no sweep features. %s", sweep_num, e.message)

    try:
        shutil.move(tmp_nwb_file, out_nwb_file)
    except OSError as e:
        logging.error("Problem renaming file: %s -> %s" % (tmp_nwb_file, out_nwb_file))
        raise e

    qc_fig_dir = os.path.join(storage_directory, 'qc_figures')
    save_qc_figures(qc_fig_dir, nwb_file, output_data, True)

    # regenerating this file
    features_json = os.path.join(storage_directory, "%d_ephys_features.json" % output_data['id'])
    json_utilities.write(features_json, output_data)
    lims_utilities.append_well_known_file(output_data['well_known_files'], features_json)

    # write output json files
    json_utilities.write(args.output_json, output_data)

            
if __name__ == "__main__": 
    main()
