# Supplementary Code 1
# Tree-aware Minimum Description Length (tMDL) calculation
# Implements Fitch's parsimony algorithm on discretized expression states

#REPRODUCIBILITY & USAGE GUIDE (README)
#--------------------------------------
#1. SYSTEM REQUIREMENTS
#   - Operating System: Tested on Windows 10/11
#   - Programming Language: Python version 3.12.8
#   - Dependencies: NumPy, Pandas, SciPy, BioPython
#   - Non-standard hardware: None
#
#2. REQUIRED INPUT FILES (DEMO DATA)
#   - Expression Matrix: A CSV/TSV file formatted similarly to the Human Protein Atlas (HPA) v24
#     transcriptomic data cited in the manuscript, containing at minimum:
#       * 'Gene name'
#       * 'Tissue' or 'Cell type' (leaf labels)
#       * Expression values (e.g., 'nTPM')
#
#3. INSTRUCTIONS FOR USE
#   - Place this script and the expression matrix file in the same working directory.
#   - The script performs the following steps:
#       (i)  Expression preprocessing: values < 1 are set to 0, then log(x + 1) transform is applied.
#       (ii) Expression-based tree construction (as described in Methods): pairwise Spearman correlations
#            between tissues/cell types are computed, converted to a correlation-based distance matrix
#            (distance = 1 − Spearman correlation), and hierarchical clustering (Ward’s method) is applied.
#            The resulting linkage matrix is converted to a rooted Newick tree.
#       (iii) Discretization of expression levels into discrete bins.
#       (iv) Fitch’s parsimony algorithm is applied on the discretized states to compute tMDL per gene.
#
#   - NOTE:
#       Expression data must not contain missing values. If missing values are present,
#       users should preprocess the data accordingly.
#
#4. EXPECTED OUTPUT
#   - The script outputs a results file containing tMDL scores per gene.
#   - A demonstration of the expected output can be verified against Supplementary Tables 1 and 2.
#
# LICENSE: MIT License
#--------------------------------------

import numpy as np
import pandas as pd
from io import StringIO
from Bio import Phylo

from scipy.spatial.distance import squareform
from scipy.cluster import hierarchy
from scipy.cluster.hierarchy import to_tree

# ----------------------------
# USER PARAMETERS
# ----------------------------
INPUT_FILE = "rna_single_cell_type.tsv"  # <-- change to your input file

SEP = "\t"  # "\t" for TSV, "," for CSV

GENE_COL = "Gene name"
LEAF_COL = "Cell type"     
EXPR_COL = "nTPM"

EXPR_CUTOFF = 1.0          # values < 1 set to 0 (as in Methods)
NUM_BINS = 6               # number of discretization bins
OUTPUT_FILE = "tMDL_scores.csv"

# ----------------------------
# 0) Load + preprocess expression (Methods: <1 -> 0, then log(x+1))
# ----------------------------
data_source = pd.read_csv(INPUT_FILE, sep=SEP)

# Ensures there aren't missing values, see NOTE above
if data_source[[GENE_COL, LEAF_COL, EXPR_COL]].isna().any().any():
    raise ValueError(
        "Missing values detected in the input expression data. "
        "Please preprocess to remove/resolve missing values before running this script.")

# Preprocess
data_source.loc[data_source[EXPR_COL] < EXPR_CUTOFF, EXPR_COL] = 0.0
data_source[EXPR_COL] = np.log(data_source[EXPR_COL] + 1.0)
data_source = data_source.drop_duplicates(subset=[GENE_COL, LEAF_COL], keep="first")
# ----------------------------
# 1) Build expression-based tree and convert to Newick
# ----------------------------
# Create genes × leaves matrix for tree construction (continuous preprocessed expression)
expr_mat = data_source.pivot(index=GENE_COL, columns=LEAF_COL, values=EXPR_COL)

if expr_mat.isna().any().any():
    raise ValueError("Expression matrix contains missing values after pivot (missing gene×leaf entries). Please preprocess.")
    
