import torch
import torch.nn.functional as F
import gradio as gr
import selfies as sf
import numpy as np
import pandas as pd
import os

from rdkit import Chem
from rdkit.Chem import QED

CKPT_PATH = "diffusion_cfg_mtr_stage2_DEPLOY.pt"

checkpoint = torch.load(CKPT_PATH, map_location="cpu")

vocab = checkpoint["vocab"]
inv_vocab = checkpoint["inv_vocab"]
max_len = checkpoint["max_len"]
raw_config = checkpoint["config"]

allowed_keys = {"dim", "depth", "heads", "dropout", "num_mtr_targets"}
model_config = {k: v for k, v in raw_config.items() if k in allowed_keys}

class PositionalEmbedding(torch.nn.Module):
    def __init__(self, max_len, dim):
        super().__init__()
        self.emb = torch.nn.Embedding(max_len, dim)

    def forward(self, x):
        positions = torch.arange(x.size(1), device=x.device)
        return self.emb(positions)[None, :, :]


class DiffusionTransformer(torch.nn.Module):
    def __init__(
        self,
        vocab_size,
        max_len,
        dim,
        depth,
        heads,
        dropout=0.1,
        num_mtr_targets=209,
    ):
        super().__init__()

        self.token_emb = torch.nn.Embedding(vocab_size, dim)
        self.pos_emb = PositionalEmbedding(max_len, dim)

        layer = torch.nn.TransformerEncoderLayer(
            d_model=dim,
            nhead=heads,
            dim_feedforward=4 * dim,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
        )

        self.encoder = torch.nn.TransformerEncoder(layer, depth)
        self.norm = torch.nn.LayerNorm(dim)

        self.head = torch.nn.Linear(dim, vocab_size)
        self.mtr_head = torch.nn.Sequential(
            torch.nn.Linear(dim, 512),
            torch.nn.ReLU(),
            torch.nn.Linear(512, num_mtr_targets),
        )

    def forward(self, x, return_mtr=False):
        h = self.token_emb(x) + self.pos_emb(x)
        h = self.encoder(h)
        h = self.norm(h)

        logits = self.head(h)

        if return_mtr:
            mtr = self.mtr_head(h[:, 0, :])
            return logits, mtr

        return logits

model = DiffusionTransformer(
    vocab_size=len(vocab),
    max_len=max_len,
    **model_config,
)
state_dict = checkpoint["model_state_dict"]

new_state_dict = {}
for k, v in state_dict.items():
    if k.startswith("pos_emb.pos_emb."):
        new_k = k.replace("pos_emb.pos_emb.", "pos_emb.emb.")
        new_state_dict[new_k] = v
    else:
        new_state_dict[k] = v

model.load_state_dict(new_state_dict)
model.eval()

@torch.no_grad()
def finalize_unmask(x):
    logits = model(x)
    return logits.argmax(dim=-1)


@torch.no_grad()
def sample_selfies_cfg(qed_bin, gamma, steps=30):
    device = "cpu"

    BOS = vocab["[BOS]"]
    EOS = vocab["[EOS]"]
    MASK = vocab["[MASK]"]
    QED_TOKEN = vocab[f"[QED_{qed_bin}]"]
    QED_NULL = vocab["[QED_NULL]"]

    x = torch.full((1, max_len), MASK,)
    x[0, 0] = BOS
    x[0, 1] = QED_TOKEN
    x[0, -1] = EOS

    x_uncond = x.clone()
    x_uncond[0, 1] = QED_NULL

    protected = (x == BOS) | (x == EOS) | (x == QED_TOKEN)

    for step in range(steps):
        logits_u = model(x_uncond)
        logits_c = model(x)

        if step >= steps // 2:
            logits = logits_u + gamma * (logits_c - logits_u)
        else:
            logits = logits_u

        probs = F.softmax(logits, dim=-1)

        for tok in ["[PAD]", "[BOS]", "[EOS]"]:
            probs[..., vocab[tok]] = 0

        probs = probs / probs.sum(dim=-1, keepdim=True)

        sampled = torch.multinomial(
            probs.view(-1, probs.size(-1)), 1
        ).view(1, max_len)

        conf = probs.max(dim=-1).values
        threshold = 0.9 * (1 - step / steps)
        remask = (conf < threshold) & (~protected)

        x = torch.where(remask, MASK, sampled)
        x_uncond = torch.where(remask, MASK, sampled)

        x[0, 0], x[0, 1], x[0, -1] = BOS, QED_TOKEN, EOS
        x_uncond[0, 0], x_uncond[0, 1], x_uncond[0, -1] = BOS, QED_NULL, EOS

    ids = finalize_unmask(x)

    selfie = "".join(
        inv_vocab[i.item()]
        for i in ids[0]
        if inv_vocab[i.item()] not in ["[PAD]", "[BOS]", "[EOS]", "[QED_NULL]"]
        and not inv_vocab[i.item()].startswith("[QED_")
    )

    try:
        smiles = sf.decoder(selfie)
        mol = Chem.MolFromSmiles(smiles)
        qed = QED.qed(mol) if mol else None
    except:
        smiles, qed = None, None

    return smiles, qed

def generate(qed_bin, gamma, n_samples):
    rows = []

    for _ in range(n_samples):
        smiles, qed = sample_selfies_cfg(qed_bin, gamma)

        if smiles is not None and qed is not None:
            rows.append([smiles, round(qed, 4)])

    if len(rows) == 0:
        return [["No valid molecules generated", None]]

    df = pd.DataFrame(rows, columns=["SMILES", "QED"])
    return df
    
def download_csv(df):
    if df is None or len(df) == 0:
        return None

    path = "/tmp/apexoracle_results.csv"
    df.to_csv(path, index=False)
    return path


with gr.Blocks(title="ApexOracle – QED-Controlled Molecular Generation") as demo:
    gr.Markdown(
        """
        # 🧪 ApexOracle – QED-Controlled Molecular Generation  
        **Discrete diffusion model with CFG + Multi-Target Regression**
        """
    )

    # 🔹 Controls (TOP)
    with gr.Column():
        qed_bin = gr.Slider(
            0, 9, step=1, value=8, label="Target QED Bin"
        )
        gamma = gr.Slider(
            0.0, 10.0, step=0.5, value=6.0, label="CFG Gamma"
        )
        n_samples = gr.Slider(
            1, 20, step=1, value=5, label="Number of Molecules"
        )

        submit_btn = gr.Button("Generate", variant="primary")
        clear_btn = gr.Button("Clear")

    gr.Markdown("### Generated Molecules")

    output_df = gr.Dataframe(
        headers=["SMILES", "QED"],
        datatype=["str", "number"],
        row_count=(5, "dynamic"),
    )

    download_btn = gr.Button("Download CSV")
    csv_file = gr.File(label="Download CSV")

    submit_btn.click(
        fn=generate,
        inputs=[qed_bin, gamma, n_samples],
        outputs=output_df,
    )

    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=output_df,
    )

    download_btn.click(
        fn=download_csv,
        inputs=output_df,
        outputs=csv_file,
    )

if __name__ == "__main__":
    demo.launch(
        auth=(os.environ.get("app_user"), os.environ.get("app_pass")),
        ssr_mode=False
    )