import time
import matplotlib.pyplot as plt
import os, sys
import matplotlib
import numpy as np
import scipy
import scipy.optimize as optim
from scipy.integrate import odeint
from fontTools.misc.cython import returns
import csv

### Data processing ###

BehavioralData = np.loadtxt(open("/Users/micah/Documents/Thesis/Data/All behavioral data.csv"), delimiter=",")
#BehavioralData = np.loadtxt(open("/Users/micah/Documents/Thesis/Data/Behavior data wo 15mo.csv"), delimiter=",")
rows = len(BehavioralData[0:-1,0])
cols = int(len(BehavioralData[0,0:])/2)

zTildeWT = np.empty(cols)
zTildeTgF = np.empty(cols)

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

def g(x, m, b):
    return m * x + b

cohen_6_month_WT = 0

# Behavioral data
for j in range(cols):
    colDone = False

    for i in range(rows):
        if(not colDone):
            if(BehavioralData[i+1,j] == 0):
                colDone = True
            elif(BehavioralData[i+1,j] == -1):
                if(j == 0):
                    line_TgF = optim.curve_fit(g,
                                               range(4),
                                               BehavioralData[1:(i+1),j],
                                               bounds=[[-np.inf, BehavioralData[1,j]-0.1], [np.inf, BehavioralData[1,j]]])
                    line_WT = optim.curve_fit(g,
                                              range(4),
                                              BehavioralData[1:(i + 1), j+cols],
                                              bounds=[[-np.inf, BehavioralData[1,j+cols]-0.1],[np.inf, BehavioralData[1,j+cols]]])
                    cohen_6_month_WT = line_WT[0][0] / np.sum(BehavioralData[1:(i + 1), j + cols])
                    zTildeTgF[j] = np.mean(np.array([(line_TgF[0][0] / np.sum(BehavioralData[1:(i+1),j])) / cohen_6_month_WT,
                                                     np.mean(BehavioralData[(i+2):-1,j])]))
                    zTildeWT[j] = 1
                    ax.scatter(BehavioralData[0, j], (line_TgF[0][0] / np.sum(BehavioralData[1:(i+1),j])) / cohen_6_month_WT, color='red', marker='v', s=50)
                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 2, j], color='red', marker='^', s=50)
                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 3, j], color='red', marker='^', s=50)

                    ax.scatter(BehavioralData[0, j], 1, color='blue', marker='v', s=50)
                    #ax.scatter(BehavioralData[0, j], BehavioralData[i + 2, j+cols], color='blue', marker='^', s=50)
                    #ax.scatter(BehavioralData[0, j], BehavioralData[i + 3, j+cols], color='blue', marker='^', s=50)
                elif(i > 1):
                    line_TgF = optim.curve_fit(g,
                                               range(4),
                                               BehavioralData[1:(i + 1), j],
                                               bounds=[[-np.inf, BehavioralData[1,j]-0.1], [np.inf, BehavioralData[1,j]]])
                    line_WT  = optim.curve_fit(g,
                                               range(4),
                                               BehavioralData[1:(i + 1), j + cols],
                                               bounds=[[-np.inf, BehavioralData[1,j+cols]-0.1], [np.inf, BehavioralData[1,j+cols]]])
                    zTildeTgF[j] = (line_TgF[0][0] / np.sum(BehavioralData[1:(i+1),j])) / cohen_6_month_WT
                    print(line_TgF[0][0])
                    zTildeWT[j]  = (line_WT[0][0] / np.sum(BehavioralData[1:(i+1),j+cols])) / cohen_6_month_WT

                    ax.scatter(BehavioralData[0, j], zTildeTgF[j], color='red', marker='v', s=50)

                    ax.scatter(BehavioralData[0, j], zTildeWT[j], color='blue', marker='v', s=50)
                elif(i==0):
                    zTildeTgF[j] = np.mean(BehavioralData[(i + 2):(i + 4), j])
                    zTildeWT[j]  = np.mean(BehavioralData[(i + 2):(i + 4), j + cols])
                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 2, j], color='red', marker='^', s=50)
                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 3, j], color='red', marker='^', s=50)

                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 2, j+cols], color='blue', marker='^', s=50)
                    ax.scatter(BehavioralData[0, j], BehavioralData[i + 3, j+cols], color='blue', marker='^', s=50)


