#!/usr/bin/python
import logging
import sys
import math
import os
import re
import copy
import json
import numpy as np
import argparse
import h5py

from allensdk.internal.core.lims_pipeline_module import PipelineModule
from allensdk.core.nwb_data_set import NwbDataSet

    
def main(jin):
    # load QC criteria and sweep table from input json file
    try:
        qc_criteria = jin['ephys_qc_criteria']
        experiment_data = jin['experiment_data']
        sweep_data = jin['sweep_data']
        nwb_file = jin["nwb_file"]
    except:
        raise IOError("Input json file is missing requisite data")

    jout = {}

    # PBS-333
    # C1NSSEED stimuli have many instances, but these instances aren't
    #   stored with the sweep. Ie, the sweep stores the stimulus value
    #   C1NSSEED while the stimulus table stores C1NSSEED_2150112.
    # To address this, check the stimulus table for any instance of
    #   C1NSSEED, and if it exists, append a plain-jane "C1NSSEED" stimulus
    #   so later checks work
    for name in jin["current_clamp_stimuli"]:
        if name.startswith("C1NSSEED_"):
            jin["current_clamp_stimuli"].append("C1NSSEED")

    # list of reasons to fail entire cell. if anything is added to this list,
    #   the cell will be tagged for failure (ie, this list also serves
    #   as a 'fail' flag)
    exp_fail_tags = []

    experiment_state = {}
    jout["experiment_data"] = experiment_state

    # blowout voltage
    experiment_state["failed_blowout"] = False
    try:
        blowout = experiment_data["blowout_mv"]
        low = qc_criteria["blowout_mv_min"]
        high = qc_criteria["blowout_mv_max"]
        if blowout is None or math.isnan(blowout):
            exp_fail_tags.append("Missing blowout value (%s)" % str(blowout))
            experiment_state["failed_blowout"] = True
        if blowout < low or blowout > high:
            exp_fail_tags.append("blowout outside of range")
            experiment_state["failed_blowout"] = True
    except Exception, e:
        exp_fail_tags.append("Error analyzing blowout. " + e.message)
        experiment_state["failed_blowout"] = True


    # "electrode 0"
    experiment_state["failed_electrode_0"] = False
    try:
        e0 = experiment_data["electrode_0_pa"]
        if e0 is None or math.isnan(e0):
            exp_fail_tags.append("e0 -- missing value (%s)" % str(e0))
            experiment_state["failed_electrode_0"] = True
        if abs(e0) > qc_criteria["electrode_0_pa_max"]:
            exp_fail_tags.append("e0 -- exceeds max")
            experiment_state["failed_electrode_0"] = True
    except Exception, e:
        exp_fail_tags.append("Error analyzing blowout. " + e.message)
        experiment_state["failed_electrode_0"] = True


    # measure clamp seal
    experiment_state["failed_seal"] = False
    try:
        seal = experiment_data["seal_gohm"]
        if seal is None or math.isnan(seal):
            exp_fail_tags.append("Invalid seal (%s)" % str(seal))
            experiment_state["failed_seal"] = True
        if seal < qc_criteria["seal_gohm_min"]:
            tgt = qc_criteria["seal_gohm_min"]
            reason = "%f versus criteria=%f" % (seal, tgt)
            exp_fail_tags.append("Seal (%s)" % reason)
            experiment_state["failed_seal"] = True
    except Exception, e:
        seal = None
        msg = "Seal 0 is not available. %s" % e.message
        logging.warning(msg)
        exp_fail_tags.append(msg)
        experiment_state["failed_seal"] = True


    # input and access resistance
    sr_tags = []

    try:
        sir_ratio = experiment_data['input_access_resistance_ratio']
        #r = experiment_data['input_resistance_mohm']
    except:
        sr_tags.append("Resistance ratio not available")

    try:
        sr = experiment_data['initial_access_resistance_mohm']
    except:
        sr_tags.append("Initial access resistance not available")

    try:
        if len(sr_tags) == 0:
            experiment_state["failed_bad_rs"] = False

            if sr < qc_criteria["access_resistance_mohm_min"]:
                experiment_state["failed_bad_rs"] = True
                tgt = qc_criteria["access_resistance_mohm_min"]
                reason = "%f versus criteria=%f" % (sr, tgt)
                sr_tags.append("access-resistance low (%s)" % reason)
            elif sr > qc_criteria["access_resistance_mohm_max"]:
                experiment_state["failed_bad_rs"] = True
                tgt = qc_criteria["access_resistance_mohm_max"]
                reason = "%f versus criteria=%f" % (sr, tgt)
                sr_tags.append("access-resistance high (%s)" % reason)

            if sir_ratio > qc_criteria["input_vs_access_resistance_min"]:
                experiment_state["failed_bad_rs"] = True
                tgt = qc_criteria["input_vs_access_resistance_min"]
                reason = "%f versus criteria=%f" % (sir_ratio, tgt)
                sr_tags.append("input/access resistance (%s)" % reason)
    except Exception, e:
        exp_fail_tags.append("Error analyzing access resistance. " + e.message)

    if len(sr_tags) > 0:
        exp_fail_tags.extend(sr_tags)


    experiment_state["fail_tags"] = exp_fail_tags


    ####################################################################
    # check features for each sweep
    sweep_state = {}
    jout["sweep_state"] = sweep_state
    for name, sweep in jin["sweep_data"].iteritems():
        try:
            # keep track of failures
            fail_tags = []

            sweep_num = sweep["sweep_number"]

            stim = sweep["ephys_stimulus"]["description"]
            if stim.endswith("_DA_0"):
                stim = stim[:-5]
            unit = sweep["stimulus_units"]
            # determine if sweep is current or voltage clamp
            # name may end in "[#]", so strip out section after open bracket
            stim_short = stim.split('[')[0]
            if stim_short in jin["voltage_clamp_stimuli"]:
                if unit != "Volts" and unit != "mV":
                    msg = "%s (%s) in wrong mode -- expected voltage clamp" % (name, stim)
                    fail_tags.append(msg)
            elif stim_short in jin["current_clamp_stimuli"]:
                if unit != "Amps" and unit != "pA":
                    msg = "%s (%s) in wrong mode -- expected current clamp" % (name, stim)
                    fail_tags.append(msg)
            else:
                fail_tags.append("%s has unrecognized stimulus (%s)" % (name, stim))

            if unit == "Volts" or unit == "mV":
                continue    # no QC on voltage clamp

            if len(fail_tags) > 0:
                sweep_state[name] = {}
                sweep_state[name]["state"] = "Fail"
                sweep_state[name]["reasons"] = fail_tags
                continue
            
            # pull data streams from file (this is for detecting truncated
            #   sweeps)
            sweep_data = NwbDataSet(nwb_file).get_sweep(sweep_num)
            volts = sweep_data['response']
            current = sweep_data['stimulus']
            hz = sweep_data['sampling_rate']
            idx_start, idx_stop = sweep_data['index_range']

            if sweep["pre_noise_rms_mv"] > qc_criteria["pre_noise_rms_mv_max"]:
                fail_tags.append("pre-noise")

            # check Vm and noise at end of recording
            # only do so if acquisition not truncated 
            # do not check for ramps, because they do not have 
            #   enough time to recover
            is_ramp = stim.startswith('C1RP')
            if is_ramp:
                logging.info("sweep %d skipping vrest criteria on ramp", sweep_num)
            else:
                # measure post-stimulus noise
                sweep_not_truncated = ( idx_stop == len(current) - 1 )
                if sweep_not_truncated:
                    post_noise_rms_mv = sweep["post_noise_rms_mv"]
                    if post_noise_rms_mv > qc_criteria["post_noise_rms_mv_max"]:
                        fail_tags.append("post-noise")
                else:
                    fail_tags.append("Truncated sweep")

            if sweep["slow_noise_rms_mv"] > qc_criteria["slow_noise_rms_mv_max"]:
                fail_tags.append("slow noise above threshold")

            if sweep["vm_delta_mv"] > qc_criteria["vm_delta_mv_max"]:
                fail_tags.append("Vm delta")


            # fail sweeps if stimulus duration is zero
            # Uncomment out hte following 3 lines to have sweeps without stimulus
            #   faile QC
            if sweep["stimulus_duration"] <= 0:
                desc = sweep["ephys_stimulus"]["description"]
                if not desc.startswith("EXTP"):
                    fail_tags.append("No stimulus detected")


            sweep_state[name] = {}
            if len(fail_tags) > 0:
                sweep_state[name]["state"] = "Fail"
                sweep_state[name]["reasons"] = fail_tags
            else:
                sweep_state[name]["state"] = "Pass"

        except:
            print("Error processing sweep %s" % name)
            raise

    ####################################
    # done - prepare and deliver results
    if len(exp_fail_tags) > 0:
        jout["qc_result"] = "failed"
    else:
        jout["qc_result"] = "passed"

    return jout

if __name__ == "__main__": 
    # read module input. PipelineModule object automatically parses the 
    #   command line to pull out input.json and output.json file names
    module = PipelineModule()
    jin = module.input_data()   # loads input.json
    jout = main(jin)
    module.write_output_data(jout)  # writes output.json

