"""Probability on discrete and continuous states, in the chapter's own numbers.

A distribution allocates one unit of mass over the states, and expectation averages a property
under that allocation. Variance and covariance record how the values fluctuate, and conditioning
and Bayes' rule revise the allocation once we observe something. The chapter works every
example below by hand, so each section prints the hand value beside the value this file computes.

Chapter 1, Sections 1.3 to 1.11 of the notes. Runs on the two-token DNA distribution of
Section 1.3.1, the three peptides of Section 1.7.1 and the motif-and-binding table of
Section 1.9.1. Prints the chapter's number, the computed number and the gap between them.

    python lecture_1/probability.py                   # every worked example
    python lecture_1/probability.py --samples 100000  # a longer Monte Carlo ladder
    python lecture_1/probability.py --quiet           # the closing report only
"""

# %% 01. Imports and configuration
from __future__ import annotations

import argparse
import math

import torch

from cis6270.runner import print_report, save_run, set_seed

torch.set_default_dtype(torch.float64)  # Double precision, because the chapter prints five decimal places.


STATES = ["AA", "AC", "CA", "CC"]  # The state space X of Section 1.3.1, with L = 2 over {A, C}.
P_X = torch.tensor([0.4, 0.1, 0.2, 0.3])  # p_X over those four states, in the order of STATES.
A_COUNT = torch.tensor([2.0, 1.0, 1.0, 0.0])  # h(x), the number of A tokens in each state.
PEPTIDES = torch.tensor([[1.0, 2.0], [2.0, 3.0], [3.0, 7.0]])  # (hydrophobicity, binding score).
JOINT = torch.tensor([[0.32, 0.08],  # Row x = 0, motif absent: columns are y = 0 and y = 1.
                      [0.24, 0.36]])  # Row x = 1, motif present: Table 1.1 of Section 1.9.1.


WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.


def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
    if not quiet:  # Under --quiet only the closing report is printed.
        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
              f"   gap {abs(by_hand - computed):.1e}")


# %% 02. A distribution over length-two sequences
def two_token_distribution(quiet: bool = False) -> dict:
    """Section 1.3.1: the four masses sum to one, and the event {x_1 = A} has probability 0.5."""
    total = float(P_X.sum())  # Normalization, sum over x in X of p_X(x).
    event = [index for index, state in enumerate(STATES) if state[0] == "A"]  # A = {x : x_1 = A}.
    first_is_a = float(P_X[event].sum())  # P(X in A), the masses of {AA, AC} added together.
    show("sum_x p_X(x)", 1.0, total, quiet)
    show("P(x_1 = A) = p(AA) + p(AC)", 0.5, first_is_a, quiet)
    return {"normalization": total, "p_first_is_a": first_is_a}


# %% 03. Densities for continuous states
def continuous_densities(quiet: bool = False) -> dict:
    """Section 1.4: a density is probability per unit length, so its height may exceed one."""
    uniform_height = 1.0 / 0.1  # The uniform density on [0, 0.1] has the value 1/0.1 there.
    narrow_mass = uniform_height * 0.02  # P(0 <= X <= 0.02) is the height times the width.
    gaussian_peak = 1.0 / (math.sqrt(2.0 * math.pi) * 0.1)  # N(0, 0.1^2) at its center.
    show("uniform[0, 0.1] density height", 10.0, uniform_height, quiet)
    show("P(0 <= X <= 0.02)", 0.2, narrow_mass, quiet)
    show("peak of N(0, 0.1^2)", 3.989, gaussian_peak, quiet)
    return {"uniform_height": uniform_height, "gaussian_peak": gaussian_peak}


# %% 04. The expected adenine count
def expected_adenine(quiet: bool = False) -> dict:
    """Section 1.5.1: E[h(X)] = sum_x p_X(x) h(x), the probability-weighted average of the count."""
    expectation = float(P_X @ A_COUNT)  # 0.4(2) + 0.1(1) + 0.2(1) + 0.3(0).
    across_thousand = 1000.0 * expectation  # The adenines expected in 1000 generated sequences.
    show("E[h(X)], expected A count", 1.1, expectation, quiet)
    show("expected A count in 1000 draws", 1100.0, across_thousand, quiet)
    return {"expectation": expectation}


