import json
import re
from pathlib import Path

import torch
import torch.nn as nn
from torch import Tensor
from jaxtyping import Float
from huggingface_hub import PyTorchModelHubMixin


class Model(nn.Module, PyTorchModelHubMixin, repo_url="https://huggingface.co/TimWalter/RAM/new/main",  paper_url="https://huggingface.co/papers/2606.09108",license="mit",):
    """
    RAM model that predicts reachability from a modified Denavit-Hartenberg morphology parametrisation and a pose.
    """

    @classmethod
    def from_id(cls, model_id: int):
        """
        Instantiate a model from its wandb ID.

        Args:
            model_id: Wandb ID of the model.

        Returns:
            model: Instantiated model.
        """

        model_dir = Path(__file__).parent.parent / "trained_models"
        pattern = rf"{model_id}-[a-z]+-[a-z]+"
        folder = next((f for f in model_dir.iterdir() if re.match(pattern, f.name)), None)
        metadata_path = model_dir / folder / 'metadata.json'
        metadata = json.load(open(metadata_path, 'r'))

        model = cls(**metadata["hyperparameter"])
        model_folder = Path(str(model_dir / folder))
        prime = model_folder / "model.pth"

        if prime.exists():
            old_state_dict =torch.load(prime)
        else:
            old_state_dict =torch.load(model_folder / "checkpoint.pth")
        new_state_dict = {}
        # 2. Map legacy state dict keys to your new flattened structure
        for key, value in old_state_dict.items():
            new_key = key
            if key.startswith("encoder.lstm."):
                new_key = key.replace("encoder.lstm.", "encoder.")
            elif key.startswith("decoder.model."):
                new_key = key.replace("decoder.model.", "decoder.")

            new_state_dict[new_key] = value

        # Load the corrected state dict into the model
        model.load_state_dict(new_state_dict)

        return model

    def __init__(self,
                 dim_encoding: int = 128,
                 num_encoder_layers: int = 1,
                 drop_prob: float = 0.0,
                 dim_decoder: int = 1792,
                 num_decoder_layer: int = 8):
        """
        Initialise the model.

        Args:
            dim_encoding: The dimension of the latent morphology encoding.
            num_encoder_layers: Number of LSTM layers.
            drop_prob: Dropout probability of the LSTM.
            dim_decoder: Hidden dimension of the MLP.
            num_decoder_layer: Number of layers of the MLP.
        """

        super().__init__()
        self.encoder = nn.LSTM(3, dim_encoding, num_encoder_layers, dropout=drop_prob, batch_first=True, bias=False)
        self.decoder = nn.Sequential(
            nn.Linear(9 + dim_encoding, dim_decoder),
            nn.ReLU(),
            *[nn.Sequential(nn.Linear(dim_decoder, dim_decoder), nn.ReLU())
              for _ in range(num_decoder_layer)],
            nn.Linear(dim_decoder, 1)
        )

    def forward(self, morph: Float[Tensor, "batch seq 3"], pose: Float[Tensor, "batch 9"]) -> Float[Tensor, "batch"]:
        """
        Predict reachability.

        Args:
            morph: Morphology description.
            pose: Pose as vector encoded.
        Returns:
            Reachability logit.
        """
        latent = self.encoder(morph)[1][0][-1]
        logit = self.decoder(torch.cat([pose, latent], dim=-1)).squeeze(-1)
        return logit

    @torch.inference_mode()
    def predict(self, *args, **kwargs):
        """
        forward with inference mode
        """
        return self.forward(*args, **kwargs)