"""Hydra-based training script for Neuracore models."""

import gc
import logging
import os
from pathlib import Path
from typing import Any

import hydra
import torch
import torch.multiprocessing as mp
from neuracore_types import ModelInitDescription
from neuracore_types.nc_data import DataType
from omegaconf import DictConfig, OmegaConf
from torch.utils.data import DataLoader, DistributedSampler, random_split

import neuracore as nc
from neuracore.api.training import _get_algorithms
from neuracore.core.utils.training_input_args_validation import (
    get_algorithm_name,
    validate_training_params,
)
from neuracore.ml import NeuracoreModel
from neuracore.ml.datasets.pytorch_single_sample_dataset import SingleSampleDataset
from neuracore.ml.datasets.pytorch_synchronized_dataset import (
    PytorchSynchronizedDataset,
)
from neuracore.ml.logging.cloud_training_logger import CloudTrainingLogger
from neuracore.ml.logging.tensorboard_training_logger import TensorboardTrainingLogger
from neuracore.ml.trainers.batch_autotuner import find_optimal_batch_size
from neuracore.ml.trainers.distributed_trainer import (
    DistributedTrainer,
    cleanup_distributed,
    setup_distributed,
)
from neuracore.ml.utils.algorithm_loader import AlgorithmLoader
from neuracore.ml.utils.algorithm_storage_handler import AlgorithmStorageHandler
from neuracore.ml.utils.device_utils import get_default_device
from neuracore.ml.utils.robot_data_spec_utils import (
    convert_str_to_robot_data_spec,
    extract_data_types,
    merge_robot_data_spec,
)
from neuracore.ml.utils.training_storage_handler import TrainingStorageHandler

# Environment setup
os.environ["PJRT_DEVICE"] = "GPU"

# Configure logging
logger = logging.getLogger(__name__)


def setup_logging(output_dir: str, rank: int = 0) -> None:
    """Setup logging configuration."""
    if rank == 0:
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s",
            handlers=[
                logging.StreamHandler(),
                logging.FileHandler(output_path / "train.log"),
            ],
        )
    else:
        # For other ranks, only log to console
        logging.basicConfig(
            level=logging.INFO,
            format=f"[Rank {rank}] %(asctime)s - %(levelname)s - %(message)s",
            handlers=[logging.StreamHandler()],
        )


def get_model_and_algorithm_config(
    cfg: DictConfig,
    model_init_description: ModelInitDescription,
) -> tuple[NeuracoreModel, dict[str, Any]]:
    """Get model and algorithm configuration."""
    algorithm_config: dict[str, Any] = {}
    if "algorithm" in cfg:
        algorithm_config = OmegaConf.to_container(cfg.algorithm, resolve=True)
        algorithm_config.pop("_target_", None)
        logger.info("Using custom algorithm parameters")
        logger.info(f"Algorithm parameters: {algorithm_config}")

        model = hydra.utils.instantiate(
            cfg.algorithm,
            model_init_description=model_init_description,
            **algorithm_config,
        )
    elif cfg.algorithm_id is not None:
        # Use algorithm_params for custom parameters
        if cfg.algorithm_params is not None:
            algorithm_config = OmegaConf.to_container(
                cfg.algorithm_params, resolve=True
            )
            logger.info("Using custom algorithm parameters")
            logger.info(f"Algorithm parameters: {algorithm_config}")

        extract_dir = Path(cfg.local_output_dir) / "algorithm"
        algorithm_loader = AlgorithmLoader(extract_dir)
        model_class = algorithm_loader.load_model()
        model = model_class(
            model_init_description=model_init_description,
            **algorithm_config,
        )
    else:
        raise ValueError(
            "Either 'algorithm' or 'algorithm_id' "
            "must be provided in the configuration"
        )
    return model, algorithm_config