cohen = ax.scatter(-1, -1, color='black', marker='v', s=50, label = "Cohen et al (2013) figure 2 H & I")
sare = ax.scatter(-1, -1, color='black', marker='^', s=50, label = "Saré et al (2020) figures 1")

ax.plot(BehavioralData[0, 0:cols], zTildeTgF, color = 'red')
ax.plot(BehavioralData[0, 0:cols], zTildeWT, color = 'blue')
ax.legend(fontsize=16, handles=[cohen, sare])
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Cognitive performance\n (relative to 6 months WT)', fontsize=22)
ax.set_ylim(0,1.5)
ax.set_xlim(5,25)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Cognitive_data.pdf")

NeuronalLossData = np.loadtxt(open("/Users/micah/Documents/Thesis/Data/Neuro Loss data.csv"), delimiter=",")

## Cortex neuronal health data
y_cTildeWT = np.empty(3)
y_cTildeTgF = np.empty(3)

# TgF data
y_cTildeTgF = NeuronalLossData[2, 0:3]

# WT data
y_cTildeWT = NeuronalLossData[2, 3:6]

## HC neuronal health data
y_hcTildeWT = np.empty(3)
y_hcTildeTgF = np.empty(3)

# TgF
y_hcTildeTgF = NeuronalLossData[4, 0:3]

# WT
y_hcTildeWT = NeuronalLossData[4, 3:6]

y_TildeTgF = np.empty(3)
y_TildeWT = np.empty(3)

# Average HC and cortex data
for i in range(3):
    y_TildeTgF[i] = (y_hcTildeTgF[i] + y_cTildeTgF[i])/2
    y_TildeWT[i] = (y_hcTildeWT[i] + y_cTildeWT[i]) / 2

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

ax.scatter(NeuronalLossData[0, 0:3], y_cTildeTgF, color='red', marker='s', s=50)
ax.scatter(NeuronalLossData[0, 0:3], y_cTildeWT, color='blue', marker='s', s=50)
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeTgF, color='red', marker='o', s=50)
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeWT, color='blue', marker='o', s=50)

ax.plot(NeuronalLossData[0, 0:3], y_TildeTgF, color='red')
ax.plot(NeuronalLossData[0, 0:3], y_TildeWT, color='blue')

C = ax.scatter(-1, -1, color='black', marker='s', s=50, label="Cortex")
HC = ax.scatter(-1, -1, color='black', marker='o', s=50, label="Hippocampus")

ax.legend(fontsize=16, handles=[HC, C])

ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Neuronal density in HC\n(relative to 6 month)', fontsize=22)
ax.set_ylim(0,1.1)
ax.set_xlim(5,27)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Neuro_dens_data.pdf")

ABData = np.loadtxt(open("/Users/micah/Documents/Thesis/Data/Amyloid-beta data.csv"), delimiter=",")

rows = len(ABData[0:,0])
cols = int(len(ABData[0,0:])/2)

x_cTilde = np.empty(cols)
x_hcTilde = np.empty(cols)

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

# HC AB data

for j in range(cols):
    s = 0
    d = 0
    for i in range(rows - 1):
        if(ABData[i+1,j] != 0):
            if(i == 0):
               # bernaud = ax.scatter(ABData[0, j], ABData[i+1, j], color='red', marker='o', s=50, label = "Bernaud et al. (2022) figure 8b")
                s += ABData[i+1, j]
                d += 1
            elif(i == 1):
                cohen = ax.scatter(ABData[0, j], ABData[i + 1, j], color='blue', marker='v', s=50, label = "Cohen et al. (2013) figure 3d")
                s += ABData[i + 1, j]
                d += 1
            else:
                sare = ax.scatter(ABData[0, j]-1, ABData[i + 1, j], color='green', marker='^', s=50, label = "Saré et al. (2020) figure 6")
                s += ABData[i + 1, j]
                d += 1
    x_hcTilde[j] = s/d

