"""Several objectives at once: MOG-DFM's cone and pCoMole's rollout.

A peptide is scored on binding, solubility and half-life together, and no dataset labels one
sequence for all of them, so the single value function of Discrete Guidance is out of reach.
MOG-DFM drops the ratio form and builds a tilt out of a rank inside each objective and a direction
against a weight vector, with a cone that removes the candidates pointing the wrong way. pCoMole
keeps the ratio and replaces the property likelihood by an exponentiated worst-objective utility,
and its value at a candidate is the average terminal reward over rollouts of the base flow.

Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).

Chapter 5, Sections 5.30 to 5.32. The single-objective tilt is in ``lecture_5/guidance.py``. The
rollout runs on the synthetic DNA of ``cis6270.data``, scored by GC fraction against AT fraction.
Prints the Pareto front, the chapter's MOG-DFM step with its cone, the chapter's pCoMole rollout,
and a rollout on DNA whose value reverses the base rates.

    python lecture_5/multi_objective.py                     # 600 rollouts per candidate on DNA
    python lecture_5/multi_objective.py --steps 4000        # a tighter estimate of the second value
    python lecture_5/multi_objective.py --steps 20 --quiet  # twenty rollouts, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import DNA, DNA_SIZE, MOTIF, decode_dna, dna_splits
from cis6270.runner import common_parser, print_report, save_run, set_seed

MASK = DNA_SIZE  # Index 4 is the mask symbol m, still held by a candidate off its own position.
OMEGA = torch.tensor([0.5, 0.5], dtype=torch.double)  # Equal weight on the two objectives of the DNA rollout.
REFERENCE = torch.tensor([0.4, 0.3], dtype=torch.double)  # The levels r_k the objectives sit above.
BETA = 2.0  # The inverse temperature of the terminal reward.


# %% 02. Several objectives at once
#
# With several objectives no single sequence is best. A candidate is on the Pareto front when no
# other candidate matches it everywhere and beats it somewhere, and a weight vector picks one point
# of that front by collapsing the scores to a single weighted number.


def dominated(scores, row):
    """True when some other candidate is at least as good everywhere and better somewhere."""
    others = torch.cat([scores[:row], scores[row + 1:]])  # Every candidate except this one.
    weak = (others >= scores[row]).all(dim=1)  # At least as good on every objective.
    strict = (others > scores[row]).any(dim=1)  # Strictly better on at least one.
    return bool((weak & strict).any())


def pareto_front(scores):
    """The indices of the candidates no other candidate dominates. scores: [M, N] -> list."""
    return [row for row in range(len(scores)) if not dominated(scores, row)]


# %% 03. MOG-DFM: rank, direction and a cone
#
# Without a joint likelihood we have no ratio to compute, so we score every candidate letter twice.
# A rank inside each objective removes the units, the inner product with the weight vector measures
# direction, and the two z-scored terms exponentiate into a tilt on the pretrained rate.


def rank_scores(changes):
    """I_n(y^i, x), the rank of each candidate's change within each objective, over the alphabet size."""
    order = changes.argsort(dim=0).argsort(dim=0) + 1  # Rank one for the smallest change.
    return order.to(changes.dtype) / changes.shape[0]  # [K, N] in {1/K, ..., 1}.


def z_score(values):
    """Norm[.], the z-score across the candidate letters at one position."""
    spread = values.std(unbiased=False).clamp(min=1e-12)  # The population deviation, held off zero.
    return (values - values.mean()) / spread  # Zero everywhere when no candidate stands out.


def combined_score(changes, importance, weights, lam=1.0):
    """Delta S = Norm[mean_n i_n I_n] + lambda Norm[Delta s . omega]. changes: [K, N] -> [K]."""
    ranks = (rank_scores(changes) * importance).mean(dim=1)  # Importance-weighted average of ranks.
    direction = changes @ weights  # [K] the directional score of each candidate.
    return z_score(ranks) + lam * z_score(direction), ranks, direction


# %% 04. Accepting a move only inside a cone
#
# A z-scored score puts about half the candidates above zero, so the tilt alone promotes the least
# bad letter at a position where every letter moves away from the requested trade-off. The angle
# between a candidate's score change and omega removes those, and a controller on the logarithm of
# the angle holds the surviving fraction near one minus the target rejection rate.


def cone_angles(changes, weights):
    """alpha^i = arccos(Delta s . omega / (||Delta s|| ||omega||)), in degrees."""
    cosine = (changes @ weights) / (changes.norm(dim=1) * weights.norm())  # [K].
    return torch.rad2deg(torch.arccos(cosine.clamp(-1.0, 1.0))), cosine


