import os
import sys
import allensdk.internal.model.biophysical.ephys_utils as ephys_utils
import check_fi_shift
import pandas as pd
import numpy as np
from collections import Counter
import subprocess

from allensdk.ephys.ephys_extractor \
    import EphysSweepFeatureExtractor, EphysSweepSetFeatureExtractor
import allensdk.core.json_utilities as ju
from allensdk.core.nwb_data_set import NwbDataSet
import allensdk.internal.model.biophysical.optimize as optimize
import logging

SEEDS = [1234, 1001, 4321, 1024, 2048]
FIT_BASE_DIR = os.path.join(os.path.dirname(__file__), "fits")
APICAL_DENDRITE_TYPE = 4
MPIEXEC = 'mpiexec'
DEFAULT_NUM_PROCESSES = 240

_fit_stage_1_log = logging.getLogger('allensdk.model.biophysical.fit_stage_1')

def find_core1_trace(data_set, all_sweeps):
    sweep_type = "C1LSCOARSE"
    _, sweeps, statuses = ephys_utils.get_sweeps_of_type(sweep_type, all_sweeps)
    sweep_status = dict(zip(sweeps, statuses))

    sweep_info = {}
    for s in sweeps:
        if sweep_status[s][-6:] == "failed":
            continue
        v, i, t = ephys_utils.get_sweep_v_i_t_from_set(data_set, s)
        if np.all(v[-100:] == 0): # Check for early termination of sweep
            continue
        stim_start, stim_dur, stim_amp, start_idx, end_idx = ephys_utils.get_step_stim_characteristics(i, t)
        swp = EphysSweepFeatureExtractor(t, v, i, start=stim_start, end=(stim_start + stim_dur))
        swp.process_spikes()
        isi_cv = swp.sweep_feature("isi_cv", allow_missing=True)
        sweep_info[s] = {"amp": stim_amp,
                         "n_spikes": len(swp.spikes()),
                         "quality": is_trace_good_quality(v, i, t),
                         "isi_cv": isi_cv}
    
    rheobase_amp = 1e12
    for s in sweep_info:
        if sweep_info[s]["amp"] < rheobase_amp and sweep_info[s]["n_spikes"] > 0:
            rheobase_amp = sweep_info[s]["amp"]
    sweep_to_use_amp = 1e12
    sweep_to_use_isi_cv = 1e121    
    sweep_to_use = -1
    for s in sweep_info:
        if sweep_info[s]["quality"] and sweep_info[s]["amp"] >= 39.0 + rheobase_amp and sweep_info[s]["isi_cv"] < 1.2 * sweep_to_use_isi_cv:
            use_new_sweep = False
            if sweep_to_use_isi_cv > 0.3 and ((sweep_to_use_isi_cv - sweep_info[s]["isi_cv"]) / sweep_to_use_isi_cv) >= 0.2:
                use_new_sweep = True
            elif sweep_info[s]["amp"] < sweep_to_use_amp:
                use_new_sweep = True
            if use_new_sweep:
                _fit_stage_1_log.info("now using sweep" + str(s))
                sweep_to_use = s
                sweep_to_use_amp = sweep_info[s]["amp"]
                sweep_to_use_isi_cv = sweep_info[s]["isi_cv"]
            
    if sweep_to_use == -1:
        _fit_stage_1_log.warn("Could not find appropriate core 1 sweep!")
        return []
    else:
        return [sweep_to_use]

