from ._keys import PyTorchLogKeys, PyTorchCheckpointKeys, MLTaskKeys


class CheckpointKeys:
    """Checkpoint callback keys."""
    TRAIN_LOSS = PyTorchLogKeys.TRAIN_LOSS
    VALIDATION_LOSS = PyTorchLogKeys.VAL_LOSS
    

class FinalizedFileKeys:
    """Keys for finalized model files."""
    MODEL_WEIGHTS = PyTorchCheckpointKeys.MODEL_STATE
    EPOCH = PyTorchCheckpointKeys.EPOCH
    TASK = PyTorchCheckpointKeys.TASK
    CLASSIFICATION_THRESHOLD = PyTorchCheckpointKeys.CLASSIFICATION_THRESHOLD
    CLASS_MAP = PyTorchCheckpointKeys.CLASS_MAP
    SEQUENCE_LENGTH = PyTorchCheckpointKeys.SEQUENCE_LENGTH
    INITIAL_SEQUENCE = PyTorchCheckpointKeys.INITIAL_SEQUENCE
    TARGET_NAME = PyTorchCheckpointKeys.TARGET_NAME
    TARGET_NAMES = PyTorchCheckpointKeys.TARGET_NAMES
    

class TaskKeys:
    """
    Task keys used in the Dragon ML pipeline:
    
    1. REGRESSION
    2. MULTITARGET_REGRESSION
    3. BINARY_CLASSIFICATION
    4. MULTICLASS_CLASSIFICATION
    5. MULTILABEL_BINARY_CLASSIFICATION
    6. BINARY_IMAGE_CLASSIFICATION
    7. MULTICLASS_IMAGE_CLASSIFICATION  
    8. BINARY_SEGMENTATION
    9. MULTICLASS_SEGMENTATION
    10. OBJECT_DETECTION
    11. SEQUENCE_SEQUENCE
    12. SEQUENCE_VALUE
    """
    REGRESSION = MLTaskKeys.REGRESSION
    MULTITARGET_REGRESSION = MLTaskKeys.MULTITARGET_REGRESSION
    BINARY_CLASSIFICATION = MLTaskKeys.BINARY_CLASSIFICATION
    MULTICLASS_CLASSIFICATION = MLTaskKeys.MULTICLASS_CLASSIFICATION
    MULTILABEL_BINARY_CLASSIFICATION = MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION
    BINARY_IMAGE_CLASSIFICATION = MLTaskKeys.BINARY_IMAGE_CLASSIFICATION
    MULTICLASS_IMAGE_CLASSIFICATION = MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION
    BINARY_SEGMENTATION = MLTaskKeys.BINARY_SEGMENTATION
    MULTICLASS_SEGMENTATION = MLTaskKeys.MULTICLASS_SEGMENTATION
    OBJECT_DETECTION = MLTaskKeys.OBJECT_DETECTION
    SEQUENCE_SEQUENCE = MLTaskKeys.SEQUENCE_SEQUENCE
    SEQUENCE_VALUE = MLTaskKeys.SEQUENCE_VALUE
