import json, numpy as np
import sys
sys.path.insert(0, 'src')
from pino.thermo.naturals import NATURAL_PROFILES

records = []
with open('data/empirical_dataset_v1.jsonl') as f:
    for line in f:
        rec = json.loads(line)
        if not rec.get('formula') or not rec.get('trajectory'):
            continue
        records.append(rec)
        if len(records) >= 200:
            break

max_variance = -1
best_rec = None
best_physics = None
for rec in records:
    physics = np.zeros((len(rec['trajectory']), len(rec['formula']), 2), dtype=np.float32)
    for t, step in enumerate(rec['trajectory']):
        for j, comp in enumerate(rec['formula']):
            name = comp.get('cas', f'ing_{j}')
            x_liq = step['x_liquid'].get(name, 0.0)
            oav = step['OAV'].get(name, 0.0)
            physics[t, j, 0] = x_liq
            physics[t, j, 1] = np.log10(max(oav, 1e-10))
    # Compute total variance across time, summed over tokens and state dims
    variance = np.sum(np.var(physics, axis=0))
    if variance > max_variance:
        max_variance = variance
        best_rec = rec
        best_physics = physics

print(f'Best sample total time variance: {max_variance:.6f}')
print(f'Formula tokens: {[i.get("cas") for i in best_rec["formula"]]}')
print(f'Physics shape: {best_physics.shape}')
print(f'x_liquid range: {best_physics[:,:,0].min():.4f} to {best_physics[:,:,0].max():.4f}')
print(f'log10OAV range: {best_physics[:,:,1].min():.4f} to {best_physics[:,:,1].max():.4f}')

# Save best record index info for diagnostic script
import pickle
with open('/tmp/best_film_record.pkl', 'wb') as f:
    pickle.dump({'record': best_rec, 'physics': best_physics}, f)