"""
FinJEPA: Financial Joint-Embedding Predictive Architecture
A JEPA-based world model for portfolio optimization over a separated action space.
"""

import math
from typing import Optional, Tuple, Dict, List
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


def off_diagonal(x):
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()


def sigreg_loss(z_a, z_b, sim_coeff=0.0, std_coeff=16.0, cov_coeff=8.0):
    sim_loss = torch.tensor(0.0, device=z_a.device)
    std_z_a = torch.sqrt(z_a.var(dim=0) + 1e-4)
    std_z_b = torch.sqrt(z_b.var(dim=0) + 1e-4)
    std_loss = torch.mean(F.relu(1 - std_z_a)) + torch.mean(F.relu(1 - std_z_b))
    N = z_a.size(0)
    z_a = z_a - z_a.mean(dim=0)
    z_b = z_b - z_b.mean(dim=0)
    cov_z_a = (z_a.T @ z_a) / (N - 1)
    cov_z_b = (z_b.T @ z_b) / (N - 1)
    cov_loss = off_diagonal(cov_z_a).pow_(2).sum() / z_a.size(1) + off_diagonal(cov_z_b).pow_(2).sum() / z_b.size(1)
    loss = sim_coeff * sim_loss + std_coeff * std_loss + cov_coeff * cov_loss
    return loss, sim_loss, std_loss, cov_loss


class TimeSeriesTokenizer(nn.Module):
    def __init__(self, in_features: int, embed_dim: int = 128, patch_size: int = 4):
        super().__init__()
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        self.proj = nn.Conv1d(in_features, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.max_patches = 1024
        self.pos_embed = nn.Parameter(torch.zeros(1, self.max_patches, embed_dim))
        nn.init.normal_(self.pos_embed, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, F = x.shape
        x = x.transpose(1, 2)
        x = self.proj(x)
        x = x.transpose(1, 2)
        N = x.size(1)
        pos = self.pos_embed[:, :N]
        x = x + pos
        return x


class RotaryPositionEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len = max_seq_len
        self.dim = dim

    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        device = x.device
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()[None, :, :]
        sin = emb.sin()[None, :, :]
        return cos, sin


def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., ::2], x[..., 1::2]
    rotated = torch.stack([-x2, x1], dim=-1).flatten(-2)
    return x * cos + rotated * sin


class AdaLN(nn.Module):
    def __init__(self, dim: int, action_dim: int):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.scale_shift = nn.Linear(action_dim, dim * 2)
        nn.init.zeros_(self.scale_shift.weight)
        nn.init.zeros_(self.scale_shift.bias)
        self.scale_shift.bias.data[:dim] = 1.0

    def forward(self, x: torch.Tensor, a_emb: torch.Tensor) -> torch.Tensor:
        x = self.norm(x)
        scale, shift = self.scale_shift(a_emb).chunk(2, dim=-1)
        return x * (1 + scale) + shift


class FinJEPATransformerBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, action_dim: int = 128, dropout: float = 0.0):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert dim % num_heads == 0
        self.adaln1 = AdaLN(dim, action_dim)
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.adaln2 = AdaLN(dim, action_dim)
        mlp_dim = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_dim), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(mlp_dim, dim), nn.Dropout(dropout),
        )
        self.rope = RotaryPositionEmbedding(self.head_dim)

    def forward(self, x: torch.Tensor, a_emb: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, mask=None):
        B, N, D = x.shape
        x_norm = self.adaln1(x, a_emb)
        qkv = self.qkv(x_norm).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q = apply_rotary(q, cos, sin)
        k = apply_rotary(k, cos, sin)
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        out = (attn @ v).transpose(1, 2).reshape(B, N, D)
        out = self.proj(out)
        x = x + out
        x_norm = self.adaln2(x, a_emb)
        x = x + self.mlp(x_norm)
        return x


