import sys
sys.path.append("../")
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

import pydaddy
import numpy as np

from HyperSINDy import Net
from baseline import Trainer
from library_utils import Library
from Datasets import SyntheticDataset
from other import init_weights, make_folder, set_random_seed
from exp_utils import get_equations, log_equations

set_random_seed(5678)

x = np.load("../data/lotkavolterra/x_train.npy")
dt = 0.01
ddsde = pydaddy.Characterize([x[:,0], x[:,1]], dt, bins=20)

with open("../results/ss_lotkavolterra.txt", "w") as f:
    # fit
    F1 = ddsde.fit('F1', order=3, threshold=0.15)
    F2 = ddsde.fit('F2', order=2, tune=True)
    G11 = ddsde.fit('G11', order=3, threshold=0.015)
    G22 = ddsde.fit('G22', order=2, threshold=0.03)
    G12 = ddsde.fit('G12', order=3, threshold=0.025)

    # print
    print("F1", file=f)
    print(F1, file=f)
    print("F2", file=f)
    print(F2, file=f)
    print("G11", file=f)
    print(G11, file=f)
    print("G22", file=f)
    print(G22, file=f)
    print("G12", file=f)
    print(G12, file=f)