"""Callbacks system for training"""

class Callback:
    """Base callback class"""
    
    def on_train_begin(self, logs=None):
        pass
    
    def on_train_end(self, logs=None):
        pass
    
    def on_epoch_begin(self, epoch, logs=None):
        pass
    
    def on_epoch_end(self, epoch, logs=None):
        pass
    
    def on_batch_begin(self, batch, logs=None):
        pass
    
    def on_batch_end(self, batch, logs=None):
        pass


class EarlyStopping(Callback):
    """Stop training when a monitored metric has stopped improving"""
    
    def __init__(self, monitor='loss', patience=10, min_delta=0.0001, mode='min', verbose=True):
        self.monitor = monitor
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.verbose = verbose
        self.wait = 0
        self.stopped_epoch = 0
        self.best = float('inf') if mode == 'min' else float('-inf')
        self.stop_training = False
    
    def on_train_begin(self, logs=None):
        self.wait = 0
        self.stopped_epoch = 0
        self.best = float('inf') if self.mode == 'min' else float('-inf')
        self.stop_training = False
    
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is None:
            return
        
        if self._monitor_improved(current):
            self.best = current
            self.wait = 0
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.stop_training = True
                if self.verbose:
                    print(f"Early stopping triggered at epoch {epoch + 1}")
    
    def _monitor_improved(self, current):
        if self.mode == 'min':
            return current < self.best - self.min_delta
        else:
            return current > self.best + self.min_delta


class ModelCheckpoint(Callback):
    """Save the model after every epoch"""
    
    def __init__(self, filepath, monitor='loss', mode='min', save_best_only=True, verbose=True):
        self.filepath = filepath
        self.monitor = monitor
        self.mode = mode
        self.save_best_only = save_best_only
        self.verbose = verbose
        self.best = float('inf') if mode == 'min' else float('-inf')
    
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is None:
            return
        
        should_save = not self.save_best_only or self._monitor_improved(current)
        
        if should_save:
            if self._monitor_improved(current):
                self.best = current
            filepath = self.filepath.replace('{epoch}', str(epoch + 1))
            logs['model'].save(filepath)
            if self.verbose:
                print(f"Epoch {epoch + 1}: {self.monitor} improved to {current:.6f}, saving model to {filepath}")
    
    def _monitor_improved(self, current):
        if self.mode == 'min':
            return current < self.best
        else:
            return current > self.best


class LearningRateScheduler(Callback):
    """Learning rate scheduler"""
    
    def __init__(self, schedule, verbose=True):
        self.schedule = schedule
        self.verbose = verbose
    
    def on_epoch_begin(self, epoch, logs=None):
        logs = logs or {}
        if callable(self.schedule):
            new_lr = self.schedule(epoch)
        elif isinstance(self.schedule, dict):
            new_lr = self.schedule.get(epoch, logs['model'].learning_rate)
        else:
            new_lr = logs['model'].learning_rate
        
        logs['model'].learning_rate = new_lr
        if self.verbose:
            print(f"Epoch {epoch + 1}: Learning rate set to {new_lr}")


class ReduceLROnPlateau(Callback):
    """Reduce learning rate when a metric has stopped improving"""
    
    def __init__(self, monitor='loss', factor=0.5, patience=5, min_delta=0.0001,
                 min_lr=1e-7, mode='min', verbose=True):
        self.monitor = monitor
        self.factor = factor
        self.patience = patience
        self.min_delta = min_delta
        self.min_lr = min_lr
        self.mode = mode
        self.verbose = verbose
        self.wait = 0
        self.best = float('inf') if mode == 'min' else float('-inf')
    
    def on_train_begin(self, logs=None):
        self.wait = 0
        self.best = float('inf') if self.mode == 'min' else float('-inf')
    
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is None:
            return
        
        if self._monitor_improved(current):
            self.best = current
            self.wait = 0
        else:
            self.wait += 1
            if self.wait >= self.patience:
                old_lr = logs['model'].learning_rate
                new_lr = max(old_lr * self.factor, self.min_lr)
                
                if new_lr < old_lr:
                    logs['model'].learning_rate = new_lr
                    if self.verbose:
                        print(f"Epoch {epoch + 1}: Reducing learning rate from {old_lr} to {new_lr}")
                    self.wait = 0
    
    def _monitor_improved(self, current):
        if self.mode == 'min':
            return current < self.best - self.min_delta
        else:
            return current > self.best + self.min_delta


class CSVLogger(Callback):
    """Callback that streams epoch results to a CSV file"""
    
    def __init__(self, filename, separator=',', append=False):
        self.filename = filename
        self.separator = separator
        self.append = append
        self.file = None
        self.keys = None
    
    def on_train_begin(self, logs=None):
        mode = 'a' if self.append else 'w'
        self.file = open(self.filename, mode)
        self.keys = None
    
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        logs_copy = logs.copy()
        logs_copy.pop('model', None)
        logs_copy['epoch'] = epoch + 1
        
        if self.keys is None:
            self.keys = list(logs_copy.keys())
            self.file.write(self.separator.join(self.keys) + '\n')
        
        values = [str(logs_copy.get(key, '')) for key in self.keys]
        self.file.write(self.separator.join(values) + '\n')
        self.file.flush()
    
    def on_train_end(self, logs=None):
        if self.file:
            self.file.close()


class ProgbarLogger(Callback):
    """Callback that prints metrics to stdout"""
    
    def __init__(self, count_mode='steps'):
        self.count_mode = count_mode
    
    def on_train_begin(self, logs=None):
        logs = logs or {}
        self.epochs = logs.get('epochs', 1)
    
    def on_epoch_begin(self, epoch, logs=None):
        logs = logs or {}
        print(f"\nEpoch {epoch + 1}/{self.epochs}")
        self.seen = 0
        self.total = logs.get('steps', 0)
    
    def on_batch_end(self, batch, logs=None):
        self.seen += 1
        progress = round(self.seen / self.total * 100, 1) if self.total > 0 else 0
        bar_length = 30
        filled = int(bar_length * self.seen / self.total) if self.total > 0 else 0
        bar = '=' * filled + '>' + '.' * (bar_length - filled - 1)
        
        print(f"\r{self.seen}/{self.total} [{bar}] {progress}%", end='', flush=True)
    
    def on_epoch_end(self, epoch, logs=None):
        print("")
