"""
MemoryCLIP-Seq — Memory-Extended CLIP-L/14 with Sequence Output

Expands MemoryCLIPModel with a sequence reconstruction head that produces
(B, 77, 768) output for direct use in SD/SDXL UNet cross-attention.

Architecture (after all segments processed):
  1. Existing memory system → memory_tokens (B, 8, 768) + anchors (B, N, 768)
  2. Collect per-segment content tokens → segment_tokens (B, total_content, 768)
  3. SequenceReconstructor:
     - 77 learned query tokens (like Q-Former)
     - Cross-attend to: cat(memory_tokens, anchors, segment_content)
     - Output: (B, 77, 768) — every position informed by full document
  4. Pooled output preserved from existing system (backward compatible)

Training targets:
  - Pooled: InfoNCE(pooled, ModernBERT_pooled) — same as v1 (unchanged)
  - Sequence: CLIP's own last_hidden_state on the same caption (truncated to 77)
    The UNet was trained on CLIP's sequence distribution — the reconstructor
    must learn to produce sequences in that same distribution, but enriched
    with the full context the memory system captured beyond 77 tokens.

At inference: produces both pooled (768,) and sequence (77, 768)
"""

import math
from typing import Optional, List, Tuple, Dict, Any
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from transformers import PretrainedConfig, PreTrainedModel, CLIPTextModel, CLIPTokenizer
from transformers.modeling_outputs import BaseModelOutput


# ══════════════════════════════════════════════════════════════════
# CONFIG
# ══════════════════════════════════════════════════════════════════

class MemoryCLIPSeqConfig(PretrainedConfig):
    model_type = "memory_clip_seq"

    def __init__(
        self,
        # CLIP text encoder
        clip_model="openai/clip-vit-large-patch14",
        clip_hidden=768,
        clip_layers=12,
        clip_max_tokens=77,
        freeze_clip=True,

        # Memory system (unchanged from v1)
        n_memory_tokens=8,
        bank_size=64,
        anchor_dim=768,
        n_bank_heads=8,
        bank_cross_layers=2,
        gate_type="gru",
        extract_layers=(1, 3, 5, 7, 9, 11),
        layer_fusion="learned",
        max_content_tokens=18,
        segment_overlap=4,
        max_segments=32,
        cv_target=0.20,

        # Sequence reconstruction (NEW)
        sequence_output=True,
        sequence_len=77,              # output length — matches SD expectation
        sequence_recon_layers=2,
        sequence_recon_heads=8,
        collect_content_tokens=True,  # store per-segment fused tokens
        max_content_positions=256,    # max total content tokens across segments

        # Teacher projection (for training)
        teacher_model="answerdotai/ModernBERT-large",
        teacher_hidden=1024,

        **kwargs,
    ):
        # Memory system params
        self.clip_model = clip_model
        self.clip_hidden = clip_hidden
        self.clip_layers = clip_layers
        self.clip_max_tokens = clip_max_tokens
        self.freeze_clip = freeze_clip
        self.n_memory_tokens = n_memory_tokens
        self.bank_size = bank_size
        self.anchor_dim = anchor_dim
        self.n_bank_heads = n_bank_heads
        self.bank_cross_layers = bank_cross_layers
        self.gate_type = gate_type
        self.extract_layers = tuple(extract_layers)
        self.layer_fusion = layer_fusion
        self.max_content_tokens = max_content_tokens
        self.segment_overlap = segment_overlap
        self.max_segments = max_segments
        self.cv_target = cv_target

        # Sequence reconstruction
        self.sequence_output = sequence_output
        self.sequence_len = sequence_len
        self.sequence_recon_layers = sequence_recon_layers
        self.sequence_recon_heads = sequence_recon_heads
        self.collect_content_tokens = collect_content_tokens
        self.max_content_positions = max_content_positions

        # Teacher
        self.teacher_model = teacher_model
        self.teacher_hidden = teacher_hidden

        super().__init__(**kwargs)

    @property
    def n_extract_layers(self):
        return len(self.extract_layers)

    @property
    def depth_profile_dim(self):
        return self.n_extract_layers * self.clip_hidden

    @property
    def effective_context(self):
        return self.max_segments * self.max_content_tokens


