syntax = "proto3";

package ml_training;

// ML Training Service provides comprehensive machine learning model training capabilities for HFT systems.
// This service manages training jobs for MAMBA-2, TLOB transformers, DQN, PPO, Liquid Networks, and TFT models
// with real-time progress monitoring, resource management, and performance tracking.
service MLTrainingService {
  // Training Job Management
  // Initiates a new training job and returns job ID immediately
  rpc StartTraining(StartTrainingRequest) returns (StartTrainingResponse);
  
  // Subscribe to real-time training progress and status updates
  rpc SubscribeToTrainingStatus(SubscribeToTrainingStatusRequest) returns (stream TrainingStatusUpdate);
  
  // Stop a running training job (idempotent operation)
  rpc StopTraining(StopTrainingRequest) returns (StopTrainingResponse);

  // Model and Job Discovery
  // List available ML models with their training parameters
  rpc ListAvailableModels(ListAvailableModelsRequest) returns (ListAvailableModelsResponse);
  
  // Get paginated list of training job history
  rpc ListTrainingJobs(ListTrainingJobsRequest) returns (ListTrainingJobsResponse);
  
  // Get comprehensive details for a specific training job
  rpc GetTrainingJobDetails(GetTrainingJobDetailsRequest) returns (GetTrainingJobDetailsResponse);
  
  // Service Health and Status
  // Check service health and resource availability
  rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse);

  // Hyperparameter Tuning Management
  // Start a new hyperparameter tuning job using Optuna
  rpc StartTuningJob(StartTuningJobRequest) returns (StartTuningJobResponse);

  // Get current status and best parameters from a tuning job
  rpc GetTuningJobStatus(GetTuningJobStatusRequest) returns (GetTuningJobStatusResponse);

  // Stop a running hyperparameter tuning job
  rpc StopTuningJob(StopTuningJobRequest) returns (StopTuningJobResponse);

  // INTERNAL: Train a single model instance with specific hyperparameters (called by Optuna subprocess)
  rpc TrainModel(TrainModelRequest) returns (TrainModelResponse);

  // Stream real-time tuning progress updates (trial completion events)
  rpc StreamTuningProgress(StreamProgressRequest) returns (stream ProgressUpdate);

  // Batch Tuning Management
  // Start batch tuning job for multiple models with automatic dependency resolution
  rpc BatchStartTuningJobs(BatchStartTuningJobsRequest) returns (BatchStartTuningJobsResponse);

  // Get batch tuning job status with per-model results
  rpc GetBatchTuningStatus(GetBatchTuningStatusRequest) returns (GetBatchTuningStatusResponse);

  // Stop a running batch tuning job
  rpc StopBatchTuningJob(StopBatchTuningJobRequest) returns (StopBatchTuningJobResponse);

  // Job Completion Callback (called by training-uploader sidecar)
  rpc ReportJobCompletion(JobCompletionReport) returns (JobCompletionAck);

  // Model Promotion Management
  rpc ListPendingPromotions(ListPendingPromotionsRequest) returns (ListPendingPromotionsResponse);
  rpc ApprovePromotion(ApprovePromotionRequest) returns (ApprovePromotionResponse);
  rpc RejectPromotion(RejectPromotionRequest) returns (RejectPromotionResponse);

  // Model approval / rejection (dashboard-facing, delegates to promotion pipeline)
  rpc ApproveModel(ApproveModelRequest) returns (ApproveModelResponse);
  rpc RejectModel(RejectModelRequest) returns (RejectModelResponse);
}

// --- Core Request/Response Messages ---

// Training mode selection
enum TrainingMode {
  TRAINING_MODE_FULL = 0;        // Full training from scratch (default, backward-compatible)
  TRAINING_MODE_FINE_TUNE = 1;   // Fine-tune from existing checkpoint
}

// Request to start a new model training job
message StartTrainingRequest {
  string model_type = 1;                // Model type ("TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT")
  DataSource data_source = 2;           // Training data source configuration
  Hyperparameters hyperparameters = 3;  // Model-specific training parameters
  bool use_gpu = 4;                     // Whether to use GPU acceleration
  string description = 5;               // Optional job description
  map<string, string> tags = 6;         // Optional categorization tags
  TrainingMode mode = 7;               // FULL (default) or FINE_TUNE
  string resume_checkpoint_path = 8;   // Path to checkpoint for fine-tune
  uint32 max_epochs = 9;              // Override epoch count (0 = use default)
}

