# training and evaluation for interpolation and fate prediction tasks
# on weinreb et al. dataset

import torch
import torch.nn.functional as F 
from torch import nn, optim

import annoy
import tqdm

from geomloss import SamplesLoss

import numpy as np
import pandas as pd
import scipy.stats

from collections import OrderedDict, Counter
from joblib import Parallel, delayed 
from types import SimpleNamespace 
from time import strftime, localtime 

import argparse
import copy
import glob
import itertools
import json 
import os
import sys

import train

FATE_DIR = "data/Klein2020_fate"
FATE_TRAIN_PATH= os.path.join(FATE_DIR, "fate_train.pt")
FATE_ANN = os.path.join(FATE_DIR, "50_20_10")
FATE_TEST_PATH = os.path.join(FATE_DIR, "fate_test.pt")
IMPUTE_DATA_PATH = "data/Klein2020_impute.pt"
WEIGHT_DIR = 'data/Klein2020_weights'

def init_config(args): 

    config = SimpleNamespace(

        seed = args.seed,
        timestamp = strftime("%a, %d %b %Y %H:%M:%S", localtime()), 

        # data parameters
        data_dir = args.data_dir, 
        data_path = args.data_path,
        weight_path = args.weight_path,
        weight = args.weight, 

        # model parameters 
        activation = args.activation,
        layers = args.layers, 
        k_dim = args.k_dim, 

        # pretraining parameters
        pretrain_burnin = 50, 
        pretrain_sd = 0.1, 
        pretrain_lr = 1e-9, 
        pretrain_epochs = args.pretrain_epochs, 

        # training parameters
        train_dt = args.train_dt, 
        train_sd = args.train_sd, 
        train_batch_size = args.train_batch, 
        ns = 2000,
        train_burnin = 100, 
        train_tau = args.train_tau, 
        train_epochs = args.train_epochs, 
        train_lr = args.train_lr, 
        train_clip = args.train_clip,
        save = args.save,

        # loss parameters
        sinkhorn_scaling = 0.7, 
        sinkhorn_blur = 0.1, 

        # file parameters
        out_dir = args.out_dir, 
        out_name = args.out_dir.split('/')[-1], 
        pretrain_pt = os.path.join(args.out_dir, 'pretrain.pt'), 
        train_pt = os.path.join(args.out_dir, 'train.{}.pt'), 
        train_log = os.path.join(args.out_dir, 'train.log'), 
        done_log = os.path.join(args.out_dir, 'done.log'),
        config_pt = os.path.join(args.out_dir, 'config.pt'),
    )

    config.train_t = []
    config.test_t = []
   
    if not os.path.exists(args.out_dir):
        print('Making directory at {}'.format(args.out_dir))
        os.makedirs(args.out_dir)
    else: 
        print('Directory exists at {}'.format(args.out_dir))

    return config

def load_data(config, base_dir = "."): 
   
    data_pt = torch.load(os.path.join(base_dir, config.data_path))
    x = data_pt['xp'] 
    y = data_pt['y']

    config.x_dim = x[0].shape[-1]
    config.t = y[-1] - y[0]
    
    y_start = y[config.start_t]
    y_ = [y_ for y_ in y if y_ > y_start]

    weight_pt = torch.load(os.path.join(base_dir, config.weight_path))
    x = [xx[m] for xx, m in zip(x, weight_pt['m'])]

    w_ = weight_pt['w'][config.start_t]
    w = {(y_start, yy): torch.from_numpy(np.exp((yy - y_start)*w_)) for yy in y_}

    return x, y, w

def train_fate(args): 

    a = copy.copy(args)

    # data

    a.data_path = os.path.join(a.data_dir, 'fate_train.pt')

    weight = os.path.basename(a.weight_path)
    weight = weight.split('.')[0].split('-')[-1]
    a.weight = weight

    # out directory

    name = (
        "{weight}-"
        "{activation}_{layers}_{k_dim}-"
        "{train_tau}"
    ).format(**a.__dict__)

    a.out_dir = os.path.join(args.out_dir, name, 'seed_{}'.format(a.seed))
    config = init_config(a)
    
    config.start_t = 0
    config.train_t = [1, 2]
    
    x, y, w = load_data(config)

    return x, y, w, config

