"""Simulation ischemia/SD in morphological detailed CA1 pyramidal neuron"""
from neuron import h, rxd, gui
from neuron.units import mM, ms, nM, sec
import math
import os
import json
import sys
import matplotlib
from matplotlib import pyplot, colors, colorbar
from math import pi, exp, inf
from analysis import mech_names, range_vars, mechs_present
import numpy as np
import pickle
from scipy.spatial import Delaunay

rxd.options.enable.extracellular = True
pyplot.ion()
h.load_file("stdrun.hoc")
h.load_file("import3d.hoc")
cvode = h.CVode()
# scale factor so the flux (Dca/dr)*Ca has units molecules/um^2/ms
mM_to_mol_per_um = rxd.constants.NA() * 1e-18

# simulation parameters
h.celsius = 37
h.v_init = -69.5
dx = 20
recDt = 1.0


class PartialShell(rxd.Shell):
    """Extention of rxd.Shell to allow fractional volume within the shell"""

    def __init__(self, lo=None, hi=None, vf=1.0):
        super(PartialShell, self).__init__(lo, hi)
        self._vf = vf

    def __repr__(self):
        return "PartialShell(lo=%r, hi=%r, vf=%r)" % (self._lo, self._hi, self._vf)

    def volumes1d(self, sec):
        return super(PartialShell, self).volumes1d(sec) - rxd.inside.volumes1d(sec) * (
            1.0 - self._vf
        )


def my_record(vec, ref, dt=None):
    if dt:
        vec.record(ref, dt)
    else:
        vec.record(ref)


all_glia = []
all_neuron = []