# %% 05. Monte Carlo estimation
def monte_carlo(samples: int, repeats: int, quiet: bool = False) -> dict:
    """Section 1.5.3: the sample average estimates E[h(X)] with an error that shrinks like 1/sqrt(N).

    The chapter's own ten-sample draw came out at 1.2 against the exact 1.1, a gap well inside the
    standard deviation 0.263 of a ten-sample average. We redraw here, so the estimate moves with
    --seed while the exact value and the error scale stay put. We repeat the ten-sample estimate
    ``repeats`` times. That turns the predicted 0.263 into a measured spread.
    """
    exact = float(P_X @ A_COUNT)  # The target of the estimator, E[h(X)] = 1.1.
    second_moment = float(P_X @ A_COUNT.square())  # E[h^2] = 0.4(4) + 0.1 + 0.2 = 1.9.
    variance = second_moment - exact ** 2  # Var(h(X)) = 1.9 - 1.21 = 0.69.
    show("E[h(X)^2]", 1.9, second_moment, quiet)
    show("Var(h(X))", 0.69, variance, quiet)
    ladder, size = [], 10  # The chapter quotes the error at N = 10 and at N = 1000.
    while size <= max(samples, 10):
        draws = torch.multinomial(P_X, size, replacement=True)  # [size] state indices from p_X.
        estimate = float(A_COUNT[draws].mean())  # (1/N) sum_n h(x^(n)), the Monte Carlo average.
        ladder.append((size, estimate, math.sqrt(variance / size)))
        size *= 10
    if not quiet:
        for size, estimate, error in ladder:
            print(f"  N = {size:<8d} estimate {estimate:.4f}   exact {exact:.4f}"
                  f"   sd of the average {error:.4f}")
    show("sd of a 10-sample average", 0.263, math.sqrt(variance / 10), quiet)
    show("sd of a 1000-sample average", 0.0263, math.sqrt(variance / 1000), quiet)
    rounds = max(repeats, 2)  # Two averages is the fewest we can take a standard deviation over.
    repeated = torch.multinomial(P_X, 10 * rounds, replacement=True)  # [10 * rounds] state indices.
    averages = A_COUNT[repeated].view(rounds, 10).mean(dim=1)  # One ten-sample average per round.
    measured = float(averages.std())  # The spread of those averages, and 0.263 predicts it.
    show(f"sd of {rounds} ten-sample averages", 0.263, measured, quiet)
    return {"exact": exact, "variance": variance, "ladder": ladder, "measured": measured,
            "rounds": rounds}


# %% 06. A worked variance
def worked_variance(quiet: bool = False) -> dict:
    """Section 1.6.2: Var(X) = E[X^2] - (E[X])^2 on the values 0, 1, 2 with masses 0.2, 0.5, 0.3."""
    values = torch.tensor([0.0, 1.0, 2.0])  # The three values X can take.
    masses = torch.tensor([0.2, 0.5, 0.3])  # p_X at those three values.
    first_moment = float(masses @ values)  # E[X] = 0 + 0.5 + 0.6 = 1.1.
    second_moment = float(masses @ values.square())  # E[X^2] = 0 + 0.5 + 1.2 = 1.7.
    from_moments = second_moment - first_moment ** 2  # The boxed identity of Section 1.6.1.
    from_definition = float(masses @ (values - first_moment).square())  # The two-pass route.
    show("E[X]", 1.1, first_moment, quiet)
    show("E[X^2]", 1.7, second_moment, quiet)
    show("Var(X) from the two moments", 0.49, from_moments, quiet)
    show("Var(X) from the definition", 0.49, from_definition, quiet)
    show("sigma_X", 0.7, math.sqrt(from_moments), quiet)
    return {"variance": from_moments, "sigma": math.sqrt(from_moments)}