def evaluate_fate(args, config): 

    # -- load data
 
    data_pt = torch.load(os.path.join(config.data_dir, 'fate_test.pt'))
    x = data_pt['x']
    y = data_pt['y']
    t = data_pt['t']

    ay_path = os.path.join(config.data_dir, '50_20_10')
    ay = annoy.AnnoyIndex(config.x_dim, 'euclidean')
    ay.load(ay_path + '.ann')
    with open(ay_path + '.txt', 'r') as f: 
        cy = np.array([line.strip() for line in f])
    
    # -- initialize

    device, kwargs = train.init(args)

    # -- model
    
    model = train.AutoGenerator(config)
    
    log_str = '{} {:.5f} {:.3e} {:.5f} {:.3e} {:d}'
    log_handle = open(os.path.join(config.out_dir, 'fate.log'), 'w') 

    names_ = []
    scores_ = []
    masks_ = []

    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        name = os.path.basename(train_pt).split('.')[1]

        checkpoint = torch.load(train_pt)
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model)

        # -- evaluate 
        torch.manual_seed(0) 

        time_elapsed = config.t
        num_steps = int(np.round(time_elapsed / config.train_dt))

        scores = []
        mask = []
        pbar = tqdm.tqdm(range(len(x)), desc = "[fate:{}]".format(name))
        for i in pbar:

            # expand data point
            x_i = x[i].expand(config.ns, -1).to(device)

            # simulate forward
            for _ in range(num_steps): 
                z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                z = z.to(device)
                x_i = model._step(x_i, dt = config.train_dt, z = z)
            x_i_ = x_i.detach().cpu().numpy()
            
            # predict
            yp = []
            for j in range(x_i_.shape[0]): 
                nn = cy[ay.get_nns_by_vector(x_i_[j], 20)]
                nn = Counter(nn).most_common(2)
                label, num = nn[0]
                if len(nn) > 1: 
                    _, num2 = nn[1]
                    if num == num2:  # deal with ties by setting it to the default class
                        label = 'Other'
                yp.append(label)
            yp = Counter(yp)

            # may want to save yp instead
            num_neu = yp['Neutrophil'] + 1 # use pseudocounts for scoring
            num_total = yp['Neutrophil'] + yp['Monocyte'] + 2
            score = num_neu / num_total
            scores.append(score)
            num_total = yp['Neutrophil'] + yp['Monocyte']
            mask.append(num_total > 0)
         
        scores = np.array(scores)
        mask = np.array(mask)

        r, pval = scipy.stats.pearsonr(y, scores)
        r_masked, pval_masked = scipy.stats.pearsonr(y[mask], scores[mask])

        log = log_str.format(name, r, pval, r_masked, pval_masked, mask.sum())
        log_handle.write(log + '\n')
        print(log)

        names_.append(name)
        scores_.append(scores)
        masks_.append(mask)

    log_handle.close()

    torch.save({
        'scores': scores_, 
        'mask': masks_, 
        'names': names_
    }, os.path.join(config.out_dir, 'fate.pt'))
   
def train_interpolate(args, data_path = IMPUTE_DATA_PATH): 

    a = copy.copy(args)
    
    weight = os.path.basename(a.weight_path)
    weight = weight.split('.')[0].split('-')[-1]
    a.weight = weight

    name = (
        "{weight}-"
        "{activation}_{layers}_{k_dim}-"
        "{train_dt}_{train_sd}_{train_tau}-"
        "{train_batch}_{train_clip}_{train_lr}"
    ).format(**a.__dict__)

    a.out_dir = os.path.join(args.out_dir, name, 'seed_{}'.format(a.seed)) 
    config = init_config(a)

    config.start_t = 0
    config.train_t = [2]
    config.test_t = [1]
    
    x, y, w = load_data(config)

    return x, y, w, config 

