!pip install biopython
from Bio import SeqIO
import numpy as np
import random
from scipy import stats

# ---------- PARAMETERS ----------
SLP_FILE = "/content/SLP.fasta"
LLP_FILE = "/content/LLP.fasta"
OUTPUT_SLP = "SLP_matched.fasta"
OUTPUT_LLP = "LLP_matched.fasta"
TOLERANCE = 0.10   # ±10% length tolerance
RANDOM_SEED = 42
random.seed(RANDOM_SEED)

# ---------- LOAD FASTA ----------
def load_sequences(fasta_file):
    records = list(SeqIO.parse(fasta_file, "fasta"))
    return records

slp_records = load_sequences(SLP_FILE)
llp_records = load_sequences(LLP_FILE)

# Compute lengths
slp_data = [(rec, len(rec.seq)) for rec in slp_records]
llp_data = [(rec, len(rec.seq)) for rec in llp_records]

# Sort by length
slp_data.sort(key=lambda x: x[1])
llp_data.sort(key=lambda x: x[1])

# ---------- GREEDY MATCHING ----------
matched_slp = []
matched_llp = []

used_slp_indices = set()

for llp_rec, llp_len in llp_data:
    lower = llp_len * (1 - TOLERANCE)
    upper = llp_len * (1 + TOLERANCE)

    candidates = [
        (i, slp_rec) for i, (slp_rec, slp_len) in enumerate(slp_data)
        if lower <= slp_len <= upper and i not in used_slp_indices
    ]

    if candidates:
        idx, slp_rec = random.choice(candidates)
        matched_llp.append(llp_rec)
        matched_slp.append(slp_rec)
        used_slp_indices.add(idx)

# ---------- SAVE MATCHED FASTA ----------
SeqIO.write(matched_slp, OUTPUT_SLP, "fasta")
SeqIO.write(matched_llp, OUTPUT_LLP, "fasta")

print(f"Matched pairs: {len(matched_slp)}")

# ---------- DISTRIBUTION CHECK ----------
slp_lengths = [len(rec.seq) for rec in matched_slp]
llp_lengths = [len(rec.seq) for rec in matched_llp]

print("SLP mean length:", np.mean(slp_lengths))
print("LLP mean length:", np.mean(llp_lengths))

ks_stat, ks_p = stats.ks_2samp(slp_lengths, llp_lengths)
print("KS test p-value (length distribution similarity):", ks_p)