message StartTrainingResponse {
  string job_id = 1;
  TrainingStatus status = 2;
  string message = 3;
}

message SubscribeToTrainingStatusRequest {
  string job_id = 1;
}

// Real-time training progress update streamed from server
message TrainingStatusUpdate {
  string job_id = 1;                    // Training job identifier
  TrainingStatus status = 2;            // Current job status
  float progress_percentage = 3;        // Training progress (0.0 to 100.0)
  uint32 current_epoch = 4;             // Current training epoch
  uint32 total_epochs = 5;              // Total epochs planned
  map<string, float> metrics = 6;       // Training metrics (loss, accuracy, sharpe_ratio, etc.)
  string message = 7;                   // Human-readable status message
  int64 timestamp = 8;                  // Update timestamp (Unix seconds)
  FinancialMetrics financial_metrics = 9; // Financial performance metrics
  ResourceUsage resource_usage = 10;    // Current resource utilization
}

message StopTrainingRequest {
  string job_id = 1;
  string reason = 2; // Optional reason for stopping
}

message StopTrainingResponse {
  bool success = 1;
  string message = 2;
}

message ListAvailableModelsRequest {}

message ListAvailableModelsResponse {
  repeated ModelDefinition models = 1;
}

message ListTrainingJobsRequest {
  uint32 page = 1;
  uint32 page_size = 2;
  TrainingStatus status_filter = 3;
  string model_type_filter = 4;
  int64 start_time = 5; // Unix timestamp in seconds
  int64 end_time = 6; // Unix timestamp in seconds
}

message ListTrainingJobsResponse {
  repeated TrainingJobSummary jobs = 1;
  uint32 total_count = 2;
  uint32 page = 3;
  uint32 page_size = 4;
}

message GetTrainingJobDetailsRequest {
  string job_id = 1;
}

message GetTrainingJobDetailsResponse {
  TrainingJobDetails job_details = 1;
}

message HealthCheckRequest {}

message HealthCheckResponse {
  bool healthy = 1;
  string message = 2;
  map<string, string> details = 3;
}

// Request to start hyperparameter tuning job
message StartTuningJobRequest {
  string model_type = 1;                // Model type to tune ("TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT")
  uint32 num_trials = 2;                // Number of tuning trials to run
  string config_path = 3;               // Path to tuning configuration file (search space, objectives)
  DataSource data_source = 4;           // Training data source for all trials
  bool use_gpu = 5;                     // Whether to use GPU acceleration
  string description = 6;               // Optional job description
  map<string, string> tags = 7;         // Optional categorization tags
}

message StartTuningJobResponse {
  string job_id = 1;                    // Unique tuning job identifier
  TuningJobStatus status = 2;           // Initial job status
  string message = 3;                   // Human-readable status message
}

// Request to query tuning job status
message GetTuningJobStatusRequest {
  string job_id = 1;                    // Tuning job identifier
}

message GetTuningJobStatusResponse {
  string job_id = 1;                    // Tuning job identifier
  TuningJobStatus status = 2;           // Current job status
  uint32 current_trial = 3;             // Current trial number (0-indexed)
  uint32 total_trials = 4;              // Total number of trials
  map<string, float> best_params = 5;   // Best hyperparameters found so far
  map<string, float> best_metrics = 6;  // Metrics for best parameters (sharpe_ratio, training_loss, etc.)
  repeated TrialResult trial_history = 7; // Complete trial history
  string message = 8;                   // Human-readable status message
  int64 started_at = 9;                 // Job start time (Unix timestamp in seconds)
  int64 updated_at = 10;                // Last update time (Unix timestamp in seconds)
}

// Request to stop a tuning job
message StopTuningJobRequest {
  string job_id = 1;                    // Tuning job identifier
  string reason = 2;                    // Optional reason for stopping
}

message StopTuningJobResponse {
  bool success = 1;                     // Whether stop was successful
  string message = 2;                   // Human-readable status message
  TuningJobStatus final_status = 3;     // Final job status after stopping
}

// INTERNAL: Request to train a model with specific hyperparameters (called by Optuna)
message TrainModelRequest {
  string model_type = 1;                // Model type ("TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT")
  map<string, float> hyperparameters = 2; // Hyperparameters to use for this trial
  DataSource data_source = 3;           // Training data source
  bool use_gpu = 4;                     // Whether to use GPU acceleration
  string trial_id = 5;                  // Optuna trial identifier for tracking
}