def evaluate_interpolate_model(args, config): 
   
    if not os.path.exists(config.done_log): 
        print(config.done_log, 'does not exist. Skipping.')
        return

    log_path = os.path.join(config.out_dir, 'interpolate.log')
    if os.path.exists(log_path): 
        print(log_path, 'exists. Skipping.')
        return

    x, y, w = load_data(config)

    # -- initialize
    device, kwargs = train.init(args) 
    model = train.AutoGenerator(config) 
    
    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)

    losses_xy = []
    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        checkpoint = torch.load(train_pt) 
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model) 

        name = os.path.basename(train_pt).split('.')[1]

        # -- evaluate

        def _evaluate_impute_model(t_cur): 
        
            torch.manual_seed(0)
            np.random.seed(0)

            t_prev = config.start_t
            y_prev = int(y[t_prev])
            y_cur = int(y[t_cur])
            time_elapsed = y_cur - y_prev
            num_steps = int(np.round(time_elapsed / config.train_dt))

            dat_prev = x[t_prev].to(device)
            dat_cur = x[t_cur].to(device)
            w_prev = train.get_weight(w[(y_prev, y_cur)], time_elapsed).cpu().numpy()

            x_s = []
            x_i_ = train.weighted_samp(dat_prev, args.evaluate_n, w_prev)

            for i in range(int(args.evaluate_n / config.ns)): 

                x_i = x_i_[i*config.ns:(i+1)*config.ns,]

                for _ in range(num_steps): 
                    z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                    z = z.to(device)
                    x_i = model._step(x_i, dt = config.train_dt, z = z)

                x_s.append(x_i.detach())

            x_s = torch.cat(x_s)

            loss_xy = ot_solver(x_s, dat_cur)
            return loss_xy
 
        for t in config.train_t:
            y_ = y[t]
            loss_xy = _evaluate_impute_model(t).item()
            losses_xy.append((name, 'train', y_, loss_xy))
        try:
            for t in config.test_t: 
                y_ = y[t]
                loss_xy = _evaluate_impute_model(t).item()
                losses_xy.append((name, 'test', y_, loss_xy))
        except AttributeError:
            continue

    losses_xy = pd.DataFrame(losses_xy, columns = ['epoch', 'eval', 't', 'loss'])
    losses_xy.to_csv(log_path, sep = '\t', index = False)
    print('Wrote results to', log_path)

def evaluate_interpolate_model_baseline(args, config): 
    
    if not os.path.exists(config.done_log): 
        print(config.done_log, 'does not exist. Skipping.')
        return

    log_path = os.path.join(config.out_dir, 'baseline.log')
    if os.path.exists(log_path): 
        print(log_path, 'exists. Skipping.')
        return

    x, y, w = load_data(config)

    # -- initialize
    device, kwargs = train.init(args) 
    model = train.AutoGenerator(config) 
    
    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)

    losses_xy = []
    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        checkpoint = torch.load(train_pt) 
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model) 

        name = os.path.basename(train_pt).split('.')[1]

        # -- evaluate
        torch.manual_seed(0)
        np.random.seed(0)

        t_cur = 1

        t_prev = config.start_t
        y_prev = int(y[t_prev])
        y_cur = int(y[t_cur])
        time_elapsed = y_cur - y_prev
        num_steps = int(np.round(time_elapsed / config.train_dt))

        dat_prev = x[t_prev].to(device)
        w_prev = train.get_weight(w[(y_prev, y_cur)], time_elapsed).cpu().numpy()

        x_s = []
        x_i_ = train.weighted_samp(dat_prev, args.evaluate_n, w_prev)

        for i in range(int(args.evaluate_n / config.ns)): 

            x_i = x_i_[i*config.ns:(i+1)*config.ns,]

            for _ in range(num_steps): 
                z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                z = z.to(device)
                x_i = model._step(x_i, dt = config.train_dt, z = z)

            x_s.append(x_i.detach())

        x_s = torch.cat(x_s)

        loss_xy = [name] + [ot_solver(x_s, x[t_].to(device)).item() for t_ in range(len(x))]
        losses_xy.append(loss_xy)
 
    losses_xy = pd.DataFrame(losses_xy, columns = ['epoch'] + y)
    losses_xy.to_csv(log_path, sep = '\t', index = False)
    print('Wrote results to', log_path) 

