import argparse
import sys
import os

import socket

from accelerate import Accelerator

from litespecformer.config import LiteSpecFormerConfig, LiteSpecFormerTrainingConfig
from litespecformer.trainer import LiteSpecFormerPreTrainer

import wandb

os.environ["WANDB_API_KEY"] = (
    "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"
)
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

parser = argparse.ArgumentParser("LiteSpecFormer Pre-training")

parser.add_argument("--team_name", type=str)
parser.add_argument("--project_name", type=str)
parser.add_argument("--experiment_name", type=str)
parser.add_argument("--scenario_name", type=str)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--resume_from_checkpoint", type=str, default=None)

args = parser.parse_args()


def obj_to_dict(obj):
    """
    通用对象转字典：读取__dict__，排除以__开头的内置属性
    """
    return {
        key: value
        for key, value in obj.__dict__.items()
        if not key.startswith("__")  # 过滤__class__、__module__等内置属性
    }


# Create the configs for the experiment
model_config = LiteSpecFormerConfig()
forecasting_config = model_config.forecasting_config
training_config = LiteSpecFormerTrainingConfig()

# Log the configs to wandb
config_dict = {
    **obj_to_dict(model_config),
    **obj_to_dict(forecasting_config),
    **obj_to_dict(training_config),
}

# Create the accelerator for distributed training
accelerator = Accelerator(
    device_placement=True,
    split_batches=True,
    mixed_precision="fp16",
    gradient_accumulation_steps=training_config.gradient_accumulation_steps,
)

if accelerator.is_main_process:
    # Initialize wandb for experiment tracking
    wandb.init(
        project=args.project_name,
        entity=args.team_name,
        notes=socket.gethostname(),
        name=args.experiment_name + "_" + str(args.seed),
        group=args.scenario_name,
        dir=str("./results"),
        job_type="training",
    )

    wandb.config.update(config_dict)

# Create the trainer for LiteSpecFormer
trainer = LiteSpecFormerPreTrainer(
    accelerator=accelerator,
    model_config=model_config,
    training_config=training_config,
    resume_from_checkpoint=args.resume_from_checkpoint,
    model_params_only=True,
)

trainer.run()