message TrainModelResponse {
  bool success = 1;                     // Whether training succeeded
  float sharpe_ratio = 2;               // Primary optimization objective (Sharpe ratio)
  float training_loss = 3;              // Final training loss
  map<string, float> validation_metrics = 4; // Additional validation metrics
  string error_message = 5;             // Error message if training failed
  int64 training_duration_seconds = 6;  // Total training time
}

// Individual trial result for tuning job history
message TrialResult {
  uint32 trial_number = 1;              // Trial index
  map<string, float> params = 2;        // Hyperparameters tested
  float objective_value = 3;            // Objective metric (e.g., Sharpe ratio)
  map<string, float> metrics = 4;       // Additional metrics
  TrialState state = 5;                 // Trial outcome state
  int64 started_at = 6;                 // Trial start time (Unix timestamp in seconds)
  int64 completed_at = 7;               // Trial completion time (Unix timestamp in seconds)
}

// Request to stream tuning progress updates
message StreamProgressRequest {
  string job_id = 1;                    // Tuning job identifier to subscribe to
}

// Real-time progress update streamed after each trial completes
message ProgressUpdate {
  string job_id = 1;                    // Tuning job identifier
  uint32 current_trial = 2;             // Current trial number (0-indexed)
  uint32 total_trials = 3;              // Total number of trials
  map<string, string> trial_params = 4; // Current trial hyperparameters (as strings for display)
  float trial_sharpe = 5;               // Current trial's Sharpe ratio (objective value)
  float best_sharpe_so_far = 6;         // Best Sharpe ratio achieved so far
  uint32 estimated_time_remaining = 7;  // Estimated seconds until completion
  TuningJobStatus status = 8;           // Current job status
  string message = 9;                   // Human-readable status message
  int64 timestamp = 10;                 // Update timestamp (Unix seconds)
  UpdateType update_type = 11;          // Type of update (trial completion, heartbeat, job complete)
}

// Type of progress update
enum UpdateType {
  UPDATE_UNKNOWN = 0;                   // Unknown/unspecified
  UPDATE_TRIAL_COMPLETE = 1;            // Trial completed
  UPDATE_HEARTBEAT = 2;                 // Keepalive heartbeat (no trial change)
  UPDATE_JOB_COMPLETE = 3;              // Job completed/stopped/failed
}

// --- Enums ---

// Current status of a training job
enum TrainingStatus {
  UNKNOWN = 0;                          // Default/unknown status
  PENDING = 1;                          // Job queued, waiting to start
  RUNNING = 2;                          // Job currently executing
  COMPLETED = 3;                        // Job finished successfully
  FAILED = 4;                           // Job failed with error
  STOPPED = 5;                          // Job manually stopped
  PAUSED = 6;                           // Job temporarily paused
}

// Status of a hyperparameter tuning job
enum TuningJobStatus {
  TUNING_UNKNOWN = 0;                   // Default/unknown status
  TUNING_PENDING = 1;                   // Job queued, waiting to start
  TUNING_RUNNING = 2;                   // Job currently executing trials
  TUNING_COMPLETED = 3;                 // Job finished all trials successfully
  TUNING_FAILED = 4;                    // Job failed with error
  TUNING_STOPPED = 5;                   // Job manually stopped before completion
}

// Outcome state of an individual trial
enum TrialState {
  TRIAL_UNKNOWN = 0;                    // Default/unknown state
  TRIAL_RUNNING = 1;                    // Trial currently executing
  TRIAL_COMPLETE = 2;                   // Trial completed successfully
  TRIAL_PRUNED = 3;                     // Trial pruned by Optuna (early stopping)
  TRIAL_FAILED = 4;                     // Trial failed with error
}

// --- Data Structures ---

message DataSource {
  oneof source {
    string historical_db_query = 1;
    string real_time_stream_topic = 2;
    string file_path = 3;
  }
  int64 start_time = 4; // Unix timestamp in seconds
  int64 end_time = 5; // Unix timestamp in seconds
}

// Provides type-safe hyperparameter configuration.
message Hyperparameters {
  oneof model_params {
    TlobParams tlob_params = 1;
    MambaParams mamba_params = 2;
    DqnParams dqn_params = 3;
    PpoParams ppo_params = 4;
    LiquidParams liquid_params = 5;
    TftParams tft_params = 6;
  }
}