def evaluate_interpolate_data(args, config): 
    
    x, y, w = load_data(config)

    device, kwargs = train.init(args) 

    pt = torch.load(args.evaluate_data) 
    x_i = torch.from_numpy(pt['sim_xp']).float().to(device)
    y_j = x[1].to(device)

    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)
    loss_xy = ot_solver(x_i, y_j) 

    import pdb; pdb.set_trace()

def main(): 
   
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--seed', type = int, default = 0)
    parser.add_argument('--no-cuda', action = 'store_true')
    parser.add_argument('--device', default = 7, type = int)
    parser.add_argument('--out_dir', default = './experiments')
    # -- data options
    parser.add_argument('--data_path')
    parser.add_argument('--data_dir')
    parser.add_argument('--weight_path', default = None)
    # -- model options
    parser.add_argument('--loss', default = 'euclidean')
    parser.add_argument('--k_dim', default = 500, type = int)
    parser.add_argument('--activation', default = 'softplus')
    parser.add_argument('--layers', default = 1, type = int)
    # -- pretrain options
    parser.add_argument('--pretrain_lr', default = 1e-9, type = float)
    parser.add_argument('--pretrain_epochs', default = 500, type = int) 
    # -- train options 
    parser.add_argument('--train_epochs', default = 5000, type = int)
    parser.add_argument('--train_lr', default = 0.01, type = float)
    parser.add_argument('--train_dt', default = 0.1, type = float)
    parser.add_argument('--train_sd', default = 0.5, type = float)
    parser.add_argument('--train_tau', default = 0, type = float)
    parser.add_argument('--train_batch', default = 0.1, type = float)
    parser.add_argument('--train_clip', default = 0.25, type = float)
    parser.add_argument('--save', default = 100, type = int)
    # -- test options
    parser.add_argument('--evaluate_n', default = 10000, type = int)
    parser.add_argument('--evaluate_data')
    parser.add_argument('--evaluate-baseline', action = 'store_true')
    # -- run options
    parser.add_argument('--task', default = 'fate')
    parser.add_argument('--train', action = 'store_true')
    parser.add_argument('--evaluate')
    parser.add_argument('--config')
    args = parser.parse_args() 

    if args.task == 'fate': 

        if args.train: 

            args.pretrain = True
            args.train = True

            train.run(args, train_fate)

        if args.evaluate == 'model': 

            if args.config:
                config = SimpleNamespace(**torch.load(args.config))
                evaluate_fate(args, config) 
            else: 
                print('Please provide a config file')

    elif args.task == 'interpolate': 
        
            if args.train: 

                args.pretrain = True
                args.train = True

                config = train.run(args, train_interpolate) 

            elif args.evaluate: 

                if args.evaluate == 'model':
                    evaluate = evaluate_interpolate_model
                elif args.evaluate == 'data': 
                    evaluate =  evaluate_interpolate_data
                elif args.evaluate == 'baseline': 
                    evaluate = evaluate_interpolate_model_baseline
                else:
                    raise NotImplementedError

                if args.config:
                    config = SimpleNamespace(**torch.load(args.config))
                else:
                    print("Please provide a config file")

                evaluate(args, config)

            


if __name__ == '__main__':
    main()

# training and evaluation for interpolation and fate prediction tasks
# on weinreb et al. dataset

import torch
import torch.nn.functional as F 
from torch import nn, optim

import annoy
import tqdm

from geomloss import SamplesLoss

import numpy as np
import pandas as pd
import scipy.stats

from collections import OrderedDict, Counter
from joblib import Parallel, delayed 
from types import SimpleNamespace 
from time import strftime, localtime 