def update_cone(phi, rejected, target=0.2, gain=0.5, low=30.0, high=90.0):
    """Phi_{t+h} = clip(Phi_t exp(eta (r - tau)), Phi_min, Phi_max), proportional control on log Phi."""
    return min(max(phi * math.exp(gain * (rejected - target)), low), high)  # Wider when it rejects more.


# %% 05. pCoMole: a reward in place of a likelihood
#
# If we score a sequence by its weakest objective, the sampler cannot trade one property away, and
# a small multiple of the sum breaks ties. The exponentiated utility times a feasibility indicator
# is the terminal reward, and the value of a candidate is the average reward over rollouts of the
# base flow started from it.


def utility(scores, weights, reference, rho=0.1):
    """U(x) = min_k omega_k (f_k(x) - r_k) + rho sum_k omega_k (f_k(x) - r_k)."""
    gaps = weights * (scores - reference)  # [N] the weighted gap on each objective.
    return float(gaps.min() + rho * gaps.sum())  # The minimum, with the tie-breaking correction.


def terminal_reward(scores, weights, reference, beta, feasible, rho=0.1):
    """G(x) = exp(beta U(x)) 1[x in F], zero for an infeasible sequence."""
    return math.exp(beta * utility(scores, weights, reference, rho)) * float(feasible)


def rollout_value(rollouts, feasible, weights, reference, beta, rho=0.1):
    """hat h_t(y_c), the average terminal reward over R continuations of the base flow."""
    rewards = [terminal_reward(s, weights, reference, beta, f, rho) for s, f in zip(rollouts, feasible)]
    return sum(rewards) / len(rewards), rewards  # A candidate with no feasible ending scores zero.


# %% 06. A rollout on the synthetic DNA
#
# The base flow of a rollout is the mask-source sampler of this chapter with the letter frequency of
# the data as its denoiser, so a continuation costs no network call. The GC fraction and the AT
# fraction of the finished sequence pull against each other. The feasible set is the sequences that
# carry the planted motif.


def continuations(partial, frequency, rollouts, generator=None):
    """Fill every masked position of a candidate from the base flow. partial: [L] -> [R, L]."""
    grid = partial[None, :].repeat(rollouts, 1)  # [R, L] one copy of the candidate per rollout.
    masked = grid == MASK  # [R, L] the positions the flow still owes a letter.
    draws = torch.multinomial(frequency, int(masked.sum()), replacement=True, generator=generator)
    return grid.masked_scatter(masked, draws)  # [R, L] the terminal sequences of the rollouts.


def dna_objectives(sequences):
    """The GC fraction and the AT fraction of each sequence. The two sum to one. [R, L] -> [R, 2]."""
    gc = ((sequences == 1) | (sequences == 2)).to(torch.get_default_dtype()).mean(dim=1)  # C and G.
    return torch.stack([gc, 1.0 - gc], dim=1)  # [R, 2] raising one lowers the other.


def carries_motif(sequences):
    """One per sequence that holds the planted motif, the hard constraint here. [R] ints."""
    return [int(MOTIF in sequence) for sequence in decode_dna(sequences)]


def dna_rollout_value(partial, frequency, rollouts, weights, reference, beta, generator=None):
    """hat h_t at one candidate, from ``rollouts`` continuations of the base flow. Returns a float."""
    finished = continuations(partial, frequency, rollouts, generator)  # [R, L].
    scores = dna_objectives(finished)  # [R, 2] the two objectives at the end of each rollout.
    value, _ = rollout_value(list(scores), carries_motif(finished), weights, reference, beta)
    return value  # Zero when no continuation of this candidate is feasible.


