import matplotlib
# Force a headless-safe backend everywhere (must be BEFORE pyplot import)
matplotlib.use("Agg", force=True)

from octopi.datasets import generators, multi_config_generator
from monai.losses import DiceLoss, FocalLoss, TverskyLoss
from octopi.models import common as builder
from typing import List, Optional, Tuple
from octopi.entry_points import common 
from octopi.utils import parsers, io
from octopi.workflows import train
import torch, os, argparse

def train_model(
    copick_config_path: str,
    target_info: Tuple[str, str, str],
    tomo_algorithm: str = 'wbp',
    voxel_size: float = 10,
    trainRunIDs: List[str] = None,
    validateRunIDs: List[str] = None,    
    model_config: str = None,
    model_weights: Optional[str] = None,
    model_save_path: str = 'results',
    num_tomo_crops: int = 16,
    tomo_batch_size: int = 15,
    lr: float = 1e-3,
    tversky_alpha: float = 0.5,
    num_epochs: int = 100,  
    val_interval: int = 5,
    best_metric: str = 'avg_f1',
    data_split: str = '0.8'
    ):
    """
    Train a 3D U-Net model using the specified CoPick configuration and target information.
    """

    # Initialize the data generator to manage training and validation datasets
    print(f'Training with {copick_config_path}\n')

    # Multi-config training
    if isinstance(copick_config_path, dict):
        data_generator = multi_config_generator.MultiConfigTrainLoaderManager(
            copick_config_path, 
            target_info[0], 
            target_session_id = target_info[2],
            target_user_id = target_info[1],
            tomo_algorithm = tomo_algorithm,
            voxel_size = voxel_size,
            tomo_batch_size = tomo_batch_size )
    else:  # Single-config training
        data_generator = generators.TrainLoaderManager(
            copick_config_path, 
            target_info[0], 
            target_session_id = target_info[2],
            target_user_id = target_info[1],
            tomo_algorithm = tomo_algorithm,
            voxel_size = voxel_size,
            tomo_batch_size = tomo_batch_size )

    # Get the data splits
    ratios = parsers.parse_data_split(data_split)
    data_generator.get_data_splits(
        trainRunIDs = trainRunIDs,
        validateRunIDs = validateRunIDs,
        train_ratio = ratios[0], val_ratio = ratios[1], test_ratio = ratios[2],
        create_test_dataset = False)
    
    # Get the reload frequency
    data_generator.get_reload_frequency(num_epochs)
    model_config['num_classes'] = data_generator.Nclasses

    # Monai Functions
    alpha = tversky_alpha
    beta = 1 - alpha
    loss_function = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True, alpha=alpha, beta=beta)  
    
    # Train the Model
    train(
        data_generator, loss_function, 
        model_config = model_config, model_weights = model_weights,
        best_metric = best_metric, num_epochs = num_epochs,
        model_save_path = model_save_path, lr0 = lr
    )

def train_model_parser(parser_description, add_slurm: bool = False):
    """
    Parse the arguments for the training model
    """
    parser = argparse.ArgumentParser(
        description=parser_description,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    # Input Arguments
    input_group = parser.add_argument_group("Input Arguments")
    common.add_config(input_group, single_config=False)
    input_group.add_argument("--target-info", type=parsers.parse_target, default="targets,octopi,1", 
                             help="Target information, e.g., 'name' or 'name,user_id,session_id'. Default is 'targets,octopi,1'.")
    input_group.add_argument("--tomo-alg", default='wbp', help="Tomogram algorithm used for training")
    input_group.add_argument("--trainRunIDs", type=parsers.parse_list, help="List of training run IDs, e.g., run1,run2,run3")
    input_group.add_argument("--validateRunIDs", type=parsers.parse_list, help="List of validation run IDs, e.g., run4,run5,run6")
    input_group.add_argument('--data-split', type=str, default='0.8', help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) "
                                "or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
    
    fine_tune_group = parser.add_argument_group("Fine-Tuning Arguments")
    fine_tune_group.add_argument('--model-config', type=str, help="Path to the model configuration file (typically used for fine-tuning)")
    fine_tune_group.add_argument('--model-weights', type=str, help="Path to the model weights file (typically used for fine-tuning)") 

    # Model Arguments
    model_group = parser.add_argument_group("UNet-Model Arguments")
    common.add_model_parameters(model_group)   
    
    # Training Arguments
    train_group = parser.add_argument_group("Training Arguments")
    common.add_train_parameters(train_group)
    
    # SLURM Arguments
    if add_slurm:
        slurm_group = parser.add_argument_group("SLURM Arguments")
        common.add_slurm_parameters(slurm_group, 'train', gpus = 1)

    args = parser.parse_args()
    return args

# Entry point with argparse
def cli():
    """
    CLI entry point for training models where results can either be saved to a local directory or a server with MLFlow.
    """

    # Parse the arguments
    parser_description = "Train 3D CNN U-Net models"
    args = train_model_parser(parser_description)

    # Parse the CoPick configuration paths
    if len(args.config) > 1:    copick_configs = parsers.parse_copick_configs(args.config)
    else:                       copick_configs = args.config[0]

    if args.model_config:
        model_config = io.load_yaml(args.model_config)
    else:
        model_config = get_model_config(args.channels, args.strides, args.res_units, args.dim_in)

    # Call the training function
    train_model(
        copick_config_path=copick_configs, 
        target_info=args.target_info,
        tomo_algorithm=args.tomo_alg,
        voxel_size=args.voxel_size,
        model_config=model_config,
        model_weights=args.model_weights,
        model_save_path=args.model_save_path,
        num_tomo_crops=args.num_tomo_crops,
        tomo_batch_size=args.tomo_batch_size,
        lr=args.lr,
        tversky_alpha=args.tversky_alpha,
        num_epochs=args.num_epochs,
        val_interval=args.val_interval,
        best_metric=args.best_metric,
        trainRunIDs=args.trainRunIDs,
        validateRunIDs=args.validateRunIDs,
        data_split=args.data_split
    )

def get_model_config(channels, strides, res_units, dim_in):
    """
        Create a model configuration dictionary if no model configuration file is provided.
    """
    model_config = {
        'architecture': 'Unet',
        'channels': channels,
        'strides': strides,
        'num_res_units': res_units, 
        'dropout': 0.1,
        'dim_in': dim_in
    }
    return model_config
    
