import functools

from jax.experimental import optimizers

import jax

import jax.config

from jax.config import config as jax_config

jax_config.update('jax_enable_x64', True) # for numerical stability, can disable if not an issue

from jax import numpy as jnp

from jax import scipy as sp

import numpy as np

import neural_tangents as nt

from neural_tangents import stax

from jax import random

from utils import *

import argparse

import design_bench

import copy

import time

parser = argparse.ArgumentParser(description="bi-level sequence learning")

parser.add_argument('--mode', choices=['distill', 'grad', 'both'], type=str, default='both')

parser.add_argument('--task', choices=['TFBind8-Exact-v0', 'Superconductor-RandomForest-v0',

'GFP-Transformer-v0', 'UTR-ResNet-v0', 'HopperController-Exact-v0',

'AntMorphology-Exact-v0', 'DKittyMorphology-Exact-v0'], type=str,

default='TFBind8-Exact-v0')

parser.add_argument('--topk', default=128, type=int)

parser.add_argument('--label', default=10.0, type=float)

parser.add_argument('--gamma', default=0.0, type=float)

parser.add_argument('--outer_lr', default=1e-1, type=float)

parser.add_argument('--Tmax', default=200, type=int)

parser.add_argument('--interval', default=200, type=int)

args = parser.parse_args()

init_fn, apply_fn, kernel_fn = stax.serial(stax.Dense(1), stax.Relu(), stax.Dense(1), stax.Relu(),

stax.Dense(1), stax.Relu(), stax.Dense(1), stax.Relu(), stax.Dense(1), stax.Relu(), stax.Dense(1))

KERNEL_FN = functools.partial(kernel_fn, get='ntk')

def distill(args):

print("After max {} median {}\n".format(max_score, median_score))

if __name__ == "__main__":

print(args)

load_d(args.task)

load_y(args.task)

distill(args)