from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Tuple, Union

import tensorflow as tf
from tensorflow import keras

from useckit.v2.dataset import Dataset
from useckit.v2.evaluator import Paradigm, ModelArchitecture

K = keras.backend
TensorLike = Union[keras.layers.Layer, tf.Tensor]

# ---- Introspection utilities --------------------------------------------------

def _clone_inputs_like(inputs: Union[tf.Tensor, List[tf.Tensor]]) -> List[keras.layers.Input]:
    """Create new Input tensors with the same shapes as `inputs`."""
    if isinstance(inputs, (list, tuple)):
        return [keras.layers.Input(shape=tuple(t.shape)[1:], dtype=t.dtype) for t in inputs]
    else:
        return [keras.layers.Input(shape=tuple(inputs.shape)[1:], dtype=inputs.dtype)]

def _l2_normalize(x: tf.Tensor) -> tf.Tensor:
    return tf.nn.l2_normalize(x, axis=-1)

# ---- Core builder -------------------------------------------------------------

@dataclass
class BuildResult:
    model: keras.Model
    callbacks: List[Callable]
    metadata: Dict

    def __getitem__(self, key):
        if key == "model":
            return self.model
        elif key == "callbacks":
            return self.callbacks
        elif key == "metadata":
            return self.metadata
        else:
            raise KeyError(f"Invalid key '{key}'. Valid keys: 'model', 'callbacks', 'metadata'.")


