from dataclasses import dataclass

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

from backbones import (
    AgeSinusoidalEncoding,
    GPTBlock,
    GaussianRBFTimeBasis,
    TimeRoPE,
    TokenAutoDiscretization,
)
from targets import PAD_IDX


@dataclass
class DeepHealthOutput:
    hidden: torch.Tensor
    time_seq: torch.Tensor
    padding_mask: torch.Tensor
    event_len: int


class OtherInfoTokenizer(nn.Module):
    PAD_KIND = 0
    CONT_KIND = 1
    CATE_KIND = 2

    def __init__(
        self,
        n_embd: int,
        n_types: int,
        n_cont_types: int,
        n_categories: int,
        cont_type_ids: list[int],
        n_value_kinds: int = 3,
        n_bins: int = 16,
    ):
        super().__init__()
        if len(cont_type_ids) != n_cont_types:
            raise ValueError(
                "cont_type_ids length must match n_cont_types, got "
                f"{len(cont_type_ids)} vs {n_cont_types}"
            )
        if n_types <= 0:
            raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
        if n_categories <= 0:
            raise ValueError(
                f"n_categories must include PAD and be > 0, got {n_categories}"
            )
        if n_value_kinds <= self.CATE_KIND:
            raise ValueError(
                f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
            )

        self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
        self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
        self.cont_value_encoder = (
            TokenAutoDiscretization(
                n_cont_types=n_cont_types,
                n_bins=n_bins,
                n_embd=n_embd,
            )
            if n_cont_types > 0
            else None
        )
        self.cate_value_emb = nn.Embedding(
            n_categories,
            n_embd,
            padding_idx=0,
        )

        cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
        for idx, type_id in enumerate(cont_type_ids):
            if type_id <= 0 or type_id >= n_types:
                raise ValueError(
                    f"continuous type id {type_id} must be in [1, {n_types})"
                )
            cont_type_index[type_id] = idx
        self.register_buffer(
            "cont_type_index",
            cont_type_index,
            persistent=False,
        )
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.type_emb.weight[0])
        nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.kind_emb.weight[0])
        nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.cate_value_emb.weight[0])

    def forward(
        self,
        other_type: torch.LongTensor,
        other_value: torch.Tensor,
        other_value_kind: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if other_type.shape != other_value.shape:
            raise ValueError(
                "other_type and other_value must have the same shape, got "
                f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
            )
        if other_type.shape != other_value_kind.shape:
            raise ValueError(
                "other_type and other_value_kind must have the same shape, got "
                f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
            )

        other_valid = other_type > 0
        type_emb = self.type_emb(other_type)
        kind_emb = self.kind_emb(other_value_kind)
        value_emb = torch.zeros_like(type_emb)

        cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
        if cont_pos.any():
            if self.cont_value_encoder is None:
                raise ValueError("continuous tokens found but n_cont_types is 0")
            cont_idx = self.cont_type_index[other_type[cont_pos]]
            if (cont_idx < 0).any():
                bad_type = other_type[cont_pos][cont_idx < 0][0].item()
                raise ValueError(
                    f"type_id={bad_type} is marked continuous but is not in "
                    "cont_type_ids"
                )
            value_emb[cont_pos] = self.cont_value_encoder(
                cont_type_idx=cont_idx,
                value=other_value[cont_pos].to(type_emb.dtype),
            )

        cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
        if cate_pos.any():
            cate_id = other_value[cate_pos].long()
            value_emb[cate_pos] = self.cate_value_emb(cate_id)

        out = type_emb + kind_emb + value_emb
        out = out * other_valid.unsqueeze(-1).to(out.dtype)
        return out, other_valid


class DeepHealth(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        n_embd: int,
        n_head: int,
        n_hist_layer: int,
        n_tab_layer: int,
        n_types: int,
        n_cont_types: int,
        n_categories: int,
        cont_type_ids: list[int],
        n_value_kinds: int = 3,
        n_bins: int = 16,
        target_mode: str = "next_token",  # "next_token" or "all_future"
        time_mode: str = "relative",  # "relative" or "absolute"
        dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
        extra_pool_reduce: str = "mean",
        dropout: float = 0.0,
    ):
        super().__init__()
        if target_mode not in ["next_token", "all_future"]:
            raise ValueError(
                "target_mode must be either 'next_token' or 'all_future'")
        if time_mode not in ["relative", "absolute"]:
            raise ValueError(
                "time_mode must be either 'relative' or 'absolute'")
        if dist_mode not in ["exponential", "weibull", "mixed"]:
            raise ValueError(
                "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
        if extra_pool_reduce not in {"mean", "sum"}:
            raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
        self.gender_embedding = nn.Embedding(
            2, n_embd)  # Assuming binary gender
        self.tokenizer = OtherInfoTokenizer(
            n_embd=n_embd,
            n_types=n_types,
            n_cont_types=n_cont_types,
            n_categories=n_categories,
            cont_type_ids=cont_type_ids,
            n_value_kinds=n_value_kinds,
            n_bins=n_bins,
        )
        self.target_mode = target_mode
        self.time_mode = time_mode
        self.dist_mode = dist_mode
        self.extra_pool_reduce = extra_pool_reduce
        self.n_embd = n_embd
        self.vocab_size = vocab_size
        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.token_embedding.weight[0])
        nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
        if dist_mode == "weibull":
            self.rho_head = nn.Linear(n_embd, vocab_size)
            nn.init.zeros_(self.rho_head.weight)
            nn.init.constant_(self.rho_head.bias, 0.5413)

        if dist_mode == "mixed":
            self.death_idx = vocab_size - 1
            self.rho_death_head = nn.Linear(n_embd, 1)
            nn.init.zeros_(self.rho_death_head.weight)
            nn.init.constant_(self.rho_death_head.bias, 0.5413)

        if time_mode == "absolute":
            self.age_encoding = AgeSinusoidalEncoding(n_embd)
            self.blocks = nn.ModuleList([
                GPTBlock(
                    n_embd=n_embd,
                    n_head=n_head,
                    use_time_rope=False,
                    use_rbf_bias=False,
                    mlp_dropout=dropout,
                ) for _ in range(n_hist_layer)
            ])
            self.rope = None
            self.rbf = None
        elif time_mode == "relative":
            self.age_encoding = None
            self.blocks = nn.ModuleList([
                GPTBlock(
                    n_embd=n_embd,
                    n_head=n_head,
                    use_time_rope=True,
                    use_rbf_bias=True,
                    mlp_dropout=dropout,
                ) for _ in range(n_hist_layer)
            ])
            self.rope = TimeRoPE(n_embd // n_head)
            self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)

        self.final_ln = nn.LayerNorm(n_embd)
        self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
        if target_mode == "next_token":
            self.risk_head.weight = self.token_embedding.weight
        self.query_token = nn.Parameter(torch.zeros(n_embd))
        nn.init.normal_(self.query_token, mean=0.0, std=0.02)

    def _make_history_attn_mask(
        self,
        padding_mask: torch.Tensor,
        time_seq: torch.Tensor,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        valid_key = padding_mask[:, None, :]  # (B, 1, L)
        visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None]
        valid = valid_key & visible_by_time
        return torch.zeros(
            valid.shape,
            device=valid.device,
            dtype=dtype,
        ).masked_fill(~valid, -1e4)[:, None, :, :]

    def _pool_other_by_time(
        self,
        h_other: torch.Tensor,
        other_time: torch.Tensor,
        other_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        batch_size, _n_other, n_embd = h_other.shape
        group_counts = [
            int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel())
            for b in range(batch_size)
        ]
        max_groups = max(group_counts, default=0)
        pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
        pooled_time = other_time.new_zeros(batch_size, max_groups)
        pooled_mask = torch.zeros(
            batch_size,
            max_groups,
            dtype=torch.bool,
            device=h_other.device,
        )
        if max_groups == 0:
            return pooled_h, pooled_time, pooled_mask

        for b in range(batch_size):
            valid_time = other_time[b, other_mask[b]]
            if valid_time.numel() == 0:
                continue
            valid_h = h_other[b, other_mask[b]]
            unique_time, inverse = torch.unique(
                valid_time,
                sorted=True,
                return_inverse=True,
            )
            n_groups = unique_time.numel()
            group_h = valid_h.new_zeros(n_groups, n_embd)
            group_h.scatter_add_(
                0,
                inverse[:, None].expand(-1, n_embd),
                valid_h,
            )
            if self.extra_pool_reduce == "mean":
                counts = valid_h.new_zeros(n_groups, 1)
                counts.scatter_add_(
                    0,
                    inverse[:, None],
                    torch.ones_like(valid_h[:, :1]),
                )
                group_h = group_h / counts.clamp_min(1.0)

            pooled_h[b, :n_groups] = group_h
            pooled_time[b, :n_groups] = unique_time
            pooled_mask[b, :n_groups] = True
        return pooled_h, pooled_time, pooled_mask

    def _forward_shared(
        self,
        event_seq: torch.LongTensor,
        time_seq: torch.FloatTensor,
        sex: torch.LongTensor,
        mode: str,
        padding_mask: torch.Tensor | None = None,
        t_query: torch.FloatTensor | None = None,
        other_type: torch.LongTensor | None = None,
        other_value: torch.Tensor | None = None,
        other_value_kind: torch.LongTensor | None = None,
        other_time: torch.FloatTensor | None = None,
        return_output: bool = False,
        **unused_kwargs,
    ) -> torch.Tensor:
        if unused_kwargs:
            unknown = ", ".join(sorted(unused_kwargs))
            raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
        if mode not in {"next_token", "all_future"}:
            raise ValueError("mode must be either 'next_token' or 'all_future'")
        if mode == "all_future" and t_query is None:
            raise ValueError("t_query is required when mode='all_future'")
        if (
            other_type is None
            or other_value is None
            or other_value_kind is None
            or other_time is None
        ):
            raise ValueError(
                "DeepHealth expects other_type, other_value, "
                "other_value_kind, and other_time."
            )

        if padding_mask is None:
            padding_mask = event_seq > PAD_IDX
        else:
            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)

        event_len = event_seq.size(1)
        h_disease = self.token_embedding(event_seq)
        t_disease = time_seq

        if other_time.shape != other_type.shape:
            raise ValueError(
                "other_time must have the same shape as other_type, got "
                f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
            )
        other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
        h_other, other_mask = self.tokenizer(
            other_type=other_type,
            other_value=other_value,
            other_value_kind=other_value_kind,
        )
        h_other = h_other.to(device=event_seq.device)
        other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
        h_other, other_time, other_mask = self._pool_other_by_time(
            h_other=h_other,
            other_time=other_time,
            other_mask=other_mask,
        )

        h_disease = torch.cat([h_disease, h_other], dim=1)
        t_disease = torch.cat([t_disease, other_time], dim=1)
        padding_mask = torch.cat([padding_mask, other_mask], dim=1)
        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)

        if mode == "all_future":
            batch_size = event_seq.size(0)
            query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
            h_disease = torch.cat([h_disease, query], dim=1)
            t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
            query_mask = torch.ones(
                batch_size,
                1,
                dtype=torch.bool,
                device=event_seq.device,
            )
            padding_mask = torch.cat([padding_mask, query_mask], dim=1)

        sex_emb = self.gender_embedding(sex)[:, None, :]
        h_disease = h_disease + sex_emb
        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)

        rope_cache = None
        rbf_cache = None
        if self.time_mode == "absolute":
            h_disease = h_disease + self.age_encoding(t_disease)
            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
        elif self.time_mode == "relative":
            rope_cache = self.rope.precompute_cache(t_disease)
            rbf_cache = self.rbf.precompute_cache(t_disease)

        attn_mask = self._make_history_attn_mask(
            padding_mask=padding_mask,
            time_seq=t_disease,
            dtype=h_disease.dtype,
        )
        for block in self.blocks:
            h_disease = block(
                h_disease,
                rope_cache=rope_cache,
                rbf_cache=rbf_cache,
                attn_mask=attn_mask,
            )
            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)

        h_disease = self.final_ln(h_disease)
        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)

        if mode == "all_future":
            hidden = h_disease[:, -1, :]
            if return_output:
                return DeepHealthOutput(
                    hidden=hidden,
                    time_seq=t_query[:, None],
                    padding_mask=torch.ones(
                        hidden.size(0),
                        1,
                        dtype=torch.bool,
                        device=hidden.device,
                    ),
                    event_len=event_len,
                )
            return hidden
        if return_output:
            return DeepHealthOutput(
                hidden=h_disease,
                time_seq=t_disease,
                padding_mask=padding_mask,
                event_len=event_len,
            )
        return h_disease[:, :event_len, :]

    def forward_next_token(self, **kwargs) -> torch.Tensor:
        return self._forward_shared(mode="next_token", **kwargs)

    def forward_all_future(self, **kwargs) -> torch.Tensor:
        return self._forward_shared(mode="all_future", **kwargs)

    def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor:
        mode = self.target_mode if target_mode is None else target_mode
        return self._forward_shared(mode=mode, **kwargs)

    def calc_risk(self, x: torch.Tensor) -> torch.Tensor:
        return self.risk_head(x)

    def calc_weibull_rho(self, x: torch.Tensor) -> torch.Tensor:
        if self.dist_mode != "weibull":
            raise RuntimeError(
                f"calc_weibull_rho called with dist_mode={self.dist_mode!r}"
            )
        return F.softplus(self.rho_head(x)) + 1e-6

    def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor:
        if self.dist_mode != "mixed":
            raise RuntimeError(
                f"calc_death_rho called with dist_mode={self.dist_mode!r}"
            )
        return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6