"""WISDM IMU Masked Encoder — self-supervised Transformer for activity recognition.

A pure-PyTorch model for encoding 10-second IMU sensor windows (6-channel
accelerometer + gyroscope @ 20 Hz) into 192-dim representations. Pretrained
with masked prediction, SupCon contrastive learning, and LMM frequency loss
on the WISDM smartphone+smartwatch dataset (18 activity classes).

Usage:
    from modeling_imu_encoder import IMUMaskedEncoder

    model = IMUMaskedEncoder.from_pretrained("NikoKKK/IMU-SelfSupEncoder-v1")
    model.eval()

    # Input: (B, 6, 200) tensor — 6 IMU channels, 200 timesteps
    with torch.no_grad():
        patch_out, intermediates, cls_out, global_freq = model(x)
    # cls_out: (B, 192) — global representation for classification
    # patch_out: (B, 20, 192) — per-window-patch embeddings
    # intermediates: {layer_idx: (B, 20, 192)} — intermediate layer outputs
"""

import math
import json
import os
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

class IMUEncoderConfig:
    """Configuration for IMUMaskedEncoder.

    Attributes:
        n_channels (int): Number of IMU sensor channels (default: 6).
        patch_size (int): Conv-stem patch size in timesteps (default: 10).
        n_patches (int): Number of patches per window (default: 20).
        embed_dim (int): Embedding dimension (default: 192).
        n_layers (int): Number of Transformer encoder layers (default: 4).
        n_heads (int): Number of attention heads (default: 6).
        mlp_ratio (float): MLP hidden/embed_dim ratio (default: 3.0).
        dropout (float): Dropout probability (default: 0.1).
        target_layers (List[int]): Layers collected as intermediate outputs.
    """

    def __init__(
        self,
        n_channels: int = 6,
        patch_size: int = 10,
        n_patches: int = 20,
        embed_dim: int = 192,
        n_layers: int = 4,
        n_heads: int = 6,
        mlp_ratio: float = 3.0,
        dropout: float = 0.1,
        target_layers: Optional[List[int]] = None,
        **kwargs,
    ):
        self.n_channels = n_channels
        self.patch_size = patch_size
        self.n_patches = n_patches
        self.embed_dim = embed_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.mlp_ratio = mlp_ratio
        self.dropout = dropout
        self.target_layers = target_layers or [2, 4]

    @classmethod
    def from_dict(cls, d: dict) -> "IMUEncoderConfig":
        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})

    def to_dict(self) -> dict:
        return {
            "n_channels": self.n_channels,
            "patch_size": self.patch_size,
            "n_patches": self.n_patches,
            "embed_dim": self.embed_dim,
            "n_layers": self.n_layers,
            "n_heads": self.n_heads,
            "mlp_ratio": self.mlp_ratio,
            "dropout": self.dropout,
            "target_layers": self.target_layers,
        }


# ---------------------------------------------------------------------------
# Sub-modules
# ---------------------------------------------------------------------------

class PatchTimeFreqEmbedding(nn.Module):
    """Conv-stem per-patch time + local frequency fusion -> embed_dim tokens."""

    def __init__(self, config: IMUEncoderConfig):
        super().__init__()
        half_dim = config.embed_dim // 2
        self.conv_stem = nn.Conv1d(
            config.n_channels, half_dim, kernel_size=config.patch_size,
            stride=config.patch_size, bias=False,
        )
        n_freq_bins = config.patch_size // 2 + 1
        self.freq_proj = nn.Linear(config.n_channels * n_freq_bins, half_dim)
        self.pos_embed = nn.Parameter(
            torch.randn(1, config.n_patches + 1, config.embed_dim) * 0.02
        )
        self.fusion = nn.Linear(config.embed_dim, config.embed_dim)
        self.norm = nn.LayerNorm(config.embed_dim)
        self.patch_size = config.patch_size
        self.n_patches = config.n_patches

        # FFT over full 200-timestep signal: rfft(200) → 101 bins, then .mean(dim=1)
        n_global_freq_bins = config.patch_size * config.n_patches // 2 + 1
        self.global_freq_proj = nn.Linear(n_global_freq_bins, config.embed_dim, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        B, C, T = x.shape
        time_feat = self.conv_stem(x)

        x_patch = x.view(B, C, self.n_patches, self.patch_size)
        x_patch_centered = x_patch - x_patch.mean(dim=-1, keepdim=True)
        freq_mag = torch.abs(torch.fft.rfft(x_patch_centered, dim=-1))
        n_freq = freq_mag.size(-1)
        freq_flat = freq_mag.permute(0, 2, 1, 3).reshape(B, self.n_patches, C * n_freq)
        freq_feat = self.freq_proj(freq_flat).transpose(1, 2)

        combined = torch.cat([time_feat, freq_feat], dim=1)
        tokens = self.fusion(combined.transpose(1, 2))

        full_freq = torch.abs(torch.fft.rfft(x, dim=-1)).mean(dim=1)
        global_freq = self.global_freq_proj(full_freq).unsqueeze(1)

        tokens = tokens + self.pos_embed[:, :self.n_patches, :]
        global_freq = global_freq + self.pos_embed[:, self.n_patches:self.n_patches + 1, :]
        tokens = self.norm(tokens)
        global_freq = self.norm(global_freq)
        return tokens, global_freq


class TransformerEncoder(nn.Module):
    """Standard Transformer encoder, collects intermediate layer outputs."""

    def __init__(self, config: IMUEncoderConfig):
        super().__init__()
        self.n_layers = config.n_layers
        self.collect_layers = config.target_layers
        self.cls_token = nn.Parameter(torch.randn(1, 1, config.embed_dim) * 0.02)
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=config.embed_dim,
                nhead=config.n_heads,
                dim_feedforward=int(config.embed_dim * config.mlp_ratio),
                dropout=config.dropout,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            for _ in range(config.n_layers)
        ])
        self.final_norm = nn.LayerNorm(config.embed_dim)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor):
        B = x.shape[0]
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, self.dropout(x)], dim=1)

        intermediates = {}
        for i, layer in enumerate(self.layers):
            x = layer(x)
            layer_idx = i + 1
            if layer_idx in self.collect_layers:
                intermediates[layer_idx] = self.final_norm(x[:, 1:, :])

        x = self.final_norm(x)
        cls_out = x[:, 0, :]
        patch_out = x[:, 1:, :]
        if self.n_layers in self.collect_layers:
            intermediates[self.n_layers] = patch_out
        return patch_out, intermediates, cls_out


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------