def determine_optimal_batch_size(
    cfg: DictConfig,
    input_robot_data_spec: dict[str, dict[DataType, list[str]]],
    output_robot_data_spec: dict[str, dict[DataType, list[str]]],
    dataset: SingleSampleDataset,
    device: torch.device | None = None,
) -> int:
    """Run batch size autotuning on a single GPU and return the result."""
    if not torch.cuda.is_available() or (
        device is not None and "cuda" not in device.type
    ):
        raise ValueError("Autotuning is only supported on GPUs.")

    if device is None:
        device = get_default_device()

    logger.info(f"Starting batch size autotuning on {device}...")

    model_init_description = ModelInitDescription(
        dataset_statistics=dataset.dataset_statistics,
        input_data_types=extract_data_types(input_robot_data_spec),
        output_data_types=extract_data_types(output_robot_data_spec),
        output_prediction_horizon=cfg.output_prediction_horizon,
    )

    model, algorithm_config = get_model_and_algorithm_config(
        cfg, model_init_description
    )

    model = model.to(device)

    max_batch_size = cfg.max_batch_size if "max_batch_size" in cfg else len(dataset)
    min_batch_size = cfg.min_batch_size if "min_batch_size" in cfg else 2
    num_workers = cfg.batch_size_autotuning_num_workers

    logger.info(
        f"using max_batch_size: {max_batch_size}, "
        f"min_batch_size: {min_batch_size}, "
        f"num_workers: {num_workers}"
    )

    # Determine per-GPU batch size
    optimal_batch_size = find_optimal_batch_size(
        dataset=dataset,
        model=model,
        model_kwargs=algorithm_config,
        min_batch_size=min_batch_size,
        max_batch_size=max_batch_size,
        dataloader_kwargs={
            "collate_fn": dataset.collate_fn,
        },
    )

    # Clean up
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

    logger.info(
        f"Autotuning complete. Optimal batch size per GPU: {optimal_batch_size}"
    )

    return optimal_batch_size


def run_training(
    rank: int,
    world_size: int,
    cfg: DictConfig,
    batch_size: int,
    input_robot_data_spec: dict[str, dict[DataType, list[str]]],
    output_robot_data_spec: dict[str, dict[DataType, list[str]]],
    dataset: PytorchSynchronizedDataset,
    device: torch.device | None = None,
) -> None:
    """Run the training process for a single GPU."""
    # Setup for distributed training
    if world_size > 1:
        nc.login()  # Ensure Neuracore is logged in on this process
        setup_distributed(rank, world_size)

    # Setup logging (different file per process)
    setup_logging(cfg.local_output_dir, rank)
    logger = logging.getLogger(__name__)

    # Set random seed (different for each process to ensure different data sampling)
    torch.manual_seed(cfg.seed + rank)

    try:
        logger.info(f"Using batch size: {batch_size}")

        # Merge data_types for synchronization
        merge_robot_data_spec(input_robot_data_spec, output_robot_data_spec)

        # Split dataset
        dataset_size = len(dataset)
        train_split = 1 - cfg.validation_split
        train_size = int(train_split * dataset_size)
        val_size = dataset_size - train_size

        # Use random split with fixed seed for deterministic behavior
        generator = torch.Generator().manual_seed(cfg.seed)
        train_dataset, val_dataset = random_split(
            dataset, [train_size, val_size], generator=generator
        )

        if world_size > 1:
            train_sampler = DistributedSampler(
                train_dataset,
                num_replicas=world_size,
                rank=rank,
                shuffle=True,
                seed=cfg.seed,
            )
            val_sampler = DistributedSampler(
                val_dataset,
                num_replicas=world_size,
                rank=rank,
                shuffle=False,
                seed=cfg.seed,
            )

            train_loader = DataLoader(
                train_dataset,
                batch_size=batch_size,
                sampler=train_sampler,
                num_workers=cfg.num_train_workers,
                pin_memory=True,
                persistent_workers=cfg.num_train_workers > 0,
                collate_fn=dataset.collate_fn,
            )

            val_loader = DataLoader(
                val_dataset,
                batch_size=batch_size,
                sampler=val_sampler,
                num_workers=cfg.num_val_workers,
                pin_memory=True,
                persistent_workers=cfg.num_val_workers > 0,
                collate_fn=dataset.collate_fn,
            )
        else:
            # Regular data loaders for single GPU training
            train_loader = DataLoader(
                train_dataset,
                batch_size=batch_size,
                shuffle=True,
                num_workers=cfg.num_train_workers,
                pin_memory=True,
                persistent_workers=cfg.num_train_workers > 0,
                collate_fn=dataset.collate_fn,
            )

            val_loader = DataLoader(
                val_dataset,
                batch_size=batch_size,
                shuffle=False,
                num_workers=cfg.num_val_workers,
                pin_memory=True,
                persistent_workers=cfg.num_val_workers > 0,
                collate_fn=dataset.collate_fn,
            )

        # Log data loader information
        logger.info(
            f"Created data loaders with {len(train_loader.dataset)} training samples "
            f"and {len(val_loader.dataset)} validation samples"
        )

        # Model doesn't need to know about ids or names, just data types
        input_data_types = extract_data_types(input_robot_data_spec)
        output_data_types = extract_data_types(output_robot_data_spec)
        model_init_description = ModelInitDescription(
            dataset_statistics=dataset.dataset_statistics,
            input_data_types=input_data_types,
            output_data_types=output_data_types,
            output_prediction_horizon=cfg.output_prediction_horizon,
        )

        model, algorithm_config = get_model_and_algorithm_config(
            cfg, model_init_description
        )

        training_storage_handler = TrainingStorageHandler(
            local_dir=cfg.local_output_dir,
            training_job_id=cfg.training_id,
            algorithm_config=algorithm_config,
        )

        logger.info(
            f"Created model with "
            f"{sum(p.numel() for p in model.parameters()):,} parameters"
        )

        training_logger: TensorboardTrainingLogger | CloudTrainingLogger
        if cfg.training_id is None:
            training_logger = TensorboardTrainingLogger(
                log_dir=Path(cfg.local_output_dir) / "tensorboard",
            )
        else:
            training_logger = CloudTrainingLogger(
                training_id=cfg.training_id,
            )

        trainer = DistributedTrainer(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            training_logger=training_logger,
            storage_handler=training_storage_handler,
            output_dir=Path(cfg.local_output_dir),
            num_epochs=cfg.epochs,
            log_freq=cfg.logging_frequency,
            keep_last_n_checkpoints=cfg.keep_last_n_checkpoints,
            clip_grad_norm=algorithm_config.get("clip_grad_norm", None),
            rank=rank,
            world_size=world_size,
            device=device,
        )

        # Resume from checkpoint if specified
        start_epoch = 0
        if cfg.resume_checkpoint_path is not None:
            try:
                checkpoint = trainer.load_checkpoint(cfg.resume_checkpoint_path)
                start_epoch = checkpoint.get("epoch", 0) + 1
                logger.info(f"Resumed from checkpoint at epoch {start_epoch}")
            except Exception:
                logger.error("Failed to load checkpoint.", exc_info=True)

        # Start training
        try:
            logger.info("Starting training...")
            trainer.train(start_epoch=start_epoch)
            logger.info("Training completed successfully!")
        except Exception:
            logger.error("Training failed.", exc_info=True)
            raise

    finally:
        # Clean up distributed process group
        if world_size > 1:
            cleanup_distributed()

        logger.info(f"Process {rank} completed")


