from __future__ import annotations

import tensorflow as tf
#import tensorflow_addons as tfa
from tensorflow.keras import layers, Model, callbacks, saving
from typing import Tuple, Dict
import numpy as np
import tqdm
from tqdm.keras import TqdmCallback

# ===================================================================
#  UTILITIES
# ===================================================================

@saving.register_keras_serializable(package='GRU')
def gaussian_nll(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:  # noqa: D401
    """Gaussian negative‑log likelihood for *scalar* targets.

    The model is assumed to predict concatenated [mu, log_sigma].

    Given targets :math:`y` and predictions :math:`\mu, \log\sigma`,
    the NLL is

    .. math::

        \mathcal{L}_{NLL} = \frac{1}{2} \exp(-2\log\sigma)(y-\mu)^2 
                             + \log\sigma.

    A small constant is added for numerical stability.
    """
    mu, log_sigma = tf.split(y_pred, 2, axis=-1)
    # Ensure y_true has the same shape as mu for the subtraction
    y_true_shaped = tf.reshape(y_true, tf.shape(mu))
    inv_var = tf.exp(-2.0 * log_sigma)  # = 1/sigma^2
    nll = 0.5 * inv_var * tf.square(y_true_shaped - mu) + log_sigma
    return tf.reduce_mean(nll)


# ===================================================================
#  MODEL FACTORY
# ===================================================================

def build_gru_model(lookback: int, n_features: int) -> Model:
    """Return the three‑head GRU model described in *revisions.txt*.

    Architecture:
      GRU(64) →
        • Dense(2, name="gauss_params") → [μ̂, log σ̂]
        • Dense(1, name="ret")          – μ̂ sliced from gauss_params
        • Dense(1, activation="sigmoid", name="dir") – probability price rises
    """
    inputs = layers.Input(shape=(lookback, n_features), name="input")
    h = layers.GRU(64, name="gru")(inputs)

    # concatenated [mu, log_sigma] for Gaussian‑NLL
    gauss_out = layers.Dense(2, name="gauss_params")(h)
    # Slice mu so we can attach a separate Huber loss on μ̂
    mu_out = layers.Lambda(lambda x: x[:, :1], name="ret")(gauss_out)
    dir_prob = layers.Dense(1, activation="sigmoid", name="dir")(h)

    model = Model(inputs, outputs=[mu_out, gauss_out, dir_prob], name="crypto_gru")

    losses: Dict[str, tf.keras.losses.Loss | str] = {
        "ret": tf.keras.losses.Huber(delta=1e-3),
        "gauss_params": gaussian_nll,
        "dir": tf.keras.losses.BinaryFocalCrossentropy(alpha=0.5, gamma=2.0, from_logits=False),
    }

    loss_weights = {"ret": 1.0, "gauss_params": 0.2, "dir": 0.4}

    # Note: The targets passed to model.fit must align with the outputs.
    # y_train should be a dictionary like:
    # { "ret": y_ret_train, "gauss_params": y_ret_train, "dir": y_dir_train }
    model.compile(
        optimizer="adam",
        loss=losses,
        loss_weights=loss_weights,
        run_eagerly=False,
    )
    return model


# ===================================================================
#  HIGH‑LEVEL TRAINING WRAPPER (REMOVED)
# ===================================================================
# Standalone train_model function removed as logic is now in GRUModelHandler.train
# def train_model(...): 
#    ...