"""
DBFT: Deflection-Basin Fusion Transformer for FWD critical strain prediction
=============================================================================
Predicts the two critical pavement responses from Falling Weight Deflectometer
(FWD) measurements:
  - AC          : horizontal tensile strain at the bottom of the asphalt layer
  - Subgrade.2  : vertical compressive strain at the top of the subgrade

Inputs (field-measurable only — no layer moduli, avoiding backcalculation):
  - Deflection basin D0..D1800  (9 geophones, treated as a spatial sequence)
  - Layer thicknesses (Asphalt, Base, Subbase)

Architecture (novel elements):
  1. Continuous sensor-offset positional encoding: Fourier features of the
     physical geophone offset (mm), so the model knows the true basin geometry
     and generalizes to arbitrary sensor arrays.
  2. Physics-informed basin tokens: SCI, BDI, BCI, AREA, AUPP indices are
     embedded as extra tokens alongside raw deflections.
  3. Thickness encoder branch producing (a) memory tokens and (b) FiLM
     (feature-wise linear modulation) parameters applied to the basin tokens.
  4. Strain-query transformer decoder: one learnable query per target strain
     cross-attends over the fused basin+thickness memory (set-to-vector
     decoding), giving per-target attention maps for interpretability.
"""

import math
import numpy as np
import pandas as pd
import torch
import torch.nn as nn

SENSOR_OFFSETS_MM = np.array([0, 200, 300, 450, 600, 900, 1200, 1500, 1800], dtype=np.float32)
DEFLECTION_COLS = ["D0", "D200", "D300", "D450", "D600", "D900", "D1200", "D1500", "D1800"]
THICKNESS_COLS = ["Asphalt", "Base", "Subbase"]
TARGET_COLS = {"AC_tensile": "AC", "Subgrade_compressive": "Subgrade.2"}


# ----------------------------------------------------------------------------
# Physics-informed deflection basin indices
# ----------------------------------------------------------------------------
def basin_indices(D: np.ndarray) -> np.ndarray:
    """D: (N, 9) deflections at SENSOR_OFFSETS_MM. Returns (N, 5) indices."""
    d0, d200, d300, d450, d600, d900, d1200, d1500, d1800 = D.T
    sci = d0 - d300                     # Surface Curvature Index (AC condition)
    bdi = d300 - d600                   # Base Damage Index
    bci = d600 - d900                   # Base Curvature Index (subbase/subgrade)
    # AASHTO AREA parameter (normalized basin area, first 4 sensors ~ 0-600 mm here)
    area = 6.0 * (1 + 2 * d300 / d0 + 2 * d600 / d0 + d900 / d0)
    # AUPP: Area Under Pavement Profile — strongly correlated with AC tensile strain
    aupp = (5 * d0 - 2 * d300 - 2 * d600 - d900) / 2.0
    return np.stack([sci, bdi, bci, area, aupp], axis=1)


# ----------------------------------------------------------------------------
# Continuous positional encoding of physical sensor offsets
# ----------------------------------------------------------------------------
class SensorOffsetEncoding(nn.Module):
    """Fourier features of the physical geophone offset (mm), projected to d_model.
    Unlike integer positional encoding, this respects the true non-uniform
    basin geometry (0,200,300,450,...) and supports arbitrary sensor arrays."""

    def __init__(self, d_model: int, n_freq: int = 16, max_offset: float = 2000.0):
        super().__init__()
        freqs = torch.exp(torch.linspace(math.log(1.0), math.log(max_offset), n_freq))
        self.register_buffer("freqs", freqs)
        self.proj = nn.Linear(2 * n_freq, d_model)

    def forward(self, offsets_mm: torch.Tensor) -> torch.Tensor:
        # offsets_mm: (S,) -> (S, d_model)
        x = offsets_mm.unsqueeze(-1) / self.freqs  # (S, n_freq)
        feats = torch.cat([torch.sin(x), torch.cos(x)], dim=-1)
        return self.proj(feats)


class FiLM(nn.Module):
    """Feature-wise linear modulation of basin tokens by the thickness code."""

    def __init__(self, cond_dim: int, d_model: int):
        super().__init__()
        self.to_gamma_beta = nn.Linear(cond_dim, 2 * d_model)

    def forward(self, tokens: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        gamma, beta = self.to_gamma_beta(cond).chunk(2, dim=-1)  # (B, d_model) each
        return tokens * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)


