acts = json.loads(Path('activities.json').read_text())
runs = [a for a in acts if a.get('type') == 'Run']
with_hr = [a for a in runs if a.get('average_heartrate')]
print(f'Total activities: {len(acts)}')
print(f'Runs: {len(runs)}')
print(f'Runs with avg HR: {len(with_hr)}')
print(f'Date range: {min(a["start_date"][:10] for a in runs)} -> {max(a["start_date"][:10] for a in runs)}')

EASY_TYPES = {None, 0}
records = []
for a in acts:
    if a.get('type') != 'Run': continue
    if a.get('workout_type') not in EASY_TYPES: continue
    hr = a.get('average_heartrate')
    spd = a.get('average_speed')   # m/s
    dist_mi = a.get('distance', 0) / 1609.34
    if not hr or not spd or hr < 90 or hr > 200: continue
    pace = 26.8224 / spd            # min/mi
    if pace < 6 or pace > 15: continue
    if dist_mi < 1.0 or dist_mi > 15: continue
    if a.get('trainer'): continue
    records.append({
        'date': datetime.fromisoformat(a['start_date'].replace('Z','+00:00')).date(),
        'pace': pace,
        'hr': hr,
    })
records.sort(key=lambda r: r['date'])
print(f'Filtered easy runs: {len(records)}')

dates = np.array([mdates.date2num(r['date']) for r in records])
paces = np.array([r['pace'] for r in records])
hrs   = np.array([r['hr'] for r in records])

# Multivariate linear: HR ~ pace + date
X = np.column_stack([paces, dates, np.ones_like(paces)])
coef, *_ = np.linalg.lstsq(X, hrs, rcond=None)
a_, b_, c_ = coef
resid = hrs - X @ coef
print(f'Pace coef: {a_:.2f} bpm per +1 min/mi slower')
print(f'Drift:     {b_*365:+.2f} bpm/year at fixed pace')
print(f'R^2:       {1 - resid.var()/hrs.var():.3f}    resid std: {resid.std():.2f} bpm')


fig, ax = plt.subplots(figsize=(9, 6))
sc = ax.scatter(paces, hrs, c=dates, cmap=gray_blue, s=22,
                edgecolor='k', linewidth=0.2, alpha=0.75)
cbar = plt.colorbar(sc, ax=ax)
loc = mdates.AutoDateLocator(); cbar.ax.yaxis.set_major_locator(loc)
cbar.ax.yaxis.set_major_formatter(mdates.ConciseDateFormatter(loc))
cbar.set_label('Date')
for q, color, lbl in [(0.10, (0.78, 0.78, 0.78), 'Early'),
                      (0.50, (0.45, 0.50, 0.55), 'Mid'),
                      (0.90, (0.10, 0.30, 0.65), 'Recent')]:
    qd = np.percentile(dates, q*100)
    xs = np.linspace(paces.min(), paces.max(), 50)
    ax.plot(xs, a_*xs + b_*qd + c_, '--', color=color, lw=2,
            label=f'{lbl} ({mdates.num2date(qd):%Y-%m})')
    
ax.invert_xaxis()
ax.set_xlabel('Pace (min/mile)'); ax.set_ylabel('HR (bpm)')
ax.set_title(f'Easy-run HR vs Pace, full Strava history (n={len(records)})')
ax.grid(alpha=0.3); ax.legend()
plt.tight_layout()
plt.show()

ref_pace = 9.5
hr_norm = hrs - a_ * (paces - ref_pace)

fig, ax = plt.subplots(figsize=(11, 5))
ax.scatter([r['date'] for r in records], hr_norm, c=dates, cmap=gray_blue,
           s=16, alpha=0.5, edgecolor='none')
order = np.argsort(dates); t_s, h_s = dates[order], hr_norm[order]
window = 90
roll_x, roll_y = [], []
for i in range(len(t_s)):
    mask = (t_s >= t_s[i]-window/2) & (t_s <= t_s[i]+window/2)
    if mask.sum() >= 5:
        roll_x.append(t_s[i]); roll_y.append(np.median(h_s[mask]))
ax.plot([mdates.num2date(x) for x in roll_x], roll_y,
        color=BLUE, lw=2, label='90-day rolling median')
ax.set_ylabel('HR @ 9:30 pace (bpm, lower = fitter)')
ax.set_xlabel('Date')
ax.xaxis.set_major_locator(mdates.YearLocator())
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
ax.set_title('Pace-normalized easy HR over 11 years')
ax.grid(alpha=0.3); ax.legend()
plt.tight_layout()
plt.show()

# Yearly medians at 9:30
print('Year   HR @ 9:30   n')
for y in sorted({r['date'].year for r in records}):
    mask = np.array([r['date'].year == y for r in records])
    if mask.sum() < 3: continue
    print(f'{y}    {np.median(hr_norm[mask]):6.1f}     {mask.sum():4d}')

year_dist = defaultdict(float)        # miles
year_time = defaultdict(float)        # seconds (moving)
year_count = defaultdict(int)
year_dist_easy = defaultdict(float)
year_time_easy = defaultdict(float)

for a in acts:
    if a.get('type') != 'Run': continue
    if a.get('manual'): continue
    dist_mi = a.get('distance', 0) / 1609.34
    mtime = a.get('moving_time', 0)
    if dist_mi <= 0 or mtime <= 0: continue
    year = datetime.fromisoformat(a['start_date'].replace('Z','+00:00')).year
    year_dist[year] += dist_mi
    year_time[year] += mtime
    year_count[year] += 1
    if a.get('workout_type') in (None, 0):
        year_dist_easy[year] += dist_mi
        year_time_easy[year] += mtime

years = sorted(year_dist)
dist = [year_dist[y] for y in years]
counts = [year_count[y] for y in years]
pace_easy = [year_time_easy[y]/60 / year_dist_easy[y] if year_dist_easy[y] > 0 else None for y in years]

fig, ax = plt.subplots(figsize=(10, 5))
ax.bar(years, dist, color=BLUE, edgecolor='k', linewidth=0.4)
for x, y_, n in zip(years, dist, counts):
    ax.text(x, y_+max(dist)*0.01, f'{y_:.0f}\n({n})', ha='center', va='bottom',
            fontsize=8.5, color='k')