// TLOB (Time-Limit Order Book) Transformer parameters
message TlobParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  uint32 sequence_length = 4;
  uint32 hidden_dim = 5;
  uint32 num_heads = 6;
  uint32 num_layers = 7;
  float dropout_rate = 8;
  bool use_positional_encoding = 9;
}

// MAMBA-2 State Space Model parameters
message MambaParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  uint32 state_dim = 4;
  uint32 hidden_dim = 5;
  uint32 num_layers = 6;
  float dt_min = 7;
  float dt_max = 8;
  bool use_cuda_kernels = 9;
}

// DQN (Deep Q-Network) parameters
message DqnParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  uint32 replay_buffer_size = 4;
  float epsilon_start = 5;
  float epsilon_end = 6;
  uint32 epsilon_decay_steps = 7;
  float gamma = 8;
  uint32 target_update_frequency = 9;
  bool use_double_dqn = 10;
  bool use_dueling = 11;
  bool use_prioritized_replay = 12;
}

// PPO (Proximal Policy Optimization) parameters
message PpoParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  float clip_ratio = 4;
  float value_loss_coef = 5;
  float entropy_coef = 6;
  uint32 rollout_steps = 7;
  uint32 minibatch_size = 8;
  float gae_lambda = 9;
}

// Liquid Network parameters
message LiquidParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  uint32 num_neurons = 4;
  float tau = 5;
  float sigma = 6;
  bool use_adaptive_tau = 7;
}

// Temporal Fusion Transformer parameters
message TftParams {
  uint32 epochs = 1;
  float learning_rate = 2;
  uint32 batch_size = 3;
  uint32 hidden_dim = 4;
  uint32 num_heads = 5;
  uint32 num_layers = 6;
  uint32 lookback_window = 7;
  uint32 forecast_horizon = 8;
  float dropout_rate = 9;
}

message ModelDefinition {
  string model_type = 1;
  string description = 2;
  Hyperparameters default_hyperparameters = 3;
  repeated string required_features = 4;
  uint32 estimated_training_time_minutes = 5;
  bool requires_gpu = 6;
}

message TrainingJobSummary {
  string job_id = 1;
  string model_type = 2;
  TrainingStatus status = 3;
  int64 created_at = 4; // Unix timestamp in seconds
  int64 started_at = 5; // Unix timestamp in seconds
  int64 completed_at = 6; // Unix timestamp in seconds
  string description = 7;
  float final_loss = 8;
  float best_validation_score = 9;
  map<string, string> tags = 10;
}

message TrainingJobDetails {
  string job_id = 1;
  string model_type = 2;
  TrainingStatus status = 3;
  int64 created_at = 4; // Unix timestamp in seconds
  int64 started_at = 5; // Unix timestamp in seconds
  int64 completed_at = 6; // Unix timestamp in seconds
  string description = 7;
  Hyperparameters hyperparameters = 8;
  DataSource data_source = 9;
  repeated TrainingStatusUpdate status_history = 10;
  FinancialMetrics final_financial_metrics = 11;
  string model_artifact_path = 12;
  map<string, string> tags = 13;
  string error_message = 14;
}

message FinancialMetrics {
  float simulated_return = 1;
  float sharpe_ratio = 2;
  float max_drawdown = 3;
  float hit_rate = 4;
  float avg_prediction_error_bps = 5;
  float risk_adjusted_return = 6;
  float var_5pct = 7;
  float expected_shortfall = 8;
}

message ResourceUsage {
  float cpu_usage_percent = 1;
  float memory_usage_gb = 2;
  float gpu_usage_percent = 3;
  float gpu_memory_usage_gb = 4;
  uint32 active_workers = 5;
}

// --- Batch Tuning Messages ---

// Request to start batch tuning for multiple models
message BatchStartTuningJobsRequest {
  repeated string model_types = 1;         // List of models to tune (DQN, PPO, MAMBA_2, TFT, etc.)
  uint32 trials_per_model = 2;             // Number of trials for each model
  string config_path = 3;                  // Path to tuning configuration file
  DataSource data_source = 4;              // Training data source for all models
  bool use_gpu = 5;                        // Whether to use GPU acceleration
  bool auto_export_yaml = 6;               // Automatically export best params to YAML (default: true)
  string yaml_export_path = 7;             // Custom YAML export path (default: ml/config/best_hyperparameters.yaml)
  string description = 8;                  // Optional batch job description
  map<string, string> tags = 9;            // Optional categorization tags
}

