"""Learned dense-to-local+algebra overlap blocks with physical export."""

from __future__ import annotations

from dataclasses import dataclass

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

from strata.modeling.algebra import HardConcreteGate
from strata.modeling.compose.lm_adapter import AlgebraGraphReadAdapter, SparseAlgebraGraphRead
from strata.modeling.ph_pat.config import PHPATConfig
from strata.modeling.ph_pat.replacement_block import (
    ChunkedLocalCausalAttention,
    DenseGlobalBlock,
    DenseGlobalCausalAttention,
)


class AlgebraOverlapBlock(nn.Module):
    """Exact dense endpoint and trainable local+graph endpoint for one block."""

    def __init__(
        self,
        dense_block: DenseGlobalBlock,
        config: PHPATConfig,
        *,
        gamma_max: float = 0.02,
    ) -> None:
        super().__init__()
        self.attn_norm = dense_block.attn_norm
        self.global_attention = dense_block.attention
        self.local_attention = ChunkedLocalCausalAttention(config)
        self.local_attention.load_state_dict(self.global_attention.state_dict(), strict=True)
        self.ffn_norm = dense_block.ffn_norm
        self.ffn = dense_block.ffn
        self.dropout = dense_block.dropout
        self.graph_adapter = AlgebraGraphReadAdapter(config.d_model, config.d_model, gamma_max=gamma_max)
        self.global_gate = HardConcreteGate(1, initial_probability=0.99)
        reference = self.global_attention.qkv.weight
        self.local_attention.to(device=reference.device, dtype=reference.dtype)
        self.graph_adapter.to(device=reference.device, dtype=reference.dtype)
        self.global_gate.to(device=reference.device)

    def forward(
        self,
        hidden: torch.Tensor,
        attention_mask: torch.Tensor,
        *,
        graph_values: torch.Tensor | None = None,
        graph_reliability: torch.Tensor | None = None,
        sparse_graph_read: SparseAlgebraGraphRead | None = None,
        graph_enabled: bool = False,
        gate_mode: str = "dense",
        gate_override: torch.Tensor | float | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        normalized = self.attn_norm(hidden)
        if gate_override is not None:
            gate = torch.as_tensor(gate_override, device=hidden.device, dtype=hidden.dtype)
            if gate.numel() != 1 or bool((gate < 0).any() or (gate > 1).any()):
                raise ValueError("gate_override must be one scalar in [0, 1]")
            gate = gate.reshape(())
            if float(gate.detach()) == 1.0:
                attention = self.global_attention(normalized, attention_mask)
            elif float(gate.detach()) == 0.0:
                attention = self.local_attention(normalized, attention_mask)
            else:
                local = self.local_attention(normalized, attention_mask)
                global_value = self.global_attention(normalized, attention_mask)
                attention = local + gate * (global_value - local)
        elif gate_mode == "dense":
            attention = self.global_attention(normalized, attention_mask)
            gate = hidden.new_tensor(1.0)
        elif gate_mode == "local":
            attention = self.local_attention(normalized, attention_mask)
            gate = hidden.new_tensor(0.0)
        elif gate_mode in ("sample", "expected"):
            local = self.local_attention(normalized, attention_mask)
            global_value = self.global_attention(normalized, attention_mask)
            gate = self.global_gate(sample=gate_mode == "sample")[0].to(hidden.dtype)
            attention = local + gate * (global_value - local)
        else:
            raise ValueError("gate_mode must be dense, local, sample, or expected")
        hidden = hidden + self.dropout(attention)
        if graph_enabled:
            if sparse_graph_read is not None:
                hidden, _update, _relative = self.graph_adapter.forward_sparse(
                    hidden, sparse_graph_read
                )
            else:
                if graph_values is None or graph_reliability is None:
                    raise ValueError("enabled graph reads require values and reliability")
                hidden, _update, _relative = self.graph_adapter(
                    hidden, graph_values, graph_reliability
                )
        hidden = hidden + self.dropout(self.ffn(self.ffn_norm(hidden)))
        return hidden, gate

    def export_local(self) -> "ExportedAlgebraBlock":
        return ExportedAlgebraBlock(self)

    def restore_global(self, config: PHPATConfig) -> DenseGlobalBlock:
        """Restore the original dense block without retaining local/chart modules."""

        restored = DenseGlobalBlock(config)
        restored.attn_norm = self.attn_norm
        restored.attention = self.global_attention
        restored.ffn_norm = self.ffn_norm
        restored.ffn = self.ffn
        restored.dropout = self.dropout
        return restored


class ExportedAlgebraBlock(nn.Module):
    """Physical local+graph block; no dense Q/K/V module remains."""

    def __init__(self, source: AlgebraOverlapBlock) -> None:
        super().__init__()
        self.attn_norm = source.attn_norm
        self.local_attention = source.local_attention
        self.ffn_norm = source.ffn_norm
        self.ffn = source.ffn
        self.dropout = source.dropout
        self.graph_adapter = source.graph_adapter

    def forward(
        self,
        hidden: torch.Tensor,
        attention_mask: torch.Tensor,
        *,
        graph_values: torch.Tensor | None = None,
        graph_reliability: torch.Tensor | None = None,
        sparse_graph_read: SparseAlgebraGraphRead | None = None,
        graph_enabled: bool = False,
    ) -> torch.Tensor:
        hidden = hidden + self.dropout(self.local_attention(self.attn_norm(hidden), attention_mask))
        if graph_enabled:
            if self.graph_adapter is not None:
                if sparse_graph_read is not None:
                    hidden, _update, _relative = self.graph_adapter.forward_sparse(
                        hidden, sparse_graph_read
                    )
                else:
                    if graph_values is None or graph_reliability is None:
                        raise ValueError("enabled graph reads require values and reliability")
                    hidden, _update, _relative = self.graph_adapter(
                        hidden, graph_values, graph_reliability
                    )
        return hidden + self.dropout(self.ffn(self.ffn_norm(hidden)))


@dataclass(frozen=True, slots=True)
class PrunableLMOutput:
    logits: torch.Tensor
    loss: torch.Tensor | None
    hidden_states: torch.Tensor
    global_gates: torch.Tensor


class PrunableAlgebraLM(nn.Module):
    """Install overlap blocks without altering unselected dense blocks."""

    def __init__(
        self,
        base_model: nn.Module,
        config: PHPATConfig,
        replacement_layers: tuple[int, ...],
        *,
        gamma_max: float = 0.02,
    ) -> None:
        super().__init__()
        self.base_model = base_model
        self.config = config
        self.replacement_layers = tuple(sorted(int(index) for index in replacement_layers))
        if len(set(self.replacement_layers)) != len(self.replacement_layers):
            raise ValueError("replacement layers must be unique")
        for index in self.replacement_layers:
            block = self.base_model.blocks[index]
            if not isinstance(block, DenseGlobalBlock):
                raise TypeError(f"layer {index} is not a dense global block")
            self.base_model.blocks[index] = AlgebraOverlapBlock(block, config, gamma_max=gamma_max)

    def install_overlap_layers(
        self,
        layer_indices: tuple[int, ...],
        *,
        gamma_max: float = 0.02,
    ) -> tuple[int, ...]:
        """Install candidate overlap paths while preserving prior physical exports."""

        requested = tuple(sorted(set(int(index) for index in layer_indices)))
        installed = []
        for index in requested:
            block = self.base_model.blocks[index]
            if isinstance(block, (AlgebraOverlapBlock, ExportedAlgebraBlock)):
                continue
            if not isinstance(block, DenseGlobalBlock):
                raise TypeError(f"layer {index} is not a dense global block")
            self.base_model.blocks[index] = AlgebraOverlapBlock(
                block,
                self.config,
                gamma_max=gamma_max,
            )
            installed.append(index)
        self.replacement_layers = tuple(sorted(set(self.replacement_layers).union(installed)))
        return tuple(installed)

    def restore_global_layers(self, layer_indices: tuple[int, ...]) -> tuple[int, ...]:
        """Remove unselected overlap paths and retain their exact dense branches."""

        restored = []
        for index in sorted(set(int(value) for value in layer_indices)):
            block = self.base_model.blocks[index]
            if isinstance(block, ExportedAlgebraBlock):
                raise ValueError(f"cannot restore physically exported layer {index}")
            if not isinstance(block, AlgebraOverlapBlock):
                continue
            self.base_model.blocks[index] = block.restore_global(self.config)
            restored.append(index)
        if restored:
            removed = set(restored)
            self.replacement_layers = tuple(
                index for index in self.replacement_layers if index not in removed
            )
        return tuple(restored)

    def export_layers(self, layer_indices: tuple[int, ...]) -> tuple[int, ...]:
        """Physically export an explicit set of overlap branches to local blocks."""

        exported = []
        for index in sorted(set(int(value) for value in layer_indices)):
            block = self.base_model.blocks[index]
            if not isinstance(block, AlgebraOverlapBlock):
                raise TypeError(f"layer {index} is not an overlap candidate")
            self.base_model.blocks[index] = block.export_local()
            exported.append(index)
        return tuple(exported)

    def forward_hidden(
        self,
        input_ids: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
        graph_values: torch.Tensor | None = None,
        graph_reliability: torch.Tensor | None = None,
        sparse_graph_read: SparseAlgebraGraphRead | None = None,
        graph_enabled: bool = False,
        graph_layer_indices: tuple[int, ...] | None = None,
        gate_mode: str = "dense",
        global_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
        else:
            attention_mask = attention_mask.to(torch.bool)
        positions = torch.arange(input_ids.shape[1], device=input_ids.device)
        hidden = self.base_model.dropout(
            self.base_model.token_embeddings(input_ids)
            + self.base_model.position_embeddings(positions).unsqueeze(0)
        )
        gates = []
        replacement_index = 0
        active_graph_layers = None if graph_layer_indices is None else set(graph_layer_indices)
        for layer_index, block in enumerate(self.base_model.blocks):
            if isinstance(block, AlgebraOverlapBlock):
                override = None
                if global_mask is not None:
                    if global_mask.shape != (len(self.replacement_layers),):
                        raise ValueError("global_mask must have one value per replacement layer")
                    override = global_mask[replacement_index]
                hidden, gate = block(
                    hidden,
                    attention_mask,
                    graph_values=graph_values,
                    graph_reliability=graph_reliability,
                    sparse_graph_read=sparse_graph_read,
                    graph_enabled=graph_enabled and (
                        active_graph_layers is None or layer_index in active_graph_layers
                    ),
                    gate_mode=gate_mode,
                    gate_override=override,
                )
                gates.append(gate.reshape(1))
                replacement_index += 1
            elif isinstance(block, ExportedAlgebraBlock):
                hidden = block(
                    hidden,
                    attention_mask,
                    graph_values=graph_values,
                    graph_reliability=graph_reliability,
                    sparse_graph_read=sparse_graph_read,
                    graph_enabled=graph_enabled and (
                        active_graph_layers is None or layer_index in active_graph_layers
                    ),
                )
                gates.append(hidden.new_zeros(1))
                replacement_index += 1
            else:
                hidden, _read = block(hidden, attention_mask)
        return hidden, attention_mask, torch.cat(gates) if gates else hidden.new_empty(0)

    def forward(
        self,
        input_ids: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        graph_values: torch.Tensor | None = None,
        graph_reliability: torch.Tensor | None = None,
        sparse_graph_read: SparseAlgebraGraphRead | None = None,
        graph_enabled: bool = False,
        graph_layer_indices: tuple[int, ...] | None = None,
        gate_mode: str = "dense",
        global_mask: torch.Tensor | None = None,
    ) -> PrunableLMOutput:
        hidden, _mask, gates = self.forward_hidden(
            input_ids,
            attention_mask=attention_mask,
            graph_values=graph_values,
            graph_reliability=graph_reliability,
            sparse_graph_read=sparse_graph_read,
            graph_enabled=graph_enabled,
            graph_layer_indices=graph_layer_indices,
            gate_mode=gate_mode,
            global_mask=global_mask,
        )
        hidden = self.base_model.final_norm(hidden)
        logits = self.base_model.lm_head(hidden)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].contiguous().view(-1, logits.shape[-1]),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )
        return PrunableLMOutput(logits, loss, hidden, gates)

    def expected_global_branches(self) -> torch.Tensor:
        values = [
            block.global_gate.expected_l0()[0]
            for block in self.base_model.blocks
            if isinstance(block, AlgebraOverlapBlock)
        ]
        return torch.stack(values) if values else next(self.parameters()).new_empty(0)

    def deterministic_global_branches(self) -> torch.Tensor:
        values = [
            block.global_gate(sample=False)[0]
            for block in self.base_model.blocks
            if isinstance(block, AlgebraOverlapBlock)
        ]
        return torch.stack(values) if values else next(self.parameters()).new_empty(0)

    def export_zero_gates(self, *, threshold: float = 0.5) -> tuple[int, ...]:
        exported = []
        for index, block in enumerate(self.base_model.blocks):
            if isinstance(block, AlgebraOverlapBlock) and not bool(
                block.global_gate.export_mask(threshold=threshold)[0]
            ):
                self.base_model.blocks[index] = block.export_local()
                exported.append(index)
        return tuple(exported)

    def dense_modules_in_replacement_layers(self) -> tuple[int, ...]:
        found = []
        for index in self.replacement_layers:
            if any(isinstance(module, DenseGlobalCausalAttention) for module in self.base_model.blocks[index].modules()):
                found.append(index)
        return tuple(found)

    def physically_removed_layers(self) -> tuple[int, ...]:
        return tuple(
            index
            for index, block in enumerate(self.base_model.blocks)
            if isinstance(block, ExportedAlgebraBlock)
        )

    def active_graph_layers(self, *, tolerance: float = 0.0) -> tuple[int, ...]:
        """Return physical graph blocks whose BRR gate is not identically zero."""

        return tuple(
            index
            for index, block in enumerate(self.base_model.blocks)
            if isinstance(block, ExportedAlgebraBlock)
            and block.graph_adapter is not None
            and bool(block.graph_adapter.brr.gate_parameter.detach().abs().max() > tolerance)
        )

    def strip_inactive_graph_adapters(self, *, tolerance: float = 0.0) -> tuple[int, ...]:
        """Physically remove adapters whose BRR contribution is exactly zero."""

        stripped = []
        for index, block in enumerate(self.base_model.blocks):
            if not isinstance(block, ExportedAlgebraBlock) or block.graph_adapter is None:
                continue
            gate = block.graph_adapter.brr.gate_parameter.detach().abs().max()
            if bool(gate <= tolerance):
                block.graph_adapter = None
                stripped.append(index)
        return tuple(stripped)


__all__ = [
    "AlgebraOverlapBlock",
    "ExportedAlgebraBlock",
    "PrunableAlgebraLM",
    "PrunableLMOutput",
]