from multiprocessing import Pool
from random import randint
from cnmodel import cells
from neuron import h
from tqdm import tqdm


def run(run_input, processes):
    results = []
    if processes == 1:
        for input in run_input:
            results.append(_run_trial(input))
    else:
        p = Pool(processes)
        for res in tqdm(p.imap_unordered(_run_trial, run_input)):
            results.append(res)
    return results


def add_pyramidal_cell():
    pyramidal = cells.Pyramidal.create(species="rat")
    pyramidal.add_dendrites()
    apical_dend = pyramidal.maindend[0]
    basal_dend = pyramidal.maindend[1]
    return pyramidal


def add_tuberculoventral_cell():
    tuberculoventral_1 = cells.Tuberculoventral.create()
    tuberculoventral_2 = cells.Tuberculoventral.create()
    return tuberculoventral_1, tuberculoventral_2


def add_dstellate_cell():
    dstel1ate = cells.DStellateEager.create()
    return dstel1ate


def _run_trial(run_input):
    seed, info, run_number = run_input
    """
    info is a dict
    """
    pyramidal = add_pyramidal_cell()
    tuberculoventral_1, tuberculoventral_2 = add_tuberculoventral_cell()
    dstellate = add_dstellate_cell()

    auditory_nerve_cells = []
    synapses = []
    inhib_synapses = []
    # auditory nerve attachments
    # attach to pyramidal cell
    for nsgc in range(48):
        auditory_nerve_cells.append(cells.DummySGC(cf=info["cf"], sr=info["sr"]))
        synapses.append(auditory_nerve_cells[-1].connect(pyramidal, type="multisite"))
        auditory_nerve_cells[-1].set_sound_stim(
            info["stim"],
            seed=seed + nsgc + randint(0, 80000),
            simulator=info["simulator"],
        )
    # attach to tuberculoventral 1
    for nsgc in range(18):
        # attach to tb cell
        auditory_nerve_cells.append(cells.DummySGC(cf=info["cf"], sr=info["sr"]))
        synapses.append(auditory_nerve_cells[-1].connect(tuberculoventral_1, type=info["synapse_type"]))
        auditory_nerve_cells[-1].set_sound_stim(
            info["stim"],
            seed=seed + nsgc + randint(0, 80000),
            simulator=info["simulator"],
        )
    # attach to tuberculoventral 2
    for nsgc in range(18):
        # attach to tb cell
        auditory_nerve_cells.append(cells.DummySGC(cf=info["cf"], sr=info["sr"]))
        synapses.append(auditory_nerve_cells[-1].connect(tuberculoventral_2, type="multisite"))
        auditory_nerve_cells[-1].set_sound_stim(
            info["stim"],
            seed=seed + nsgc + randint(0, 80000),
            simulator=info["simulator"],
        )
    for nsgc in range(24):
        # attach to dstellate cell
        auditory_nerve_cells.append(cells.DummySGC(cf=info["cf"], sr=info["sr"]))
        synapses.append(auditory_nerve_cells[-1].connect(dstellate, type="multisite"))
        auditory_nerve_cells[-1].set_sound_stim(
            info["stim"],
            seed=seed + nsgc + randint(0, 80000),
            simulator=info["simulator"],
        )
    # Connections between network cells
    # for _ in range(5):
    #     inhib_synapses.append(cartwheel.connect(pyramidal, type="simple"))
    for _ in range(21):
        inhib_synapses.append(tuberculoventral_1.connect(pyramidal, type="simple"))
        inhib_synapses.append(tuberculoventral_2.connect(pyramidal, type="simple"))
    for _ in range(15):
        inhib_synapses.append(dstellate.connect(pyramidal, type="simple"))
        inhib_synapses.append(dstellate.connect(tuberculoventral_1, type='simple'))
        inhib_synapses.append(dstellate.connect(tuberculoventral_2, type='simple'))
    for _ in range(3):
        inhib_synapses.append(dstellate.connect(dstellate, type="simple"))
    stim = insert_current_clamp(pyramidal.soma)

    # set up our recording vectors for each cell
    Vm = h.Vector()
    Vm.record(pyramidal.soma(0.5)._ref_v)
    Vmtb = h.Vector()
    Vmtb.record(tuberculoventral_1.soma(0.5)._ref_v)
    Vmds = h.Vector()
    Vmds.record(dstellate.soma(0.5)._ref_v)
    rtime = h.Vector()
    rtime.record(h._ref_t)
    # hoc trial run
    h.tstop = 1e3 * info["run_duration"]  # duration of a run
    h.celsius = info["temp"]
    h.dt = info["dt"]
    init_cells([pyramidal, tuberculoventral_1, tuberculoventral_2, dstellate])
    info["init"]()
    h.t = 0.0
    h.run()
    # dtime = time.time() - start
    # print(f"Trial {run_number} completed after {dtime} secs")

    return {
        "time": list(rtime),
        "vm": list(Vm),
        "auditory_nerve_cells": [x._spiketrain.tolist() for x in auditory_nerve_cells],
        "vmtb": list(Vmtb),
        "vmds": list(Vmds),
    }


def insert_current_clamp(sec):
    """
    :param sec: to attach too
    dur: ms
    amp: nA
    delay: ms
    :return: stim needs to be put in a variable to stay alive
    """
    stim = h.IClamp(0.5, sec=sec)
    stim.dur = 30
    stim.amp = -0.2
    stim.delay = 120
    return stim


def init_cells(cell: list):
    for x in cell:
        x.cell_initialize()