message BatchStartTuningJobsResponse {
  string batch_id = 1;                     // Unique batch job identifier
  repeated string execution_order = 2;     // Model execution order (after dependency resolution)
  string message = 3;                      // Human-readable status message
  BatchTuningStatus status = 4;            // Initial batch status
}

// Request to get batch tuning job status
message GetBatchTuningStatusRequest {
  string batch_id = 1;                     // Batch job identifier
}

message GetBatchTuningStatusResponse {
  string batch_id = 1;                     // Batch job identifier
  BatchTuningStatus status = 2;            // Current batch status
  uint32 current_model_index = 3;          // Index of currently executing model (0-based)
  uint32 total_models = 4;                 // Total number of models in batch
  repeated ModelTuningResult results = 5;  // Results for completed models
  string current_model = 6;                // Currently tuning model type
  int64 started_at = 7;                    // Batch start time (Unix timestamp)
  int64 updated_at = 8;                    // Last update time (Unix timestamp)
  int64 estimated_completion_time = 9;     // Estimated completion time (Unix timestamp)
  string yaml_export_path = 10;            // Path where YAML will be exported
}

// Individual model tuning result within batch
message ModelTuningResult {
  string model_type = 1;                   // Model type (DQN, PPO, etc.)
  string job_id = 2;                       // Individual tuning job ID
  TuningJobStatus status = 3;              // Model tuning status
  map<string, float> best_params = 4;      // Best hyperparameters found
  map<string, float> best_metrics = 5;     // Best metrics achieved
  uint32 trials_completed = 6;             // Number of trials completed
  int64 started_at = 7;                    // Model tuning start time
  int64 completed_at = 8;                  // Model tuning completion time
  string error_message = 9;                // Error message if failed
}

// Request to stop batch tuning job
message StopBatchTuningJobRequest {
  string batch_id = 1;                     // Batch job identifier
  string reason = 2;                       // Optional reason for stopping
}

message StopBatchTuningJobResponse {
  bool success = 1;                        // Whether stop was successful
  string message = 2;                      // Human-readable status message
  BatchTuningStatus final_status = 3;      // Final batch status
  repeated ModelTuningResult completed_results = 4; // Results for completed models
}

// Batch tuning job status
enum BatchTuningStatus {
  BATCH_UNKNOWN = 0;                       // Default/unknown status
  BATCH_PENDING = 1;                       // Batch queued, waiting to start
  BATCH_RUNNING = 2;                       // Batch currently executing models
  BATCH_COMPLETED = 3;                     // All models completed successfully
  BATCH_PARTIALLY_COMPLETED = 4;           // Some models succeeded, some failed
  BATCH_FAILED = 5;                        // Batch failed (all models failed or critical error)
  BATCH_STOPPED = 6;                       // Batch manually stopped
}

// --- Job Completion Callback (from training-uploader sidecar) ---

message JobCompletionReport {
  string job_id = 1;
  string s3_path = 2;
  bool success = 3;
  string error_message = 4;
  map<string, double> metrics = 5;
}

message JobCompletionAck {
  bool accepted = 1;
  string promotion_status = 2;
}

// --- Model Promotion ---

message ListPendingPromotionsRequest {}

message ListPendingPromotionsResponse {
  repeated PendingPromotion promotions = 1;
}

message PendingPromotion {
  string model_id = 1;
  string model_type = 2;
  string symbol = 3;
  string s3_path = 4;
  map<string, double> new_metrics = 5;
  map<string, double> current_metrics = 6;
  int64 trained_at = 7;
  string job_id = 8;
}

message ApprovePromotionRequest {
  string model_id = 1;
}

message ApprovePromotionResponse {
  bool success = 1;
  string message = 2;
}

message RejectPromotionRequest {
  string model_id = 1;
  string reason = 2;
}

message RejectPromotionResponse {
  bool success = 1;
  string message = 2;
}

// --- Model Approval / Rejection (dashboard-facing) ---

message ApproveModelRequest {
  string model_id = 1;
  string promoted_to = 2;  // e.g. "production", "staging", "canary"
}

message ApproveModelResponse {
  bool success = 1;
  string message = 2;
}

message RejectModelRequest {
  string model_id = 1;
  string reason = 2;
}

message RejectModelResponse {
  bool success = 1;
  string message = 2;
}