import os
import sys
import numpy as np
import pandas as pd
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from .src.gnnwr.datasets import init_dataset
from .src.gnnwr.models import GTNNWR

data = pd.read_csv('E:/CODE/gnnwr/data/demo_data_gtnnwr.csv')
data["id"] = np.arange(len(data))
train_dataset, val_dataset, test_dataset = init_dataset(data=data,
                                                        test_ratio=0.15,
                                                        valid_ratio=0.1,
                                                        x_columns=['refl_b01', 'refl_b02',\
                                                        'refl_b03','refl_b04','refl_b05',\
                                                        'refl_b07'],
                                                        y_column=['SiO3'],
                                                        spatial_columns=['proj_x', 'proj_y'],
                                                        temporal_column=['day'],
                                                        id_column=['id'],
                                                        model_type="gtnnwr",
                                                        sample_seed=48,
                                                        batch_size=1024)
optimizer_params = {
    "scheduler":"MultiStepLR",
    "scheduler_milestones":[1000, 2000, 3000, 4000],
    "scheduler_gamma":0.8,
}
gtnnwr = GTNNWR(train_dataset, val_dataset, test_dataset, [[3], [512,256,64]],drop_out=0.4,optimizer='Adadelta',optimizer_params=optimizer_params,
                write_path = "../demo_result/gtnnwr_runs",
                model_name="GTNNWR_DSi")
gtnnwr.add_graph()

gtnnwr.run(1500,1000)

gtnnwr.result()