import os
import sys

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(PROJECT_ROOT)

import yaml
import argparse
from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger
from trainer import XrayReg
import logging
import wandb

from lightning.pytorch.callbacks import LearningRateMonitor


def parse_args():
    parser = argparse.ArgumentParser(description="Train Xray Model")
    parser.add_argument(
        "--config",
        type=str,
        default="configs/config.yaml",
        help="Path to the config file",
    )
    return parser.parse_args()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    try:
        args = parse_args()

        with open(args.config, "r") as ymlfile:
            config = yaml.safe_load(ymlfile)

        wandb.init(
            project=config.get("wandb_project", "xray_regression_noaug"))

        wandb_logger = WandbLogger(
            project=config.get("wandb_project", "xray_regression_noaug"))

        lr_monitor = LearningRateMonitor(logging_interval="step")

        trainer = Trainer(
            max_epochs=config["training"]["max_epochs"],
            log_every_n_steps=config["logging"]["log_every_n_steps"],
            logger=wandb_logger,
            callbacks=[lr_monitor])

        model = XrayReg(config)

        logger.info("Starting training...")
        trainer.fit(model)

        logger.info("Training completed. Starting testing...")
        trainer.test(model)

        logger.info("Testing completed. Logging test results...")
        model.save_test_results_to_wandb()
        logger.info("Test results saved to Wandb")

        wandb.finish()
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        sys.exit(1)