# %% 07. Covariance, correlation and the covariance matrix
def peptide_covariance(quiet: bool = False) -> dict:
    """Sections 1.7.1 and 1.7.2: three equally likely peptides, their correlation and their Sigma."""
    mu = PEPTIDES.mean(dim=0)  # (mu_X, mu_Y) = (2, 4), each coordinate mean; [2].
    centered = PEPTIDES - mu  # The deviations (-1, -2), (0, -1), (1, 3); [3, 2].
    sigma_matrix = centered.T @ centered / PEPTIDES.shape[0]  # E[(X-mu)(X-mu)^T]; [2, 2].
    cov_xy = float(sigma_matrix[0, 1])  # Cov(X, Y) = (2 + 0 + 3)/3 = 5/3.
    sd_x = math.sqrt(float(sigma_matrix[0, 0]))  # sigma_X = sqrt(2/3).
    sd_y = math.sqrt(float(sigma_matrix[1, 1]))  # sigma_Y = sqrt(14/3).
    rho = cov_xy / (sd_x * sd_y)  # Correlation, and it divides out the units of both variables.
    show("mu_X", 2.0, float(mu[0]), quiet)
    show("mu_Y", 4.0, float(mu[1]), quiet)
    show("Cov(X, Y)", 5.0 / 3.0, cov_xy, quiet)
    show("Var(X) = Sigma_11", 2.0 / 3.0, float(sigma_matrix[0, 0]), quiet)
    show("Var(Y) = Sigma_22", 14.0 / 3.0, float(sigma_matrix[1, 1]), quiet)
    show("sigma_X across the peptides", 0.8165, sd_x, quiet)
    show("sigma_Y across the peptides", 2.160, sd_y, quiet)
    show("rho_XY", 0.9449, rho, quiet)
    show("Sigma symmetry, Sigma_21 - Sigma_12", 0.0, float(sigma_matrix[1, 0] - cov_xy), quiet)
    return {"sigma_matrix": sigma_matrix, "rho": rho}


# %% 08. Gaussian noise and the noising equation
def noising_covariance(quiet: bool = False) -> dict:
    """Section 1.8: Cov(X_t) = alpha_t^2 Cov(X_0) + sigma_t^2 I at alpha_t = 0.6, sigma_t = 0.8."""
    alpha_t, sigma_t = 0.6, 0.8  # A variance-preserving pair, since 0.36 + 0.64 = 1.
    identity = torch.eye(2)  # The isotropic noise covariance; [2, 2].
    isotropic = alpha_t ** 2 * identity + sigma_t ** 2 * identity  # Data covariance I stays I.
    wide_data = torch.diag(torch.tensor([4.0, 1.0]))  # Cov(X_0) = diag(4, 1); [2, 2].
    anisotropic = alpha_t ** 2 * wide_data + sigma_t ** 2 * identity  # diag(2.08, 1.00).
    show("alpha_t^2 + sigma_t^2", 1.0, alpha_t ** 2 + sigma_t ** 2, quiet)
    show("Cov(X_t)_11 with Cov(X_0) = I", 1.0, float(isotropic[0, 0]), quiet)
    show("Cov(X_t)_11 with diag(4, 1)", 2.08, float(anisotropic[0, 0]), quiet)
    show("Cov(X_t)_22 with diag(4, 1)", 1.00, float(anisotropic[1, 1]), quiet)
    return {"anisotropic": anisotropic}


# %% 09. Joint, marginal and conditional distributions
def joint_marginal_conditional(quiet: bool = False) -> dict:
    """Sections 1.9.1 to 1.9.3: the motif-and-binding table, its marginals and its conditionals."""
    total = float(JOINT.sum())  # 0.36 + 0.24 + 0.08 + 0.32 = 1.
    p_x = JOINT.sum(dim=1)  # p(x) = sum_y p(x, y), motif present or absent; [2].
    p_y = JOINT.sum(dim=0)  # p(y) = sum_x p(x, y), high or low binding; [2].
    p_y_given_x = JOINT / p_x[:, None]  # p(y | x) = p(x, y) / p(x), renormalized row by row; [2, 2].
    show("sum over the four cells", 1.0, total, quiet)
    show("p(X = 1), motif present", 0.60, float(p_x[1]), quiet)
    show("p(X = 0), motif absent", 0.40, float(p_x[0]), quiet)
    show("p(Y = 1), high binding", 0.44, float(p_y[1]), quiet)
    show("p(Y = 0), low binding", 0.56, float(p_y[0]), quiet)
    show("p(Y = 1 | X = 1)", 0.6, float(p_y_given_x[1, 1]), quiet)
    show("p(Y = 0 | X = 1)", 0.4, float(p_y_given_x[1, 0]), quiet)
    show("p(Y = 1 | X = 0)", 0.2, float(p_y_given_x[0, 1]), quiet)
    show("p(Y = 0 | X = 0)", 0.8, float(p_y_given_x[0, 0]), quiet)
    return {"p_x": p_x, "p_y": p_y, "p_y_given_x": p_y_given_x}


