gfp/pcomole.py

import argparse
from typing import List, Callable, Optional, Tuple, Dict, Any
import torch
import torch.nn.functional as F
import yaml
from easydict import EasyDict as edict
from tqdm import tqdm
import time
import math

from pathlib import Path
import os
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from generate import build_model_and_stuff, tokenize_input_str, detokenize_output
from objectives import GFPExcitationPred, GFPBrightPred, GFPLength
from constraints import GFP, Length, GFPEmissionPred

import pdb

# ---------------------------------------------------------------------------
# small utilities
# ---------------------------------------------------------------------------

def extract_objective_vector(seqs, objective_models, device):
    values = []

    for obj in objective_models:
        scores = obj(seqs) # list of shape B
        values.append(torch.tensor(scores, device=device, dtype=torch.float32))

    return torch.stack(values, dim=1)  # (B,m)


def compute_scores_print(seqs, objective_models, constraint_models, device, return_scores=False):
    objective_scores = extract_objective_vector(seqs, objective_models, device).squeeze(0)
    scores = [score.item() for score in objective_scores]
    for constraint in constraint_models:
        scores.append(constraint(seqs)[0])
    print(scores)

    if return_scores:
        return scores


# ---------------------------------------------------------------------------
# edit utilities
# ---------------------------------------------------------------------------
@torch.no_grad()
def _sample_multiple_edits_batch(
    x: torch.Tensor,                 # (B, Lmax) padded
    lam_ins: torch.Tensor,           # (B, Lmax)
    logits_ins: torch.Tensor,        # (B, Lmax, V)
    lam_del: torch.Tensor,           # (B, Lmax)
    lam_sub: torch.Tensor,           # (B, Lmax)
    logits_sub: torch.Tensor,        # (B, Lmax, V)
    pad_id: int,
    bos_id: int,
    eos_id: int,
    allowed_tokens: Optional[torch.Tensor] = None,  # 1D LongTensor of vocab ids
    delta: float = 1.0,
    max_len_cap: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Multi-edit small-step proposal:
      - per position i: total rate λ_i = λ_ins + λ_del + λ_sub (after masking invalid ops)
      - fire with p_i = 1 - exp(-delta * λ_i)  (independently per position)
      - if fired: pick op ~ proportional to (λ_ins, λ_del, λ_sub)
      - if op is ins/sub: draw token from softmax(logits_{ins/sub}[i]) (with allowed_tokens masking)
      - apply all fired edits "simultaneously" using a left-to-right scan on the original tokens:
          del: skip token
          sub: replace token
          ins: insert *after* the token

    Returns:
      x_out: (B, Lout) padded
      base_rate: (B,)   relative proposal weight (safe vs underflow): exp(sum_fired log_ratio)
    """
    assert x.dim() == 2, f"x must be (B,Lmax), got {tuple(x.shape)}"
    device = x.device
    B, Lmax = x.shape
    V = logits_ins.shape[-1]
    eps = 1e-30

    if allowed_tokens is not None:
        if not torch.is_tensor(allowed_tokens):
            allowed_tokens = torch.tensor(allowed_tokens, device=device, dtype=torch.long)
        else:
            allowed_tokens = allowed_tokens.to(device=device, dtype=torch.long)

    # masks
    nonpad = (x != pad_id)
    lengths = nonpad.sum(dim=1)  # (B,)
    is_bos = (x == bos_id)
    is_eos = (x == eos_id)

    # mask rates on invalid positions (match your single-edit masking rules)
    lam_ins[:] = 0.0
    ins_rate = lam_ins.clone()
    ins_rate = ins_rate.masked_fill(~nonpad, 0.0)
    ins_rate = ins_rate.masked_fill(is_eos, 0.0)              # no insertion at eos

    del_rate = lam_del.clone()
    del_rate = del_rate.masked_fill(~nonpad, 0.0)
    del_rate = del_rate.masked_fill(is_bos | is_eos, 0.0)      # no delete bos/eos

    sub_rate = lam_sub.clone()
    sub_rate = sub_rate.masked_fill(~nonpad, 0.0)
    sub_rate = sub_rate.masked_fill(is_bos | is_eos, 0.0)      # no sub bos/eos

    # if at cap, disallow insertions
    if max_len_cap is not None:
        at_cap = lengths >= max_len_cap
        if at_cap.any():
            ins_rate = ins_rate.masked_fill(at_cap.unsqueeze(1), 0.0)

    lam_total = ins_rate + del_rate + sub_rate  # (B, Lmax)
    # pdb.set_trace()
    del_rate *= 500

    # pdb.set_trace()
    # fire prob: p = 1 - exp(-delta*lam_total)  (use expm1 for stability)
    a = (delta * lam_total).clamp_min(0.0)
    p_fire = (-torch.expm1(-a)).masked_fill(~nonpad, 0.0)  # (B, Lmax)
    fired = (torch.rand_like(p_fire) < p_fire) & (lam_total > 1e-12) & nonpad

    # op probs per fired position: proportional to rates
    rates3 = torch.stack([ins_rate, del_rate, sub_rate], dim=-1)           # (B,Lmax,3)
    denom = lam_total.unsqueeze(-1).clamp_min(1e-12)
    op_probs = rates3 / denom                                              # (B,Lmax,3)

    # sample op only where fired
    fired_flat = fired.view(-1)
    idx_fired = fired_flat.nonzero(as_tuple=True)[0]                       # (K,)
    op_idx_flat = torch.zeros((B * Lmax,), device=device, dtype=torch.long)  # default 0

    if idx_fired.numel() > 0:
        op_p = op_probs.view(-1, 3)[idx_fired]                             # (K,3)
        op_p = op_p / op_p.sum(dim=1, keepdim=True).clamp_min(1e-12)
        op_idx_flat[idx_fired] = torch.multinomial(op_p, 1).squeeze(1)     # (K,)

    op_idx = op_idx_flat.view(B, Lmax)  # 0=ins,1=del,2=sub

    ins_mask = fired & (op_idx == 0)
    del_mask = fired & (op_idx == 1)
    sub_mask = fired & (op_idx == 2)

    # helper: mask logits to allowed_tokens
    def _mask_logits_full(logits_2d: torch.Tensor) -> torch.Tensor:
        # logits_2d: (K, V)
        if allowed_tokens is None:
            return logits_2d
        add = torch.full_like(logits_2d, -1e9)
        add[:, allowed_tokens] = 0.0
        return logits_2d + add

    # sample tokens for ins/sub at masked positions
    ins_tok = torch.full((B, Lmax), pad_id, device=device, dtype=torch.long)
    sub_tok = torch.full((B, Lmax), pad_id, device=device, dtype=torch.long)

    if ins_mask.any():
        idx_ins = ins_mask.view(-1).nonzero(as_tuple=True)[0]
        logits_sel = logits_ins.view(-1, V)[idx_ins]
        logits_sel = _mask_logits_full(logits_sel)
        q = F.softmax(logits_sel, dim=-1)
        samp = torch.multinomial(q, 1).squeeze(1)
        ins_tok.view(-1)[idx_ins] = samp

    if sub_mask.any():
        idx_sub = sub_mask.view(-1).nonzero(as_tuple=True)[0]
        logits_sel = logits_sub.view(-1, V)[idx_sub]
        logits_sel = _mask_logits_full(logits_sel)
        q = F.softmax(logits_sel, dim=-1)
        samp = torch.multinomial(q, 1).squeeze(1)
        sub_tok.view(-1)[idx_sub] = samp

    # -------------------------
    # base_rate: (B,) relative weight to avoid underflow
    # For each fired position:
    #   ratio = ((1-exp(-a)) / exp(-a)) * P(op | fired) * P(token | op)
    #         = (exp(a)-1) * (rate/total) * token_prob
    # log_ratio = log(expm1(a)) + log(op_prob) + log(token_prob)
    # -------------------------
    base_log = torch.zeros((B,), device=device, dtype=torch.float32)

    if idx_fired.numel() > 0:
        b_idx = (idx_fired // Lmax).to(torch.long)                # (K,)
        op_choice = op_idx_flat[idx_fired].to(torch.long)         # (K,)

        a_sel = a.view(-1)[idx_fired].to(torch.float32)           # (K,)
        log_expm1 = torch.log(torch.expm1(a_sel).clamp_min(eps))   # (K,)

        op_p_sel = op_probs.view(-1, 3)[idx_fired].to(torch.float32)
        op_p_sel = op_p_sel / op_p_sel.sum(dim=1, keepdim=True).clamp_min(1e-12)
        op_prob_sel = op_p_sel.gather(1, op_choice.view(-1, 1)).squeeze(1).clamp_min(eps)
        log_op = torch.log(op_prob_sel)

        log_tok = torch.zeros_like(log_op)

        # token prob for ins
        ins_k = (op_choice == 0)
        if ins_k.any():
            idx_ins_k = idx_fired[ins_k]
            tok_sel = ins_tok.view(-1)[idx_ins_k]
            logits_sel = logits_ins.view(-1, V)[idx_ins_k]
            logits_sel = _mask_logits_full(logits_sel)
            logq = F.log_softmax(logits_sel, dim=-1)
            log_tok[ins_k] = logq.gather(1, tok_sel.view(-1, 1)).squeeze(1)

        # token prob for sub
        sub_k = (op_choice == 2)
        if sub_k.any():
            idx_sub_k = idx_fired[sub_k]
            tok_sel = sub_tok.view(-1)[idx_sub_k]
            logits_sel = logits_sub.view(-1, V)[idx_sub_k]
            logits_sel = _mask_logits_full(logits_sel)
            logq = F.log_softmax(logits_sel, dim=-1)
            log_tok[sub_k] = logq.gather(1, tok_sel.view(-1, 1)).squeeze(1)

        log_ratio = log_expm1 + log_op + log_tok
        base_log.scatter_add_(0, b_idx, log_ratio)

    base_rate = torch.exp(base_log).clamp_min(0.0)  # (B,)

    # -------------------------
    # apply edits to build new padded batch
    # -------------------------
    new_seqs = []
    new_lens = []

    for b in range(B):
        seq = x[b]
        valid = (seq != pad_id)
        tokens = seq[valid].tolist()
        Lb = len(tokens)

        if Lb == 0:
            out_tokens = [eos_id]
        else:
            out_tokens = []
            for i in range(Lb):
                t_i = tokens[i]

                if i < Lmax and bool(del_mask[b, i].item()):
                    continue

                if i < Lmax and bool(sub_mask[b, i].item()):
                    out_tokens.append(int(sub_tok[b, i].item()))
                else:
                    out_tokens.append(int(t_i))

                if i < Lmax and bool(ins_mask[b, i].item()):
                    out_tokens.append(int(ins_tok[b, i].item()))

            if len(out_tokens) == 0 or out_tokens[-1] != eos_id:
                out_tokens.append(eos_id)

        if max_len_cap is not None and len(out_tokens) > max_len_cap:
            out_tokens = out_tokens[:max_len_cap]
            if out_tokens[-1] != eos_id:
                out_tokens[-1] = eos_id

        new_seqs.append(torch.tensor(out_tokens, device=device, dtype=torch.long))
        new_lens.append(len(out_tokens))

    Lout = max(1, max(new_lens) if new_lens else 1)
    x_out = torch.full((B, Lout), pad_id, device=device, dtype=x.dtype)
    for b, s in enumerate(new_seqs):
        x_out[b, : s.numel()] = s

    return x_out, base_rate

# ---------------------------------------------------------------------------
# ATC + G_T
# ---------------------------------------------------------------------------
def _augmented_tchebycheff(
    f_vals: torch.Tensor,
    w: torch.Tensor,
    rho: float,
    z: torch.Tensor,
) -> torch.Tensor:
    diff = f_vals - z
    term1 = torch.min(w * diff, dim=1).values
    term2 = rho * torch.sum(w * diff, dim=1)
    return term1 + term2


def _G_T(
    x: torch.Tensor,
    objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]],
    constraint_models: List[Callable[[torch.Tensor], torch.Tensor]],
    w: torch.Tensor,
    rho: float,
    z: torch.Tensor,
    beta: float,
    tokenizer, ws_for_invalid=False
):
    device = x.device
    seqs = [seq.replace(' ', '') for seq in tokenizer.batch_decode(x, skip_special_tokens=True)]

    constraint_results = []
    for constraint in constraint_models:
        res = constraint(seqs)
        constraint_results.append(res)
    
    # if ws_for_invalid is False:
    #     pdb.set_trace()
    constraint_results = torch.tensor(constraint_results, device=device)

    survived_seq_indices = (constraint_results == 1).all(dim=0).nonzero(as_tuple=True)[0]
    survived_seqs = [seqs[idx] for idx in survived_seq_indices.tolist()] # (B')

    weighted_sum_full = torch.full((len(seqs),), float("-inf"), device=device)
    G_full = torch.full((len(seqs),), float("-inf"), device=device)

    # objectives
    if ws_for_invalid:
        f_vals = extract_objective_vector(seqs, objective_models, x.device)
        weighted_sum_full = torch.sum(w * f_vals, dim=1) 
        u_atc = _augmented_tchebycheff(f_vals, w, rho, z)
        G = beta * u_atc
        G_full[survived_seq_indices] = G[survived_seq_indices]
    else:
        if survived_seq_indices.numel() > 0:
            f_vals = extract_objective_vector(survived_seqs, objective_models, x.device)    # (B', m)
            u_atc = _augmented_tchebycheff(f_vals, w, rho, z)   # (B',)
            G = beta * u_atc # (B',)
            weighted_sum = torch.sum(w * f_vals, dim=1)  # (B',)
            G_full[survived_seq_indices] = G
            weighted_sum_full[survived_seq_indices] = weighted_sum

    # return full-size tensors (B,)
    return G_full, weighted_sum_full



# ---------------------------------------------------------------------------
# rollout
# ---------------------------------------------------------------------------
@torch.no_grad()
def short_rollout_batch(
    model,
    x0: torch.Tensor,            # (B, Lmax) padded
    time_grid: torch.Tensor,
    start_idx: int,
    pad_id: int,
    bos_id: int,
    eos_id: int,
    allowed_tokens: Optional[torch.Tensor],
    max_len_cap: Optional[int],
    num_rollouts: int = 1,
    num_steps: int =32
) -> torch.Tensor:
    """
    Returns:
      xT: (B*num_rollouts, Lmax)
    Grouping:
      xT[i*num_rollouts:(i+1)*num_rollouts] corresponds to candidate i.
    """
    device = x0.device
    B, Lmax = x0.shape

    # repeat each candidate num_rollouts times (grouped)
    x = x0.repeat_interleave(num_rollouts, dim=0)  # (B*num_rollouts, Lmax)

    # rollout in batch
    for j in range(start_idx + 1, time_grid.numel()):
        t_j = time_grid[j].view(1).to(device)

        mask = (x != pad_id) 

        lam_ins, logits_ins, lam_del, lam_sub, logits_sub, *_ = model(x_t=x, mask=mask, t=t_j)

        x, _ = _sample_multiple_edits_batch(
            x,
            lam_ins, logits_ins,
            lam_del, lam_sub, logits_sub,
            pad_id, bos_id, eos_id,
            allowed_tokens,
            delta=float(1/(num_steps-1)),
            max_len_cap=max_len_cap,
        )

    return x

# ---------------------------------------------------------------------------
# finalizer
# ---------------------------------------------------------------------------
def _finalize_from_last(
    model,
    x_last: torch.Tensor,
    time_grid: torch.Tensor,
    last_step: int,
    pad_id: int,
    bos_id: int,
    eos_id: int,
    allowed_tokens: Optional[torch.Tensor],
    objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]],
    constraint_models: List[Callable[[torch.Tensor], torch.Tensor]],
    w: torch.Tensor,
    rho: float,
    ref_z: torch.Tensor,
    beta_final: float,
    max_len_cap: Optional[int] = None,
    num_final_rollouts: int = 16,
    num_steps: int = 32,
    tokenizer=None
) -> torch.Tensor:

    logG_last, _ = _G_T(x_last, objective_models, constraint_models, w, rho, ref_z, beta_final, tokenizer, ws_for_invalid=False)

    x_Ts = short_rollout_batch(model, x_last, time_grid, last_step, pad_id, bos_id, eos_id, allowed_tokens, max_len_cap, num_final_rollouts, num_steps)
    logG, _, = _G_T(x_Ts, objective_models, constraint_models, w, rho, ref_z, beta_final, tokenizer, ws_for_invalid=False)

    idx = torch.isfinite(logG).nonzero(as_tuple=True)[0].tolist()
    
    if len(idx) == 0 or torch.max(logG) < logG_last:
        return x_last, logG_last
    else:
        best_idx = torch.argmax(logG).item()
        best_seq = x_Ts[best_idx].unsqueeze(0)
        return best_seq, logG[best_idx]

def cope_strict(
    model,
    x0: torch.Tensor,
    *,
    pad_id: int,
    bos_id: int,
    eos_id: int,
    allowed_tokens: Optional[torch.Tensor],
    objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]],
    constraint_models: List[Callable[[torch.Tensor], torch.Tensor]],
    w: torch.Tensor,
    rho: float,
    ref_z: torch.Tensor,
    beta_start: float = 1.0,
    beta_end: float = 3.0,
    num_steps: int = 32,
    num_candidates: int = 8,
    num_rollouts: int = 4,
    max_len_cap: Optional[int] = None,
    device: Optional[torch.device] = None,
    num_final_rollouts: int = 16,
    cfg, tokenizer,
    diagnostics=None,
) -> torch.Tensor:
    if device is None:
        device = x0.device
    x = x0.clone().to(device)
    x_input = x0.clone().to(device)
    time_grid = torch.linspace(0.0, 1.0, steps=num_steps, device=device)
    last_timestep = 0

    best_terminal = None
    best_terminal_logG = float("-inf")

    def _sync():
        if diagnostics is not None and torch.cuda.is_available():
            torch.cuda.synchronize()

    def _tic():
        _sync()
        return time.perf_counter()

    def _toc(bucket, t0):
        if diagnostics is None:
            return
        _sync()
        getattr(diagnostics, bucket).add(time.perf_counter() - t0)

    with torch.no_grad():
        for step in tqdm(range(num_steps - 1)):
            if diagnostics is not None:
                diagnostics.num_executed_steps += 1
            t = time_grid[step].view(1)
            frac = step / max(1, (num_steps - 1))
            beta_t = beta_start + (beta_end - beta_start) * frac

            # model forward + candidate proposal
            t0 = _tic()
            mask = (x != pad_id)
            lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type = model(x_t=x, mask=mask, t=t)

            candidates = [x.squeeze(0)] # compute the scores of current sequence with the candidates
            base_rates = []
            for _ in range(num_candidates):
                cand_seq, base_rate = _sample_multiple_edits_batch(
                    x,
                    lam_ins, logits_ins,
                    lam_del, lam_sub, logits_sub,
                    pad_id, bos_id, eos_id,
                    allowed_tokens,
                    delta=float(1/(num_steps-1)),
                    max_len_cap=max_len_cap,
                )
                if not torch.equal(cand_seq, x):
                    candidates.append(cand_seq.squeeze(0))
                    base_rates.append(base_rate)
            candidates = list(set(candidates))
            batch_candidates = torch.nn.utils.rnn.pad_sequence(candidates, batch_first=True, padding_value=pad_id)
            _toc("proposal", t0)
            print("Initial Candidates: ", len(candidates))
            # We only want the survived candidates to improve the objective weights
            t0 = _tic()
            cand_logG, cand_ws = _G_T(batch_candidates, objective_models, constraint_models, w, rho, ref_z, beta_t, tokenizer, ws_for_invalid=True)
            _toc("screening", t0)

            curr_logG = cand_logG[0]
            curr_ws = cand_ws[0]
            cand_logG = cand_logG[1:]
            cand_ws = cand_ws[1:]
            batch_candidates = batch_candidates[1:, :]

            if len(batch_candidates) == 0:
                if diagnostics is not None:
                    diagnostics.num_steps_no_selectable_candidate += 1
                continue

            improve_idx = (cand_ws > curr_ws).nonzero(as_tuple=True)[0]
            survived_candidates = batch_candidates[improve_idx, :]
            base_rates = [base_rates[i] for i in improve_idx] # (num_survived_candidates,)
            print("Num Candidates Survived: ", len(improve_idx))
            if len(improve_idx) == 0:
                if diagnostics is not None:
                    diagnostics.num_steps_no_selectable_candidate += 1
                continue
            
            # Keep all the rollout terminal sequences in one batch
            t0 = _tic()
            x_Ts = short_rollout_batch(model, survived_candidates, time_grid, step, pad_id, bos_id, eos_id, allowed_tokens, max_len_cap, num_rollouts, num_steps)
            _toc("rollout", t0)
            if diagnostics is not None:
                diagnostics.num_rollouts_simulated += int(survived_candidates.shape[0] * num_rollouts)
                diagnostics.num_evaluated_candidates += int(survived_candidates.shape[0])

            # Constraints are taken into account for the terminal sequences
            t0 = _tic()
            logG, _, = _G_T(x_Ts, objective_models, constraint_models, w, rho, ref_z, beta_t, tokenizer, ws_for_invalid=False)
            _toc("terminal_oracle", t0)

            if diagnostics is not None:
                logG_reshaped_tmp = logG.reshape(survived_candidates.shape[0], num_rollouts)
                n_all_inf = int((~torch.isfinite(logG_reshaped_tmp).any(dim=1)).sum().item())
                diagnostics.num_candidates_all_rollouts_infeasible += n_all_inf
                if torch.isfinite(logG).any():
                    diagnostics.observed_any_feasible_terminal = True

            curr_best_terminal_logG = torch.max(logG)
            if best_terminal_logG <= curr_best_terminal_logG:
                best_terminal_idx = torch.argmax(logG)
                best_terminal = x_Ts[best_terminal_idx].unsqueeze(0)
                best_terminal_logG = curr_best_terminal_logG

                best_terminal_smiles_seq = tokenizer.batch_decode(best_terminal, skip_special_tokens=True)[0].replace(' ', '')
                print("\nSaved Best Terminal: ", best_terminal_smiles_seq)
                print("Saved Best Terminal Length: ", best_terminal.shape[1])
                print("Saved Best Terminal logG: ", best_terminal_logG)

            
            t0 = _tic()
            logG = logG.reshape(survived_candidates.shape[0], num_rollouts)
            log_h_hat = torch.logsumexp(logG, dim=1) - math.log(num_rollouts)   # (num_survived_candidates,)
            idx = (logG.max(dim=1).values > curr_logG).nonzero(as_tuple=True)[0]
            final_survived_candidates = survived_candidates[idx, :]
            if len(final_survived_candidates) == 0:
                _toc("selection", t0)
                if diagnostics is not None:
                    diagnostics.num_steps_no_selectable_candidate += 1
                continue

            # Doob-like transform
            log_h_hat = log_h_hat[idx]
            base_rates_t = torch.tensor([base_rates[i] for i in idx.tolist()], device=device, dtype=torch.float32)
            log_base = 0.5 * torch.log(base_rates_t.clamp_min(1e-30))
            log_weights = log_base + log_h_hat
            probs = torch.softmax(log_weights, dim=0) 
            if torch.isnan(probs).any():
                pdb.set_trace()
            selected_idx = torch.multinomial(probs, 1).item()
            x = final_survived_candidates[selected_idx].unsqueeze(0)
            _toc("selection", t0)
            seq = tokenizer.batch_decode(x, skip_special_tokens=True)[0].replace(' ', '')
            print(seq)
            print("Current Length: ", len(seq))
            compute_scores_print([seq], objective_models, constraint_models, device)
            last_timestep = step

        # finalize
        t0 = _tic()
        x_final_rollout, logG_final_rollout = _finalize_from_last(
            model,
            x,
            time_grid,
            last_timestep,
            pad_id,
            bos_id,
            eos_id,
            allowed_tokens,
            objective_models,
            constraint_models,
            w,
            rho,
            ref_z,
            beta_end,
            max_len_cap=max_len_cap,
            num_final_rollouts=num_final_rollouts,
            num_steps=num_steps,
            tokenizer=tokenizer,
        )
        # Finalize spans rollout+oracle; attribute to those buckets proportionally by calling pattern:
        # keep finalize time in terminal_oracle + rollout jointly via selection bucket unused —
        # attribute whole finalize block to terminal_oracle for accounting simplicity is wrong.
        # Split: short_rollout is inside _finalize; count finalize wall time under rollout+terminal
        # by putting it in terminal_oracle only if we cannot split without changing _finalize.
        # Prefer attributing to rollout (contains short_rollout_batch) + terminal via half split.
        if diagnostics is not None:
            _sync()
            elapsed = time.perf_counter() - t0
            diagnostics.rollout.add(elapsed * 0.5)
            diagnostics.terminal_oracle.add(elapsed * 0.5)
            if torch.isfinite(logG_final_rollout):
                diagnostics.observed_any_feasible_terminal = True

        if logG_final_rollout >= best_terminal_logG:
            best_terminal = x_final_rollout

    if diagnostics is not None:
        if best_terminal is None or not diagnostics.observed_any_feasible_terminal:
            diagnostics.run_failure = True
            diagnostics.used_original_input_fallback = True
            if best_terminal is None:
                best_terminal = x_input
        if best_terminal is not None:
            diagnostics.output_sequence = tokenizer.batch_decode(
                best_terminal, skip_special_tokens=True
            )[0].replace(' ', '')

    return best_terminal

# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root_dir", type=str, default=str(REPO_ROOT / "gfp" / "FPredX"))
    parser.add_argument("--config", type=str, default=str(REPO_ROOT / "configs" / "config_gfp.yaml"))
    parser.add_argument("--ckpt", type=str, required=True)
    parser.add_argument("--input", type=str, required=True)
    parser.add_argument("--num_steps", type=int, default=32)
    parser.add_argument("--max_len_cap", type=int, default=None)
    parser.add_argument("--num_candidates", type=int, default=10)
    parser.add_argument("--num_rollouts", type=int, default=5)
    parser.add_argument("--beta_start", type=float, default=1.0)
    parser.add_argument("--beta_end", type=float, default=3.0)
    parser.add_argument("--alpha_start", type=float, default=0.8)
    parser.add_argument("--alpha_end", type=float, default=0.1)
    parser.add_argument("--num_final_rollouts", type=int, default=16)
    parser.add_argument("--objective_weights", type=float, nargs='+')
    parser.add_argument("--ref_z", type=float, nargs='+')
    parser.add_argument("--rho", type=float, default=1)
    parser.add_argument("--laser", type=float, default=488, help="Keep excitation maximum near the laser you actually have")
    parser.add_argument("--output_file", type=str, default='sample.csv')
    args = parser.parse_args()
    if args.output_file:
        out_dir = os.path.dirname(os.path.abspath(args.output_file))
        if out_dir:
            os.makedirs(out_dir, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    with open(args.config, "r") as f:
        cfg = edict(yaml.safe_load(f))

    editflow, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id = build_model_and_stuff(cfg, device)

    ckpt = torch.load(args.ckpt, map_location=device)
    editflow.load_state_dict(ckpt["state_dict"], strict=False)
    model = editflow.model.to(device)
    model.eval()

    x0 = tokenize_input_str(args.input, tokenizer, bos_id, eos_id, device)

    allowed_tokens = torch.tensor(
        [tok for tok in source_dist._allowed_tokens if tok not in (eps_id,) and tok not in range(24,33)],
        device=device,
        dtype=torch.long,
    )

    length = GFPLength(args.input)
    excitation = GFPExcitationPred(root_dir=args.root_dir, laser=args.laser)
    brightness = GFPBrightPred(root_dir=args.root_dir)

    file_name = args.output_file.split('/')[-1].split('.')[0]
    if file_name == 'length_excitation_brightness':
        objective_models = [length, excitation, brightness]
        # args.objective_weights = [3,1,1]
    elif file_name == 'length_brightness':
        objective_models = [length, brightness]
        # args.objective_weights = [3,1]
    elif file_name == 'length_excitation':
        objective_models = [length, excitation]
        # args.objective_weights = [3,1]
    elif file_name == 'length':
        objective_models = [length]
    else:
        objective_models = [length, excitation, brightness]
    num_objectives = len(objective_models)
    if not args.objective_weights:
        objective_weights = torch.tensor([1.0 / num_objectives] * num_objectives).to(device)
    else:
        objective_weights = torch.tensor(args.objective_weights).to(device)

    if not args.ref_z:
        ref_z = torch.zeros(num_objectives).to(device)
    else:
        ref_z = torch.tensor(args.ref_z).to(device)

    gfp_hard_constraint = GFP(device)
    emission_soft_constraint = GFPEmissionPred(root_dir=args.root_dir)
    length_soft_constraint = Length(args.input)
    constraint_models = [length_soft_constraint, gfp_hard_constraint, emission_soft_constraint]
    # pdb.set_trace()
    for _ in range(100):
        x_T = cope_strict(
            model=model,
            x0=x0,
            pad_id=pad_id,
            bos_id=bos_id,
            eos_id=eos_id,
            allowed_tokens=allowed_tokens,
            objective_models=objective_models,
            constraint_models=constraint_models,
            w=objective_weights,
            rho=0.5,
            ref_z=ref_z,
            beta_start=args.beta_start,
            beta_end=args.beta_end,
            num_steps=args.num_steps,
            num_candidates=args.num_candidates,
            num_rollouts=args.num_rollouts,
            max_len_cap=args.max_len_cap,
            num_final_rollouts=args.num_final_rollouts,
            cfg=cfg, tokenizer=tokenizer
        )
        
        out_str = tokenizer.batch_decode(x_T, skip_special_tokens=True)[0].replace(' ', '')
        print("----------------------------")
        # print(f"Initial Sequence: {args.input}\n")
        # print(f"Initial Scores:")
        # compute_scores_print([args.input], objective_models, constraint_models, device)

        print(f"\nDesigned Sequence: {out_str}\n")
        print("Final scores:")
        scores = compute_scores_print([out_str], objective_models, constraint_models, device, return_scores=True)

        with open(args.output_file, 'a') as f:
            f.write(f"{out_str},{len(out_str)}")
            for score in scores:
                f.write(f",{score}")
            f.write('\n')


if __name__ == "__main__":
    main()