# Cortex AB data

for j in range(cols):
    s = 0
    d = 0
    for i in range(rows - 1):
        if(ABData[i + 1, j + cols] != 0):
            if(i <= 1):
                #bernaud = ax.scatter(ABData[0, j], ABData[i+1, j + cols],
                 #                    color='red', marker='o', s=50, label = "Bernaud et al. (2022) figure 8b")
                s += ABData[i+1, j + cols]
                d += 1
            elif(i == 2):
                cohen = ax.scatter(ABData[0, j], ABData[i + 1, j + cols],
                                   color='blue', marker='v', s=50, label = "Cohen et al. (2013) figure 3d")
                s += ABData[i + 1, j + cols]
                d += 1
            else:
                sare = ax.scatter(ABData[0, j] - 1, ABData[i + 1, j + cols],
                                  color='green', marker='^', s=50, label = "Saré et al. (2020) figure 6")
                s += ABData[i + 1, j + cols]
                d += 1
    x_cTilde[j] = s/d


x_Tilde = np.empty(cols)
for i in range(cols):
    x_Tilde[i] = (x_cTilde[i] + x_hcTilde[i])/2
ind = ABData[0, 0:5]
ind = np.delete(ind, 1)
ind = np.delete(ind, 1)

dat = x_Tilde
dat = np.delete(dat, 1)
dat = np.delete(dat, 1)

#ax.scatter(ABData[0, 0], ABData[1, 0], color='black', marker='s', s=50)
plt.plot(ind, dat, color='red')
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('A$\\beta$ n-fold increase in HC and cortex\n(TgF, relative to 6 month)', fontsize=22)

#ax.legend(fontsize=16, handles=[bernaud, cohen, sare])
ax.legend(fontsize=16, handles=[cohen, sare])

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/AB_data.pdf")

### Model ###
L = 20
Nx = 102
dx = L / (Nx - 1)
time_sp = np.linspace(0, L, Nx) + 6

def f (x, a):
    return a*x

def g (y, alpha):
    return alpha*(y-1)

def logistic_func (t, k, lm, x0, t0):
    return (x0 * k * np.e **(lm * (t - t0)))/(k + x0 * np.e**(lm * (t - t0)) - x0)

def dx_dt (x, t, lm, k):
    return lm * x * (1 - x/k)

def dy_dt (y, t, x, a, b):
    return -(a*x + b)*y

def dz_dt (z, t, y, c):
    return c*(y-1)*z

def AB_model (lm, k, x0):
    sol = odeint(dx_dt, x0, time_sp - 6, args=(lm, k))
    return sol

def Neuron_model(y0, x, a, b):
    sol = np.empty(Nx)
    h = dx
    t = 0
    w = y0

    sol[0] = w
    i = 1
    while(i < Nx):
        K1 = h * dy_dt(w, t, x[i - 1], a, b)
        K2 = h * dy_dt(w + K1 / 2, t + h / 2, (x[i - 1] + x[i]) / 2, a, b)
        K3 = h * dy_dt(w + K2 / 2, t + h / 2, (x[i - 1] + x[i]) / 2, a, b)
        K4 = h * dy_dt(w + K3, t + h, x[i], a, b)

        w += (K1 + 2*K2 + 2*K3 + K4)/6
        t = i*h

        sol[i] = w
        i += 1
    return sol