def find_core2_trace(data_set, all_sweeps):
    sweep_type = "C2SQRHELNG"
    _, sweeps, statuses = ephys_utils.get_sweeps_of_type(sweep_type, all_sweeps)
    sweep_status = dict(zip(sweeps, statuses))
    amp_list = []
    core2_amps = {}
    for s in sweeps:
        if sweep_status[s][-6:] == "failed":
            continue
        v, i, t = ephys_utils.get_sweep_v_i_t_from_set(data_set, s)
        if np.all(v[-100:] == 0):
            continue
        stim_start, stim_dur, stim_amp, start_idx, end_idx = ephys_utils.get_step_stim_characteristics(i, t)
        if stim_start is np.nan:
            sweep_status[s] = "manual_failed"
            continue
        core2_amps[s] = stim_amp
        amp_list.append(stim_amp)
    core2_amp_counter = Counter(amp_list)
    common_amps = core2_amp_counter.most_common(3)
    best_amp = 0
    best_n_good = -1
    for amp, _ in common_amps:
        n_good = 0
        for k in core2_amps:
            if core2_amps[k] == amp and sweep_status[k][-6:] == "passed":
                v, i, t = ephys_utils.get_sweep_v_i_t_from_set(data_set, k)
                if is_trace_good_quality(v, i, t):
                    n_good += 1
        if n_good > best_n_good:
            best_n_good = n_good
            best_amp = amp
        elif n_good == best_n_good and amp < best_amp:
            best_amp = amp
    if best_n_good <= 1:
        return []
    else:
        sweeps_to_fit = []
        for k in core2_amps:
            if core2_amps[k] == best_amp and sweep_status[k][-6:] == "passed":
                sweeps_to_fit.append(k)
        return sweeps_to_fit

def is_trace_good_quality(v, i, t):
    stim_start, stim_dur, stim_amp, start_idx, end_idx = ephys_utils.get_step_stim_characteristics(i, t)
    swp = EphysSweepFeatureExtractor(t, v, i, start=stim_start, end=(stim_start + stim_dur))
    swp.process_spikes()

    spikes = swp.spikes()
    rate = swp.sweep_feature("avg_rate")

    if rate < 5.0:
        return False

    time_to_end = stim_start + stim_dur - spikes[-1]["threshold_t"]
    avg_end_isi = (((spikes[-1]["threshold_t"] - spikes[-2]["threshold_t"]) +
                   (spikes[-2]["threshold_t"] - spikes[-3]["threshold_t"])) / 2.0)

    if time_to_end > 2 * avg_end_isi:
        return False

    isis = np.diff([spk["threshold_t"] for spk in spikes])
    if check_for_pause(isis):
        return False

    return True


def check_for_pause(isis):
    if len(isis) <= 2:
        return False
    
    for i, isi in enumerate(isis[1:-1]):
        if isi > 3 * isis[i + 1 - 1] and isi > 3 * isis[i + 1 + 1]:
            return True
    return False


def collect_target_features(ft):
    min_std_dict = {
        'avg_rate': 0.5,
        'adapt': 0.001,
        'peak_v': 2.0,
        'trough_v': 2.0,
        'fast_trough_v': 2.0,
        'slow_trough_delta_v': 2.0,
        'slow_trough_delta_time': 0.05,
        'latency': 5.0,
        'isi_cv': 0.1,
        'mean_isi': 0.5,
        'first_isi': 1.0,
        'time_to_end': 50.0,
        'v_baseline': 2.0,
        'width': 0.0001,
        'upstroke': 50.0,
        'downstroke': 50.0,
        'upstroke_v': 2.0,
        'downstroke_v': 2.0,
        'threshold_v': 2.0,
        'peak_to_fast_tr_time': 0.0005,
        'phase_slope': 5.0,
    }

    target_features = []
    for k in ft:
        t = {"name": k, "mean": ft[k]["mean"], "stdev": ft[k]["stdev"]}
        if k in min_std_dict and min_std_dict[k] > ft[k]["stdev"]:
            t["stdev"] = min_std_dict[k]
        target_features.append(t)
    return target_features


