# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import torch
import numpy as np
import mdtraj

import boltzgen.zmatrix as zmatrix
import boltzgen.internal as ics
import boltzgen.mixed as mixed

import normflow as nf
from boltzgen.flows import CoordinateTransform
from boltzgen.distributions import Boltzmann

from tqdm import tqdm
from matplotlib import pyplot as plt

from autograd import grad
from autograd import numpy as np
from openmmtools.constants import kB
from simtk import openmm as mm
from simtk import unit
from simtk.openmm import app
from openmmtools.testsystems import AlanineDipeptideVacuum

import boltzgen.openmm_interface as omi

# Load the alanine dipeptide trajectory
aldp_traj = mdtraj.load('alanine_dipeptide/trajectory/aldp100000.h5')

# +
# Set up coordinate transformation

z_matrix = [
    (1, [4, 5, 6]),
    (0, [1, 4, 5]),
    (2, [1, 0, 4]),
    (3, [1, 0, 2]),
    (7, [6, 4, 5]),
    (9, [8, 6, 7]),
    (10, [8, 6, 9]),
    (11, [10, 8, 9]),
    (12, [10, 8, 11]),
    (13, [10, 11, 12]),
    (17, [16, 14, 15]),
    (19, [18, 16, 17]),
    (20, [18, 19, 16]),
    (21, [18, 19, 20])
]

backbone_indices = [4, 5, 6, 8, 14, 15, 16, 18]

aldp_traj.center_coordinates()

# superpose on the backbone
ind = aldp_traj.top.select("backbone")

aldp_traj.superpose(aldp_traj, 0, atom_indices=ind, ref_atom_indices=ind)

# Gather the training data into a pytorch Tensor with the right shape
training_data = aldp_traj.xyz
n_atoms = training_data.shape[1]
n_dim = n_atoms * 3
training_data_npy = training_data.reshape(-1, n_dim)
training_data = torch.from_numpy(training_data_npy.astype("float32"))

# +
# Set up simulation object for energy computation

temperature = 1000
kT = kB * temperature

testsystem = AlanineDipeptideVacuum
vacuum_sim = app.Simulation(testsystem.topology, testsystem.system,
                            mm.LangevinIntegrator(temperature * unit.kelvin , 1.0 / unit.picosecond, 
                                                  1.0 * unit.femtosecond),
                            platform=mm.Platform.getPlatformByName('CPU'))
vacuum_sim.context.setPositions(testsystem.positions)

# Energy function
openmm_energy = omi.OpenMMEnergyInterface.apply

# +
# Set up flow model

# Define flows
K = 16
torch.manual_seed(0)

# Device to be used
enable_cuda = True
device = torch.device('cuda' if torch.cuda.is_available() and enable_cuda else 'cpu')

latent_size = 60
b = torch.Tensor([1 if i % 2 == 0 else 0 for i in range(latent_size)])
flows = []
for i in range(K):
    s = nf.nets.MLP([latent_size, 4 * latent_size, 4 * latent_size, latent_size],
                    output_fn='tanh', output_scale=2, init_zeros=True)
    t = nf.nets.MLP([latent_size, 4 * latent_size, 4 * latent_size, latent_size],
                    init_zeros=True)
    if i % 2 == 0:
        flows += [nf.flows.MaskedAffineFlow(b, s, t)]
    else:
        flows += [nf.flows.MaskedAffineFlow(1 - b, s, t)]
    #flows += [nf.flows.Planar(latent_size)]
    flows += [nf.flows.ActNorm(latent_size)]
flows += [CoordinateTransform(training_data, 66, z_matrix, backbone_indices)]

# Set prior and q0
prior = Boltzmann(vacuum_sim.context, temperature, energy_cut=1e2, energy_max=1e20)
q0 = nf.distributions.DiagGaussian(latent_size)

# Construct flow model
nfm = nf.NormalizingFlow(q0=q0, flows=flows, p=prior)

# Move model on GPU if available
nfm = nfm.to(device)
nfm = nfm.double()

ind = torch.randint(len(training_data), (128, ))
x = training_data[ind, :].double().to(device)
kld = nfm.forward_kld(x)

# +
# Train model
batch_size = 128
max_iter = 50000
n_data = len(training_data)
eval_rkld = 50


loss_hist = np.array([])
fkld_hist = np.array([])
rkld_hist = np.array([])

optimizer = torch.optim.AdamW(nfm.parameters(), lr=2e-4, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
    optimizer.zero_grad()
    ind = torch.randint(n_data, (batch_size, ))
    x = training_data[ind, :].double().to(device)
    fkld = nfm.forward_kld(x)
    loss = fkld
    if not torch.isnan(loss) and loss < 0:
        loss.backward()
        torch.nn.utils.clip_grad_value_(nfm.parameters(), .01)
        gradient_norm = torch.nn.utils.clip_grad.clip_grad_norm_(nfm.parameters(), 500.)
        optimizer.step()
    
    loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
    fkld_hist = np.append(fkld_hist, fkld.to('cpu').data.numpy())
    if (it + 1) % eval_rkld == 0:
        rkld = nfm.reverse_kld(num_samples=batch_size)
        rkld_hist = np.append(rkld_hist, rkld.to('cpu').data.numpy())
# -

loss_hist[loss_hist > 0] = np.nan
plt.plot(loss_hist)
plt.show()
plt.plot(fkld_hist)
plt.show()
plt.plot(rkld_hist)
plt.show()

nfm.save('models/real_nvp_ml_50000.pt')

z, log_q = nfm.sample(10000)
z_d = training_data[::10, :].double().to(device)
log_p = nfm.p.log_prob(z)
log_p_d = nfm.p.log_prob(z_d)
log_q_d = nfm.log_prob(z_d)
E_np = -log_p.data.numpy()
E_q_np = -log_q.data.numpy()
E_q_d_np = -log_q_d.data.numpy()
E_d_np = -log_p_d.data.numpy()
w = torch.exp(log_p - log_q)
w_np = w.data.numpy()

# +
plt.hist(E_d_np, bins=100, alpha=0.5, label='data', range=[-0.03, -0.005])
plt.hist(E_np, bins=100, alpha=0.5, label='samples, -log(p)', range=[-0.03, -0.005])
plt.legend(loc='upper right')
plt.savefig('E_dist_p.png', dpi=300)
plt.show()

plt.hist(E_q_d_np, bins=100, alpha=0.5, label='data, -log(q)', range=[-400, -300])
plt.hist(E_q_np, bins=100, alpha=0.5, label='samples, -log(q)', range=[-400, -300])
plt.legend(loc='upper right')
plt.savefig('E_dist_q.png', dpi=300)
plt.show()

# +
z, _ = nfm.sample(100000)
z, _ = nfm.flows[-1].inverse(z)
z_d, _ = nfm.flows[-1].inverse(training_data.double())
z_np = z.data.numpy()
z_d_np = z_d.data.numpy()

for i in range(60):
    print(i)
    print(np.mean(z_d_np[:, i]))
    plt.hist(z_d_np[:, i], bins=100, alpha=0.5, label='data')#, range=[-3.1, 3.1])
    plt.hist(z_np[:, i], bins=100, alpha=0.5, label='samples')#, range=[-3.1, 3.1])
    plt.legend(loc='upper right')
    plt.show()

StarcoderdataJupyterScriptsDedupFiltered