@hydra.main(version_base=None, config_path="config", config_name="config")
def main(cfg: DictConfig) -> None:
    """Main function to run the training script."""
    # Resolve the configuration
    OmegaConf.resolve(cfg)

    # Print configuration
    logger.info("Training configuration:")
    logger.info(OmegaConf.to_yaml(cfg, resolve=True))

    if "algorithm" in cfg and cfg.algorithm_id is not None:
        raise ValueError(
            "Both 'algorithm' and 'algorithm_id' are provided. "
            "Please specify only one."
        )
    if "algorithm" not in cfg and cfg.algorithm_id is None:
        raise ValueError(
            "Neither 'algorithm' nor 'algorithm_id' is provided. " "Please specify one."
        )

    if cfg.dataset_id is None and cfg.dataset_name is None:
        raise ValueError("Either 'dataset_id' or 'dataset_name' must be provided.")
    if cfg.dataset_id is not None and cfg.dataset_name is not None:
        raise ValueError(
            "Both 'dataset_id' and 'dataset_name' are provided. "
            "Please specify only one."
        )

    # Login and get dataset
    nc.login()
    if cfg.org_id is not None:
        nc.set_organization(cfg.org_id)

    if cfg.dataset_id is not None:
        dataset = nc.get_dataset(id=cfg.dataset_id)
    elif cfg.dataset_name is not None:
        dataset = nc.get_dataset(name=cfg.dataset_name)
    else:
        raise ValueError("Either 'dataset_id' or 'dataset_name' must be provided.")

    # Sort out data specs
    if cfg.input_robot_data_spec is not None:
        if not isinstance(cfg.input_robot_data_spec, DictConfig):
            raise ValueError(
                "'input_robot_data_spec' must be a dictionary "
                "mapping robot IDs to dictionary of data types to lists of data names."
            )
        input_robot_data_spec = convert_str_to_robot_data_spec(
            cfg.input_robot_data_spec
        )
    else:
        input_data_types = [DataType(dt) for dt in cfg.input_data_types]
        input_robot_data_spec = {
            robot_id: {dt: [] for dt in input_data_types}
            for robot_id in dataset.robot_ids
        }
    if cfg.output_robot_data_spec is not None:
        if not isinstance(cfg.output_robot_data_spec, DictConfig):
            raise ValueError(
                "'output_robot_data_spec' must either be None or a dictionary "
                "mapping robot IDs to dictions of data types to lists of data names."
            )
        output_robot_data_spec = convert_str_to_robot_data_spec(
            cfg.output_robot_data_spec
        )
    else:
        output_data_types = [DataType(dt) for dt in cfg.output_data_types]
        output_robot_data_spec = {
            robot_id: {dt: [] for dt in output_data_types}
            for robot_id in dataset.robot_ids
        }

    batch_size = cfg.batch_size

    algorithms_jsons = _get_algorithms()
    if cfg.algorithm_id is not None:
        algorithm_name = get_algorithm_name(
            algorithm_id=cfg.algorithm_id,
            algorithm_jsons=algorithms_jsons,
        )
    else:
        algorithm_name = cfg.algorithm._target_.rsplit(".", 1)[-1]

    validate_training_params(
        dataset,
        dataset_name=cfg.dataset_name if cfg.dataset_name is not None else "",
        algorithm_name=algorithm_name,
        input_robot_data_spec=input_robot_data_spec,
        output_robot_data_spec=output_robot_data_spec,
        algorithm_jsons=algorithms_jsons,
    )

    # Prepare data types for synchronization
    robot_data_spec = merge_robot_data_spec(
        input_robot_data_spec, output_robot_data_spec
    )

    synchronized_dataset = dataset.synchronize(
        frequency=cfg.frequency,
        robot_data_spec=robot_data_spec,
        prefetch_videos=True,
        max_prefetch_workers=cfg.max_prefetch_workers,
    )

    # Setup logging for main process
    setup_logging(cfg.local_output_dir)

    # Check if distributed training is enabled and multiple GPUs are available
    world_size = torch.cuda.device_count()

    if cfg.algorithm_id is not None:
        # Download the algorithm so that it can be processed later
        logger.info(f"Downloading algorithm from cloud with ID: {cfg.algorithm_id}")
        storage_handler = AlgorithmStorageHandler(algorithm_id=cfg.algorithm_id)
        extract_dir = Path(cfg.local_output_dir) / "algorithm"
        storage_handler.download_algorithm(extract_dir=extract_dir)
        logger.info(f"Algorithm extracted to {extract_dir}")

    device = None
    if cfg.device is not None:
        device = torch.device(cfg.device)
    else:
        device = get_default_device()

    # Create a pytorch synchronized dataset
    # NOTE: we are creating it here, and not in training to access the first sample
    # for batch size autotuning, if used.
    pytorch_dataset = PytorchSynchronizedDataset(
        synchronized_dataset=synchronized_dataset,
        input_robot_data_spec=input_robot_data_spec,
        output_robot_data_spec=output_robot_data_spec,
        output_prediction_horizon=cfg.output_prediction_horizon,
    )

    # Handle batch size configuration
    if isinstance(batch_size, str) and batch_size.lower() == "auto":
        sample = pytorch_dataset.load_sample(0).to(device)
        single_sample_dataset = SingleSampleDataset(
            sample=sample,
            input_robot_data_spec=input_robot_data_spec,
            output_robot_data_spec=output_robot_data_spec,
            output_prediction_horizon=cfg.output_prediction_horizon,
            dataset_statistics=pytorch_dataset.dataset_statistics,
            num_recordings=len(pytorch_dataset),
        )

        optimal_batch_size = determine_optimal_batch_size(
            cfg=cfg,
            input_robot_data_spec=input_robot_data_spec,
            output_robot_data_spec=output_robot_data_spec,
            dataset=single_sample_dataset,
            device=device,
        )
        batch_size = optimal_batch_size

    else:
        batch_size = int(batch_size)

    if world_size > 1:
        # Use multiprocessing to launch multiple processes
        mp.spawn(
            run_training,
            args=(
                world_size,
                cfg,
                batch_size,
                input_robot_data_spec,
                output_robot_data_spec,
                pytorch_dataset,
                device,
            ),
            nprocs=world_size,
            join=True,
        )
    else:
        # Single GPU or CPU training
        run_training(
            0,
            1,
            cfg,
            batch_size,
            input_robot_data_spec,
            output_robot_data_spec,
            pytorch_dataset,
            device,
        )


if __name__ == "__main__":
    main()