import argparse
import copy
import glob
import itertools
import json 
import os
import sys

import train

FATE_DIR = "data/Klein2020_fate"
FATE_TRAIN_PATH= os.path.join(FATE_DIR, "fate_train.pt")
FATE_ANN = os.path.join(FATE_DIR, "50_20_10")
FATE_TEST_PATH = os.path.join(FATE_DIR, "fate_test.pt")
IMPUTE_DATA_PATH = "data/Klein2020_impute.pt"
WEIGHT_DIR = 'data/Klein2020_weights'

def init_config(args): 

    config = SimpleNamespace(

        seed = args.seed,
        timestamp = strftime("%a, %d %b %Y %H:%M:%S", localtime()), 

        # data parameters
        data_dir = args.data_dir, 
        data_path = args.data_path,
        weight_path = args.weight_path,
        weight = args.weight, 

        # model parameters 
        activation = args.activation,
        layers = args.layers, 
        k_dim = args.k_dim, 

        # pretraining parameters
        pretrain_burnin = 50, 
        pretrain_sd = 0.1, 
        pretrain_lr = 1e-9, 
        pretrain_epochs = args.pretrain_epochs, 

        # training parameters
        train_dt = args.train_dt, 
        train_sd = args.train_sd, 
        train_batch_size = args.train_batch, 
        ns = 2000,
        train_burnin = 100, 
        train_tau = args.train_tau, 
        train_epochs = args.train_epochs, 
        train_lr = args.train_lr, 
        train_clip = args.train_clip,
        save = args.save,

        # loss parameters
        sinkhorn_scaling = 0.7, 
        sinkhorn_blur = 0.1, 

        # file parameters
        out_dir = args.out_dir, 
        out_name = args.out_dir.split('/')[-1], 
        pretrain_pt = os.path.join(args.out_dir, 'pretrain.pt'), 
        train_pt = os.path.join(args.out_dir, 'train.{}.pt'), 
        train_log = os.path.join(args.out_dir, 'train.log'), 
        done_log = os.path.join(args.out_dir, 'done.log'),
        config_pt = os.path.join(args.out_dir, 'config.pt'),
    )

    config.train_t = []
    config.test_t = []
   
    if not os.path.exists(args.out_dir):
        print('Making directory at {}'.format(args.out_dir))
        os.makedirs(args.out_dir)
    else: 
        print('Directory exists at {}'.format(args.out_dir))

    return config

def load_data(config, base_dir = "."): 
   
    data_pt = torch.load(os.path.join(base_dir, config.data_path))
    x = data_pt['xp'] 
    y = data_pt['y']

    config.x_dim = x[0].shape[-1]
    config.t = y[-1] - y[0]
    
    y_start = y[config.start_t]
    y_ = [y_ for y_ in y if y_ > y_start]

    weight_pt = torch.load(os.path.join(base_dir, config.weight_path))
    x = [xx[m] for xx, m in zip(x, weight_pt['m'])]

    w_ = weight_pt['w'][config.start_t]
    w = {(y_start, yy): torch.from_numpy(np.exp((yy - y_start)*w_)) for yy in y_}

    return x, y, w

def train_fate(args): 

    a = copy.copy(args)

    # data

    a.data_path = os.path.join(a.data_dir, 'fate_train.pt')

    weight = os.path.basename(a.weight_path)
    weight = weight.split('.')[0].split('-')[-1]
    a.weight = weight

    # out directory

    name = (
        "{weight}-"
        "{activation}_{layers}_{k_dim}-"
        "{train_tau}"
    ).format(**a.__dict__)

    a.out_dir = os.path.join(args.out_dir, name, 'seed_{}'.format(a.seed))
    config = init_config(a)
    
    config.start_t = 0
    config.train_t = [1, 2]
    
    x, y, w = load_data(config)

    return x, y, w, config