# ══════════════════════════════════════════════════════════════════
# EXISTING COMPONENTS (unchanged from memory_clip.py)
# ══════════════════════════════════════════════════════════════════

class GeometricMemoryBank(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.max_size = config.bank_size
        self.dim = config.anchor_dim
        self.depth_compressor = nn.Sequential(
            nn.Linear(config.depth_profile_dim, config.clip_hidden * 2),
            nn.GELU(),
            nn.LayerNorm(config.clip_hidden * 2),
            nn.Linear(config.clip_hidden * 2, config.anchor_dim),
        )
        self.temporal_proj = nn.Linear(1, config.anchor_dim, bias=False)
        self.cross_attn = nn.ModuleList([
            nn.MultiheadAttention(config.clip_hidden, config.n_bank_heads,
                                  batch_first=True, dropout=0.1)
            for _ in range(config.bank_cross_layers)
        ])
        self.cross_norms = nn.ModuleList([
            nn.LayerNorm(config.clip_hidden)
            for _ in range(config.bank_cross_layers)
        ])
        self.cross_ffns = nn.ModuleList([
            nn.Sequential(
                nn.Linear(config.clip_hidden, config.clip_hidden * 2),
                nn.GELU(),
                nn.Linear(config.clip_hidden * 2, config.clip_hidden))
            for _ in range(config.bank_cross_layers)
        ])
        self.ffn_norms = nn.ModuleList([
            nn.LayerNorm(config.clip_hidden)
            for _ in range(config.bank_cross_layers)
        ])

    def init_bank(self, batch_size, device):
        return {"anchors": torch.zeros(batch_size, 0, self.dim, device=device),
                "n_written": 0}

    def write(self, bank, depth_cls, segment_idx=0):
        B = depth_cls.shape[0]
        anchor = self.depth_compressor(depth_cls.reshape(B, -1))
        anchor = F.normalize(anchor, dim=-1)
        t = torch.tensor([[segment_idx]], dtype=anchor.dtype, device=anchor.device)
        anchor = anchor + 0.1 * self.temporal_proj(t / max(self.max_size, 1))
        anchor = F.normalize(anchor, dim=-1)
        anchors = torch.cat([bank["anchors"], anchor.unsqueeze(1)], dim=1)
        if anchors.shape[1] > self.max_size:
            anchors = anchors[:, -self.max_size:]
        return {"anchors": anchors, "n_written": bank["n_written"] + 1,
                "live_anchor": anchor}

    def read(self, memory_tokens, bank):
        anchors = bank["anchors"]
        if anchors.shape[1] == 0:
            return memory_tokens
        x = memory_tokens
        for attn, norm, ffn, ffn_norm in zip(
            self.cross_attn, self.cross_norms, self.cross_ffns, self.ffn_norms):
            residual = x
            x, _ = attn(norm(x), anchors, anchors)
            x = residual + x
            residual = x
            x = residual + ffn(ffn_norm(x))
        return x


class DeltaMemoryGate(nn.Module):
    def __init__(self, config):
        super().__init__()
        H = config.clip_hidden
        self.reset_proj = nn.Linear(H * 2, H)
        self.update_proj = nn.Linear(H * 2, H)
        self.candidate_proj = nn.Linear(H * 2, H)
        self.norm = nn.LayerNorm(H)

    def forward(self, old, new):
        cat = torch.cat([old, new], dim=-1)
        r = torch.sigmoid(self.reset_proj(cat))
        z = torch.sigmoid(self.update_proj(cat))
        h = torch.tanh(self.candidate_proj(torch.cat([r * old, new], dim=-1)))
        return self.norm(z * old + (1 - z) * h)


class LayerFusion(nn.Module):
    def __init__(self, config):
        super().__init__()
        n = config.n_extract_layers
        self.weights = nn.Parameter(torch.ones(n) / n)
        self.proj = nn.Linear(config.clip_hidden, config.clip_hidden)
        self.norm = nn.LayerNorm(config.clip_hidden)

    def forward(self, layer_outputs):
        w = F.softmax(self.weights, dim=0)
        stacked = torch.stack(layer_outputs)
        fused = (stacked * w.view(-1, 1, 1, 1)).sum(0)
        return self.norm(self.proj(fused))


class TeacherProjector(nn.Module):
    def __init__(self, student_dim, teacher_dim):
        super().__init__()
        self.proj = nn.Linear(student_dim, teacher_dim, bias=True)

    def forward(self, x):
        return self.proj(x)


# ══════════════════════════════════════════════════════════════════
# NEW: SEQUENCE RECONSTRUCTOR
# ══════════════════════════════════════════════════════════════════

class SequenceReconstructor(nn.Module):
    """
    Learned query tokens that cross-attend into the full memory state
    to produce a fixed-length output sequence.

    Similar to Q-Former / Perceiver: N queries → N output positions,
    each informed by the full document through memory + anchors + content.

    Input context (K, V):
        cat(memory_tokens, bank_anchors, [content_tokens])
        = (B, 8 + N_segs + [content_len], 768)

    Output:
        (B, sequence_len, 768) — e.g. (B, 77, 768) for SD compatibility
    """
    def __init__(self, config):
        super().__init__()
        H = config.clip_hidden
        S = config.sequence_len

        # Learned queries — one per output position
        self.query_tokens = nn.Parameter(torch.randn(1, S, H) * 0.02)

        # Positional encoding for queries (so the model knows position 0 vs 76)
        self.query_pos = nn.Parameter(torch.randn(1, S, H) * 0.02)

        # Cross-attention layers: queries attend to memory context
        self.layers = nn.ModuleList()
        for _ in range(config.sequence_recon_layers):
            self.layers.append(nn.ModuleDict({
                "cross_attn": nn.MultiheadAttention(
                    H, config.sequence_recon_heads,
                    batch_first=True, dropout=0.1),
                "cross_norm": nn.LayerNorm(H),
                "self_attn": nn.MultiheadAttention(
                    H, config.sequence_recon_heads,
                    batch_first=True, dropout=0.1),
                "self_norm": nn.LayerNorm(H),
                "ffn": nn.Sequential(
                    nn.Linear(H, H * 4),
                    nn.GELU(),
                    nn.Linear(H * 4, H)),
                "ffn_norm": nn.LayerNorm(H),
            }))

        self.out_norm = nn.LayerNorm(H)

    def forward(self, context):
        """
        context: (B, context_len, 768) — concatenated memory + anchors + content
        returns: (B, sequence_len, 768)
        """
        B = context.shape[0]
        queries = (self.query_tokens + self.query_pos).expand(B, -1, -1)

        for layer in self.layers:
            # Cross-attend to memory context
            residual = queries
            q_normed = layer["cross_norm"](queries)
            queries, _ = layer["cross_attn"](q_normed, context, context)
            queries = residual + queries

            # Self-attend among output positions
            residual = queries
            q_normed = layer["self_norm"](queries)
            queries, _ = layer["self_attn"](q_normed, q_normed, q_normed)
            queries = residual + queries

            # FFN
            residual = queries
            queries = residual + layer["ffn"](layer["ffn_norm"](queries))

        return self.out_norm(queries)


# ══════════════════════════════════════════════════════════════════
# SEGMENTATION
# ══════════════════════════════════════════════════════════════════

def segment_text(text, clip_tokenizer, max_content=18, overlap=4, max_segments=32):
    full_tokens = clip_tokenizer.encode(text, add_special_tokens=False)
    segments = []
    stride = max_content - overlap
    pos = 0
    while pos < len(full_tokens) and len(segments) < max_segments:
        end = min(pos + max_content, len(full_tokens))
        chunk = full_tokens[pos:end]
        sos = clip_tokenizer.bos_token_id or 49406
        eos = clip_tokenizer.eos_token_id or 49407
        input_ids = [sos] + chunk + [eos]
        n_real = len(chunk) + 2  # SOS + content + EOS
        n_pad = 77 - len(input_ids)
        if n_pad > 0:
            input_ids = input_ids + [0] * n_pad
        else:
            input_ids = input_ids[:77]
        mask = [1] * min(n_real, 77) + [0] * max(n_pad, 0)
        mask = mask[:77]
        segments.append({
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(mask, dtype=torch.long),
            "n_content": len(chunk),  # track how many real content tokens
        })
        if end >= len(full_tokens):
            break
        pos += stride
    return segments


# ══════════════════════════════════════════════════════════════════
# MODEL
# ══════════════════════════════════════════════════════════════════

class MemoryCLIPSeqModel(PreTrainedModel):
    """
    Memory-Extended CLIP-L/14 with Sequence Output.

    Produces both:
      - pooled: (B, 768) — backward compatible with v1
      - sequence: (B, 77, 768) — for SD/SDXL cross-attention

    The sequence head uses learned queries (Q-Former style) that
    cross-attend into the full memory state after all segments
    are processed. Every output position is informed by the
    complete document context.
    """
    config_class = MemoryCLIPSeqConfig
    supports_gradient_checkpointing = False

    @classmethod
    def _can_set_experts_implementation(cls):
        return False

    def __init__(self, config):
        super().__init__(config)

        # ── Existing memory system (unchanged, loadable from v1 weights) ──
        self.memory_embeddings = nn.Parameter(
            torch.randn(1, config.n_memory_tokens, config.clip_hidden) * 0.02)
        self.layer_fusion = LayerFusion(config)
        self.bank = GeometricMemoryBank(config)
        self.gate = DeltaMemoryGate(config)

        self.output_proj = nn.Sequential(
            nn.Linear(config.clip_hidden, config.clip_hidden),
            nn.GELU(), nn.LayerNorm(config.clip_hidden))
        self.memory_output_fusion = nn.Sequential(
            nn.Linear(config.clip_hidden * 2, config.clip_hidden),
            nn.GELU(),
            nn.Linear(config.clip_hidden, config.clip_hidden))

        self.clip_cross_attn = nn.ModuleList([
            nn.MultiheadAttention(config.clip_hidden, config.n_bank_heads,
                                  batch_first=True, dropout=0.1)
            for _ in range(config.bank_cross_layers)
        ])
        self.clip_cross_norms = nn.ModuleList([
            nn.LayerNorm(config.clip_hidden)
            for _ in range(config.bank_cross_layers)
        ])
        self.clip_cross_ffns = nn.ModuleList([
            nn.Sequential(
                nn.Linear(config.clip_hidden, config.clip_hidden * 2),
                nn.GELU(),
                nn.Linear(config.clip_hidden * 2, config.clip_hidden))
            for _ in range(config.bank_cross_layers)
        ])
        self.clip_cross_ffn_norms = nn.ModuleList([
            nn.LayerNorm(config.clip_hidden)
            for _ in range(config.bank_cross_layers)
        ])

        self.proj_modern = TeacherProjector(config.clip_hidden, config.teacher_hidden)

        # ── NEW: Sequence reconstruction ──
        if config.sequence_output:
            self.sequence_reconstructor = SequenceReconstructor(config)

        # ── CLIP (lazy load) ──
        self._clip_text = None
        self._clip_tokenizer = None

        self.post_init()

    @property
    def clip_text(self):
        if self._clip_text is None:
            self._clip_text = CLIPTextModel.from_pretrained(self.config.clip_model)
            self._clip_text.config.output_hidden_states = True
            for p in self._clip_text.parameters():
                p.requires_grad = False
            self._clip_text = self._clip_text.to(self.memory_embeddings.device)
        return self._clip_text

    @property
    def clip_tokenizer(self):
        if self._clip_tokenizer is None:
            self._clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model)
        return self._clip_tokenizer

    def init_state(self, batch_size, device=None):
        if device is None:
            device = self.memory_embeddings.device
        state = {
            "memory": self.memory_embeddings.expand(batch_size, -1, -1).clone(),
            "bank": self.bank.init_bank(batch_size, device),
            "segment_idx": 0,
        }
        # Collect content tokens across segments for sequence reconstruction
        if self.config.collect_content_tokens:
            state["content_tokens"] = []
        return state

    def forward_segment(self, input_ids, attention_mask, state):
        """
        Process one segment. Identical to v1 except:
          - Collects content tokens from fused output into state
          - Returns more outputs for training
        """
        B = input_ids.shape[0]
        memory_state = state["memory"]
        bank = state["bank"]
        seg_idx = state["segment_idx"]

        # Bank read
        memory_tokens = self.bank.read(memory_state, bank)

        # CLIP forward (frozen)
        max_len = self.config.clip_max_tokens
        with torch.no_grad():
            clip_out = self.clip_text(
                input_ids=input_ids[:, :max_len],
                attention_mask=attention_mask[:, :max_len],
                output_hidden_states=True, return_dict=True)

        all_hiddens = clip_out.hidden_states
        selected = [all_hiddens[i + 1] for i in self.config.extract_layers]
        fused = self.layer_fusion(selected)  # (B, 77, 768)

        # Memory cross-attend to CLIP
        mem_enriched = memory_tokens
        for attn, norm, ffn, ffn_norm in zip(
            self.clip_cross_attn, self.clip_cross_norms,
            self.clip_cross_ffns, self.clip_cross_ffn_norms):
            residual = mem_enriched
            mem_enriched, _ = attn(norm(mem_enriched), fused, fused)
            mem_enriched = residual + mem_enriched
            residual = mem_enriched
            mem_enriched = residual + ffn(ffn_norm(mem_enriched))

        # Depth profile + gate + bank write
        depth_cls = torch.stack([h[:, 1, :] for h in selected], dim=1)
        new_memory = self.gate(memory_state, mem_enriched)
        new_bank = self.bank.write(bank, depth_cls, seg_idx)

        # Pooled output (unchanged from v1)
        clip_pooled = clip_out.pooler_output
        if clip_pooled is None:
            clip_pooled = clip_out.last_hidden_state[:, -1, :]
        cls_output = self.output_proj(clip_pooled)
        memory_delta = self.memory_output_fusion(
            torch.cat([cls_output, new_memory.mean(dim=1)], dim=-1))
        fused_output = cls_output + memory_delta

        # Collect content tokens (skip SOS at 0 and padding)
        # Real content is positions 1..n_content (between SOS and EOS)
        new_state = {
            "memory": new_memory,
            "bank": {"anchors": new_bank["anchors"],
                     "n_written": new_bank["n_written"]},
            "segment_idx": seg_idx + 1,
        }
        if self.config.collect_content_tokens:
            # Extract real content tokens from fused (skip SOS at pos 0)
            n_real = attention_mask.sum(dim=1).max().item()  # max real tokens
            content = fused[:, 1:n_real-1, :]  # skip SOS and EOS
            new_state["content_tokens"] = state.get("content_tokens", []) + [content]

        return fused_output, new_state

    def reconstruct_sequence(self, state):
        """
        After all segments processed, build the full-context sequence.

        Assembles context from:
          - memory_tokens: (B, 8, 768) — fully contextualized
          - bank_anchors: (B, N, 768) — depth-profile per segment
          - content_tokens: (B, total, 768) — per-token from all segments

        Returns: (B, 77, 768)
        """
        memory = state["memory"]          # (B, 8, 768)
        anchors = state["bank"]["anchors"]  # (B, N, 768)

        # Build context for reconstruction
        context_parts = [memory, anchors]

        if self.config.collect_content_tokens and "content_tokens" in state:
            content_list = state["content_tokens"]
            if content_list:
                # Concatenate all segments' content tokens
                all_content = torch.cat(content_list, dim=1)  # (B, total, 768)
                # Truncate to max to prevent OOM
                max_c = self.config.max_content_positions
                if all_content.shape[1] > max_c:
                    all_content = all_content[:, :max_c, :]
                context_parts.append(all_content)

        context = torch.cat(context_parts, dim=1)  # (B, 8+N+content, 768)

        return self.sequence_reconstructor(context)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        texts: Optional[List[str]] = None,
        return_dict: bool = True,
        output_sequence: Optional[bool] = None,
        **kwargs,
    ) -> BaseModelOutput:
        """
        Accepts either:
          - input_ids + attention_mask (single 77-token segment)
          - texts (list of strings, auto-segmented)

        Returns BaseModelOutput:
          - last_hidden_state: (B, 77, 768) sequence if output_sequence else (B, 1, 768) pooled
          - hidden_states: dict with pooled, sequence, state accessible during training
        """
        device = self.memory_embeddings.device
        do_sequence = output_sequence if output_sequence is not None else self.config.sequence_output

        if texts is not None:
            pooled_list = []
            seq_list = []
            state_list = []

            for text in texts:
                p, s, st = self._encode_single(text, device, do_sequence)
                pooled_list.append(p)
                if s is not None:
                    seq_list.append(s)
                state_list.append(st)

            pooled = torch.stack(pooled_list)  # (B, 768)

            if do_sequence and seq_list:
                sequence = torch.stack(seq_list)  # (B, 77, 768)
            else:
                sequence = None

        elif input_ids is not None:
            # Single segment — no sequence reconstruction possible
            state = self.init_state(input_ids.shape[0], device)
            pooled, state = self.forward_segment(
                input_ids.to(device), attention_mask.to(device), state)
            sequence = None

        else:
            raise ValueError("Provide either input_ids or texts")

        # Return format
        if do_sequence and sequence is not None:
            output_tensor = sequence  # (B, 77, 768)
        else:
            output_tensor = pooled.unsqueeze(1)  # (B, 1, 768)

        if return_dict:
            return BaseModelOutput(
                last_hidden_state=output_tensor,
                hidden_states=(pooled, sequence),
                attentions=None,
            )
        return (output_tensor,)

    def _encode_single(self, text, device, do_sequence=True):
        """Encode a single text through segmented memory + optional sequence."""
        segments = segment_text(
            text, self.clip_tokenizer,
            self.config.max_content_tokens,
            self.config.segment_overlap,
            self.config.max_segments)

        state = self.init_state(1, device)
        output = None
        for seg in segments:
            ids = seg["input_ids"].unsqueeze(0).to(device)
            mask = seg["attention_mask"].unsqueeze(0).to(device)
            output, state = self.forward_segment(ids, mask, state)

        pooled = output.squeeze(0)  # (768,)

        sequence = None
        if do_sequence and hasattr(self, "sequence_reconstructor"):
            seq = self.reconstruct_sequence(state)  # (1, 77, 768)
            sequence = seq.squeeze(0)  # (77, 768)

        return pooled, sequence, state

    def encode(self, texts, batch_size=32, show_progress=False,
               return_sequence=False):
        """
        Convenience: encode text(s) to embeddings.

        return_sequence=False → (N, 768) pooled (backward compatible)
        return_sequence=True  → (N, 77, 768) sequence (for diffusion)
        """
        device = self.memory_embeddings.device
        single = isinstance(texts, str)
        if single:
            texts = [texts]

        all_pooled = []
        all_seq = []
        iterator = range(0, len(texts), batch_size)
        if show_progress:
            from tqdm import tqdm
            iterator = tqdm(iterator, desc="Encoding")

        with torch.no_grad():
            for i in iterator:
                batch = texts[i:i + batch_size]
                for text in batch:
                    p, s, _ = self._encode_single(text, device, return_sequence)
                    all_pooled.append(p)
                    if s is not None:
                        all_seq.append(s)

        pooled = torch.stack(all_pooled)

        if return_sequence and all_seq:
            sequence = torch.stack(all_seq)
            result = sequence if not single else sequence.squeeze(0)
        else:
            result = pooled if not single else pooled.squeeze(0)

        return result

    def encode_for_diffusion(self, texts):
        """
        Direct interface for SD/SDXL: returns (B, 77, 768).
        Drop-in replacement for CLIP text encoder output.
        """
        return self.encode(texts, return_sequence=True)