"""
CORAL loader -- Cohort B, protocol v5 section 5.3.1.

Implements exactly what the frozen protocol specifies, nothing more:

  * rates = raw count / callable trinucleotide abundance
    (the deposit's pre-normalised columns are NOT used -- they are scaled to
     10,000 and rounded to integers, discarding precision)
  * aggregation to species by SUMMING raw counts and SUMMING callable
    abundances across all branches of the same focal species, then computing
    rates from the sums -- never by averaging normalised spectra
  * exclusion of branches with fewer than 1,000 assigned substitutions
    (the source study's own criterion)
"""

import re
import numpy as np
import pandas as pd

from context3 import from_coral_key

MIN_BRANCH_MUTATIONS = 1000
CAT_RE = re.compile(r"^([ACGT])\[([CT])>([ACGT])\]([ACGT])$")


def load(path, sheet="CORALMutationsDB"):
    return pd.read_excel(path, sheet_name=sheet)


def category_columns(df):
    """The 96 raw count columns, in file order."""
    return [c for c in df.columns if isinstance(c, str) and CAT_RE.match(c.strip())]


def abundance_columns(df):
    """The 32 callable trinucleotide abundance columns (middle base C or T)."""
    out = []
    for c in df.columns:
        s = str(c).strip()
        if len(s) == 3 and s[1] in "CT" and all(b in "ACGT" for b in s):
            out.append(c)
    return out


def context_of(category):
    """'A[C>T]G' -> 'ACG' (the trinucleotide whose abundance is the denominator)."""
    m = CAT_RE.match(category.strip())
    return m.group(1) + m.group(2) + m.group(4)


def aggregate_by_species(df, cats, abunds, species_col="taxa1"):
    """Sum raw counts and callable abundances per focal species.

    Branches below the source study's threshold are dropped first.
    Returns (counts_df, abundance_df, dropped_branches, species_branch_counts).
    """
    totals = df[cats].sum(axis=1)
    keep = totals >= MIN_BRANCH_MUTATIONS
    dropped = int((~keep).sum())
    d = df.loc[keep]
    counts = d.groupby(species_col)[cats].sum()
    abundance = d.groupby(species_col)[abunds].sum()
    nbranch = d.groupby(species_col).size()
    return counts, abundance, dropped, nbranch


def spectra_from_sums(counts, abundance):
    """Per-species rate spectrum: count / callable abundance, then normalised.

    Returns dict {species: {(canonical RNA 3-mer, alt): rate}}, each species'
    rates scaled to sum to 1 so that Phi is a weighted mean and the overall
    mutation-rate level (which this estimand does not use) drops out.
    """
    out = {}
    cats = list(counts.columns)
    denom_col = {c: context_of(c) for c in cats}
    for sp in counts.index:
        row = counts.loc[sp]
        ab = abundance.loc[sp]
        raw = {}
        for c in cats:
            a = ab.get(denom_col[c], 0.0)
            if a and a > 0:
                raw[c] = float(row[c]) / float(a)
        tot = sum(raw.values())
        if tot <= 0:
            continue
        spec = {}
        for c, v in raw.items():
            ctx, alt = from_coral_key(c)
            spec[(ctx, alt)] = v / tot
        out[sp] = spec
    return out


def multinomial_bootstrap(counts_row, abundance_row, cats, n, rng):
    """Bootstrap replicates of one species' spectrum, for sigma^2_{R_s}.

    Resamples the raw category counts multinomially at the observed total,
    holding callable abundances fixed (they are not a sampled quantity).
    """
    p = counts_row.values.astype(float)
    total = int(p.sum())
    p = p / p.sum()
    denom = np.array([float(abundance_row.get(context_of(c), 0.0)) for c in cats])
    keys = [from_coral_key(c) for c in cats]
    reps = []
    for draw in rng.multinomial(total, p, size=n):
        with np.errstate(divide="ignore", invalid="ignore"):
            r = np.where(denom > 0, draw / denom, 0.0)
        s = r.sum()
        if s <= 0:
            continue
        reps.append({k: v for k, v in zip(keys, r / s)})
    return reps