def evaluate_fate(args, config): 

    # -- load data
 
    data_pt = torch.load(os.path.join(config.data_dir, 'fate_test.pt'))
    x = data_pt['x']
    y = data_pt['y']
    t = data_pt['t']

    ay_path = os.path.join(config.data_dir, '50_20_10')
    ay = annoy.AnnoyIndex(config.x_dim, 'euclidean')
    ay.load(ay_path + '.ann')
    with open(ay_path + '.txt', 'r') as f: 
        cy = np.array([line.strip() for line in f])
    
    # -- initialize

    device, kwargs = train.init(args)

    # -- model
    
    model = train.AutoGenerator(config)
    
    log_str = '{} {:.5f} {:.3e} {:.5f} {:.3e} {:d}'
    log_handle = open(os.path.join(config.out_dir, 'fate.log'), 'w') 

    names_ = []
    scores_ = []
    masks_ = []

    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        name = os.path.basename(train_pt).split('.')[1]

        checkpoint = torch.load(train_pt)
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model)

        # -- evaluate 
        torch.manual_seed(0) 

        time_elapsed = config.t
        num_steps = int(np.round(time_elapsed / config.train_dt))

        scores = []
        mask = []
        pbar = tqdm.tqdm(range(len(x)), desc = "[fate:{}]".format(name))
        for i in pbar:

            # expand data point
            x_i = x[i].expand(config.ns, -1).to(device)

            # simulate forward
            for _ in range(num_steps): 
                z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                z = z.to(device)
                x_i = model._step(x_i, dt = config.train_dt, z = z)
            x_i_ = x_i.detach().cpu().numpy()
            
            # predict
            yp = []
            for j in range(x_i_.shape[0]): 
                nn = cy[ay.get_nns_by_vector(x_i_[j], 20)]
                nn = Counter(nn).most_common(2)
                label, num = nn[0]
                if len(nn) > 1: 
                    _, num2 = nn[1]
                    if num == num2:  # deal with ties by setting it to the default class
                        label = 'Other'
                yp.append(label)
            yp = Counter(yp)

            # may want to save yp instead
            num_neu = yp['Neutrophil'] + 1 # use pseudocounts for scoring
            num_total = yp['Neutrophil'] + yp['Monocyte'] + 2
            score = num_neu / num_total
            scores.append(score)
            num_total = yp['Neutrophil'] + yp['Monocyte']
            mask.append(num_total > 0)
         
        scores = np.array(scores)
        mask = np.array(mask)

        r, pval = scipy.stats.pearsonr(y, scores)
        r_masked, pval_masked = scipy.stats.pearsonr(y[mask], scores[mask])

        log = log_str.format(name, r, pval, r_masked, pval_masked, mask.sum())
        log_handle.write(log + '\n')
        print(log)

        names_.append(name)
        scores_.append(scores)
        masks_.append(mask)

    log_handle.close()

    torch.save({
        'scores': scores_, 
        'mask': masks_, 
        'names': names_
    }, os.path.join(config.out_dir, 'fate.pt'))
   
def train_interpolate(args, data_path = IMPUTE_DATA_PATH): 

    a = copy.copy(args)
    
    weight = os.path.basename(a.weight_path)
    weight = weight.split('.')[0].split('-')[-1]
    a.weight = weight

    name = (
        "{weight}-"
        "{activation}_{layers}_{k_dim}-"
        "{train_dt}_{train_sd}_{train_tau}-"
        "{train_batch}_{train_clip}_{train_lr}"
    ).format(**a.__dict__)

    a.out_dir = os.path.join(args.out_dir, name, 'seed_{}'.format(a.seed)) 
    config = init_config(a)

    config.start_t = 0
    config.train_t = [2]
    config.test_t = [1]
    
    x, y, w = load_data(config)

    return x, y, w, config 