ax.set_xlabel('Year'); ax.set_ylabel('Distance (miles)')
ax.set_title('Annual running distance (label = miles, n = runs)')
ax.set_xticks(years); ax.grid(alpha=0.3, axis='y')
ax.set_ylim(0, max(dist)*1.15)
plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(10, 5))
easy_pace = np.array([p if p is not None else np.nan for p in pace_easy], dtype=float)
ax.plot(years, easy_pace, '-o', color=BLUE, lw=2, ms=8, mec='k', mew=0.5)
for y, p, n in zip(years, easy_pace, counts):
    if np.isnan(p): continue
    m = int(p); s = int(round((p-m)*60))
    ax.annotate(f'{m}:{s:02d}', xy=(y, p), xytext=(0, -14),
                textcoords='offset points', ha='center', fontsize=9)
ax.set_xlabel('Year'); ax.set_ylabel('Average pace (min/mile)')
ax.set_title('Annual average pace, easy/default runs (distance-weighted; up = faster)')
ax.set_xticks(years)
ymin = np.nanmin(easy_pace) - 0.4
ymax = np.nanmax(easy_pace) + 0.4
ax.set_ylim(ymax, ymin)  # invert so up = faster
ax.grid(alpha=0.3)
plt.tight_layout()
plt.show()

def log_curve(d, a, b):
    return a + b * np.log10(d)

# Races where I was pacing someone else, not my own effort
PACING_RUNS = {
    '2026-03-01', '2026-03-08', '2025-11-27', '2025-07-27',
    '2025-04-27', '2025-03-15', '2025-05-17',
}

races = []
for a in acts:
    if a.get('type') != 'Run' or a.get('workout_type') != 1: continue
    spd = a.get('average_speed')
    dist_mi = a.get('distance', 0) / 1609.34
    if not spd or dist_mi <= 0: continue
    pace = 26.8224 / spd
    date_str = a['start_date'][:10]
    if date_str in PACING_RUNS: continue
    if pace > 13: continue
    races.append({
        'date': datetime.fromisoformat(a['start_date'].replace('Z','+00:00')).date(),
        'year': datetime.fromisoformat(a['start_date'].replace('Z','+00:00')).year,
        'dist': dist_mi,
        'pace': pace,
        'name': a.get('name',''),
    })
races.sort(key=lambda r: r['date'])
print(f'Races kept: {len(races)}')

year_fits = {}
print('\nPer-year log fits: pace = a + b*log10(distance)')
for yr in sorted({r['year'] for r in races}):
    rs = [r for r in races if r['year'] == yr]
    if len(rs) < 2: continue
    d = np.array([r['dist'] for r in rs])
    p = np.array([r['pace'] for r in rs])
    po, _ = curve_fit(log_curve, d, p, maxfev=20000)
    pred = log_curve(d, *po)
    r2 = 1 - (p - pred).var() / p.var() if len(p) > 2 else 1.0
    year_fits[yr] = (po, r2, len(rs))
    print(f'  {yr} (n={len(rs)}): a={po[0]:.2f}  b={po[1]:+.2f}  '
          f'pace@5K={log_curve(3.107,*po):.2f}  pace@HM={log_curve(13.11,*po):.2f}  '
          f'pace@M={log_curve(26.22,*po):.2f}  R^2={r2:.2f}')

years_set = sorted({r['year'] for r in races})
y_min, y_max = min(years_set), max(years_set)
def year_color(y):
    return gray_blue((y - y_min) / max(1, y_max - y_min))

fig, ax = plt.subplots(figsize=(10, 6.5))
for r in races:
    ax.scatter(r['dist'], r['pace'], color=year_color(r['year']), s=80,
               edgecolor='k', linewidth=0.6, zorder=3)
    ax.annotate(f"{r['date']:%y-%m}", xy=(r['dist'], r['pace']),
                xytext=(4, 4), textcoords='offset points', fontsize=7, alpha=0.8)
xs = np.linspace(2, max(r['dist'] for r in races)*1.05, 400)
for yr, (po, r2y, n) in year_fits.items():
    label = f'{yr} fit (n={n}'
    if n > 2: label += f', R^2={r2y:.2f}'
    label += ')'
    ax.plot(xs, log_curve(xs, *po), '-', color=year_color(yr), lw=2.2, label=label)
for yr in years_set:
    ax.scatter([], [], color=year_color(yr), s=60, edgecolor='k', label=f'{yr} races')
ax.set_xlabel('Distance (miles)')
ax.set_ylabel('Pace (min/mile, lower = faster)')
ax.set_title(f'Race pace vs distance, log fit per year (n={len(races)})')
ax.grid(alpha=0.3)
ax.legend(loc='lower right', fontsize=9)
plt.tight_layout()
plt.show()

def draw_circles(circles, circle_chain=[], show_label=False):
    fig, ax = plt.subplots(figsize=(8, 8)) # note we must use plt.subplots, not plt.subplot
    xmin, xmax = 999, -999
    ymin, ymax = 999, -999

    for i, c in enumerate(circles):
        patch = plt.Circle(c.center, c.r, color='gray')
        if c.x - c.r < xmin: xmin = c.x - c.r
        if c.x + c.r > xmax: xmax = c.x + c.r
        if c.y - c.r < ymin: ymin = c.y - c.r
        if c.y + c.r > ymax: ymax = c.y + c.r
        ax.add_patch(patch)
        if show_label:
            ax.text(c.x - 0.1, c.y - 0.1, f'C_{i}')
        
    line_x = [circles[idx].x for idx in circle_chain]
    line_y = [circles[idx].y for idx in circle_chain]
    ax.plot(line_x, line_y, marker='o')

    xwidth = xmax - xmin
    ywidth = ymax - ymin
    xcenter = xmin + xwidth / 2
    ycenter = ymin + ywidth / 2
    xwidth = ywidth = max(xwidth, ywidth)
    
    ax.set_xlim(xcenter - xwidth / 2, xcenter + xwidth / 2)
    ax.set_ylim(ycenter - ywidth / 2, ycenter + ywidth / 2)
    plt.show()

def circle_intersect(circle1, circle2, test_only=False):
    """
    Intersection points of two circles using the construction of triangles
    as proposed by Paul Bourke, 1997.
    http://paulbourke.net/geometry/circlesphere/
    """

    x0, y0 = circle1.x, circle1.y
    x1, y1 = circle2.x, circle2.y
    r0 = circle1.r
    r1 = circle2.r

    d = sqrt((x0 - x1)**2 + (y0 - y1)**2)
    if d > (r0 + r1):
        return -1, None # do not intersect
    elif d <= abs(r1 - r0):
        return -2, None # one circle is contained within the circle

    if d == (r0 + r1) or d == (r0 - r1):
        CASE = 1
    else:
        CASE = 2
    
    if test_only:
        return CASE, None

    a = (r0*r0 - r1*r1 + d*d) / (2 * d)
    h = sqrt(r0*r0 - a*a)

    x2 = x0 + a * (x1 - x0) / d
    y2 = y0 + a * (y1 - y0) / d

    I1 = (x2 + h * (y1 - y0) / d, y2 - h * (x1 - x0) / d)
    if CASE == 1:
        return CASE, (I1,)

    I2 = (x2 - h * (y1 - y0) / d, y2 + h * (x1 - x0) / d)

    theta1 = math.degrees(math.atan2(I1[1]-y0, I1[0]-x0))
    theta2 = math.degrees(math.atan2(I2[1]-y0, I2[0]-x0))
    if theta2 > theta1:
        I1, I2 = I2, I1
    return CASE, (I1, I2)
    