# %% 07. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    args = parser.parse_args()
    set_seed(args.seed)

    # Section 5.30: four candidates on two objectives, and the two weight vectors.
    candidates = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]], dtype=torch.double)
    front = pareto_front(candidates)
    report = {
        "pareto_front": front,  # Candidates 1 and 2; 3 and 4 are dominated.
        "weighted_at_0.6_0.4": [round(float(s), 3) for s in candidates @ torch.tensor(
            [0.6, 0.4], dtype=torch.double)],  # 0.60 for the first and 0.64 for the second.
        "weighted_at_0.8_0.2": [round(float(s), 3) for s in candidates @ torch.tensor(
            [0.8, 0.2], dtype=torch.double)],  # 0.70 and 0.62, a reversal of the order.
    }

    # Section 5.31.2: one MOG-DFM step at a position holding A, with three candidate letters.
    changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]], dtype=torch.double)
    importance = torch.tensor([1.2, 0.8], dtype=torch.double)
    weights = torch.tensor([0.6, 0.4], dtype=torch.double)
    delta_s, ranks, direction = combined_score(changes, importance, weights)
    report["rank_average"] = [round(float(r), 3) for r in ranks]  # (0.500, 0.700, 0.850, 0.450).
    report["directional_scores"] = [round(float(d), 3) for d in direction]  # (0, 0.100, 0.090, -0.028).
    report["combined_score"] = [round(float(s), 3) for s in delta_s]  # (-1.511, 1.541, 2.297, -2.327).
    posterior = torch.tensor([0.3, 0.3, 0.3, 0.1], dtype=torch.double)  # The denoiser at t = 0.7.
    base = posterior[1:] / (1.0 - 0.7)  # The pretrained rates out of A toward C, G and T.
    report["base_rates"] = [round(float(r), 4) for r in base]  # (1.000, 1.000, 0.333).
    tilted = base * delta_s[1:].exp()
    report["guided_rates"] = [round(float(r), 3) for r in tilted]  # (4.667, 9.948, 0.033).
    angles, cosines = cone_angles(changes[1:], weights)
    report["cone_cosines"] = [round(float(c), 4) for c in cosines]  # (0.6727, 0.7894, -0.3032).
    report["cone_angles_deg"] = [round(float(a), 1) for a in angles]  # (47.7, 37.9, 107.7).
    inside = angles <= 60.0  # The cone test at Phi = 60 degrees.
    survivors = tilted[inside]
    report["after_cone"] = [round(float(r), 3) for r in survivors / survivors.sum()]  # (0.319, 0.681).
    rejected = float((~inside).to(torch.get_default_dtype()).mean())  # One candidate of three.
    report["next_cone_angle"] = round(update_cone(60.0, rejected), 1)  # 64.1.

    # Section 5.32.2: three candidate edits, each rolled out four times.
    omega, reference = torch.tensor([0.5, 0.5], dtype=torch.double), torch.zeros(2, dtype=torch.double)
    edits = [
        ([(0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5)], [1, 1, 0, 1], 0.6),
        ([(1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1)], [1, 1, 1, 1], 1.2),
        ([(0.4, 0.4)] * 4, [0, 0, 0, 0], 2.0),
    ]
    values = [rollout_value([torch.tensor(s, dtype=torch.double) for s in rollouts],
                            feasible, omega, reference, 2.0)[0] for rollouts, feasible, _ in edits]
    base_rates = torch.tensor([rate for _, _, rate in edits], dtype=torch.double)
    products = torch.tensor(values) * base_rates  # The base rate times the continuation value.
    report["rollout_values"] = [round(v, 3) for v in values]  # (1.390, 1.181, 0).
    report["rollout_products"] = [round(float(p), 3) for p in products]  # (0.834, 1.417, 0).
    report["rollout_selection"] = [round(float(p), 3) for p in products / products.sum()]
    report["rollout_base_selection"] = [round(float(p), 3) for p in base_rates / base_rates.sum()]

    # The same rollout on the shared DNA, from a partly revealed sequence with ACG already in place.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
    frequency = counts / counts.sum()  # [4] the denoiser of the base flow, GC-rich like the data.
    # The base rates below are kappa_dot/(1 - kappa) times that denoiser, read at t = 0.3.
    revealed = torch.full((args.length,), MASK, dtype=torch.long)
    revealed[2:5] = torch.tensor([0, 1, 2])  # Positions 3, 4 and 5 hold A, C and G.
    finish = {letter: revealed.clone() for letter in "TC"}
    for letter, candidate in finish.items():
        candidate[5] = DNA.index(letter)  # Position 6 takes T, completing ACGT, or C, and C does not.
    seen = {letter: dna_rollout_value(candidate, frequency, args.steps, OMEGA, REFERENCE, BETA)
            for letter, candidate in finish.items()}
    rates = torch.tensor([float(frequency[DNA.index(letter)]) for letter in finish]) / (1.0 - 0.3)
    guided = rates * torch.tensor([seen[letter] for letter in finish])  # The base rate times hat h.
    report["dna_rollout_value_T"] = round(seen["T"], 4)  # Every continuation carries the motif.
    report["dna_rollout_value_C"] = round(seen["C"], 4)  # Only the ones that plant it elsewhere.
    report["dna_base_rates_T_C"] = [round(float(r), 4) for r in rates]  # C has the higher base rate.
    report["dna_base_selection"] = [round(float(r), 4) for r in rates / rates.sum()]
    report["dna_guided_selection"] = [round(float(r), 4) for r in guided / guided.sum()]
    print_report("several objectives at once", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()