def Behavior_model(z0, y, c):
    #a = 0, b = 20, N = Nx-1
    sol = np.empty(Nx)
    h = dx
    t = 0
    w = z0
    sol[0] = w
    i = 1
    while(i < Nx):
        K1 = h * dz_dt(w, t, y[i-1], c)
        K2 = h * dz_dt(w + K1 / 2, t + h / 2, (y[i-1]+y[i]) / 2, c)
        K3 = h * dz_dt(w + K2 / 2, t + h / 2, (y[i-1]+y[i]) / 2, c)
        K4 = h * dz_dt(w + K3, t + h, y[i], c)

        w += (K1 + 2*K2 + 2*K3 + K4)/6
        t = i*h

        sol[i] = w
        i += 1
    return sol

def ls_AB (lmk, x0):
    lm = np.abs(lmk[0])
    k = np.abs(lmk[1])
    s = 0
    #logistic_func (t, k, lm, x0, t0)
    u = logistic_func(time_sp, k, lm, x0,6)
    s = (s + (u[0] - x_Tilde[0]) ** 2
         #+ (u[int(dx ** (-1) * 3)] - x_Tilde[1]) ** 2
         #+ (u[int(dx ** (-1) * 6)] - x_Tilde[2]) ** 2
         + (u[int(dx ** (-1) * 10)] - x_Tilde[3]) ** 2
         + (u[int(dx ** (-1) * 19)] - x_Tilde[4]) ** 2)
    return s

def ls_NH (ab_lm, k, x_TgF, x_init_WT):
    s = 0
    a = np.abs(ab_lm[0])
    b = np.abs(ab_lm[1])
    lm_WT = np.abs(ab_lm[2])

    y0_TgF = y_hcTildeTgF[0]
    y0_WT = y_hcTildeWT[0]

    x0_WT = (k * x_init_WT)/(x_init_WT + (k - x_init_WT) * np.e**(lm_WT * -2))
    x_WT = logistic_func(time_sp, k, lm_WT, x0_WT, 6)

    #Neuron_model(y0, x, a, b)
    u_TgF = Neuron_model(y0_TgF, x_TgF, a, b)
    u_WT  = Neuron_model(y0_WT, x_WT, a, b)
    s += ((u_TgF[0] - y_hcTildeTgF[0]) ** 2
        + (u_TgF[int(dx ** (-1) * 10)] - y_TildeTgF[1]) ** 2
        + (u_TgF[int(dx ** (-1) * 20)] - y_TildeTgF[2]) ** 2
        + (u_WT[0] - y_hcTildeWT[0]) ** 2
        + (u_WT[int(dx ** (-1) * 10)] - y_TildeWT[1]) ** 2
        + (u_WT[int(dx ** (-1) * 20)] - y_TildeWT[2]) ** 2)
    return s

def ls_B (cd, y_TgF, y_WT):
    s = 0
    c = np.abs(cd[0])

    u_TgF = Behavior_model(zTildeTgF[0], y_TgF, c)
    u_WT = Behavior_model(zTildeWT[0], y_WT, c)
    s += ((u_TgF[0] - zTildeTgF[0])**2
          #+ (u_TgF[int(dx ** (-1) * 3)] - (zTildeTgF[0]+zTildeTgF[1])/2)**2
          + (u_TgF[int(dx ** (-1) * 6)] - zTildeTgF[1])**2
          + (u_TgF[int(dx ** (-1) * 9)] - zTildeTgF[2])**2
          + (u_TgF[int(dx ** (-1) * 12)] - zTildeTgF[3])**2
          + (u_TgF[int(dx ** (-1) * 18)] - zTildeTgF[4])**2
          + (u_WT[0] - zTildeWT[0])**2
          #+ (u_WT[int(dx ** (-1) * 3)] - (zTildeWT[0] + zTildeWT[1]) / 2) ** 2
          + (u_WT[int(dx ** (-1) * 6)] - zTildeWT[1])**2
          + (u_WT[int(dx ** (-1) * 9)] - zTildeWT[2])**2
          + (u_WT[int(dx ** (-1) * 12)] - zTildeWT[3])**2
          + (u_WT[int(dx ** (-1) * 18)] - zTildeWT[4])**2)
    return s

