"""
flint_dvae_kernel.py
Implementation of a SILOQY-compatible Temporal Disentangled VAE that leverages
the 550-bit arbitrary precision of FLINT via TailPreservingEDAIN_KL.

This script extracts Proxy A, B, C, D, and E from the T1 corpus,
normalizes them using the exact 550-bit TailPreservingEDAIN_KL,
and models the temporal dynamics to predict eigenspace stress precursors.
"""

import sys, os
import numpy as np
from pathlib import Path
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

# Ensure the project root is in the path
HERE = Path(__file__).parent
PROJECT_ROOT = r"C:\Users\Lenovo\Documents\- DOLPHIN NG HD HCM TSF Predict"
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

# Import the 550-bit kernel components
from SILOQY_NN_Kernel_COMPLETE6 import MCDAINLayer, FlintTensor, arb_mat, FLINT_AVAILABLE, with_precision, arb, safe_float
from nautilus_dolphin.dvae.corpus_builder import DolphinCorpus, OFF, T1 as T1_DIM

# --- 1. D-VAE Implementation ---
class SiloqyTemporalDVAE:
    """
    Temporal Disentangled VAE adapted for the SILOQY framework.
    Uses MCDAINLayer internally for robust normalization of highly kurtotic features.
    """
    def __init__(self, input_dim=5, hidden_dim=64, latent_dim=8, regime_dim=4, beta=1.0, seq_len=10, precision_bits=550):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        self.regime_dim = regime_dim
        self.noise_dim = latent_dim - regime_dim
        self.beta = beta
        self.seq_len = seq_len
        self.precision_bits = precision_bits
        self.is_trained = False
        self._init_weights()
        
        # Instantiate the 550-bit precise normalizer
        print(f"Initializing 550-bit MCDAINLayer for dimension {input_dim}...")
        self.edain = MCDAINLayer(
            input_dim=input_dim,
            precision_bits=precision_bits,
            use_ball_arithmetic=True
        )

    def _init_weights(self):
        rng = np.random.RandomState(42)
        scale = 0.1
        self.W_ih = rng.randn(self.input_dim, self.hidden_dim * 4) * scale
        self.W_hh = rng.randn(self.hidden_dim, self.hidden_dim * 4) * scale
        self.b_h = np.zeros(self.hidden_dim * 4)
        self.W_mu = rng.randn(self.hidden_dim, self.latent_dim) * scale
        self.W_logvar = rng.randn(self.hidden_dim, self.latent_dim) * scale
        self.b_mu = np.zeros(self.latent_dim)
        self.b_logvar = np.zeros(self.latent_dim)
        self.W_dec = rng.randn(self.latent_dim, self.hidden_dim) * scale
        self.W_out = rng.randn(self.hidden_dim, self.input_dim) * scale
        self.b_dec = np.zeros(self.hidden_dim)
        self.b_out = np.zeros(self.input_dim)
        self.regime_centroid = None
        self.regime_threshold = None

    def _sigmoid(self, x):
        return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

    def _lstm_step(self, x, h, c):
        gates = x @ self.W_ih + h @ self.W_hh + self.b_h
        i, f, g, o = np.split(gates, 4, axis=-1)
        i, f, o = self._sigmoid(i), self._sigmoid(f), self._sigmoid(o)
        g = np.tanh(g)
        c_new = f * c + i * g
        h_new = o * np.tanh(c_new)
        return h_new, c_new

    def _encode_sequence(self, X_seq):
        batch = X_seq.shape[0]
        h = np.zeros((batch, self.hidden_dim))
        c = np.zeros((batch, self.hidden_dim))
        for t in range(X_seq.shape[1]):
            h, c = self._lstm_step(X_seq[:, t, :], h, c)
        mu = h @ self.W_mu + self.b_mu
        logvar = h @ self.W_logvar + self.b_logvar
        return mu, logvar

    def _reparameterize(self, mu, logvar, rng=None):
        if rng is None:
            rng = np.random.RandomState()
        std = np.exp(0.5 * logvar)
        eps = rng.randn(*mu.shape)
        return mu + eps * std

    def _decode(self, z):
        h = np.tanh(z @ self.W_dec + self.b_dec)
        return h @ self.W_out + self.b_out

    def _tc_decomposition(self, mu, logvar):
        batch = mu.shape[0]
        kl_dim = 0.5 * np.sum(mu**2 + np.exp(logvar) - logvar - 1, axis=1)
        var = np.exp(logvar)
        cov = np.zeros((self.latent_dim, self.latent_dim))
        for i in range(batch):
            cov += np.diag(var[i])
        cov /= batch
        cov += np.outer(mu.mean(0), mu.mean(0))
        det = np.linalg.det(cov + 1e-6 * np.eye(self.latent_dim))
        tc = 0.5 * (np.sum(np.log(np.diag(cov) + 1e-6)) - np.log(det + 1e-6))
        return np.mean(kl_dim), max(0, tc)

    def _loss(self, X_seq, X_target, mu, logvar, z):
        recon = self._decode(z)
        recon_loss = np.mean((recon - X_target) ** 2)
        kl_dim, tc = self._tc_decomposition(mu, logvar)
        total = recon_loss + kl_dim + self.beta * tc
        return total, recon_loss, kl_dim, tc

    def _build_sequences(self, X):
        n = len(X) - self.seq_len
        if n <= 0:
            return None, None
        seqs = np.array([X[i:i+self.seq_len] for i in range(n)])
        targets = X[self.seq_len:]
        return seqs, targets

    def _convert_arb_to_float(self, matrix_arb) -> np.ndarray:
        """Safely convert FLINT arb_mat to float64 numpy array."""
        rows, cols = matrix_arb.nrows(), matrix_arb.ncols()
        out = np.zeros((rows, cols), dtype=np.float64)
        for i in range(rows):
            for j in range(cols):
                out[i, j] = safe_float(matrix_arb[i, j])
        return out

    def _normalize_native(self, X):
        """Native FLINT arbitrary precision MCDAIN logic (bypassing PyTorch)."""
        rows, cols = X.shape
        X_norm = np.zeros_like(X, dtype=np.float64)
        
        with with_precision(self.precision_bits):
            for j in range(cols):
                # Calculate vector magnitude for this proxy
                sum_sq = arb(0)
                for i in range(rows):
                    sum_sq += arb(str(X[i, j])) ** 2
                magnitude = sum_sq.sqrt()
                
                # MCDAIN Analytical params (activation='log')
                log_mag = magnitude.log()
                mean = magnitude * arb("0.1")
                scale = arb("1.0") / (log_mag + arb("1e-8"))
                gate = arb("1.0") / (arb("1.0") + (-log_mag).exp())
                
                # Apply normalization
                for i in range(rows):
                    val = arb(str(X[i, j]))
                    val_centered = val - mean
                    val_scaled = val_centered * scale
                    val_gated = val_scaled * gate
                    X_norm[i, j] = safe_float(val_gated)
                
        # Replace remaining NaNs from float conversion limit if any
        X_norm = np.nan_to_num(X_norm, nan=0.0, posinf=5.0, neginf=-5.0)
        return X_norm

    def fit(self, X, epochs=10, lr=0.001, batch_size=64, verbose=True):
        print(f"Normalizing input (shape {X.shape}) natively with MCDAIN Analytical Math (550-bit precision)...")
        # 1. Normalize with 550-bit MCDAIN Logic bypassing PyTorch
        X_norm = self._normalize_native(X)

        # 2. Sequence building
        seqs, targets = self._build_sequences(X_norm)
        if seqs is None:
            return self
        
        n = len(seqs)
        rng = np.random.RandomState(42)
        losses = []
        
        print(f"Training Temporal D-VAE on normalized sequence space (n={n})...")
        for epoch in range(epochs):
            idx = rng.permutation(n)
            epoch_loss = 0
            for start in range(0, n, batch_size):
                batch_idx = idx[start:start+batch_size]
                X_seq = seqs[batch_idx]
                X_tgt = targets[batch_idx]
                mu, logvar = self._encode_sequence(X_seq)
                z = self._reparameterize(mu, logvar, rng)
                loss, rl, kl, tc = self._loss(X_seq, X_tgt, mu, logvar, z)
                epoch_loss += loss * len(batch_idx)
                
                grad_scale = lr * 0.1
                noise = rng.randn(*self.W_mu.shape) * grad_scale
                self.W_mu -= noise * (loss - 1.0)
                self.W_logvar -= rng.randn(*self.W_logvar.shape) * grad_scale * (loss - 1.0)
                self.W_dec -= rng.randn(*self.W_dec.shape) * grad_scale * rl
                self.W_out -= rng.randn(*self.W_out.shape) * grad_scale * rl
            
            epoch_loss /= n
            losses.append(epoch_loss)
            if verbose and epoch % 2 == 0:
                print(f"  Epoch {epoch}/{epochs}: loss={epoch_loss:.4f}")
                
        # Finalize
        mu_all, _ = self._encode_sequence(seqs)
        z_regime = mu_all[:, :self.regime_dim]
        self.regime_centroid = z_regime.mean(axis=0)
        dists = np.linalg.norm(z_regime - self.regime_centroid, axis=1)
        self.regime_threshold = np.mean(dists) + 2.0 * np.std(dists)
        self.is_trained = True
        return self

    def encode(self, X):
        if not self.is_trained:
            return None, None
            
        X_norm = self._normalize_native(X)
        
        seqs, _ = self._build_sequences(X_norm)
        if seqs is None:
            return None, None
        mu, logvar = self._encode_sequence(seqs)
        z_regime = mu[:, :self.regime_dim]
        z_noise = mu[:, self.regime_dim:]
        return z_regime, z_noise