# ----------------------------------------------------------------------------
# DBFT model
# ----------------------------------------------------------------------------
class DBFT(nn.Module):
    def __init__(
        self,
        d_model: int = 64,
        n_heads: int = 4,
        n_encoder_layers: int = 3,
        n_decoder_layers: int = 2,
        d_ff: int = 128,
        dropout: float = 0.10,
        n_targets: int = 2,
        n_thickness: int = 3,
        n_indices: int = 5,
        use_film: bool = True,
        use_indices: bool = True,
        use_thickness_memory: bool = True,
    ):
        super().__init__()
        self.use_film = use_film
        self.use_indices = use_indices
        self.use_thickness_memory = use_thickness_memory

        # --- basin branch ---
        self.deflection_embed = nn.Linear(1, d_model)
        self.pos_enc = SensorOffsetEncoding(d_model)
        self.register_buffer("offsets", torch.tensor(SENSOR_OFFSETS_MM))

        # physics-informed index tokens
        self.index_embed = nn.Linear(1, d_model)
        self.index_type_embed = nn.Parameter(torch.randn(n_indices, d_model) * 0.02)

        enc_layer = nn.TransformerEncoderLayer(
            d_model, n_heads, d_ff, dropout, activation="gelu",
            batch_first=True, norm_first=True,
        )
        self.basin_encoder = nn.TransformerEncoder(enc_layer, n_encoder_layers)

        # --- thickness branch ---
        self.thickness_encoder = nn.Sequential(
            nn.Linear(n_thickness, d_model), nn.GELU(),
            nn.Linear(d_model, d_model), nn.GELU(),
        )
        self.film = FiLM(d_model, d_model)
        self.thickness_token_proj = nn.Linear(d_model, d_model)

        # --- fusion decoder: learnable strain queries ---
        self.strain_queries = nn.Parameter(torch.randn(n_targets, d_model) * 0.02)
        dec_layer = nn.TransformerDecoderLayer(
            d_model, n_heads, d_ff, dropout, activation="gelu",
            batch_first=True, norm_first=True,
        )
        self.fusion_decoder = nn.TransformerDecoder(dec_layer, n_decoder_layers)

        self.head = nn.Sequential(
            nn.LayerNorm(d_model), nn.Linear(d_model, d_model), nn.GELU(),
            nn.Linear(d_model, 1),
        )

    def forward(self, deflections, indices, thickness):
        """deflections (B,9), indices (B,5), thickness (B,3) — all standardized."""
        B = deflections.shape[0]

        # basin tokens with continuous offset encoding
        tok = self.deflection_embed(deflections.unsqueeze(-1))        # (B,9,d)
        tok = tok + self.pos_enc(self.offsets).unsqueeze(0)           # (B,9,d)

        if self.use_indices:
            idx_tok = self.index_embed(indices.unsqueeze(-1)) + self.index_type_embed
            tok = torch.cat([tok, idx_tok], dim=1)                    # (B,14,d)

        # thickness conditioning
        t_code = self.thickness_encoder(thickness)                    # (B,d)
        if self.use_film:
            tok = self.film(tok, t_code)

        memory = self.basin_encoder(tok)                              # (B,S,d)
        if self.use_thickness_memory:
            memory = torch.cat([memory, self.thickness_token_proj(t_code).unsqueeze(1)], dim=1)

        queries = self.strain_queries.unsqueeze(0).expand(B, -1, -1)  # (B,2,d)
        decoded = self.fusion_decoder(queries, memory)                # (B,2,d)
        return self.head(decoded).squeeze(-1)                        # (B,2)

    @torch.no_grad()
    def attention_map(self, deflections, indices, thickness):
        """Cross-attention weights of each strain query over memory tokens
        (last decoder layer), for interpretability."""
        self.eval()
        B = deflections.shape[0]
        tok = self.deflection_embed(deflections.unsqueeze(-1))
        tok = tok + self.pos_enc(self.offsets).unsqueeze(0)
        if self.use_indices:
            idx_tok = self.index_embed(indices.unsqueeze(-1)) + self.index_type_embed
            tok = torch.cat([tok, idx_tok], dim=1)
        t_code = self.thickness_encoder(thickness)
        if self.use_film:
            tok = self.film(tok, t_code)
        memory = self.basin_encoder(tok)
        if self.use_thickness_memory:
            memory = torch.cat([memory, self.thickness_token_proj(t_code).unsqueeze(1)], dim=1)

        x = self.strain_queries.unsqueeze(0).expand(B, -1, -1)
        attn_out = None
        for i, layer in enumerate(self.fusion_decoder.layers):
            q = layer.norm1(x)
            x = x + layer.dropout1(layer.self_attn(q, q, q, need_weights=False)[0])
            q2 = layer.norm2(x)
            out, w = layer.multihead_attn(q2, memory, memory,
                                          need_weights=True, average_attn_weights=True)
            if i == len(self.fusion_decoder.layers) - 1:
                attn_out = w  # (B, n_targets, S_mem)
            x = x + layer.dropout2(out)
            x = x + layer._ff_block(layer.norm3(x))
        return attn_out


# ----------------------------------------------------------------------------
# Data loading / preprocessing
# ----------------------------------------------------------------------------
def load_dataset(path="strain_result.xlsx"):
    df = pd.read_csv(path) if str(path).endswith(".csv") else pd.read_excel(path)
    D = df[DEFLECTION_COLS].to_numpy(np.float32)
    H = df[THICKNESS_COLS].to_numpy(np.float32)
    Y = df[list(TARGET_COLS.values())].to_numpy(np.float32)
    I = basin_indices(D).astype(np.float32)
    return D, I, H, Y, df


class Standardizer:
    def fit(self, x):
        self.mean = x.mean(axis=0, keepdims=True)
        self.std = x.std(axis=0, keepdims=True) + 1e-8
        return self

    def transform(self, x):
        return (x - self.mean) / self.std

    def inverse(self, x):
        return x * self.std + self.mean