"""
dheeren.models — architecture definitions (pure PyTorch).
Used ONLY as the fallback backend when neither ONNX nor OpenVINO is available.
For normal usage, ONNX/OpenVINO backends are used instead via inference_backend.py
"""

import math, re
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from .config import CFG


# ── Tokenizer ─────────────────────────────────────────────────────────────────

class SimpleTokenizer:
    def __init__(self, word2idx, idx2word, max_length=77):
        self.word2idx   = word2idx
        self.idx2word   = idx2word
        self.max_length = max_length

    def encode(self, text: str) -> torch.Tensor:
        words = re.findall(r'\b\w+\b', text.lower())
        ids   = [self.word2idx["<SOS>"]]
        for w in words[:self.max_length - 2]:
            ids.append(self.word2idx.get(w, self.word2idx["<UNK>"]))
        ids.append(self.word2idx["<EOS>"])
        ids += [self.word2idx["<PAD>"]] * (self.max_length - len(ids))
        return torch.tensor(ids[:self.max_length], dtype=torch.long)


# ── VAE ───────────────────────────────────────────────────────────────────────

class _Decoder(nn.Module):
    def __init__(self, latent_dim=4, hidden_dims=[128, 256, 512], out_channels=3):
        super().__init__()
        hd = list(reversed(hidden_dims))
        self.decoder_input = nn.Conv2d(latent_dim, hd[0], 3, padding=1)
        self.decoder = nn.Sequential(*[
            nn.Sequential(
                nn.ConvTranspose2d(hd[i], hd[i+1], 3, 2, 1, 1),
                nn.BatchNorm2d(hd[i+1]),
                nn.LeakyReLU(0.2)
            )
            for i in range(len(hd) - 1)
        ])
        self.final_layer = nn.Sequential(
            nn.ConvTranspose2d(hd[-1], hd[-1], 3, 2, 1, 1),
            nn.BatchNorm2d(hd[-1]),
            nn.LeakyReLU(0.2),
            nn.Conv2d(hd[-1], out_channels, 3, padding=1),
            nn.Tanh()
        )

    def forward(self, z):
        return self.final_layer(self.decoder(self.decoder_input(z)))


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Conv2d(3, 64, 3, 2, 1))  # stub for key compat
        self.decoder = _Decoder()

    def decode(self, z):
        return self.decoder(z)


# ── CLIP Text Encoder ─────────────────────────────────────────────────────────

class _MHA(nn.Module):
    def __init__(self, dim, heads, drop=0.1):
        super().__init__()
        self.heads = heads
        self.hd    = dim // heads
        self.qkv   = nn.Linear(dim, dim * 3)
        self.proj  = nn.Linear(dim, dim)
        self.drop  = nn.Dropout(drop)

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.hd).permute(2,0,3,1,4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        a = (q @ k.transpose(-2,-1)) * (self.hd ** -0.5)
        if mask is not None:
            a = a.masked_fill(mask == 0, float('-inf'))
        a = self.drop(F.softmax(a, dim=-1))
        return self.proj((a @ v).transpose(1,2).reshape(B, N, C))


class _TBlock(nn.Module):
    def __init__(self, dim, heads, mlp_ratio, drop):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn  = _MHA(dim, heads, drop)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp   = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Dropout(drop),
            nn.Linear(dim * mlp_ratio, dim),
            nn.Dropout(drop),
        )

    def forward(self, x, mask=None):
        x = x + self.attn(self.norm1(x), mask)
        return x + self.mlp(self.norm2(x))


class TextEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(CFG.clip_vocab_size, CFG.clip_embed_dim)
        self.pos_embedding   = nn.Parameter(torch.randn(1, CFG.clip_max_seq_length,
                                                         CFG.clip_embed_dim))
        self.transformer     = nn.ModuleList([
            _TBlock(CFG.clip_embed_dim, CFG.clip_num_heads,
                    CFG.clip_mlp_ratio, CFG.clip_dropout)
            for _ in range(CFG.clip_num_layers)
        ])
        self.norm = nn.LayerNorm(CFG.clip_embed_dim)

    def forward(self, ids):
        x = self.token_embedding(ids) + self.pos_embedding[:, :ids.shape[1]]
        for block in self.transformer:
            x = block(x)
        return F.normalize(self.norm(x)[:, 0, :], dim=-1)


# ── UNet ──────────────────────────────────────────────────────────────────────

class _TSEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        hd = self.dim // 2
        e  = math.log(10000) / (hd - 1)
        e  = torch.exp(torch.arange(hd, device=t.device) * -e)
        e  = t[:, None] * e[None, :]
        return torch.cat([e.sin(), e.cos()], dim=-1)


class _CrossAttn(nn.Module):
    def __init__(self, qdim, cdim, heads):
        super().__init__()
        self.heads  = heads
        self.hd     = qdim // heads
        self.to_q   = nn.Linear(qdim, qdim)
        self.to_k   = nn.Linear(cdim, qdim)
        self.to_v   = nn.Linear(cdim, qdim)
        self.to_out = nn.Linear(qdim, qdim)

    def forward(self, x, ctx):
        B, N, C = x.shape
        def sp(t): return t.reshape(B, -1, self.heads, self.hd).permute(0,2,1,3)
        q, k, v = sp(self.to_q(x)), sp(self.to_k(ctx)), sp(self.to_v(ctx))
        a = F.softmax((q @ k.transpose(-2,-1)) * (self.hd ** -0.5), dim=-1)
        return self.to_out((a @ v).permute(0,2,1,3).reshape(B, N, C))