def evaluate_interpolate_model(args, config): 
   
    if not os.path.exists(config.done_log): 
        print(config.done_log, 'does not exist. Skipping.')
        return

    log_path = os.path.join(config.out_dir, 'interpolate.log')
    if os.path.exists(log_path): 
        print(log_path, 'exists. Skipping.')
        return

    x, y, w = load_data(config)

    # -- initialize
    device, kwargs = train.init(args) 
    model = train.AutoGenerator(config) 
    
    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)

    losses_xy = []
    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        checkpoint = torch.load(train_pt) 
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model) 

        name = os.path.basename(train_pt).split('.')[1]

        # -- evaluate

        def _evaluate_impute_model(t_cur): 
        
            torch.manual_seed(0)
            np.random.seed(0)

            t_prev = config.start_t
            y_prev = int(y[t_prev])
            y_cur = int(y[t_cur])
            time_elapsed = y_cur - y_prev
            num_steps = int(np.round(time_elapsed / config.train_dt))

            dat_prev = x[t_prev].to(device)
            dat_cur = x[t_cur].to(device)
            w_prev = train.get_weight(w[(y_prev, y_cur)], time_elapsed).cpu().numpy()

            x_s = []
            x_i_ = train.weighted_samp(dat_prev, args.evaluate_n, w_prev)

            for i in range(int(args.evaluate_n / config.ns)): 

                x_i = x_i_[i*config.ns:(i+1)*config.ns,]

                for _ in range(num_steps): 
                    z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                    z = z.to(device)
                    x_i = model._step(x_i, dt = config.train_dt, z = z)

                x_s.append(x_i.detach())

            x_s = torch.cat(x_s)

            loss_xy = ot_solver(x_s, dat_cur)
            return loss_xy
 
        for t in config.train_t:
            y_ = y[t]
            loss_xy = _evaluate_impute_model(t).item()
            losses_xy.append((name, 'train', y_, loss_xy))
        try:
            for t in config.test_t: 
                y_ = y[t]
                loss_xy = _evaluate_impute_model(t).item()
                losses_xy.append((name, 'test', y_, loss_xy))
        except AttributeError:
            continue

    losses_xy = pd.DataFrame(losses_xy, columns = ['epoch', 'eval', 't', 'loss'])
    losses_xy.to_csv(log_path, sep = '\t', index = False)
    print('Wrote results to', log_path)

def evaluate_interpolate_model_baseline(args, config): 
    
    if not os.path.exists(config.done_log): 
        print(config.done_log, 'does not exist. Skipping.')
        return

    log_path = os.path.join(config.out_dir, 'baseline.log')
    if os.path.exists(log_path): 
        print(log_path, 'exists. Skipping.')
        return

    x, y, w = load_data(config)

    # -- initialize
    device, kwargs = train.init(args) 
    model = train.AutoGenerator(config) 
    
    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)

    losses_xy = []
    train_pts = sorted(glob.glob(config.train_pt.format('*')))
    for train_pt in train_pts: 

        checkpoint = torch.load(train_pt) 
        print('Loading model from {}'.format(train_pt))
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        print(model) 

        name = os.path.basename(train_pt).split('.')[1]

        # -- evaluate
        torch.manual_seed(0)
        np.random.seed(0)

        t_cur = 1

        t_prev = config.start_t
        y_prev = int(y[t_prev])
        y_cur = int(y[t_cur])
        time_elapsed = y_cur - y_prev
        num_steps = int(np.round(time_elapsed / config.train_dt))

        dat_prev = x[t_prev].to(device)
        w_prev = train.get_weight(w[(y_prev, y_cur)], time_elapsed).cpu().numpy()

        x_s = []
        x_i_ = train.weighted_samp(dat_prev, args.evaluate_n, w_prev)

        for i in range(int(args.evaluate_n / config.ns)): 

            x_i = x_i_[i*config.ns:(i+1)*config.ns,]

            for _ in range(num_steps): 
                z = torch.randn(x_i.shape[0], x_i.shape[1]) * config.train_sd
                z = z.to(device)
                x_i = model._step(x_i, dt = config.train_dt, z = z)

            x_s.append(x_i.detach())

        x_s = torch.cat(x_s)

        loss_xy = [name] + [ot_solver(x_s, x[t_].to(device)).item() for t_ in range(len(x))]
        losses_xy.append(loss_xy)
 
    losses_xy = pd.DataFrame(losses_xy, columns = ['epoch'] + y)
    losses_xy.to_csv(log_path, sep = '\t', index = False)
    print('Wrote results to', log_path) 