def prepare_stage_1(description, passive_fit_data):
    output_directory = description.manifest.get_path('WORKDIR')
    neuronal_model_data = ju.read(description.manifest.get_path('neuronal_model_data'))
    specimen_data = neuronal_model_data['specimen']
    specimen_id = neuronal_model_data['specimen_id']
    is_spiny = not any(t['name'] == u'dendrite type - aspiny' for t in specimen_data['specimen_tags'])
    all_sweeps = specimen_data['ephys_sweeps']
    data_set = NwbDataSet(description.manifest.get_path('stimulus_path'))
    swc_path = description.manifest.get_path('MORPHOLOGY')
    
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    ra = passive_fit_data['ra']
    cm1 = passive_fit_data['cm1']
    cm2 = passive_fit_data['cm2']

    # Check for fi curve shift to decide to use core1 or core2
    fi_shift, n_core2 = check_fi_shift.estimate_fi_shift(data_set, all_sweeps)
    fi_shift_threshold = 30.0
    sweeps_to_fit = []
    if abs(fi_shift) > fi_shift_threshold:
        _fit_stage_1_log.info("FI curve shifted; using Core 1")
        sweeps_to_fit = find_core1_trace(data_set, all_sweeps)
    else:
        sweeps_to_fit = find_core2_trace(data_set, all_sweeps)

        if sweeps_to_fit == []:
            _fit_stage_1_log.info("Not enough good Core 2 traces; using Core 1")
            sweeps_to_fit = find_core1_trace(data_set, all_sweeps)

    _fit_stage_1_log.debug("will use sweeps: " + str(sweeps_to_fit))

    jxn = -14.0

    t_set = []
    v_set = []
    i_set = []
    for s in sweeps_to_fit:
        v, i, t = ephys_utils.get_sweep_v_i_t_from_set(data_set, s)
        v += jxn
        stim_start, stim_dur, stim_amp, start_idx, end_idx = ephys_utils.get_step_stim_characteristics(i, t)
        t_set.append(t)
        v_set.append(v)
        i_set.append(i)
    ext = EphysSweepSetFeatureExtractor(t_set, v_set, i_set, start=stim_start, end=(stim_start + stim_dur))
    ext.process_spikes()

    ft = {}
    blacklist = ["isi_type"]
    for k in ext.sweeps()[0].spike_feature_keys():
        if k in blacklist:
            continue
        pair = {}
        pair["mean"] = float(ext.spike_feature_averages(k).mean())
        pair["stdev"] = float(ext.spike_feature_averages(k).std())
        ft[k] = pair

    # "Delta" features
    sweep_avg_slow_trough_delta_time = []
    sweep_avg_slow_trough_delta_v = []
    sweep_avg_peak_trough_delta_time = []
    for swp in ext.sweeps():
        threshold_t = swp.spike_feature("threshold_t")
        fast_trough_t = swp.spike_feature("fast_trough_t")
        slow_trough_t = swp.spike_feature("slow_trough_t")

        delta_t = slow_trough_t - fast_trough_t
        delta_t[np.isnan(delta_t)] = 0.
        sweep_avg_slow_trough_delta_time.append(np.mean(delta_t[:-1] / np.diff(threshold_t)))

        fast_trough_v = swp.spike_feature("fast_trough_v")
        slow_trough_v = swp.spike_feature("slow_trough_v")
        delta_v = fast_trough_v - slow_trough_v
        delta_v[np.isnan(delta_v)] = 0.
        sweep_avg_slow_trough_delta_v.append(delta_v.mean())

    ft["slow_trough_delta_time"] = {"mean": float(np.mean(sweep_avg_slow_trough_delta_time)),
                                    "stdev": float(np.std(sweep_avg_slow_trough_delta_time))}
    ft["slow_trough_delta_v"] = {"mean": float(np.mean(sweep_avg_slow_trough_delta_v)),
                                 "stdev": float(np.std(sweep_avg_slow_trough_delta_v))}

    baseline_v = float(ext.sweep_features("v_baseline").mean())
    passive_fit_data["e_pas"] = baseline_v
    for k in ext.sweeps()[0].sweep_feature_keys():
        pair = {}
        pair["mean"] = float(ext.sweep_features(k).mean())
        pair["stdev"] = float(ext.sweep_features(k).std())
        ft[k] = pair

    # Determine highest step to check for depolarization block
    noise_1_sweeps, _, _ = ephys_utils.get_sweeps_of_type("C1NSSEED_1", all_sweeps)
    noise_2_sweeps, _, _ = ephys_utils.get_sweeps_of_type("C1NSSEED_2", all_sweeps)
    step_sweeps, _, _ = ephys_utils.get_sweeps_of_type("C1LSCOARSE", all_sweeps)
    all_sweeps = noise_1_sweeps + noise_2_sweeps + step_sweeps
    max_i = 0
    for s in all_sweeps:
        try:
            v, i, t = ephys_utils.get_sweep_v_i_t_from_set(data_set, s['sweep_number'])
        except:
            pass
        if np.max(i) > max_i:
            max_i = np.max(i)
    max_i += 10 # add 10 pA
    max_i *= 1e-3 # convert to nA

    # ----------- Generate output and submit jobs ---------------

    # Set up directories
    # Decide which fit(s) we are doing
    if (is_spiny and ft["width"]["mean"] < 0.8) or (not is_spiny and ft["width"]["mean"] > 0.8):
        fit_types = ["f6", "f12"]
    elif is_spiny:
        fit_types = ["f6"]
    else:
        fit_types = ["f12"]

    for fit_type in fit_types:
        fit_type_dir = os.path.join(output_directory, fit_type)
        if not os.path.exists(fit_type_dir):
            os.makedirs(fit_type_dir)
        for seed in SEEDS:
            seed_dir = "{:s}/s{:d}".format(fit_type_dir, seed)
            if not os.path.exists(seed_dir):
                os.makedirs(seed_dir)

    # Collect and save data for target.json file
    target_dict = {}
    target_dict["passive"] = [{
        "ra": ra,
        "cm": { "soma": cm1, "axon": cm1, "dend": cm2 },
        "e_pas": baseline_v
    }]

    swc_data = pd.read_table(swc_path, sep='\s', comment='#', header=None)
    has_apic = False
    if APICAL_DENDRITE_TYPE in pd.unique(swc_data[1]):
        has_apic = True
        _fit_stage_1_log.info("Has apical dendrite")
    else:
        _fit_stage_1_log.info("Does not have apical dendrite")

    if has_apic:
        target_dict["passive"][0]["cm"]["apic"] = cm2

    target_dict["fitting"] = [{
        "junction_potential": jxn,
        "sweeps": sweeps_to_fit,
        "passive_fit_info": passive_fit_data,
        "max_stim_test_na": max_i,        
    }]

    target_dict["stimulus"] = [{
        "amplitude": 1e-3 * stim_amp,
        "delay": 1000.0,
        "duration": 1e3 * stim_dur
    }]

    target_dict["manifest"] = []
    target_dict["manifest"].append({"type": "file", "spec": swc_path, "key": "MORPHOLOGY"})

    target_dict["target_features"] = collect_target_features(ft)

    target_file = os.path.join(output_directory, 'target.json')
    ju.write(target_file, target_dict)

    # Create config.json for each fit type
    config_base_data = ju.read(os.path.join(FIT_BASE_DIR,
                                            'config_base.json'))


    jobs = []
    for fit_type in fit_types:
        config = config_base_data.copy()
        fit_type_dir = os.path.join(output_directory, fit_type)
        config_path = os.path.join(fit_type_dir, "config.json")

        config["biophys"][0]["model_file"] = [ target_file, config_path]
        if has_apic:
            fit_style_file = os.path.join(FIT_BASE_DIR, 'fit_styles', '%s_fit_style.json' % (fit_type))
        else:
            fit_style_file = os.path.join(FIT_BASE_DIR, "fit_styles", "%s_noapic_fit_style.json" % (fit_type))

        config["biophys"][0]["model_file"].append(fit_style_file)
        config["manifest"].append({"type": "dir", "spec": fit_type_dir, "key": "FITDIR"})
        ju.write(config_path, config)

        for seed in SEEDS:
            logfile = os.path.join(output_directory, fit_type, 's%d' % seed, 'stage_1.log')
            jobs.append({
                    'config_path': os.path.abspath(config_path),
                    'fit_type': fit_type,
                    'log': os.path.abspath(logfile),
                    'seed': seed,
                    'num_processes': DEFAULT_NUM_PROCESSES
                    })
    return jobs


def run_stage_1(jobs):
    for job in jobs:
        args = [MPIEXEC,
                '-np',
                str(job['num_processes']),
                sys.executable,
                '-m',
                optimize.__name__,
                str(job['seed']),
                job['config_path'],
                str(optimize.DEFAULT_NGEN),
                str(optimize.DEFAULT_MU)]
        _fit_stage_1_log.debug(args)
        with open(job['log'], "w") as outfile:
            subprocess.call(args, stdout=outfile)