STATUS_SYMBOLS = {
    'ok':        '✓',
    'vanishing': '↓',
    'exploding': '↑',
    'stagnant':  '■',
}

STATUS_LABELS = {
    'ok':        'OK',
    'vanishing': 'VANISHING GRADIENT',
    'exploding': 'EXPLODING GRADIENT',
    'stagnant':  'STAGNANT LAYER',
}


class Reporter:
    def __init__(self, log_every=10, verbose=True):
        self.log_every = log_every
        self.verbose = verbose
        self._alerts = []

    def alert(self, layer_name, status, norm, step):
        entry = {'step': step, 'layer': layer_name, 'status': status, 'norm': norm}
        self._alerts.append(entry)
        if self.verbose:
            sym = STATUS_SYMBOLS.get(status, '?')
            label = STATUS_LABELS.get(status, status.upper())
            parts = [
                '[torch-surgeon] step', str(step).rjust(6),
                '|', sym, label,
                '|', layer_name,
                '| norm=' + format(norm, '.4e')
            ]
            print(' '.join(parts))

    def summary(self, stats, step):
        if not self.verbose or not stats:
            return
        if step % self.log_every != 0:
            return

        ranked = sorted(stats.items(), key=lambda x: x[1].get('norm', 0), reverse=True)
        print()
        print('[torch-surgeon] Step ' + str(step) + ' - gradient attribution ranking')
        print('  ' + 'Layer'.ljust(35) + ' ' + 'Norm'.rjust(12) + ' Status')
        print('  ' + '-' * 35 + ' ' + '-' * 12 + ' ' + '-' * 12)
        top = ranked[:5]
        bot = ranked[-3:] if len(ranked) > 5 else []
        for name, s in top:
            sym = STATUS_SYMBOLS.get(s['status'], '?')
            print('  ' + name.ljust(35) + ' ' + format(s['norm'], '12.4e') + ' ' + sym + ' ' + s['status'])
        if bot:
            print('  ...(' + str(len(ranked) - 5) + ' layers omitted)...')
            for name, s in bot:
                sym = STATUS_SYMBOLS.get(s['status'], '?')
                print('  ' + name.ljust(35) + ' ' + format(s['norm'], '12.4e') + ' ' + sym + ' ' + s['status'])
        print()

    def get_alerts(self):
        return list(self._alerts)

    def clear_alerts(self):
        self._alerts.clear()