class ModelBuilder:
    """
    Wraps the MODEL_DESCRIPTIONS (dl-4-tsc family) and constructs:
      - Multiclass classifiers
      - Binary (verification) classifiers
      - Metric-learning models (contrastive or triplet)
      - Autoencoders for one-class detection (RMSE thresholding)
    without changing the original description callables.

    Assumptions:
      - Each description callable returns (compiled_model, callbacks).
      - The penultimate layer is a sensible feature tensor for a new head.
      - TWIESN is ignored.
    """

    def __init__(
        self,
        model_descriptions: Dict[str, Callable[[TensorLike, int], Tuple[keras.Model, List[Callable]]]],
        ignored: Optional[set[str]] = None,
    ):
        self.model_descriptions = dict(model_descriptions)
        self.ignored = set() if ignored is None else set(ignored)

    # --- Public APIs ----------------------------------------------------------

    def build_model(self, paradigm: Paradigm, dataset: Dataset, model_architecture: ModelArchitecture, **kwargs):
        """Builds and returns models as defined by the respective paradigm and architecture.
        kwargs can contain additional parameters passed to the model (e.g., an optimizer)."""
        #print("#TODO: Implement the passing of the Optimizer via optimizer=, here! -- kwargs:", kwargs)

        if paradigm == Paradigm.CSC:
            return self._build_multiclass_classification(arch_key=model_architecture,
                                                         input_shape=dataset.get_dl_training_shape(),
                                                         num_classes=len(dataset.get_unique_classes()),
                                                         **kwargs)
        elif paradigm == Paradigm.BV:
            return self._build_binary_verification_classifier(arch_key=model_architecture,
                                                              input_shape=dataset.get_dl_training_shape(),
                                                              **kwargs)
        elif paradigm == Paradigm.AD:
            return self._build_autoencoder_oneclass_classifier(arch_key=model_architecture,
                                                               input_shape=dataset.get_dl_training_shape(),
                                                               **kwargs)
        elif paradigm == Paradigm.DL:
            assert "distance_learning_loss" in kwargs, "distance_learning_loss must be specified for DL paradigm."
            if kwargs["distance_learning_loss"] == "contrastive":
                return self._build_metriclearning_contrastive(arch_key=model_architecture,
                                                              input_shape=dataset.get_dl_training_shape(),
                                                              **kwargs)
            elif kwargs["distance_learning_loss"] == "triplet":
                return self._build_metriclearning_triplet(arch_key=model_architecture,
                                                          input_shape=dataset.get_dl_training_shape(),
                                                          **kwargs)
            else:
                raise ValueError(
                    f"No or unknown `distance_learning_loss` specified: "
                    f"{'None' if not hasattr(self, 'distance_learning_loss') else self.distance_learning_loss}."
                )

    def _build_multiclass_classification(
        self,
        arch_key: ModelArchitecture,
        input_shape: Tuple[int, int],
        num_classes: int,
        freeze_backbone: bool = False,
        optimizer: Optional[keras.optimizers.Optimizer] = None,
        **kwargs
    ) -> BuildResult:
        backbone, callbacks = self._make_backbone(arch_key, input_shape)
        if freeze_backbone:
            backbone.trainable = False

        x = backbone.outputs[0]
        y = keras.layers.Dense(num_classes, activation="softmax", name="cls_head")(x)
        model = keras.Model(inputs=backbone.inputs, outputs=y, name=f"{arch_key}_multiclass")

        model.compile(
            optimizer=optimizer or keras.optimizers.Adam(),
            loss="categorical_crossentropy",
            metrics=["accuracy"],
        )
        return BuildResult(model, callbacks, {"paradigm": "multiclass", "arch": arch_key})

    def _build_binary_verification_classifier(
        self,
        arch_key: ModelArchitecture,
        input_shape: Tuple[int, int],
        freeze_backbone: bool = False,
        class_weight: Optional[dict] = None,
        optimizer: Optional[keras.optimizers.Optimizer] = None,
        **kwargs
    ) -> BuildResult:
        """
        Classic 1-head sigmoid classifier (useful if you frame verification as binary classification).
        """
        backbone, callbacks = self._make_backbone(arch_key, input_shape)
        if freeze_backbone:
            backbone.trainable = False

        x = backbone.outputs[0]
        y = keras.layers.Dense(1, activation="sigmoid", name="bin_head")(x)
        model = keras.Model(inputs=backbone.inputs, outputs=y, name=f"{arch_key}_binary")

        model.compile(
            optimizer=optimizer or keras.optimizers.Adam(),
            loss="binary_crossentropy",
            metrics=["accuracy", keras.metrics.AUC(name="auc")],
        )
        # Note: pass class_weight to model.fit if desired
        return BuildResult(model, callbacks, {"paradigm": "binary", "arch": arch_key})

    def _build_metriclearning_contrastive(
        self,
        arch_key: ModelArchitecture,
        input_shape: Tuple[int, int],
        embedding_dim: Optional[int] = None,
        l2_normalize: bool = True,
        margin: float = 1.0,
        optimizer: Optional[keras.optimizers.Optimizer] = None,
        **kwargs
    ) -> BuildResult:
        """
        Siamese (two-tower) model with contrastive loss.
        y_true ∈ {0,1} with 1 = same, 0 = different, model outputs per-pair distances.
        See: Keras Siamese/contrastive example and TFA ContrastiveLoss.
        """
        backbone, callbacks = self._make_backbone(arch_key, input_shape)

        # Two sets of inputs matching the backbone
        inp_a = _clone_inputs_like(backbone.inputs)
        inp_b = _clone_inputs_like(backbone.inputs)

        # Shared backbone call (reuses weights)
        feat_a = backbone(inp_a if len(inp_a) > 1 else inp_a[0])
        feat_b = backbone(inp_b if len(inp_b) > 1 else inp_b[0])

        # Optional projection head (shared layers across both towers)
        embedding_dim = keras.layers.Dense(embedding_dim, activation=None, name="proj") if embedding_dim is not None else None
        l2_layer = keras.layers.Lambda(_l2_normalize, name="l2_norm") if l2_normalize else None

        def proj(x):
            if embedding_dim is not None:
                x = embedding_dim(x)
            if l2_layer is not None:
                x = l2_layer(x)
            return x

        za = proj(feat_a)
        zb = proj(feat_b)

        # L2 distance
        dist = keras.layers.Lambda(lambda t: tf.norm(t[0] - t[1], axis=-1), name="pair_distance")([za, zb])

        siamese = keras.Model(
            inputs=(inp_a + inp_b),
            outputs=dist,
            name=f"{arch_key}_siamese_contrastive"
        )

        from useckit.v2.utils.contrastive_loss import contrastive_loss
        siamese.compile(
            optimizer=optimizer or keras.optimizers.Adam(),
            loss=contrastive_loss,
        )
        meta = {
            "paradigm": "verification_contrastive",
            "arch": arch_key,
            "embedding_dim": embedding_dim,
            "margin": margin,
        }
        return BuildResult(siamese, callbacks, meta)

    def _build_metriclearning_triplet(
        self,
        arch_key: ModelArchitecture,
        input_shape: Tuple[int, int],
        embedding_dim: Optional[int] = None,
        l2_normalize: bool = True,
        margin: float = 0.2,
        optimizer: Optional[keras.optimizers.Optimizer] = None,
        **kwargs
    ) -> BuildResult:
        """
        Single-tower embedding model compiled with TripletSemiHardLoss.
        During training, pass integer class labels as y; the loss internally mines triplets per batch.
        See: TensorFlow Addons triplet tutorial.
        """
        backbone, callbacks = self._make_backbone(arch_key, input_shape)
        z = backbone.outputs[0]
        if embedding_dim is not None:
            z = keras.layers.Dense(embedding_dim, activation=None, name="proj")(z)
        if l2_normalize:
            z = keras.layers.Lambda(_l2_normalize, name="l2_norm")(z)

        embedder = keras.Model(inputs=backbone.inputs, outputs=z, name=f"{arch_key}_embedder_triplet")

        from useckit.paradigms.distance_learning.prediction_models.online_triplet_loss.triplet_semihard_loss_function import \
            TripletSemiHardLoss
        embedder.compile(
            optimizer=optimizer or keras.optimizers.Adam(),
            loss=TripletSemiHardLoss(),
        )
        meta = {
            "paradigm": "metric_triplet",
            "arch": arch_key,
            "embedding_dim": embedding_dim,
            "margin": margin,
            "tfa": True,
        }
        return BuildResult(embedder, callbacks, meta)

    def _build_autoencoder_oneclass_classifier(
        self,
        arch_key: ModelArchitecture,
        input_shape: Tuple[int, int],
        embedding_dim: Optional[int] = 128,
        decoder_channels: Optional[int] = None,
        optimizer: Optional[keras.optimizers.Optimizer] = None,
        **kwargs
    ) -> BuildResult:
        """
        Uses the classification backbone as an encoder, adds a small bottleneck, and a generic 1D decoder
        to reconstruct the input (shape (T, C)). Train with MSE; decide via RMSE threshold on recon error.
        Inspired by the Keras time-series autoencoder example.
        """
        T, C = input_shape
        dec_ch = decoder_channels or min(64, max(8, C * 4))

        backbone, callbacks = self._make_backbone(arch_key, input_shape)
        x = backbone.outputs[0]

        # Bottleneck
        z = x
        if embedding_dim is not None:
            z = keras.layers.Dense(embedding_dim, activation="relu", name="bottleneck")(z)

        # Project back to T * dec_ch, reshape to (T, dec_ch), then a few Conv1D + final Conv1D(C, 1)
        y = keras.layers.Dense(T * dec_ch, activation="relu")(z)
        y = keras.layers.Reshape((T, dec_ch))(y)
        y = keras.layers.Conv1D(dec_ch, kernel_size=5, padding="same", activation="relu")(y)
        y = keras.layers.Conv1D(dec_ch, kernel_size=3, padding="same", activation="relu")(y)
        recon = keras.layers.Conv1D(C, kernel_size=1, padding="same", activation=None, name="reconstruction")(y)

        ae = keras.Model(inputs=backbone.inputs, outputs=recon, name=f"{arch_key}_autoencoder")

        ae.compile(
            optimizer=optimizer or keras.optimizers.Adam(),
            loss="mse",
            metrics=[keras.metrics.RootMeanSquaredError(name="rmse")],
        )
        meta = {
            "paradigm": "autoencoder_oneclass",
            "arch": arch_key,
            "embedding_dim": embedding_dim,
            "decoder_channels": dec_ch,
        }
        return BuildResult(ae, callbacks, meta)

    # --- Internals ------------------------------------------------------------

    def _make_backbone(
        self, arch_key: str, input_shape: Tuple[int, int], nb_classes_hint: int = 2
    ) -> Tuple[keras.Model, List[Callable]]:
        """
        Calls the original description callable once, then drops the last layer to expose features.
        """
        if arch_key in self.ignored:
            raise ValueError(f"Architecture {arch_key} is ignored in this builder.")
        if arch_key not in self.model_descriptions:
            raise KeyError(f"Unknown architecture key: {arch_key}")

        desc_fn = self.model_descriptions[arch_key]

        # Create an "input_layer" that matches the expected shape (T, C).
        # Many dl4tsc models accept a single Input, but MCNN/MCDCNN internally
        # create separate inputs; they still infer shape from this Input.
        input_layer = keras.layers.Input(shape=input_shape, name="input")

        base_model, callbacks = desc_fn(input_layer, nb_classes_hint)  # original compiled model

        # The penultimate layer's output is our feature tensor.
        if not base_model.layers:
            raise RuntimeError(f"Empty model for {arch_key}")

        feature_tensor = base_model.layers[-2].output  # robust default across these defs
        backbone = keras.Model(
            inputs=base_model.inputs, outputs=feature_tensor, name=f"{arch_key}_backbone"
        )
        return backbone, callbacks


if __name__ == "__main__":
    from useckit.v2.model_descriptions import MODEL_DESCRIPTIONS

    builder = ModelBuilder(MODEL_DESCRIPTIONS, ignored={"TWIESN"})

    # Shapes: (timesteps, channels)
    input_shape = (128, 3)

    # Multiclass (e.g., 5 classes)
    multiclass = builder.build_multiclass_classification("RESNET", input_shape, num_classes=5)
    multiclass.model.summary()

    # Binary verification classifier
    binary = builder.build_binary_verification_classifier("FCN", input_shape)
    binary.model.summary()

    # Contrastive verification (siamese)
    siamese = builder.build_metriclearning_contrastive("INCEPTION", input_shape, embedding_dim=64, margin=1.0)
    siamese.model.summary()

    # Autoencoder (one-class)
    ae = builder.build_autoencoder_oneclass("ENCODER", input_shape, embedding_dim=128)
    ae.model.summary()