# %% 10. Bayes' rule and the law of total probability
def bayes(quiet: bool = False) -> dict:
    """Section 1.10: the motif probability after observing high binding, with its own denominator."""
    p_x = JOINT.sum(dim=1)  # The prior p(x) over motif presence; [2].
    p_y_given_x = JOINT / p_x[:, None]  # The likelihood p(y | x); [2, 2].
    evidence = float(p_y_given_x[:, 1] @ p_x)  # p(y) = sum_x p(y | x) p(x), summed over the cause.
    posterior = float(p_y_given_x[1, 1] * p_x[1]) / evidence  # p(X = 1 | Y = 1) = 0.36 / 0.44.
    show("evidence p(Y = 1)", 0.44, evidence, quiet)
    show("posterior p(X = 1 | Y = 1)", 0.818, posterior, quiet)
    show("prior p(X = 1)", 0.60, float(p_x[1]), quiet)
    show("likelihood ratio 0.6 / 0.2", 3.0, float(p_y_given_x[1, 1] / p_y_given_x[0, 1]), quiet)
    return {"evidence": evidence, "posterior": posterior}


# %% 11. Independence
def independence(quiet: bool = False) -> dict:
    """Section 1.11: p(x, y) equals p(x) p(y) under independence, and the table fails that test."""
    p_x = JOINT.sum(dim=1)  # Marginal over motif presence; [2].
    p_y = JOINT.sum(dim=0)  # Marginal over binding activity; [2].
    product = float(p_x[1] * p_y[1])  # 0.60 times 0.44 = 0.264 under a factorized model.
    joint_entry = float(JOINT[1, 1])  # The table's own entry, 0.36.
    show("p(X = 1) p(Y = 1)", 0.264, product, quiet)
    show("p(X = 1, Y = 1)", 0.36, joint_entry, quiet)
    show("association missed by factorizing", 0.096, joint_entry - product, quiet)
    return {"product": product, "joint_entry": joint_entry, "gap": joint_entry - product}


# %% 12. The report
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--samples", type=int, default=1000, help="largest Monte Carlo sample size")
    parser.add_argument("--steps", type=int, default=5000, help="repeats of the ten-sample estimate")
    parser.add_argument("--seed", type=int, default=0, help="random seed for the Monte Carlo draws")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    set_seed(args.seed)  # Seed first, so the Monte Carlo ladder repeats under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.3 to 1.11, worked example by worked example\n")
    two_token_distribution(args.quiet)
    continuous_densities(args.quiet)
    adenine = expected_adenine(args.quiet)
    estimates = monte_carlo(args.samples, args.steps, args.quiet)
    spread = worked_variance(args.quiet)
    peptides = peptide_covariance(args.quiet)
    noising = noising_covariance(args.quiet)
    table = joint_marginal_conditional(args.quiet)
    posterior = bayes(args.quiet)
    association = independence(args.quiet)

    report = {
        "E[h(X)], expected A count": adenine["expectation"],
        "Monte Carlo estimate at N": estimates["ladder"][-1][0],
        "Monte Carlo estimate": estimates["ladder"][-1][1],
        "Var(h(X))": estimates["variance"],
        "ten-sample averages taken": estimates["rounds"],
        "sd of those averages": estimates["measured"],
        "Var(X) on 0, 1, 2": spread["variance"],
        "rho_XY across three peptides": peptides["rho"],
        "Cov(X_t)_11 from diag(4, 1)": float(noising["anisotropic"][0, 0]),
        "p(Y = 1), high binding": float(table["p_y"][1]),
        "p(X = 1 | Y = 1)": posterior["posterior"],
        "p(X=1,Y=1) - p(X=1)p(Y=1)": association["gap"],
    }
    print_report("Probability, Sections 1.3 to 1.11", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work every example of Sections 1.3 to 1.11 and print the report.