"""

Tap-Withdrawal circuit still under development - it does not produce the correct behavior!

"""

import c302
import sys
import importlib

import neuroml.writers as writers

range_incl = lambda start, end: range(start, end + 1)


def setup(
    parameter_set,
    generate=False,
    duration=400,
    dt=0.05,
    target_directory="examples",
    data_reader=c302.FW_DATA_READER,
    param_overrides={},
    config_param_overrides={},
    verbose=True,
):
    ParameterisedModel = getattr(
        importlib.import_module("c302.parameters_%s" % parameter_set),
        "ParameterisedModel",
    )
    params = ParameterisedModel()

    params.set_bioparameter(
        "unphysiological_offset_current", "0pA", "Testing TapWithdrawal", "0"
    )
    params.set_bioparameter(
        "unphysiological_offset_current_del", "0 ms", "Testing TapWithdrawal", "0"
    )
    params.set_bioparameter(
        "unphysiological_offset_current_dur", "2000 ms", "Testing TapWithdrawal", "0"
    )

    VA_motors = ["VA%s" % c for c in range_incl(1, 12)]
    VB_motors = ["VB%s" % c for c in range_incl(1, 11)]
    DA_motors = ["DA%s" % c for c in range_incl(1, 9)]
    DB_motors = ["DB%s" % c for c in range_incl(1, 7)]
    DD_motors = ["DD%s" % c for c in range_incl(1, 6)]
    VD_motors = ["VD%s" % c for c in range_incl(1, 13)]
    AS_motors = ["AS%s" % c for c in range_incl(1, 11)]
    AS_motors = []
    TW_cells = [
        "AVAL",
        "AVAR",
        "AVBL",
        "AVBR",
        "PVCL",
        "PVCR",
        "AVDL",
        "AVDR",
        "DVA",
        "PVDL",
        "PVDR",
        "PLML",
        "PLMR",
        "AVM",
        "ALML",
        "ALMR",
    ]
    # TW_sensory = ["PLML", "PLMR", "AVM", "ALML", "ALMR"]
    all_motors = list(
        VA_motors
        + VB_motors
        + DA_motors
        + DB_motors
        + DD_motors
        + VD_motors
        + AS_motors
    )
    # all_motors = []

    # neurons = ['AVAL', 'AVAR', 'AVBL', 'AVBR', 'PVCL', 'PVCR', 'AVDL', 'AVDR', 'DVA', 'PVDL', 'PVDR', 'PLML', 'PLMR', 'AVM', 'ALML', 'ALMR']

    muscles_to_include = False
    # muscles_to_include = True # ALL muscles
    # muscles_to_include = ['MVL01', 'MVL10']

    cells = list(TW_cells + all_motors)

    cells_to_plot = list(cells)
    # cells_to_plot = ["VB5", "VA4"]

    cells += [  # 'AS1', 'AS10', 'AS11', 'AS2', 'AS3', 'AS4', 'AS5', 'AS6', 'AS7', 'AS8', 'AS9',
        # 'AVFL', 'AVFR', 'AVKR', 'AVL',
        # 'CEPVL', 'CEPVR',
        #
        #
        # 'DVB',
        # 'HSNL', 'HSNR',
        # 'IL1DL', 'IL1DR', 'IL1L', 'IL1R', 'IL1VL', 'IL1VR',
        # 'PDA', 'PDB',
        # 'PVNL', 'PVNR',
        # 'RID', 'RIML', 'RIMR', 'RIVL', 'RIVR',
        # 'RMDDL', 'RMDDR', 'RMDL', 'RMDR', 'RMDVL', 'RMDVR', 'RMED', 'RMEL', 'RMER', 'RMEV', 'RMFL', 'RMGL', 'RMGR',
        # 'RMHL', 'RMHR',
        # 'SMBDL', 'SMBDR', 'SMBVL', 'SMBVR', 'SMDDL', 'SMDDR', 'SMDVL', 'SMDVR',
        # 'URADL', 'URADR', 'URAVL', 'URAVR',
        # 'VC1', 'VC2', 'VC3', 'VC4', 'VC5', 'VC6',
    ]

    # cells = None
    # cells_to_stimulate = ['PLML', 'PLMR', 'AVM']
    cells_to_stimulate = []

    # cells_to_plot = ['AVAL', 'AVAR', 'AVBL', 'AVBR', 'PLML', 'PLMR', 'AVM', 'ALML', 'ALMR', 'PVDL', 'PVDR', 'AVDL', 'AVDR']
    # cells_to_plot = ['PLML', 'AVM', 'AVBL', 'AVAL']
    # cells_to_plot = cells
    # cells_to_plot += motors
    reference = "c302_%s_TapWithdrawal" % parameter_set

    conn_polarity_override = {
        "ALML-ALML": "inh",
        "ALML-PVCL": "inh",
        "ALML-PVCR": "inh",
        "ALML-AVDR": "inh",  ##
        "ALMR-PVCR": "inh",
        "AVM-PVCL": "inh",
        "AVM-PVCR": "inh",
        "AVM-AVBL": "inh",
        "AVM-AVBR": "inh",
        "AVM-AVDL": "inh",  ##
        "AVM-AVDR": "inh",  ##
        "PVDL-PVDR": "inh",
        "PVDL-PVCL": "exc",
        "PVDL-PVCR": "exc",
        "PVDL-AVAL": "inh",
        "PVDL-AVAR": "inh",
        "PVDL-AVDL": "inh",
        "PVDL-AVDR": "inh",
        "PVDR-PVDL": "inh",
        "PVDR-DVA": "exc",
        "PVDR-PVCL": "exc",
        "PVDR-PVCR": "exc",
        "PVDR-AVAL": "inh",
        "PVDR-AVAR": "inh",
        "PVDR-AVDL": "inh",
        "DVA-PVCL": "inh",  #
        "DVA-PVCR": "inh",  #
        "DVA-AVAL": "inh",
        "DVA-AVAR": "inh",
        "DVA-AVBL": "inh",  #
        "DVA-AVBR": "inh",  #
        "DVA-AVDR": "inh",
        "PVCL-DVA": "exc",
        "PVCL-PVCL": "inh",
        "PVCL-PVCR": "inh",
        "PVCL-AVAL": "inh",
        "PVCL-AVAR": "inh",
        "PVCL-AVBL": "exc",
        "PVCL-AVBR": "exc",
        "PVCL-AVDL": "inh",
        "PVCL-AVDR": "inh",
        "PVCR-PVDL": "inh",
        "PVCR-PVDR": "inh",
        "PVCR-DVA": "exc",
        "PVCR-PVCL": "inh",
        "PVCR-AVAL": "inh",
        "PVCR-AVAR": "inh",
        "PVCR-AVBL": "exc",
        "PVCR-AVBR": "exc",
        "PVCR-AVDL": "inh",
        "PVCR-AVDR": "inh",
        "AVAL-PVCL": "inh",
        "AVAL-PVCR": "inh",
        "AVAL-AVAR": "inh",
        "AVAL-AVBL": "inh",
        "AVAL-AVDL": "inh",
        "AVAL-AVDR": "inh",
        "AVAR-PVCL": "inh",
        "AVAR-PVCR": "inh",
        "AVAR-AVAL": "inh",
        "AVAR-AVBL": "inh",
        "AVAR-AVBR": "inh",
        "AVAR-AVDL": "inh",
        "AVAR-AVDR": "inh",
        "AVBL-DVA": "inh",
        "AVBL-PVCR": "inh",
        "AVBL-AVAL": "inh",
        "AVBL-AVAR": "inh",
        "AVBL-AVBR": "inh",
        "AVBL-AVDR": "inh",
        "AVBR-AVAL": "inh",
        "AVBR-AVAR": "inh",
        "AVBR-AVBL": "inh",
        "AVBR-AVDL": "inh",
        "AVDL-PVCL": "inh",
        "AVDL-AVAL": "exc",
        "AVDL-AVAR": "exc",
        "AVDL-AVDR": "exc",
        "AVDR-PVCR": "inh",
        "AVDR-AVAL": "exc",
        "AVDR-AVAR": "exc",
        "AVDR-AVBL": "inh",
        "AVDR-AVDL": "exc",
        "DA9-DVA": "inh",
        "DVA-DA2": "inh",
        "PVCL-DA5": "inh",
        "PVCL-DA8": "inh",
        "PVCR-DA2": "inh",
        "PVCR-DA4": "inh",
        "PVCR-DA5": "inh",
        "PVCR-DA7": "inh",
        "AVBL-DA5": "inh",
        "AVBL-DA7": "inh",
        "AVBR-DA3": "inh",
        "AVBR-DA4": "inh",
        "AVBR-DA7": "inh",
        "AVDL-DA1": "inh",
        "AVDL-DA2": "inh",
        "AVDL-DA3": "inh",
        "AVDL-DA4": "inh",
        "AVDL-DA5": "inh",
        "AVDL-DA8": "inh",
        "AVDR-DA1": "inh",
        "AVDR-DA2": "inh",
        "AVDR-DA3": "inh",
        "AVDR-DA4": "inh",
        "AVDR-DA5": "inh",
        "AVDR-DA8": "inh",
        "DB1-DA1": "inh",
        "DB1-DA2": "inh",
        "DB2-DA2": "inh",
        "DB2-DA3": "inh",
        "DB2-DA4": "inh",
        "DB3-DA4": "inh",
        "DB3-DA5": "inh",
        "DB4-DA5": "inh",
        "DB5-DA6": "inh",
        "DB5-DA7": "inh",
        "DB5-DA8": "inh",
        "DB6-DA8": "inh",
        "DB6-DA9": "inh",
        "DB7-DA9": "inh",
        "AVAR-DB2": "inh",
        "AVAR-DB3": "inh",
        "AVAL-DB7": "inh",
        "DA1-DB1": "inh",
        "DA2-DA3": "inh",
        "DA2-DB1": "inh",
        "DA3-DB3": "inh",
        "DA4-DB2": "inh",
        "DA5-DB4": "inh",
        "DA6-DB5": "inh",
        "DA7-DB6": "inh",
        "DA8-DB7": "inh",
        "DA9-DB7": "inh",
        "DVA-VA2": "inh",
        "DVA-VA6": "inh",
        "DVA-VA8": "inh",
        "DVA-VA12": "inh",
        "PVCL-VA11": "inh",
        "PVDR-VA9": "inh",
        "PVDR-VA12": "inh",
        "AVBL-VA2": "inh",
        "AVBL-VA10": "inh",
        "AVBR-VA3": "inh",
        "AVBR-VA4": "inh",
        "AVDL-VA3": "inh",
        "AVDL-VA5": "inh",
        "AVDL-VA10": "inh",
        "AVDL-VA12": "inh",
        "AVDR-VA2": "inh",
        "AVDR-VA3": "inh",
        "AVDR-VA5": "inh",
        "AVDR-VA11": "inh",
        "VB1-VA1": "inh",
        "VB1-VA2": "inh",
        "VB1-VA3": "inh",
        "VB1-VA4": "inh",
        "VB2-VA2": "inh",
        "VB2-VA3": "inh",
        "VB3-VA4": "inh",
        "VB3-VA5": "inh",
        "VB4-VA4": "inh",
        "VB4-VA5": "inh",
        "VB5-VA6": "inh",
        "VB6-VA7": "inh",
        "VB6-VA8": "inh",
        "VB7-VA9": "inh",
        "VB7-VA10": "inh",
        "VB8-VA11": "inh",
        "VB9-VA11": "inh",
        "VB10-VA11": "inh",
        "VB11-VA12": "inh",
        "VB11-PVCR": "inh",
        "VB4-VB5": "inh",
        "VA2-VB1": "inh",
        "VA2-VB2": "inh",
        "VA3-VB2": "inh",
        "VA3-VB3": "inh",
        "VA4-AVDL": "inh",
        "VA4-VB3": "inh",
        "VA4-VB4": "inh",
        "VA5-VB4": "inh",
        "VA6-VB4": "inh",
        "VA6-VB5": "inh",
        "VA7-VB6": "inh",
        "VA8-VB6": "inh",
        "VA9-VB7": "inh",
        "VA9-VB8": "inh",
        "VA10-VB8": "inh",
        "VA10-VB9": "inh",
        "VA11-VB10": "inh",
        "VA12-PVCL": "inh",
        "VA12-PVCR": "inh",
        "VA12-DB7": "inh",
        "VA12-VB11": "inh",
        "AVM-VB3": "inh",
        #'VB4-VB5':'exc',
        #'DD1-DA2':'inh',
        #'DD1-VB2':'inh',
        #'DD2-DA3':'inh',
        #'DD3-DA5':'inh',
    }

    conn_number_override = {
        #'PVCL-AVDL':7*0.1,
        #'PVCL-AVDR':11*0.1,
        #'PVCR-AVDL':16*0.1,
        #'PVCR-AVDR':6*0.1,
        #'PVCR-AVDR_GJ':2 * 0.01,
        #'AVDR-PVCR_GJ':2 * 0.01,
        "PVCL-AVAL_GJ": 5 * 0.01,
        "AVAL-PVCL_GJ": 5 * 0.01,
        "PVCL-AVAR_GJ": 10 * 0.01,
        "AVAR-PVCL_GJ": 10 * 0.01,
        "PVCR-AVAL_GJ": 15 * 0.01,
        "AVAL-PVCR_GJ": 15 * 0.01,
        "PVCR-AVAR_GJ": 22 * 0.01,
        "AVAR-PVCR_GJ": 22 * 0.01,
        "PVCL-PLML_GJ": 4 * 0.01,  ##
        "PVCR-PLMR_GJ": 8 * 0.01,  ##
        #'AVDL-AVM_GJ':8 * 0.01,
        #'ALML-AVM_GJ':1 * 0.01,
        #'ALMR-AVM_GJ':1 * 0.01,
        #'AVDR-ALMR_GJ':1 * 0.01,
        # 'AVDL-AVAL':37*0.1,
        # 'AVDL-AVAR':37*0.1,
        # 'AVDR-AVAL':41*0.1,
        # 'AVDR-AVAR':52*0.1,
        # 'AVDL-AVAL_GJ':7*0.1,
        # 'AVAL-AVDL_GJ':7*0.1,
        # 'AVDL-AVAR_GJ':2*0.1,
        # 'AVAR-AVDL_GJ':2*0.1,
        # 'AVDR-AVAL_GJ':9*0.1,
        # 'AVAL-AVDR_GJ':9*0.1,
        # 'AVDR-AVAR_GJ':15*0.1,
        # 'AVAR-AVDR_GJ':15*0.1,
        # 'ALMR-AVDR_GJ':2*5,
        # 'AVDR-ALMR_GJ':2*5,
        "AVAR-AVBL_GJ": 3 * 0.01,
        "AVBL-AVAR_GJ": 3 * 0.01,
        # 'AVAR-AVAL_GJ':18*2,
        # 'AVAL-AVAR_GJ':18*2,
        "PVDL-AVAR_GJ": 4 * 0.01,
        "AVAR-PVDL_GJ": 4 * 0.01,
        "PVDR-AVAL_GJ": 6 * 0.01,
        "AVAL-PVDR_GJ": 6 * 0.01,
        "AVBL-VA11_GJ": 1 * 0.01,
        "VA11-AVBL_GJ": 1 * 0.01,
        "AVBR-VA11_GJ": 3 * 0.01,
        "VA11-AVBR_GJ": 3 * 0.01,
        "PVCR-VA11_GJ": 3 * 0.01,
        "VA11-PVCR_GJ": 3 * 0.01,
        "DVA-VA11_GJ": 1 * 0.01,
        "VA11-DVA_GJ": 1 * 0.01,
        "PVCL-VA12_GJ": 18 * 0.01,
        "VA12-PVCL_GJ": 18 * 0.01,
        "PVCR-VA12_GJ": 8 * 0.01,
        "VA12-PVCR_GJ": 8 * 0.01,
        "AVAL-VB11_GJ": 2 * 0.01,
        "VB11-AVAL_GJ": 2 * 0.01,
        "PVCL-DA4_GJ": 1 * 0.01,
        "DA4-PVCL_GJ": 1 * 0.01,
        "PVCL-DA7_GJ": 1 * 0.01,
        "DA7-PVCL_GJ": 1 * 0.01,
        "PVCR-DA7_GJ": 3 * 0.01,
        "DA7-PVCR_GJ": 3 * 0.01,
        "PVCL-DA8_GJ": 17 * 0.01,
        "DA8-PVCL_GJ": 17 * 0.01,
        "PVCR-DA8_GJ": 1 * 0.01,
        "DA8-PVCR_GJ": 1 * 0.01,
        "DVA-DA9_GJ": 3 * 0.01,
        "DA9-DVA_GJ": 3 * 0.01,
        "PVCR-DA9_GJ": 3 * 0.01,
        "DA9-PVCR_GJ": 3 * 0.01,
        "DB7-VA10_GJ": 1 * 0.01,
        "VA10-DB7_GJ": 1 * 0.01,
        "VA4-VB3_GJ": 1 * 0.01,
        "VB3-VA4_GJ": 1 * 0.01,
        "VA11-VB10_GJ": 3 * 0.01,
        "VB10-VA11_GJ": 3 * 0.01,
        "VA12-VB11_GJ": 7 * 0.01,
        "VB11-VA12_GJ": 7 * 0.01,
        "VB11-DA9_GJ": 7 * 0.01,
        "DA9-VB11_GJ": 7 * 0.01,
    }

    nml_doc = None

    if generate:
        nml_doc = c302.generate(
            reference,
            params,
            cells=cells,
            cells_to_plot=cells_to_plot,
            cells_to_stimulate=cells_to_stimulate,
            conn_polarity_override=conn_polarity_override,
            conn_number_override=conn_number_override,
            muscles_to_include=muscles_to_include,
            duration=duration,
            dt=dt,
            target_directory=target_directory,
            data_reader=data_reader,
            param_overrides=param_overrides,
            verbose=verbose,
        )

        # stim_amplitude = "6pA"
        # stim_amplitude = "5.135697186048022pA"

        for vb in VB_motors:
            c302.add_new_sinusoidal_input(
                nml_doc,
                cell=vb,
                delay="0ms",
                duration="1000ms",
                amplitude="3pA",
                period="150ms",
                params=params,
            )

        for db in DB_motors:
            c302.add_new_sinusoidal_input(
                nml_doc,
                cell=db,
                delay="0ms",
                duration="1000ms",
                amplitude="3pA",
                period="150ms",
                params=params,
            )

        # c302.add_new_input(nml_doc, "AVM", "10ms", "700ms", stim_amplitude, params)
        # c302.add_new_input(nml_doc, "ALML", "10ms", "700ms", stim_amplitude, params)
        # c302.add_new_input(nml_doc, "ALMR", "10ms", "700ms", stim_amplitude, params)
        # c302.add_new_input(nml_doc, "PLML", "10ms", "700ms", stim_amplitude, params)
        # c302.add_new_input(nml_doc, "PLMR", "10ms", "700ms", stim_amplitude, params)

        nml_file = target_directory + "/" + reference + ".net.nml"
        writers.NeuroMLWriter.write(
            nml_doc, nml_file
        )  # Write over network file written above...

        c302.print_("(Re)written network file to: " + nml_file)

    return cells, cells_to_stimulate, params, muscles_to_include, nml_doc


if __name__ == "__main__":
    parameter_set = sys.argv[1] if len(sys.argv) == 2 else "C2"

    setup(parameter_set, generate=True, data_reader=c302.FW_DATA_READER)