def evaluate_interpolate_data(args, config): 
    
    x, y, w = load_data(config)

    device, kwargs = train.init(args) 

    pt = torch.load(args.evaluate_data) 
    x_i = torch.from_numpy(pt['sim_xp']).float().to(device)
    y_j = x[1].to(device)

    ot_solver = SamplesLoss("sinkhorn", p = 2, blur = config.sinkhorn_blur, 
        scaling = config.sinkhorn_scaling)
    loss_xy = ot_solver(x_i, y_j) 

    import pdb; pdb.set_trace()

def main(): 
   
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--seed', type = int, default = 0)
    parser.add_argument('--no-cuda', action = 'store_true')
    parser.add_argument('--device', default = 7, type = int)
    parser.add_argument('--out_dir', default = './experiments')
    # -- data options
    parser.add_argument('--data_path')
    parser.add_argument('--data_dir')
    parser.add_argument('--weight_path', default = None)
    # -- model options
    parser.add_argument('--loss', default = 'euclidean')
    parser.add_argument('--k_dim', default = 500, type = int)
    parser.add_argument('--activation', default = 'softplus')
    parser.add_argument('--layers', default = 1, type = int)
    # -- pretrain options
    parser.add_argument('--pretrain_lr', default = 1e-9, type = float)
    parser.add_argument('--pretrain_epochs', default = 500, type = int) 
    # -- train options 
    parser.add_argument('--train_epochs', default = 5000, type = int)
    parser.add_argument('--train_lr', default = 0.01, type = float)
    parser.add_argument('--train_dt', default = 0.1, type = float)
    parser.add_argument('--train_sd', default = 0.5, type = float)
    parser.add_argument('--train_tau', default = 0, type = float)
    parser.add_argument('--train_batch', default = 0.1, type = float)
    parser.add_argument('--train_clip', default = 0.25, type = float)
    parser.add_argument('--save', default = 100, type = int)
    # -- test options
    parser.add_argument('--evaluate_n', default = 10000, type = int)
    parser.add_argument('--evaluate_data')
    parser.add_argument('--evaluate-baseline', action = 'store_true')
    # -- run options
    parser.add_argument('--task', default = 'fate')
    parser.add_argument('--train', action = 'store_true')
    parser.add_argument('--evaluate')
    parser.add_argument('--config')
    args = parser.parse_args() 

    if args.task == 'fate': 

        if args.train: 

            args.pretrain = True
            args.train = True

            train.run(args, train_fate)

        if args.evaluate == 'model': 

            if args.config:
                config = SimpleNamespace(**torch.load(args.config))
                evaluate_fate(args, config) 
            else: 
                print('Please provide a config file')

    elif args.task == 'interpolate': 
        
            if args.train: 

                args.pretrain = True
                args.train = True

                config = train.run(args, train_interpolate) 

            elif args.evaluate: 

                if args.evaluate == 'model':
                    evaluate = evaluate_interpolate_model
                elif args.evaluate == 'data': 
                    evaluate =  evaluate_interpolate_data
                elif args.evaluate == 'baseline': 
                    evaluate = evaluate_interpolate_model_baseline
                else:
                    raise NotImplementedError

                if args.config:
                    config = SimpleNamespace(**torch.load(args.config))
                else:
                    print("Please provide a config file")

                evaluate(args, config)

            


if __name__ == '__main__':
    main()

# training and evaluation for interpolation and fate prediction tasks # on weinreb et al. dataset # data parameters # model parameters # pretraining parameters # training parameters # loss parameters # file parameters # data # out directory # -- load data # -- initialize # -- model # -- evaluate # expand data point # simulate forward # predict # deal with ties by setting it to the default class # may want to save yp instead # use pseudocounts for scoring # -- initialize # -- evaluate # -- initialize # -- evaluate # -- data options # -- model options # -- pretrain options # -- train options # -- test options # -- run options