def tangent_circle_center(circle1, circle2, radius3):
    """
    find center of circle tangent to both circle1 and circle2
    
    circle1 and circle2 is assumed to be tangent
    """

    x0, y0 = circle1.x, circle1.y
    x1, y1 = circle2.x, circle2.y
    r0 = circle1.r + radius3
    r1 = circle2.r + radius3

    d = sqrt((x0 - x1)**2 + (y0 - y1)**2)
    a = (r0*r0 - r1*r1 + d*d) / (2 * d)
    h = sqrt(r0*r0 - a*a)

    x2 = x0 + a * (x1 - x0) / d
    y2 = y0 + a * (y1 - y0) / d

    I1 = (x2 + h * (y1 - y0) / d, y2 - h * (x1 - x0) / d)
    I2 = (x2 - h * (y1 - y0) / d, y2 + h * (x1 - x0) / d)

    torque = (x1-x0)*(I1[1]-y0)-(y1-y0)*(I1[0]-x0)
    if torque > 0:
        I1, I2 = I2, I1
    return I1, I2

"""
A simple example of an animated plot
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

fig, ax = plt.subplots(figsize=(8, 8)) # note we must use plt.subplots, not plt.subplot
xmin, xmax = 999, -999
ymin, ymax = 999, -999

line, = ax.plot([], [])

c3 = Circle(p3[0][0], p3[0][1], radii[2])
circles = [c1, c2, c3]
chain = [0, 1, 2, 0]
index = 0

ax.set_xlim(-20, 20)
ax.set_ylim(-20, 20)

def animate(index):
    global circles, chain
    if index > len(chain) - 2:
        index = 0
        
    circles, chain = add_circle(circles, chain, index, radii[3 + index])

    #line.set_xdata(circles[-1].x)
    #line.set_ydata(circles[-1].y)
    
    patch = plt.Circle(circles[-1].center, circles[-1].r, color='gray')
    ax.add_patch(patch)
    
    x = []
    y = []
    
    for idx in chain:
        x.append(circles[idx].x)
        y.append(circles[idx].y)
        
    x.append(circles[chain[0]].x)
    y.append(circles[chain[0]].y)
        
    line.set_data(x, y)
    return line,

#Init only required for blitting to give a clean slate.
def init():
    x = []
    y = []
    
    for i in range(3):
        patch = plt.Circle(circles[i].center, circles[i].r, color='gray')
        ax.add_patch(patch)
        x.append(circles[i].x)
        y.append(circles[i].y)
        
    x.append(circles[0].x)
    y.append(circles[0].y)
        
    line.set_data(x, y)
    return line,

#ani = animation.FuncAnimation(fig, animate, np.arange(1, 200), init_func=init,
#    interval=25, blit=False)

ani = animation.FuncAnimation(fig, animate, np.arange(1, 100), interval=350, 
                              init_func=init, blit=True)
plt.close()

from IPython.display import HTML
HTML(ani.to_html5_video())

from rdkit.Chem.Draw import rdMolDraw2D
from io import BytesIO

ms = [Chem.RemoveHs(m) for m in product_library]
for m in ms: tmp=AllChem.Compute2DCoords(m)
legends=["%s" % (x.GetProp("_Name"), ) for i, x in enumerate(ms)]

molsPerRow = 4
subImgSize = (300, 300)
width = subImgSize[0] * molsPerRow
height = subImgSize[1] * int(len(ms) / molsPerRow)
d2d = rdMolDraw2D.MolDraw2DCairo(width, height, subImgSize[0], subImgSize[1])
d2d.drawOptions().legendFontSize=24

d2d.DrawMolecules(ms,legends=legends)
d2d.FinishDrawing()
img = BytesIO()
img.write(d2d.GetDrawingText())

from PIL import Image
Image.open(img)

%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib

from io import BytesIO
import pandas as pd
import numpy as np
from IPython.display import SVG

# RDKit 
import rdkit
from rdkit.Chem import PandasTools
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import DataStructs
from rdkit.Chem import rdMolDescriptors
from rdkit.Chem import rdRGroupDecomposition
from rdkit.Chem.Draw import IPythonConsole #Needed to show molecules
from rdkit.Chem import Draw
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D
from rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions #Only needed if modifying defaults

DrawingOptions.bondLineWidth=1.8
IPythonConsole.ipython_useSVG=True
from rdkit import RDLogger
RDLogger.DisableLog('rdApp.warning')
print(rdkit.__version__)

# misc
from tqdm.notebook import tqdm
from math import pi, sqrt

mean = np.mean(A, axis=0)
std = np.std(A, axis=0)
A = (A - mean) / std
C = np.dot(A.T, A) / A.shape[0] # correlation matrix
w, v = np.linalg.eig(C) # w = eigenvalues, v = eigenvectors

gamma = A.shape[1] / A.shape[0]
gamma_p = (1 + sqrt(gamma))**2
gamma_m = (1 - sqrt(gamma))**2
rho = lambda x: sqrt(np.clip(gamma_p - x, 0, None) * np.clip(x - gamma_m, 0, None)) / (2*pi*gamma*x)

x = np.arange(0., np.max(np.real(w)), 0.1)
y = np.array([rho(_) for _ in x])
plt.hist(np.real(w), density=True)
plt.plot(x, y)
plt.xlabel('Eigenvalue')
plt.ylabel('Prob.')
plt.show()

p = 1024
fps_active = [np.array(AllChem.GetMorganFingerprintAsBitVect(m,3,p)) for m in mols_active]
fps_decoy = [np.array(AllChem.GetMorganFingerprintAsBitVect(m,3,p)) for m in mols_decoy]

plt.figure(figsize=(12, 8))
plt.imshow(fps_active, interpolation=None, cmap=plt.cm.gray)
plt.ylabel('Ligand')
plt.xlabel('Fingerprint')
plt.show()

def MP_distribution(N, p):
    """return MP distribution function based on the shape of input matrix"""
    gamma = p / N
    gamma_p = (1 + sqrt(gamma))**2
    gamma_m = (1 - sqrt(gamma))**2
    rho = lambda x: sqrt(np.clip(gamma_p - x, 0, None) * np.clip(x - gamma_m, 0, None)) / (2*pi*gamma*x)
    return rho

plt.hist(np.real(w), bins=50, density=True) # plot eigenvalue distribution

N, p = A_normed.shape
rho = MP_distribution(N, p) # MP distribution
w = np.real(w) # real values of eigenvalues
x = np.arange(0., np.max(w), 0.1)
y = np.array([rho(_) for _ in x])
plt.plot(x, y) # plot MP distribution

plt.ylabel('Probability')
plt.xlabel('Eigenvalues')
plt.ylim(0, 0.2)
plt.show()

# active
prob_active, bins_active = np.histogram(dist_active, bins=12, density=True)
x_active = (bins_active[1:] + bins_active[:-1]) / 2
plt.plot(x_active, prob_active)

# decoy
prob_decoy, bins_decoy = np.histogram(dist_decoy, bins=25, density=True)
x_decoy = (bins_decoy[1:] + bins_decoy[:-1]) / 2
plt.plot(x_decoy, prob_decoy)

plt.xlabel('Distance from V')
plt.ylabel('Probability')
plt.show()

train_size = 0.1
np.random.seed(1904234)
train_indices = np.array([np.random.random() > train_size for i in range(len(A_active))])
train_mat = A_active[train_indices]
test_mat = np.concatenate([A_active[~train_indices], A_decoy])
test_target = np.concatenate([[True] * np.sum(~train_indices), [False] * len(A_decoy)])

model = RMTClassifier()
model.fit(train_mat)
score = model.predict_proba(test_mat)

from sklearn import metrics
fpr, tpr, threshold = metrics.roc_curve(test_target, score)
auc = metrics.roc_auc_score(test_target, score)

plt.plot(fpr, tpr)
plt.xlabel("False Positive")
plt.ylabel("True Positive")
plt.text(x=0.8, y=0.05, s="AUC = %.2f" % auc)
plt.show()

from sklearn import metrics
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC

%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib

from io import BytesIO
import pandas as pd
import numpy as np
from IPython.display import SVG

# RDKit 
import rdkit
from rdkit.Chem import PandasTools
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import DataStructs
from rdkit.Chem import rdMolDescriptors
from rdkit.Chem import rdRGroupDecomposition
from rdkit.Chem.Draw import IPythonConsole #Needed to show molecules
from rdkit.Chem import Draw
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D
from rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions #Only needed if modifying defaults

DrawingOptions.bondLineWidth=1.8
IPythonConsole.ipython_useSVG=True
from rdkit import RDLogger
RDLogger.DisableLog('rdApp.warning')
print(rdkit.__version__)

# pytorch
import torch
from torch.utils.data import DataLoader, Dataset, Sampler
from torch import nn

# misc
from typing import Dict, Iterator, List, Optional, Union, OrderedDict, Tuple
from tqdm.notebook import tqdm
from functools import reduce

# we will define a class which holds various parameter for D-MPNN
class TrainArgs:
    smiles_column = None
    no_cuda = False
    gpu = None
    num_workers = 8
    batch_size = 50
    no_cache_mol = False
    dataset_type = 'regression'
    task_names = []
    seed = 0
    hidden_size = 300
    bias = False
    depth = 3
    dropout = 0.0
    undirected = False
    aggregation = 'mean'
    aggregation_norm = 100
    ffn_num_layers = 2
    ffn_hidden_size = 300
    init_lr = 1e-4
    max_lr = 1e-3
    final_lr = 1e-4
    num_lrs = 1
    warmup_epochs = 2.0
    epochs = 30

    @property
    def device(self) -> torch.device:
        """The :code:`torch.device` on which to load and process data and models."""
        if not self.cuda:
            return torch.device('cpu')

        return torch.device('cuda', self.gpu)

    @device.setter
    def device(self, device: torch.device) -> None:
        self.cuda = device.type == 'cuda'
        self.gpu = device.index

    @property
    def cuda(self) -> bool:
        """Whether to use CUDA (i.e., GPUs) or not."""
        return not self.no_cuda and torch.cuda.is_available()

    @cuda.setter
    def cuda(self, cuda: bool) -> None:
        self.no_cuda = not cuda

from random import Random

# Cache of graph featurizations
CACHE_GRAPH = True
SMILES_TO_GRAPH = {}

def cache_graph():
    return CACHE_GRAPH

def set_cache_graph(cache_graph):
    global CACHE_GRAPH
    CACHE_GRAPH = cache_graph

# Cache of RDKit molecules
CACHE_MOL = True
SMILES_TO_MOL: Dict[str, Chem.Mol] = {}

def cache_mol() -> bool:
    r"""Returns whether RDKit molecules will be cached."""
    return CACHE_MOL

def set_cache_mol(cache_mol: bool) -> None:
    r"""Sets whether RDKit molecules will be cached."""
    global CACHE_MOL
    CACHE_MOL = cache_mol
    
# Atom feature sizes
MAX_ATOMIC_NUM = 100
ATOM_FEATURES = {
    'atomic_num': list(range(MAX_ATOMIC_NUM)),
    'degree': [0, 1, 2, 3, 4, 5],
    'formal_charge': [-1, -2, 1, 2, 0],
    'chiral_tag': [0, 1, 2, 3],
    'num_Hs': [0, 1, 2, 3, 4],
    'hybridization': [
        Chem.rdchem.HybridizationType.SP,
        Chem.rdchem.HybridizationType.SP2,
        Chem.rdchem.HybridizationType.SP3,
        Chem.rdchem.HybridizationType.SP3D,
        Chem.rdchem.HybridizationType.SP3D2
    ],
}

# Distance feature sizes
PATH_DISTANCE_BINS = list(range(10))
THREE_D_DISTANCE_MAX = 20
THREE_D_DISTANCE_STEP = 1
THREE_D_DISTANCE_BINS = list(range(0, THREE_D_DISTANCE_MAX + 1, THREE_D_DISTANCE_STEP))

# len(choices) + 1 to include room for uncommon values; + 2 at end for IsAromatic and mass
ATOM_FDIM = sum(len(choices) + 1 for choices in ATOM_FEATURES.values()) + 2
EXTRA_ATOM_FDIM = 0
BOND_FDIM = 14


def get_atom_fdim():
    """Gets the dimensionality of the atom feature vector."""
    return ATOM_FDIM + EXTRA_ATOM_FDIM

def get_bond_fdim(atom_messages=False):
    """Gets the dimensionality of the bond feature vector.
    """
    return BOND_FDIM + (not atom_messages) * get_atom_fdim()

def onek_encoding_unk(value: int, choices: List[int]):
    encoding = [0] * (len(choices) + 1)
    index = choices.index(value) if value in choices else -1
    encoding[index] = 1

    return encoding

def atom_features(atom: Chem.rdchem.Atom, functional_groups: List[int] = None):
    """Builds a feature vector for an atom.
    """
    features = onek_encoding_unk(atom.GetAtomicNum() - 1, ATOM_FEATURES['atomic_num']) + \
               onek_encoding_unk(atom.GetTotalDegree(), ATOM_FEATURES['degree']) + \
               onek_encoding_unk(atom.GetFormalCharge(), ATOM_FEATURES['formal_charge']) + \
               onek_encoding_unk(int(atom.GetChiralTag()), ATOM_FEATURES['chiral_tag']) + \
               onek_encoding_unk(int(atom.GetTotalNumHs()), ATOM_FEATURES['num_Hs']) + \
               onek_encoding_unk(int(atom.GetHybridization()), ATOM_FEATURES['hybridization']) + \
               [1 if atom.GetIsAromatic() else 0] + \
               [atom.GetMass() * 0.01]  # scaled to about the same range as other features
    if functional_groups is not None:
        features += functional_groups
    return features

def bond_features(bond: Chem.rdchem.Bond):
    """Builds a feature vector for a bond.
    """
    if bond is None:
        fbond = [1] + [0] * (BOND_FDIM - 1)
    else:
        bt = bond.GetBondType()
        fbond = [
            0,  # bond is not None
            bt == Chem.rdchem.BondType.SINGLE,
            bt == Chem.rdchem.BondType.DOUBLE,
            bt == Chem.rdchem.BondType.TRIPLE,
            bt == Chem.rdchem.BondType.AROMATIC,
            (bond.GetIsConjugated() if bt is not None else 0),
            (bond.IsInRing() if bt is not None else 0)
        ]
        fbond += onek_encoding_unk(int(bond.GetStereo()), list(range(6)))
    return fbond
    
class MoleculeDatapoint:
    def __init__(self,
                 smiles: str,
                 targets: List[Optional[float]] = None,
                 row: OrderedDict = None):
        
        self.smiles = smiles
        self.targets = targets
        self.features = []
        self.row = row

    @property
    def mol(self) -> Chem.Mol:
        """Gets the corresponding list of RDKit molecules for the corresponding SMILES."""
        mol = SMILES_TO_MOL.get(self.smiles, Chem.MolFromSmiles(self.smiles))
        if cache_mol():
            SMILES_TO_MOL[self.smiles] = mol
        return mol

    def set_features(self, features: np.ndarray) -> None:
        """Sets the features of the molecule.
        """
        self.features = features

    def extend_features(self, features: np.ndarray) -> None:
        """Extends the features of the molecule.
        """
        self.features = np.append(self.features, features) if self.features is not None else features

    def num_tasks(self) -> int:
        """Returns the number of prediction tasks.
        """
        return len(self.targets)

    def set_targets(self, targets: List[Optional[float]]):
        """Sets the targets of a molecule.
        """
        self.targets = targets

    def reset_features_and_targets(self) -> None:
        """Resets the features and targets to their raw values."""
        self.features, self.targets = self.raw_features, self.raw_targets
        
        
class MoleculeDataset(Dataset):
    def __init__(self, data: List[MoleculeDatapoint]):
        self._data = data
        self._scaler = None
        self._batch_graph = None
        self._random = Random()

    def smiles(self) -> List[str]:
        return [d.smiles for d in self._data]

    def mols(self) -> List[Chem.Mol]:
        return [d.mol for d in self._data]

    def targets(self) -> List[List[Optional[float]]]:
        return [d.targets for d in self._data]

    def num_tasks(self) -> int:
        return self._data[0].num_tasks() if len(self._data) > 0 else None

    def set_targets(self, targets: List[List[Optional[float]]]) -> None:
        assert len(self._data) == len(targets)
        for i in range(len(self._data)):
            self._data[i].set_targets(targets[i])

    def reset_features_and_targets(self) -> None:
        for d in self._data:
            d.reset_features_and_targets()

    def __len__(self) -> int:
        return len(self._data)

    def __getitem__(self, item) -> Union[MoleculeDatapoint, List[MoleculeDatapoint]]:
        return self._data[item]
    
    def batch_graph(self):
        if self._batch_graph is None:
            self._batch_graph = []

            mol_graphs = []
            for d in self._data:
                mol_graphs_list = []
                if d.smiles in SMILES_TO_GRAPH:
                    mol_graph = SMILES_TO_GRAPH[d.smiles]
                else:
                    mol_graph = MolGraph(d.mol)
                    if cache_graph():
                        SMILES_TO_GRAPH[d.smiles] = mol_graph
                mol_graphs.append([mol_graph])

            self._batch_graph = [BatchMolGraph([g[i] for g in mol_graphs]) for i in range(len(mol_graphs[0]))]

        return self._batch_graph
    
    def features(self) -> List[np.ndarray]:
        """
        Returns the features associated with each molecule (if they exist).

        :return: A list of 1D numpy arrays containing the features for each molecule or None if there are no features.
        """
        if len(self._data) == 0 or self._data[0].features is None:
            return None

        return [d.features for d in self._data]
    

def index_select_ND(source: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
    """Selects the message features from source corresponding to the atom or bond indices in index.
    """
    index_size = index.size()             # (num_atoms/num_bonds, max_num_bonds)
    suffix_dim = source.size()[1:]        # (hidden_size,)
    final_size = index_size + suffix_dim  # (num_atoms/num_bonds, max_num_bonds, hidden_size)

    target = source.index_select(dim=0, index=index.view(-1)) # (num_atoms/num_bonds * max_num_bonds, hidden_size)
    target = target.view(final_size)                          # (num_atoms/num_bonds, max_num_bonds, hidden_size)
    return target

class MolGraph:
    def __init__(self, mol, atom_descriptors=None):
        # Convert SMILES to RDKit molecule if necessary
        if type(mol) == str:
            mol = Chem.MolFromSmiles(mol)

        self.n_atoms = 0  # number of atoms
        self.n_bonds = 0  # number of bonds
        self.f_atoms = []  # mapping from atom index to atom features
        self.f_bonds = []  # mapping from bond index to concat(in_atom, bond) features
        self.a2b = []  # mapping from atom index to incoming bond indices
        self.b2a = []  # mapping from bond index to the index of the atom the bond is coming from
        self.b2revb = []  # mapping from bond index to the index of the reverse bond

        # Get atom features
        self.f_atoms = [atom_features(atom) for atom in mol.GetAtoms()]
        if atom_descriptors is not None:
            self.f_atoms = [f_atoms + descs.tolist() for f_atoms, descs in zip(self.f_atoms, atom_descriptors)]

        self.n_atoms = len(self.f_atoms)

        # Initialize atom to bond mapping for each atom
        for _ in range(self.n_atoms):
            self.a2b.append([])

        # Get bond features
        for a1 in range(self.n_atoms):
            for a2 in range(a1 + 1, self.n_atoms):
                bond = mol.GetBondBetweenAtoms(a1, a2)

                if bond is None:
                    continue

                f_bond = bond_features(bond)
                self.f_bonds.append(self.f_atoms[a1] + f_bond)
                self.f_bonds.append(self.f_atoms[a2] + f_bond)

                # Update index mappings
                b1 = self.n_bonds
                b2 = b1 + 1
                self.a2b[a2].append(b1)  # b1 = a1 --> a2
                self.b2a.append(a1)
                self.a2b[a1].append(b2)  # b2 = a2 --> a1
                self.b2a.append(a2)
                self.b2revb.append(b2)
                self.b2revb.append(b1)
                self.n_bonds += 2


class BatchMolGraph:
    """A `BatchMolGraph` represents the graph structure and featurization of a batch of molecules.
    """

    def __init__(self, mol_graphs: List[MolGraph]):
        self.atom_fdim = get_atom_fdim()
        self.bond_fdim = get_bond_fdim()

        # Start n_atoms and n_bonds at 1 b/c zero padding
        self.n_atoms = 1  # number of atoms (start at 1 b/c need index 0 as padding)
        self.n_bonds = 1  # number of bonds (start at 1 b/c need index 0 as padding)
        self.a_scope = []  # list of tuples indicating (start_atom_index, num_atoms) for each molecule
        self.b_scope = []  # list of tuples indicating (start_bond_index, num_bonds) for each molecule

        # All start with zero padding so that indexing with zero padding returns zeros
        f_atoms = [[0] * self.atom_fdim]  # atom features
        f_bonds = [[0] * self.bond_fdim]  # combined atom/bond features
        a2b = [[]]  # mapping from atom index to incoming bond indices
        b2a = [0]  # mapping from bond index to the index of the atom the bond is coming from
        b2revb = [0]  # mapping from bond index to the index of the reverse bond
        for mol_graph in mol_graphs:
            f_atoms.extend(mol_graph.f_atoms)
            f_bonds.extend(mol_graph.f_bonds)

            for a in range(mol_graph.n_atoms):
                a2b.append([b + self.n_bonds for b in mol_graph.a2b[a]])

            for b in range(mol_graph.n_bonds):
                b2a.append(self.n_atoms + mol_graph.b2a[b])
                b2revb.append(self.n_bonds + mol_graph.b2revb[b])

            self.a_scope.append((self.n_atoms, mol_graph.n_atoms))
            self.b_scope.append((self.n_bonds, mol_graph.n_bonds))
            self.n_atoms += mol_graph.n_atoms
            self.n_bonds += mol_graph.n_bonds

        self.max_num_bonds = max(1, max(len(in_bonds) for in_bonds in a2b))  # max with 1 to fix a crash in rare case of all single-heavy-atom mols

        self.f_atoms = torch.FloatTensor(f_atoms)
        self.f_bonds = torch.FloatTensor(f_bonds)
        self.a2b = torch.LongTensor([a2b[a] + [0] * (self.max_num_bonds - len(a2b[a])) for a in range(self.n_atoms)])
        self.b2a = torch.LongTensor(b2a)
        self.b2revb = torch.LongTensor(b2revb)
        self.b2b = None  # try to avoid computing b2b b/c O(n_atoms^3)
        self.a2a = None  # only needed if using atom messages

    def get_components(self, atom_messages: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor,
                                                                   torch.LongTensor, torch.LongTensor, torch.LongTensor,
                                                                   List[Tuple[int, int]], List[Tuple[int, int]]]:
        return self.f_atoms, self.f_bonds, self.a2b, self.b2a, self.b2revb, self.a_scope, self.b_scope

    def get_b2b(self) -> torch.LongTensor:
        """Computes (if necessary) and returns a mapping from each bond index to all the incoming bond indices.
        """
        if self.b2b is None:
            b2b = self.a2b[self.b2a]  # num_bonds x max_num_bonds
            # b2b includes reverse edge for each bond so need to mask out
            revmask = (b2b != self.b2revb.unsqueeze(1).repeat(1, b2b.size(1))).long()  # num_bonds x max_num_bonds
            self.b2b = b2b * revmask

        return self.b2b

    def get_a2a(self) -> torch.LongTensor:
        """Computes (if necessary) and returns a mapping from each atom index to all neighboring atom indices.
        """
        if self.a2a is None:
            # b = a1 --> a2
            # a2b maps a2 to all incoming bonds b
            # b2a maps each bond b to the atom it comes from a1
            # thus b2a[a2b] maps atom a2 to neighboring atoms a1
            self.a2a = self.b2a[self.a2b]  # num_atoms x max_num_bonds

        return self.a2a

# Atom feature sizes
MAX_ATOMIC_NUM = 100
ATOM_FEATURES = {
    'atomic_num': list(range(MAX_ATOMIC_NUM)),
    'degree': [0, 1, 2, 3, 4, 5],
    'formal_charge': [-1, -2, 1, 2, 0],
    'chiral_tag': [0, 1, 2, 3],
    'num_Hs': [0, 1, 2, 3, 4],
    'hybridization': [
        Chem.rdchem.HybridizationType.SP,
        Chem.rdchem.HybridizationType.SP2,
        Chem.rdchem.HybridizationType.SP3,
        Chem.rdchem.HybridizationType.SP3D,
        Chem.rdchem.HybridizationType.SP3D2
    ],
}

# Distance feature sizes
PATH_DISTANCE_BINS = list(range(10))
THREE_D_DISTANCE_MAX = 20
THREE_D_DISTANCE_STEP = 1
THREE_D_DISTANCE_BINS = list(range(0, THREE_D_DISTANCE_MAX + 1, THREE_D_DISTANCE_STEP))

# len(choices) + 1 to include room for uncommon values; + 2 at end for IsAromatic and mass
ATOM_FDIM = sum(len(choices) + 1 for choices in ATOM_FEATURES.values()) + 2
EXTRA_ATOM_FDIM = 0
BOND_FDIM = 14


def get_atom_fdim() -> int:
    """Gets the dimensionality of the atom feature vector."""
    return ATOM_FDIM + EXTRA_ATOM_FDIM

def get_bond_fdim() -> int:
    """Gets the dimensionality of the bond feature vector.
    """
    return BOND_FDIM + get_atom_fdim()


def onek_encoding_unk(value: int, choices: List[int]) -> List[int]:
    encoding = [0] * (len(choices) + 1)
    index = choices.index(value) if value in choices else -1
    encoding[index] = 1

    return encoding

def atom_features(atom: Chem.rdchem.Atom, functional_groups: List[int] = None) -> List[Union[bool, int, float]]:
    """Builds a feature vector for an atom.
    """
    features = onek_encoding_unk(atom.GetAtomicNum() - 1, ATOM_FEATURES['atomic_num']) + \
               onek_encoding_unk(atom.GetTotalDegree(), ATOM_FEATURES['degree']) + \
               onek_encoding_unk(atom.GetFormalCharge(), ATOM_FEATURES['formal_charge']) + \
               onek_encoding_unk(int(atom.GetChiralTag()), ATOM_FEATURES['chiral_tag']) + \
               onek_encoding_unk(int(atom.GetTotalNumHs()), ATOM_FEATURES['num_Hs']) + \
               onek_encoding_unk(int(atom.GetHybridization()), ATOM_FEATURES['hybridization']) + \
               [1 if atom.GetIsAromatic() else 0] + \
               [atom.GetMass() * 0.01]  # scaled to about the same range as other features
    if functional_groups is not None:
        features += functional_groups
    return features


def initialize_weights(model: nn.Module) -> None:
    """Initializes the weights of a model in place.
    """
    for param in model.parameters():
        if param.dim() == 1:
            nn.init.constant_(param, 0)
        else:
            nn.init.xavier_normal_(param)

class MPNEncoder(nn.Module):
    def __init__(self, args, atom_fdim, bond_fdim):
        super(MPNEncoder, self).__init__()
        self.atom_fdim = atom_fdim
        self.bond_fdim = bond_fdim
        self.hidden_size = args.hidden_size
        self.bias = args.bias
        self.depth = args.depth
        self.dropout = args.dropout
        self.layers_per_message = 1
        self.undirected = False
        self.atom_messages = False
        self.device = args.device
        self.aggregation = args.aggregation
        self.aggregation_norm = args.aggregation_norm

        self.dropout_layer = nn.Dropout(p=self.dropout)
        self.act_func = nn.ReLU()
        self.cached_zero_vector = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)

        # Input
        input_dim = self.bond_fdim
        self.W_i = nn.Linear(input_dim, self.hidden_size, bias=self.bias)
        w_h_input_size = self.hidden_size

        # Shared weight matrix across depths (default)
        self.W_h = nn.Linear(w_h_input_size, self.hidden_size, bias=self.bias)
        self.W_o = nn.Linear(self.atom_fdim + self.hidden_size, self.hidden_size)

    def forward(self, mol_graph):
        """Encodes a batch of molecular graphs.
        """
        f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope = mol_graph.get_components()
        f_atoms, f_bonds, a2b, b2a, b2revb = f_atoms.to(self.device), f_bonds.to(self.device), a2b.to(self.device), b2a.to(self.device), b2revb.to(self.device)

        input = self.W_i(f_bonds)  # num_bonds x hidden_size
        message = self.act_func(input)  # num_bonds x hidden_size

        # Message passing
        for depth in range(self.depth - 1):
            # m(a1 -> a2) = [sum_{a0 \in nei(a1)} m(a0 -> a1)] - m(a2 -> a1)
            # message      a_message = sum(nei_a_message)      rev_message
            nei_a_message = index_select_ND(message, a2b)  # num_atoms x max_num_bonds x hidden
            a_message = nei_a_message.sum(dim=1)  # num_atoms x hidden
            rev_message = message[b2revb]  # num_bonds x hidden
            message = a_message[b2a] - rev_message  # num_bonds x hidden

            message = self.W_h(message)
            message = self.act_func(input + message)  # num_bonds x hidden_size
            message = self.dropout_layer(message)  # num_bonds x hidden

        a2x = a2b
        nei_a_message = index_select_ND(message, a2x)  # num_atoms x max_num_bonds x hidden
        a_message = nei_a_message.sum(dim=1)  # num_atoms x hidden
        a_input = torch.cat([f_atoms, a_message], dim=1)  # num_atoms x (atom_fdim + hidden)
        atom_hiddens = self.act_func(self.W_o(a_input))  # num_atoms x hidden
        atom_hiddens = self.dropout_layer(atom_hiddens)  # num_atoms x hidden

        # Readout
        mol_vecs = []
        for i, (a_start, a_size) in enumerate(a_scope):
            if a_size == 0:
                mol_vecs.append(self.cached_zero_vector)
            else:
                cur_hiddens = atom_hiddens.narrow(0, a_start, a_size)
                mol_vec = cur_hiddens  # (num_atoms, hidden_size)
                if self.aggregation == 'mean':
                    mol_vec = mol_vec.sum(dim=0) / a_size
                elif self.aggregation == 'sum':
                    mol_vec = mol_vec.sum(dim=0)
                elif self.aggregation == 'norm':
                    mol_vec = mol_vec.sum(dim=0) / self.aggregation_norm
                mol_vecs.append(mol_vec)

        mol_vecs = torch.stack(mol_vecs, dim=0)  # (num_molecules, hidden_size)

        return mol_vecs  # num_molecules x hidden
    

class MPN(nn.Module):
    def __init__(self, args, atom_fdim=None, bond_fdim=None):
        super(MPN, self).__init__()
        self.atom_fdim = atom_fdim or get_atom_fdim()
        self.bond_fdim = bond_fdim or get_bond_fdim()
        self.device = args.device
        self.encoder = MPNEncoder(args, self.atom_fdim, self.bond_fdim)

    def forward(self, batch):
        """Encodes a batch of molecules.
        """
        if type(batch[0]) != BatchMolGraph:
            batch = [mol2graph(b) for b in batch]

        encodings = [self.encoder(batch[0])]
        output = reduce(lambda x, y: torch.cat((x, y), dim=1), encodings)
        return output
    

class MoleculeModel(nn.Module):
    def __init__(self, args, featurizer=False):
        super(MoleculeModel, self).__init__()

        self.classification = args.dataset_type == 'classification'
        self.featurizer = featurizer

        self.output_size = args.num_tasks

        if self.classification:
            self.sigmoid = nn.Sigmoid()
            
        self.create_encoder(args)
        self.create_ffn(args)

        initialize_weights(self)

    def create_encoder(self, args):
        self.encoder = MPN(args)

    def create_ffn(self, args):
        first_linear_dim = args.hidden_size
        dropout = nn.Dropout(args.dropout)
        activation = nn.ReLU()

        # Create FFN layers
        if args.ffn_num_layers == 1:
            ffn = [
                dropout,
                nn.Linear(first_linear_dim, self.output_size)
            ]
        else:
            ffn = [
                dropout,
                nn.Linear(first_linear_dim, args.ffn_hidden_size)
            ]
            for _ in range(args.ffn_num_layers - 2):
                ffn.extend([
                    activation,
                    dropout,
                    nn.Linear(args.ffn_hidden_size, args.ffn_hidden_size),
                ])
            ffn.extend([
                activation,
                dropout,
                nn.Linear(args.ffn_hidden_size, self.output_size),
            ])

        # Create FFN model
        self.ffn = nn.Sequential(*ffn)

    def featurize(self, batch, features_batch=None, atom_descriptors_batch=None):
        """Computes feature vectors of the input by running the model except for the last layer.
        """
        return self.ffn[:-1](self.encoder(batch, features_batch, atom_descriptors_batch))

    def forward(self, batch):
        output = self.ffn(self.encoder(batch))

        # Don't apply sigmoid during training b/c using BCEWithLogitsLoss
        if self.classification and not self.training:
            output = self.sigmoid(output)
            
        return output

from torch.optim.lr_scheduler import _LRScheduler
from torch.optim import Adam, Optimizer

class NoamLR(_LRScheduler):
    """
    Noam learning rate scheduler with piecewise linear increase and exponential decay.

    The learning rate increases linearly from init_lr to max_lr over the course of
    the first warmup_steps (where :code:`warmup_steps = warmup_epochs * steps_per_epoch`).
    Then the learning rate decreases exponentially from :code:`max_lr` to :code:`final_lr` over the
    course of the remaining :code:`total_steps - warmup_steps` (where :code:`total_steps =
    total_epochs * steps_per_epoch`). This is roughly based on the learning rate
    schedule from `Attention is All You Need <https://arxiv.org/abs/1706.03762>`_, section 5.3.
    """
    def __init__(self,
                 optimizer: Optimizer,
                 warmup_epochs: List[Union[float, int]],
                 total_epochs: List[int],
                 steps_per_epoch: int,
                 init_lr: List[float],
                 max_lr: List[float],
                 final_lr: List[float]):

        assert len(optimizer.param_groups) == len(warmup_epochs) == len(total_epochs) == len(init_lr) == \
               len(max_lr) == len(final_lr)

        self.num_lrs = len(optimizer.param_groups)

        self.optimizer = optimizer
        self.warmup_epochs = np.array(warmup_epochs)
        self.total_epochs = np.array(total_epochs)
        self.steps_per_epoch = steps_per_epoch
        self.init_lr = np.array(init_lr)
        self.max_lr = np.array(max_lr)
        self.final_lr = np.array(final_lr)

        self.current_step = 0
        self.lr = init_lr
        self.warmup_steps = (self.warmup_epochs * self.steps_per_epoch).astype(int)
        self.total_steps = self.total_epochs * self.steps_per_epoch
        self.linear_increment = (self.max_lr - self.init_lr) / self.warmup_steps

        self.exponential_gamma = (self.final_lr / self.max_lr) ** (1 / (self.total_steps - self.warmup_steps))

        super(NoamLR, self).__init__(optimizer)

    def get_lr(self) -> List[float]:
        return list(self.lr)

    def step(self, current_step: int = None):
        if current_step is not None:
            self.current_step = current_step
        else:
            self.current_step += 1

        for i in range(self.num_lrs):
            if self.current_step <= self.warmup_steps[i]:
                self.lr[i] = self.init_lr[i] + self.current_step * self.linear_increment[i]
            elif self.current_step <= self.total_steps[i]:
                self.lr[i] = self.max_lr[i] * (self.exponential_gamma[i] ** (self.current_step - self.warmup_steps[i]))
            else:  # theoretically this case should never be reached since training should stop at total_steps
                self.lr[i] = self.final_lr[i]

            self.optimizer.param_groups[i]['lr'] = self.lr[i]

import threading

def construct_molecule_batch(data):
    data = MoleculeDataset(data)
    data.batch_graph()  # Forces computation and caching of the BatchMolGraph for the molecules
    return data

class MoleculeSampler(Sampler):
    def __init__(self, dataset, shuffle=False, seed=0):
        super(Sampler, self).__init__()

        self.dataset = dataset
        self.shuffle = shuffle
        self._random = Random(seed)
        self.positive_indices = self.negative_indices = None
        self.length = len(self.dataset)

    def __iter__(self):
        indices = list(range(len(self.dataset)))
        if self.shuffle:
            self._random.shuffle(indices)
        return iter(indices)

    def __len__(self):
        return self.length
    

class MoleculeDataLoader(DataLoader):
    def __init__(self,
                 dataset: MoleculeDataset,
                 batch_size: int = 50,
                 num_workers: int = 8,
                 shuffle: bool = False,
                 seed: int = 0):

        self._dataset = dataset
        self._batch_size = batch_size
        self._num_workers = num_workers
        self._shuffle = shuffle
        self._seed = seed
        self._context = None
        self._class_balance = False
        self._timeout = 0
        is_main_thread = threading.current_thread() is threading.main_thread()
        
        if not is_main_thread and self._num_workers > 0:
            self._context = 'forkserver'  # In order to prevent a hanging
            self._timeout = 3600  # Just for sure that the DataLoader won't hang

        self._sampler = MoleculeSampler(
            dataset=self._dataset,
            shuffle=self._shuffle,
            seed=self._seed
        )

        super(MoleculeDataLoader, self).__init__(
            dataset=self._dataset,
            batch_size=self._batch_size,
            sampler=self._sampler,
            num_workers=self._num_workers,
            collate_fn=construct_molecule_batch,
            multiprocessing_context=self._context,
            timeout=self._timeout
        )

    @property
    def targets(self) -> List[List[Optional[float]]]:
        if self._class_balance or self._shuffle:
            raise ValueError('Cannot safely extract targets when class balance or shuffle are enabled.')

        return [self._dataset[index].targets for index in self._sampler]

    @property
    def iter_size(self) -> int:
        return len(self._sampler)

    def __iter__(self) -> Iterator[MoleculeDataset]:
        return super(MoleculeDataLoader, self).__iter__()