class TrainableModule(Module):
    def __init__(
        self,
        hyperparameters: Hyperparams | None = None,
        name: str | None = None,
    ):
        super().__init__(hyperparameters=hyperparameters, name=name)
        self.trackers = NoDependency({k: {} for k in REQUIRED_KEYS})
        self.tracked = NoDependency({})
        self.targets = NoDependency(lambda: {})
        self._step_var = tf.Variable(0, shape=(), dtype=tf.int64, trainable=False)
        self._epoch_var = tf.Variable(0, shape=(), dtype=tf.int64, trainable=False)
        self.batch_indices: list[int] | np.ndarray | None = None
        self.optimizer = None
        self._compiled = False
        self._set_default_trackers()

    def _set_default_trackers(self):
        for key in REQUIRED_KEYS:
            self.trackers.setdefault(key, {})

    @property
    def compiled(self) -> bool:
        return self._compiled

    @property
    def step(self) -> int:
        return int(self._step_var)

    @property
    def epoch(self) -> int:
        return int(self._epoch_var)

    def compile(  # noqa: A003
        self,
        optimizer: tf.keras.optimizers.Optimizer | None = None,
        logged_keys: Literal["main", "all"] | list[str] = "all",
    ):
        def outputs():
            losses = flatten_dict(self.losses, sep="/")
            all_losses = self.get_all_losses()
            by_category = defaultdict(dict)
            for key, loss in losses.items():
                by_category[loss.category][key] = all_losses[key]
            self.add_metric("loss", tf.math.add_n(list(all_losses.values()), name="loss"))
            for category, cat_losses in by_category.items():
                if len(cat_losses) > 1:
                    self.add_metric(
                        f"{category}_loss",
                        tf.math.add_n(list(cat_losses.values()), name=f"{category}_loss")
                        / self._get_loss_lambda(category),
                    )
            metrics = {
                "/".join(k.split(":")[0].split("/")[-3:]): l.result()
                for k, l in flatten_dict(self.metrics, sep="/").items()
            }
            if logged_keys == "all":
                logged_key_set = set(metrics.keys())
            elif logged_keys == "main":
                logged_key_set = {f"{category}_loss" for category in by_category.keys()} | {"loss"}
            else:
                logged_key_set = set(logged_keys)
            return {k: v for k, v in metrics.items() if k in logged_key_set}

        self.targets = outputs
        if optimizer is None:
            optimizer = make_optimizer(self.hyperparameters)
        self.optimizer = optimizer
        # we create the compiled function externally, once so that we can easily decompile it when required.
        self.train_step: TrainStep = tf.function(self.python_train_step)
        self._compiled = True

    def decompile(self):
        del self.targets
        del self.optimizer
        del self.train_step
        self.targets = NoDependency(lambda **_: {})
        self.optimizer = None
        self._compiled = False

    def _get_loss_lambda(self, loss_category: str) -> float:  # noqa: U100
        return 1.0

    def get_all_losses(self) -> dict[str, tf.Tensor]:
        return {
            name: loss.value()
            * self._get_loss_lambda(loss.category)
            * (tf.constant(1.0, dtype=tf.float32) if loss.mult is None else loss.mult)
            for name, loss in flatten_dict(self.losses, sep="/").items()
        }

    def get_total_loss(self) -> tf.Tensor:
        return tf.math.add_n(
            list(self.get_all_losses().values()),
            name="loss",
        )

    def wipe_clean(self, reset_counters: bool = True, reset_trackers: bool = True):
        """re-initialize all variables"""
        if reset_trackers:
            self.trackers = NoDependency({k: {} for k in self.trackers})
            self._set_default_trackers()
        if reset_counters:
            self._step_var.assign(0)
            self._epoch_var.assign(0)

    def get_placeholder(
        self,
        name: str,
        data: pd.DataFrame | np.ndarray,
        batch: bool = False,
        last_is_sim: bool | None = None,
        batch_indices: list[int] | None = None,
        dtype: tf.DType = tf.float32,
    ) -> tf.Tensor:
        name = f"{name}_ph"
        if isinstance(data, pd.DataFrame):
            data = data.ext.values_1
        data_np = np.array(data, dtype=dtype.as_numpy_dtype)
        shape = list(data_np.shape)
        if last_is_sim or (last_is_sim is None and len(shape) > 0 and shape[-1] == 1):
            shape[-1] = None
        ph = None
        if batch:
            shape[0] = None
            batch_indices = self.batch_indices if batch_indices is not None else None
            if batch_indices is not None:
                data_np = data_np[batch_indices]
            ph = tf.convert_to_tensor(data_np, dtype=dtype)
        if ph is None:
            ph = tf.constant(data_np, dtype=dtype, shape=data_np.shape, name=name)
        return tf.ensure_shape(ph, shape=shape, name=name)

    def _make_aligned_string(self, values: dict[str, np.ndarray | "ScalarLike"]) -> str:
        max_k = max(len(k) for k in values.keys())
        return "\n".join(
            [
                f"{k: <{max_k}}: {np.average(v):013.8f}"
                for k, v in sorted(
                    values.items(),
                    key=lambda p: (np.inf if np.isnan(p[1]) else p[1]),
                )
            ]
        )

    def python_train_step(
        self, batch, optimizer: tf.optimizers.Optimizer, return_grads: bool = False
    ) -> tuple[dict[str, tf.Tensor], dict[str, tf.Tensor], tf.Tensor]:
        from wt_ml.layers.layer_utils import to_dense

        self.clear()
        for child in self.submodules:
            if isinstance(child, Module):
                child.clear()
        with tf.GradientTape() as tape:
            self(batch, training=True)
            total_loss = self.get_total_loss()
        trn_vars = self.trn_vars
        gradients = tape.gradient(total_loss, trn_vars)
        optimizer.apply_gradients(zip(gradients, trn_vars))
        step = self._step_var.assign_add(1)
        if return_grads:
            gradients_tracker = {
                variable.name: to_dense(grad) for grad, variable in zip(gradients, trn_vars) if grad is not None
            }
            return (self.targets(), self.get_all_losses() | {"loss": total_loss}, step, gradients_tracker)
        else:
            return (self.targets(), self.get_all_losses() | {"loss": total_loss}, step)

    @tf.function
    def calculate_grad_magnitudes(self, batch):
        from wt_ml.layers.layer_utils import to_dense

        self.clear()
        for child in self.submodules:
            if isinstance(child, Module):
                child.clear()
        with tf.GradientTape(persistent=True) as tape:
            self(batch, training=True)  # noqa: F841
            loss = self.get_total_loss()
            losses = self.get_all_losses()
        trn_vars = self.trn_vars
        gradients = tape.gradient(loss, trn_vars)
        gradients_tracker = {variable.name: to_dense(grad) for grad, variable in zip(gradients, trn_vars)}
        grad_magnitude_tracker = {k: tf.norm(v) for k, v in gradients_tracker.items() if v is not None}
        grad_magnitude_tracker["loss"] = tf.norm(
            tf.concat([tf.reshape(t, (-1,)) for t in gradients_tracker.values() if t is not None], axis=0)
        )
        for k, l in losses.items():
            l_grads = [tf.reshape(g, (-1,)) for g in tape.gradient(l, self.trn_vars) if g is not None]
            if len(l_grads) == 0:
                continue
            elif len(l_grads) == 1:
                flattened = l_grads[0]
            else:
                flattened = tf.concat(l_grads, axis=0)
            grad_magnitude_tracker[k] = tf.norm(flattened)
        return grad_magnitude_tracker

    def train(
        self,
        dataset_factory: Callable[[], Iterable],
        num_steps: int,
        epochs: int = 1,
        verbosity: bool | int = 1,
        print_keys: str | Sequence[str] = "all",
        callbacks: Sequence[Callback] | CallbacksList = (),
        track_grads: bool | int = False,
        smoothing: float = 0.0,
        min_interval: float = 0.25,
        unit_scale: bool = True,
        position: int | None = None,
        tqdm_args: dict[str, Any] = {},
        **kwargs,
    ):
        # kwargs only supports things we actually use.
        # Right now, only things that fall into track_kwarg_names are allowed. Else raise error.
        track_grads_frequency = to_frequency(track_grads)
        track_kwarg_names = [f"track_{key}" for key in self.tracked.keys()]
        track_frequency_values = {key: to_frequency(kwargs.get(f"track_{key}", 0)) for key in self.tracked.keys()}
        unrecognized = [k for k in kwargs.keys() if k not in track_kwarg_names]
        if len(unrecognized) > 0:
            raise TypeError(f"{unrecognized} are not valid keyword arguments for train.")
        verbosity_frequency = to_frequency(verbosity, default=1)

        assert self.optimizer is not None, "You must have an optimizer to train. Did you forget to compile?"

        for track_name in track_frequency_values.keys():
            self.trackers.setdefault(f"all_{track_name}", {})

        epoch_progress_bar = None
        if verbosity_frequency > 0:
            smoothing = 0.03 ** (1 / verbosity_frequency) if smoothing == 0 else smoothing
            epoch_progress_bar = self.get_epoch_progress(
                epochs, smoothing, unit_scale, min_interval, position, tqdm_args
            )

        if num_steps is None:
            raise ValueError("Refusing to train forever, you must provide num_steps if supplying an infinite dataset.")

        epoch_writer = None if epoch_progress_bar is None else epoch_progress_bar.write
        callbacks = callbacks if isinstance(callbacks, CallbacksList) else CallbacksList(callbacks)
        callbacks.register_model(self)
        callbacks.register_writer(epoch_writer)
        callbacks.on_train_start()
        for _ in range(epochs):
            batch_progress_bar = (
                self.get_batch_progress(num_steps, smoothing, unit_scale, min_interval, tqdm_args)
                if verbosity_frequency > 0 and num_steps > 1
                else None
            )
            finished = self.process_batches(
                dataset_factory,
                track_grads_frequency,
                track_frequency_values,
                callbacks,
                verbosity_frequency,
                print_keys,
                batch_progress_bar,
                epoch_progress_bar,
            )
            self._epoch_var.assign_add(1)
            if finished:
                if verbosity:
                    epoch_progress_bar.close()
                break
        callbacks.on_train_end()
        return self.trackers["all_losses"]

    def get_epoch_progress(self, epochs, smoothing, unit_scale, min_interval, position, tqdm_args):
        return tqdm(
            total=self.epoch + epochs,
            smoothing=smoothing,
            unit="epoch",
            dynamic_ncols=True,
            initial=self.epoch,
            unit_scale=unit_scale,
            mininterval=min_interval,
            position=position,
            **tqdm_args,
        )

    def get_batch_progress(self, num_steps, smoothing, unit_scale, min_interval, tqdm_args):
        return tqdm(
            total=num_steps,
            smoothing=smoothing,
            unit="step",
            dynamic_ncols=True,
            leave=False,
            unit_scale=unit_scale,
            mininterval=min_interval,
            **tqdm_args,
        )

    def process_batches(
        self,
        dataset_factory: Callable[[], Iterable],
        track_grads_frequency: int,
        track_frequency_values: dict[str, int],
        callbacks: CallbacksList,
        verbosity_frequency: int,
        print_keys: str | Sequence[str],
        batch_progress_bar: tqdm,
        epoch_progress_bar: tqdm,
        force_update: bool = False,
    ) -> bool:
        batch_write = batch_progress_bar.write if batch_progress_bar is not None else logger.info
        self.reset_metrics(num_metrics=len(self._local_metrics))
        finished = False
        outputs = {}
        for batch_index, batch in enumerate(dataset_factory()):
            current_step = self.step
            with tf.profiler.experimental.Trace("Train", step_num=current_step, _r=1):
                self.perform_tracking_operations(
                    current_step, track_grads_frequency, batch, track_frequency_values, force=force_update
                )
                outputs, losses, *_ = self.train_step(batch, optimizer=self.optimizer)
                outputs = {k: to_numpy(v) for k, v in outputs.items()}
                losses = {k: to_numpy(v) for k, v in losses.items()}
            if not callbacks.on_batch_end(outputs, batch, write=batch_write):
                finished = True
                break
            self.trackers["all_losses"][current_step] = losses
            self.trackers["all_metrics"][current_step] = outputs
            if verbosity_frequency > 0 and batch_index % verbosity_frequency == 0:
                _ = self.handle_verbose_operations(
                    verbosity_frequency, batch_index, print_keys, batch_progress_bar, outputs
                )
            if any(np.isnan(value) for key, value in outputs.items() if "loss" in key):
                raise ValueError(f"NaN Loss at step {current_step}!,\n{self._make_aligned_string(outputs)}")
        if verbosity_frequency > 0:
            self.handle_epoch_progress(verbosity_frequency, print_keys, epoch_progress_bar, outputs)
        if batch_progress_bar is not None:
            batch_progress_bar.close()
        self.trackers["epoch_metrics"][self.epoch] = outputs
        if not callbacks.on_epoch_end(self.trackers["epoch_metrics"][self.epoch]):
            finished = True
        self.reset_metrics(num_metrics=len(self._local_metrics))
        return finished

    def run_all_trackers(self):
        current_step = self.step
        for track_key, tracked_func in self.tracked.items():
            with tf.profiler.experimental.Trace(track_key):
                self.reset_metrics(num_metrics=len(self._local_metrics))
                self.trackers.setdefault(f"all_{track_key}", {})[current_step] = tracked_func()
                self.reset_metrics(num_metrics=len(self._local_metrics))

    def perform_tracking_operations(
        self, current_step, track_grads_frequency, batch, track_frequency_values, force: bool = False
    ):
        if track_grads_frequency and (force or current_step % track_grads_frequency == 0):
            with tf.profiler.experimental.Trace("Grads"):
                grad_magnitude_tracker = self.calculate_grad_magnitudes(batch)
            self.trackers["all_grad_magnitudes"][current_step] = to_numpy(grad_magnitude_tracker, skip_none=True)
        for track_key, track_value in track_frequency_values.items():
            if track_value and (force or current_step % track_value == 0):
                with tf.profiler.experimental.Trace(track_key):
                    self.reset_metrics(num_metrics=len(self._local_metrics))
                    self.trackers[f"all_{track_key}"][current_step] = self.tracked[track_key]()
                    self.reset_metrics(num_metrics=len(self._local_metrics))

    def handle_verbose_operations(self, verbosity_frequency, batch_index, print_keys, batch_progress_bar, outputs):
        output_keys = outputs.keys() if print_keys == "all" else print_keys
        if batch_progress_bar is not None and batch_index % verbosity_frequency == 0:
            batch_progress_bar.set_postfix(
                {key: f"{np.average(value):09.4f}" for key, value in outputs.items() if key in output_keys}
            )
            batch_progress_bar.update(verbosity_frequency)
        return output_keys

    def handle_epoch_progress(self, verbosity_frequency, print_keys, epoch_progress_bar, outputs):
        output_keys = outputs.keys() if print_keys == "all" else print_keys
        if verbosity_frequency > 0:
            epoch_progress_bar.set_postfix(
                {key: f"{metric:09.4f}" for key, metric in outputs.items() if key in output_keys}
            )
            epoch_progress_bar.update(1)

    def get_kwargs_to_save(self, **kwargs):  # noqa: U100
        return {
            "name": self.name,
            "hyperparameters": self.hyperparameter_tree,
        }

    def save_trackers(self, filepath: str | Path, include_trackers=False):
        with filepath.open("wb") as fp:
            if include_trackers:
                pickle.dump({k: list(v.items()) for k, v in self.trackers.items()}, fp)
            else:
                pickle.dump({k: [] for k in self.trackers.keys()}, fp)

    def save(self, folder: str | Path, include_trackers=False, **kwargs):
        saver = tf.train.Checkpoint(self)
        folder = Path(folder)
        folder.mkdir(parents=True, exist_ok=True)
        with (folder / self.KWARGS_FILE_NAME).open("w") as fp:
            yaml.safe_dump(self.get_kwargs_to_save(**kwargs), fp)
        saver.write(str(folder / "model"))
        self.save_trackers(folder / self.TRACKERS_FILE_NAME, include_trackers)
        with (folder / "meta.json").open("w+") as fp:
            # useful for debugging issues with model caching.
            json.dump({"id": str(id(self)), "name": self.name, "name_scope": str(self.name_scope.name)}, fp)
        self.hyperparameters.save(folder / "hyperparameters.yml")

    def restore_trackers(self, filepath: str | Path):
        from wt_ml.tuning.utils import in_cpu

        if filepath.exists():
            with filepath.open("rb") as fp:
                model_tracker = in_cpu(pickle.load)(fp)
                # We are looping across trackers as (v.items() if isinstance(v,dict) else v) to
                # handle backward compatibility with old trackers after the feature that converts
                # all trackers and intermediaries to dictionaries
                self.trackers = NoDependency(
                    {
                        k: {k1: v1 for k1, v1 in (v.items() if isinstance(v, dict) else v)}
                        for k, v in model_tracker.items()
                    }
                )
        self._set_default_trackers()

    def restore(
        self, path: str | Path, no_trackers: bool = False, partial_restore: bool = False, catch_exceptions: bool = True
    ):
        path = Path(path)
        if partial_restore:
            restore.partial_restore(self, path)
        else:
            saver = tf.train.Checkpoint(self)
            try:
                saver.read(str(path / "model")).expect_partial()
            except (ValueError, AssertionError) as err:
                if catch_exceptions:
                    logger.exception(err, f"Model unrestored. A {type(err).__name__} exception has occured.")
                else:
                    raise

        if path.is_dir():
            self.restore_trackers(path / self.TRACKERS_FILE_NAME) if not no_trackers else None
        else:
            raise FileNotFoundError(f"Path doesn't exist - {path}.")

    @classmethod
    def load_kwargs(cls, folder: str | Path, **kwargs):  # noqa: U100
        folder = Path(folder)
        with (folder / cls.KWARGS_FILE_NAME).open("r") as fp:
            loaded = yaml.safe_load(fp)
        return loaded

    @classmethod
    def from_save(
        cls,
        folder: str | Path,
        no_trackers: bool = False,
        override_kwargs: dict[str, Any] | None = None,
        build_argument=None,
        **kwargs,
    ):
        folder = Path(folder)
        constructor_kwargs = cls.load_kwargs(folder, **kwargs)
        if override_kwargs is not None:
            constructor_kwargs = constructor_kwargs | override_kwargs
        logger.info(f"Loading model {constructor_kwargs.get('name', cls.__name__)} from path {folder}.")
        model = cls(**constructor_kwargs)
        if build_argument is not None:
            model(build_argument, training=True, debug=False, skip_metrics=False)
        model.compile()
        model.restore(folder, no_trackers=no_trackers)
        return model

    @overload
    def get_tracker_at(self, name: str, epoch: Sequence[int] | Literal["all"]) -> list[Any]:  # noqa: U100
        ...

    @overload
    def get_tracker_at(self, name: str, epoch: int | None) -> Any:  # noqa: U100
        ...

    def get_tracker_at(self, name: str, epoch: EpochSpec = None) -> list[Any] | Any:
        all_name = f"all_{name}"
        if all_name not in self.trackers:
            self.trackers[all_name] = {}
        if epoch is None:
            if name in ("losses", "grads", "grad_magnitudes"):
                epoch = np.max(list(self.trackers[all_name].keys()))
            else:
                epoch = int(self.epoch)
            if epoch not in self.trackers[all_name]:
                if name not in self.tracked:
                    return None
                self.trackers[all_name][epoch] = self.tracked[name]()
        elif epoch == "all":
            cur_epoch = int(self.epoch)
            if cur_epoch not in self.trackers[all_name] and name not in ("losses", "grads", "grad_magnitudes"):
                self.trackers[all_name][cur_epoch] = self.tracked[name]()
            epoch = tuple(self.trackers[all_name].keys())
        return map_possible_seq(epoch, lambda ep: self.trackers[all_name][ep])

    def viz_tracker(
        self,
        name: str,
        label: str,
        epoch_min: int = 0,
        epoch_max: int | None = None,
        height: int = 800,
        smoothing: float = 0,
        is_epoch: bool = False,
    ):
        loss_df = pd.DataFrame(to_numpy(self.trackers[name])).T
        loss_df.columns.names = [label]
        loss_df.index.names = ["epochs"] if is_epoch else ["step"]
        loss_df.sort_values(label, axis=1, inplace=True)
        if smoothing > 0:
            if smoothing > 1:
                raise ValueError("Cannot have smoothing greater than one.")
            elif smoothing == 1:
                loss_df = loss_df.expanding(1).mean()
            else:
                loss_df = loss_df.ewm(alpha=1 - smoothing).mean()
        elif smoothing < 0:
            raise ValueError("Cannot have a smoothing less than zero.")
        idxs = loss_df.index >= epoch_min
        if epoch_max is not None:
            idxs = idxs & (loss_df.index <= epoch_max)
        fig = loss_df.loc[idxs].plot(backend="plotly", height=height)
        fig.update_layout(legend=dict(orientation="h", yanchor="bottom", y=-2, xanchor="left", x=0))
        return fig

    def viz_losses(self, **kwargs):
        return self.viz_tracker(name="all_losses", label="loss", is_epoch=False, **kwargs)

    def viz_metrics(self, epoch: bool = False, **kwargs):
        if epoch:
            return self.viz_tracker(name="epoch_metrics", label="loss", is_epoch=True, **kwargs)
        else:
            return self.viz_tracker(name="all_metrics", label="loss", is_epoch=False, **kwargs)

    def viz_grads(self, **kwargs):
        return self.viz_tracker(name="all_grad_magnitudes", label="gradient", is_epoch=False, **kwargs)

    TRACKERS_FILE_NAME = "trackers.pkl"
    KWARGS_FILE_NAME = "kwargs.yaml"