### Fit to AB data ###

lmk = [1,50]

rslt = optim.minimize(ls_AB, x0=lmk, args = (1), method='nelder-mead',
                      options={'xatol': 1e-8, 'disp': True})
print("[TgF lambda, K]")
print(np.abs(rslt["x"]))

lm_TgF = rslt["x"][0]
k = rslt["x"][1]

x_TgF = logistic_func(time_sp, k, lm_TgF, 1, 6)

### Fit to Neuronal Density data ###

ab_lm = [0.1, 0.2, 0.1]
x_init_WT = 0.290439 * x_TgF[0] # logistic_func(4, k, lm_TgF, 1, 6)
#ls_HC_NH (ab_lm, k, x_TgF, x_init_WT)
rslt = optim.minimize(ls_NH, x0=ab_lm, args = (k, x_TgF, x_init_WT), method='nelder-mead',
                      options={'xatol': 1e-8, 'disp': True})
print("[a, b, WT lambda]")
print(np.abs(rslt["x"]))

a = np.abs(rslt["x"][0])
b = np.abs(rslt["x"][1])
lm_WT = np.abs(rslt["x"][2])

x0_WT = x_init_WT# (k * x_init_WT)/(x_init_WT + (k - x_init_WT) * np.e**(lm_WT * -2))
x_WT = logistic_func(time_sp, k, lm_WT, x0_WT, 6)

y_TgF = Neuron_model(y_TildeTgF[0], x_TgF, a, b)
y_WT = Neuron_model(y_TildeWT[0], x_WT, a, b)

## Fit to Cognitive data ##

cd = [0.1]

rslt = optim.minimize(ls_B, x0=cd, args = (y_TgF, y_WT), method='nelder-mead',
                       options={'xatol': 1e-8, 'disp': True})
print("[c]")
print(np.abs(rslt["x"]))
c = np.abs(rslt["x"][0])

z_TgF = Behavior_model(zTildeTgF[0], y_TgF, c)
z_WT = Behavior_model(zTildeWT[0], y_WT, c)

### Plots ###

## AB ##

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)
plt.plot(time_sp, x_WT, color = 'blue', linestyle='dashed')
ax.plot(time_sp, x_TgF, color = 'red')
#ax.scatter(ABData[0, 0:5], x_Tilde, color='red', marker='o', s=50)
ax.scatter(ABData[0, 0], x_Tilde[0], color='red', marker='o', s=50)
ax.scatter(ABData[0, 3:5], x_Tilde[3:5], color='red', marker='o', s=50)
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('N-fold increase of A$\\beta$\nRelative to 6 month TgF', fontsize=22)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/AB_sim.pdf")

## Neuron density ##
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

# TgF #
ax.plot(time_sp, y_TgF, color ='red')
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeTgF, color='red', marker='o', s=50)
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Neuronal density in HC and cortex\n(relative to 6 month)', fontsize=22)
ax.set_ylim(0, 1.2)
# WT #
plt.plot(time_sp, y_WT, color ='blue')
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeWT, color='blue', marker='o', s=50)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Neuro_dens_sim.pdf")

## Behavior ##
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

ax.set_ylim(0,1.3)

#TgF
ax.plot(time_sp, z_TgF, color ='red')
ax.plot(time_sp, z_WT, color ='blue')


ax.scatter(BehavioralData[0, 0:2], zTildeTgF[0:2], color = 'red')
ax.scatter(BehavioralData[0, 3:cols], zTildeTgF[3:cols], color = 'red')
ax.scatter(BehavioralData[0, 2], zTildeTgF[2], color = 'red')

ax.scatter(BehavioralData[0, 0:2], zTildeWT[0:2], color = 'blue')
ax.scatter(BehavioralData[0, 3:cols], zTildeWT[3:cols], color = 'blue')
ax.scatter(BehavioralData[0, 2], zTildeWT[2], color = 'blue')
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Cognitive performance\n (relative to 6 months WT)', fontsize=22)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Cognitive_sim.pdf")