# Spearman correlation across leaves (columns)
correlation_matrix = expr_mat.corr(method="spearman")

# Correlation-based distance
distance_matrix = 1 - correlation_matrix

# Hierarchical clustering (Ward’s method) on the correlation-based distance matrix
condensed_distance_matrix = squareform(distance_matrix.values)
linkage_matrix = hierarchy.linkage(condensed_distance_matrix, method="ward")

# Convert linkage matrix to rooted Newick format (root = final merge)
def linkage_to_newick(node, parent_dist, leaf_names, newick="") -> str:
    """Convert a scipy.cluster.hierarchy.to_tree() node to Newick format."""
    if node.is_leaf():
        return f"'{leaf_names[node.id]}':{parent_dist - node.dist}{newick}"
    else:
        if newick:
            newick = f"):{parent_dist - node.dist}{newick}"
        else:
            newick = ");"
        newick = linkage_to_newick(node.get_left(), node.dist, leaf_names, newick)
        newick = linkage_to_newick(node.get_right(), node.dist, leaf_names, f",{newick}")
        newick = f"({newick}"
        return newick

tree_root = to_tree(linkage_matrix)
leaf_names = list(distance_matrix.columns)  # must match the linkage input order
newick_tree = linkage_to_newick(tree_root, 0, leaf_names)

# Parse the tree using Biopython (used for downstream tMDL calculation)
tree = Phylo.read(StringIO(newick_tree), "newick")

# ----------------------------
# 2) Expression discretization
# ----------------------------
def bin_expression_nonzero_quantiles(values, num_bins):
    """
    Discretize continuous expression values into bins {1..num_bins} using quantiles
    of the global non-zero distribution. Zeros map to the lowest bin.
    """
    x = np.asarray(values, dtype=float)
    out = np.full_like(x, np.nan, dtype=float)

    finite = np.isfinite(x)
    xf = x[finite]
    nz = xf > 0

    edges = np.quantile(xf[nz], np.linspace(0, 1, num_bins + 1)[1:])
    out[finite] = np.clip(np.digitize(xf, edges) + 1, 1, num_bins)
    return out

# Apply discretization on the long-format expression values, then pivot to genes × leaves
data_source["ExpressionBins"] = bin_expression_nonzero_quantiles(data_source[EXPR_COL].values, NUM_BINS)
binned_data = data_source.pivot(index=GENE_COL, columns=LEAF_COL, values="ExpressionBins")

# ----------------------------
# 3) Fitch parsimony (tMDL)
# ----------------------------
def calculate_tmdl(tree, binned_data, num_bins):
    """
    Compute tMDL (Fitch parsimony score) for each gene.

    Parameters
    ----------
    tree : Bio.Phylo tree
        Rooted, strictly-binary tree.
    binned_data : pd.DataFrame
        genes × leaves matrix of discrete states in {1..num_bins}.
    num_bins : int
        Number of bins (states 1..num_bins).

    Returns
    -------
    pd.Series
        tMDL scores indexed by gene.
    """
    all_states = set(range(1, num_bins + 1))
    col_index = {c: i for i, c in enumerate(binned_data.columns)}

    def fitch_score(gene_states):
        score = 0

        def traverse(clade):
            nonlocal score
            if clade.is_terminal():
                v = gene_states[col_index[clade.name]]
                if not np.isfinite(v):
                    return set(all_states)
                return {int(v)}

            left = traverse(clade.clades[0])
            right = traverse(clade.clades[1])

            inter = left & right
            if inter:
                return inter

            score += 1
            return left | right

        traverse(tree.root)
        return score

    scores = [fitch_score(row.values) for _, row in binned_data.iterrows()]
    return pd.Series(scores, index=binned_data.index, name="tMDL")

tmdl_scores = calculate_tmdl(tree, binned_data, num_bins=NUM_BINS)

# Save output
tmdl_scores.to_csv(OUTPUT_FILE, index=True)