class _ST(nn.Module):
    def __init__(self, ch, cdim, heads):
        super().__init__()
        self.norm = nn.GroupNorm(32, ch)
        self.proj_in = nn.Conv2d(ch, ch, 1)
        self.transformer_blocks = nn.ModuleList([nn.ModuleDict({
            'norm1': nn.LayerNorm(ch),
            'attn1': _MHA(ch, heads),
            'norm2': nn.LayerNorm(ch),
            'attn2': _CrossAttn(ch, cdim, heads),
            'norm3': nn.LayerNorm(ch),
            'mlp':   nn.Sequential(
                nn.Linear(ch, ch * 4), nn.GELU(), nn.Linear(ch * 4, ch))
        })])
        self.proj_out = nn.Conv2d(ch, ch, 1)

    def forward(self, x, ctx):
        b, c, h, w = x.shape
        xi = x
        x  = rearrange(self.proj_in(self.norm(x)), 'b c h w -> b (h w) c')
        for bl in self.transformer_blocks:
            x = x + bl['attn1'](bl['norm1'](x))
            x = x + bl['attn2'](bl['norm2'](x), ctx)
            x = x + bl['mlp'](bl['norm3'](x))
        return xi + self.proj_out(rearrange(x, 'b (h w) c -> b c h w', h=h, w=w))


class _RB(nn.Module):
    def __init__(self, ic, oc, tdim, drop=0.1):
        super().__init__()
        self.norm1    = nn.GroupNorm(32, ic)
        self.conv1    = nn.Conv2d(ic, oc, 3, padding=1)
        self.time_emb = nn.Sequential(nn.SiLU(), nn.Linear(tdim, oc))
        self.norm2    = nn.GroupNorm(32, oc)
        self.drop     = nn.Dropout(drop)
        self.conv2    = nn.Conv2d(oc, oc, 3, padding=1)
        self.skip     = nn.Conv2d(ic, oc, 1) if ic != oc else nn.Identity()
        self.act      = nn.SiLU()

    def forward(self, x, t):
        h = self.conv1(self.act(self.norm1(x)))
        h = h + self.time_emb(t)[:, :, None, None]
        h = self.conv2(self.drop(self.act(self.norm2(h))))
        return h + self.skip(x)


class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        tdim = CFG.unet_model_channels * 4
        self.time_embed = nn.Sequential(
            _TSEmb(CFG.unet_model_channels),
            nn.Linear(CFG.unet_model_channels, tdim),
            nn.SiLU(),
            nn.Linear(tdim, tdim),
        )
        self.conv_in = nn.Conv2d(CFG.unet_in_channels,
                                 CFG.unet_model_channels, 3, padding=1)

        self.down_blocks = nn.ModuleList()
        chans = [CFG.unet_model_channels]
        nc    = CFG.unet_model_channels

        for lv, m in enumerate(CFG.unet_channel_mult):
            oc = CFG.unet_model_channels * m
            for _ in range(CFG.unet_num_res_blocks):
                use_attn = (CFG.unet_image_size // (2**lv)) in CFG.unet_attention_resolutions
                self.down_blocks.append(nn.ModuleList([
                    _RB(nc, oc, tdim, CFG.unet_dropout),
                    _ST(oc, CFG.unet_context_dim, CFG.unet_num_heads) if use_attn else None
                ]))
                nc = oc
                chans.append(nc)
            if lv != len(CFG.unet_channel_mult) - 1:
                self.down_blocks.append(nn.ModuleList([nn.Conv2d(nc, nc, 3, 2, 1), None]))
                chans.append(nc)

        self.mid_block1 = _RB(nc, nc, tdim, CFG.unet_dropout)
        self.mid_attn   = _ST(nc, CFG.unet_context_dim, CFG.unet_num_heads)
        self.mid_block2 = _RB(nc, nc, tdim, CFG.unet_dropout)

        self.up_blocks = nn.ModuleList()
        for lv, m in enumerate(reversed(CFG.unet_channel_mult)):
            oc = CFG.unet_model_channels * m
            for _ in range(CFG.unet_num_res_blocks + 1):
                rl = len(CFG.unet_channel_mult) - 1 - lv
                use_attn = (CFG.unet_image_size // (2**rl)) in CFG.unet_attention_resolutions
                self.up_blocks.append(nn.ModuleList([
                    _RB(nc + chans.pop(), oc, tdim, CFG.unet_dropout),
                    _ST(oc, CFG.unet_context_dim, CFG.unet_num_heads) if use_attn else None
                ]))
                nc = oc
            if lv != len(CFG.unet_channel_mult) - 1:
                self.up_blocks.append(nn.ModuleList([nn.ConvTranspose2d(nc, nc, 4, 2, 1), None]))

        self.out = nn.Sequential(
            nn.GroupNorm(32, nc),
            nn.SiLU(),
            nn.Conv2d(nc, CFG.unet_out_channels, 3, padding=1),
        )

    def forward(self, x, t, ctx):
        te = self.time_embed(t)
        h  = self.conv_in(x)
        skips = [h]

        for blk, attn in self.down_blocks:
            if isinstance(blk, nn.Conv2d) and blk.stride[0] == 2:
                h = blk(h)
            else:
                h = blk(h, te)
                h = attn(h, ctx) if attn else h
            skips.append(h)

        h = self.mid_block1(h, te)
        h = self.mid_attn(h, ctx)
        h = self.mid_block2(h, te)

        for blk, attn in self.up_blocks:
            if isinstance(blk, nn.ConvTranspose2d):
                h = blk(h)
            else:
                h = blk(torch.cat([h, skips.pop()], 1), te)
                h = attn(h, ctx) if attn else h

        return self.out(h)