import sys
from pathlib import Path
import json
import numpy as np
import torch

sys.path.insert(0, "src")
sys.path.insert(0, "scripts")

from predict_shalimar import (
    build_model_inputs,
    aggregate_physics_for_tokens,
    expand_naturals_for_vle,
    map_ingredients_to_cas,
    parse_formula,
    load_pimt_checkpoint,
    HARDCODED_CAS,
)
from pino.registry import AromaRegistry
from pino.embeddings import OlfactoryEmbeddingEngine
from pino.verifier import FragrancePipelineVerifier


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pdf_path = "/home/hermes/.hermes/cache/documents/doc_fa4254a05dd7_SHALIMAR_1990.pdf"

    import fitz

    doc = fitz.open(pdf_path)
    text = "\n".join(page.get_text() for page in doc)
    doc.close()

    formula = parse_formula(text)
    registry = AromaRegistry()
    mapped = map_ingredients_to_cas(formula, registry)

    semantic_recipe = []
    total = 0.0
    for name, (cas, qty) in mapped.items():
        if "dipropylene glycol" in name.lower() or not cas:
            continue
        semantic_recipe.append({"cas": cas, "weight_fraction": qty})
        total += qty
    for item in semantic_recipe:
        item["weight_fraction"] /= total

    vle_recipe = expand_naturals_for_vle(semantic_recipe)
    verifier = FragrancePipelineVerifier()
    result = verifier.run_sim(vle_recipe, duration_seconds=8 * 3600, interval_seconds=600, skip_ifra=True)
    if result["status"] != "passed":
        print("Simulation failed")
        return

    engine = OlfactoryEmbeddingEngine()
    tokens = build_model_inputs(semantic_recipe, registry, engine).to(device)
    physics, _ = aggregate_physics_for_tokens(result["trajectory"], semantic_recipe)
    physics_t = torch.from_numpy(physics).unsqueeze(0).float().to(device)
    mask = torch.zeros(1, tokens.size(1), dtype=torch.bool, device=device)

    model, heads, _ = load_pimt_checkpoint()
    model.to(device)
    heads.to(device)

    activations = {}

    def make_hook(name):
        def fn(m, i, o):
            activations[name] = o.detach() if not isinstance(o, tuple) else o[0].detach()
        return fn

    model.gating.register_forward_hook(make_hook("film"))
    model.encoder.register_forward_hook(make_hook("encoder"))

    with torch.no_grad():
        latent = model(tokens, physics_t, src_key_padding_mask=mask)
        output = heads(latent)

    print(f"tokens shape: {tokens.shape}")
    print(f"physics shape: {physics_t.shape}")
    print(f"physics per-channel mean time-variance: {physics_t[0].var(dim=0).mean(dim=0).cpu().numpy()}")
    print(f"latent shape: {latent.shape}")
    print(f"latent mean per-token time-variance: {latent[0].var(dim=0).mean().item():.6e}")
    print(f"latent max  per-token time-variance: {latent[0].var(dim=0).max().item():.6e}")
    for k in ["film", "encoder"]:
        v = activations[k]
        print(f"{k} shape: {v.shape}")
        print(f"  mean per-token time-variance: {v[0].var(dim=0).mean().item():.6e}")
        print(f"  max  per-token time-variance: {v[0].var(dim=0).max().item():.6e}")
    obj = output["objective"]
    print(f"objective shape: {obj.shape}")
    print(f"  mean per-dim time-variance: {obj[0].var(dim=0).mean().item():.6e}")
    print(f"  max  per-dim time-variance: {obj[0].var(dim=0).max().item():.6e}")


if __name__ == "__main__":
    main()