class Glia:
    """Glia model with simplified morphology for clearance of extracellular K+.
    Based on Somjen et al 2008 https://modeldb.science/113446
    """

    def __init__(self, x, y, z, L=100, r=20, recV=False, recI=False, recO=False):
        """
        Create a simple astrocyte model

        Args:
            x (float): x-position of the cell.
            y (float): y-position of the cell.
            z (float): z-position of the cell.
            recV (float, optional): Option to record voltages at given timestep.
            recI (float, optional): Option to record intracellular concentrations at given timestep.
            recO (float, optional): Option to record extracellular concentrations at given timestep.
        """
        global all_glia
        self.name = "glia%i" % len(all_glia)
        all_glia.append(self)
        self.x = x
        self.y = y
        self.z = z
        self.sec = h.Section(name=self.name)
        self.sec.nseg = 101
        self.allsec = [self.sec]
        for i in range(101):
            self.sec.pt3dadd(x, y + i * L / 101, z, 2.0 * r)
        for mech in [
            "k_ion",
            "cl_ion",
            "ca_ion",
            "na_ion",
            "ATP_ion",
            "ip3_ion",
            "glu_ion",
            "nakpump",
            "clleak",
            "naleak",
            "kleak",
            "capump",
            "nacax",
            "kir",
            "kdr",
            "nkcc1",
            "kcc2",
            "caleak",
            "mglur",
        ]:
            print("insert", mech)
            self.sec.insert(mech)

        for seg in self.sec:
            seg.kir.gbar = 0.025  # Somjen
            seg.kdr.gkdrbar = 0.2  # "
            seg.nkcc1.g = 0.15  # Increased due to higher Cli then neuron
            seg.clleak.gcl = 1e-4  # Decreased " "
            seg.nakpump.Kk = 20  # Glia pump doesn't saturate as easily

        if recV:  # record membrane potential (shown in figure 1C)
            dt = recV if isinstance(recV, float) else recDt
            self.gliaV = [
                h.Vector().record(self.sec(x / 101)._ref_v, dt)
                for x in range(1, 101, 10)
            ]

        if recI:
            dt = recI if isinstance(recI, float) else recDt
            self.gliaKi = [
                h.Vector().record(self.sec(x / 101)._ref_ki, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaNai = [
                h.Vector().record(self.sec(x / 101)._ref_nai, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaCai = [
                h.Vector().record(self.sec(x / 101)._ref_cai, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaCli = [
                h.Vector().record(self.sec(x / 101)._ref_cli, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaGlui = [
                h.Vector().record(self.sec(x / 101)._ref_glui, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaATPi = [
                h.Vector().record(self.sec(x / 101)._ref_ATPi, dt)
                for x in range(1, 101, 10)
            ]

        if recO:
            dt = recO if isinstance(recO, float) else recDt
            self.gliaKo = [
                h.Vector().record(self.sec(x / 101)._ref_ko, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaNao = [
                h.Vector().record(self.sec(x / 101)._ref_nao, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaCao = [
                h.Vector().record(self.sec(x / 101)._ref_cao, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaClo = [
                h.Vector().record(self.sec(x / 101)._ref_clo, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaGluo = [
                h.Vector().record(self.sec(x / 101)._ref_gluo, dt)
                for x in range(1, 101, 10)
            ]
            self.gliaATPo = [
                h.Vector().record(self.sec(x / 101)._ref_ATPo, dt)
                for x in range(1, 101, 10)
            ]

    def __repr__(self):
        return self.name


# additional mechanisms not in original MM model
homeostatic_mechs = [
    "capump",
    "caleak",
    "nkcc1",
    "kcc2",
    "nap",
    "nacax",
    "kleak",
    "naleak",
    "clleak",
    "nakpump",
]
syn_mechs = ["nmda", "ampa", "mglur", "glurelease"]


class Neuron:
    """Python copy of MiglioreEtAl2018PLOSCompBiol2018/cell_seed3_0-pyr-08.hoc
    with additional mechanisms required to maintain ion homeostasis.
    """

    def __init__(self, x, y, z, recV=False, recI=False, recO=False):
        """
        Create a neuron

        Args:
            x (float): Shift the cell in the x-direction.
            y (float): Shift the cell in the y-direction.
            z (float): Shift the cell in the z-direction.
            recV (float, optional): Option to record voltages at given timestep.
            recI (float, optional): Option to record intracellular concentrations at given timestep.
            recO (float, optional): Option to record extracellular concentrations at given timestep.
        """
        global all_neuron
        self.name = "cell%i" % len(all_neuron)
        all_neuron.append(self)
        self.x = x
        self.y = y
        self.z = z
        self.recV = recV
        self.recI = recI
        self.recO = recO
        self.dend = []
        if x != 0 or y != 0 or z != 0:
            self._shift(x, y, z)

    def _shift(self, x, y, z):
        for sec in self.all:
            n = int(h.n3d(sec=sec))
            xs = [h.x3d(i, sec=sec) for i in range(n)]
            ys = [h.y3d(i, sec=sec) for i in range(n)]
            zs = [h.z3d(i, sec=sec) for i in range(n)]
            ds = [h.diam3d(i, sec=sec) for i in range(n)]
            for i, a, b, c, d in zip(range(len(xs)), xs, ys, zs, ds):
                sec.pt3dchange(i, a + x, b + y, c + z, d)

    def replace_axon(self, L_target=60.0, nseg0=5):
        """replace_axon method based on original HOC model"""

        nseg_total = 2 * nseg0
        chunk = L_target / nseg_total
        nSec = len(self.axon)
        if nSec < 1:
            print(
                "Less than two axon sections are present! Add an axon to the morphology and try again!"
            )
            return
        diams = []
        lens = []
        for sec in self.axon:
            sec.nseg = 1 + int(sec.L / chunk / 2.0) * 2
            for seg in sec:
                diams.append(seg.diam)
                lens.append(sec.L / sec.nseg)
                if len(diams) > nseg_total:
                    break
            if len(diams) > nseg_total:
                break

        # delete the old axon
        for sec in self.axon:
            h.disconnect(sec=sec)
            cell.all.remove(sec)

        # create the new axon
        self.axon = [h.Section(name="axon0"), h.Section(name="axon1")]

        L_real = 0
        diams.reverse()
        lens.reverse()
        for sec in self.axon:
            sec.L = L_target / 2.0
            sec.nseg = int(nseg_total / 2)
            cell.all.append(sec)
            for seg in sec:
                seg.diam = diams.pop()
                L_real += lens.pop()
        self.axon[0].connect(cell.soma[0])
        self.axon[1].connect(self.axon[0])
        print(
            "Target stub axon length:",
            L_target,
            "um, equivalent length: ",
            L_real,
            "um",
        )

    def __repr__(self):
        return self.name

    def setup(self):
        # insert the same mechanisms with the same parameters in both the soma
        # and the dendrite
        for sec in self.all:
            sec.nseg = 1 + 2 * int(sec.L / 40)
        self.replace_axon()
        # python import gives a slightly different soma diam than in HOC
        cell.soma[0].diam = 6.686145576954842
        for sec in self.all:
            for mechanims in ["kdr", "na3"] + homeostatic_mechs:
                sec.insert(mechanims)

        self.dend = self.apic[0:83]
        self.apic = self.apic[83:]

        h.distance(0, cell.soma[0](0.5))
        for sec in self.soma:
            for mechanism in ["kmb", "kap", "hd", "can", "cal", "cat", "kca", "cagk"]:
                sec.insert(mechanism)
                for seg in sec:
                    seg.na3.ar = 0.8
        for sec in self.apic + self.dend:
            for mechanism in [
                "kad",
                "hd",
                "can",
                "cal",
                "cat",
                "kca",
                "cagk",
            ] + syn_mechs:
                sec.insert(mechanism)
        for sec in self.apic:
            for seg in sec:
                seg.na3.ar = 0.5
        for sec in self.axon:
            for mechanism in ["kmb", "kap"]:
                sec.insert(mechanism)

    def record(self):
        if self.recV:  # record membrane potential
            dt = self.recV if isinstance(self.recV, float) else None
            self.somaV = [h.Vector().record(self.soma[0](0.5)._ref_v, dt)]
            self.apicalV = [h.Vector().record(sec(0.5)._ref_v, dt) for sec in self.apic]
            self.basalV = [h.Vector().record(sec(0.5)._ref_v, dt) for sec in self.dend]

        if self.recI:
            dt = self.recI if isinstance(self.recI, float) else None
            # record from soma
            self.somaCai = [
                record(sec(0.5), ca, dt, [sur, cyt, er]) for sec in self.soma
            ]
            self.somaCli = [record(sec(0.5), cl, dt) for sec in self.soma]
            self.somaKi = [record(sec(0.5), k, dt) for sec in self.soma]
            self.somaNai = [record(sec(0.5), na, dt) for sec in self.soma]
            self.somaGlui = [
                h.Vector().record(sec(0.5)._ref_glui, dt) for sec in self.soma
            ]
            self.somaATPi = [record(sec(0.5), ATP, dt) for sec in self.soma]
            self.somaO2i = [record(sec(0.5), O2, dt) for sec in self.soma]

            # record from apical dendrites
            self.apicCai = [
                record(sec(0.5), ca, dt, [sur, cyt, er]) for sec in self.apic
            ]
            self.apicCli = [record(sec(0.5), cl, dt) for sec in self.apic]
            self.apicKi = [record(sec(0.5), k, dt) for sec in self.apic]
            self.apicNai = [record(sec(0.5), na, dt) for sec in self.apic]
            self.apicGlui = [
                h.Vector().record(sec(0.5)._ref_glui, dt) for sec in self.apic
            ]
            self.apicATPi = [record(sec(0.5), ATP, dt) for sec in self.apic]
            self.apicO2i = [record(sec(0.5), O2, dt) for sec in self.apic]
            self.apicV = [h.Vector().record(sec(0.5)._ref_v, dt) for sec in self.apic]

            # record from basal dendrites
            self.dendCai = [
                record(sec(0.5), ca, dt, [sur, cyt, er]) for sec in self.dend
            ]
            self.dendCli = [record(sec(0.5), cl, dt) for sec in self.dend]
            self.dendKi = [record(sec(0.5), k, dt) for sec in self.dend]
            self.dendNai = [record(sec(0.5), na, dt) for sec in self.dend]
            self.dendGlui = [
                h.Vector().record(sec(0.5)._ref_glui, dt) for sec in self.dend
            ]
            self.dendATPi = [record(sec(0.5), ATP, dt) for sec in self.dend]
            self.dendO2i = [record(sec(0.5), O2, dt) for sec in self.dend]
            self.dendV = [h.Vector().record(sec(0.5)._ref_v, dt) for sec in self.apic]

        if self.recO:
            dt = self.recO if isinstance(self.recO, float) else None
            # record from apical dendrites
            self.apicKo = [h.Vector().record(sec(0.5)._ref_ko, dt) for sec in self.apic]
            self.apicNao = [
                h.Vector().record(sec(0.5)._ref_nao, dt) for sec in self.apic
            ]
            self.apicClo = [
                h.Vector().record(sec(0.5)._ref_clo, dt) for sec in self.apic
            ]
            self.apicCao = [
                h.Vector().record(sec(0.5)._ref_cao, dt) for sec in self.apic
            ]
            self.apicGluo = [
                h.Vector().record(sec(0.5)._ref_gluo, dt) for sec in self.apic
            ]

            # record from basal dendrites
            self.dendKo = [h.Vector().record(sec(0.5)._ref_ko, dt) for sec in self.dend]
            self.dendNao = [
                h.Vector().record(sec(0.5)._ref_nao, dt) for sec in self.dend
            ]
            self.dendClo = [
                h.Vector().record(sec(0.5)._ref_clo, dt) for sec in self.dend
            ]
            self.dendCao = [
                h.Vector().record(sec(0.5)._ref_cao, dt) for sec in self.dend
            ]
            self.dendGluo = [
                h.Vector().record(sec(0.5)._ref_gluo, dt) for sec in self.dend
            ]

            # record from soma
            self.somaKo = h.Vector().record(self.soma[0](0.5)._ref_ko, dt)
            self.somaNao = h.Vector().record(self.soma[0](0.5)._ref_nao, dt)
            self.somaCao = h.Vector().record(self.soma[0](0.5)._ref_cao, dt)
            self.somaGluo = h.Vector().record(self.soma[0](0.5)._ref_gluo, dt)
            self.somaClo = h.Vector().record(self.soma[0](0.5)._ref_clo, dt)
            self.somaO2o = h.Vector().record(self.soma[0](0.5)._ref_oxygeno, dt)

    def extrema(self, dx=None):
        h.define_shape()
        xs = [sec.x3d(pt) for sec in cell.all for pt in range(sec.n3d())]
        ys = [sec.y3d(pt) for sec in cell.all for pt in range(sec.n3d())]
        zs = [sec.z3d(pt) for sec in cell.all for pt in range(sec.n3d())]
        if dx:
            if hasattr(dx, "__len__"):
                dx, dy, dz = dx
            else:
                dx, dy, dz = dx, dx, dx
            xlo, xhi = min(xs) // dx, max(xs) // dx + 1
            ylo, yhi = min(ys) // dy, max(ys) // dy + 1
            zlo, zhi = min(zs) // dz, max(zs) // dz + 1
            return xlo * dx, ylo * dy, zlo * dz, xhi * dx, yhi * dy, zhi * dz
        xlo, xhi = min(xs), max(xs)
        ylo, yhi = min(ys), max(ys)
        zlo, zhi = min(zs), max(zs)
        return xlo, ylo, zlo, xhi, yhi, zhi


# load the morphology
morph = h.Import3d_Neurolucida3()
morph.input("mpg_mergedend.asc")
i3d = h.Import3d_GUI(morph, 0)

# create the CA1 neuron
cell = Neuron(0, 0, 0, recV=recDt, recI=recDt, recO=recDt)
i3d.instantiate(cell)
cell.setup()


# Add rxd components to the model
def initcon(intra, extra, gintra=None):
    """
    Provide a function to initialize rxd concentrations

    Parameters:
    - intra (float): The intracellular concentration (mM)
    - extra (float): The extracellular concentration (mM)
    - gintra (float, optional): The glia intracellular concentratrion. Defaults to None.

    Returns:
    A function that takes a rxd.node as an argument and returns a concentration.
    """
    return (
        lambda nd: extra
        if isinstance(nd, rxd.node.NodeExtracellular)
        else gintra
        if nd.sec in all_glia_sec
        else intra
    )


def initca(cacyt, caer, caecs):
    """
    Provide a function to initialize calcium concentrations

    Parameters:
    - cacyt (float): The cytosolic calcium concentration (mM)
    - caer (float): The ER calcium concentration (mM)
    - caecs (float, optional): The extracellular calcium concentration (mM)

    Returns:
    A function that takes a rxd.node as an argument and returns a concentration.
    """
    return (
        lambda nd: caecs
        if isinstance(nd, rxd.node.NodeExtracellular)
        else caer
        if nd.region.name == "er"
        else cacyt
    )


# neuronal volume
Vn = sum([seg.volume() for sec in cell.all for seg in sec])

# set glia volume = neuronal volume/2
Vg = Vn / 2

# intracellular volume
Vc = Vg + Vn

# total volume
xlo, ylo, zlo, xhi, yhi, zhi = cell.extrema(dx)
Vt = (xhi - xlo) * (yhi - ylo) * (zhi - zlo)

# convex hull around the cell
nx, ny, nz = int((xhi - xlo) // dx), int((yhi - ylo) // dx), int((zhi - zlo) // dx)
bbox = np.array((nx, ny, nz))
xs = [sec.x3d(pt) for sec in cell.all for pt in range(sec.n3d())]
ys = [sec.y3d(pt) for sec in cell.all for pt in range(sec.n3d())]
zs = [sec.z3d(pt) for sec in cell.all for pt in range(sec.n3d())]
pts = np.array([xs, ys, zs]).T
hull = Delaunay(pts)
bbox = np.zeros((nx, ny, nz))


def tort(x, y, z):
    return 1.6 if hull.find_simplex((x, y, z)) >= 0 else inf


Vtbox = bbox.sum() * dx**3

# glia position and morphology
glia_x = np.median(xs)
glia_z = np.mean(zs)
glia_y0, glia_y1 = min(ys), max(ys)
glia_L = glia_y1 - glia_y0
glia_r = (0.5 * Vn / (glia_L * pi)) ** 0.5

# create glia
glia = Glia(glia_x, glia_y0, glia_x, glia_L, glia_r, recV=True, recI=True, recO=True)
all_glia_sec = [g.sec for g in all_glia]


def fi0():
    for sec in all_glia_sec:
        sec.v = -80


fih = h.FInitializeHandler(0, fi0)

# effective volume fraction to give alpha=0.2, beta=0.7, gamma=0.1
x1 = Vc * (0.2 / 0.9) / ((1 - (0.2 / 0.9)) * (Vtbox - Vc))
Vecs = (Vtbox - Vc) * x1
vol_frac = 0.2  # Vecs/Vtbox
scale_factor = 5
ecs = rxd.Extracellular(
    xlo=xlo,
    ylo=ylo,
    zlo=zlo,
    xhi=xhi,
    yhi=yhi,
    zhi=zhi,
    dx=dx,
    volume_fraction=vol_frac / scale_factor,
    tortuosity=tort,
)
fc, fe = (
    0.83,
    0.17,
)  # cytoplasmic, er volume fractions

import itertools

all_ex_axon = list(itertools.chain(*[cell.apic, cell.dend, cell.soma]))
sur = rxd.Region(
    all_ex_axon, name="surf", nrn_region="i", geometry=rxd.Shell(0.85, 1.0)
)
cyt = rxd.Region(all_ex_axon, name="cyt", geometry=PartialShell(0, 0.85, fc))
gcyt = rxd.Region(glia.allsec, name="glia", nrn_region="i")

mem = rxd.Region(all_ex_axon, name="mem", geometry=rxd.membrane)
border = rxd.Region(
    all_ex_axon, name="border", geometry=rxd.ScalableBorder(diam_scale=0.85)
)
er = rxd.Region(all_ex_axon, name="er", geometry=rxd.FractionalVolume(fe))
cyt_er_membrane = rxd.Region(all_ex_axon, geometry=rxd.ScalableBorder(1))

# glu_store = rxd.Region([glia.sec], name='store', nrn_region='i')

# Initial glu store so no transport at baseline
# seg.gluo*exp((h.FARADAY/(1e3*h.R*(h.celsius+273.15)))*(2.0*seg.v - 3*seg.ena + seg.ek) - log(0.5))
init_glu_store, init_glu = 2.383441120793437e-10, 0
cli0, clo0 = 6.6, 119.0
nai0, nao0 = 10.0, 147.0
ki0, ko0 = 125.0, 3.4
cai0, ca_er0, cao0 = 60e-6, 1.0588835294117647, 1.4
CB0 = 0.2
glu = rxd.Species([ecs], d=2.1, name="glu", charge=-1, initial=init_glu, atolscale=1e-3)
na = rxd.Species(
    [ecs, sur, cyt, gcyt], d=1.78, name="na", charge=1, initial=initcon(10, 147, 55)
)
k = rxd.Species(
    [ecs, sur, cyt, gcyt], d=2.62, name="k", charge=1, initial=initcon(125, 2.9, 80)
)
cl = rxd.Species(
    [ecs, sur, cyt, gcyt], d=1.0, name="cl", charge=-1, initial=initcon(6.6, 119.0, 30)
)
dr = rxd.Parameter([border], name="dr", value=lambda nd: 0.425 * nd.segment.diam / 2.0)

ATPss = 2.59 * mM
tauATP = 3.8610038610038613 * ms  # 10 * ms
O2ss = 0.05 * mM

epsilon_o2 = 0.17 / sec  # mM/sec
# diffusion coeffishent interpolated from;
# https://doi.org/10.1016/0304-4165(96)00053-0
ATP = rxd.Species([sur, cyt, gcyt], d=0.445, name="ATP", charge=1, initial=ATPss)
ATP_cyt = ATP[cyt]
O2Dc = 1.54
O2 = rxd.Species(
    [sur, cyt, gcyt, ecs], d=O2Dc * 1.6**2, name="oxygen", charge=1, initial=O2ss
)
Vatp = (1 - ATP / ATPss) / tauATP
ATPrestore = rxd.Reaction(6 * O2[cyt], 30 * ATP[cyt], Vatp * O2[cyt], mass_action=False)
O2restore = rxd.Rate(O2[ecs], epsilon_o2 * (O2ss - O2[ecs]))


ca_cyt0, ca_er0, ca_ecs0 = 60e-6, 1.0588835294117647, 1.4
caDiff = 0.233
ca = rxd.Species(
    [ecs, sur, cyt, er, gcyt],
    d=caDiff,
    name="ca",
    charge=2,
    initial=initca(ca_cyt0, ca_er0, ca_ecs0),
    atolscale=1e-3,
)
ip3degTau = 1000  # ms

initGlumGluR = 280 * init_glu * 30e-3 / (0.016 + 280 * init_glu)

init_ip3 = ip3degTau * 0.016 * 0.2083e-3 * initGlumGluR
ip3 = rxd.Species(cyt, d=0.283, name="ip3", charge=1, initial=init_ip3)


ca_cyt = ca[cyt]
ca_er = ca[er]
ip_cyt = ip3[cyt]

kf = 5.5
kb = 0.0026

CB = rxd.Species(
    [cyt, er], d=0.043, name="CB", charge=0, initial=CB0
)  # CalBindin (Anwar)
caCB = rxd.Species(
    [cyt, er], d=0.043, name="caCB", charge=0, initial=kf * cai0 * CB0 / kb
)
CB_cyt = CB[cyt]
caCB_cyt = caCB[cyt]
# Calcium-CB complex
cabuf = rxd.Reaction(ca_cyt + CB_cyt, caCB_cyt, kf, kb)
# caextrude = rxd.Rate(ca, (0.0-ca[cyt])/5.0, regions=cyt, membrane_flux=False)
Kserca = 0.1  # Michaelis constant for SERCA pump
gserca = 4.0
Katp = 1.61  # (Mahmmoud 2008)
gserca = 4.0 * (Katp + ATPss) / ATPss  # rescale so Vmax is the same
serca = rxd.MultiCompartmentReaction(
    2 * ca_cyt + ATP_cyt > 2 * ca_er,
    0.5
    * gserca
    * ATP_cyt
    / (Katp + ATP_cyt)
    * (1e3 * ca_cyt) ** 2
    / (Kserca**2 + (1e3 * ca_cyt) ** 2),
    membrane=cyt_er_membrane,
    mass_action=False,
)

# action of IP3 receptor
Kip3 = 0.13
Kact = 0.4
minf = ip_cyt * 1000.0 * ca_cyt / (ip_cyt + Kip3) / (1000.0 * ca_cyt + Kact)
minf_init = init_ip3 * 1000 * 60e-6 / (init_ip3 + Kip3) / (1000.0 * 60e-3 + Kact)
ip3r_gate_state = rxd.State(cyt_er_membrane, initial=0.8)
h_gate = ip3r_gate_state[cyt_er_membrane]
kgate = 120400.0 * (minf * h_gate) ** 3
kgate_init = 120400.0 * (minf_init * 0.8) ** 3

ip3r = rxd.MultiCompartmentReaction(
    ca_er, ca_cyt, kgate, kgate, membrane=cyt_er_membrane
)


# IP3 receptor gating
ip3rg = rxd.Rate(h_gate, (1.0 / (1 + (0.1 / 60e-6) * ca_cyt / (0.4)) - h_gate) / 400.0)
# IP3 degradation - moves towards baseline level (ip3_init)
ip3deg = rxd.Rate(ip_cyt, (init_ip3 - ip_cyt) / ip3degTau, membrane_flux=False)

degGluRate = 1.0  # should lookup value (taken from m1dyst)
degGlu = rxd.Rate(glu, -degGluRate * (glu - init_glu) * 1e3, regions=[ecs])

### RYR - based on Sneyd et al, 2003
k_a_pos = 1500000000000.0  # mM^-4/ms
k_a_neg = 0.0288  # /ms
k_b_pos = 1500000000.0  # mM^-3/ms
k_b_neg = 0.3859  # /ms
k_c_pos = 0.00175  # /ms
k_c_neg = 0.0001  # /ms
v1ryr = 100  # /ms
Ka_4 = k_a_neg / k_a_pos  # Ka**4
Kb_3 = k_b_neg / k_b_pos  # Kb**3
Kc = k_c_neg / k_c_pos
# w_state is fraction of RYR not in C2 state (closed state), ie fraction of RYR that is open
c3ryr = (ca_cyt**3) / Kb_3
c4ryr = Ka_4 / (ca_cyt**4)
c3ryr_init = (cai0**3) / Kb_3
c4ryr_init = Ka_4 / (cai0**4)

w_inf = (1.0 + c4ryr + c3ryr) / (1.0 + (1.0 / Kc) + c4ryr + c3ryr)
w_init = (1.0 + c4ryr_init + c3ryr_init) / (1.0 + (1.0 / Kc) + c4ryr_init + c3ryr_init)
w_state = rxd.State(cyt_er_membrane, initial=w_init)
w_state_mem = w_state[cyt_er_membrane]
w_rate = rxd.Rate(w_state, k_c_neg * (w_inf - w_state_mem) / w_inf)
ryr_gate = w_state_mem * (1.0 + c3ryr) / (1.0 + c4ryr + c3ryr)
k_ryr = v1ryr * ryr_gate
ryr = rxd.MultiCompartmentReaction(
    ca_er, ca_cyt, k_ryr, k_ryr, membrane=cyt_er_membrane
)

ryr_gate0 = w_init * (1.0 + c3ryr_init) / (1.0 + c4ryr_init + c3ryr_init)
rate_serca0 = gserca * (cai0 * 1e3) ** 2 / (Kserca**2 + (1e3 * cai0) ** 2)
rate_ryr0 = v1ryr * ryr_gate0 * (ca_er0 - cai0)
rate_ip3r0 = kgate_init * ca_er0 - kgate_init * cai0
# k_serca0 - ( k_ryr0 + k_ip3r0 + k_leak0) * (1.25 - 60e-6)
# gleak = 3.0   # leak channel: bidirectional ca flow btwn cyt <> ER
gleak = (-rate_serca0 + (rate_ip3r0 - rate_ryr0) * (ca_er0 - cai0)) / (cai0 - ca_er0)
leak = rxd.MultiCompartmentReaction(
    ca_er, ca_cyt, gleak, gleak, membrane=cyt_er_membrane
)


def set_synaptic_density(cell, total=30634.5, apic_syn=18653.5, basal_syn=11981):
    print("set_synaptic_density")
    # set synaptic density -- from https://doi.org/10.1016/S0306-4522(00)00496-6
    apic = apic_syn / total
    basal = basal_syn / total
    apic_area = sum([sum(k[sur].nodes(sec).surface_area) for sec in cell.apic])
    basal_area = sum([sum(k[sur].nodes(sec).surface_area) for sec in cell.dend])
    total_area = apic_area + basal_area

    apic_vol = sum([sum(k[sur].nodes(sec).volume) for sec in cell.apic])
    basal_vol = sum([sum(k[sur].nodes(sec).volume) for sec in cell.dend])
    total_vol = apic_vol + basal_vol

    for sec in cell.all:
        if sec in cell.apic:
            for seg in sec:
                seg.glurelease.dsyn = apic_syn / apic_area  # syns per um^2
                # scale condutance by synaptic density
                # assume gbar from unifrom distribution of total synapses
                if hasattr(seg, "nmda"):
                    seg.nmda.gbar *= apic * (total_area / apic_area)
                seg.ampa.gbar *= apic * (total_area / apic_area)
                if hasattr(seg, "mglur"):
                    seg.mglur.totalmGluR *= apic * (total_vol / apic_vol)

        elif sec in cell.dend:
            for seg in sec:
                seg.glurelease.dsyn = basal_syn / basal_area
                if hasattr(seg, "nmda"):
                    seg.nmda.gbar *= basal * (total_area / basal_area)
                seg.ampa.gbar *= basal * (total_area / basal_area)
                if hasattr(seg, "mglur"):
                    seg.mglur.totalmGluR *= basal * (total_vol / basal_vol)


def record(sec, sp, dt=h.dt, reg=[sur, cyt]):
    """
    Record rxd concentrations from a given rxd.Regions.

    Parameters:
        sec: NEURON section object
            The section from which to record.
        sp: rxd.Species or rxd.State
            The species to record from
        dt: float, optional
            Time step for recording (default is h.dt).
        reg: list of rxd regions, optional
            List of NEURON regions to record from (default is [sur, cyt]).

    Returns:
        list of h.Vector
            List of h.Vector objects containing that record the concentration.
    """
    recs = []
    for r in reg:
        recs.append(h.Vector().record(sp[r].nodes(sec)[0]._ref_value, dt))
    return recs


# setup recording
cell.record()

# set pointer for glurelease
for sec in h.allsec():
    for seg in sec:
        if hasattr(seg, "glurelease"):
            h.setpointer(cell.soma[0](0.5)._ref_v, "vsrc", seg.glurelease)

h.finitialize(-69.5)
# Initial concentrations
h.clo0_cl_ion = clo0
h.cli0_cl_ion = cli0
h.nao0_na_ion = nao0
h.nai0_na_ion = nai0
h.ko0_k_ion = ko0
h.ki0_k_ion = ki0
h.cai0_ca_ion = cai0
h.cao0_ca_ion = cao0
h.glui0_glu_ion = init_glu
h.gluo0_glu_ion = init_glu  # 25nM


# diffusions
diffusions = []
for sp in [ca, k, cl, na, ATP]:
    diffusions.append(
        rxd.MultiCompartmentReaction(
            sp[cyt],
            sp[sur],
            mM_to_mol_per_um * sp._d / dr[border],
            mM_to_mol_per_um * sp._d / dr[border],
            border=border,
        )
    )
diffusions.append(
    rxd.MultiCompartmentReaction(
        O2[ecs],
        O2[sur],
        (0.5 / 0.425) * mM_to_mol_per_um * O2Dc / dr[border],
        (0.5 / 0.425) * mM_to_mol_per_um * O2Dc / dr[border],
        membrane=mem,
    )
)
diffusions.append(
    rxd.MultiCompartmentReaction(
        O2[cyt],
        O2[sur],
        mM_to_mol_per_um * O2Dc / dr[border],
        mM_to_mol_per_um * O2Dc / dr[border],
        border=border,
    )
)


def getnode(seg, s, r):
    return [nd for nd in s[r].nodes if nd.segment == seg][0]


t_vec = h.Vector()
t_vec.record(h._ref_t, recDt)

apical_dist = []
for sec in cell.apic:
    apical_dist.append(h.distance(cell.soma[0](0.5), sec(0.5)))

basal_dist = []
for sec in cell.dend:
    basal_dist.append(h.distance(cell.soma[0](0.5), sec(0.5)))


# mechanisms and principle parameter taken from mod files
params = {
    "ampa": "gbar",
    "cagk": "gbar",
    "cal": "gcalbar",
    "caleak": "gca",
    "can": "gcanbar",
    "capacitance": "",
    "capump": "totpump",
    "cat": "gcatbar",
    "glurelease": "vmaxN",
    "hd": "P",
    "kad": "gkabar",
    "kap": "gkabar",
    "kca": "gbar",
    "kccmerge": "gkcc2",
    "nkcc1": "g",
    "kcc2": "g",
    "kdb": "gkdbar",
    "kdr": "gkdrbar",
    "kdrb": "gkdrbar",
    "kmb": "gbar",
    "mglur": "totalmGluR",
    "na3": "gbar",
    "nacax": "imax",
    "nakpump": "imax",
    "nap": "gnabar",
    "nax": "gbar",
    "nmda": "gbar",
    "kleak": "gk",
    "naleak": "gna",
    "clleak": "gcl",
    "ampa": "gbar",
    "kad": "gkabar",
    "kdb": "gkd",
    "kdrb": "gkdrbar",
    "mglur": "totalmGluR",
    "na3": "gbar",
    "nmda": "gbar",
}


# match mechanisms to ion species involved
def mech_dict(species, include=True):
    res = {sp: {} for sp in species}
    rvs = range_vars()
    for mech in rvs:
        for sp in species:
            if "i%s" % sp in rvs[mech]:
                if include:
                    res[sp][mech] = params[mech] if mech in params else None
                elif mech in params:
                    res[sp][mech] = params[mech]
    return res


def save_params(filename):
    from neuron import h

    data = {}
    sp = ["ca", "cl", "k", "na", "ATP"]
    mechs = mech_dict(sp, include=False)
    for sec in h.allsec():
        data[repr(sec)] = get_params(sec(0.5), mechs)
        data[repr(sec)]["conc"] = {
            s: (getattr(sec(0.5), f"{s}i"), getattr(sec(0.5), f"{s}o")) for s in sp
        }
    with open(filename, "w") as f:
        json.dump(data, f, indent=4)


# ion specific leaks to be modified to give zero-flux at RMP.
param_mod = {"ca": "caleak", "na": "naleak", "k": "kleak", "cl": "clleak"}


def param_init(sp_mechs, secs=None):
    """Set the ion-specific leaks to balance all other ion-specific membrane currents"""
    for sec in secs:
        mechs = mechs_present(sec)
        for seg in sec:
            for sp in sp_mechs:
                sp_sum = 0
                sp_mod = 0
                p_mod = param_mod[sp]
                gbar = 0
                m = getattr(seg, p_mod)
                sp_mod = getattr(m, "i%s" % sp)
                gbar = getattr(m, params[p_mod])
                sp_sum = getattr(seg, "i%s" % sp) - sp_mod
                # assuming the mechanism being modified is linearly dependent
                # on the parameter
                gbar = -gbar * sp_sum / sp_mod
                if gbar < 0:
                    raise Exception(
                        "Negative conductance required: %s.%s.%s %g"
                        % (seg, p_mod, params[p_mod], gbar)
                    )
                setattr(m, params[p_mod], gbar)


def get_params(seg, ions):
    params = dict()
    for sp in ions:
        ion = "i%s" % sp
        im = ions[sp]
        params[ion] = dict()
        for mech_name in im:
            if hasattr(seg, mech_name):
                mech = seg.__getattribute__(mech_name)
                g = im[mech_name]
                g_val = mech.__getattribute__(g)
                if g_val == 0:
                    params[ion][mech_name] = {g: [0, 0, 0]}
                else:
                    i = mech.__getattribute__(ion)
                    params[ion][mech_name] = {g: [g_val, i, i / g_val]}
    return params


stim = h.IClamp(cell.soma[0](0.5))
stim.delay = 100
stim.dur = 500
stim.amp = 1.0

t_vec = h.Vector()
t_vec.record(h._ref_t)


with open("optimized_parameters.json", "r") as fin:
    BPOparams = json.load(fin)  # [param_load]

# include basal dendrites in 'apical' so parameters are the same
seclookup = {
    "all": cell.all,
    "basal": cell.dend,
    "apical": cell.dend + cell.apic,
    "somatic": cell.soma,
    "axonal": cell.axon,
}


def load_BPO(BPOparam):
    with open("BPOparameters.json", "r") as f:
        param_config = json.load(f)
    for param in param_config:
        # replace kccmerge
        if "kccmerge" in param["param_name"]:
            param["param_name"] = (
                "g_nkcc1" if "gplus" in param["param_name"] else "g_kcc2"
            )
        if "sectionlist" in param:
            if isinstance(param["sectionlist"], str):
                secll = [param["sectionlist"]]
            else:
                secll = param["sectionlist"]
        else:
            secll = ["all"]
        param["value"] = {}
        for sl in secll:
            key = param["param_name"] + "." + sl
            if key in BPOparam:
                value = BPOparam[key]
                param["value"][sl] = value
            else:
                param["value"][sl] = 0
    return param_config


param_config = load_BPO(BPOparams)
for param in param_config:
    if param["type"] == "global":
        continue
    pname = param["param_name"]
    if "sectionlist" in param:
        if isinstance(param["sectionlist"], str):
            secll = [param["sectionlist"]]
        else:
            secll = param["sectionlist"]
    else:
        secll = ["all"]
    if "_" in pname:
        var, mech = pname.split("_")
    else:
        var = pname
        mech = None
    for r in secll:
        for sec in seclookup[r]:
            if mech:
                for seg in sec:
                    if "dist" in param:
                        dist = param["dist"]
                        dist = dist.replace("{distance}", "h.distance(seg)").replace(
                            "{value}", "%1.30e" % param["value"][r]
                        )
                        value = eval(dist)
                    else:
                        value = param["value"][r]
                    if hasattr(seg, mech):
                        setattr(getattr(seg, mech), var, value)
            else:
                setattr(sec, var, param["value"][r])

rebalance = True


def rxd_init():
    global rebalance
    if rebalance:
        # calculate leaks
        d = mech_dict(["ca", "cl", "k", "na"])
        if cvode.active():
            cvode.re_init()
        else:
            h.fcurrent()
        param_init(d, cell.all)
        if cvode.active():
            cvode.re_init()
        else:
            h.fcurrent()
        param_init(d, cell.all)
        if cvode.active():
            cvode.re_init()
        h.fcurrent()


fih_rxd = h.FInitializeHandler(2, rxd_init)


def record_currents(seg, Dt=h.dt):
    species = ["ca", "cl", "k", "na", "glu"]
    mechs = mech_dict(species)
    vecs = {}
    for sp in species:
        vecs[sp] = {}
        for m in mechs[sp]:
            if hasattr(seg, m) and str(m)[-3:] != "ion":
                vecs[sp][m] = h.Vector().record(
                    getattr(getattr(seg, m), f"_ref_acc{sp}"), Dt
                )

    return vecs


def all_currents(Dt):
    currents = {}
    voltages = {}
    for sec in cell.all:
        currents[repr(sec)] = record_currents(sec(0.5), Dt=Dt)
        voltages[repr(sec)] = h.Vector().record(sec(0.5)._ref_v, Dt)
    return voltages, currents


voltages, currents = all_currents(2.5)


def saveall(path, filename):
    """
    Save all the recorded data to a specified file.

    This function saves a dictionary of recorded data a pickle file specified by the 'path' and 'filename' parameters, then clears the vectors.

    Args:
        path (str): The directory path where the file will be saved.
        filename (str): The name of the file to be saved.

    Returns:
        None
    """
    from neuron import hoc

    data = {}
    for obj in [cell, glia]:
        for x in dir(obj):
            item = getattr(obj, x)
            if isinstance(item, hoc.HocObject) and hasattr(item, "size"):  # vector
                data[repr(obj) + "." + x] = item
            elif isinstance(item, list) and len(item) > 0:  # list
                if isinstance(item[0], hoc.HocObject):
                    data[repr(obj) + "." + x] = [
                        y.as_numpy()
                        for y in item
                        if isinstance(y, hoc.HocObject) and hasattr(y, "size")
                    ]
                elif isinstance(item[0], list):
                    data[repr(obj) + "." + x] = [
                        [
                            z.as_numpy()
                            for z in y
                            if (z, hoc.HocObject) and hasattr(z, "size")
                        ]
                        for y in item
                    ]
    data["basal_dist"] = basal_dist
    data["apical_dist"] = apical_dist
    data["t_vec"] = t_vec

    with open(os.path.join(path, filename), "wb") as fout:
        pickle.dump(data, fout)

    data = {}
    # reset vectors
    for obj in [cell, glia]:
        for x in dir(obj):
            item = getattr(obj, x)
            if isinstance(item, hoc.HocObject) and hasattr(item, "size"):  # vector
                item.resize(0)
            elif isinstance(item, list) and len(item) > 0:
                for o in item:
                    if isinstance(o, hoc.HocObject) and hasattr(o, "size"):  # vector
                        o.resize(0)

    data = {}
    data["voltages"] = voltages
    data["currents"] = currents
    data["states"] = rxd.node._states
    with open(os.path.join(path, "current_" + filename), "wb") as fout:
        pickle.dump(data, fout)
    data = {}

    for sec, v in voltages.items():
        v.resize(0)
    for sec, sps in currents.items():
        for sp, mechs in sps.items():
            for v in mechs.values():
                v.resize(0)


def loadall(cell, filenames):
    """
    Load data from multiple pickle files created by `saveall`

    This function loads data from multiple files specified by the 'filenames' parameter. This overwrites the members of the

    Args:
        cell (Neuron object): An instatance of Neuron to set the members to the values loaded from the filenames.
        filenames (list of str): A list containing the names of the files to be loaded.

    Returns:
        None
    """
    from neuron import hoc

    print("Loading from", filenames)
    global basal_dist, apical_dist, voltages, currents, t_vec
    if isinstance(filenames, list):
        filename = filenames[0]
    else:
        filename = filenames

    with open(filename, "rb") as f:
        data = pickle.load(f)
    for x in dir(cell):
        if "cell0." + x in data:
            item = getattr(cell, x)
            val = data["cell0." + x]
            if isinstance(item, hoc.HocObject) and hasattr(item, "size"):
                setattr(cell, x, val)
            elif isinstance(val, list):
                if len(val) > 0 and isinstance(val[0], list):
                    setattr(cell, x, [[z for z in y] for y in val])
                else:
                    setattr(cell, x, [y for y in val])
    if "basal_dist" in data:
        basal_dist = data["basal_dist"]
    if "apical_dist" in data:
        apical_dist = data["apical_dist"]
    if "voltages" in data:
        voltages = data["voltages"]
    if "currents" in data:
        currents = data["currents"]
    if "states" in data:
        rxd.node._states = data["states"]
    if "t_vec" in data:
        t_vec = data["t_vec"].as_numpy()
    if filename != filenames:
        for fn in filenames[1:]:
            with open(fn, "rb") as f:
                data = pickle.load(f)
            if "t_vec" in data:
                print(data["t_vec"][-1] * 1e-3)
            for x in dir(cell):
                if "cell0." + x in data:
                    item = getattr(cell, x)
                    val = data["cell0." + x]
                    if isinstance(item, hoc.HocObject) and hasattr(item, "size"):
                        item.append(val)
                    elif isinstance(val, list):
                        if len(val) > 0 and isinstance(val[0], list):
                            """newitem = []
                            for a,b in zip(item,val):
                                subitem = []
                                for y,z in zip(a,b):
                                    subitem.append(np.append(y,z))
                                newitem.append(subitem)"""
                            setattr(cell, x, val)
                        else:
                            setattr(
                                cell, x, [np.append(y, z) for y, z in zip(item, val)]
                            )
                    else:
                        setattr(cell, x, np.concatenate(lst, val))
            if "currents" in data:
                for sec in currents:
                    for ion in currents[sec]:
                        for mech in currents[sec][ion]:
                            currents[sec][ion][mech].append(
                                data["currents"][sec][ion][mech]
                            )


def homeo(O20=O2ss, nmda=True, glu=True, mglur=True):
    """
    Apply homeostatic regulation to maintain cellular balance.

    This function applies homeostatic regulation to maintain cellular balance based on the provided parameters.

    Args:
        O20 (float, optional): Initial oxygen level. Defaults to 0.05mM.
        nmda (bool, optional): Whether to include NMDA receptors. Defaults to True.
        glu (bool, optional): Whether to include glutamate receptors. Defaults to True.
        mglur (bool, optional): Whether to include metabotropic glutamate receptors. Defaults to True.

    Returns:
        None
    """
    global serca, O2ss, O2restore, rebalance
    if not nmda:
        remove_mech("nmda")
    if not glu:
        remove_mech("glurelease")
    if not mglur:
        remove_mech("mglur")
    stim.amp = 0
    O2ss = O20
    O2restore = rxd.Rate(O2[ecs], epsilon_o2 * (O2ss - O2[ecs]))
    rebalance = False
    h.finitialize(-69.5)
    O2.nodes.value = O20


def isch(O20=0, nmda=True, glu=True, mglur=True):
    """
    Simulate ischemic conditions by changing the baseline oxygen level.

    Args:
        O20 (float, optional): Initial oxygen level. Defaults to 0.
        nmda (bool, optional): Whether to include NMDA receptors. Defaults to True.
        glu (bool, optional): Whether to include glutamate receptors. Defaults to True.
        mglur (bool, optional): Whether to include metabotropic glutamate receptors. Defaults to True.
    Returns:
        None

    """
    global serca, O2ss, O2restore, rebalance
    # ATPss = ATP0
    # serca = rxd.MultiCompartmentReaction(2*ca_cyt + ATP_cyt > 2*ca_er, 0.5*gserca*ATP_cyt/(Katp+ATP_cyt)*(1e3*ca_cyt)**2/(Kserca**2+(1e3*ca_cyt)**2)/ca_er, membrane=cyt_er_membrane, mass_action=False)
    # ATPrestore = rxd.Rate(ATP, kATP*(ATP0-ATP))
    O2ss = O20
    O2restore = rxd.Rate(O2[ecs], epsilon_o2 * (O2ss - O2[ecs]))
    if not nmda:
        remove_mech("nmda")
    if not glu:
        remove_mech("glurelease")
    if not mglur:
        remove_mech("mglur")
    rebalance = False
    h.finitialize(-69.5)
    O2.nodes.value = O20


def penumbra(O20=0, nmda=True, glu=True, mglur=True):
    """
    Simulate penumbra conditions by changing the baseline oxygen level and removing the initial current clamp stimulus.

    Args:
        O20 (float, optional): Initial oxygen level. Defaults to 0.
        nmda (bool, optional): Whether to include NMDA receptors. Defaults to True.
        glu (bool, optional): Whether to include glutamate receptors. Defaults to True.
        mglur (bool, optional): Whether to include metabotropic glutamate receptors. Defaults to True.
    Returns:
        None

    """
    global stim, serca, gserca, rebalance, O2ss, O2restore
    # ATPss = ATP0
    # serca = rxd.MultiCompartmentReaction(2*ca_cyt + ATP_cyt > 2*ca_er,
    # 0.5*gserca*ATP_cyt/(Katp+ATP_cyt)*(1e3*ca_cyt)**2/(Kserca**2+(1e3*ca_cyt)**2)/ca_er,membrane=cyt_er_membrane, mass_action=False)
    # ATPrestore = rxd.Rate(ATP, kATP*(ATP0-ATP))
    O2ss = O20
    O2restore = rxd.Rate(O2[ecs], epsilon_o2 * (O2ss - O2[ecs]))
    stim.amp = 0
    rebalance = False
    if not nmda:
        remove_mech("nmda")
    if not glu:
        remove_mech("glurelease")
    if not mglur:
        remove_mech("mglur")
    h.finitialize(-69.5)
    O2.nodes.value = O20


def remove_mech(mech):
    print(f"removing {mech}")
    for sec in h.allsec():
        if hasattr(sec(0.5), mech):
            sec.uninsert(mech)


def setTitle(title, ax=None):
    if not ax:
        ax = pyplot.gca()
    bbox = ax.get_yaxis().label.get_window_extent()
    x, _ = ax.transAxes.inverted().transform([bbox.x0, bbox.y0])
    ax.set_title(title, ha="left", x=x, fontweight="bold")


def boxoff(ax=None):
    ax = ax if ax else pyplot.gca()
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()


def plotall(t=10, reg=0, ko=True):
    fig = pyplot.figure()
    fig.set_size_inches([12, 8])
    apic_sec = 60
    dend_sec = 40
    ax1 = pyplot.subplot(5, 3, (1, 13))
    t_vec = np.linspace(0, t, len(cell.somaV[0]))
    t_vec2 = np.linspace(0, t, len(cell.somaKo))
    pyplot.plot(t_vec, cell.somaV[0], label="soma")
    pyplot.plot(t_vec, cell.apicV[apic_sec], label="apical")
    pyplot.plot(t_vec, cell.dendV[dend_sec], label="basal")
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$V_m$ (mV)")
    pyplot.legend(frameon=False)
    boxoff(ax1)
    t_vec = np.linspace(0, t, len(cell.somaKi[0][reg]))
    ax2 = pyplot.subplot(5, 3, 2)
    pyplot.plot(t_vec, cell.somaKi[0][reg])
    pyplot.plot(t_vec, cell.apicKi[apic_sec][reg])
    pyplot.plot(t_vec, cell.dendKi[dend_sec][reg])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$K_i$ (mM)")
    boxoff(ax2)

    if ko:
        ax3 = pyplot.subplot(5, 3, 3)
        pyplot.plot(t_vec2, cell.somaKo)
        pyplot.plot(t_vec2, cell.apicKo[apic_sec])
        pyplot.plot(t_vec2, cell.dendKo[dend_sec])
        pyplot.xlabel("time (s)")
        pyplot.ylabel("$K_o$ (mM)")
        boxoff(ax3)

    ax4 = pyplot.subplot(5, 3, 5)
    pyplot.plot(t_vec, cell.somaNai[0][reg])
    pyplot.plot(t_vec, cell.apicNai[apic_sec][reg])
    pyplot.plot(t_vec, cell.dendNai[dend_sec][reg])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Na_i$ (mM)")
    boxoff(ax4)

    ax5 = pyplot.subplot(5, 3, 6)
    pyplot.plot(t_vec2, cell.somaNao)
    pyplot.plot(t_vec2, cell.apicNao[apic_sec])
    pyplot.plot(t_vec2, cell.dendNao[dend_sec])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Na_o$ (mM)")
    boxoff(ax5)

    ax6 = pyplot.subplot(5, 3, 8)
    pyplot.plot(t_vec, cell.somaCli[0][reg])
    pyplot.plot(t_vec, cell.apicCli[apic_sec][reg])
    pyplot.plot(t_vec, cell.dendCli[dend_sec][reg])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Cl_i$ (mM)")
    boxoff(ax6)

    ax7 = pyplot.subplot(5, 3, 9)
    pyplot.plot(t_vec2, cell.somaClo)
    pyplot.plot(t_vec2, cell.apicClo[apic_sec])
    pyplot.plot(t_vec2, cell.dendClo[dend_sec])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Cl_o$ (mM)")
    boxoff(ax7)

    ax8 = pyplot.subplot(10, 3, 20)
    pyplot.plot(t_vec, 1e3 * cell.somaCai[0][reg])
    pyplot.plot(t_vec, 1e3 * cell.apicCai[apic_sec][reg])
    pyplot.plot(t_vec, 1e3 * cell.dendCai[dend_sec][reg])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Ca_i$ ($\mu$M)")
    boxoff(ax8)

    ax9 = pyplot.subplot(10, 3, 23)
    pyplot.plot(t_vec, 1e3 * cell.somaCai[0][2])
    pyplot.plot(t_vec, 1e3 * cell.apicCai[apic_sec][2])
    pyplot.plot(t_vec, 1e3 * cell.dendCai[dend_sec][2])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Ca_{ER}$ ($\mu$M)")
    boxoff(ax9)

    ax10 = pyplot.subplot(5, 3, 12)
    pyplot.plot(t_vec2, 1e3 * cell.somaCao)
    pyplot.plot(t_vec2, 1e3 * cell.apicCao[apic_sec])
    pyplot.plot(t_vec2, 1e3 * cell.dendCao[dend_sec])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Cao_i$ ($\mu$M)")
    boxoff(ax10)

    ax11 = pyplot.subplot(5, 3, 14)
    pyplot.plot(t_vec2, cell.somaGlui[0])
    pyplot.plot(t_vec2, cell.apicGlui[apic_sec])
    pyplot.plot(t_vec2, cell.dendGlui[dend_sec])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Glu_i$ (mM)")
    boxoff(ax11)

    ax12 = pyplot.subplot(5, 3, 15)
    pyplot.plot(t_vec2, 1e3 * cell.somaGluo)
    pyplot.plot(t_vec2, 1e3 * cell.apicGluo[apic_sec])
    pyplot.plot(t_vec2, 1e3 * cell.dendGluo[dend_sec])
    pyplot.xlabel("time (s)")
    pyplot.ylabel("$Glu_o$ ($\mu$M)")
    boxoff(ax12)


if __name__ == "__main__":
    import argparse

    # Create argument parser
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--tstop",
        required=False,
        type=float,
        default=60,
        help="time (in seconds) to run the simulation. Default 60s",
    )
    # Add named arguments
    parser.add_argument(
        "--O2",
        required=False,
        type=float,
        default=O2ss,
        help="O2 'steady-state' concentration in mM. Default 0.05mM",
    )
    parser.add_argument(
        "--saveint",
        type=float,
        default=10,
        help="Interval (in seconds) at which to save the results. Default 10s",
    )
    parser.add_argument(
        "--mode", type=str, default="isch", help="Run 'isch' or 'penumbra' simulation"
    )
    parser.add_argument(
        "--path", type=str, default="", help="path of the directory to save the data"
    )
    parser.add_argument(
        "--cao",
        type=float,
        default=ca_ecs0,
        help="extracellular Ca2+ concentration in mM",
    )
    parser.add_argument(
        "--ko", type=float, default=ko0, help="extracellular K+ concentration in mM"
    )
    parser.add_argument(
        "--nmda", action="store_false", default=True, help="block NMDA receptors"
    )
    parser.add_argument(
        "--glu", action="store_false", default=True, help="block all glutamate"
    )
    parser.add_argument(
        "--mglur", action="store_false", default=True, help="block mGluR"
    )
    parser.add_argument("--nax", action="store_false", default=True, help="block mGluR")
    parser.add_argument(
        "--load",
        action="store_true",
        help="load previous results instead of running the simulation. Default False",
    )
    parser.add_argument("--plot", action="store_true", help="Plot the results.")

    args = parser.parse_args()
    savename = f"CA1Pyr_{args.mode}_{args.O2}"
    savename += "_glu" if not args.glu else ""
    savename += "_nmda" if not args.nmda else ""
    savename += "_mglur" if not args.mglur else ""
    savename += "_nax" if not args.nax else ""

    savename += f"_cao{args.cao}" if args.cao != ca_ecs0 else ""
    savename += f"_ko{args.ko}" if args.ko != ko0 else ""
    if args.nax:
        for sec in h.allsec():
            for seg in sec:
                if hasattr(seg, "na3"):
                    seg.na3.ar = 1.0
    print("initialize")
    h.finitialize(-69.5)
    if args.mode == "penumbra":
        penumbra(args.O2, args.nmda, args.glu, args.mglur)
    elif args.mode == "isch":
        isch(args.O2, args.nmda, args.glu, args.mglur)
    if args.glu:
        set_synaptic_density(cell)
    # reduce extracellular ca
    if args.cao != ca_ecs0:
        for nd in ca[ecs].nodes:
            nd.value = args.cao
    if args.ko != ko0:
        for nd in k[ecs].nodes:
            nd.value = ko0
    print("Savename:", savename)
    if args.load:
        idx = 1
        files = []
        while True:
            file_path = os.path.join(args.path, savename + str(idx) + ".pkl")
            if os.path.exists(file_path):
                files.append(file_path)
                idx += 1
            else:
                break
        loadall(cell, files)
    else:
        files = []
        rng = int(args.tstop / args.saveint)
        for i in range(rng):
            h.continuerun((i + 1) * args.saveint * 1000)
            saveall(args.path, savename + str(i + 1) + ".pkl")
            files.append(savename + str(i + 1) + ".pkl")
        if h.t < args.tstop:
            h.continuerun(args.tstop * 1000)
            saveall(args.path, savename + str(rng + 1) + ".pkl")
            files.append(savename + str(rng + 1) + ".pkl")
    if args.plot:
        if not args.load:
            loadall(cell, files)
        plotall(args.tstop)
        pyplot.savefig(f"plot_{savename}.png")