import os
import argparse
from pathlib import Path
import glob
import sys

import pandas as pd
from realtabformer import REaLTabFormer
import logging
import wandb
from relsyndgb.metadata import Metadata
from relsyndgb.data import load_tables, save_tables, remove_sdv_columns

MODEL_NAME = "REALTABFORMER"
BATCH_SIZE_PARENT = 1024
BATCH_SIZE_CHILD = 1

args = argparse.ArgumentParser()
args.add_argument("--dataset-name", type=str, default="rossmann_subsampled")
args.add_argument("--real_data_path", type=str, default="/d/hpc/projects/FRI/mj5835/relsyndgb/data/downloads")
args.add_argument("--synthetic_data_path", type=str, default="/d/hpc/projects/FRI/mj5835/relsyndgb/data/synthetic/")
args.add_argument("--full_sensitivity", type=bool, default=True)
args.add_argument("--retrain", type=bool, default=True)
args.add_argument("--run_id", type=str, default="1")
args = args.parse_args()
dataset_name = args.dataset_name
full_sensitivity = args.full_sensitivity
retrain = args.retrain
real_data_path = args.real_data_path
synthetic_data_path = args.synthetic_data_path
run_id = args.run_id

logger = logging.getLogger(f'{MODEL_NAME}_logger')

# logger.basicConfig(
#     level=logging.DEBUG,  # Set the desired log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
#     format='%(asctime)s - %(levelname)s - %(message)s',
#     # filename=f"{MODEL_NAME}_{dataset_name}.log",  # Specify the file where logs should be stored
#     stream=sys.stdout,
# )

logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler = logging.StreamHandler(stream=sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)

logger.info(f"START LOGGING...")

os.environ["WANDB_PROJECT"]=f"{MODEL_NAME}_{dataset_name}"

# save your trained model checkpoint to wandb
os.environ["WANDB_LOG_MODEL"]="true"

# turn off watch to log faster
os.environ["WANDB_WATCH"]="false"

# os.environ["WANDB_DISABLED"] = "true"

os.environ["PYTORCH_CUDA_ALLOC_CONF"]= "expandable_segments:True"


logger.info(f"START {MODEL_NAME}...")


logger.debug("Loading real data...")
metadata = Metadata().load_from_json(Path(real_data_path) / f'{dataset_name}/metadata.json')
real_data = load_tables(Path(real_data_path) / f'{dataset_name}', metadata)
real_data, metadata = remove_sdv_columns(real_data, metadata)
metadata.validate_data(real_data)
logger.debug("Real data loaded")

synthetic_data = {}

# GENERATE SYNTHETIC DATA ---------------------------------

root_table_name = metadata.get_root_tables()[0]
parent_df = real_data[root_table_name]

join_on = metadata.get_primary_key(root_table_name)
parent_model_path = f'models/{MODEL_NAME}/{dataset_name}/checkpoint_{root_table_name}' 

if retrain or not(os.path.exists(parent_model_path) and len(os.listdir(parent_model_path)) > 0):
    logger.debug("Starting parent table generation...")
    logger.debug("Fitting parent model...")
    parent_model = REaLTabFormer(model_type="tabular", batch_size=BATCH_SIZE_PARENT, report_to="wandb", epochs=5)
    os.makedirs(parent_model_path, exist_ok=True)
    # fit and save parent model
    parent_model.fit(parent_df.drop(join_on, axis=1), full_sensitivity=full_sensitivity)
    parent_model.save(parent_model_path)
    logger.debug("Parent fitted and saved")
else:
    logger.debug("Skipping parent table generation and reading last trained model...")

# load trained parent model
directories = list(filter(os.path.isdir, glob.glob(f"{parent_model_path}/id*")))
directories.sort(key=lambda x: os.path.getmtime(x))
parent_model_path = directories[-1]
parent_model = REaLTabFormer(
    model_type="tabular",
    parent_realtabformer_path=parent_model_path,
)

parent_model = parent_model.load_from_dir(parent_model_path)
# sample from parent model
logger.debug("Sampling from parent model...")
parent_samples = parent_model.sample(real_data[root_table_name].shape[0])
parent_samples.index.name = join_on
parent_samples = parent_samples.reset_index()
synthetic_data[root_table_name] = parent_samples
logger.debug("Done.")

#GENERATE CHILD TABLES
for child_table_name in metadata.get_children(root_table_name):
    logger.info(f"Starting child table generation for {child_table_name}...")
    child_model_path = f'models/{MODEL_NAME}/{dataset_name}/checkpoint_{child_table_name}'
    child_df = real_data[child_table_name]
    assert ((join_on in parent_df.columns) and
            (join_on in child_df.columns))
    if retrain or not(os.path.exists(child_model_path) and len(os.listdir(child_model_path)) > 0):
        child_model = REaLTabFormer(
            model_type="relational",
            parent_realtabformer_path=parent_model_path,
            output_max_length=None,
            train_size=0.95,
            batch_size=BATCH_SIZE_CHILD,
            report_to="wandb",
            epochs=5)
        logger.debug(f"Fitting child model for table {child_table_name}...")
        # fit child model
        child_model.fit(
            df = child_df,
            in_df=parent_df,
            join_on=join_on)
        # save child model
        os.makedirs(child_model_path, exist_ok=True)
        child_model.save(child_model_path)
        logger.debug("Child fitted and saved")

    directories = list(filter(os.path.isdir, glob.glob(f"{child_model_path}/id*")))
    directories.sort(key=lambda x: os.path.getmtime(x))
    child_model_path = directories[-1]
    child_model = REaLTabFormer(
        model_type="relational",
            parent_realtabformer_path=parent_model_path,
            output_max_length=None,
            train_size=0.95,
            batch_size=BATCH_SIZE_CHILD,
            report_to="wandb"
    )
    child_model = child_model.load_from_dir(child_model_path)
    
    logger.debug("Sampling from child model...")
    child_samples = child_model.sample(
        input_unique_ids=parent_samples[join_on],
        input_df=parent_samples.drop(join_on, axis=1),
        gen_batch=64)
    
    child_samples.index.name = join_on
    child_samples = child_samples.reset_index()
    
    synthetic_data[child_table_name] = child_samples
    logger.debug("Done.")


# SAVE DATA ---------------------------------
logger.info("Saving synthetic data...")
save_data_path = Path(synthetic_data_path) / dataset_name / MODEL_NAME / run_id
save_tables(synthetic_data, save_data_path)
logger.info("COMPLETE GENERATION DONE.")