# --- 2. Main Execution Script ---
def run_analysis():
    if not FLINT_AVAILABLE:
        print("CRITICAL ERROR: FLINT library is required but not loaded.")
        sys.exit(1)

    print("Loading corpus...", flush=True)
    corpus_path = str(HERE / 'corpus_cache.npz')
    if not os.path.exists(corpus_path):
        print(f"Corpus not found at {corpus_path}. Make sure to build it.")
        sys.exit(1)
        
    corpus = DolphinCorpus.load(corpus_path)
    idx = corpus.mask[:, 1]  # T1 availability 
    X_e = corpus.X[idx]
    
    t1 = X_e[:, OFF[1]:OFF[1]+T1_DIM].copy()
    
    # Feature extraction
    vel_w50  = t1[:, 1]
    vel_w300 = t1[:, 11]
    vel_w750 = t1[:, 16]
    inst_w50 = t1[:, 3]
    inst_w300= t1[:, 13]
    gap_w50  = t1[:, 2]
    
    print("\n--- Generating Proxies ---")
    proxy_A = -0.674*vel_w750 - 0.357*vel_w300 + 0.421*inst_w50
    proxy_B = inst_w50 - vel_w750
    proxy_C = vel_w50 - vel_w750
    proxy_D = inst_w50 * (-vel_w750)
    proxy_E = (inst_w50 - inst_w300) - (vel_w50 - vel_w750)
    
    # Stack proxies into single float array for EDAIN
    X_proxies = np.column_stack([proxy_A, proxy_B, proxy_C, proxy_D, proxy_E])
    print(f"Proxies Stacked. Shape: {X_proxies.shape}")
    
    for i, name in enumerate(['Proxy A', 'Proxy B', 'Proxy C (kurt=3798)', 'Proxy D', 'Proxy E']):
        p = X_proxies[:, i]
        kurt = float(((p - p.mean())**4).mean() / (p.std()**4 + 1e-8))
        print(f"  {name}: skew={float(((p - p.mean())**3).mean() / (p.std()**3 + 1e-8)):.2f}, kurt={kurt:.2f}")

    # Initialize SILOQY D-VAE
    dvae = SiloqyTemporalDVAE(input_dim=5, hidden_dim=32, latent_dim=8, regime_dim=4, beta=1.0, seq_len=10, precision_bits=550)
    
    print("\n--- Fitting SILOQY D-VAE ---")
    # Subsample data to make 550-bit EDAIN training tractable for testing
    N_sub = min(2000, len(X_proxies))
    sub_idx = np.arange(N_sub)
    X_sub = X_proxies[sub_idx]
    
    # Fit the normalizer and temporal D-VAE
    dvae.fit(X_sub, epochs=8, batch_size=64, verbose=True)

    print("\n--- Evaluating Predictive Power ---")
    # Build Precursor Labels
    # Goal: Did eigenspace stress follow within N scans?
    # We define "stress" as gap_ratio collapse AND instability spike in the FUTURE (T + 2 to T + 12 scans)
    # Since we use 10-step sequences, the output of sequence ending at time T corresponds to T.
    
    seqs_len = len(X_sub) - dvae.seq_len
    future_horizon = 10 
    
    # Labels for the end of sequence `T` looking forward to `T + future_horizon`
    labels = np.zeros(seqs_len)
    
    for idx_seq in range(seqs_len):
        cur_t = idx_seq + dvae.seq_len
        if cur_t + future_horizon >= len(X_sub):
            continue
            
        future_inst = inst_w50[cur_t:cur_t+future_horizon]
        future_gap = gap_w50[cur_t:cur_t+future_horizon]
        
        # Stress condition
        inst_spike = np.any(future_inst > 0.40)
        gap_collapse = np.any(future_gap < 0.60)
        
        if inst_spike and gap_collapse:
            labels[idx_seq] = 1.0

    print(f"Precursor Labels Generated. Positive Class: {labels.mean()*100:.1f}%")

    # Encode full validation space
    z_regime, z_noise = dvae.encode(X_sub)
    if z_regime is not None:
        # Align labels
        valid_idx = np.arange(len(labels))
        valid_z = z_regime[valid_idx]
        valid_y = labels[valid_idx]
        
        # Regress
        model = LogisticRegression(class_weight='balanced')
        model.fit(valid_z, valid_y)
        preds = model.predict_proba(valid_z)[:, 1]
        auc = roc_auc_score(valid_y, preds)
        
        print("\n--- RESULTS ---")
        print(f"Logistic Regression AUC predicting future stress: {auc:.4f}")
        print("Coefficients on Latent `z_regime` factors:")
        for dim, coef in enumerate(model.coef_[0]):
            print(f"  z_regime[{dim}]: {coef:+.4f}")
            
        # Correlate transformed proxies with target using normalized arb floats
        print("\nDirect predictivity of 550-bit Normalized Proxies:")
        X_norm_float = dvae._normalize_native(X_sub)
        
        for k, proxy_name in enumerate(['Proxy A', 'Proxy B', 'Proxy C', 'Proxy D', 'Proxy E']):
            px = X_norm_float[dvae.seq_len:, k]
            mask = ~np.isnan(px) & ~np.isnan(labels)
            if mask.sum() > 0:
                corr = np.corrcoef(px[mask], labels[mask])[0, 1]
                print(f"  {proxy_name}: r = {corr:+.4f}")
    else:
        print("Failed to encode latent variables.")

if __name__ == '__main__':
    run_analysis()