class FinJEPATransformerEncoder(nn.Module):
    def __init__(self, dim=128, depth=4, num_heads=4, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.blocks = nn.ModuleList([
            FinJEPATransformerBlock(dim, num_heads, mlp_ratio, action_dim=dim, dropout=dropout)
            for _ in range(depth)
        ])
        self.register_buffer("dummy_action", torch.zeros(1, 1, dim))
        self.norm = nn.LayerNorm(dim)
        self.rope = RotaryPositionEmbedding(dim // num_heads)

    def forward(self, x, mask=None):
        B, N, D = x.shape
        cos, sin = self.rope(x, N)
        a = self.dummy_action.expand(B, N, D)
        for block in self.blocks:
            x = block(x, a, cos, sin, mask)
        x = self.norm(x)
        return x


class ActionEmbedder(nn.Module):
    def __init__(self, n_assets=10, signal_vocab_size=3, hidden_dim=128, out_dim=128):
        super().__init__()
        self.n_assets = n_assets
        self.weight_proj = nn.Sequential(
            nn.Linear(n_assets, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, hidden_dim),
        )
        self.signal_embed = nn.Embedding(signal_vocab_size, hidden_dim // n_assets)
        self.hedge_embed = nn.Embedding(2, hidden_dim // 2)
        fusion_in = hidden_dim + n_assets * (hidden_dim // n_assets) + hidden_dim // 2
        self.fusion = nn.Sequential(
            nn.Linear(fusion_in, out_dim), nn.GELU(), nn.Linear(out_dim, out_dim),
        )

    def forward(self, weights, signals, hedge=None):
        w_emb = self.weight_proj(weights)
        s_emb = self.signal_embed(signals).flatten(1)
        if hedge is None:
            h_emb = torch.zeros(weights.size(0), self.fusion[0].in_features - w_emb.size(1) - s_emb.size(1), device=weights.device)
        else:
            h_emb = self.hedge_embed(hedge).squeeze(1)
        a_emb = self.fusion(torch.cat([w_emb, s_emb, h_emb], dim=-1))
        return a_emb


class FinJEPAPredictor(nn.Module):
    def __init__(self, dim=128, depth=6, num_heads=4, mlp_ratio=4.0, action_dim=128, dropout=0.0, max_target_tokens=64):
        super().__init__()
        self.dim = dim
        self.max_target_tokens = max_target_tokens
        self.target_queries = nn.Parameter(torch.zeros(1, max_target_tokens, dim))
        nn.init.normal_(self.target_queries, std=0.02)
        self.blocks = nn.ModuleList([
            FinJEPATransformerBlock(dim, num_heads, mlp_ratio, action_dim, dropout)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(dim)
        self.action_proj = nn.Linear(action_dim, dim)
        self.rope = RotaryPositionEmbedding(dim // num_heads)

    def forward(self, x, action_emb, n_target_tokens=None, mask=None):
        B, N_ctx, D = x.shape
        n_target = n_target_tokens if n_target_tokens is not None else 1
        queries = self.target_queries[:, :n_target].expand(B, -1, -1)
        x = torch.cat([x, queries], dim=1)
        N_total = x.size(1)
        a_seq = self.action_proj(action_emb).unsqueeze(1).expand(B, N_total, D)
        cos, sin = self.rope(x, N_total)
        for block in self.blocks:
            x = block(x, a_seq, cos, sin, mask)
        x = self.norm(x)
        target_pred = x[:, N_ctx:]
        return target_pred


class InverseDynamicsModel(nn.Module):
    def __init__(self, z_dim, n_assets, signal_vocab_size=3, hidden_dim=256):
        super().__init__()
        self.n_assets = n_assets
        self.mlp = nn.Sequential(
            nn.Linear(z_dim * 2, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
        )
        self.weight_head = nn.Linear(hidden_dim, n_assets)
        self.signal_head = nn.Linear(hidden_dim, n_assets * signal_vocab_size)

    def forward(self, z_t, z_tp1):
        h = self.mlp(torch.cat([z_t, z_tp1], dim=-1))
        weights = self.weight_head(h)
        signals_logits = self.signal_head(h).reshape(-1, self.n_assets, 3)
        return {"weights": weights, "signals_logits": signals_logits}


class FinJEPA(nn.Module):
    def __init__(
        self, in_features=7, n_assets=10, patch_size=4, embed_dim=128,
        encoder_depth=4, encoder_heads=4, predictor_depth=6, predictor_heads=4,
        action_hidden_dim=128, signal_vocab_size=3, mlp_ratio=4.0,
        dropout=0.0, ema_decay=0.996, use_idm=True, use_td_branch=False,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.ema_decay = ema_decay
        self.use_idm = use_idm
        self.use_td_branch = use_td_branch
        self.n_assets = n_assets

        self.tokenizer = TimeSeriesTokenizer(in_features, embed_dim, patch_size)
        self.context_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout)
        self.target_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout)
        for p in self.target_encoder.parameters():
            p.requires_grad = False

        self.action_embedder = ActionEmbedder(n_assets, signal_vocab_size, action_hidden_dim, embed_dim)
        self.predictor = FinJEPAPredictor(embed_dim, predictor_depth, predictor_heads, mlp_ratio, embed_dim, dropout)

        if use_idm:
            self.idm = InverseDynamicsModel(embed_dim, n_assets, signal_vocab_size, hidden_dim=256)
        else:
            self.idm = None

        if use_td_branch:
            self.task_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout)
            for p in self.task_encoder.parameters():
                p.requires_grad = False
            self.task_predictor = FinJEPAPredictor(embed_dim, predictor_depth, predictor_heads, mlp_ratio, embed_dim, dropout)
        else:
            self.task_encoder = None
            self.task_predictor = None

        self._init_target_encoder()

    def _init_target_encoder(self):
        self.target_encoder.load_state_dict(self.context_encoder.state_dict())
        for p in self.target_encoder.parameters():
            p.requires_grad = False

    def update_target(self):
        with torch.no_grad():
            for param_s, param_t in zip(self.context_encoder.parameters(), self.target_encoder.parameters()):
                param_t.data.mul_(self.ema_decay).add_(param_s.data, alpha=1 - self.ema_decay)

    def encode_context(self, x, mask=None):
        tokens = self.tokenizer(x)
        z = self.context_encoder(tokens, mask)
        return z

    @torch.no_grad()
    def encode_target(self, x, mask=None):
        tokens = self.tokenizer(x)
        z = self.target_encoder(tokens, mask)
        return z

    def forward(self, context_series, target_series, weights, signals, hedge=None, context_mask=None, target_mask=None):
        z_context = self.encode_context(context_series, context_mask)
        with torch.no_grad():
            z_target = self.encode_target(target_series, target_mask)
        action_emb = self.action_embedder(weights, signals, hedge)
        n_target = z_target.size(1)
        z_pred = self.predictor(z_context, action_emb, n_target_tokens=n_target, mask=context_mask)
        idm_out = None
        if self.use_idm and self.idm is not None:
            z_t = z_context[:, -1]
            z_next = z_target[:, 0]
            idm_out = self.idm(z_t, z_next)
        return {"z_pred": z_pred, "z_target": z_target, "action_emb": action_emb, "idm": idm_out}

    def predict_next_state(self, state_series, weights, signals, hedge=None, mask=None):
        z = self.encode_context(state_series, mask)
        action_emb = self.action_embedder(weights, signals, hedge)
        z_next = self.predictor(z, action_emb, n_target_tokens=1, mask=mask)
        return z_next


class FinJEPALoss(nn.Module):
    def __init__(self, pred_loss="l1", alpha=2.0, beta=1.0, delta=4.0, omega=0.5, gamma=0.5, use_sigreg=True):
        super().__init__()
        self.pred_loss = pred_loss
        self.alpha = alpha
        self.beta = beta
        self.delta = delta
        self.omega = omega
        self.gamma = gamma
        self.use_sigreg = use_sigreg

    def forward(self, outputs, actions_gt=None, rollout_outputs=None):
        z_pred = outputs["z_pred"]
        z_target = outputs["z_target"]
        idm_out = outputs.get("idm")

        if self.pred_loss == "l1":
            l_pred = F.l1_loss(z_pred, z_target)
        else:
            l_pred = F.mse_loss(z_pred, z_target)

        B, N, D = z_pred.shape
        z_pred_flat = z_pred.reshape(B * N, D)
        z_target_flat = z_target.reshape(B * N, D)
        l_reg, l_sim, l_std, l_cov = sigreg_loss(
            z_pred_flat, z_target_flat, sim_coeff=0.0, std_coeff=self.alpha, cov_coeff=self.beta
        )

        if z_pred.size(1) > 1:
            z_diff = (z_pred[:, 1:] - z_pred[:, :-1]).pow(2).mean()
        else:
            z_diff = torch.tensor(0.0, device=z_pred.device)
        l_temporal = self.delta * z_diff

        l_idm = torch.tensor(0.0, device=z_pred.device)
        if idm_out is not None and actions_gt is not None:
            l_w = F.mse_loss(idm_out["weights"], actions_gt["weights"])
            l_s = F.cross_entropy(idm_out["signals_logits"].reshape(-1, 3), actions_gt["signals"].reshape(-1))
            l_idm = self.omega * (l_w + l_s)

        l_rollout = torch.tensor(0.0, device=z_pred.device)
        if rollout_outputs is not None and len(rollout_outputs) > 0:
            for ro in rollout_outputs:
                if self.pred_loss == "l1":
                    l_rollout += F.l1_loss(ro["z_pred"], ro["z_target"])
                else:
                    l_rollout += F.mse_loss(ro["z_pred"], ro["z_target"])
            l_rollout = self.gamma * (l_rollout / len(rollout_outputs))

        total = l_pred + l_reg + l_temporal + l_idm + l_rollout
        return {
            "loss": total,
            "loss_pred": l_pred.item(),
            "loss_reg": l_reg.item(),
            "loss_sim": l_sim.item(),
            "loss_std": l_std.item(),
            "loss_cov": l_cov.item(),
            "loss_temporal": l_temporal.item(),
            "loss_idm": l_idm.item() if isinstance(l_idm, torch.Tensor) else 0.0,
            "loss_rollout": l_rollout.item() if isinstance(l_rollout, torch.Tensor) else 0.0,
        }