import torch, json, numpy as np
from pathlib import Path
import sys
sys.path.insert(0, 'src')

from pino.embeddings import OlfactoryEmbeddingEngine
from pino.heads import PIMTHeads
from pino.pimt_model import DEFAULT_EMBEDDING_DIM, PhysicsInformedMixtureTransformer
from pino.registry import AromaRegistry
from pino.thermo.naturals import NATURAL_PROFILES

# Use a fresh, randomly initialized model to verify the FiLM block propagates time variance
device = torch.device('cpu')
model = PhysicsInformedMixtureTransformer(embedding_dim=DEFAULT_EMBEDDING_DIM, state_dim=2, hidden_dim=256, num_heads=4, num_layers=4).to(device)
model.eval()

# Load a sample record
with open('data/empirical_dataset_v1.jsonl') as f:
    for line in f:
        rec = json.loads(line)
        if any(str(i.get('cas','')).startswith('NATURAL:') for i in rec.get('formula',[])):
            sample = rec
            break

engine = OlfactoryEmbeddingEngine()
registry = AromaRegistry()

tokens = []
for item in sample['formula']:
    cas = str(item['cas']).strip()
    bare = cas.replace('NATURAL:', '')
    rec = registry.get(bare)
    smiles = rec.get('smiles', '') if rec else ''
    z = engine.get_embedding(smiles, cas=cas)
    tokens.append(torch.from_numpy(z))
tokens = torch.stack(tokens, dim=0).unsqueeze(0).float().to(device)

physics = np.zeros((len(sample['trajectory']), len(sample['formula']), 2), dtype=np.float32)
const_to_token = {}
for idx, item in enumerate(sample['formula']):
    raw_cas = str(item['cas']).strip()
    bare = raw_cas.replace('NATURAL:', '')
    profile = NATURAL_PROFILES.get(bare) or NATURAL_PROFILES.get(f'NATURAL:{bare}')
    if profile:
        for const_cas in profile['constituents']:
            const_to_token[const_cas] = idx
    else:
        const_to_token[bare] = idx
for t, step in enumerate(sample['trajectory']):
    for const_cas, token_idx in const_to_token.items():
        physics[t, token_idx, 0] += step['x_liquid'].get(const_cas, 0.0)
        physics[t, token_idx, 1] += step['OAV'].get(const_cas, 0.0)
physics[:,:,1] = np.log10(np.maximum(physics[:,:,1], 1e-10))
physics_t = torch.from_numpy(physics).unsqueeze(0).float().to(device)
mask = torch.zeros(1, tokens.size(1), dtype=torch.bool, device=device)

with torch.no_grad():
    gated = model.gating(tokens, physics_t)
    print('FiLM gated variance across time:', gated.var(dim=1).mean().item())
    print('FiLM gated first vs last timestep mean diff:', (gated[0,0] - gated[0,-1]).abs().mean().item())
    print('FiLM gated first vs last timestep max diff:', (gated[0,0] - gated[0,-1]).abs().max().item())

    proj = model.input_proj(gated)
    b, t, s, _ = proj.shape
    x = proj.reshape(b * t, s, 256)
    enc_out = model.encoder(x, src_key_padding_mask=mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s))
    enc_out = enc_out.reshape(b, t, s, 256)
    print('Encoder output variance across time:', enc_out.var(dim=1).mean().item())
    print('Encoder output first vs last timestep mean diff:', (enc_out[0,0] - enc_out[0,-1]).abs().mean().item())
    print('Encoder output first vs last timestep max diff:', (enc_out[0,0] - enc_out[0,-1]).abs().max().item())

    # Compare to zero physics
    physics0 = torch.zeros_like(physics_t)
    gated0 = model.gating(tokens, physics0)
    proj0 = model.input_proj(gated0)
    x0 = proj0.reshape(b * t, s, 256)
    enc_out0 = model.encoder(x0, src_key_padding_mask=mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s))
    enc_out0 = enc_out0.reshape(b, t, s, 256)
    print('Diff encoder vs zero physics (mean):', (enc_out - enc_out0).abs().mean().item())
    print('Diff encoder vs zero physics (max):', (enc_out - enc_out0).abs().max().item())