plt.show()


### TPPU treatment from Fang et al (2025) ###

# Figure 2 (B & C) indicates ~0.15579 of nontreatment TgF @ 9-10 months

def ls_AB_TPPU (lm):
    lm = np.abs(lm[0])
    s = 0

    #logistic_func (t, k, lm, x0, t0)
    u = logistic_func(time_sp, k, lm, x_Tilde[0],6)
    s = (s + (u[int(dx ** (-1) * 6.5)] - 0.15579 * x_TgF[int(dx ** (-1) * 6.5)]) ** 2)
    return s

lm = 0.1
rslt = optim.minimize(ls_AB_TPPU, x0=lm, method='nelder-mead',
                       options={'xatol': 1e-8, 'disp': True})
print("[TgF TPPU lambda]")
print(np.abs(rslt["x"]))

lm_TgF_TPPU = rslt["x"][0]

x_TgF_TPPU = logistic_func(time_sp, k, lm_TgF_TPPU, 1, 6)

y_TgF_TPPU = Neuron_model(y_TildeTgF[0], x_TgF_TPPU, a, b)

z_TgF_TPPU = Behavior_model(zTildeTgF[0], y_TgF_TPPU, c)

zDiff = z_TgF_TPPU/z_WT
print(x_TgF[int(dx ** (-1) * 6.5)])
print(0.15579 * x_TgF[int(dx ** (-1) * 6.5)])
### Plots ###

## AB ##

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

ax.set_ylim(0, 1.3)
plt.plot(time_sp, x_WT, color = 'blue', linestyle='dashed')
ax.plot(time_sp, x_TgF, color = 'red')
ax.plot(time_sp, x_TgF_TPPU, color = 'darkred', linestyle= 'dashed')
#ax.scatter(ABData[0, 0:5], x_Tilde, color='red', marker='o', s=50)
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Average n-fold increase of\nA$\\beta$ in HC and cortex', fontsize=22)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/AB_sim_TPPU.pdf")

## Neuron density ##
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

# TgF #
ax.plot(time_sp, y_TgF, color ='red')
ax.plot(time_sp, y_TgF_TPPU, color ='darkred', linestyle= 'dashed')
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeTgF, color='red', marker='o', s=50)
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Neuronal density in HC and cortex\n(relative to 6 month)', fontsize=22)
ax.set_ylim(0, 1.2)

# WT #
plt.plot(time_sp, y_WT, color ='blue')
ax.scatter(NeuronalLossData[0, 0:3], y_hcTildeWT, color='blue', marker='o', s=50)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Neuro_dens_sim_TPPU.pdf")

## Behavior ##
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111)

ax.set_ylim(0,1.3)

#TgF
ax.plot(time_sp, z_TgF, color ='red')
ax.plot(time_sp, z_TgF_TPPU, color ='darkred', linestyle= 'dashed')
ax.plot(time_sp, z_WT, color ='blue')


ax.scatter(BehavioralData[0, 0:2], zTildeTgF[0:2], color = 'red')
ax.scatter(BehavioralData[0, 3:cols], zTildeTgF[3:cols], color = 'red')
ax.scatter(BehavioralData[0, 2], zTildeTgF[2], color = 'red')

ax.scatter(BehavioralData[0, 0:2], zTildeWT[0:2], color = 'blue')
ax.scatter(BehavioralData[0, 3:cols], zTildeWT[3:cols], color = 'blue')
ax.scatter(BehavioralData[0, 2], zTildeWT[2], color = 'blue')
ax.set_xlabel('Age (months)', fontsize=22)
ax.set_ylabel('Cognitive performance\n (relative to 6 months WT)', fontsize=22)

fig.savefig("/Users/micah/Documents/Thesis/Thesis plots/Cognitive_sim_TPPU.pdf")

plt.show()
