import torch.nn as nn
from .detector import GradientDetector
from .reporter import Reporter
from .hooks import make_backward_hook, make_forward_hook
from .plotter import Plotter

DEFAULT_RULES = {
    'vanishing_threshold': 0.01,
    'exploding_threshold': 100.0,
    'stagnant_steps': 50,
    'log_every': 10,
    'plot': False,
    'verbose': True,
    'ema_alpha': 0.1,
}


class Surgeon:
    """
    Real-time gradient pathology detector for PyTorch models.

    Usage::

        from torch_surgeon import Surgeon
        surgeon = Surgeon(model, rules='default')
        surgeon.attach()
        # ... training loop ...
        surgeon.detach()
        report = surgeon.report()
    """

    def __init__(self, model: nn.Module, rules='default'):
        self.model = model
        self._hooks = []
        self._attached = False
        self._step = [0]  # mutable container so hooks can increment

        # Resolve rules
        if rules == 'default':
            cfg = DEFAULT_RULES.copy()
        elif isinstance(rules, dict):
            cfg = {**DEFAULT_RULES, **rules}
        else:
            raise ValueError("rules must be 'default' or a dict")

        self._log_every = cfg['log_every']

        self.detector = GradientDetector(
            vanishing_threshold=cfg['vanishing_threshold'],
            exploding_threshold=cfg['exploding_threshold'],
            stagnant_steps=cfg['stagnant_steps'],
            ema_alpha=cfg['ema_alpha'],
        )
        self.reporter = Reporter(
            log_every=cfg['log_every'],
            verbose=cfg['verbose'],
        )
        self.plotter = Plotter(
            log_every=cfg['log_every'],
        ) if cfg['plot'] else None

    def attach(self):
        """Register hooks on all named layers. Call before training loop."""
        if self._attached:
            return self

        for name, module in self.model.named_modules():
            if not name:  # skip root module
                continue
            if not list(module.children()):  # leaf modules only
                bwd = module.register_full_backward_hook(
                    make_backward_hook(name, self.detector,
                                       self.reporter, self._step)
                )
                fwd = module.register_forward_hook(
                    make_forward_hook(name, self.detector)
                )
                self._hooks.extend([bwd, fwd])

        # Step counter hook on model forward
        def _step_hook(module, input, output):
            self._step[0] += 1
            stats = self.detector.get_stats()
            self.reporter.summary(stats, self._step[0])
            if self.plotter:
                self.plotter.update(stats, self._step[0])

        root_hook = self.model.register_forward_hook(_step_hook)
        self._hooks.append(root_hook)
        self._attached = True
        n_layers = len(self._hooks) - 1
        print('[torch-surgeon] Attached to ' + str(n_layers) + ' layer(s). Monitoring vanishing, exploding, and stagnant gradients.')
        return self

    def detach(self):
        """Remove all hooks cleanly. Call after training loop."""
        for h in self._hooks:
            h.remove()
        self._hooks.clear()
        self._attached = False
        print(f'[torch-surgeon] All hooks removed.')
        return self

    def report(self):
        """
        Return per-layer stats dict.

        Returns::

            {
                'layer_name': {
                    'mean': float,
                    'std': float,
                    'norm': float,
                    'ratio': float,
                    'status': 'ok' | 'vanishing' | 'exploding' | 'stagnant',
                },
                ...
            }
        """
        return self.detector.get_stats()

    def save_plot(self, path='gradient_norms.png'):
        """Save the gradient norm plot to disk."""
        if self.plotter:
            self.plotter.save(path)
        else:
            print('[torch-surgeon] Plotting not enabled. '
                  "Pass rules={'plot': True} to enable.")

    def reset(self):
        """Reset all stats and step counter."""
        self.detector.reset()
        self.reporter.clear_alerts()
        self._step[0] = 0
        return self

    @property
    def is_attached(self):
        return self._attached

    @property
    def step(self):
        return self._step[0]