class IMUMaskedEncoder(nn.Module):
    """Self-supervised encoder for IMU-based human activity recognition.

    Input:  (B, 6, 200) — 6-channel IMU window (accel_x/y/z, gyro_x/y/z) @ 20 Hz.
    Output: patch_out (B, 20, 192), intermediates dict, cls_out (B, 192),
            global_freq (B, 1, 192).

    The CLS token (`cls_out`) is the primary global representation for
    downstream classification. Intermediate layer outputs can be used for
    multi-level feature extraction or distillation.

    Usage:
        model = IMUMaskedEncoder.from_pretrained("NikoKKK/IMU-SelfSupEncoder-v1")
        model.eval()

        x = torch.randn(8, 6, 200)  # 8 windows of 10 seconds each
        with torch.no_grad():
            patch_out, intermediates, cls_out, global_freq = model(x)

        print(cls_out.shape)  # (8, 192)
    """

    _HUB_URL = "https://huggingface.co/NikoKKK/IMU-SelfSupEncoder-v1"

    def __init__(self, config: IMUEncoderConfig):
        super().__init__()
        self.config = config
        self.embed = PatchTimeFreqEmbedding(config)
        self.transformer = TransformerEncoder(config)
        self.mask_token = nn.Parameter(torch.randn(1, 1, config.embed_dim) * 0.02)

    @classmethod
    def from_pretrained(cls, model_id: str = "NikoKKK/IMU-SelfSupEncoder-v1",
                        force_download: bool = False) -> "IMUMaskedEncoder":
        """Load pretrained model from Hugging Face Hub or local path.

        Args:
            model_id: Hugging Face model ID (e.g. "NikoKKK/IMU-SelfSupEncoder-v1")
                      or local directory path.
            force_download: Force re-download from Hub.

        Returns:
            IMUMaskedEncoder with pretrained weights loaded.
        """
        from huggingface_hub import hf_hub_download
        from safetensors.torch import load_file

        # Determine if model_id is a local path
        if os.path.isdir(model_id):
            config_path = os.path.join(model_id, "config.json")
            weights_path = os.path.join(model_id, "model.safetensors")
        else:
            config_path = hf_hub_download(
                model_id, "config.json", force_download=force_download,
            )
            weights_path = hf_hub_download(
                model_id, "model.safetensors", force_download=force_download,
            )

        with open(config_path, "r") as f:
            config_dict = json.load(f)
        config = IMUEncoderConfig.from_dict(config_dict)

        model = cls(config)
        state_dict = load_file(weights_path)
        model.load_state_dict(state_dict, strict=True)
        return model

    def forward(self, x: torch.Tensor, mask_matrix=None):
        """Forward pass.

        Args:
            x: (B, C, T) tensor — IMU sensor window.
               Expected C=6 (accel_x,y,z + gyro_x,y,z), T=200 @ 20 Hz.
            mask_matrix: Optional (B, N_patches) boolean tensor. When given,
                         masked positions are replaced with [MASK] tokens.
                         Used during pretraining only.

        Returns:
            patch_out:    (B, N_patches, embed_dim)
            intermediates: {layer_idx: (B, N_patches, embed_dim)}
            cls_out:      (B, embed_dim) — global representation
            global_freq:  (B, 1, embed_dim) — global frequency token
        """
        B, C, T = x.shape
        tokens, global_freq = self.embed(x)

        if mask_matrix is not None:
            mask_tok = self.mask_token.expand(B, tokens.size(1), -1)
            tokens = torch.where(mask_matrix.unsqueeze(-1), mask_tok, tokens)

        input_tokens = torch.cat([global_freq, tokens], dim=1)
        full_out, intermediates, cls_out = self.transformer(input_tokens)

        # Remove global_freq position from patch outputs
        patch_out = full_out[:, 1:, :]
        trimmed_intermediates = {k: v[:, 1:, :] for k, v in intermediates.items()}

        return patch_out, trimmed_intermediates, cls_out, global_freq

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Extract CLS token embedding for downstream tasks.

        Args:
            x: (B, 6, 200) IMU sensor window.

        Returns:
            (B, embed_dim) global representation.
        """
        _, _, cls_out, _ = self.forward(x)
        return cls_out

    def encode_with_layers(self, x: torch.Tensor) -> Dict[int, torch.Tensor]:
        """Extract multi-level intermediate representations.

        Args:
            x: (B, 6, 200) IMU sensor window.

        Returns:
            {layer_idx: (B, N_patches, embed_dim)} for configured target layers.
        """
        _, intermediates, _, _ = self.forward(x)
        return intermediates


# For HF AutoModel compatibility
IMUMaskedEncoder.config_class = IMUEncoderConfig